Add Xarray DataArray values as dimension to another DataArray

Question

I'm trying to complete an analysis using xarray.DataArray.polyfit starting with two DataArrays, both with dimensions x, y, and time. What I want to do is replace the time dimension with unique values of the first array and the matching values of the second array.

So, starting with something like this:

import xarray as xr
import numpy as np

temperature = xr.DataArray(
    np.random.rand(5,5,30),
    coords=dict(x=range(5), y=range(5), time=range(30)),
    dims=('x','y','time')
)

size = xr.DataArray(
    np.random.rand(5,5,30),
    coords=dict(x=range(5), y=range(5), time=range(30)),
    dims=('x','y','time')
)

I want to create a single DataArray with dimensions x, y, and temperature where the temperature dimension would be unique values of temperature (and time would be dropped). Then I could run e.g.new_array.polyfit('temperature', deg=1) and get, for each x/y value, regression coefficients for temperature vs size.

That is, it would look like the result of

new_array = xr.DataArray(
    <reshaped version of size that is 5 x 5 x number of unique values of temperature>,
    coords=dict(x=range(5), y=range(5), temperature=<unique values of temperature>),
    dims=('x','y','temperature')
)

I've tried several things. The closest I've gotten is probably size.expand_dims({"temperature": temperature}) but that doesn't work with >1-dimensional DataArrays.

EDIT: I've considered there may be a lot of unique values of temperature. I could also see binning the values before creating that dimension.

Jesse Anderson · Accepted Answer · 2023-10-27 00:32:08Z

0

The context here is important, as the goal is to use .polyfit. One way may be to use .curvefit instead, which allows using a DataArray as coordinates. So for a 1-d polynomial (linear) fit, we could do

def lr(x, m, b):
  return m * x + b

fit = temperature.fit(size, func=lr, reduce_dims="time")

However I'm guessing that since this is a generic optimization that it is slower than polyfit.

answered Oct 27, 2023 at 0:32

Jesse Anderson

4,62328 silver badges36 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Add Xarray DataArray values as dimension to another DataArray

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related