diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 993ce7c8..1458a8ec 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``/``read_netcdf`` keep the auxiliary coordinates of dense variables, constraints, expressions and ``Model.parameters``, such as the levels on a v1 flat dimension. ``read_netcdf`` now drops only the scalar coordinates xarray attaches from other containers. * ``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/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..a5055d3f 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1230,16 +1230,24 @@ def parse_multiindex_attr(value: str | Iterable[str]) -> list[str]: return [str(n) for n in json.loads(value)] return [str(n) for n in value] + # Selecting a container's data variables attaches the scalar coords of all + # containers, so each is mapped to its owner: the longest matching prefix, + # as names may contain dashes. Coords with dimensions only attach to the + # container whose dimensions they share. + prefixes = {str(k).rsplit("-", 1)[0] for k in ds} + scalar_owner = { + c: max( + (p for p in prefixes if str(c).startswith(f"{p}-")), key=len, default=None + ) + for c in ds.coords + if ds[c].ndim == 0 + } + def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ds = ds[[k for k in ds if has_prefix(str(k), prefix)]] - multiindexes = [] - for dim in ds.dims: - attr = ds.attrs.get(f"{dim}_multiindex") - if attr is None: - continue - for name in parse_multiindex_attr(attr): - multiindexes.append(prefix + "-" + name) - ds = ds.drop_vars(set(ds.coords) - set(ds.dims) - set(multiindexes)) + ds = ds.drop_vars( + [c for c in ds.coords if scalar_owner.get(c, prefix) != prefix] + ) to_rename = set([*ds.dims, *ds.coords, *ds]) ds = ds.rename({d: remove_prefix(d, prefix) for d in to_rename}) ds.attrs = { diff --git a/test/test_io.py b/test/test_io.py index dabee6f5..9b07fe05 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -416,6 +416,32 @@ def test_model_to_netcdf_with_multiindex_scipy_engine( assert_model_equal(m, read_netcdf(fn)) +def test_model_to_netcdf_aux_coords(tmp_path: Path) -> None: + # Flat dimension with levels as auxiliary coordinates (the v1 alternative + # to a MultiIndex), on a variable, a constraint, a named expression and a + # parameter. + m = Model() + coords = xr.Coordinates( + { + "snapshot": np.arange(4), + "time": ("snapshot", pd.date_range("2020-01-01", periods=4, freq="h")), + "period": ("snapshot", [2020, 2020, 2030, 2030]), + } + ) + x = m.add_variables(lower=0, coords=coords, name="x") + m.add_constraints(x >= 1, name="c") + m.add_expressions(2 * x, name="e") + # Scalar coords are attached to every container on read; the dash name + # makes this container's prefix extend the prefix of "c". + m.add_constraints(x.sel(snapshot=0) >= 2, name="c-first") + m.add_objective(x.sum()) + m.parameters["demand"] = xr.DataArray(np.ones(4), coords=coords) + + fn = tmp_path / "test_aux_coords.nc" + m.to_netcdf(fn) + assert_model_equal(m, read_netcdf(fn)) + + def test_model_to_netcdf_with_expressions( model_with_expressions: Model, tmp_path: Path ) -> None: