Efficiently interpolating a small xarray.DataArray into coordinates of a larger array?

I have a large, high-resolution 3D (time: 200, y: 2000, x: 2000) xarray.DataArray similar to this:

import xarray as xr
import numpy as np
import pandas as pd

time = pd.date_range("2019-01-01", "2021-12-30", periods=200)
y_large = np.linspace(-1000000, -1032000, 2000)
x_large = np.linspace(-1968000, -2000000, 2000)
data_large = np.random.randint(low=0, high=10, size=(200, 2000, 2000))

da_large = xr.DataArray(
    data=data_large,
    coords={"time": time, "y": y_large, "x": x_large},
    dims=("time", "y", "x"),
)
da_large

I also have a smaller, low-resolution (time: 200, y: 100, x: 100) xarray.DataArray containing different data, but covering the same x and y extents:

y_small = np.linspace(-1000000, -1032000, 100)
x_small = np.linspace(-1968000, -2000000, 100)
data_small = np.random.randint(low=100, high=110, size=(200, 100, 100))

da_small = xr.DataArray(
    data=data_small,
    coords={"time": time, "y": y_small, "x": x_small},
    dims=("time", "y", "x"),
)
da_small

I need to interpolate my small low-resolution array (da_small) into the higher resolution grid of my larger array (da_large), so that I end up with a time: 200, y: 2000, x: 2000 array containing values resampled from da_small.

I thought I'd be able to do this using xarray's .interp() method, by passing in my higher-res coordinates to sample and interpolate values from da_small into each pixel of da_large:

da_small.interp(x=da_large.x, y=da_large.y, method="linear")

However, my challenge is that this operation causes an extremely large spike in memory, which crashes my kernel. This is a blocker for me, as my actual data can be even larger than this example (up to several thousand pixels tall/wide and several hundred timesteps deep).

My question: How can I perform this kind of operation (re-scaling or interpolating a small array into the grid of a larger array) in a more efficient way, avoiding such a large peak memory usage?

(If possible, I'd prefer a solution compatible with xarray so it can slot into my existing workflows.)

> (If possible, I'd prefer a solution compatible with xarray so it can slot into my existing workflows

So you might consider using xarray with Dask.
That would be better than xarray's .interp() method, which is trying to hold all the computed results in memory

> ## Dask
>
> Xarray is an open source project and Python package that extends the labeled data functionality of Pandas to N-dimensional array-like datasets. It shares a similar API to NumPy and Pandas and supports both Dask and NumPy arrays under the hood.

And, if you can use Dask, then you have access to parallelization or chunking

> Dask arrays are composed of many NumPy (or NumPy-like) arrays.
How these arrays are arranged can significantly affect performance. For example, for a square array you might arrange your chunks along rows, along columns, or in a more square-like fashion. Different arrangements of NumPy arrays will be faster or slower for different algorithms.

Here, I am less interested in raw performance, and more about memory management using chunking, which divides the data into smaller chunks that are more manageable.
You can specify the chunk size according to your memory capacity:

import xarray as xr
import numpy as np
import pandas as pd

# Defining the large high-resolution DataArray
time = pd.date_range("2019-01-01", "2021-12-30", periods=200)
y_large = np.linspace(-1000000, -1032000, 2000)
x_large = np.linspace(-1968000, -2000000, 2000)
data_large = np.random.randint(low=0, high=10, size=(200, 2000, 2000))

da_large = xr.DataArray(
    data=data_large,
    coords={"time": time, "y": y_large, "x": x_large},
    dims=("time", "y", "x"),
)

# Defining the small low-resolution DataArray
y_small = np.linspace(-1000000, -1032000, 100)
x_small = np.linspace(-1968000, -2000000, 100)
data_small = np.random.randint(low=100, high=110, size=(200, 100, 100))

da_small = xr.DataArray(
    data=data_small,
    coords={"time": time, "y": y_small, "x": x_small},
    dims=("time", "y", "x"),
)

# Convert to Dask arrays using chunking
da_large_dask = da_large.chunk({'time': 10, 'y': 500, 'x': 500})
da_small_dask = da_small.chunk({'time': 10, 'y': 25, 'x': 25})

# Perform interpolation operation
da_interp_dask = da_small_dask.interp(x=da_large_dask.x, y=da_large_dask.y, method="linear")

# Compute the result using Dask
with dask.config.set(scheduler='threads'):  # or 'processes', depending on your needs
    da_interp = da_interp_dask.compute()

You would use Dask to divide your data into chunks and performs computations on these chunks instead of the whole array at once.
That should reduce memory usage by only keeping the necessary data in memory at any given time.

The chunk sizes do not directly specify the amount of memory to use, but they do impact memory usage. The larger the chunk, the more memory each computation will need, as more data is loaded at once.

You can change the chunk sizes by modifying the numbers in the dictionary passed to the .chunk() method. For instance, if you want to halve the memory usage, you can reduce each dimension of the chunk size by about sqrt(2) (as memory usage is proportional to the product of the three dimensions):

da_large_dask = da_large.chunk({'time': 7, 'y': 350, 'x': 350})
da_small_dask = da_small.chunk({'time': 7, 'y': 18, 'x': 18})

Or you can use dask.array.core.normalize_chunks to try and calculate chunk sizes based on array size and the memory limit.

Interpolating a small array into a larger grid using xarray's .interp() can indeed lead to significant memory consumption, especially for large datasets. To perform this operation more efficiently and avoid the large memory spike, you can use xarray's .reindex() method along with Dask for lazy computation and chunking. Dask allows you to perform computations on smaller, manageable chunks of the data, reducing memory overhead.

Here's a step-by-step guide on how to achieve this:

1. Install the required libraries if you haven't already:

pip install xarray dask

2. Import the necessary modules:

import xarray as xr
import dask.array as da

3. Convert your DataArrays to Dask arrays using da.from_array():

da_large = da.from_array(da_large)
da_small = da.from_array(da_small)

4. Use da.reindex() to perform the interpolation:

da_interp = da_small.reindex(y=da_large.y, x=da_large.x, method="linear")

5. Convert the resulting Dask array back to an xarray DataArray:

da_interp = xr.DataArray(da_interp, coords={"time": time, "y": da_large.y, "x": da_large.x}, dims=("time", "y", "x"))

Now, da_interp will be a Dask-based xarray DataArray containing the interpolated values from da_small in the grid of da_large. This operation will not immediately consume a large amount of memory since Dask computes the results lazily on smaller chunks, making it more memory-efficient.

If you need to perform further computations on da_interp, remember to use Dask-aware functions, as they can optimize the computation by breaking it down into smaller tasks.

Additionally, you can control the chunk size by specifying the chunks parameter in da.from_array(). For example, setting chunks={"time": 1, "y": 100, "x": 100} will create chunks of size 1 timestep by 100 rows by 100 columns, which can further reduce memory overhead.

By leveraging Dask and xarray together, you can efficiently handle large datasets while maintaining compatibility with xarray's functionality and workflow.