Docs, Examples, and Validation for the XESMF extension (#4823)

* permutedims only before giving to python

* fix docstring example

* fix docstring example

* fix docstring example

* add docstrings in docs

* try add XESMF.Regridder in Docs

* fix syntax

* fix type

* add XESMF

* it's a Regridder

* simplify

* add docs for OceananigansXESMFExt.regrid!

* cleanup

* cleaner

* two gaussians

* align

* fix docstring

* exponential -> gaussian

* no need for amplitude kwarg

* add validation scripts

---------

Co-authored-by: Simone Silvestri <[email protected]>

Author: navidcy

Project.toml

@@ -101,7 +101,7 @@ StaticArrays = "1"
 Statistics = "1.9"
 StructArrays = "0.4, 0.5, 0.6, 0.7"
 TimesDates = "0.3"
-XESMF = "0.1.4"
+XESMF = "0.1.5"
 julia = "1.9"
 oneAPI = "2.0.1"
 

docs/Project.toml

@@ -13,6 +13,7 @@ Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
 SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"
+XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"
 
 [compat]
 CairoMakie = "0.11, 0.12, 0.13, 0.14"

docs/make.jl

@@ -11,15 +11,16 @@ Distributed.addprocs(2)
     set_theme!(Theme(fontsize=20))
     CairoMakie.activate!(type = "svg")
 
-    using Oceananigans
     using NCDatasets
+    using XESMF
+
+    using Oceananigans
+    using Oceananigans.AbstractOperations
     using Oceananigans.Operators
     using Oceananigans.Diagnostics
     using Oceananigans.OutputWriters
-    using Oceananigans.TurbulenceClosures
     using Oceananigans.TimeSteppers
-    using Oceananigans.AbstractOperations
-
+    using Oceananigans.TurbulenceClosures
     using Oceananigans.BoundaryConditions: Flux, Value, Gradient, Open
 
     bib_filepath = joinpath(dirname(@__FILE__), "oceananigans.bib")
@@ -173,18 +174,20 @@ format = Documenter.HTML(collapselevel = 1,
 
 DocMeta.setdocmeta!(Oceananigans, :DocTestSetup, :(using Oceananigans); recursive=true)
 
-OceananigansNCDatasetsExt = if isdefined(Base, :get_extension)
-    Base.get_extension(Oceananigans, :OceananigansNCDatasetsExt)
-else
-    Oceananigans.OceananigansNCDatasetsExt
+modules = Module[]
+OceananigansNCDatasetsExt = isdefined(Base, :get_extension) ? Base.get_extension(Oceananigans, :OceananigansNCDatasetsExt) : Oceananigans.OceananigansNCDatasetsExt
+OceananigansXESMFExt = isdefined(Base, :get_extension) ? Base.get_extension(Oceananigans, :OceananigansXESMFExt) : Oceananigans.OceananigansXESMFExt
+
+for m in [Oceananigans, XESMF, OceananigansNCDatasetsExt, OceananigansXESMFExt]
+    if !isnothing(m)
+        push!(modules, m)
+    end
 end
 
-makedocs(sitename = "Oceananigans.jl",
+makedocs(; sitename = "Oceananigans.jl",
          authors = "Climate Modeling Alliance and contributors",
-         format = format,
-         pages = pages,
+         format, pages, modules,
          plugins = [bib],
-         modules = [Oceananigans, OceananigansNCDatasetsExt],
          warnonly = [:cross_references],
          draft = false,        # set to true to speed things up
          doctest = true,       # set to false to speed things up

docs/src/appendix/library.md

@@ -71,6 +71,10 @@ Private = false
 Modules = [Oceananigans.Fields]
 Private = false
 ```
+```@docs
+XESMF.Regridder
+OceananigansXESMFExt.regrid!
+```
 
 ## Forcings
 

ext/OceananigansXESMFExt.jl

@@ -9,24 +9,28 @@ using Oceananigans.Grids: AbstractGrid, λnodes, φnodes, Center, Face, total_le
 import Oceananigans.Fields: regrid!
 import XESMF: Regridder
 
-# 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
+node_array(ξ::AbstractMatrix, Nx, Ny) = view(ξ, 1:Nx, 1:Ny)
 
-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::AbstractVector, Nx, Ny) = repeat(view(x, 1:Nx), 1, Ny)
 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::AbstractVector, Nx, Ny) = repeat(transpose(view(y, 1:Ny)), Nx, 1)
 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))
+vertex_array(ξ::AbstractMatrix, Nx, Ny) = view(ξ, 1:Nx+1, 1:Ny+1)
+
+x_vertex_array(x::AbstractVector, Nx, Ny) = repeat(view(x, 1:Nx+1), 1, Ny+1)
 x_vertex_array(x::AbstractMatrix, Nx, Ny) = 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::AbstractVector, Nx, Ny) = repeat(transpose(view(y, 1:Ny+1)), Nx+1, 1)
 y_vertex_array(y::AbstractMatrix, Nx, Ny) = vertex_array(y, Nx, Ny)
 
+"""
+    xesmf_coordinates(grid::AbstractGrid, ℓx, ℓy, ℓz)
+
+Extract the coordinates (latitude/longitude) and the coordinates' bounds from
+`grid` at locations `ℓx, ℓy, ℓz`.
+"""
 function xesmf_coordinates(grid::AbstractGrid, ℓx, ℓy, ℓz)
     Nx, Ny, Nz = size(grid)
 
@@ -44,23 +48,26 @@ function xesmf_coordinates(grid::AbstractGrid, ℓx, ℓy, ℓz)
     λv = x_vertex_array(λv, Nx, Ny)
     φv = y_vertex_array(φv, Nx, Ny)
 
-    return Dict("lat"   => φ,  # φ is latitude
-                "lon"   => λ,  # λ is longitude
-                "lat_b" => φv,
-                "lon_b" => λv)
-end
-
-function xesmf_coordinates(dst_field::AbstractField, src_field::AbstractField)
-
-    ℓx, ℓy, ℓz = Oceananigans.Fields.instantiated_location(src_field)
+    # Python's xESMF expects 2D arrays with (x, y) coordinates
+    # in which y varies in dim=1 and x varies in dim=2
+    # therefore we transpose the coordinate matrices
+    coords_dictionary = Dict("lat"   => permutedims(φ, (2, 1)),  # φ is latitude
+                             "lon"   => permutedims(λ, (2, 1)),  # λ is longitude
+                             "lat_b" => permutedims(φv, (2, 1)),
+                             "lon_b" => permutedims(λv, (2, 1)))
 
-    dst_grid = dst_field.grid
-    src_grid = src_field.grid
+    return coords_dictionary
+end
 
-    dst_coordinates = xesmf_coordinates(dst_grid, ℓx, ℓy, ℓz)
-    src_coordinates = xesmf_coordinates(src_grid, ℓx, ℓy, ℓz)
+"""
+    xesmf_coordinates(field::AbstractField)
 
-    return dst_coordinates, src_coordinates
+Extract the coordinates (latitude/longitude) and the coordinates' bounds from
+the `field`'s grid.
+"""
+function xesmf_coordinates(field::AbstractField)
+    ℓx, ℓy, ℓz = Oceananigans.Fields.instantiated_location(field)
+    return xesmf_coordinates(field.grid, ℓx, ℓy, ℓz)
 end
 
 """
@@ -89,20 +96,24 @@ For more information, see the Python xESMF documentation at:
 Example
 =======
 
-```@example
+To create a regridder for two fields that live on different grids.
+
+```@example regridding
 using Oceananigans
 using XESMF
 
 z = (-1, 0)
-tg = TripolarGrid(; size=(360, 170, 1), z, southernmost_latitude = -80)
-llg = LatitudeLongitudeGrid(; size=(360, 180, 1), z,
+tg = TripolarGrid(; size=(180, 85, 1), z, southernmost_latitude = -80)
+llg = LatitudeLongitudeGrid(; size=(170, 80, 1), z,
                             longitude=(0, 360), latitude=(-82, 90))
 
 src_field = CenterField(tg)
 dst_field = CenterField(llg)
 
-regridder = Oceananigans.Fields.Regridder(dst_field, src_field, method="conservative")
+regridder = XESMF.Regridder(dst_field, src_field, method="conservative")
 ```
+
+We can use the above regridder to regrid via [`regrid!`](@ref).
 """
 function Regridder(dst_field::AbstractField, src_field::AbstractField; method="conservative")
 
@@ -117,7 +128,8 @@ function Regridder(dst_field::AbstractField, src_field::AbstractField; method="c
     dst_Nz = size(dst_field)[3]
     @assert src_field.grid.z.cᵃᵃᶠ[1:src_Nz+1] == dst_field.grid.z.cᵃᵃᶠ[1:dst_Nz+1]
 
-    dst_coordinates, src_coordinates = xesmf_coordinates(dst_field, src_field)
+    dst_coordinates = xesmf_coordinates(dst_field)
+    src_coordinates = xesmf_coordinates(src_field)
     periodic = Oceananigans.Grids.topology(src_field.grid, 1) === Periodic ? true : false
 
     regridder = XESMF.Regridder(src_coordinates, dst_coordinates; method, periodic)
@@ -155,7 +167,7 @@ llg = LatitudeLongitudeGrid(; size=(360, 180, 1), z,
 src_field = CenterField(tg)
 dst_field = CenterField(llg)
 
-λ₀, φ₀ = 150, 30.  # degrees
+λ₀, φ₀ = 150, 30   # degrees
 width = 12         # degrees
 set!(src_field, (λ, φ, z) -> exp(-((λ - λ₀)^2 + (φ - φ₀)^2) / 2width^2))
 

test/test_xesmf.jl

@@ -4,6 +4,8 @@ using XESMF
 using SparseArrays
 using LinearAlgebra
 
+gaussian_bump(λ, φ; λ₀=0, φ₀=0, width=10) = exp(-((λ - λ₀)^2 + (φ - φ₀)^2) / 2width^2)
+
 for arch in archs
     @testset "XESMF extension [$(typeof(arch))]" begin
         @info "Testing XESMF regridding [$(typeof(arch))]..."
@@ -34,9 +36,9 @@ for arch in archs
             src_field = CenterField(src_grid)
             dst_field = CenterField(dst_grid)
 
-            λ₀, φ₀ = 150, 30.  # degrees
             width = 12         # degrees
-            set!(src_field, (λ, φ, z) -> exp(-((λ - λ₀)^2 + (φ - φ₀)^2) / 2width^2))
+            set!(src_field,
+                 (λ, φ, z) -> gaussian_bump(λ, φ; λ₀=150, φ₀=30, width) - 2gaussian_bump(λ, φ; λ₀=270, φ₀=-20, width))
 
             regridder = XESMF.Regridder(dst_field, src_field)
 

validation/regridding/latitude_longitude_regridding.jl

@@ -56,4 +56,3 @@ scatter!(ax5, Array(interior(c_xy, 180, 60, :)), z)
 scatter!(ax6, Array(interior(c_xyz, 180, 60, :)), z′)
 
 display(fig)
-