Add extension to ConservativeRegridding.jl

.github/workflows/ci.yml

@@ -26,19 +26,22 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 
 jobs:
-  sharding_tests:
-    name: Sharding Tests - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+  reactant_tests:
+    name: ${{ matrix.test_name }} - Julia ${{ matrix.version }} - ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     timeout-minutes: 120
     strategy:
       fail-fast: false
       matrix:
-        version:
-          - '1.10'
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
+        include:
+          - test_name: "Serial Reactant"
+            test_group: "reactant"
+            version: '1.10'
+            os: 'ubuntu-latest'
+          - test_name: "Sharding with Reactant"
+            test_group: "sharding"
+            version: '1.10'
+            os: 'ubuntu-latest'
     steps:
       - run: |
           touch LocalPreferences.toml
@@ -51,112 +54,51 @@ jobs:
       - 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:
           JULIA_DEBUG: "Reactant, Reactant_jll"
           REACTANT_TEST: true
-          TEST_GROUP: "sharding"
+          TEST_GROUP: ${{ matrix.test_group }}
 
-  mpi_tripolar:
-    name: Distributed MPI Tripolar Grid - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
-    runs-on: ${{ matrix.os }}
-    timeout-minutes: 120
-    strategy:
-      fail-fast: false
-      matrix:
-        version:
-          - '1.10'
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-    steps:
-      - uses: actions/checkout@v4
-      - 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: "mpi_tripolar"
-
-  turbulent_closures:
-    name: Turbulence closures - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
-    runs-on: ${{ matrix.os }}
-    timeout-minutes: 120
-    strategy:
-      fail-fast: false
-      matrix:
-        version:
-          - '1.10'
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-    steps:
-      - uses: actions/checkout@v4
-      - 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: "turbulence_closures"
-
-  reactant:
-    name: Reactant - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
-    runs-on: ${{ matrix.os }}
-    timeout-minutes: 120
-    strategy:
-      fail-fast: false
-      matrix:
-        version:
-          - '1.10'
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-    steps:
-      - uses: actions/checkout@v4
-      - 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: "reactant"
-
-  metal:
-    name: Metal - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+  additional_tests:
+    name: ${{ matrix.test_name }} - Julia ${{ matrix.version }} - ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     timeout-minutes: 60
     strategy:
       fail-fast: false
       matrix:
-        version:
-          - '1.10'
-        os:
-          - macos-latest
-        arch:
-          - arm64
+        os: ['ubuntu-latest']
+        include:
+          - test_name: "Distributed MPI Tripolar Grid"
+            test_group: "mpi_tripolar"
+            version: '1.10'
+          - test_name: "Turbulence Closures"
+            test_group: "turbulence_closures"
+            version: '1.10'
+          - test_name: "Metal"
+            test_group: "metal"
+            os: macos-latest
+            version: '1.10'
+          - test_name: "Conservative Regridding"
+            test_group: "conservative_regridding"
+            version: '1.10'
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
+      - name: Instantiate environment
+        timeout-minutes: 90
+        shell: julia --color=yes --project {0}
+        run: |
+          using Pkg
+          Pkg.Registry.update()
+          Pkg.add([PackageSpec(; name="ConservativeRegridding", url="https://github.com/JuliaGeo/ConservativeRegridding.jl.git", rev="mk/interface")])
+          Pkg.instantiate()
       - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         env:
-          TEST_GROUP: "metal"
-
+          TEST_GROUP: ${{ matrix.test_group }}

Project.toml

@@ -40,6 +40,7 @@ TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
+ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 MakieCore = "20f20a25-4f0e-4fdf-b5d1-57303727442b"
@@ -56,11 +57,13 @@ OceananigansMetalExt = "Metal"
 OceananigansNCDatasetsExt = "NCDatasets"
 OceananigansOneAPIExt = "oneAPI"
 OceananigansReactantExt = ["Reactant", "KernelAbstractions", "ConstructionBase"]
+OceananigansConservativeRegriddingExt = ["ConservativeRegridding"]
 
 [compat]
 AMDGPU = "1.2.7"
 Adapt = "4.1.1"
 CUDA = "5.7"
+ConservativeRegridding = "0.1.0"
 ConstructionBase = "1"
 Crayons = "4"
 CubedSphere = "0.2, 0.3"
@@ -103,6 +106,7 @@ oneAPI = "2.0.1"
 [extras]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
+ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
@@ -114,4 +118,4 @@ TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"
 oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 
 [targets]
-test = ["AMDGPU", "oneAPI", "DataDeps", "SafeTestsets", "Test", "Enzyme", "Reactant", "Metal", "CUDA_Runtime_jll", "MPIPreferences", "TimesDates", "NCDatasets"]
+test = ["AMDGPU", "oneAPI", "DataDeps", "SafeTestsets", "Test", "Enzyme", "Reactant", "Metal", "CUDA_Runtime_jll", "MPIPreferences", "TimesDates", "NCDatasets", "ConservativeRegridding"]

ext/OceananigansConservativeRegriddingExt.jl

@@ -0,0 +1,147 @@
+module OceananigansConservativeRegriddingExt
+
+using ConservativeRegridding
+import ConservativeRegridding: Regridder, regrid!
+
+using Oceananigans
+using Oceananigans.Grids: ξnode, ηnode, architecture, AbstractGrid
+using Oceananigans.Fields: AbstractField, location
+using KernelAbstractions: @index, @kernel
+
+# TODO: extend regridding to more cases
+const RegriddableField{LX, LY} = Field{LX, LY, Nothing}
+
+compute_cell_matrix(field::AbstractField) = 
+    compute_cell_matrix(field.grid, location(field), field.indices)
+
+function compute_cell_matrix(grid::AbstractGrid, (LX, LY, LZ), indices=Oceananigans.Grids.default_indices(3))
+
+    ℓx = LX()
+    ℓy = LY()
+
+    if isnothing(ℓx) || isnothing(ℓy)
+        throw(ArgumentError("cell_matrix can only be computed for fields with non-nothing horizontal location."))
+    end
+
+    arch = architecture(grid)
+    FT = eltype(grid)
+
+    Fx, Fy, _ = size(grid, (LX, LY, LZ), indices)
+
+    vertices_per_cell = 5 # convention: [sw, nw, ne, se, sw]
+    ArrayType = Oceananigans.Architectures.array_type(arch)
+    cell_matrix = ArrayType{Tuple{FT, FT}}(undef, vertices_per_cell, Fx*Fy)
+
+    Oceananigans.Utils.launch!(arch, grid, (Fx, Fy), _compute_cell_matrix!, cell_matrix, Fx, ℓx, ℓy, grid)
+
+    return cell_matrix
+end
+
+flip(::Face) = Center()
+flip(::Center) = Face()
+
+left_index(i, ::Center) = i
+left_index(i, ::Face) = i-1
+right_index(i, ::Center) = i + 1
+right_index(i, ::Face) = i
+
+@kernel function _compute_cell_matrix!(cell_matrix, Fx, ℓx, ℓy, grid)
+    i, j = @index(Global, NTuple)
+
+    vx = flip(ℓx)
+    vy = flip(ℓy)
+
+    isw = left_index(i, ℓx)
+    jsw = left_index(j, ℓy)
+
+    inw = left_index(i, ℓx)
+    jnw = right_index(j, ℓy)
+
+    ine = right_index(i, ℓx)
+    jne = right_index(j, ℓy)
+
+    ise = right_index(i, ℓx)
+    jse = left_index(j, ℓy)
+
+    xsw = ξnode(isw, jsw, 1, grid, vx, vy, nothing)
+    ysw = ηnode(isw, jsw, 1, grid, vx, vy, nothing)
+
+    xnw = ξnode(inw, jnw, 1, grid, vx, vy, nothing)
+    ynw = ηnode(inw, jnw, 1, grid, vx, vy, nothing)
+
+    xne = ξnode(ine, jne, 1, grid, vx, vy, nothing)
+    yne = ηnode(ine, jne, 1, grid, vx, vy, nothing)
+
+    xse = ξnode(ise, jse, 1, grid, vx, vy, nothing)
+    yse = ηnode(ise, jse, 1, grid, vx, vy, nothing)
+
+    linear_idx = i + (j - 1) * Fx
+    @inbounds begin
+        cell_matrix[1, linear_idx] = (xsw, ysw)
+        cell_matrix[2, linear_idx] = (xnw, ynw)
+        cell_matrix[3, linear_idx] = (xne, yne)
+        cell_matrix[4, linear_idx] = (xse, yse)
+        cell_matrix[5, linear_idx] = (xsw, ysw)
+    end
+end
+
+function Regridder(dst_field::RegriddableField, src_field::RegriddableField)
+    src_cells = compute_cell_matrix(src_field)
+    dst_cells = compute_cell_matrix(dst_field)
+
+    return ConservativeRegridding.Regridder(src_cells, dst_cells)
+end
+
+"""
+    regrid!(a, b)
+
+    regrid!(a, regridder, b)
+
+Regrid field `b` onto the grid of field `a`. 
+
+Example
+=======
+
+Generate a tracer field on a vertically stretched grid and regrid it on a regular grid.
+
+```jldoctest
+using Oceananigans
+
+Nz, Lz = 2, 1.0
+topology = (Flat, Flat, Bounded)
+
+input_grid = RectilinearGrid(size=Nz, z = [0, Lz/3, Lz], topology=topology, halo=1)
+input_field = CenterField(input_grid)
+input_field[1, 1, 1:Nz] = [2, 3]
+
+output_grid = RectilinearGrid(size=Nz, z=(0, Lz), topology=topology, halo=1)
+output_field = CenterField(output_grid)
+
+regrid!(output_field, input_field)
+
+output_field[1, 1, :]
+
+# output
+4-element OffsetArray(::Vector{Float64}, 0:3) with eltype Float64 with indices 0:3:
+ 0.0
+ 2.333333333333333
+ 3.0
+ 0.0
+```
+"""
+function regrid!(dst_field::RegriddableField,
+                 regridder::ConservativeRegridding.Regridder,
+                 src_field::RegriddableField)
+
+    dst_vec = vec(interior(dst_field))
+    src_vec = vec(interior(src_field))
+
+    return regrid!(dst_vec, regridder, src_vec)
+end
+
+function regrid!(dst_field::RegriddableField, src_field::RegriddableField)
+    regridder = ConservativeRegridding.Regridder(dst_field, src_field)
+    return regrid!(dst_field, regridder, src_field)
+end
+
+end # OceananigansConservativeRegriddingExt