Further enhancements of the XESMF extension (#4819)

* tweaks

* generalize xesmf regridder for GPUs

* fix example

* cleanup

* add comment

Author: navidcy

.buildkite/pipeline.yml

@@ -87,6 +87,7 @@ steps:
           - "🧅 multi_region"
           - "🦧 scripts"
           - "👺 enzyme"
+          - "🍱 xesmf"
           - "👹 reactant_1"
           - "🎭 reactant_2"
     retry:

.github/workflows/ci.yml

@@ -159,33 +159,3 @@ jobs:
       - uses: julia-actions/julia-runtest@v1
         env:
           TEST_GROUP: "metal"
-
-
-  xesmf:
-    name: XESMF - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
-    runs-on: ${{ matrix.os }}
-    timeout-minutes: 60
-    strategy:
-      fail-fast: false
-      matrix:
-        version:
-          - '1.10'
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-        include:
-          - os: macOS-latest
-            arch: aarch64
-            version: '1.10'
-    steps:
-      - uses: actions/checkout@v5
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
-      - uses: julia-actions/cache@v2
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
-        env:
-          TEST_GROUP: "xesmf"

ext/OceananigansXESMFExt.jl

@@ -9,24 +9,23 @@ using Oceananigans.Grids: λnodes, φnodes, Center, Face, total_length
 import Oceananigans.Fields: regrid!
 import XESMF: Regridder, extract_xesmf_coordinates_structure
 
-function x_node_array(x::AbstractVector, Nx, Ny)
-    return Array(repeat(view(x, 1:Nx), 1, Ny))'
-end
-function  y_node_array(x::AbstractVector, Nx, Ny)
-    return Array(repeat(view(x, 1:Ny)', Nx, 1))'
-end
-x_node_array(x::AbstractMatrix, Nx, Ny) = Array(view(x, 1:Nx, 1:Ny))'
+# permutedims below is used because Python's xESMF expects
+# 2D arrays with (x, y) coordinates with y varying in dim=1 and x varying in dim=2
 
-function x_vertex_array(x::AbstractVector, Nx, Ny)
-    return Array(repeat(view(x, 1:Nx+1), 1, Ny+1))'
-end
-function y_vertex_array(x::AbstractVector, Nx, Ny)
-    return Array(repeat(view(x, 1:Ny+1)', Nx+1, 1))'
-end
-x_vertex_array(x::AbstractMatrix, Nx, Ny) = Array(view(x, 1:Nx+1, 1:Ny+1))'
+node_array(ξ::AbstractMatrix, Nx, Ny) = permutedims(view(ξ, 1:Nx, 1:Ny), (2, 1))
+vertex_array(ξ::AbstractMatrix, Nx, Ny) = permutedims(view(ξ, 1:Nx+1, 1:Ny+1), (2, 1))
+
+x_node_array(x::AbstractVector, Nx, Ny) = permutedims(repeat(view(x, 1:Nx), 1, Ny), (2, 1))
+x_node_array(x::AbstractMatrix, Nx, Ny) = node_array(x, Nx, Ny)
+
+y_node_array(y::AbstractVector, Nx, Ny) = repeat(view(y, 1:Ny), 1, Nx)
+y_node_array(y::AbstractMatrix, Nx, Ny) = node_array(y, Nx, Ny)
+
+x_vertex_array(x::AbstractVector, Nx, Ny) = permutedims(repeat(view(x, 1:Nx+1), 1, Ny+1), (2, 1))
+x_vertex_array(x::AbstractMatrix, Nx, Ny) = vertex_array(x, Nx, Ny)
 
-y_node_array(x::AbstractMatrix, Nx, Ny) = x_node_array(x, Nx, Ny)
-y_vertex_array(x::AbstractMatrix, Nx, Ny) = x_vertex_array(x, Nx, Ny)
+y_vertex_array(y::AbstractVector, Nx, Ny) = repeat(view(y, 1:Ny+1), 1, Nx+1)
+y_vertex_array(y::AbstractMatrix, Nx, Ny) = vertex_array(y, Nx, Ny)
 
 function extract_xesmf_coordinates_structure(dst_field::AbstractField, src_field::AbstractField)
 

test/runtests.jl

@@ -252,6 +252,13 @@ CUDA.allowscalar() do
         end
     end
 
+    # Tests for XESMF extension
+    if group == :xesmf || group == :all
+        @testset "XESMF extension tests" begin
+            include("test_xesmf.jl")
+        end
+    end
+
     if group == :sharding || group == :all
         @testset "Sharding Reactant extension tests" begin
             # Broken for the moment (trying to fix them in https://github.com/CliMA/Oceananigans.jl/pull/4293)
@@ -281,13 +288,6 @@ CUDA.allowscalar() do
         end
     end
 
-    # Tests for XESMF extension
-    if group == :xesmf || group == :all
-        @testset "XESMF extension tests" begin
-            include("test_xesmf.jl")
-        end
-    end
-
     if group == :convergence
         include("test_convergence.jl")
     end

test/test_xesmf.jl

@@ -4,44 +4,53 @@ using XESMF
 using SparseArrays
 using LinearAlgebra
 
-z = (-1, 0)
-southernmost_latitude = -80
-radius = Oceananigans.Grids.R_Earth
+for arch in archs
+    @testset "XESMF extension [$(typeof(arch))]" begin
+        @info "Testing XESMF regridding [$(typeof(arch))]..."
 
-llg_coarse = LatitudeLongitudeGrid(; size=(176, 88, 1),
-                                   longitude=(0, 360),
-                                   latitude=(southernmost_latitude, 90),
-                                   z, radius)
+        z = (-1, 0)
+        southernmost_latitude = -80
+        radius = Oceananigans.Grids.R_Earth
 
-llg_fine = LatitudeLongitudeGrid(; size=(360, 180, 1),
-                                 longitude=(0, 360),
-                                 latitude=(southernmost_latitude, 90),
-                                 z, radius)
+        llg_coarse = LatitudeLongitudeGrid(arch; size=(176, 88, 1),
+                                        longitude=(0, 360),
+                                        latitude=(southernmost_latitude, 90),
+                                        z, radius)
 
-tg = TripolarGrid(; size=(360, 170, 1), z, southernmost_latitude, radius)
+        llg_fine = LatitudeLongitudeGrid(arch; size=(360, 170, 1),
+                                        longitude=(0, 360),
+                                        latitude=(southernmost_latitude, 90),
+                                        z, radius)
 
-@testset "XESMF extension" begin
+        tg = TripolarGrid(arch; size=(360, 170, 1), z, southernmost_latitude, radius)
 
-    for (src_grid, dst_grid) in ((llg_coarse, llg_fine),
-                                 (llg_fine, llg_coarse),
-                                 (tg, llg_fine))
 
-        @info "  Regridding from $(nameof(typeof(src_grid))) to $(nameof(typeof(dst_grid)))"
+        for (src_grid, dst_grid) in ((llg_coarse, llg_fine),
+                                     (llg_fine, llg_coarse),
+                                     (tg, llg_fine))
 
-        src_field = CenterField(src_grid)
-        dst_field = CenterField(dst_grid)
+            @info "  Regridding from $(nameof(typeof(src_grid))) to $(nameof(typeof(dst_grid)))"
 
-        λ₀, φ₀ = 150, 30.  # degrees
-        width = 12         # degrees
-        set!(src_field, (λ, φ, z) -> exp(-((λ - λ₀)^2 + (φ - φ₀)^2) / 2width^2))
+            src_field = CenterField(src_grid)
+            dst_field = CenterField(dst_grid)
 
-        R = XESMF.Regridder(dst_field, src_field)
-        @test R.weights isa SparseMatrixCSC
+            λ₀, φ₀ = 150, 30.  # degrees
+            width = 12         # degrees
+            set!(src_field, (λ, φ, z) -> exp(-((λ - λ₀)^2 + (φ - φ₀)^2) / 2width^2))
 
-        regrid!(dst_field, R, src_field)
+            regridder = XESMF.Regridder(dst_field, src_field)
 
-        # ∫ dst_field dA = ∫ src_field dA
-        @test isapprox(first(Field(Integral(dst_field, dims=(1, 2)))),
-                       first(Field(Integral(src_field, dims=(1, 2)))), rtol=1e-4)
+            if arch isa CPU
+                @test regridder.weights isa SparseMatrixCSC
+            elseif arch isa GPU{CUDABackend}
+                @test regridder.weights isa CUDA.CUSPARSE.CuSparseMatrixCSC
+            end
+
+            regrid!(dst_field, regridder, src_field)
+
+            # ∫ dst_field dA ≈ ∫ src_field dA
+            @test @allowscalar isapprox(first(Field(Integral(dst_field, dims=(1, 2)))),
+                                        first(Field(Integral(src_field, dims=(1, 2)))), rtol=1e-4)
+        end
     end
 end