try out XESMF regridder

Author: juliasloan25

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"

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"

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)