xarray: заменить значения массива, соответствующие определенным датам в datetime

Question

У меня есть пример массива нулей:

time = np.arange('2000', '2005', dtype='datetime64[D]')
test_array = xr.DataArray(np.zeros(len(time)), coords={'time': time}, dims=['time'])

Теперь, если у меня есть некоторые данные, например test_data = np.ones(365), которые я хочу поместить в массив, соответствующий году 2001 (в котором 365 дней)как мне это сделать?

Я хочу сделать что-то вроде: test_array[test_array.where(time='2001')] = test_data, но .where() здесь не работает.

Answer 1

Функция с тремя аргументами xarray.where может быть более элегантной альтернативой:

import pandas as pd
import xarray as xr

times = pd.date_range('2000', '2002')
da = xr.DataArray(range(len(times)), [('time', times)])
result = xr.where(da.time.dt.year == 2001, 1, da)

Она также работает с массивами значений:

ones = xr.ones_like(da)
result = xr.where(da.time.dt.year == 2001, ones, da)

Если вы начинаете с чистого массива NumPy, вам нужно привести его к DataArray и убедиться, что его временная координата точно совпадает с временной координатой da;если начальная длина массива NumPy отличается от длины da, вам необходимо добавить шаг переиндексации.Вот один из способов сделать это:

import numpy as np

year_2001_times = da.time.sel(time=da.time.dt.year == 2001)
arr = np.random.random(len(year_2001_times))

random_da = xr.DataArray(arr, [('time', year_2001_times)])
reindexed_random_da = random_da.reindex_like(da)

result = xr.where(da.time.dt.year == 2001, reindexed_random_da, da)

Answer 2

Следующее решение работает, но если есть более элегантный способ, я бы хотел знать.

ind_start = (test_array.indexes['time'] == pd.Timestamp('2001-01-01')).argmax()
ind_end = (test_array.indexes['time'] == pd.Timestamp('2001-12-31')).argmax()
test_array[ind_start:ind_end + 1] = test_data