Mask area at each timestep for area-weighting

[bradyrx]

In e.g. the following function (and area_weight), the area array should be masked by the data. But it should also be evaluated/masked at each time step in case there is varying gaps in the data.

E.g., area = area.where(ds.notnull().all('time')) as a conservative method.

def cos_weight(da, lat_coord="lat", lon_coord="lon", one_dimensional=True):
    non_spatial = [i for i in get_dims(da) if i not in [lat_coord, lon_coord]]
    filter_dict = {}
    while len(non_spatial) > 0:
        filter_dict.update({non_spatial[0]: 0})
        non_spatial.pop(0)
    if one_dimensional:
        lon, lat = np.meshgrid(da[lon_coord], da[lat_coord])
    else:
        lat = da[lat_coord]
    # NaN out land to not go into area-weighting
    lat = lat.astype("float")
    nan_mask = np.asarray(da.isel(filter_dict).isnull())
    lat[nan_mask] = np.nan
    cos_lat = np.cos(np.deg2rad(lat))
    aw_da = (da * cos_lat).sum(lat_coord).sum(lon_coord) / np.nansum(cos_lat)
    return aw_da

[bradyrx]

Or better yet:

area, ds = xr.broadcast(area, ds)
area = area.where(ds.notnull())

[bradyrx]

Closing since I deprecated area_weight given its inclusion in xarray now..