try to make arrays F_CONTIGUOUS

Author: juliasloan25

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"

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"

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