From 73d54b696f4a712cc565dcc0bce404f19b9a6c64 Mon Sep 17 00:00:00 2001 From: Mayk Thewessen Date: Wed, 7 Oct 2026 13:55:17 +0200 Subject: [PATCH 1/4] fix: write timezone-aware coordinates in Model.to_netcdf netCDF has no timezone-aware datetime type, so to_netcdf raised on tz-aware coords (dimension indexes, MultiIndex levels, CSR-backed constraints and data variables in Model.parameters). Write them as naive UTC with the zone names in a `_linopy_tz` attribute and restore them in read_netcdf. Co-Authored-By: Claude Opus 5.5 --- doc/release_notes.rst | 1 + linopy/common.py | 74 ++++++++++++++++++++++++++++++++++++++++--- linopy/io.py | 38 +++++++++++++++++++++- test/test_io.py | 40 +++++++++++++++++++++++ 4 files changed, 148 insertions(+), 5 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 4c8bb82f1..998c1168c 100644 --- a/doc/release_notes.rst +++ b/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, 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 `__) * ``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 `__) * ``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 `__) diff --git a/linopy/common.py b/linopy/common.py index 16b1da0c2..6fc8ecbcc 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -1316,6 +1316,57 @@ 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 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 +1380,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 +1390,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 +1415,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 +1433,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/io.py b/linopy/io.py index 80568f185..a60c17c0f 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,6 +28,9 @@ from linopy import solvers from linopy.common import ( + TZ_ATTR, + index_from_naive_utc, + index_to_naive_utc, sos_weights, to_polars, ) @@ -1048,6 +1051,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 +1107,22 @@ 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.data_vars.items(): + if isinstance(var.dtype, pd.DatetimeTZDtype): + zones_by_name[str(name)] = [str(var.dtype.tz)] + naive = pd.DatetimeIndex(var.data).tz_convert("UTC").tz_localize(None) + ds = ds.assign({name: var.copy(data=naive.values)}) + 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 +1278,17 @@ 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) + else: + var = ds[name] + utc = pd.DatetimeIndex(var.values).tz_localize("UTC") + tz_values = utc.tz_convert(zones[0]) + ds = ds.assign({name: (var.dims, tz_values, var.attrs)}) + return ds def container_names(kind: str) -> list[str]: diff --git a/test/test_io.py b/test/test_io.py index dabee6f50..17f44dccd 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -197,6 +197,46 @@ 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(p.parameters["snapshots"].dtype.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.skipif(not HAS_NETCDF4, reason="netCDF4 not installed") def test_model_to_netcdf_compression(tmp_path: Path) -> None: m = Model() From 99364baf0fca693f92e2b6faa682ab87adafdcaf Mon Sep 17 00:00:00 2001 From: Mayk Thewessen Date: Wed, 7 Oct 2026 16:05:17 +0200 Subject: [PATCH 2/4] test: read timezone via DatetimeIndex for mypy Co-Authored-By: Claude Opus 5.5 --- test/test_io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_io.py b/test/test_io.py index 17f44dccd..2365e32b8 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -228,7 +228,7 @@ def test_model_to_netcdf_tz_aware_coords( p = read_netcdf(fn) assert_model_equal(m, p) assert p.parameters.equals(m.parameters) - assert str(p.parameters["snapshots"].dtype.tz) == tz + 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]] From 2b1f231a2edcee6126d6c95157024f63395e52b0 Mon Sep 17 00:00:00 2001 From: Mayk Thewessen Date: Wed, 7 Oct 2026 23:17:31 +0200 Subject: [PATCH 3/4] fix: also write timezone-aware auxiliary coordinates Auxiliary coordinates with a timezone-aware dtype (the flat-dimension pattern of v1 semantics) still made to_netcdf raise, in dense datasets and in the aux coords of CSR constraints. Handle every non-index variable through one pair of value helpers, and store CSR aux zones as an attribute on the aux variable. Co-Authored-By: Claude Opus 5.5 --- doc/release_notes.rst | 2 +- linopy/common.py | 34 ++++++++++++++++++++++++++++++++++ linopy/csr.py | 37 +++++++++++++++++++++++++++++-------- linopy/io.py | 28 ++++++++++++++++++---------- test/test_io.py | 28 +++++++++++++++++++++++++++- 5 files changed, 109 insertions(+), 20 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 998c1168c..8a54ebc8f 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -94,7 +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, 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.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 `__) * ``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 `__) * ``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 `__) diff --git a/linopy/common.py b/linopy/common.py index 6fc8ecbcc..fe3c89ea8 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -1367,6 +1367,40 @@ def _with_levels(idx: pd.Index, levels: list[pd.Index]) -> pd.Index: 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. diff --git a/linopy/csr.py b/linopy/csr.py index 7c27463ad..f97629c02 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/io.py b/linopy/io.py index a60c17c0f..4d1f6c6ec 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -33,6 +33,8 @@ 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 @@ -1115,11 +1117,14 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: if any(zones): zones_by_name[str(dim)] = zones ds = _assign_index(ds, dim, idx) - for name, var in ds.data_vars.items(): - if isinstance(var.dtype, pd.DatetimeTZDtype): - zones_by_name[str(name)] = [str(var.dtype.tz)] - naive = pd.DatetimeIndex(var.data).tz_convert("UTC").tz_localize(None) - ds = ds.assign({name: var.copy(data=naive.values)}) + 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)}) @@ -1283,11 +1288,14 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: if name in ds.indexes: idx = index_from_naive_utc(ds.indexes[name], zones) ds = _assign_index(ds, name, idx) - else: - var = ds[name] - utc = pd.DatetimeIndex(var.values).tz_localize("UTC") - tz_values = utc.tz_convert(zones[0]) - ds = ds.assign({name: (var.dims, tz_values, var.attrs)}) + 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 diff --git a/test/test_io.py b/test/test_io.py index 2365e32b8..a885d1b12 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 @@ -237,6 +237,32 @@ def test_model_to_netcdf_tz_aware_coords( 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) + + assert p.variables["x"].labels.equals(m.variables["x"].labels.drop_vars("time")) + assert p.variables["x"].lower.equals(m.variables["x"].lower.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() From a2d28460d0dd6dc8d6ac54f149d27fd15aa439b9 Mon Sep 17 00:00:00 2001 From: Mayk Thewessen Date: Thu, 8 Oct 2026 15:49:57 +0200 Subject: [PATCH 4/4] test: compare dense aux-coord round trip independent of #1024 Co-Authored-By: Claude Opus 5.5 --- test/test_io.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/test_io.py b/test/test_io.py index a885d1b12..01a4fd601 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -253,8 +253,10 @@ def test_model_to_netcdf_tz_aware_aux_coords( m.to_netcdf(fn) p = read_netcdf(fn) - assert p.variables["x"].labels.equals(m.variables["x"].labels.drop_vars("time")) - assert p.variables["x"].lower.equals(m.variables["x"].lower.drop_vars("time")) + # 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"])