312 lines
9.5 KiB
Python
312 lines
9.5 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from xarray import CFTimeIndex, DataArray, infer_freq
|
|
from xarray.coding.calendar_ops import convert_calendar, interp_calendar
|
|
from xarray.coding.cftime_offsets import date_range
|
|
from xarray.testing import assert_identical
|
|
from xarray.tests import requires_cftime
|
|
|
|
cftime = pytest.importorskip("cftime")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source, target, use_cftime, freq",
|
|
[
|
|
("standard", "noleap", None, "D"),
|
|
("noleap", "proleptic_gregorian", True, "D"),
|
|
("noleap", "all_leap", None, "D"),
|
|
("all_leap", "proleptic_gregorian", False, "4h"),
|
|
],
|
|
)
|
|
def test_convert_calendar(source, target, use_cftime, freq):
|
|
src = DataArray(
|
|
date_range("2004-01-01", "2004-12-31", freq=freq, calendar=source),
|
|
dims=("time",),
|
|
name="time",
|
|
)
|
|
da_src = DataArray(
|
|
np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
|
|
)
|
|
|
|
conv = convert_calendar(da_src, target, use_cftime=use_cftime)
|
|
|
|
assert conv.time.dt.calendar == target
|
|
|
|
if source != "noleap":
|
|
expected_times = date_range(
|
|
"2004-01-01",
|
|
"2004-12-31",
|
|
freq=freq,
|
|
use_cftime=use_cftime,
|
|
calendar=target,
|
|
)
|
|
else:
|
|
expected_times_pre_leap = date_range(
|
|
"2004-01-01",
|
|
"2004-02-28",
|
|
freq=freq,
|
|
use_cftime=use_cftime,
|
|
calendar=target,
|
|
)
|
|
expected_times_post_leap = date_range(
|
|
"2004-03-01",
|
|
"2004-12-31",
|
|
freq=freq,
|
|
use_cftime=use_cftime,
|
|
calendar=target,
|
|
)
|
|
expected_times = expected_times_pre_leap.append(expected_times_post_leap)
|
|
np.testing.assert_array_equal(conv.time, expected_times)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source,target,freq",
|
|
[
|
|
("standard", "360_day", "D"),
|
|
("360_day", "proleptic_gregorian", "D"),
|
|
("proleptic_gregorian", "360_day", "4h"),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("align_on", ["date", "year"])
|
|
def test_convert_calendar_360_days(source, target, freq, align_on):
|
|
src = DataArray(
|
|
date_range("2004-01-01", "2004-12-30", freq=freq, calendar=source),
|
|
dims=("time",),
|
|
name="time",
|
|
)
|
|
da_src = DataArray(
|
|
np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
|
|
)
|
|
|
|
conv = convert_calendar(da_src, target, align_on=align_on)
|
|
|
|
assert conv.time.dt.calendar == target
|
|
|
|
if align_on == "date":
|
|
np.testing.assert_array_equal(
|
|
conv.time.resample(time="ME").last().dt.day,
|
|
[30, 29, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30],
|
|
)
|
|
elif target == "360_day":
|
|
np.testing.assert_array_equal(
|
|
conv.time.resample(time="ME").last().dt.day,
|
|
[30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 29],
|
|
)
|
|
else:
|
|
np.testing.assert_array_equal(
|
|
conv.time.resample(time="ME").last().dt.day,
|
|
[30, 29, 30, 30, 31, 30, 30, 31, 30, 31, 29, 31],
|
|
)
|
|
if source == "360_day" and align_on == "year":
|
|
assert conv.size == 360 if freq == "D" else 360 * 4
|
|
else:
|
|
assert conv.size == 359 if freq == "D" else 359 * 4
|
|
|
|
|
|
def test_convert_calendar_360_days_random():
|
|
da_std = DataArray(
|
|
np.linspace(0, 1, 366),
|
|
dims=("time",),
|
|
coords={
|
|
"time": date_range(
|
|
"2004-01-01",
|
|
"2004-12-31",
|
|
freq="D",
|
|
calendar="standard",
|
|
use_cftime=False,
|
|
)
|
|
},
|
|
)
|
|
da_360 = DataArray(
|
|
np.linspace(0, 1, 360),
|
|
dims=("time",),
|
|
coords={
|
|
"time": date_range("2004-01-01", "2004-12-30", freq="D", calendar="360_day")
|
|
},
|
|
)
|
|
|
|
conv = convert_calendar(da_std, "360_day", align_on="random")
|
|
conv2 = convert_calendar(da_std, "360_day", align_on="random")
|
|
assert (conv != conv2).any()
|
|
|
|
conv = convert_calendar(da_360, "standard", use_cftime=False, align_on="random")
|
|
assert np.datetime64("2004-02-29") not in conv.time
|
|
conv2 = convert_calendar(da_360, "standard", use_cftime=False, align_on="random")
|
|
assert (conv2 != conv).any()
|
|
|
|
# Ensure that added days are evenly distributed in the 5 fifths of each year
|
|
conv = convert_calendar(da_360, "noleap", align_on="random", missing=np.nan)
|
|
conv = conv.where(conv.isnull(), drop=True)
|
|
nandoys = conv.time.dt.dayofyear[:366]
|
|
assert all(nandoys < np.array([74, 147, 220, 293, 366]))
|
|
assert all(nandoys > np.array([0, 73, 146, 219, 292]))
|
|
|
|
|
|
@requires_cftime
|
|
@pytest.mark.parametrize(
|
|
"source,target,freq",
|
|
[
|
|
("standard", "noleap", "D"),
|
|
("noleap", "proleptic_gregorian", "4h"),
|
|
("noleap", "all_leap", "ME"),
|
|
("360_day", "noleap", "D"),
|
|
("noleap", "360_day", "D"),
|
|
],
|
|
)
|
|
def test_convert_calendar_missing(source, target, freq):
|
|
src = DataArray(
|
|
date_range(
|
|
"2004-01-01",
|
|
"2004-12-31" if source != "360_day" else "2004-12-30",
|
|
freq=freq,
|
|
calendar=source,
|
|
),
|
|
dims=("time",),
|
|
name="time",
|
|
)
|
|
da_src = DataArray(
|
|
np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
|
|
)
|
|
out = convert_calendar(da_src, target, missing=np.nan, align_on="date")
|
|
|
|
expected_freq = freq
|
|
assert infer_freq(out.time) == expected_freq
|
|
|
|
expected = date_range(
|
|
"2004-01-01",
|
|
"2004-12-31" if target != "360_day" else "2004-12-30",
|
|
freq=freq,
|
|
calendar=target,
|
|
)
|
|
np.testing.assert_array_equal(out.time, expected)
|
|
|
|
if freq != "ME":
|
|
out_without_missing = convert_calendar(da_src, target, align_on="date")
|
|
expected_nan = out.isel(time=~out.time.isin(out_without_missing.time))
|
|
assert expected_nan.isnull().all()
|
|
|
|
expected_not_nan = out.sel(time=out_without_missing.time)
|
|
assert_identical(expected_not_nan, out_without_missing)
|
|
|
|
|
|
@requires_cftime
|
|
def test_convert_calendar_errors():
|
|
src_nl = DataArray(
|
|
date_range("0000-01-01", "0000-12-31", freq="D", calendar="noleap"),
|
|
dims=("time",),
|
|
name="time",
|
|
)
|
|
# no align_on for conversion to 360_day
|
|
with pytest.raises(ValueError, match="Argument `align_on` must be specified"):
|
|
convert_calendar(src_nl, "360_day")
|
|
|
|
# Standard doesn't support year 0
|
|
with pytest.raises(
|
|
ValueError, match="Source time coordinate contains dates with year 0"
|
|
):
|
|
convert_calendar(src_nl, "standard")
|
|
|
|
# no align_on for conversion from 360 day
|
|
src_360 = convert_calendar(src_nl, "360_day", align_on="year")
|
|
with pytest.raises(ValueError, match="Argument `align_on` must be specified"):
|
|
convert_calendar(src_360, "noleap")
|
|
|
|
# Datetime objects
|
|
da = DataArray([0, 1, 2], dims=("x",), name="x")
|
|
with pytest.raises(ValueError, match="Coordinate x must contain datetime objects."):
|
|
convert_calendar(da, "standard", dim="x")
|
|
|
|
|
|
def test_convert_calendar_same_calendar():
|
|
src = DataArray(
|
|
date_range("2000-01-01", periods=12, freq="6h", use_cftime=False),
|
|
dims=("time",),
|
|
name="time",
|
|
)
|
|
out = convert_calendar(src, "proleptic_gregorian")
|
|
assert src is out
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"source,target",
|
|
[
|
|
("standard", "noleap"),
|
|
("noleap", "proleptic_gregorian"),
|
|
("standard", "360_day"),
|
|
("360_day", "proleptic_gregorian"),
|
|
("noleap", "all_leap"),
|
|
("360_day", "noleap"),
|
|
],
|
|
)
|
|
def test_interp_calendar(source, target):
|
|
src = DataArray(
|
|
date_range("2004-01-01", "2004-07-30", freq="D", calendar=source),
|
|
dims=("time",),
|
|
name="time",
|
|
)
|
|
tgt = DataArray(
|
|
date_range("2004-01-01", "2004-07-30", freq="D", calendar=target),
|
|
dims=("time",),
|
|
name="time",
|
|
)
|
|
da_src = DataArray(
|
|
np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
|
|
)
|
|
conv = interp_calendar(da_src, tgt)
|
|
|
|
assert_identical(tgt.time, conv.time)
|
|
|
|
np.testing.assert_almost_equal(conv.max(), 1, 2)
|
|
assert conv.min() == 0
|
|
|
|
|
|
@requires_cftime
|
|
def test_interp_calendar_errors():
|
|
src_nl = DataArray(
|
|
[1] * 100,
|
|
dims=("time",),
|
|
coords={
|
|
"time": date_range("0000-01-01", periods=100, freq="MS", calendar="noleap")
|
|
},
|
|
)
|
|
tgt_360 = date_range("0001-01-01", "0001-12-30", freq="MS", calendar="standard")
|
|
|
|
with pytest.raises(
|
|
ValueError, match="Source time coordinate contains dates with year 0"
|
|
):
|
|
interp_calendar(src_nl, tgt_360)
|
|
|
|
da1 = DataArray([0, 1, 2], dims=("x",), name="x")
|
|
da2 = da1 + 1
|
|
|
|
with pytest.raises(
|
|
ValueError, match="Both 'source.x' and 'target' must contain datetime objects."
|
|
):
|
|
interp_calendar(da1, da2, dim="x")
|
|
|
|
|
|
@requires_cftime
|
|
@pytest.mark.parametrize(
|
|
("source_calendar", "target_calendar", "expected_index"),
|
|
[("standard", "noleap", CFTimeIndex), ("all_leap", "standard", pd.DatetimeIndex)],
|
|
)
|
|
def test_convert_calendar_produces_time_index(
|
|
source_calendar, target_calendar, expected_index
|
|
):
|
|
# https://github.com/pydata/xarray/issues/9138
|
|
time = date_range("2000-01-01", "2002-01-01", freq="D", calendar=source_calendar)
|
|
temperature = np.ones(len(time))
|
|
da = DataArray(
|
|
data=temperature,
|
|
dims=["time"],
|
|
coords=dict(
|
|
time=time,
|
|
),
|
|
)
|
|
converted = da.convert_calendar(target_calendar)
|
|
assert isinstance(converted.indexes["time"], expected_index)
|