Applying a mask from one 2D xarray (lat, lon) to a 3D xarray (time, lat lon)

[pecos27]

Hi!
Is there a way to apply a masked 2d xarary array (lat and lon) to another 3D array of the same spatial shape (but now there's time, lat, lon)?
I am basically applying a mask based on positive or negative values from a linear trend analysis, basically masking all positive trends/negative trends individually.

The linear trend function isn't that relevant, but basically I would like to apply the 2D (lat, lon) mask made from the linear trend plot to the air temp anomalies (time, lat, lon). Is that possible in xarray?

I first used xarray.where to just isolate any values below 0 (using air temp rasm as reproducible example):

ds = xr.tutorial.open_dataset('rasm').load()

Tair = ds.Tair

Tair_anom = Tair.groupby('time.month') - Tair.groupby('time.month').mean(dim='time')

Tair_anom[2].plot()

# can ignore this part... this is just to get the linear trend plots 
def linear_trend_regression(data):
    msk = np.ma.masked_invalid(data[0])
    arr = np.array(data) #dropna was not working for some reason, workaround is to convert the nans to zeros and then mask
    # convert all nans to zeros and then remove zeros
    where_are_NaNs = np.isnan(arr)
    arr[where_are_NaNs] = 0
    #mask all values where 0 is
    data_mask = np.ma.masked_equal(arr,0) #instead of nans, now there is '--'
    time = range(0, data.shape[0])
    vals2 = data_mask.reshape(len(time), -1)
    regressions = np.polyfit(time, vals2, 1)
    trends = regressions[0,:].reshape(data_mask.shape[1], data_mask.shape[2])
    masked_trends = np.ma.masked_where(np.ma.getmask(msk),trends)
    return xr.DataArray(data=masked_trends,
                            coords=[('y',data.y.values), ('x',data.x.values)],
                            name='trends',
                            attrs=dict(
                                description="Linear trends",
                                units="unit / time"
                            ))

# apply linear trend 
linear_trend = linear_trend_regression(Tair_anom)

#MAKING THRESHOLD (any trend less than zero)
threshold_mask_lower_limit = linear_trend.where(linear_trend < 0)

threshold_mask_lower_limit.plot()

# now when you plot the threshold mask lower limit plot, the top region is masked out 

# question is now how to apply this mask to Tair_anoms 

I originally tried using numpy masking arrays by copying the mask and applying it to the air temps anomaly data array, but then there was a disagreement with the dimension sizes since the air temp anomaly contains the time axis..
Thank you so much!

[dcherian]

Have you tried Tair_anoms.where(linear_trend > 0)? where broadcasts so you can apply a mask with a different number of dimensions.