Test if xugrid's and iMOD Python's Regridder behave similar enough

[JoerivanEngelen]

It would be good to test iMOD Python's regridder and xugrid's behave similar enough.
If there are any noticeable differences these should be documented.

[Huite]

The most important thing is that the xugrid one is missing the Voxelizer and LayerRegridder, essentially.

This requires top and bottom data, or bounds on a z coordinate. I've implemented the logic for computing the weights: https://github.com/Deltares/xugrid/blob/fc0bdcb0eaeff62887f1e62fadd61957629dff61/xugrid/regrid/overlap_1d.py#L158

But it's not exposed yet here (class exists: https://github.com/Deltares/xugrid/blob/fc0bdcb0eaeff62887f1e62fadd61957629dff61/xugrid/regrid/structured.py#L736)

[JoerivanEngelen]

I just modified some unittests of the legacy regridder to verify if results are the same, which they are. Furthermore I did a simple performance test and it seems xugrid's is faster:

# %%
import imod
import xugrid as xu
import numpy as np

xmin = 0.0
xmax = 1000.0
ymin = 0.0
ymax = 1000.0
dcell_src = 1.0
dcell_dst = 2.5
layer = np.arange(20) + 1

src = imod.util.empty_3d(dcell_src, xmin, xmax, dcell_src, ymin, ymax, layer)
like = imod.util.empty_2d(dcell_dst, xmin, xmax, dcell_dst, ymin, ymax)

rng = np.random.Generator(np.random.PCG64(12345))
src[:, :, :] = rng.random(src.shape)

# %%
dst = imod.prepare.Regridder(method="mean", use_relative_weights=True).regrid(src, like)

# %timeit prints on my machine:
# 1.58 s ± 10.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# %%
dst_new = xu.OverlapRegridder(source=src, target=like, method="mean").regrid(src)

# %timeit prints on my machine:
# 90.9 ms ± 1.38 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[Huite]

That speedup would be satisfying!

Note that both rely on Numba so there may be some JIT-overhead.

The xugrid regridder should be a lot more memory hungry since it's computing and storing all the weights.
On the other hand, the imod regridder recomputes the same weights over and over (just multiplying floats).

[JoerivanEngelen]

For completeness, I also tested this (continuation of previous script), xugrid wins again:

...
# %%
regridder = imod.prepare.Regridder(method="mean", use_relative_weights=True)
# %%
dst = regridder.regrid(src, like)
# %timeit prints on my machine:
# 393 ms ± 5.68 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# %%
regridder_new = xu.OverlapRegridder(source=src, target=like, method="mean")
# %%
dst_new = regridder_new.regrid(src)
# %timeit prints on my machine:
# 15.9 ms ± 285 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

[JoerivanEngelen]

I noticed some slight differences in the location of NaN values when downscaling with the BarycentricInterpolator. It seems the BarycentricInterpolator activates one line of cells more than imod.prepare.Regridder. I prefer this over the previous behavior of the BarycentricInterpolator where it deactivated more cells than the imod.prepare.Regridder. It is probably something to document though.

I see for safety, iMOD Python now masks everything with the smallest possible idomain based on all regridders, so then xu.OverlapRegridder determines what is active or not instead of xu.BarycentricInterpolator. (Granted that xu.CentroidLocator is not used).

# %%
import imod
import xugrid as xu
import numpy as np

xmin = 0.0
xmax = 1000.0
ymin = 0.0
ymax = 1000.0
dcell_fine = 1.0
dcell_coarse = 2.5
layer = np.arange(20) + 1

# %%
## TEST CASE: UPSCALING GRID WITH SQUARE OF NAN CELLS IN CENTER
src = imod.util.empty_3d(dcell_fine, xmin, xmax, dcell_fine, ymin, ymax, layer)
like = imod.util.empty_2d(dcell_coarse, xmin, xmax, dcell_coarse, ymin, ymax)

rng = np.random.Generator(np.random.PCG64(12345))
src[:, :, :] = rng.random(src.shape)
src[:, 400:600, 400:600] = np.nan

# %%
dst = imod.prepare.Regridder(method="mean", use_relative_weights=True).regrid(src, like)
dst_new = xu.OverlapRegridder(source=src, target=like, method="mean").regrid(src)
# Test if nodata values are at the same location
assert not np.any(np.isnan(dst_new) != np.isnan(dst))

# %%
dst_new = xu.BarycentricInterpolator(source=src, target=like).regrid(src)
# Test if nodata values are at the same location
assert not np.any(np.isnan(dst_new) != np.isnan(dst))

# %%
## TEST CASE: DOWNSCALING GRID WITH SQUARE OF NAN CELLS IN CENTER
src = imod.util.empty_3d(dcell_coarse, xmin, xmax, dcell_coarse, ymin, ymax, layer)
like = imod.util.empty_2d(dcell_fine, xmin, xmax, dcell_fine, ymin, ymax)

rng = np.random.Generator(np.random.PCG64(12345))
src[:, :, :] = rng.random(src.shape)
src[:, 150:250, 150:250] = np.nan

# %%
dst = imod.prepare.Regridder(method="multilinear", use_relative_weights=True).regrid(src, like)
dst_new = xu.OverlapRegridder(source=src, target=like, method="mean").regrid(src)
assert not np.any(np.isnan(dst_new) != np.isnan(dst))
# Doesn't fail, so NaN at same location
# %%
dst_new = xu.BarycentricInterpolator(source=src, target=like).regrid(src)
assert not np.any(np.isnan(dst_new) != np.isnan(dst))
# AssertionError, so NaN at different location. Turns out is one single strip of
# cells at the edge of the inactive area 
# %%
# Print number of NaN rows in the regridded arrays
print(len(np.unique(np.nonzero((np.isnan(dst)).values)[1])))
# 248
print(len(np.unique(np.nonzero((np.isnan(dst_new)).values)[1])))
# 250

# So new interpolator has 2 more NaN rows than the new one
# %%
## TEST CASE: DOWNSCALING GRID WITH LINE OF INCREASING VALUES IN CENTER
src = imod.util.empty_3d(dcell_coarse, xmin, xmax, dcell_coarse, ymin, ymax, layer)
like = imod.util.empty_2d(dcell_fine, xmin, xmax, dcell_fine, ymin, ymax)
src[:, 200, :] = np.arange(src.shape[1]) + 1

# %%
dst = imod.prepare.Regridder(method="multilinear", use_relative_weights=True).regrid(src, like)
dst_new = xu.OverlapRegridder(source=src, target=like, method="mean").regrid(src)
assert not np.any(np.isnan(dst_new) != np.isnan(dst))
# %%
dst_new = xu.BarycentricInterpolator(source=src, target=like).regrid(src)
# Print number of non-NaN rows in the regridded arrays
print(np.unique(np.nonzero((~np.isnan(dst)).values)[1]))
# array([500, 501, 502])
print(np.unique(np.nonzero((~np.isnan(dst_new)).values)[1]))
# array([499, 500, 501, 502, 503])

# So new interpolator has 2 more NaN rows than the new one