Docs, Examples, and Validation for the XESMF extension

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]
 CUDA = "5.4 - 5.8.2"

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,22 @@ 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 = []
+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, 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,
          plugins = [bib],
-         modules = [Oceananigans, OceananigansNCDatasetsExt],
+         modules,
          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,9 @@ Private = false
 Modules = [Oceananigans.Fields]
 Private = false
 ```
+```@docs
+XESMF.NetCDFWriter
+```
 
 ## Forcings
 

ext/OceananigansXESMFExt.jl

@@ -9,24 +9,28 @@ using Oceananigans.Grids: λnodes, φnodes, Center, Face, total_length
 import Oceananigans.Fields: regrid!
 import XESMF: Regridder, extract_xesmf_coordinates_structure
 
-# 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)
 
+"""
+    extract_xesmf_coordinates_structure(dst_field::AbstractField, src_field::AbstractField)
+
+Extract the coordinates (latitude/longitude) and the coordinates' bounds from the
+`src_field` and `dst_field`.
+"""
 function extract_xesmf_coordinates_structure(dst_field::AbstractField, src_field::AbstractField)
 
     ℓx, ℓy, ℓz = Oceananigans.Fields.instantiated_location(src_field)
@@ -60,15 +64,19 @@ function extract_xesmf_coordinates_structure(dst_field::AbstractField, src_field
     λvˢ = x_vertex_array(λvˢ, Nˢx, Nˢy)
     φvˢ = y_vertex_array(φvˢ, Nˢx, Nˢy)
 
-    dst_coordinates = Dict("lat"   => φᵈ,  # φ is latitude
-                           "lon"   => λᵈ,  # λ is longitude
-                           "lat_b" => φvᵈ,
-                           "lon_b" => λvᵈ)
+    # 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
+
+    dst_coordinates = 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)))
 
-    src_coordinates = Dict("lat"   => φˢ,  # φ is latitude
-                           "lon"   => λˢ,  # λ is longitude
-                           "lat_b" => φvˢ,
-                           "lon_b" => λvˢ)
+    src_coordinates = 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)))
 
     return dst_coordinates, src_coordinates
 end
@@ -99,19 +107,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")
@@ -165,7 +178,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))