From 6d95541179808a7317c4aed9bd9da1f9ea63b844 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Mon, 23 Jun 2025 07:27:05 -0700 Subject: [PATCH 01/17] use TripolarGrid with OceananigansSimulation --- NEWS.md | 3 ++ .../components/ocean/oceananigans.jl | 39 +++++++------------ 2 files changed, 17 insertions(+), 25 deletions(-) diff --git a/NEWS.md b/NEWS.md index 8585a03277..f2173a36e9 100644 --- a/NEWS.md +++ b/NEWS.md @@ -6,6 +6,9 @@ ClimaCoupler.jl Release Notes ### ClimaCoupler features +#### Use TripolarGrid with OceananigansSimulation PR[#1409](https://github.com/CliMA/ClimaCoupler.jl/pull/1409) +Switch from using the Oceananigans.jl `LatitudeLongitudeGrid` to `TripolarGrid`. + #### Use `update_turbulent_fluxes!` instead of `update_field!` for atmosphere PR[#1511](https://github.com/CliMA/ClimaCoupler.jl/pull/1511) Instead of using an `update_field!` method that dispatches on `::Val{:turbulent_fluxes}` to update turbulent fluxes in the atmosphere, we switch to using a function `update_turbulent_fluxes!`. diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index fd44886bb8..3eb92e3599 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -55,33 +55,20 @@ function OceananigansSimulation( download_dataset(en4_temperature) download_dataset(en4_salinity) - # Set up ocean grid (1 degree) - resolution_points = (360, 160, 32) - Nz = last(resolution_points) + # Set up tripolar ocean grid (1 degree) + Nx = 360 + Ny = 180 + Nz = 40 depth = 4000 # meters z = OC.ExponentialDiscretization(Nz, -depth, 0; scale = 0.85 * depth) - # Regular LatLong because we know how to do interpolation there - - # TODO: When moving to TripolarGrid, note that we need to be careful about - # ensuring the coordinate systems align (ie, rotate vectors on the OC grid) - - underlying_grid = OC.LatitudeLongitudeGrid( - arch; - size = resolution_points, - longitude = (-180, 180), - latitude = (-80, 80), # NOTE: Don't goo to high up when using LatLongGrid, or the cells will be too small - z, - halo = (7, 7, 7), - ) - + underlying_grid = OC.TripolarGrid(arch; size = (Nx, Ny, Nz), halo = (7, 7, 4), z) bottom_height = CO.regrid_bathymetry( underlying_grid; minimum_depth = 30, interpolation_passes = 20, major_basins = 1, ) - grid = OC.ImmersedBoundaryGrid( underlying_grid, OC.GridFittedBottom(bottom_height); @@ -113,17 +100,19 @@ function OceananigansSimulation( # Set initial condition to EN4 state estimate at start_date OC.set!(ocean.model, T = en4_temperature[1], S = en4_salinity[1]) + # Construct a remapper from the exchange grid to `Center, Center` fields long_cc = OC.λnodes(grid, OC.Center(), OC.Center(), OC.Center()) lat_cc = OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center()) - # TODO: Go from 0 to Nx+1, Ny+1 (for halos) (for LatLongGrid) + # Create a 2D matrix containing each lat/long combination as a LatLongPoint + # Note this must be done on CPU since the CC.Remapper module is not GPU-compatible + target_points_cc = Array(CC.Geometry.LatLongPoint.(lat_cc, long_cc)) - # Construct a remapper from the exchange grid to `Center, Center` fields - long_cc = reshape(long_cc, length(long_cc), 1) - lat_cc = reshape(lat_cc, 1, length(lat_cc)) - target_points_cc = @. CC.Geometry.LatLongPoint(lat_cc, long_cc) - # TODO: We can remove the `nothing` after CC > 0.14.33 - remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc, nothing) + if pkgversion(CC) >= v"0.14.34" + remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc) + else + remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc, nothing) + end # Construct two 2D Center/Center fields to use as scratch space while remapping scratch_cc1 = OC.Field{OC.Center, OC.Center, Nothing}(grid) From 0dd0acbc316d7da1a1664295e37465d5c855a3e3 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Wed, 10 Sep 2025 09:30:56 -0700 Subject: [PATCH 02/17] use JLD2Writer --- .../components/ocean/oceananigans.jl | 29 ++++++++----------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 3eb92e3599..09eda450e1 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -134,23 +134,18 @@ function OceananigansSimulation( ocean_fresh_water_density = 999.8, ) - # Before version 0.96.22, the NetCDFWriter was broken on GPU - if arch isa OC.CPU || pkgversion(OC) >= v"0.96.22" - # TODO: Add more diagnostics, make them dependent on simulation duration, take - # monthly averages - # Save all tracers and velocities to a NetCDF file at daily frequency - outputs = merge(ocean.model.tracers, ocean.model.velocities) - netcdf_writer = OC.NetCDFWriter( - ocean.model, - outputs; - schedule = OC.TimeInterval(86400), # Daily output - filename = joinpath(output_dir, "ocean_diagnostics.nc"), - indices = (:, :, grid.Nz), - overwrite_existing = true, - array_type = Array{Float32}, - ) - ocean.output_writers[:diagnostics] = netcdf_writer - end + # Save all tracers and velocities to a JLD2 file at daily frequency + outputs = merge(ocean.model.tracers, ocean.model.velocities) + jld2_writer = OC.JLD2Writer( + ocean.model, + outputs; + schedule = OC.TimeInterval(86400), # Daily output + filename = joinpath(output_dir, "ocean_diagnostics"), + indices = (:, :, grid.Nz), + overwrite_existing = true, + array_type = Array{Float32}, + ) + ocean.output_writers[:diagnostics] = jld2_writer sim = OceananigansSimulation(ocean, area_fraction, ocean_properties, remapping) return sim From ab0aa67c9bfbd329e9340748ed8eee2351b81bef Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 30 Sep 2025 10:29:56 -0700 Subject: [PATCH 03/17] use Oceananigans v0.99 --- experiments/ClimaEarth/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/experiments/ClimaEarth/Project.toml b/experiments/ClimaEarth/Project.toml index ff292ea7e6..c9e6137cc0 100644 --- a/experiments/ClimaEarth/Project.toml +++ b/experiments/ClimaEarth/Project.toml @@ -50,6 +50,7 @@ CUDA = "5" EnsembleKalmanProcesses = "2" Insolation = "0.10.2" Interpolations = "0.14, 0.15" +JLD2 = "0.4, 0.5, 0.6" Oceananigans = "0.100" StaticArrays = "1" YAML = "0.4" From e16f3113ce306bcfadfd18bde943f71a0c6bb264 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 30 Sep 2025 15:01:28 -0700 Subject: [PATCH 04/17] use XESMF (WIP) --- experiments/ClimaEarth/components/ocean/oceananigans.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 09eda450e1..11effaa286 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -6,6 +6,7 @@ import ClimaCore as CC import Thermodynamics as TD import ClimaOcean.EN4: download_dataset using KernelAbstractions: @kernel, @index, @inbounds +using XESMF # to load Oceananigans regridding extension """ OceananigansSimulation{SIM, A, OPROP, REMAP} From f5c917d43db5e749183c6d8fd98c1c01012685d6 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Thu, 2 Oct 2025 12:02:32 -0700 Subject: [PATCH 05/17] try out XESMF regridder --- experiments/ClimaEarth/Manifest-v1.11.toml | 55 ++++++++++++ experiments/ClimaEarth/Project.toml | 3 +- .../components/ocean/oceananigans.jl | 87 +++++++++++++++++-- 3 files changed, 138 insertions(+), 7 deletions(-) diff --git a/experiments/ClimaEarth/Manifest-v1.11.toml b/experiments/ClimaEarth/Manifest-v1.11.toml index 07f41c1930..39f9e76ef5 100644 --- a/experiments/ClimaEarth/Manifest-v1.11.toml +++ b/experiments/ClimaEarth/Manifest-v1.11.toml @@ -689,6 +689,12 @@ git-tree-sha1 = "d9d26935a0bcffc87d2613ce14c527c99fc543fd" uuid = "f0e56b4a-5159-44fe-b623-3e5288b988bb" version = "2.5.0" +[[deps.CondaPkg]] +deps = ["JSON3", "Markdown", "MicroMamba", "Pidfile", "Pkg", "Preferences", "Scratch", "TOML", "pixi_jll"] +git-tree-sha1 = "bd491d55b97a036caae1d78729bdb70bf7dababc" +uuid = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" +version = "0.2.33" + [[deps.ConstructionBase]] git-tree-sha1 = "b4b092499347b18a015186eae3042f72267106cb" uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" @@ -2088,6 +2094,12 @@ version = "1.5.0" LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d" SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622" +[[deps.MicroMamba]] +deps = ["Pkg", "Scratch", "micromamba_jll"] +git-tree-sha1 = "011cab361eae7bcd7d278f0a7a00ff9c69000c51" +uuid = "0b3b1443-0f03-428d-bdfb-f27f9c1191ea" +version = "0.1.14" + [[deps.MicrosoftMPI_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] git-tree-sha1 = "bc95bf4149bf535c09602e3acdf950d9b4376227" @@ -2413,6 +2425,12 @@ git-tree-sha1 = "7d2f8f21da5db6a806faf7b9b292296da42b2810" uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" version = "2.8.3" +[[deps.Pidfile]] +deps = ["FileWatching", "Test"] +git-tree-sha1 = "2d8aaf8ee10df53d0dfb9b8ee44ae7c04ced2b03" +uuid = "fa939f87-e72e-5be4-a000-7fc836dbe307" +version = "1.3.0" + [[deps.Pixman_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "LLVMOpenMP_jll", "Libdl"] git-tree-sha1 = "db76b1ecd5e9715f3d043cec13b2ec93ce015d53" @@ -2535,6 +2553,20 @@ git-tree-sha1 = "1d36ef11a9aaf1e8b74dacc6a731dd1de8fd493d" uuid = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d" version = "1.3.0" +[[deps.PythonCall]] +deps = ["CondaPkg", "Dates", "Libdl", "MacroTools", "Markdown", "Pkg", "Serialization", "Tables", "UnsafePointers"] +git-tree-sha1 = "34510e11cabd7964291f32f14d28b367e9960e6e" +uuid = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +version = "0.9.28" + + [deps.PythonCall.extensions] + CategoricalArraysExt = "CategoricalArrays" + PyCallExt = "PyCall" + + [deps.PythonCall.weakdeps] + CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" + PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" + [[deps.QOI]] deps = ["ColorTypes", "FileIO", "FixedPointNumbers"] git-tree-sha1 = "8b3fc30bc0390abdce15f8822c889f669baed73d" @@ -3316,6 +3348,11 @@ weakdeps = ["LLVM"] [deps.UnsafeAtomics.extensions] UnsafeAtomicsLLVM = ["LLVM"] +[[deps.UnsafePointers]] +git-tree-sha1 = "c81331b3b2e60a982be57c046ec91f599ede674a" +uuid = "e17b2a0c-0bdf-430a-bd0c-3a23cae4ff39" +version = "1.0.0" + [[deps.VectorInterface]] deps = ["LinearAlgebra"] git-tree-sha1 = "cea8abaa6e43f72f97a09cf95b80c9eb53ff75cf" @@ -3340,6 +3377,12 @@ git-tree-sha1 = "c1a7aa6219628fcd757dede0ca95e245c5cd9511" uuid = "efce3f68-66dc-5838-9240-27a6d6f5f9b6" version = "1.0.0" +[[deps.XESMF]] +deps = ["CondaPkg", "LinearAlgebra", "PythonCall", "SparseArrays"] +git-tree-sha1 = "f7c53612764f77438c23743d20c926169cb57b45" +uuid = "2e0b0046-e7a1-486f-88de-807ee8ffabe5" +version = "0.1.5" + [[deps.XML2_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"] git-tree-sha1 = "80d3930c6347cfce7ccf96bd3bafdf079d9c0390" @@ -3494,6 +3537,12 @@ git-tree-sha1 = "86addc139bca85fdf9e7741e10977c45785727b7" uuid = "337d8026-41b4-5cde-a456-74a10e5b31d1" version = "1.11.3+0" +[[deps.micromamba_jll]] +deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl"] +git-tree-sha1 = "2ca2ac0b23a8e6b76752453e08428b3b4de28095" +uuid = "f8abcde7-e9b7-5caa-b8af-a437887ae8e4" +version = "1.5.12+0" + [[deps.nghttp2_jll]] deps = ["Artifacts", "Libdl"] uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d" @@ -3510,6 +3559,12 @@ deps = ["Artifacts", "Libdl"] uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0" version = "17.4.0+2" +[[deps.pixi_jll]] +deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl"] +git-tree-sha1 = "f349584316617063160a947a82638f7611a8ef0f" +uuid = "4d7b5844-a134-5dcd-ac86-c8f19cd51bed" +version = "0.41.3+0" + [[deps.x264_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] git-tree-sha1 = "14cc7083fc6dff3cc44f2bc435ee96d06ed79aa7" diff --git a/experiments/ClimaEarth/Project.toml b/experiments/ClimaEarth/Project.toml index c9e6137cc0..2c8fa28891 100644 --- a/experiments/ClimaEarth/Project.toml +++ b/experiments/ClimaEarth/Project.toml @@ -32,10 +32,12 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f" Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" +XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5" YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6" [compat] ArgParse = "1.1" +CUDA = "5" ClimaAnalysis = "0.5.10" ClimaAtmos = "0.27, 0.28, 0.29, 0.30, 0.31" ClimaCalibrate = "0.1" @@ -46,7 +48,6 @@ ClimaParams = "1.0" ClimaSeaIce = "0.3" ClimaTimeSteppers = "0.7, 0.8" ClimaUtilities = "0.1" -CUDA = "5" EnsembleKalmanProcesses = "2" Insolation = "0.10.2" Interpolations = "0.14, 0.15" diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 11effaa286..fcbb1ebd53 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -1,5 +1,6 @@ import Oceananigans as OC import ClimaOcean as CO +import ClimaAtmos as CA import ClimaCoupler: Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Utilities import ClimaComms import ClimaCore as CC @@ -46,6 +47,9 @@ function OceananigansSimulation( output_dir, comms_ctx = ClimaComms.context(), ) + # using Dates + # start_date = Dates.DateTime(2008) + # stop_date = Dates.DateTime(2008, 1, 2) arch = comms_ctx.device isa ClimaComms.CUDADevice ? OC.GPU() : OC.CPU() # Retrieve EN4 data (monthly) @@ -102,13 +106,84 @@ function OceananigansSimulation( OC.set!(ocean.model, T = en4_temperature[1], S = en4_salinity[1]) # Construct a remapper from the exchange grid to `Center, Center` fields - long_cc = OC.λnodes(grid, OC.Center(), OC.Center(), OC.Center()) - lat_cc = OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center()) + + # Tripolar to cubed sphere (fails currently) + # long_oc = OC.λnodes(grid, OC.Center(), OC.Center(), OC.Center()) + # lat_oc = OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center()) + + T = OC.CenterField(grid) # field on tripolar grid + OC.set!(T, en4_temperature[1]) + coords_tripolar, _ = XESMF.extract_xesmf_coordinates_structure(T, T) # Create a 2D matrix containing each lat/long combination as a LatLongPoint # Note this must be done on CPU since the CC.Remapper module is not GPU-compatible - target_points_cc = Array(CC.Geometry.LatLongPoint.(lat_cc, long_cc)) + # target_points_oc = Array(CC.Geometry.LatLongPoint.(lat_oc, long_oc)) + + # TODO double check these coords are on centers, get lat/lon bounds for each element + # boundary_space = axes(area_fraction) + boundary_space = CC.CommonSpaces.CubedSphereSpace( + Float32; + radius = 1.0, + n_quad_points = 2, + h_elem = 10, + ) + boundary_coords = CC.Fields.coordinate_field(boundary_space) + boundary_lat_arr = CC.Fields.field2array(boundary_coords.lat) + boundary_long_arr = CC.Fields.field2array(boundary_coords.long) + coords_cubedsphere = Dict("lat" => boundary_lat_arr, "lon" => boundary_long_arr) + + regridder_tripolar_to_cubedsphere = + XESMF.Regridder(coords_cubedsphere, coords_tripolar; method = "conservative") + regridder_cubedsphere_to_tripolar = + XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "conservative") + + # Tripolar to LatLon + T = OC.CenterField(grid) # field on tripolar grid + OC.set!(T, en4_temperature[1]) + + grid_latlon = OC.LatitudeLongitudeGrid( + arch; + size = (Nx, Ny, Nz), + longitude = (0, 360), + latitude = (-81, 90), + z, + ) + field_latlon = OC.CenterField(grid_latlon) + remapper_tripolar_to_latlon = XESMF.Regridder(T, field_latlon; method = "conservative") + remapper_tripolar_to_latlon( + vec(OC.interior(field_latlon, :, :, Nz)), + vec(OC.interior(T, :, :, Nz)), + ) + OC.fill_halo_regions!(field_latlon) + # heatmap(field_latlon) # using GeoMakie, using CairoMakie + + # LatLon to cubed sphere + boundary_space = CC.CommonSpaces.CubedSphereSpace( + Float32; + radius = 1.0, + n_quad_points = 2, + h_elem = 10, + ) + field_cubedsphere = Interfacer.remap(field_latlon, boundary_space) + # fieldheatmap(field_cubedsphere) # using ClimaCoreMakie + + # Cubed sphere to LatLon + long_oc = OC.λnodes(grid_latlon, OC.Center(), OC.Center(), OC.Center()) + lat_oc = OC.φnodes(grid_latlon, OC.Center(), OC.Center(), OC.Center()) + long_oc = reshape(long_oc, length(long_oc), 1) + lat_oc = reshape(lat_oc, 1, length(lat_oc)) + target_points_oc = @. CC.Geometry.LatLongPoint(lat_oc, long_oc) + + remapper_cubedsphere_to_latlon = CC.Remapping.Remapper(boundary_space, target_points_oc) + field_cubedsphere = CC.Fields.zeros(boundary_space) + CC.Remapping.interpolate!( + field_cubedsphere, + remapper_cubedsphere_to_latlon, + field_latlon, + ) + + # Previous remapper for LatLon if pkgversion(CC) >= v"0.14.34" remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc) else @@ -194,14 +269,14 @@ Interfacer.step!(sim::OceananigansSimulation, t) = # We always want the surface, so we always set zero(pt.lat) for z """ - to_node(pt::CA.ClimaCore.Geometry.LatLongPoint) + to_node(pt::CCGeometry.LatLongPoint) Transform `LatLongPoint` into a tuple (long, lat, 0), where the 0 is needed because we only care about the surface. """ -@inline to_node(pt::CA.ClimaCore.Geometry.LatLongPoint) = pt.long, pt.lat, zero(pt.lat) +@inline to_node(pt::CC.Geometry.LatLongPoint) = pt.long, pt.lat, zero(pt.lat) # This next one is needed if we have "LevelGrid" -@inline to_node(pt::CA.ClimaCore.Geometry.LatLongZPoint) = pt.long, pt.lat, zero(pt.lat) +@inline to_node(pt::CC.Geometry.LatLongZPoint) = pt.long, pt.lat, zero(pt.lat) """ map_interpolate(points, oc_field::OC.Field) From a8d783378e7574078e502e9a3bb550da98e15d3a Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Fri, 31 Oct 2025 15:48:50 -0700 Subject: [PATCH 06/17] use XESMF bilinear [skip ci] --- .../components/ocean/oceananigans.jl | 45 ++++++++++++------- 1 file changed, 30 insertions(+), 15 deletions(-) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index fcbb1ebd53..4dd75e26af 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -9,6 +9,12 @@ import ClimaOcean.EN4: download_dataset using KernelAbstractions: @kernel, @index, @inbounds using XESMF # to load Oceananigans regridding extension +OceananigansXESMFExt = + Base.get_extension( + OC, + :OceananigansXESMFExt, + ).OceananigansXESMFExt; + """ OceananigansSimulation{SIM, A, OPROP, REMAP} @@ -113,29 +119,38 @@ function OceananigansSimulation( T = OC.CenterField(grid) # field on tripolar grid OC.set!(T, en4_temperature[1]) - coords_tripolar, _ = XESMF.extract_xesmf_coordinates_structure(T, T) + coords_tripolar = OceananigansXESMFExt.xesmf_coordinates(T) + # Put everything on CPU + coords_tripolar = Dict(k => Array(v) for (k, v) in coords_tripolar) # Create a 2D matrix containing each lat/long combination as a LatLongPoint # Note this must be done on CPU since the CC.Remapper module is not GPU-compatible # target_points_oc = Array(CC.Geometry.LatLongPoint.(lat_oc, long_oc)) - # TODO double check these coords are on centers, get lat/lon bounds for each element - # boundary_space = axes(area_fraction) - boundary_space = CC.CommonSpaces.CubedSphereSpace( - Float32; - radius = 1.0, - n_quad_points = 2, - h_elem = 10, - ) - boundary_coords = CC.Fields.coordinate_field(boundary_space) - boundary_lat_arr = CC.Fields.field2array(boundary_coords.lat) - boundary_long_arr = CC.Fields.field2array(boundary_coords.long) - coords_cubedsphere = Dict("lat" => boundary_lat_arr, "lon" => boundary_long_arr) + # Get the latitude and longitude of each node on the boundary space + boundary_space = axes(area_fraction) + # CC.CommonSpaces.CubedSphereSpace( + # Float32; + # radius = 1.0, + # n_quad_points = 2, + # h_elem = 10, + # ) + cubedsphere_coords = CC.Fields.coordinate_field(boundary_space) + # Put everything on CPU + cubedsphere_lat_arr = Array(CC.Fields.field2array(cubedsphere_coords.lat)) + cubedsphere_long_arr = Array(CC.Fields.field2array(cubedsphere_coords.long)) + + # lon2d varies along cols (dim 2) — x direction + cubedsphere_lon2d = repeat(cubedsphere_long_arr', length(cubedsphere_lat_arr), 1) # transpose lon to make it row vector + # lat2d varies along rows (dim 1) — y direction + cubedsphere_lat2d = repeat(cubedsphere_lat_arr, 1, length(cubedsphere_long_arr)) # column vector repeated across + + coords_cubedsphere = Dict("lat" => cubedsphere_lat2d, "lon" => cubedsphere_lon2d) regridder_tripolar_to_cubedsphere = - XESMF.Regridder(coords_cubedsphere, coords_tripolar; method = "conservative") + XESMF.Regridder(coords_cubedsphere, coords_tripolar; method = "bilinear") regridder_cubedsphere_to_tripolar = - XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "conservative") + XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "bilinear") # Tripolar to LatLon T = OC.CenterField(grid) # field on tripolar grid From d965eeed02f91401dcbb09917cef0a37dfd68c9c Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Mon, 3 Nov 2025 16:20:01 -0800 Subject: [PATCH 07/17] xesmf works tp->cs --- .../components/ocean/oceananigans.jl | 100 ++++++++++++++---- 1 file changed, 81 insertions(+), 19 deletions(-) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 4dd75e26af..4a101d473d 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -9,11 +9,7 @@ import ClimaOcean.EN4: download_dataset using KernelAbstractions: @kernel, @index, @inbounds using XESMF # to load Oceananigans regridding extension -OceananigansXESMFExt = - Base.get_extension( - OC, - :OceananigansXESMFExt, - ).OceananigansXESMFExt; +OceananigansXESMFExt = Base.get_extension(OC, :OceananigansXESMFExt).OceananigansXESMFExt; """ OceananigansSimulation{SIM, A, OPROP, REMAP} @@ -53,9 +49,6 @@ function OceananigansSimulation( output_dir, comms_ctx = ClimaComms.context(), ) - # using Dates - # start_date = Dates.DateTime(2008) - # stop_date = Dates.DateTime(2008, 1, 2) arch = comms_ctx.device isa ClimaComms.CUDADevice ? OC.GPU() : OC.CPU() # Retrieve EN4 data (monthly) @@ -111,6 +104,11 @@ function OceananigansSimulation( # Set initial condition to EN4 state estimate at start_date OC.set!(ocean.model, T = en4_temperature[1], S = en4_salinity[1]) + + # Set up the XESMF regridder from tripolar to cubedsphere + ## TODO keep using ClimaCore remapper for cubedsphere to tripolar because it uses + # all the information we have about the grid + # Construct a remapper from the exchange grid to `Center, Center` fields # Tripolar to cubed sphere (fails currently) @@ -123,6 +121,7 @@ function OceananigansSimulation( # Put everything on CPU coords_tripolar = Dict(k => Array(v) for (k, v) in coords_tripolar) + # Create a 2D matrix containing each lat/long combination as a LatLongPoint # Note this must be done on CPU since the CC.Remapper module is not GPU-compatible # target_points_oc = Array(CC.Geometry.LatLongPoint.(lat_oc, long_oc)) @@ -136,21 +135,84 @@ function OceananigansSimulation( # h_elem = 10, # ) cubedsphere_coords = CC.Fields.coordinate_field(boundary_space) - # Put everything on CPU - cubedsphere_lat_arr = Array(CC.Fields.field2array(cubedsphere_coords.lat)) - cubedsphere_long_arr = Array(CC.Fields.field2array(cubedsphere_coords.long)) - # lon2d varies along cols (dim 2) — x direction - cubedsphere_lon2d = repeat(cubedsphere_long_arr', length(cubedsphere_lat_arr), 1) # transpose lon to make it row vector - # lat2d varies along rows (dim 1) — y direction - cubedsphere_lat2d = repeat(cubedsphere_lat_arr, 1, length(cubedsphere_long_arr)) # column vector repeated across + # (Ni, Nj, _, Nh) = size(parent(cubedsphere_coords.lat)) + # cubedsphere_lat2d = Array(reshape(parent(cubedsphere_coords.lat), Ni*Nj, Nh)) + # cubedsphere_lon2d = Array(reshape(parent(cubedsphere_coords.long), Ni*Nj, Nh)) + + flat_lat = Array(reshape(vec(parent(cubedsphere_coords.lat)), :, 1)) + flat_lon = Array(reshape(vec(parent(cubedsphere_coords.long)), :, 1)) + # squareish_lat_dim_i = floor(Int, sqrt(length(flat_lat))) + # squareish_lat_dim_j = cld(length(flat_lat), squareish_lat_dim_i) + # lat_arr = similar(flat_lat, squareish_lat_dim_i, squareish_lat_dim_j) + + # for i in 1:length(flat_lat) + # lat_arr[i] = flat_lat[i] + # end + # for i in (length(flat_lat) + 1):length(lat_arr) + # lat_arr[i] = 0 + # end + + # flat_lon = vec(parent(cubedsphere_coords.long)) + # squareish_lon_dim_i = floor(Int, sqrt(length(flat_lon))) + # squareish_lon_dim_j = cld(length(flat_lon), squareish_lon_dim_i) + # lon_arr = similar(flat_lon, squareish_lon_dim_i, squareish_lon_dim_j) + + # for i in 1:length(flat_lon) + # lon_arr[i] = flat_lon[i] + # end + # for i in (length(flat_lon) + 1):length(lon_arr) + # lon_arr[i] = 0 + # end + + + # # Put everything on CPU + # cubedsphere_lat_arr = Array(CC.Fields.field2array(cubedsphere_coords.lat)) + # cubedsphere_long_arr = Array(CC.Fields.field2array(cubedsphere_coords.long)) + + # # lon2d varies along cols (dim 2) — x direction + # cubedsphere_lon2d = collect(repeat(cubedsphere_long_arr', length(cubedsphere_lat_arr), 1)) # transpose lon to make it row vector + # # lat2d varies along rows (dim 1) — y direction + # cubedsphere_lat2d = collect(repeat(cubedsphere_lat_arr, 1, length(cubedsphere_long_arr))) # column vector repeated across + + coords_cubedsphere = Dict("lat" => flat_lat, "lon" => flat_lon) + + GC.gc() + + # ## Contents of XESMF.Regridder constructor + # method = "bilinear" + # regridder_tripolar_to_cubedsphere = XESMF.xesmf.Regridder(coords_cubedsphere, coords_tripolar, method; periodic = false) + # method = uppercasefirst(string(method)) + + # weights = XESMF.sparse_regridder_weights(regridder) + + # Ndst, Nsrc = size(weights) + + # temp_src = zeros(Nsrc) + # temp_dst = zeros(Ndst) + + # return Regridder(method, weights, temp_src, temp_dst) - coords_cubedsphere = Dict("lat" => cubedsphere_lat2d, "lon" => cubedsphere_lon2d) regridder_tripolar_to_cubedsphere = - XESMF.Regridder(coords_cubedsphere, coords_tripolar; method = "bilinear") - regridder_cubedsphere_to_tripolar = - XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "bilinear") + XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "bilinear") # src, dest, method + + src_vec = vec(OC.interior(T, :, :, Nz)) # tripolar source + dst_vec = vec(parent(CC.Fields.zeros(boundary_space))) # cubedsphere dest + + regridder_tripolar_to_cubedsphere(dst_vec, src_vec) # REGRID MATRIX MULT + + # LinearAlgebra.mul!(dst_vec, regridder_tripolar_to_cubedsphere.weights, src_vec) + + dst_arr = reshape(dst_vec, size(parent(cubedsphere_coords.lat))...) + dst_field = CC.Fields.array2field(dst_arr, boundary_space) + + using ClimaCoreMakie + fieldheatmap(dst_field) + + + # regridder_cubedsphere_to_tripolar = + # XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "bilinear") # Tripolar to LatLon T = OC.CenterField(grid) # field on tripolar grid From 3464f2fb0321634166b635cb3e8b5f61f9d4db7a Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 4 Nov 2025 13:15:10 -0800 Subject: [PATCH 08/17] add function to construct remappers --- .../components/ocean/oceananigans.jl | 283 ++++++++---------- 1 file changed, 117 insertions(+), 166 deletions(-) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 4a101d473d..482d7a3460 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -22,13 +22,16 @@ It contains the following objects: - `ocean::SIM`: The Oceananigans simulation object. - `area_fraction::A`: A ClimaCore Field representing the surface area fraction of this component model on the exchange grid. - `ocean_properties::OPROP`: A NamedTuple of ocean properties and parameters -- `remapping::REMAP`: Objects needed to remap from the exchange (spectral) grid to Oceananigans spaces. +- `remapper_oc_to_cs::REMAP1`: Objects needed to remap from the Oceananigans space to the exchange (spectral) grid. +- `remapper_cs_to_oc::REMAP2`: Objects needed to remap from the exchange (spectral) grid to the Oceananigans space. """ -struct OceananigansSimulation{SIM, A, OPROP, REMAP} <: Interfacer.OceanModelSimulation +struct OceananigansSimulation{SIM, A, OPROP, REMAP1, REMAP2} <: + Interfacer.OceanModelSimulation ocean::SIM area_fraction::A ocean_properties::OPROP - remapping::REMAP + remapper_oc_to_cs::REMAP1 + remapper_cs_to_oc::REMAP2 end """ @@ -104,204 +107,152 @@ function OceananigansSimulation( # Set initial condition to EN4 state estimate at start_date OC.set!(ocean.model, T = en4_temperature[1], S = en4_salinity[1]) - - # Set up the XESMF regridder from tripolar to cubedsphere - ## TODO keep using ClimaCore remapper for cubedsphere to tripolar because it uses - # all the information we have about the grid - - # Construct a remapper from the exchange grid to `Center, Center` fields - - # Tripolar to cubed sphere (fails currently) - # long_oc = OC.λnodes(grid, OC.Center(), OC.Center(), OC.Center()) - # lat_oc = OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center()) - - T = OC.CenterField(grid) # field on tripolar grid - OC.set!(T, en4_temperature[1]) - coords_tripolar = OceananigansXESMFExt.xesmf_coordinates(T) - - # Put everything on CPU - coords_tripolar = Dict(k => Array(v) for (k, v) in coords_tripolar) - - # Create a 2D matrix containing each lat/long combination as a LatLongPoint - # Note this must be done on CPU since the CC.Remapper module is not GPU-compatible - # target_points_oc = Array(CC.Geometry.LatLongPoint.(lat_oc, long_oc)) - - # Get the latitude and longitude of each node on the boundary space + # Get the remapper objects to go in both directions between the Oceananigans and Cubed sphere grids boundary_space = axes(area_fraction) - # CC.CommonSpaces.CubedSphereSpace( - # Float32; - # radius = 1.0, - # n_quad_points = 2, - # h_elem = 10, - # ) - cubedsphere_coords = CC.Fields.coordinate_field(boundary_space) + remapper_oc_to_cs, remapper_cs_to_oc = construct_remappers(grid, boundary_space) - # (Ni, Nj, _, Nh) = size(parent(cubedsphere_coords.lat)) - # cubedsphere_lat2d = Array(reshape(parent(cubedsphere_coords.lat), Ni*Nj, Nh)) - # cubedsphere_lon2d = Array(reshape(parent(cubedsphere_coords.long), Ni*Nj, Nh)) - - flat_lat = Array(reshape(vec(parent(cubedsphere_coords.lat)), :, 1)) - flat_lon = Array(reshape(vec(parent(cubedsphere_coords.long)), :, 1)) - # squareish_lat_dim_i = floor(Int, sqrt(length(flat_lat))) - # squareish_lat_dim_j = cld(length(flat_lat), squareish_lat_dim_i) - # lat_arr = similar(flat_lat, squareish_lat_dim_i, squareish_lat_dim_j) - - # for i in 1:length(flat_lat) - # lat_arr[i] = flat_lat[i] - # end - # for i in (length(flat_lat) + 1):length(lat_arr) - # lat_arr[i] = 0 - # end - - # flat_lon = vec(parent(cubedsphere_coords.long)) - # squareish_lon_dim_i = floor(Int, sqrt(length(flat_lon))) - # squareish_lon_dim_j = cld(length(flat_lon), squareish_lon_dim_i) - # lon_arr = similar(flat_lon, squareish_lon_dim_i, squareish_lon_dim_j) - - # for i in 1:length(flat_lon) - # lon_arr[i] = flat_lon[i] - # end - # for i in (length(flat_lon) + 1):length(lon_arr) - # lon_arr[i] = 0 - # end + ocean_properties = (; + ocean_reference_density = 1020, + ocean_heat_capacity = 3991, + ocean_fresh_water_density = 999.8, + ) + # Save all tracers and velocities to a JLD2 file at daily frequency + outputs = merge(ocean.model.tracers, ocean.model.velocities) + jld2_writer = OC.JLD2Writer( + ocean.model, + outputs; + schedule = OC.TimeInterval(86400), # Daily output + filename = joinpath(output_dir, "ocean_diagnostics"), + indices = (:, :, grid.Nz), + overwrite_existing = true, + array_type = Array{Float32}, + ) + ocean.output_writers[:diagnostics] = jld2_writer - # # Put everything on CPU - # cubedsphere_lat_arr = Array(CC.Fields.field2array(cubedsphere_coords.lat)) - # cubedsphere_long_arr = Array(CC.Fields.field2array(cubedsphere_coords.long)) + sim = OceananigansSimulation( + ocean, + area_fraction, + ocean_properties, + remapper_oc_to_cs, + remapper_cs_to_oc, + ) + return sim +end - # # lon2d varies along cols (dim 2) — x direction - # cubedsphere_lon2d = collect(repeat(cubedsphere_long_arr', length(cubedsphere_lat_arr), 1)) # transpose lon to make it row vector - # # lat2d varies along rows (dim 1) — y direction - # cubedsphere_lat2d = collect(repeat(cubedsphere_lat_arr, 1, length(cubedsphere_long_arr))) # column vector repeated across +""" + _construct_remappers(grid, boundary_space) - coords_cubedsphere = Dict("lat" => flat_lat, "lon" => flat_lon) +Construct the remapper objects to go in both directions between the Oceananigans and Cubed sphere grids. +Both objects contain a remapper object and relevant scratch space. - GC.gc() +- Oceananigans to ClimaCore +In this direction we use XESMF bilinear interpolation. - # ## Contents of XESMF.Regridder constructor - # method = "bilinear" - # regridder_tripolar_to_cubedsphere = XESMF.xesmf.Regridder(coords_cubedsphere, coords_tripolar, method; periodic = false) - # method = uppercasefirst(string(method)) + Example: remap the Oceananigans field `T` to the ClimaCore field `T_cubedsphere` - # weights = XESMF.sparse_regridder_weights(regridder) + # Convert the Oceananigans field to a flat vector + remapper_oc_to_cs.src_vec .= Array(vec(OC.interior(T, :, :, Nz))) # Oceananigans source vector - # Ndst, Nsrc = size(weights) + # Apply the XESMF regridder (matrix multiply) + remapper_oc_to_cs.remapper(remapper_oc_to_cs.dest_vec, remapper_oc_to_cs.src_vec) - # temp_src = zeros(Nsrc) - # temp_dst = zeros(Ndst) + # Convert the output vector to a 2D array and then to a ClimaCore field + remapper_oc_to_cs.dest_arr .= reshape(remapper_oc_to_cs.dest_vec, size(remapper_oc_to_cs.dest_arr)...) - # return Regridder(method, weights, temp_src, temp_dst) + # Copy the remapped data to the ClimaCore field + # Note: in general we avoid accessing the parent of a field, but here we make an exception + # since we're using the underlying array to remap. + parent(T_cubedsphere) .= remapper_oc_to_cs.dest_arr +- ClimaCore to Oceananigans +In this direction we use the ClimaCore remapper for this because it uses all the information +we have about the spectral element cubed sphere grid, which XESMF does not support. - regridder_tripolar_to_cubedsphere = - XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "bilinear") # src, dest, method + Example: remap the ClimaCore field `F_turb_ρτxz_cs` to the Oceananigans field `F_turb_ρτxz_oc`: - src_vec = vec(OC.interior(T, :, :, Nz)) # tripolar source - dst_vec = vec(parent(CC.Fields.zeros(boundary_space))) # cubedsphere dest + # Remap the ClimaCore momentum flux to a Oceananigans field using scratch arrays and fields + CC.Remapping.interpolate!( + remapper_cs_to_oc.scratch_arr1, + remapper_cs_to_oc.remapper, + F_turb_ρτxz_cs, # ClimaCore field + ) + OC.set!(remapper_cs_to_oc.scratch_oc1, remapper_cs_to_oc.scratch_arr1) # zonal momentum flux + F_turb_ρτxz_oc = remapper_cs_to_oc.scratch_oc1 # Oceananigans field - regridder_tripolar_to_cubedsphere(dst_vec, src_vec) # REGRID MATRIX MULT +Arguments: +- `grid`: The Oceananigans grid (TripolarGrid or LatitudeLongitudeGrid). +- `boundary_space`: The boundary space (ClimaCore SpectralElementSpace2D). - # LinearAlgebra.mul!(dst_vec, regridder_tripolar_to_cubedsphere.weights, src_vec) +Returns: +- `remapper_oc_to_cs`: The remapper object to go from the Oceananigans grid to the Cubed sphere nodes. +""" +function _construct_remappers(grid, boundary_space) + ## Remapper: Oceananigans `Center, Center` to Cubed sphere nodes + # Get the Oceananigans coordinates and put them on CPU + coords_oc = OceananigansXESMFExt.xesmf_coordinates(grid) + coords_oc = Dict(k => Array(v) for (k, v) in coords_oc) - dst_arr = reshape(dst_vec, size(parent(cubedsphere_coords.lat))...) - dst_field = CC.Fields.array2field(dst_arr, boundary_space) + # Get the latitude and longitude of each node on the boundary space + cubedsphere_coords = CC.Fields.coordinate_field(boundary_space) - using ClimaCoreMakie - fieldheatmap(dst_field) + # Get the cubed sphere latitude and longitude, each as an Nx1 Matrix + cubedsphere_lat = Array(reshape(vec(parent(cubedsphere_coords.lat)), :, 1)) + cubedsphere_lon = Array(reshape(vec(parent(cubedsphere_coords.long)), :, 1)) + coords_cubedsphere = Dict("lat" => cubedsphere_lat, "lon" => cubedsphere_lon) - # regridder_cubedsphere_to_tripolar = - # XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "bilinear") + # Construct the XESMF regridder object + regridder_oceananigans_to_cubedsphere = + XESMF.Regridder(coords_oc, coords_cubedsphere; method = "bilinear") - # Tripolar to LatLon - T = OC.CenterField(grid) # field on tripolar grid - OC.set!(T, en4_temperature[1]) + # Allocate space for source an destination vectors to use as intermediate storage + src_vec_oc = Array(vec(OC.Field{OC.Center, OC.Center, Nothing}(grid))) # 2D field on Center/Center + field_cubedsphere = CC.Fields.zeros(boundary_space) # 2D field on boundary space (cubed sphere) + dest_vec_cubedsphere = vec(parent(field_cubedsphere)) + dest_arr_cubedsphere = + deepcopy(reshape(dest_vec_cubedsphere, size(parent(field_cubedsphere))...)) - grid_latlon = OC.LatitudeLongitudeGrid( - arch; - size = (Nx, Ny, Nz), - longitude = (0, 360), - latitude = (-81, 90), - z, - ) - field_latlon = OC.CenterField(grid_latlon) - remapper_tripolar_to_latlon = XESMF.Regridder(T, field_latlon; method = "conservative") - remapper_tripolar_to_latlon( - vec(OC.interior(field_latlon, :, :, Nz)), - vec(OC.interior(T, :, :, Nz)), + remapper_oc_to_cs = (; + remapper = regridder_oceananigans_to_cubedsphere, + src_vec = src_vec_oc, + dest_vec = dest_vec_cubedsphere, + dest_arr = dest_arr_cubedsphere, ) - OC.fill_halo_regions!(field_latlon) - # heatmap(field_latlon) # using GeoMakie, using CairoMakie - - # LatLon to cubed sphere - boundary_space = CC.CommonSpaces.CubedSphereSpace( - Float32; - radius = 1.0, - n_quad_points = 2, - h_elem = 10, - ) - field_cubedsphere = Interfacer.remap(field_latlon, boundary_space) - # fieldheatmap(field_cubedsphere) # using ClimaCoreMakie - # Cubed sphere to LatLon - long_oc = OC.λnodes(grid_latlon, OC.Center(), OC.Center(), OC.Center()) - lat_oc = OC.φnodes(grid_latlon, OC.Center(), OC.Center(), OC.Center()) + ## Remapper: Cubed sphere nodes to Oceananigans grid `Center, Center` + long_oc = OC.λnodes(grid, OC.Center(), OC.Center(), OC.Center()) + lat_oc = OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center()) long_oc = reshape(long_oc, length(long_oc), 1) lat_oc = reshape(lat_oc, 1, length(lat_oc)) target_points_oc = @. CC.Geometry.LatLongPoint(lat_oc, long_oc) - remapper_cubedsphere_to_latlon = CC.Remapping.Remapper(boundary_space, target_points_oc) - field_cubedsphere = CC.Fields.zeros(boundary_space) - CC.Remapping.interpolate!( - field_cubedsphere, - remapper_cubedsphere_to_latlon, - field_latlon, - ) - - - # Previous remapper for LatLon - if pkgversion(CC) >= v"0.14.34" - remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc) - else - remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc, nothing) - end + # Construct the ClimaCore remapper object + remapper_cubedsphere_to_oceananigans = + CC.Remapping.Remapper(boundary_space, target_points_oc) # Construct two 2D Center/Center fields to use as scratch space while remapping - scratch_cc1 = OC.Field{OC.Center, OC.Center, Nothing}(grid) - scratch_cc2 = OC.Field{OC.Center, OC.Center, Nothing}(grid) + scratch_oc1 = OC.Field{OC.Center, OC.Center, Nothing}(grid) + scratch_oc2 = OC.Field{OC.Center, OC.Center, Nothing}(grid) # Construct two scratch arrays to use while remapping # We get the array type, float type, and dimensions from the remapper object to maintain consistency - ArrayType = ClimaComms.array_type(remapper_cc.space) - FT = CC.Spaces.undertype(remapper_cc.space) - interpolated_values_dim..., _buffer_length = size(remapper_cc._interpolated_values) + ArrayType = ClimaComms.array_type(remapper_cs_to_oc.space) + FT = CC.Spaces.undertype(remapper_cs_to_oc.space) + interpolated_values_dim..., _buffer_length = + size(remapper_cs_to_oc._interpolated_values) scratch_arr1 = ArrayType(zeros(FT, interpolated_values_dim...)) scratch_arr2 = ArrayType(zeros(FT, interpolated_values_dim...)) - remapping = (; remapper_cc, scratch_cc1, scratch_cc2, scratch_arr1, scratch_arr2) - - ocean_properties = (; - ocean_reference_density = 1020, - ocean_heat_capacity = 3991, - ocean_fresh_water_density = 999.8, - ) - - # Save all tracers and velocities to a JLD2 file at daily frequency - outputs = merge(ocean.model.tracers, ocean.model.velocities) - jld2_writer = OC.JLD2Writer( - ocean.model, - outputs; - schedule = OC.TimeInterval(86400), # Daily output - filename = joinpath(output_dir, "ocean_diagnostics"), - indices = (:, :, grid.Nz), - overwrite_existing = true, - array_type = Array{Float32}, + remapper_cs_to_oc = (; + remapper = remapper_cubedsphere_to_oceananigans, + scratch_oc1, + scratch_oc2, + scratch_arr1, + scratch_arr2, ) - ocean.output_writers[:diagnostics] = jld2_writer - sim = OceananigansSimulation(ocean, area_fraction, ocean_properties, remapping) - return sim + return remapper_oc_to_cs, remapper_cs_to_oc end """ @@ -456,17 +407,17 @@ function FluxCalculator.update_turbulent_fluxes!(sim::OceananigansSimulation, fi sim.remapping.remapper_cc, F_turb_ρτxz, ) - OC.set!(sim.remapping.scratch_cc1, sim.remapping.scratch_arr1) # zonal momentum flux + OC.set!(sim.remapping.scratch_oc1, sim.remapping.scratch_arr1) # zonal momentum flux CC.Remapping.interpolate!( sim.remapping.scratch_arr2, sim.remapping.remapper_cc, F_turb_ρτyz, ) - OC.set!(sim.remapping.scratch_cc2, sim.remapping.scratch_arr2) # meridional momentum flux + OC.set!(sim.remapping.scratch_oc2, sim.remapping.scratch_arr2) # meridional momentum flux # Rename for clarity; these are now Center, Center Oceananigans fields - F_turb_ρτxz_cc = sim.remapping.scratch_cc1 - F_turb_ρτyz_cc = sim.remapping.scratch_cc2 + F_turb_ρτxz_cc = sim.remapping.scratch_oc1 + F_turb_ρτyz_cc = sim.remapping.scratch_oc2 # Set the momentum flux BCs at the correct locations using the remapped scratch fields oc_flux_u = surface_flux(sim.ocean.model.velocities.u) From 1ea3d7ae2048bc6d8700c85deb10629ec02fcdfc Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 4 Nov 2025 14:30:00 -0800 Subject: [PATCH 09/17] document remap, add get_remapper --- docs/src/interfacer.md | 65 ++++++ .../components/ocean/oceananigans.jl | 186 ++++++++++-------- src/Interfacer.jl | 47 ++++- 3 files changed, 204 insertions(+), 94 deletions(-) diff --git a/docs/src/interfacer.md b/docs/src/interfacer.md index 91918366ba..ecee3088fe 100644 --- a/docs/src/interfacer.md +++ b/docs/src/interfacer.md @@ -303,6 +303,71 @@ function update_field!(sim::AbstractSurfaceStub, ::Val{:surface_diffuse_albedo}, end ``` +## Remapping functions + +For component models that don't use ClimaCore Fields, some additional functions +must be extended to enable remapping between the component model's grid +and the boundary space of the coupled simulation. + +### `remap(field, target_space, remapper)` + +Remap the given `field` onto the `target_space`. If the field is already +on the target space or a compatible one, it is returned unchanged. + +For ClimaCore Fields, this function is implemented by default and does not require +a remapper object. Component models that use non-ClimaCore fields (such as Oceananigans) +must extend this function to provide a method that handles remapping from their +native field type to a ClimaCore Field on the target space. + +!!! note "Performance" + The `remap` method allocates a new field and is not efficient for + performance-critical code. For better performance, use the in-place option + `remap!` instead. + +**Signature:** +- `field`: The source field to be remapped +- `target_space`: The target space (typically the boundary space) onto which the field should be remapped +- `remapper`: An optional remapper object returned by `get_remapper_to_cc(sim)`. For ClimaCore Fields, this can be `nothing`. + +**Returns:** A new field remapped onto the target space + +### `remap!(target_field, source, remapper)` + +Remap the given `source` field onto the `target_field` in place. This is +the preferred method for remapping when performance is important, as it avoids +allocating a new field. + +For ClimaCore Fields, this function is implemented by default. Component models +that use non-ClimaCore fields must extend this function to provide a method that +handles remapping from their native field type into the provided `target_field`. + +**Signature:** +- `target_field`: The destination field (must be a ClimaCore Field) where remapped data will be stored +- `source`: The source field to be remapped (can be a ClimaCore Field or a non-ClimaCore field type) +- `remapper`: An optional remapper object returned by `get_remapper_to_cc(sim)`. For ClimaCore Fields, this can be `nothing`. + +**Returns:** `nothing` (updates `target_field` in place) + +### `get_remapper_to_cc(sim::ComponentModelSimulation)` + +Return the remapper object used to remap quantities from this component model's grid +onto the boundary space. + +By default, this function returns `nothing`, which is intended for use with components that +use the default remapping functions (i.e. components using ClimaCore Fields). +Components that require an alternative remapper (such as the XESMF regridder for Oceananigans) +should extend this function to return their remapper object. If this function is extended, +the component model will also need to extend `remap` and `remap!` to use the remapper object. + +**Signature:** +- `sim`: The component model simulation + +**Returns:** A remapper object (or `nothing` for ClimaCore-based models) + +**Example:** For an Oceananigans simulation, this returns a NamedTuple containing an +XESMF regridder object with an initialized weight matrix, as well as scratch space +to be used during remapping. + ## Interfacer API ```@docs ClimaCoupler.Interfacer.CoupledSimulation diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 482d7a3460..4beecac7cc 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -22,16 +22,16 @@ It contains the following objects: - `ocean::SIM`: The Oceananigans simulation object. - `area_fraction::A`: A ClimaCore Field representing the surface area fraction of this component model on the exchange grid. - `ocean_properties::OPROP`: A NamedTuple of ocean properties and parameters -- `remapper_oc_to_cs::REMAP1`: Objects needed to remap from the Oceananigans space to the exchange (spectral) grid. -- `remapper_cs_to_oc::REMAP2`: Objects needed to remap from the exchange (spectral) grid to the Oceananigans space. +- `remapper_oc_to_cc::REMAP1`: Objects needed to remap from the Oceananigans space to the exchange (spectral) grid. +- `remapper_cc_to_oc::REMAP2`: Objects needed to remap from the exchange (spectral) grid to the Oceananigans space. """ struct OceananigansSimulation{SIM, A, OPROP, REMAP1, REMAP2} <: Interfacer.OceanModelSimulation ocean::SIM area_fraction::A ocean_properties::OPROP - remapper_oc_to_cs::REMAP1 - remapper_cs_to_oc::REMAP2 + remapper_oc_to_cc::REMAP1 + remapper_cc_to_oc::REMAP2 end """ @@ -109,7 +109,7 @@ function OceananigansSimulation( # Get the remapper objects to go in both directions between the Oceananigans and Cubed sphere grids boundary_space = axes(area_fraction) - remapper_oc_to_cs, remapper_cs_to_oc = construct_remappers(grid, boundary_space) + remapper_oc_to_cc, remapper_cc_to_oc = construct_remappers(grid, boundary_space) ocean_properties = (; ocean_reference_density = 1020, @@ -134,8 +134,8 @@ function OceananigansSimulation( ocean, area_fraction, ocean_properties, - remapper_oc_to_cs, - remapper_cs_to_oc, + remapper_oc_to_cc, + remapper_cc_to_oc, ) return sim end @@ -149,43 +149,43 @@ Both objects contain a remapper object and relevant scratch space. - Oceananigans to ClimaCore In this direction we use XESMF bilinear interpolation. - Example: remap the Oceananigans field `T` to the ClimaCore field `T_cubedsphere` + Example: remap the Oceananigans field `T` to the ClimaCore field `T_climacore` # Convert the Oceananigans field to a flat vector - remapper_oc_to_cs.src_vec .= Array(vec(OC.interior(T, :, :, Nz))) # Oceananigans source vector + remapper_oc_to_cc.src_vec .= Array(vec(OC.interior(T, :, :, Nz))) # Oceananigans source vector # Apply the XESMF regridder (matrix multiply) - remapper_oc_to_cs.remapper(remapper_oc_to_cs.dest_vec, remapper_oc_to_cs.src_vec) + remapper_oc_to_cc.remapper(remapper_oc_to_cc.dest_vec, remapper_oc_to_cc.src_vec) - # Convert the output vector to a 2D array and then to a ClimaCore field - remapper_oc_to_cs.dest_arr .= reshape(remapper_oc_to_cs.dest_vec, size(remapper_oc_to_cs.dest_arr)...) + # Convert the output vector to a 2D array + remapper_oc_to_cc.dest_arr .= reshape(remapper_oc_to_cc.dest_vec, size(remapper_oc_to_cc.dest_arr)...) # Copy the remapped data to the ClimaCore field # Note: in general we avoid accessing the parent of a field, but here we make an exception # since we're using the underlying array to remap. - parent(T_cubedsphere) .= remapper_oc_to_cs.dest_arr + parent(T_climacore) .= remapper_oc_to_cc.dest_arr - ClimaCore to Oceananigans In this direction we use the ClimaCore remapper for this because it uses all the information we have about the spectral element cubed sphere grid, which XESMF does not support. - Example: remap the ClimaCore field `F_turb_ρτxz_cs` to the Oceananigans field `F_turb_ρτxz_oc`: + Example: remap the ClimaCore field `F_turb_ρτxz_cc` to the Oceananigans field `F_turb_ρτxz_oc`: # Remap the ClimaCore momentum flux to a Oceananigans field using scratch arrays and fields CC.Remapping.interpolate!( - remapper_cs_to_oc.scratch_arr1, - remapper_cs_to_oc.remapper, - F_turb_ρτxz_cs, # ClimaCore field + remapper_cc_to_oc.scratch_arr1, + remapper_cc_to_oc.remapper, + F_turb_ρτxz_cc, # ClimaCore field ) - OC.set!(remapper_cs_to_oc.scratch_oc1, remapper_cs_to_oc.scratch_arr1) # zonal momentum flux - F_turb_ρτxz_oc = remapper_cs_to_oc.scratch_oc1 # Oceananigans field + OC.set!(remapper_cc_to_oc.scratch_oc1, remapper_cc_to_oc.scratch_arr1) # zonal momentum flux + F_turb_ρτxz_oc = remapper_cc_to_oc.scratch_oc1 # Oceananigans field Arguments: - `grid`: The Oceananigans grid (TripolarGrid or LatitudeLongitudeGrid). - `boundary_space`: The boundary space (ClimaCore SpectralElementSpace2D). Returns: -- `remapper_oc_to_cs`: The remapper object to go from the Oceananigans grid to the Cubed sphere nodes. +- `remapper_oc_to_cc`: The remapper object to go from the Oceananigans grid to the Cubed sphere nodes. """ function _construct_remappers(grid, boundary_space) ## Remapper: Oceananigans `Center, Center` to Cubed sphere nodes @@ -194,30 +194,30 @@ function _construct_remappers(grid, boundary_space) coords_oc = Dict(k => Array(v) for (k, v) in coords_oc) # Get the latitude and longitude of each node on the boundary space - cubedsphere_coords = CC.Fields.coordinate_field(boundary_space) + climacore_coords = CC.Fields.coordinate_field(boundary_space) # Get the cubed sphere latitude and longitude, each as an Nx1 Matrix - cubedsphere_lat = Array(reshape(vec(parent(cubedsphere_coords.lat)), :, 1)) - cubedsphere_lon = Array(reshape(vec(parent(cubedsphere_coords.long)), :, 1)) + climacore_lat = Array(reshape(vec(parent(climacore_coords.lat)), :, 1)) + climacore_lon = Array(reshape(vec(parent(climacore_coords.long)), :, 1)) - coords_cubedsphere = Dict("lat" => cubedsphere_lat, "lon" => cubedsphere_lon) + coords_climacore = Dict("lat" => climacore_lat, "lon" => climacore_lon) # Construct the XESMF regridder object - regridder_oceananigans_to_cubedsphere = - XESMF.Regridder(coords_oc, coords_cubedsphere; method = "bilinear") + regridder_oceananigans_to_climacore = + XESMF.Regridder(coords_oc, coords_climacore; method = "bilinear") # Allocate space for source an destination vectors to use as intermediate storage src_vec_oc = Array(vec(OC.Field{OC.Center, OC.Center, Nothing}(grid))) # 2D field on Center/Center - field_cubedsphere = CC.Fields.zeros(boundary_space) # 2D field on boundary space (cubed sphere) - dest_vec_cubedsphere = vec(parent(field_cubedsphere)) - dest_arr_cubedsphere = - deepcopy(reshape(dest_vec_cubedsphere, size(parent(field_cubedsphere))...)) + field_climacore = CC.Fields.zeros(boundary_space) # 2D field on boundary space (cubed sphere) + dest_vec_climacore = vec(parent(field_climacore)) + dest_arr_climacore = + deepcopy(reshape(dest_vec_climacore, size(parent(field_climacore))...)) - remapper_oc_to_cs = (; - remapper = regridder_oceananigans_to_cubedsphere, + remapper_oc_to_cc = (; + remapper = regridder_oceananigans_to_climacore, src_vec = src_vec_oc, - dest_vec = dest_vec_cubedsphere, - dest_arr = dest_arr_cubedsphere, + dest_vec = dest_vec_climacore, + dest_arr = dest_arr_climacore, ) ## Remapper: Cubed sphere nodes to Oceananigans grid `Center, Center` @@ -228,7 +228,7 @@ function _construct_remappers(grid, boundary_space) target_points_oc = @. CC.Geometry.LatLongPoint(lat_oc, long_oc) # Construct the ClimaCore remapper object - remapper_cubedsphere_to_oceananigans = + remapper_climacore_to_oceananigans = CC.Remapping.Remapper(boundary_space, target_points_oc) # Construct two 2D Center/Center fields to use as scratch space while remapping @@ -237,22 +237,22 @@ function _construct_remappers(grid, boundary_space) # Construct two scratch arrays to use while remapping # We get the array type, float type, and dimensions from the remapper object to maintain consistency - ArrayType = ClimaComms.array_type(remapper_cs_to_oc.space) - FT = CC.Spaces.undertype(remapper_cs_to_oc.space) + ArrayType = ClimaComms.array_type(remapper_cc_to_oc.space) + FT = CC.Spaces.undertype(remapper_cc_to_oc.space) interpolated_values_dim..., _buffer_length = - size(remapper_cs_to_oc._interpolated_values) + size(remapper_cc_to_oc._interpolated_values) scratch_arr1 = ArrayType(zeros(FT, interpolated_values_dim...)) scratch_arr2 = ArrayType(zeros(FT, interpolated_values_dim...)) - remapper_cs_to_oc = (; - remapper = remapper_cubedsphere_to_oceananigans, + remapper_cc_to_oc = (; + remapper = remapper_climacore_to_oceananigans, scratch_oc1, scratch_oc2, scratch_arr1, scratch_arr2, ) - return remapper_oc_to_cs, remapper_cs_to_oc + return remapper_oc_to_cc, remapper_cc_to_oc end """ @@ -295,46 +295,6 @@ end Interfacer.step!(sim::OceananigansSimulation, t) = OC.time_step!(sim.ocean, float(t) - sim.ocean.model.clock.time) -# We always want the surface, so we always set zero(pt.lat) for z -""" - to_node(pt::CCGeometry.LatLongPoint) - -Transform `LatLongPoint` into a tuple (long, lat, 0), where the 0 is needed because we only -care about the surface. -""" -@inline to_node(pt::CC.Geometry.LatLongPoint) = pt.long, pt.lat, zero(pt.lat) -# This next one is needed if we have "LevelGrid" -@inline to_node(pt::CC.Geometry.LatLongZPoint) = pt.long, pt.lat, zero(pt.lat) - -""" - map_interpolate(points, oc_field::OC.Field) - -Interpolate the given 3D field onto the target points. - -If the underlying grid does not contain a given point, return 0 instead. - -TODO: Use a non-allocating version of this function (simply replace `map` with `map!`) -""" -function map_interpolate(points, oc_field::OC.Field) - loc = map(L -> L(), OC.Fields.location(oc_field)) - grid = oc_field.grid - data = oc_field.data - - # TODO: There has to be a better way - min_lat, max_lat = extrema(OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center())) - - map(points) do pt - FT = eltype(pt) - - # The oceananigans grid does not cover the entire globe, so we should not - # interpolate outside of its latitude bounds. Instead we return 0 - min_lat < pt.lat < max_lat || return FT(0) - - fᵢ = OC.Fields.interpolate(to_node(pt), data, loc, grid) - convert(FT, fᵢ)::FT - end -end - """ surface_flux(f::OC.AbstractField) @@ -349,8 +309,23 @@ function surface_flux(f::OC.AbstractField) end end -function Interfacer.remap(field::OC.Field, target_space) - return map_interpolate(CC.Fields.coordinate_field(target_space), field) +""" + Interfacer.remap(field::OC.Field, target_space, remapper_oc_to_cc) + Interfacer.remap(operation::OC.AbstractOperation, target_space, remapper_oc_to_cc) + +Remap the given Oceananigans field onto the target space using the remapper object. +If an operation is provided, it is evaluated and the resulting field is remapped. + +Arguments: +- `field/operation`: The Oceananigans field or operation to remap. +- `target_space`: The target space (ClimaCore SpectralElementSpace2D). +- `remapper_oc_to_cc`: The remapper object to go from the Oceananigans grid to cubed sphere nodes. +""" +function Interfacer.remap(field::OC.Field, target_space, remapper_oc_to_cc) + # Allocate a new ClimaCore field and remap the data to it + target_field = CC.Fields.zeros(target_space) + Interfacer.remap!(target_field, field, target_space, remapper_oc_to_cc) + return target_field end function Interfacer.remap(operation::OC.AbstractOperations.AbstractOperation, target_space) @@ -359,7 +334,49 @@ function Interfacer.remap(operation::OC.AbstractOperations.AbstractOperation, ta return Interfacer.remap(evaluated_field, target_space) end -Interfacer.get_field(sim::OceananigansSimulation, ::Val{:area_fraction}) = sim.area_fraction +function Interfacer.remap!( + target_field::CC.Fields.Field, + field::OC.Field, + remapper_oc_to_cc, +) + # Since we always regrid from the surface of the ocean, take only the top layer of the field + Nz = field.grid.underlying_grid.Nz + + # Convert the Oceananigans field to a flat vector + remapper_oc_to_cc.src_vec .= Array(vec(OC.interior(field, :, :, Nz))) + + # Apply the XESMF regridder (matrix multiply) + remapper_oc_to_cc.remapper(remapper_oc_to_cc.dest_vec, remapper_oc_to_cc.src_vec) + + # Convert the output vector to a 2D array + remapper_oc_to_cc.dest_arr .= + reshape(remapper_oc_to_cc.dest_vec, size(remapper_oc_to_cc.dest_arr)...) + + # Allocate a new ClimaCore field and copy the remapped data to it + # Note: in general we avoid accessing the parent of a field, but here we make an exception + # since we're already remapping with the underlying array. + parent(target_field) .= remapper_oc_to_cc.dest_arr + return nothing +end + +function Interfacer.remap!( + target_field::CC.Fields.Field, + operation::OC.AbstractOperations.AbstractOperation, + target_space, + remapper_oc_to_cc, +) + evaluated_field = OC.Field(operation) + OC.compute!(evaluated_field) + return Interfacer.remap!(target_field, evaluated_field, target_space, remapper_oc_to_cc) +end + +""" + Interfacer.get_remapper_to_cc(sim::OceananigansSimulation) + +Return the remapper object used to remap quantities from the Oceananigans grid +to the ClimaCore boundary space. +""" +Interfacer.get_remapper_to_cc(sim::OceananigansSimulation) = sim.remapper_oc_to_cc # TODO: Better values for this @@ -367,6 +384,7 @@ Interfacer.get_field(sim::OceananigansSimulation, ::Val{:area_fraction}) = sim.a # Oceananingans with Float64, so we have no way to know the float type here. Sticking with # Float32 ensures that nothing is accidentally promoted to Float64. We will need to change # this anyway. +Interfacer.get_field(sim::OceananigansSimulation, ::Val{:area_fraction}) = sim.area_fraction Interfacer.get_field(sim::OceananigansSimulation, ::Val{:roughness_buoyancy}) = Float32(5.8e-5) Interfacer.get_field(sim::OceananigansSimulation, ::Val{:roughness_momentum}) = diff --git a/src/Interfacer.jl b/src/Interfacer.jl index f9c3a496be..361c86d210 100644 --- a/src/Interfacer.jl +++ b/src/Interfacer.jl @@ -21,6 +21,7 @@ export CoupledSimulation, LandModelSimulation, OceanModelSimulation, get_field, + get_remapper_to_cc, update_field!, AbstractSurfaceStub, SurfaceStub, @@ -241,26 +242,42 @@ get_field(sim::SurfaceModelSimulation, ::Val{:height_disp}) = """ - get_field(sim, what, target_space) + get_field(sim, quantity, target_space) Return `quantity` in `sim` remapped onto the `target_space` This is equivalent to calling `get_field`, and then `remap`. """ function get_field(sim, quantity, target_space) - return remap(get_field(sim, quantity), target_space) + return remap(get_field(sim, quantity), target_space, get_remapper_to_cc(sim)) end """ get_field!(target_field, sim, quantity) -Remap `quantity` in `sim` remapped onto the `target_field`. +Remap `quantity` in `sim` remapped onto the `target_field. +If this component model uses a remapper object to remap quantities onto +the boundary space, `Interfacer.get_remapper_to_cc` should be extended to return +the remapper object. """ function get_field!(target_field, sim, quantity) - remap!(target_field, get_field(sim, quantity)) + remap!(target_field, get_field(sim, quantity), get_remapper_to_cc(sim)) return nothing end +""" + get_remapper_to_cc(::ComponentModelSimulation) + +Return the remapper object used to remap quantities from this component +model onto the boundary space. + +Components that use the default remapping functions (i.e. components using +ClimaCore Fields) do not need to extend this function. Components that require +an alternative remapper should extend this function and likely `remap` and +`remap!` as well. +""" +get_remapper_to_cc(::ComponentModelSimulation) = nothing + """ update_field!(::AtmosModelSimulation, ::Val, _...) @@ -440,7 +457,13 @@ Non-ClimaCore fields should provide a method to this function. """ function remap end -function remap(field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace) +# We don't need a remapper object to remap a ClimaCore field onto a ClimaCore space +remap(field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace) = + remap(field, target_space, nothing) +remap!(target_field::CC.Fields.Field, source::Union{CC.Fields.Field, Number}) = + remap!(target_field, source, nothing) + +function remap(field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace, _) source_space = axes(field) comms_ctx = ClimaComms.context(source_space) @@ -495,20 +518,24 @@ function remap(field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace) end end -function remap(num::Number, target_space::CC.Spaces.AbstractSpace) +function remap(num::Number, target_space::CC.Spaces.AbstractSpace, _) return num end """ - remap!(target_field, source) + remap!(target_field::CC.Fields.Field, source::Union{CC.Fields.Field, Number}, remapper) -Remap the given `source` onto the `target_field`. +Remap the given ClimaCore `source` field onto the ClimaCore `target_field`. Non-ClimaCore fields should provide a method to [`Interfacer.remap`](@ref), or directly to this function. """ -function remap!(target_field, source) - target_field .= remap(source, axes(target_field)) +function remap!( + target_field::CC.Fields.Field, + source::Union{CC.Fields.Field, Number}, + remapper, +) + target_field .= remap(source, axes(target_field), remapper) return nothing end From c8cd77cffa5e0204773cacbf7a2475d1b0f2ac30 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 4 Nov 2025 14:32:11 -0800 Subject: [PATCH 10/17] rename remapper_cc --- .../components/ocean/oceananigans.jl | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 4beecac7cc..e4e1736f74 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -422,13 +422,13 @@ function FluxCalculator.update_turbulent_fluxes!(sim::OceananigansSimulation, fi # Remap momentum fluxes onto reduced 2D Center, Center fields using scratch arrays and fields CC.Remapping.interpolate!( sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, + sim.remapping.remapper_cc_to_oc, F_turb_ρτxz, ) OC.set!(sim.remapping.scratch_oc1, sim.remapping.scratch_arr1) # zonal momentum flux CC.Remapping.interpolate!( sim.remapping.scratch_arr2, - sim.remapping.remapper_cc, + sim.remapping.remapper_cc_to_oc, F_turb_ρτyz, ) OC.set!(sim.remapping.scratch_oc2, sim.remapping.scratch_arr2) # meridional momentum flux @@ -451,8 +451,16 @@ function FluxCalculator.update_turbulent_fluxes!(sim::OceananigansSimulation, fi sim.ocean_properties # Remap the latent and sensible heat fluxes using scratch arrays - CC.Remapping.interpolate!(sim.remapping.scratch_arr1, sim.remapping.remapper_cc, F_lh) # latent heat flux - CC.Remapping.interpolate!(sim.remapping.scratch_arr2, sim.remapping.remapper_cc, F_sh) # sensible heat flux + CC.Remapping.interpolate!( + sim.remapping.scratch_arr1, + sim.remapping.remapper_cc_to_oc, + F_lh, + ) # latent heat flux + CC.Remapping.interpolate!( + sim.remapping.scratch_arr2, + sim.remapping.remapper_cc_to_oc, + F_sh, + ) # sensible heat flux # Rename for clarity; recall F_turb_energy = F_lh + F_sh remapped_F_lh = sim.remapping.scratch_arr1 @@ -470,7 +478,7 @@ function FluxCalculator.update_turbulent_fluxes!(sim::OceananigansSimulation, fi # add the component due to precipitation (that was done with the radiative fluxes) CC.Remapping.interpolate!( sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, + sim.remapping.remapper_cc_to_oc, F_turb_moisture, ) moisture_fresh_water_flux = sim.remapping.scratch_arr1 ./ ocean_fresh_water_density @@ -573,7 +581,7 @@ function FieldExchanger.update_sim!(sim::OceananigansSimulation, csf) # Remap radiative flux onto scratch array; rename for clarity CC.Remapping.interpolate!( sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, + sim.remapping.remapper_cc_to_oc, csf.F_radiative, ) remapped_F_radiative = sim.remapping.scratch_arr1 @@ -587,12 +595,12 @@ function FieldExchanger.update_sim!(sim::OceananigansSimulation, csf) # Remap precipitation fields onto scratch arrays; rename for clarity CC.Remapping.interpolate!( sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, + sim.remapping.remapper_cc_to_oc, csf.P_liq, ) CC.Remapping.interpolate!( sim.remapping.scratch_arr2, - sim.remapping.remapper_cc, + sim.remapping.remapper_cc_to_oc, csf.P_snow, ) remapped_P_liq = sim.remapping.scratch_arr1 From 34d5e903df0a565fa50b3e262dfd575656fba858 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 4 Nov 2025 14:50:59 -0800 Subject: [PATCH 11/17] improve docs, typo --- docs/src/interfacer.md | 13 +++++++++++-- .../ClimaEarth/components/ocean/oceananigans.jl | 6 +++--- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/docs/src/interfacer.md b/docs/src/interfacer.md index ecee3088fe..98fda41448 100644 --- a/docs/src/interfacer.md +++ b/docs/src/interfacer.md @@ -303,11 +303,20 @@ function update_field!(sim::AbstractSurfaceStub, ::Val{:surface_diffuse_albedo}, end ``` -## Remapping functions +## Remapping interface For component models that don't use ClimaCore Fields, some additional functions must be extended to enable remapping between the component model's grid -and the boundary space of the coupled simulation. +and the boundary space of the coupled simulation. These are described below. + +To regrid from a component model's grid to the boundary space, we can typically +use ClimaCore's Remapping module. Users may want to create a remapping object +containing both the ClimaCore.Remapping.Remapper object and scratch space to +reduce allocations during the remapping. +This has been done for the OceananigansSimulation in +`experiments/ClimaEarth/components/ocean/oceananigans.jl`. +This direction simply requires supplying a matrix of ClimaCore.Geometry.LatLongPoint +objects containing latitude/longitude pairs at each point of the source grid. ### `remap(field, target_space, remapper)` diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index e4e1736f74..a8e91639b9 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -237,10 +237,10 @@ function _construct_remappers(grid, boundary_space) # Construct two scratch arrays to use while remapping # We get the array type, float type, and dimensions from the remapper object to maintain consistency - ArrayType = ClimaComms.array_type(remapper_cc_to_oc.space) - FT = CC.Spaces.undertype(remapper_cc_to_oc.space) + ArrayType = ClimaComms.array_type(boundary_space) + FT = CC.Spaces.undertype(boundary_space) interpolated_values_dim..., _buffer_length = - size(remapper_cc_to_oc._interpolated_values) + size(remapper_climacore_to_oceananigans._interpolated_values) scratch_arr1 = ArrayType(zeros(FT, interpolated_values_dim...)) scratch_arr2 = ArrayType(zeros(FT, interpolated_values_dim...)) From 1b7de546ef6740bca536a95ca10b18fcf0e94a99 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 4 Nov 2025 14:51:59 -0800 Subject: [PATCH 12/17] update gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 4f4dd88cd6..02dd40f24f 100644 --- a/.gitignore +++ b/.gitignore @@ -30,6 +30,7 @@ docs/src/generated/ # Experiments !experiments/ClimaEarth/**/Manifest.toml !experiments/ClimaCore/**/Manifest.toml +experiments/ClimaEarth/.CondaPkg/* # Output output/ From fedd3cfb696cebb530d7fc96ca78cdb9e8db743d Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 4 Nov 2025 15:26:05 -0800 Subject: [PATCH 13/17] fix --- experiments/ClimaEarth/components/ocean/oceananigans.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index a8e91639b9..9dafa7967e 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -141,7 +141,7 @@ function OceananigansSimulation( end """ - _construct_remappers(grid, boundary_space) + construct_remappers(grid, boundary_space) Construct the remapper objects to go in both directions between the Oceananigans and Cubed sphere grids. Both objects contain a remapper object and relevant scratch space. @@ -187,7 +187,7 @@ Arguments: Returns: - `remapper_oc_to_cc`: The remapper object to go from the Oceananigans grid to the Cubed sphere nodes. """ -function _construct_remappers(grid, boundary_space) +function construct_remappers(grid, boundary_space) ## Remapper: Oceananigans `Center, Center` to Cubed sphere nodes # Get the Oceananigans coordinates and put them on CPU coords_oc = OceananigansXESMFExt.xesmf_coordinates(grid) From 4ea2ae7a1f7a48c3a8f6db1c9ac4ec92d8b44a59 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Wed, 5 Nov 2025 09:01:58 -0800 Subject: [PATCH 14/17] call xesmf_coordinates correctly [skip ci] --- experiments/ClimaEarth/components/ocean/oceananigans.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 9dafa7967e..83957bd884 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -148,6 +148,7 @@ Both objects contain a remapper object and relevant scratch space. - Oceananigans to ClimaCore In this direction we use XESMF bilinear interpolation. +Note that we assume the Oceananigans Field is on Center, Center, Center. Example: remap the Oceananigans field `T` to the ClimaCore field `T_climacore` @@ -190,7 +191,8 @@ Returns: function construct_remappers(grid, boundary_space) ## Remapper: Oceananigans `Center, Center` to Cubed sphere nodes # Get the Oceananigans coordinates and put them on CPU - coords_oc = OceananigansXESMFExt.xesmf_coordinates(grid) + coords_oc = + OceananigansXESMFExt.xesmf_coordinates(grid, OC.Center(), OC.Center(), OC.Center()) coords_oc = Dict(k => Array(v) for (k, v) in coords_oc) # Get the latitude and longitude of each node on the boundary space From 1c2d5b6310c2bbd32916ace294d2a6969e74d25d Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Wed, 5 Nov 2025 15:28:46 -0800 Subject: [PATCH 15/17] try to make arrays F_CONTIGUOUS --- experiments/ClimaEarth/Manifest-v1.11.toml | 2 +- experiments/ClimaEarth/Project.toml | 1 + .../components/ocean/oceananigans.jl | 27 ++++++++++++++++--- 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/experiments/ClimaEarth/Manifest-v1.11.toml b/experiments/ClimaEarth/Manifest-v1.11.toml index 39f9e76ef5..49dbbde0a4 100644 --- a/experiments/ClimaEarth/Manifest-v1.11.toml +++ b/experiments/ClimaEarth/Manifest-v1.11.toml @@ -2,7 +2,7 @@ julia_version = "1.11.6" manifest_format = "2.0" -project_hash = "bf902fbe4cdfe69d6bdc0c99b12a9c00c5aa28ee" +project_hash = "4aefcc1c61d8db91f115838c0af8a131cde73b99" [[deps.ADTypes]] git-tree-sha1 = "27cecae79e5cc9935255f90c53bb831cc3c870d7" diff --git a/experiments/ClimaEarth/Project.toml b/experiments/ClimaEarth/Project.toml index 2c8fa28891..c3ac1b5651 100644 --- a/experiments/ClimaEarth/Project.toml +++ b/experiments/ClimaEarth/Project.toml @@ -27,6 +27,7 @@ NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" Poppler_jll = "9c32591e-4766-534b-9725-b71a8799265b" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" +PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index 83957bd884..fd89f6a5ca 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -8,6 +8,7 @@ import Thermodynamics as TD import ClimaOcean.EN4: download_dataset using KernelAbstractions: @kernel, @index, @inbounds using XESMF # to load Oceananigans regridding extension +using PythonCall OceananigansXESMFExt = Base.get_extension(OC, :OceananigansXESMFExt).OceananigansXESMFExt; @@ -204,9 +205,20 @@ function construct_remappers(grid, boundary_space) coords_climacore = Dict("lat" => climacore_lat, "lon" => climacore_lon) + # Make arrays F-contiguous (column major) using `np.asfortranarray` to reduce memory footprint and + # prevent segfaults since XESMF works more efficiently with F-contiguous arrays. + np = pyimport("numpy") + coords_oc_contiguous = + Dict(k => pyconvert(Array, np.asfortranarray(v)) for (k, v) in coords_oc) + coords_climacore_contiguous = + Dict(k => pyconvert(Array, np.asfortranarray(v)) for (k, v) in coords_climacore) + # Construct the XESMF regridder object - regridder_oceananigans_to_climacore = - XESMF.Regridder(coords_oc, coords_climacore; method = "bilinear") + regridder_oceananigans_to_climacore = XESMF.Regridder( + coords_oc_contiguous, + coords_climacore_contiguous; + method = "bilinear", + ) # Allocate space for source an destination vectors to use as intermediate storage src_vec_oc = Array(vec(OC.Field{OC.Center, OC.Center, Nothing}(grid))) # 2D field on Center/Center @@ -223,10 +235,17 @@ function construct_remappers(grid, boundary_space) ) ## Remapper: Cubed sphere nodes to Oceananigans grid `Center, Center` + # For a TripolarGrid, latitude and longitude are already 2D arrays, + # so we broadcast over them directly. long_oc = OC.λnodes(grid, OC.Center(), OC.Center(), OC.Center()) lat_oc = OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center()) - long_oc = reshape(long_oc, length(long_oc), 1) - lat_oc = reshape(lat_oc, 1, length(lat_oc)) + + # For a LatitudeLongitudeGrid, latitude and longitude are 1D collections, + # so we reshape them to 2D arrays. + if grid isa OC.LatitudeLongitudeGrid + long_oc = reshape(long_oc, length(long_oc), 1) + lat_oc = reshape(lat_oc, 1, length(lat_oc)) + end target_points_oc = @. CC.Geometry.LatLongPoint(lat_oc, long_oc) # Construct the ClimaCore remapper object From b8f06273b986a42cc23dfc05f26d3d3a0d2b9b4c Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 11 Nov 2025 17:04:30 -0800 Subject: [PATCH 16/17] put lon on [-180, 180] for TripolarGrid --- .../components/ocean/oceananigans.jl | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index fd89f6a5ca..321896fc4b 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -8,7 +8,6 @@ import Thermodynamics as TD import ClimaOcean.EN4: download_dataset using KernelAbstractions: @kernel, @index, @inbounds using XESMF # to load Oceananigans regridding extension -using PythonCall OceananigansXESMFExt = Base.get_extension(OC, :OceananigansXESMFExt).OceananigansXESMFExt; @@ -195,6 +194,10 @@ function construct_remappers(grid, boundary_space) coords_oc = OceananigansXESMFExt.xesmf_coordinates(grid, OC.Center(), OC.Center(), OC.Center()) coords_oc = Dict(k => Array(v) for (k, v) in coords_oc) + if grid.underlying_grid isa OC.TripolarGrid + # TripolarGrid is defined on [0, 360], so we to convert to [-180, 180] to match ClimaCore + coords_oc["lon"] .-= 180 + end # Get the latitude and longitude of each node on the boundary space climacore_coords = CC.Fields.coordinate_field(boundary_space) @@ -205,20 +208,9 @@ function construct_remappers(grid, boundary_space) coords_climacore = Dict("lat" => climacore_lat, "lon" => climacore_lon) - # Make arrays F-contiguous (column major) using `np.asfortranarray` to reduce memory footprint and - # prevent segfaults since XESMF works more efficiently with F-contiguous arrays. - np = pyimport("numpy") - coords_oc_contiguous = - Dict(k => pyconvert(Array, np.asfortranarray(v)) for (k, v) in coords_oc) - coords_climacore_contiguous = - Dict(k => pyconvert(Array, np.asfortranarray(v)) for (k, v) in coords_climacore) - # Construct the XESMF regridder object - regridder_oceananigans_to_climacore = XESMF.Regridder( - coords_oc_contiguous, - coords_climacore_contiguous; - method = "bilinear", - ) + regridder_oceananigans_to_climacore = + XESMF.Regridder(coords_oc, coords_climacore; method = "bilinear") # Allocate space for source an destination vectors to use as intermediate storage src_vec_oc = Array(vec(OC.Field{OC.Center, OC.Center, Nothing}(grid))) # 2D field on Center/Center @@ -242,9 +234,12 @@ function construct_remappers(grid, boundary_space) # For a LatitudeLongitudeGrid, latitude and longitude are 1D collections, # so we reshape them to 2D arrays. - if grid isa OC.LatitudeLongitudeGrid + if grid.underlying_grid isa OC.LatitudeLongitudeGrid long_oc = reshape(long_oc, length(long_oc), 1) lat_oc = reshape(lat_oc, 1, length(lat_oc)) + else + # TripolarGrid is defined on [0, 360], so we to convert to [-180, 180] to match ClimaCore + long_oc .-= 180 end target_points_oc = @. CC.Geometry.LatLongPoint(lat_oc, long_oc) From 0e186a1f3c4068c6f4bc7e6f76885fd08bf77391 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 11 Nov 2025 17:27:37 -0800 Subject: [PATCH 17/17] only run CMIP on buildkite --- .buildkite/pipeline.yml | 748 ++++++++++++++++++++-------------------- 1 file changed, 374 insertions(+), 374 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index a5927400f5..d84e7456c3 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -26,10 +26,10 @@ steps: rm -rf ${JULIA_DEPOT_PATH} fi - - echo "--- Instantiate package env" - - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - - "julia --project -e 'using Pkg; Pkg.precompile()'" - - "julia --project -e 'using Pkg; Pkg.status()'" + # - echo "--- Instantiate package env" + # - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + # - "julia --project -e 'using Pkg; Pkg.precompile()'" + # - "julia --project -e 'using Pkg; Pkg.status()'" - echo "--- Instantiate ClimaEarth experiments env" - "julia --project=experiments/ClimaEarth/ -e 'using Pkg; Pkg.develop(path=\".\")'" @@ -38,19 +38,19 @@ steps: - "julia --project=experiments/ClimaEarth/ -e 'using Pkg; Pkg.precompile()'" - "julia --project=experiments/ClimaEarth/ -e 'using Pkg; Pkg.status()'" - - echo "--- Instantiate ClimaCore experiments env" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.develop(path=\".\")'" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.add(\"MPI\")'" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.precompile()'" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.status()'" + # - echo "--- Instantiate ClimaCore experiments env" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.develop(path=\".\")'" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.add(\"MPI\")'" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.precompile()'" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.status()'" - - echo "--- Instantiate test env" - - "julia --project=test/ -e 'using Pkg; Pkg.develop(path=\".\")'" - - "julia --project=test/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - - "julia --project=test/ -e 'using Pkg; Pkg.add(\"MPI\")'" - - "julia --project=test/ -e 'using Pkg; Pkg.precompile()'" - - "julia --project=test/ -e 'using Pkg; Pkg.status()'" + # - echo "--- Instantiate test env" + # - "julia --project=test/ -e 'using Pkg; Pkg.develop(path=\".\")'" + # - "julia --project=test/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + # - "julia --project=test/ -e 'using Pkg; Pkg.add(\"MPI\")'" + # - "julia --project=test/ -e 'using Pkg; Pkg.precompile()'" + # - "julia --project=test/ -e 'using Pkg; Pkg.status()'" concurrency: 1 concurrency_group: 'depot/climacoupler-ci' @@ -63,339 +63,339 @@ steps: - wait - - group: "Unit Tests" - steps: - - - label: "MPI Utilities unit tests" - key: "utilities_mpi_tests" - command: "srun julia --color=yes --project=test/ test/utilities_tests.jl" - timeout_in_minutes: 5 - env: - CLIMACOMMS_CONTEXT: "MPI" - agents: - slurm_ntasks: 2 - slurm_mem: 16GB - - - label: "MPI Interfacer unit tests" - key: "interfacer_mpi_tests" - command: "srun julia --color=yes --project=test/ test/interfacer_tests.jl" - timeout_in_minutes: 5 - env: - CLIMACOMMS_CONTEXT: "MPI" - agents: - slurm_ntasks: 2 - slurm_mem: 16GB - - - group: "GPU: unit tests" - steps: - - label: "GPU runtests" - command: "julia --color=yes --project=test/ test/runtests.jl" - timeout_in_minutes: 10 - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_ntasks: 1 - slurm_gres: "gpu:1" - slurm_mem: 24GB - - - group: "ClimaEarth tests" - steps: - - label: "ClimaEarth runtests" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/runtests.jl" - agents: - slurm_mem: 16GB - - - label: "MPI restarts" - key: "mpi_restarts" - command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart.jl" - env: - CLIMACOMMS_CONTEXT: "MPI" - timeout_in_minutes: 50 - soft_fail: - - exit_status: -1 - - exit_status: 255 - agents: - slurm_ntasks: 2 - slurm_mem: 32GB - - - label: "GPU restarts" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart.jl" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_ntasks: 1 - slurm_gres: "gpu:1" - slurm_mem: 32GB - - - group: "Integration Tests" - steps: - # SLABPLANET EXPERIMENTS - - # Slabplanet default: - # - this is the most lightweight example with conservation and visual checks, with CLI specification as follows - # - numerics: dt = dt_cpl = 200s, nelem = 4 - # - physics: bulk aerodynamic surface fluxes, gray radiation, idealized insolation, equil moisture model, 0-moment microphysics - # - input data: monotonous remapping (land mask, SST, SIC) - # - slurm: unthreaded, 1 ntask - # - diagnostics: check and plot energy conservation, output plots after 9 days - - label: "Slabplanet: default" - key: "slabplanet_default" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_default.yml --job_id slabplanet_default" - artifact_paths: "experiments/ClimaEarth/output/slabplanet_default/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "Slabplanet: dry, no radiation" - key: "slabplanet_dry_norad" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_dry_norad.yml --job_id slabplanet_dry_norad" - artifact_paths: "experiments/ClimaEarth/output/slabplanet_dry_norad/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "Slabplanet: extra atmos diagnostics" - key: "slabplanet_atmos_diags" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_atmos_diags.yml --job_id slabplanet_atmos_diags" - artifact_paths: "experiments/ClimaEarth/output/slabplanet_atmos_diags/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "Slabplanet: timevarying insolation + rayleigh sponge" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_realinsol_rayleigh.yml --job_id slabplanet_realinsol_rayleigh" - artifact_paths: "experiments/ClimaEarth/output/slabplanet_realinsol_rayleigh/artifacts/total_energy*.png" - agents: - slurm_mem: 20GB - - - label: "Slabplanet terra: atmos and bucket" - key: "slabplanet_terra" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_terra.yml --job_id slabplanet_terra" - artifact_paths: "experiments/ClimaEarth/output/slabplanet_terra/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "Slabplanet aqua: atmos and slab ocean" - key: "slabplanet_aqua" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_aqua.yml --job_id slabplanet_aqua" - artifact_paths: "experiments/ClimaEarth/output/slabplanet_aqua/artifacts/*" - agents: - slurm_mem: 20GB - - # AMIP EXPERIMENTS - - # Test default behavior with no config file or job ID provided - - label: "AMIP: default" - key: "amip_default" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl" - artifact_paths: "experiments/ClimaEarth/output/amip_default/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "AMIP: ED only + integrated land" - key: "amip_edonly_integrated_land" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_integrated_land.yml --job_id amip_edonly_integrated_land" - artifact_paths: "experiments/ClimaEarth/output/amip_edonly_integrated_land/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "AMIP: ED only + bucket" - key: "amip_edonly_bucket" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_bucket.yml --job_id amip_edonly_bucket" - artifact_paths: "experiments/ClimaEarth/output/amip_edonly_bucket/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "AMIP: bucket initial condition test" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_bucket_ic.yml --job_id amip_bucket_ic" - artifact_paths: "experiments/ClimaEarth/output/amip_bucket_ic/artifacts/*" - agents: - slurm_ntasks: 1 - slurm_mem: 20GB - - - label: "AMIP: integrated land non-spun up initial condition test" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_land_ic.yml --job_id amip_land_ic" - artifact_paths: "experiments/ClimaEarth/output/amip_land_ic/artifacts/*" - agents: - slurm_ntasks: 1 - slurm_mem: 20GB - - - label: "AMIP - Float64 + hourly checkpoint" - key: "amip" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_ft64_hourly_checkpoints.yml --job_id amip_coarse_ft64_hourly_checkpoints" - artifact_paths: "experiments/ClimaEarth/output/amip_coarse_ft64_hourly_checkpoints/artifacts/*" - env: - FLAME_PLOT: "" - BUILD_HISTORY_HANDLE: "" - agents: - slurm_ntasks: 1 - slurm_mem: 20GB - - - label: "AMIP - Float32 test" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_ft32.yml --job_id amip_coarse_ft32" - artifact_paths: "experiments/ClimaEarth/output/amip_coarse_ft32/artifacts/*" - agents: - slurm_ntasks: 1 - slurm_mem: 20GB - - - label: "AMIP - Component dts test" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_component_dts.yml --job_id target_amip_component_dts" - artifact_paths: "experiments/ClimaEarth/output/target_amip_component_dts/artifacts/*" - agents: - slurm_ntasks: 1 - slurm_mem: 20GB - - - label: "MPI AMIP" - command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_mpi.yml --job_id amip_coarse_mpi" - artifact_paths: "experiments/ClimaEarth/output/amip_coarse_mpi/artifacts/*" - timeout_in_minutes: 30 - env: - CLIMACOMMS_CONTEXT: "MPI" - agents: - slurm_ntasks: 4 - slurm_mem_per_cpu: 12GB - - # short high-res performance test - - label: "Unthreaded AMIP FINE" # also reported by longruns with a flame graph - key: "unthreaded_amip_fine" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_n1_shortrun.yml --job_id target_amip_n1_shortrun" - artifact_paths: "experiments/ClimaEarth/output/target_amip_n1_shortrun/artifacts/*" - env: - BUILD_HISTORY_HANDLE: "" - agents: - slurm_mem: 20GB - - # CLIMACORE EXPERIMENTS - - - label: "sea_breeze" - command: "julia --color=yes --project=experiments/ClimaCore experiments/ClimaCore/sea_breeze/run.jl" - artifact_paths: "experiments/ClimaCore/sea_breeze/output/*" - agents: - slurm_mem: 20GB - - - label: "heat-diffusion" - command: "julia --color=yes --project=experiments/ClimaCore/ experiments/ClimaCore/heat-diffusion/run.jl" - artifact_paths: "experiments/ClimaCore/output/heat-diffusion/artifacts/*" - agents: - slurm_mem: 20GB - - - group: "GPU integration tests" - steps: - # GPU RUNS: slabplanet - - label: "GPU Slabplanet: default" - key: "gpu_slabplanet_default" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_default.yml --job_id gpu_slabplanet_default" - artifact_paths: "experiments/ClimaEarth/output/gpu_slabplanet_default/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "GPU Slabplanet: albedo from function" - key: "gpu_slabplanet_albedo_function" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_albedo_function.yml --job_id gpu_slabplanet_albedo_function" - artifact_paths: "experiments/ClimaEarth/output/gpu_slabplanet_albedo_function/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU Slabplanet: extra atmos diagnostics" - key: "gpu_slabplanet_atmos_diags" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_atmos_diags.yml --job_id gpu_slabplanet_atmos_diags" - artifact_paths: "experiments/ClimaEarth/output/gpu_slabplanet_atmos_diags/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - # GPU RUNS: AMIP - - label: "GPU AMIP: default" - key: "gpu_amip_default" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_default.yml --job_id gpu_amip_default" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_default/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: ED only + integrated land" - key: "gpu_amip_edonly_integrated_land" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_integrated_land.yml --job_id gpu_amip_edonly_integrated_land" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_edonly_integrated_land/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: ED only + bucket" - key: "gpu_amip_edonly_bucket" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_bucket.yml --job_id gpu_amip_edonly_bucket" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_edonly_bucket/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: diag. EDMF + integrated land" - key: "gpu_amip_diagedmf_integrated_land" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_diagedmf_integrated_land.yml --job_id gpu_amip_diagedmf_integrated_land" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_diagedmf_integrated_land/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: diag. EDMF + bucket" - key: "gpu_amip_diagedmf_bucket" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_diagedmf_bucket.yml --job_id gpu_amip_diagedmf_bucket" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_diagedmf_bucket/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP test: albedo from function" - key: "gpu_amip_albedo_function" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_function.yml --job_id gpu_amip_albedo_function" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_albedo_function/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP target: topography and diagnostic EDMF" - key: "gpu_amip_target_topo_diagedmf_shortrun" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_target_topo_diagedmf_shortrun.yml --job_id gpu_amip_target_topo_diagedmf_shortrun" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_target_topo_diagedmf_shortrun/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: albedo from temporal map + 0M" - key: "gpu_amip_albedo_temporal_map" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_temporal_map.yml --job_id gpu_amip_albedo_temporal_map" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_albedo_temporal_map/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: albedo from temporal map + 1M" - key: "gpu_amip_albedo_temporal_map_1M" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_temporal_map_1M.yml --job_id gpu_amip_albedo_temporal_map_1M" - artifact_paths: "experiments/ClimaEarth/output/gpu_amip_albedo_temporal_map_1M/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 + # - group: "Unit Tests" + # steps: + + # - label: "MPI Utilities unit tests" + # key: "utilities_mpi_tests" + # command: "srun julia --color=yes --project=test/ test/utilities_tests.jl" + # timeout_in_minutes: 5 + # env: + # CLIMACOMMS_CONTEXT: "MPI" + # agents: + # slurm_ntasks: 2 + # slurm_mem: 16GB + + # - label: "MPI Interfacer unit tests" + # key: "interfacer_mpi_tests" + # command: "srun julia --color=yes --project=test/ test/interfacer_tests.jl" + # timeout_in_minutes: 5 + # env: + # CLIMACOMMS_CONTEXT: "MPI" + # agents: + # slurm_ntasks: 2 + # slurm_mem: 16GB + + # - group: "GPU: unit tests" + # steps: + # - label: "GPU runtests" + # command: "julia --color=yes --project=test/ test/runtests.jl" + # timeout_in_minutes: 10 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_ntasks: 1 + # slurm_gres: "gpu:1" + # slurm_mem: 24GB + + # - group: "ClimaEarth tests" + # steps: + # - label: "ClimaEarth runtests" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/runtests.jl" + # agents: + # slurm_mem: 16GB + + # - label: "MPI restarts" + # key: "mpi_restarts" + # command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart.jl" + # env: + # CLIMACOMMS_CONTEXT: "MPI" + # timeout_in_minutes: 50 + # soft_fail: + # - exit_status: -1 + # - exit_status: 255 + # agents: + # slurm_ntasks: 2 + # slurm_mem: 32GB + + # - label: "GPU restarts" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart.jl" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_ntasks: 1 + # slurm_gres: "gpu:1" + # slurm_mem: 32GB + + # - group: "Integration Tests" + # steps: + # # SLABPLANET EXPERIMENTS + + # # Slabplanet default: + # # - this is the most lightweight example with conservation and visual checks, with CLI specification as follows + # # - numerics: dt = dt_cpl = 200s, nelem = 4 + # # - physics: bulk aerodynamic surface fluxes, gray radiation, idealized insolation, equil moisture model, 0-moment microphysics + # # - input data: monotonous remapping (land mask, SST, SIC) + # # - slurm: unthreaded, 1 ntask + # # - diagnostics: check and plot energy conservation, output plots after 9 days + # - label: "Slabplanet: default" + # key: "slabplanet_default" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_default.yml --job_id slabplanet_default" + # artifact_paths: "experiments/ClimaEarth/output/slabplanet_default/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet: dry, no radiation" + # key: "slabplanet_dry_norad" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_dry_norad.yml --job_id slabplanet_dry_norad" + # artifact_paths: "experiments/ClimaEarth/output/slabplanet_dry_norad/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet: extra atmos diagnostics" + # key: "slabplanet_atmos_diags" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_atmos_diags.yml --job_id slabplanet_atmos_diags" + # artifact_paths: "experiments/ClimaEarth/output/slabplanet_atmos_diags/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet: timevarying insolation + rayleigh sponge" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_realinsol_rayleigh.yml --job_id slabplanet_realinsol_rayleigh" + # artifact_paths: "experiments/ClimaEarth/output/slabplanet_realinsol_rayleigh/artifacts/total_energy*.png" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet terra: atmos and bucket" + # key: "slabplanet_terra" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_terra.yml --job_id slabplanet_terra" + # artifact_paths: "experiments/ClimaEarth/output/slabplanet_terra/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet aqua: atmos and slab ocean" + # key: "slabplanet_aqua" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_aqua.yml --job_id slabplanet_aqua" + # artifact_paths: "experiments/ClimaEarth/output/slabplanet_aqua/artifacts/*" + # agents: + # slurm_mem: 20GB + + # # AMIP EXPERIMENTS + + # # Test default behavior with no config file or job ID provided + # - label: "AMIP: default" + # key: "amip_default" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl" + # artifact_paths: "experiments/ClimaEarth/output/amip_default/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "AMIP: ED only + integrated land" + # key: "amip_edonly_integrated_land" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_integrated_land.yml --job_id amip_edonly_integrated_land" + # artifact_paths: "experiments/ClimaEarth/output/amip_edonly_integrated_land/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "AMIP: ED only + bucket" + # key: "amip_edonly_bucket" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_bucket.yml --job_id amip_edonly_bucket" + # artifact_paths: "experiments/ClimaEarth/output/amip_edonly_bucket/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "AMIP: bucket initial condition test" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_bucket_ic.yml --job_id amip_bucket_ic" + # artifact_paths: "experiments/ClimaEarth/output/amip_bucket_ic/artifacts/*" + # agents: + # slurm_ntasks: 1 + # slurm_mem: 20GB + + # - label: "AMIP: integrated land non-spun up initial condition test" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_land_ic.yml --job_id amip_land_ic" + # artifact_paths: "experiments/ClimaEarth/output/amip_land_ic/artifacts/*" + # agents: + # slurm_ntasks: 1 + # slurm_mem: 20GB + + # - label: "AMIP - Float64 + hourly checkpoint" + # key: "amip" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_ft64_hourly_checkpoints.yml --job_id amip_coarse_ft64_hourly_checkpoints" + # artifact_paths: "experiments/ClimaEarth/output/amip_coarse_ft64_hourly_checkpoints/artifacts/*" + # env: + # FLAME_PLOT: "" + # BUILD_HISTORY_HANDLE: "" + # agents: + # slurm_ntasks: 1 + # slurm_mem: 20GB + + # - label: "AMIP - Float32 test" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_ft32.yml --job_id amip_coarse_ft32" + # artifact_paths: "experiments/ClimaEarth/output/amip_coarse_ft32/artifacts/*" + # agents: + # slurm_ntasks: 1 + # slurm_mem: 20GB + + # - label: "AMIP - Component dts test" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_component_dts.yml --job_id target_amip_component_dts" + # artifact_paths: "experiments/ClimaEarth/output/target_amip_component_dts/artifacts/*" + # agents: + # slurm_ntasks: 1 + # slurm_mem: 20GB + + # - label: "MPI AMIP" + # command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_mpi.yml --job_id amip_coarse_mpi" + # artifact_paths: "experiments/ClimaEarth/output/amip_coarse_mpi/artifacts/*" + # timeout_in_minutes: 30 + # env: + # CLIMACOMMS_CONTEXT: "MPI" + # agents: + # slurm_ntasks: 4 + # slurm_mem_per_cpu: 12GB + + # # short high-res performance test + # - label: "Unthreaded AMIP FINE" # also reported by longruns with a flame graph + # key: "unthreaded_amip_fine" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_n1_shortrun.yml --job_id target_amip_n1_shortrun" + # artifact_paths: "experiments/ClimaEarth/output/target_amip_n1_shortrun/artifacts/*" + # env: + # BUILD_HISTORY_HANDLE: "" + # agents: + # slurm_mem: 20GB + + # # CLIMACORE EXPERIMENTS + + # - label: "sea_breeze" + # command: "julia --color=yes --project=experiments/ClimaCore experiments/ClimaCore/sea_breeze/run.jl" + # artifact_paths: "experiments/ClimaCore/sea_breeze/output/*" + # agents: + # slurm_mem: 20GB + + # - label: "heat-diffusion" + # command: "julia --color=yes --project=experiments/ClimaCore/ experiments/ClimaCore/heat-diffusion/run.jl" + # artifact_paths: "experiments/ClimaCore/output/heat-diffusion/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - group: "GPU integration tests" + # steps: + # # GPU RUNS: slabplanet + # - label: "GPU Slabplanet: default" + # key: "gpu_slabplanet_default" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_default.yml --job_id gpu_slabplanet_default" + # artifact_paths: "experiments/ClimaEarth/output/gpu_slabplanet_default/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "GPU Slabplanet: albedo from function" + # key: "gpu_slabplanet_albedo_function" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_albedo_function.yml --job_id gpu_slabplanet_albedo_function" + # artifact_paths: "experiments/ClimaEarth/output/gpu_slabplanet_albedo_function/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU Slabplanet: extra atmos diagnostics" + # key: "gpu_slabplanet_atmos_diags" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_atmos_diags.yml --job_id gpu_slabplanet_atmos_diags" + # artifact_paths: "experiments/ClimaEarth/output/gpu_slabplanet_atmos_diags/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # # GPU RUNS: AMIP + # - label: "GPU AMIP: default" + # key: "gpu_amip_default" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_default.yml --job_id gpu_amip_default" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_default/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: ED only + integrated land" + # key: "gpu_amip_edonly_integrated_land" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_integrated_land.yml --job_id gpu_amip_edonly_integrated_land" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_edonly_integrated_land/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: ED only + bucket" + # key: "gpu_amip_edonly_bucket" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_bucket.yml --job_id gpu_amip_edonly_bucket" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_edonly_bucket/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: diag. EDMF + integrated land" + # key: "gpu_amip_diagedmf_integrated_land" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_diagedmf_integrated_land.yml --job_id gpu_amip_diagedmf_integrated_land" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_diagedmf_integrated_land/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: diag. EDMF + bucket" + # key: "gpu_amip_diagedmf_bucket" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_diagedmf_bucket.yml --job_id gpu_amip_diagedmf_bucket" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_diagedmf_bucket/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP test: albedo from function" + # key: "gpu_amip_albedo_function" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_function.yml --job_id gpu_amip_albedo_function" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_albedo_function/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP target: topography and diagnostic EDMF" + # key: "gpu_amip_target_topo_diagedmf_shortrun" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_target_topo_diagedmf_shortrun.yml --job_id gpu_amip_target_topo_diagedmf_shortrun" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_target_topo_diagedmf_shortrun/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: albedo from temporal map + 0M" + # key: "gpu_amip_albedo_temporal_map" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_temporal_map.yml --job_id gpu_amip_albedo_temporal_map" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_albedo_temporal_map/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: albedo from temporal map + 1M" + # key: "gpu_amip_albedo_temporal_map_1M" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_temporal_map_1M.yml --job_id gpu_amip_albedo_temporal_map_1M" + # artifact_paths: "experiments/ClimaEarth/output/gpu_amip_albedo_temporal_map_1M/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 - group: "CMIP" steps: @@ -410,28 +410,28 @@ steps: slurm_mem: 20GB slurm_gpus: 1 - - group: "Calibration experiments" - steps: - - label: "Perfect model calibration test" - key: "amip_pm_calibration" - command: - - "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/run_calibration.jl" - artifact_paths: "experiments/calibration/output/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - CLIMACOMMS_CONTEXT: "SINGLETON" - SHORT_RUN: "" - agents: - slurm_mem: 64GB - slurm_ntasks: 3 - slurm_gpus_per_task: 1 - slurm_cpus_per_task: 4 - - - wait - - # plot job performance history - - label: ":chart_with_downwards_trend: build history" - command: - - build_history staging # name of branch to plot - artifact_paths: - - "build_history.html" + # - group: "Calibration experiments" + # steps: + # - label: "Perfect model calibration test" + # key: "amip_pm_calibration" + # command: + # - "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/run_calibration.jl" + # artifact_paths: "experiments/calibration/output/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # CLIMACOMMS_CONTEXT: "SINGLETON" + # SHORT_RUN: "" + # agents: + # slurm_mem: 64GB + # slurm_ntasks: 3 + # slurm_gpus_per_task: 1 + # slurm_cpus_per_task: 4 + + # - wait + + # # plot job performance history + # - label: ":chart_with_downwards_trend: build history" + # command: + # - build_history staging # name of branch to plot + # artifact_paths: + # - "build_history.html"