diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 993ce7c8..bd6a33c6 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -103,6 +103,7 @@ Enable it with ``linopy.options["semantics"] = "v1"``. Legacy stays the default * ``to_netcdf``/``read_netcdf`` keep a quadratic objective quadratic. (`#903 `__) * ``read_netcdf`` restores containers in insertion order, so a round trip no longer permutes ``matrices.A``. ``assert_model_equal`` now compares the order. (`#934 `__) * ``read_netcdf`` no longer fails when a frozen constraint has a dimension or level name with ``-``. +* ``to_netcdf`` no longer raises on timezone-aware datetimes in indexes, ``MultiIndex`` levels, auxiliary coordinates, frozen constraints and ``Model.parameters``. They are written as naive UTC with the zone in an attribute, and ``read_netcdf`` restores them. * ``Model.copy`` keeps frozen constraints frozen; ``deep=True`` copies their arrays and a CSR-backed objective. (`#981 `__) Files from earlier versions load as before. diff --git a/linopy/common.py b/linopy/common.py index 16b1da0c..fe3c89ea 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -1316,6 +1316,91 @@ def __call__(self) -> bool: return self.value +TZ_ATTR = "_linopy_tz" + + +def index_to_naive_utc(idx: pd.Index) -> tuple[pd.Index, list[str | None]]: + """ + Convert the timezone-aware datetime levels of an index to naive UTC. + + netCDF has no timezone-aware datetime type, so these levels are written as + naive UTC. Levels are those of a MultiIndex, or the index itself otherwise. + Reverse with :func:`index_from_naive_utc`. + + Returns + ------- + tuple[pd.Index, list[str | None]] + The converted index and the timezone name per level, None for levels + that were not timezone-aware. + """ + levels = list(idx.levels) if isinstance(idx, pd.MultiIndex) else [idx] + zones = [ + None if getattr(lvl, "tz", None) is None else str(lvl.tz) for lvl in levels + ] + converted = [ + lvl if zone is None else lvl.tz_convert("UTC").tz_localize(None) + for lvl, zone in zip(levels, zones) + ] + return _with_levels(idx, converted), zones + + +def index_from_naive_utc(idx: pd.Index, zones: list[str | None]) -> pd.Index: + """ + Localize the naive UTC levels of an index to their timezones. + + Reverses :func:`index_to_naive_utc`. + """ + levels = list(idx.levels) if isinstance(idx, pd.MultiIndex) else [idx] + converted = [ + lvl if zone is None else lvl.tz_localize("UTC").tz_convert(zone) + for lvl, zone in zip(levels, zones) + ] + return _with_levels(idx, converted) + + +def _with_levels(idx: pd.Index, levels: list[pd.Index]) -> pd.Index: + if not isinstance(idx, pd.MultiIndex): + return levels[0] + # set_levels drops the custom ``name`` attribute of a MultiIndex + new = idx.set_levels(levels) + new.name = getattr(idx, "name", None) + return new + + +def values_to_naive_utc(values: Any) -> tuple[np.ndarray, str | None]: + """ + Convert timezone-aware datetime values of any shape to naive UTC. + + ``values`` holds timezone-aware datetimes if it has a + ``pd.DatetimeTZDtype`` or is an object array of timestamps sharing one + timezone; other values are returned as an array, unconverted. Reverse + with :func:`values_from_naive_utc`. + + Returns + ------- + tuple[np.ndarray, str | None] + The converted values and their timezone name, None if the values + were not timezone-aware. + """ + arr = np.asarray(values) + if arr.dtype != object: + return arr, None + idx = pd.Index(arr.ravel()) + if not isinstance(idx, pd.DatetimeIndex) or idx.tz is None: + return arr, None + naive = idx.tz_convert("UTC").tz_localize(None) + return naive.to_numpy().reshape(arr.shape), str(idx.tz) + + +def values_from_naive_utc(values: np.ndarray, tz: str) -> pd.DatetimeIndex: + """ + Localize naive UTC datetime values to ``tz``, flattened to an index. + + Reverses :func:`values_to_naive_utc`. + """ + return pd.DatetimeIndex(np.ravel(values)).tz_localize("UTC").tz_convert(tz) + + def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]: """ Serialize a list of pd.Index (including MultiIndex) to a DataArray dict. @@ -1329,6 +1414,8 @@ def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]: """ data_vars: dict[str, DataArray] = {} for i, c in enumerate(coords): + c, zones = index_to_naive_utc(c) + tz_attrs = {TZ_ATTR: json.dumps(zones)} if any(zones) else {} if isinstance(c, pd.MultiIndex): for j, level_values in enumerate(c.levels): data_vars[f"_index{i}_level{j}"] = DataArray( @@ -1337,10 +1424,15 @@ def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]: data_vars[f"_index{i}_codes"] = DataArray( np.array(c.codes).T, dims=[f"_indexdim{i}", f"_indexdim{i}_nlevels"], - attrs={"level_names": json.dumps([str(n) for n in c.names])}, + attrs={ + "level_names": json.dumps([str(n) for n in c.names]), + **tz_attrs, + }, ) else: - data_vars[f"_index{i}"] = DataArray(np.array(c), dims=[f"_indexdim{i}"]) + data_vars[f"_index{i}"] = DataArray( + np.array(c), dims=[f"_indexdim{i}"], attrs=tz_attrs + ) return data_vars @@ -1357,9 +1449,11 @@ def coords_from_dataset(ds: Dataset, coord_dims: list[str]) -> list[pd.Index]: codes = ds[f"_index{i}_codes"] level_names = json.loads(codes.attrs["level_names"]) level_keys = [f"_index{i}_level{j}" for j in range(len(level_names))] - coords.append(_multiindex(ds, codes.values.T, level_keys, level_names, d)) + idx = _multiindex(ds, codes.values.T, level_keys, level_names, d) + coords.append(_restore_tz(idx, codes)) elif f"_index{i}" in ds: - coords.append(pd.Index(ds[f"_index{i}"].values, name=d)) + values = ds[f"_index{i}"] + coords.append(_restore_tz(pd.Index(values.values, name=d), values)) elif f"_coord_{d}_codes" in ds: prefix = f"_coord_{d}_level_" level_names = [ @@ -1373,6 +1467,12 @@ def coords_from_dataset(ds: Dataset, coord_dims: list[str]) -> list[pd.Index]: return coords +def _restore_tz(idx: pd.Index, var: DataArray) -> pd.Index: + if TZ_ATTR not in var.attrs: + return idx + return index_from_naive_utc(idx, json.loads(var.attrs[TZ_ATTR])) + + def _multiindex( ds: Dataset, codes_2d: np.ndarray, diff --git a/linopy/csr.py b/linopy/csr.py index 7c27463a..f97629c0 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -39,7 +39,13 @@ import scipy.sparse from xarray import DataArray, Dataset -from linopy.common import coords_from_dataset, coords_to_dataset_vars +from linopy.common import ( + TZ_ATTR, + coords_from_dataset, + coords_to_dataset_vars, + values_from_naive_utc, + values_to_naive_utc, +) from linopy.config import options from linopy.constants import HELPER_DIMS, TERM_DIM, PerformanceWarning from linopy.semantics import ( @@ -99,7 +105,7 @@ def from_dataset(cls, ds: Dataset | DataArray, dims: Iterable[str]) -> Grid: def from_netcdf_vars(cls, ds: Dataset, dims: Iterable[str]) -> Grid: """Read back a grid written by :meth:`to_netcdf_vars`.""" aux: AuxCoords = { - da.attrs["name"]: (da.attrs.get("dim", ()), da.to_numpy()) + da.attrs["name"]: (da.attrs.get("dim", ()), _aux_from_netcdf(da)) for k, da in ds.data_vars.items() if str(k).startswith(_AUX_PREFIX) } @@ -110,14 +116,17 @@ def to_netcdf_vars(self) -> dict[str, DataArray]: The indexes and auxiliary coordinates as plain data variables for netcdf, named by position with the coordinate names as attributes. """ - aux = { - f"{_AUX_PREFIX}{j}": DataArray( - v, + aux = {} + for j, (n, (d, v)) in enumerate(self.aux.items()): + # netCDF has no timezone-aware datetime type, write them as naive UTC + values, zone = values_to_naive_utc(v) + aux[f"{_AUX_PREFIX}{j}"] = DataArray( + values, dims=[f"{_AUX_PREFIX}dim{j}"] if isinstance(d, str) else [], - attrs={"name": n} | ({"dim": d} if isinstance(d, str) else {}), + attrs={"name": n} + | ({"dim": d} if isinstance(d, str) else {}) + | ({TZ_ATTR: zone} if zone is not None else {}), ) - for j, (n, (d, v)) in enumerate(self.aux.items()) - } return coords_to_dataset_vars(self.coords) | aux @property @@ -850,6 +859,18 @@ def _aux_coords(ds: Dataset | DataArray, dims: set[str]) -> AuxCoords: return aux +def _aux_from_netcdf(da: DataArray) -> np.ndarray: + """ + The values of an auxiliary coordinate written by :meth:`Grid.to_netcdf_vars`, + timezone-aware datetimes restored as the object array ``_aux_coords`` holds. + """ + values = da.to_numpy() + if TZ_ATTR not in da.attrs: + return values + aware = values_from_naive_utc(values, da.attrs[TZ_ATTR]) + return np.asarray(aware, dtype=object).reshape(values.shape) + + def _member_rows( grid: Grid, scatter_codes: dict[str, np.ndarray], n_members: int ) -> np.ndarray: diff --git a/linopy/expressions.py b/linopy/expressions.py index 197482ad..e0414483 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1275,7 +1275,9 @@ def _add_constant_legacy( if is_nan_scalar(other): check_user_nan() warn_legacy_absence(self.const) - return self.assign(const=self.const.fillna(0) + other) + return self.assign( + const=self.const.fillna(0) + cast("float | np.number", other) + ) self_const, da, needs_reindex = self._broadcast_and_align( other, fill_value, join ) diff --git a/linopy/io.py b/linopy/io.py index 80568f18..4d1f6c6e 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -11,7 +11,7 @@ import shutil import time import warnings -from collections.abc import Callable, Iterable +from collections.abc import Callable, Hashable, Iterable from dataclasses import replace from importlib.metadata import version from io import BufferedWriter @@ -28,8 +28,13 @@ from linopy import solvers from linopy.common import ( + TZ_ATTR, + index_from_naive_utc, + index_to_naive_utc, sos_weights, to_polars, + values_from_naive_utc, + values_to_naive_utc, ) from linopy.constants import CONCAT_DIM, FACTOR_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR from linopy.objective import Objective, linear_part @@ -1048,6 +1053,12 @@ def non_bool_dict( DEFAULT_NETCDF_COMPRESSION: dict[str, Any] = {"zlib": True, "complevel": 3} +def _assign_index(ds: xr.Dataset, dim: Hashable, idx: pd.Index) -> xr.Dataset: + if isinstance(idx, pd.MultiIndex): + return ds.assign_coords(xr.Coordinates.from_pandas_multiindex(idx, dim)) + return ds.assign_coords({dim: idx}) + + def to_netcdf( m: Model, *args: Any, compression: dict[str, Any] | bool = True, **kwargs: Any ) -> None: @@ -1098,6 +1109,25 @@ def to_netcdf( ) def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: + # netCDF cannot store timezone-aware datetimes, write them as naive UTC + zones_by_name = {} + for dim in ds.dims: + if dim in ds.indexes: + idx, zones = index_to_naive_utc(ds.indexes[dim]) + if any(zones): + zones_by_name[str(dim)] = zones + ds = _assign_index(ds, dim, idx) + for name, var in ds.variables.items(): + if name in ds.indexes: + continue + values, zone = values_to_naive_utc(var.data) + if zone is not None: + zones_by_name[str(name)] = [zone] + naive = {name: var.copy(data=values)} + ds = ds.assign_coords(naive) if name in ds.coords else ds.assign(naive) + if zones_by_name: + ds = ds.assign_attrs({TZ_ATTR: json.dumps(zones_by_name)}) + to_rename = set([*ds.dims, *ds.coords, *ds]) ds = ds.rename({d: f"{prefix}-{d}" for d in to_rename}) ds.attrs = {f"{prefix}-{k}": v for k, v in ds.attrs.items()} @@ -1253,6 +1283,20 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: names = parse_multiindex_attr(ds.attrs.pop(f"{dim}_multiindex")) ds = ds.set_index({dim: names}) # type: ignore[dict-item] + if TZ_ATTR in ds.attrs: + for name, zones in json.loads(ds.attrs.pop(TZ_ATTR)).items(): + if name in ds.indexes: + idx = index_from_naive_utc(ds.indexes[name], zones) + ds = _assign_index(ds, name, idx) + elif name in ds.variables: + var = ds.variables[name] + values = values_from_naive_utc(var.values, zones[0]) + aware = {name: (var.dims, values, var.attrs)} + if name in ds.coords: + ds = ds.assign_coords(aware) + else: + ds = ds.assign(aware) + return ds def container_names(kind: str) -> list[str]: diff --git a/test/test_io.py b/test/test_io.py index dabee6f5..01a4fd60 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -23,7 +23,7 @@ from linopy.constraints import Constraint from linopy.expressions import LinearExpression, QuadraticExpression from linopy.io import CONTAINER_ORDER_ATTR, signed_number -from linopy.testing import assert_exprequal, assert_model_equal +from linopy.testing import assert_conequal, assert_exprequal, assert_model_equal HAS_NETCDF4 = importlib.util.find_spec("netCDF4") is not None @@ -197,6 +197,74 @@ def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None: assert_model_equal(m, p) +@pytest.mark.parametrize( + ("multiindex", "freeze"), [(False, False), (False, True), (True, False)] +) +@pytest.mark.parametrize("tz", ["UTC", "Europe/Amsterdam"]) +def test_model_to_netcdf_tz_aware_coords( + tz: str, multiindex: bool, freeze: bool, tmp_path: Path +) -> None: + # Spans the spring DST change in Europe/Amsterdam (02:00 does not exist) + time = pd.date_range("2020-03-29", periods=4, freq="h", tz=tz, name="time") + m = Model() + if multiindex: + from linopy.semantics import is_v1 + + if is_v1(): + pytest.skip("v1 rejects MultiIndex; this model only builds under legacy") + index = pd.MultiIndex.from_product( + [[2030, 2040], time], names=["period", "time"] + ) + x = m.add_variables(0, pd.Series(10.0, index=index), name="x") + else: + x = m.add_variables(lower=0, coords=[time], name="x") + m.add_constraints(x >= 1, name="c", freeze=freeze) + m.add_objective(x.sum()) + m.parameters = m.parameters.assign(snapshots=("snapshot", time.rename(None))) + assert "snapshots" in m.parameters.data_vars + + fn = tmp_path / "tz.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + assert_model_equal(m, p) + assert p.parameters.equals(m.parameters) + assert str(pd.DatetimeIndex(p.parameters["snapshots"].data).tz) == tz + + for obj in (p.variables["x"], p.constraints["c"]): + idx = obj.indexes[x.dims[0]] + values = idx.get_level_values("time") if multiindex else idx + assert values.unique().equals(time) + assert str(values.tz) == tz + + +@pytest.mark.parametrize("freeze", [False, True]) +@pytest.mark.parametrize("tz", ["UTC", "Europe/Amsterdam"]) +def test_model_to_netcdf_tz_aware_aux_coords( + tz: str, freeze: bool, tmp_path: Path +) -> None: + time = pd.date_range("2020-03-29", periods=4, freq="h", tz=tz) + coords = xr.Coordinates({"snapshot": np.arange(4), "time": ("snapshot", time)}) + m = Model() + x = m.add_variables(lower=0, coords=coords, name="x") + m.add_constraints(x >= 1, name="c", freeze=freeze) + m.add_objective(x.sum()) + + fn = tmp_path / "tz_aux.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + # Dense aux coords are dropped on read until #1024 lands; compare without them + for attr in ("labels", "lower"): + read = getattr(p.variables["x"], attr).drop_vars("time", errors="ignore") + assert read.equals(getattr(m.variables["x"], attr).drop_vars("time")) + assert (p.constraints["c"].rhs.values == m.constraints["c"].rhs.values).all() + if freeze: + assert_conequal(m.constraints["c"], p.constraints["c"]) + aux = pd.DatetimeIndex(p.constraints["c"].data.coords["time"].values) + assert aux.equals(time) + assert str(aux.tz) == tz + + @pytest.mark.skipif(not HAS_NETCDF4, reason="netCDF4 not installed") def test_model_to_netcdf_compression(tmp_path: Path) -> None: m = Model()