fix: keep auxiliary coordinates of dense containers in netcdf round trip

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 <https://github.com/PyPSA/linopy/issues/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 <https://github.com/PyPSA/linopy/issues/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 <https://github.com/PyPSA/linopy/issues/981>`__)
 
 Files from earlier versions load as before.

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
         )

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 = {

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: