fix: write timezone-aware coordinates in Model.to_netcdf

doc/release_notes.rst

@@ -94,6 +94,7 @@ Upcoming Version
 * ``sum()`` over a dimension no longer raises when another dimension of the expression has size 0; it returns an expression without terms over the kept coordinates, as summing over the empty dimension itself already did. (https://github.com/PyPSA/linopy/issues/906)
 * A multi-key ``groupby`` now returns its groups sorted by key tuple, like the single-key path. The key combinations were numbered by iterating a ``set``, so the group order was arbitrary and changed between processes with ``PYTHONHASHSEED``.
 * The ``linopy.options`` context manager now restores the values that were active on entry instead of resetting all options to their defaults.
+* ``Model.to_netcdf`` no longer raises on timezone-aware datetime coordinates (also as a ``MultiIndex`` level, as an auxiliary coordinate, in CSR-backed constraints and in ``Model.parameters``). netCDF has no timezone-aware datetime type, so they are written as naive UTC with the timezone in an attribute, and ``read_netcdf`` restores them.
 * ``Model.copy`` (and the ``copy.copy``/``copy.deepcopy`` protocols) no longer downgrades a quadratic objective to a linear one. The objective was rebuilt as a ``LinearExpression`` regardless of its type, so the copy silently solved a different problem. ``linopy.testing.assert_model_equal`` now compares the objective by expression type as well, which it could not do before. (`#903 <https://github.com/PyPSA/linopy/issues/903>`__)
 * ``Model.to_netcdf``/``linopy.read_netcdf`` downgraded a quadratic objective the same way; the expression type is now stored alongside the objective and restored on read. Files written by earlier versions are read as before. (`#903 <https://github.com/PyPSA/linopy/issues/903>`__)
 * ``linopy.read_netcdf`` now restores variables, expressions and constraints in their original insertion order instead of alphabetically, so ``matrices.A`` of a round-tripped model is no longer a row/column permutation of the original. The order is stored in the file; files written by earlier versions still load in sorted order. ``linopy.testing.assert_model_equal`` now also compares container order. (`#934 <https://github.com/PyPSA/linopy/issues/934>`__)

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,

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:

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]:

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()