def _to_quarter(
pr: Optional[xarray.DataArray] = None,
tas: Optional[xarray.DataArray] = None,
) -> xarray.DataArray:
"""Convert daily, weekly or monthly time series to quarterly time series according to ANUCLIM specifications."""
if tas is not None and pr is not None:
raise ValueError(
"Supply only one variable, 'tas' (exclusive) or 'pr'.")
freq = xarray.infer_freq((tas if tas is not None else pr).time)
if freq is None:
raise ValueError("Can't infer sampling frequency of the input data.")
if freq.upper().startswith("D"):
if tas is not None:
tas = tg_mean(tas, freq="7D")
if pr is not None:
# Accumulate on a week
# Ensure units are back to a "rate" for rate2amount below
pr = convert_units_to(precip_accumulation(pr, freq="7D"), "mm")
pr.attrs["units"] = "mm/week"
freq = "W"
if freq.upper().startswith("W"):
window = 13
elif freq.upper().startswith("M"):
window = 3
else:
raise NotImplementedError(
f'Unknown input time frequency "{freq}": must be one of "D", "W" or "M".'
)
if tas is not None:
tas = ensure_chunk_size(tas, time=np.ceil(window / 2))
out = tas.rolling(time=window, center=False).mean(skipna=False)
out.attrs = tas.attrs
elif pr is not None:
pr = ensure_chunk_size(pr, time=np.ceil(window / 2))
pram = rate2amount(pr)
out = pram.rolling(time=window, center=False).sum()
out.attrs = pr.attrs
out.attrs["units"] = pram.units
else:
raise ValueError("No variables supplied.")
out = ensure_chunk_size(out, time=-1)
return out
def test_ensure_chunk_size():
da = xr.DataArray(np.zeros((20, 21, 20)), dims=("x", "y", "z"))
out = ensure_chunk_size(da, x=10, y=-1)
assert da is out
dac = da.chunk({"x": (1, ) * 20, "y": (10, 10, 1), "z": (10, 10)})
out = ensure_chunk_size(dac, x=3, y=5, z=-1)
assert out.chunks[0] == (3, 3, 3, 3, 3, 5)
assert out.chunks[1] == (10, 11)
assert out.chunks[2] == (20, )
def add_endpoints(
da: xr.DataArray,
left: List[Union[int, float, xr.DataArray, List[int], List[float]]],
right: List[Union[int, float, xr.DataArray, List[int], List[float]]],
dim: str = "quantiles",
):
"""Add left and right endpoints to a DataArray.
Parameters
----------
da : DataArray
Source array.
left : [x, y]
Values to prepend
right : [x, y]
Values to append.
dim : str
Dimension along which to add endpoints.
"""
elems = []
for (x, y) in (left, right):
if isinstance(y, xr.DataArray):
if "quantiles" not in y.dims:
y = y.expand_dims("quantiles")
y = y.assign_coords(quantiles=x)
else:
y = xr.DataArray(y, coords={dim: x}, dims=(dim, ))
elems.append(y)
l, r = elems # pylint: disable=unbalanced-tuple-unpacking
out = xr.concat((l, da, r), dim=dim)
return ensure_chunk_size(out, **{dim: -1})
def add_cyclic_bounds(da: xr.DataArray, att: str, cyclic_coords: bool = True):
"""Reindex an array to include the last slice at the beginning and the first at the end.
This is done to allow interpolation near the end-points.
Parameters
----------
da : Union[xr.DataArray, xr.Dataset]
An array
att : str
The name of the coordinate to make cyclic
cyclic_coords : bool
If True, the coordinates are made cyclic as well,
if False, the new values are guessed using the same step as their neighbour.
Returns
-------
Union[xr.DataArray, xr.Dataset]
da but with the last element along att prepended and the last one appended.
"""
qmf = da.pad({att: (1, 1)}, mode="wrap")
if not cyclic_coords:
vals = qmf.coords[att].values
diff = da.coords[att].diff(att)
vals[0] = vals[1] - diff[0]
vals[-1] = vals[-2] + diff[-1]
qmf = qmf.assign_coords({att: vals})
qmf[att].attrs.update(da.coords[att].attrs)
return ensure_chunk_size(qmf, **{att: -1})
def _to_quarter(
freq: str,
pr: Optional[xarray.DataArray] = None,
tas: Optional[xarray.DataArray] = None,
) -> xarray.DataArray:
"""Convert daily, weekly or monthly time series to quarterly time series according to ANUCLIM specifications."""
if freq.upper().startswith("D"):
if tas is not None:
tas = tg_mean(tas, freq="7D")
if pr is not None:
pr = precip_accumulation(pr, freq="7D")
pr.attrs["units"] = "mm/week"
freq = "W"
if freq.upper().startswith("W"):
window = 13
u = units.week
elif freq.upper().startswith("M"):
window = 3
u = units.month
else:
raise NotImplementedError(
f'Unknown input time frequency "{freq}": must be one of "D", "W" or "M".'
)
if tas is not None:
tas = ensure_chunk_size(tas, time=np.ceil(window / 2))
if pr is not None:
pr = ensure_chunk_size(pr, time=np.ceil(window / 2))
with xarray.set_options(keep_attrs=True):
if pr is not None:
pr = pint_multiply(pr, 1 * u, "mm")
out = pr.rolling(time=window, center=False).sum()
out.attrs = pr.attrs
out.attrs["units"] = "mm"
if tas is not None:
out = tas.rolling(time=window, center=False).mean(skipna=False)
out.attrs = tas.attrs
out = ensure_chunk_size(out, time=-1)
return out
def _to_quarter(
freq: str,
pr: Optional[xarray.DataArray] = None,
tas: Optional[xarray.DataArray] = None,
) -> xarray.DataArray:
"""Convert daily, weekly or monthly time series to quarterly time series according to ANUCLIM specifications."""
if freq.upper().startswith("D"):
if tas is not None:
tas = tg_mean(tas, freq="7D")
if pr is not None:
# Accumulate on a week
# Ensure units are back to a "rate" for rate2amount below
pr = convert_units_to(precip_accumulation(pr, freq="7D"), "mm")
pr.attrs["units"] = "mm/week"
freq = "W"
if freq.upper().startswith("W"):
window = 13
elif freq.upper().startswith("M"):
window = 3
else:
raise NotImplementedError(
f'Unknown input time frequency "{freq}": must be one of "D", "W" or "M".'
)
if tas is not None:
tas = ensure_chunk_size(tas, time=np.ceil(window / 2))
if pr is not None:
pr = ensure_chunk_size(pr, time=np.ceil(window / 2))
if pr is not None:
pram = rate2amount(pr)
out = pram.rolling(time=window, center=False).sum()
out.attrs = pr.attrs
out.attrs["units"] = pram.units
if tas is not None:
out = tas.rolling(time=window, center=False).mean(skipna=False)
out.attrs = tas.attrs
out = ensure_chunk_size(out, time=-1)
return out