How to align xarray dimensions without calculation?

[ConstantinVasilev]

Hi, I already have detailed problem description here so I hope I don't need to re-describe.

It seems like any calculation auto triggers such alignment, however, I'm not able to efficiently obtain the same result with any of the align methods.

Edit: for completeness, here's the description posted on StackOverflow:

I have 2 DataArrays with partially overlapping dimensions. If I perform any calculation the dimensions of the result are auto aligned according to the exact matches between the overlapping dimensions.

How can I achieve the same result without performing the redundant calculation below:y_da.y * x_da.x/x_da.x?

I tried: align, broadcast, merge methods without any success...

import xarray as xr
import pandas as pd

x_da = pd.DataFrame({
    'reporting_date': pd.to_datetime([
        '2024-01-01','2024-02-01','2024-03-01',
        '2024-01-01','2024-02-01','2024-03-01',
        '2024-01-01','2024-02-01','2024-03-01',
    ]),
    'user': [
        'A','A','A',
        'B','B','B',
        'C','C','C',
    ],
    'rating': [
        1,1,1,
        1,2,3,
        3,2,1
    ],
    'x': [
        10,20,30,
        20,30,40,
        30,40,50
    ]
})

y_da = pd.DataFrame({
    'rating': [
        1,1,1,1,1,1,
        2,2,2,2,2,2,
        3,3,3,3,3,3
    ],
    'horizon': [
        1,2,3,4,5,6,
        1,2,3,4,5,6,
        1,2,3,4,5,6
    ],
    'y': [
        0.1,0.2,0.3,0.4,0.5,0.6,
        0.2,0.3,0.4,0.5,0.6,0.7,
        0.3,0.4,0.5,0.6,0.7,0.8,
    ]
})

x_da = x_da.set_index(['user','reporting_date','rating']).to_xarray()
y_da = y_da.set_index(['rating','horizon']).to_xarray()


ds = x_da.assign(y=y_da.y * x_da.x/x_da.x)


display(ds.sel({'user':'A'}).y)

[keewis]

First things first, while I did spend the time to look at the StackOverflow question, I do believe that by deliberately not copying over the actual question (which was not that hard, I just did it for you) you reduce the probability that someone will spend the time to help you – maintainer time is very scarce, and any hurdle you put in the way (however small it may seem) may be one too many.

With that out of the way, I believe what you're missing is that nan divided by nan is nan, not 1, so what you believe to be a no-op besides the broadcasting is, in fact, not a no-op. You can align the results by either using fillna(1) on the result of x/x:

In [6]: y, x = xr.broadcast(y_da, x_da)
   ...: expected1 = y_da["y"] * (x_da["x"] / x_da["x"]).fillna(1)
   ...: xr.testing.assert_equal(y["y"], expected1)

or by masking the broadcasted values:

In [7]: expected2 = y_da["y"] * x_da["x"] / x_da["x"]
   ...: xr.testing.assert_equal(y["y"].where(x["x"].notnull()), expected2)