将所有 xarray 数据集值替换为常量

Question

I have an xarray dataset. I want to make a copy of that so it has the same dimensions/coordinates/shape as the original. That's easy.

import xarray as xr

n_segs = 4
n_dates = 5
num_vars = 4
dims = (n_segs, n_dates)

das = [xr.DataArray(np.random.rand(*dims), dims=['seg_id', 'date'])
       for i in range(num_vars)]

ds_orig = xr.Dataset({'a': das[0], 'b': das[1], 'c': das[2], 'd': das[3]})
ds_copy = ds_orig.copy(deep=True)

Then I want to assign all the values in the copy a constant value (let's say 1). I've figured out how to do this with where :

ds_copy.where(ds_copy == np.nan, other=1)

but this assumes that none of my values will be nan and is a little counter-intuitive IMO. Is there a more robust way?

I suppose I can also loop through the data variables (which is what this suggests for Pandas)...:

for v in ds_copy.data_vars:
    ds_copy[v].loc[:, :] = 1

Maybe what I'm looking for here is a replace method.

Answer 1

I would recommend the loop approach because it will preserve dtypes from the original values. Only one ellipsis in the loc is enough, and the .data_vars can be omitted (datasets have a dictionary interface ):

for v in ds_copy:
    ds_copy[v].loc[:] = 1

To get a more robust version of the where version, you can pass False directly to make sure other will always be used:

ds_copy.where(False, 1)

When storing ints and floats, keeping or not the dtype will probably not have any effect, however, if there are also string or boolean variables, results may change drastically.