diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 993ce7c8..1b7407d6 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -4,6 +4,8 @@ Release Notes Upcoming Version ---------------- +* ``read_netcdf`` no longer drops parameters, parameter attributes or ``MultiIndex`` dimensions whose name contains ``-``. (`#1025 `__) + Version 0.10.0 -------------- 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..35a79b5c 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -15,6 +15,7 @@ from dataclasses import replace from importlib.metadata import version from io import BufferedWriter +from itertools import chain from pathlib import Path from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, Any, TypeVar @@ -1219,7 +1220,7 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: m = Model(sparse=bool(ds.attrs.get("sparse", False))) def has_prefix(k: str, prefix: str) -> bool: - return k.rsplit("-", 1)[0] == prefix + return owner(k) == prefix def remove_prefix(k: str, prefix: str) -> str: return k[len(prefix) + 1 :] @@ -1269,15 +1270,32 @@ def container_names(kind: str) -> list[str]: ) return ordered + containers = { + kind: container_names(kind) + for kind in ("variables", "expressions", "constraints") + } + # Data variables and attributes are stored as ``-``, with + # prefix ``-``, ``objective`` or ``parameters``. Names and keys + # (parameter names, attribute names, dimensions) may contain dashes, so a + # key belongs to the longest known prefix. + prefixes = {"objective", "parameters", *chain.from_iterable(containers.values())} + + def owner(k: str) -> str | None: + parts = k.split("-") + for i in range(len(parts) - 1, 0, -1): + if (prefix := "-".join(parts[:i])) in prefixes: + return prefix + return None + variables = {} - for k in container_names("variables"): + for k in containers["variables"]: name = remove_prefix(k, "variables") variables[name] = Variable(get_prefix(ds, k), m, name) m._variables = Variables(variables, m) expressions: dict[str, LinearExpression | QuadraticExpression] = {} - for k in container_names("expressions"): + for k in containers["expressions"]: name = remove_prefix(k, "expressions") expr_ds = get_prefix(ds, k) expr_type = expr_ds.attrs.pop(EXPR_TYPE_ATTR, None) @@ -1295,7 +1313,7 @@ def container_names(kind: str) -> list[str]: m._expressions = Expressions(expressions, m) constraints: dict[str, ConstraintBase] = {} - for k in container_names("constraints"): + for k in containers["constraints"]: name = remove_prefix(k, "constraints") con_ds = get_prefix(ds, k) if con_ds.attrs.get("_linopy_format") == "csr": diff --git a/test/test_io.py b/test/test_io.py index dabee6f5..fbbbeb0c 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -360,6 +360,44 @@ def test_model_to_netcdf_with_dash_names( assert_model_equal(m, p) +def test_model_to_netcdf_with_dash_named_parameters( + model: Model, tmp_path: Path +) -> None: + m = model + m.parameters["my-param"] = xr.DataArray( + [5, 6], dims=["my-dim"], coords={"my-dim": ["a", "b"]} + ) + m.parameters.attrs["my-attr"] = "kept" + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert_model_equal(m, p) + assert p.parameters.attrs == m.parameters.attrs + + +def test_model_to_netcdf_with_dash_named_multiindex_dim(tmp_path: Path) -> None: + from linopy.semantics import is_v1 + + if is_v1(): + pytest.skip("v1 rejects MultiIndex; this model only builds under legacy") + index = pd.MultiIndex.from_tuples( + [(1, "a"), (1, "b"), (2, "a")], names=["first", "second"] + ) + coords = xr.Coordinates.from_pandas_multiindex(index, "multi-dim") + m = Model() + x = m.add_variables(coords=coords, name="x") + m.add_objective(x.sum()) + m.parameters["param"] = xr.DataArray([1, 2, 3], coords=coords) + fn = tmp_path / "test.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert_model_equal(m, p) + for index in (p.variables["x"].indexes, p.parameters.indexes): + assert isinstance(index["multi-dim"], pd.MultiIndex) + + def test_model_to_netcdf_with_status_and_condition( model_with_dash_names: Model, tmp_path: Path ) -> None: