Chunked xarray: load only 1 cell in memory efficiently

Context:

I have a datacube with 3 variables (3D arrays, dims:time,y,x). The datacube is too big to fit in memory so I chunk it with xarray/dask. I want to apply a function to every cell in x,y of every variable in my datacube.

Problem:

My method takes a long time to load only one cell (1 minute) and I have to do that 112200 times. I use a for loop with dataset.variable.isel(x=i, y=j).values to load a single 1D array from my variables. Is there a better way to do that ? Also, knowing my dataset is chunked, is there a way to do that in parallel for all the chunks at once ?

Code example:

# Setup
import xarray as xr
import numpy as np

# Create the dimensions
x = np.linspace(0, 99, 100)
y = np.linspace(0, 349, 350)
time = np.linspace(0, 299, 300)

# Create the dataset
xrds= xr.Dataset()

# Add the dimensions to the dataset
xrds['time'] = time
xrds['y'] = y
xrds['x'] = x



# Create the random data variables with chunking
chunksize = (10, 100, 100)  # Chunk size for the variables
data_var1 = np.random.rand(len(time), len(y), len(x))
data_var2 = np.random.rand(len(time), len(y), len(x))
data_var3 = np.random.rand(len(time), len(y), len(x))

xrds['data_var1'] = (('time', 'y', 'x'), data_var1, {'chunks': chunksize})
xrds['data_var2'] = (('time', 'y', 'x'), data_var2, {'chunks': chunksize})
xrds['data_var3'] = (('time', 'y', 'x'), data_var3, {'chunks': chunksize})

#### ---- My Attempt ---- ####

# Iterate through all the variables in my dataset
for var_name, var_data in xrds.data_vars.items():

    # if variable is 3D
    if var_data.shape == (xrds.dims['time'], xrds.dims['y'], xrds.dims['x']):

        # Iterate through every cell of the variable along the x and y axis only
        for i in range(xrds.dims['y']):
            for j in range(xrds.dims['x']):

                # Load a single 1D cell into memory (len(cell) = len(time))
                print(xrds.v.isel(y=i,x=j).values)

I find that explicitly iterating over the xarray is faster than isel(), by about 10%.

Example:

    for var_name, var_data in xrds.data_vars.items():

        # if variable is 3D
        if var_data.shape == (xrds.dims['time'], xrds.dims['y'], xrds.dims['x']):

            # Iterate through every cell of the variable along the x and y axis only
            for i_array in xrds['data_var1'].transpose('x', 'y', 'time'):
                x_coordinate = i_array.x.item()
                for cell in i_array.transpose('y', 'time'):
                    y_coordinate = cell.y.item()
                    # Do something with cell

This takes 17.38s, versus 20.47s for the original.

PS: The line chunksize = (10, 100, 100) seems very suspicious to me. It seems like if you want to load an array corresponding to the entire time axis at once, the chunks should be set so that this doesn't require looking at multiple chunks. It seems like chunksize = (len(time), 100, 100) would be more efficient. However, I benchmarked this both ways and it doesn't make a difference for this data size. May make a difference on your larger problem.