Save time by reusing regridder

There is an important reason why the regridding is broken into two steps (making the regridder and perform regridding). For high-resolution grids, making the regridder (i.e. “computing regridding weights”, explained later) is quite computationally expensive, but performing regridding on data (“applying regridding weights”) is still pretty fast.

[1]:

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
import xesmf as xe

Prepare data

The grids in previous examples were all quite small and the regridding was almost instantaneous. Let’s try a large-ish grid here.

[2]:

ds_in = xe.util.grid_2d(
    -120, 120, 0.4, -60, 60, 0.3  # longitude range and resolution
)  # latitude range and resolution
ds_in

[3]:

ds_out = xe.util.grid_2d(-120, 120, 0.6, -60, 60, 0.4)
ds_out

Also make a large-ish 4D data, with multiple time frames and vertical levels.

[4]:

ds_in.coords["time"] = np.arange(1, 11)
ds_in.coords["lev"] = np.arange(1, 51)
ds_in["data2D"] = xe.data.wave_smooth(ds_in["lon"], ds_in["lat"])
ds_in["data4D"] = ds_in["time"] * ds_in["lev"] * ds_in["data2D"]
ds_in

It is almost 1GB!

[5]:

ds_in["data4D"].nbytes / 1e9  # Byte -> GB

[5]:

0.96

The data itself is not too interesting… We only focus on performance here.

[6]:

plt.figure(figsize=[12, 3])

plt.subplot(121)
ds_in["data4D"].isel(time=0, lev=0).plot()
plt.title("2D field")

plt.subplot(122)
ds_in["data4D"].mean(dim=["x", "y"]).plot()
plt.title("extra dimensions to test broadcasting")

[6]:

Text(0.5, 1.0, 'extra dimensions to test broadcasting')

../_images/notebooks_Reuse_regridder_12_1.png

Build Regridder

Making a bilinear regridder takes ~7s on my Mac! ('conservative' would take even longer. Try it yourself.)

[7]:

%%time
regridder = xe.Regridder(ds_in, ds_out, 'bilinear')

CPU times: user 4.19 s, sys: 228 ms, total: 4.41 s
Wall time: 4.43 s

[8]:

regridder

[8]:

xESMF Regridder
Regridding algorithm:       bilinear
Input grid shape:           (400, 600)
Output grid shape:          (300, 400)
Output grid dimension name: ('y', 'x')
Periodic in longitude?      False

Apply regridding

However, applying the regridder to 1GB of data only takes ~0.5s

[9]:

%%time
dr_out = regridder(ds_in['data4D'])

CPU times: user 408 ms, sys: 205 ms, total: 613 ms
Wall time: 612 ms

Why applying regridding is so fast?

Most regridding algorithms (including all 5 algorithms in ESMF) are linear, i.e. the output data field is linearly dependent on the input data field. Any linear transform can be viewed as a matrix-vector multiplication \(y = Ax\), where \(A\) is a matrix containing regridding weights, and \(x\), \(y\) are input and output data fields flatten to 1D.

Computing the weight matrix \(A\) is expensive, but \(A\) only depends on input and output grids, not on input data. That means we can use the same \(A\) on different input fields \(x\), as long as the grid structure is not changed.

An xESMF regridder has an attribute weights, i.e. the weight matrix.

[10]:

regridder.weights

[10]:

<120000x240000 sparse matrix of type '<class 'numpy.float64'>'
        with 480000 stored elements in COOrdinate format>

It is typically very sparse, because a single destination point will only receive contribution from a small number of source points.

[11]:

plt.spy(regridder.weights)
plt.xlabel("input grid indices")
plt.ylabel("output grid indices")

[11]:

Text(0, 0.5, 'output grid indices')

../_images/notebooks_Reuse_regridder_26_1.png

Retrieve regridder

To avoid recomputing the same weights later, you can write the weights to disk and reload them later. First write the weights using the to_netcdf method. A default file name will be chosen if none is given.

[12]:

fn = regridder.to_netcdf()
print(fn)

bilinear_400x600_300x400.nc

[13]:

%%bash
du -sh bilinear_400x600_300x400.nc

7.4M    bilinear_400x600_300x400.nc

When you open the notebook next time, simply set the weights argument to the path to the netCDF file. The weight file is typically pretty small (due to sparsity), so reading it is almost instantaneous.

[14]:

%%time
regridder2 = xe.Regridder(ds_in, ds_out, 'bilinear', weights=fn)

CPU times: user 33.7 ms, sys: 2.3 ms, total: 36 ms
Wall time: 33.1 ms

The second-step, applying those weights to data, is just a matrix multiplication \(y=Ax\). With highly-optimized sparse matrix multiplication library, it is blazingly fast.

[15]:

%%time
dr_out2 = regridder2(ds_in['data4D'])

CPU times: user 739 ms, sys: 191 ms, total: 929 ms
Wall time: 926 ms

The retrieved regridder gives the same result as the first regridder.

[16]:

xr.testing.assert_identical(dr_out, dr_out2)  # they are equal

For even larger grids, you might spend several minutes computing the weights. But once they are computed, you don’t have to do it again.