Docs, Examples, and Validation for the XESMF extension

[navidcy]

• Further cleanup of:
  

  Oceananigans.jl/ext/OceananigansXESMFExt.jl

  Lines 12 to 28 in cbfb6f2

  	# 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) = 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_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)

  
  eg use permutedims instead of '
• Does the regridding reproduces the Python xESMF regridding?
• (Perhaps related to the above) Can the integral before and after regridding be better than rtol = 1e-4?
• Documentation/Example for how to do a regridding using XESMF.jl (docstrings with examples should suffice for now; convert some example(s) to doctest(s)).

Closes #4818

[navidcy]

@simone-silvestri, @glwagner:

Could we optimise the for loop for k at:

Oceananigans.jl/ext/OceananigansXESMFExt.jl

Lines 179 to 191 in cbfb6f2

	function regrid!(dst_field, regridder::XESMF.Regridder, src_field)
	Nz = size(src_field.grid)[3]
	topo_z = topology(src_field)[3]()
	ℓz = location(src_field)[3]()
	for k in 1:total_length(ℓz, topo_z, Nz)
	src = vec(interior(src_field, :, :, k))
	dst = vec(interior(dst_field, :, :, k))
	regridder(dst, src)
	end
	return dst_field
	end

?

[glwagner]

@simone-silvestri, @glwagner:

Could we optimise the for loop for k at:

Oceananigans.jl/ext/OceananigansXESMFExt.jl

Lines 179 to 191 in cbfb6f2

	function regrid!(dst_field, regridder::XESMF.Regridder, src_field)
	Nz = size(src_field.grid)[3]
	topo_z = topology(src_field)[3]()
	ℓz = location(src_field)[3]()
	for k in 1:total_length(ℓz, topo_z, Nz)
	src = vec(interior(src_field, :, :, k))
	dst = vec(interior(dst_field, :, :, k))
	regridder(dst, src)
	end
	return dst_field
	end

?

I can't think of anything

[simone-silvestri]

We could leverage batched multiplies from the CUBLAS and BLAS libraries.
Something like this https://github.com/FluxML/NNlib.jl/blob/13d901d130d29662d81bbc6c44c67a4513182002/ext/NNlibCUDAExt/batchedmul.jl#L3

[navidcy]

Related to reproducing python's xESMF results: https://forum.access-hive.org.au/t/conservative-regridding-with-xesmf-puzzled-because-the-integrals-dont-quite-agree/5343

[glwagner]

@juliasloan25 once this PR is in, you will probably be ready to use XESMF for tripolar -> lat lon interpolation. (@navidcy can say more maybe). You might be able to use it now.

[navidcy]

You can use it now! This PR cleans up some pf the code and adds up examples. But doesn’t add functionality.

I wanted to confirm that we reproduce pythons xesmf results. I’m pretty sure we do. Im gonna wrap that up and include it.

bottom line: no need to wait for this PR @juliasloan25

[navidcy]

I confirmed that the regridder here reproduces python's results.

In particular:

using Oceananigans
using XESMF

arch = CPU()

z = (-1, 0)

radius = Oceananigans.defaults.planet_radius

llg_coarse = LatitudeLongitudeGrid(arch; z, radius,
                                   size = (180, 90, 1),
                                   longitude = (0, 360),
                                   latitude = (-90, 90))

llg_fine = LatitudeLongitudeGrid(arch; z, radius,
                                 size = (360, 180, 1),
                                 longitude = (0, 360),
                                 latitude = (-90, 90))

src_field = CenterField(llg_coarse)
dst_field = CenterField(llg_fine)

width = 12         # degrees
set!(src_field, (λ, φ, z) -> exp(-((λ - 150)^2 + (φ - 30)^2) / 2width^2) - 2exp(-((λ - 270)^2 + (φ + 20)^2) / 2width^2))

regridder = XESMF.Regridder(dst_field, src_field, method="conservative")

regrid!(dst_field, regridder, src_field)

println("integral src_field = ", first(Field(Integral(src_field, dims=(1, 2)))))
println("integral dst_field = ", first(Field(Integral(dst_field, dims=(1, 2)))))

gives

integral src_field = -1.1089929769375445e13
integral dst_field = -1.109066523206981e13

And in Python, here's a Jupyter notebook getting exactly the same numbers:

python-xesmf.pdf

[navidcy]

This is ready to review/merge. Could somebody have a look?

[simone-silvestri]

wow, not so conservative after all! (with zstar at #4811 I am stressing over a relative error in tracer conservation of 1e-13 😅, maybe I don't need to be so pedantic...)

[simone-silvestri]

Shall we add the validation somewhere in the XESMF.jl package?

[navidcy]

You mean the Python validation?
Because otherwise it's part of the test. Perhaps not exactly...

I can add both in a validation script.

I'm amazed how it's not giving the same integral to machine precision. But I don't know how xESMF computes their weights. Do they compute them in a way to ensure that the integrals are the same?

[simone-silvestri]

I would have suspected so. In general, if they compute intersections between surface areas it should be quite simple to ensure that the integral is conserved. Maybe they assume the areas are planes rather than spherical sectors? This could introduce an error. If this is the case, a rectilinear grid should lead to exact conservation.

[navidcy]

The grid I tried didn't have overlapping areas. I used a lat-lon grid of 1deg and 2deg.
But I also computed the spherical area of the two grids and I got exactly 4πR^2 both for the source and for the destination grid... So I'm confused.

[navidcy]

I added the validation scripts; I'll merge when CI passes.