feat(csr): constraint and export boundaries of the sparse path (PR 2 of #972)

doc/release_notes.rst

@@ -60,7 +60,9 @@ Upcoming Version
 * Multiplying, dividing, adding or subtracting a non-scalar constant (``DataArray``, ``Series``, ``ndarray``, …) keeps a CSR-backed ``LinearExpression`` sparse: the coefficients are scaled per cell and only the constant changes, instead of expanding the dense rectangle. The constant is aligned exactly as on the dense path (NaN, label mismatch, joins and auxiliary coordinates); an operand that adds dimensions still densifies. ``expr.const`` is now served from the sparse backing as well. (`#965 <https://github.com/PyPSA/linopy/issues/965>`__)
 * ``sum(dim=...)`` and a further ``groupby(...).sum()`` keep a CSR-backed ``LinearExpression`` sparse: summed rows are merged in the sparse backing instead of stacking the summed dimensions into the term dimension of the dense rectangle. Auxiliary coordinates, absent cells and the constant follow the dense path; ``groupby(...).sum(sparse=False)`` still densifies. (`#964 <https://github.com/PyPSA/linopy/issues/964>`__)
 * ``where``, ``sel``, ``isel``, ``loc`` and ``[]`` keep a CSR-backed ``LinearExpression`` sparse: the selection is evaluated on the grid's row numbers with the dense (xarray) semantics, including ``drop=``, ``method=``, scalar coordinates left by a scalar selection and auxiliary coordinates, and the selected rows are gathered from the sparse backing; masked cells become absent. A condition or indexer that introduces dimensions still densifies. A merge of differing grids now stays sparse when operands carry auxiliary coordinates, which are checked for conflicts as on the dense path. (`#966 <https://github.com/PyPSA/linopy/issues/966>`__)
-* ``Model.add_constraints(..., mask=..., freeze=True)`` with a CSR-backed left-hand side applies the mask sparsely and returns a ``CSRConstraint`` instead of expanding the dense rectangle. (`#970 <https://github.com/PyPSA/linopy/issues/970>`__)
+* ``Model.add_constraints(..., mask=..., freeze=True)`` with a CSR-backed left-hand side, or with a ``CSRConstraint`` built beforehand, applies the mask sparsely and returns a ``CSRConstraint`` instead of expanding the dense rectangle. An expression on the right-hand side (``lhs <= other_expr`` or ``add_constraints(lhs, "<=", other_expr)``) no longer densifies a CSR-backed left-hand side: it is moved to the left-hand side, its constant becoming the right-hand side, as on the dense path. ``penalty=`` stays unsupported together with ``freeze=True`` and raises. (`#970 <https://github.com/PyPSA/linopy/issues/970>`__)
+* ``LinearExpression.flat`` and ``LinearExpression.to_polars`` on a CSR-backed expression are emitted directly from the sparse backing instead of expanding the dense rectangle; the rows equal those of the dense path, with absent cells and zero coefficients dropped. (`#968 <https://github.com/PyPSA/linopy/issues/968>`__)
+* A CSR-backed linear objective stays sparse when set with ``Model.add_objective`` and when the model is exported or solved: the objective vector ``matrices.c``, LP/MPS files, netcdf output, ``Model.copy``, persistent snapshots and direct solver APIs all read it without expanding the dense rectangle, with results equal to the dense objective. The objective name is now owned by the ``Objective`` itself (``Objective.name`` and ``Objective.attrs``) instead of being written into the expression's attributes. (`#967 <https://github.com/PyPSA/linopy/issues/967>`__)
 * Persistent snapshots of tz-aware ``DatetimeIndex`` coordinates no longer materialise an object array of ``Timestamp`` per container per capture and diff. Coordinates are stored as UTC-ns arrays with the timezone identity carried alongside, making snapshot capture ~24x and warm-start diffs ~33x faster on tz-aware models, while naive and tz-aware coordinates — and differing timezones — stay correctly unequal. (`#960 <https://github.com/PyPSA/linopy/pull/960>`__)
 
 **Bug fixes**
@@ -73,7 +75,9 @@ Upcoming Version
 * ``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>`__)
 * Adding or removing variables after a constraint was added with ``freeze=True`` no longer breaks ``model.matrices``. The frozen constraint stored dense variable positions as its matrix columns, so blocks frozen at different times disagreed on their width and stacking them raised ``ValueError: inconsistent shapes``. Raw variable labels are stored instead and mapped to positions when the matrix is assembled. Frozen constraints in netCDF files written by earlier versions are read as before. (`#926 <https://github.com/PyPSA/linopy/issues/926>`__)
-
+* A ``CSRConstraint`` built from a CSR-backed expression now keeps its auxiliary coordinates, e.g. the key levels of a multi-key ``groupby(..., observed=True).sum(sparse=True)`` or a scalar coordinate left by a scalar selection, so its ``coords`` equal those of the dense constraint. They are kept by ``to_dense()``/``mutable()`` and written to and read from netcdf, and ``Constraint.freeze()`` keeps them as well. (`#941 <https://github.com/PyPSA/linopy/issues/941>`__)
+* ``linopy.read_netcdf`` no longer fails for a whole model when a frozen constraint has a dimension or MultiIndex level whose name contains ``-``: the coordinate metadata of a ``CSRConstraint`` is now stored under positional variable names with the names as attributes. Files written by earlier versions are read as before.
+* ``CSRConstraint.loc``, ``update``, ``from_rule`` and ``soften``, and assigning to ``coeffs``, ``vars`` or ``sign``, now raise an ``AttributeError`` saying the operation is not supported on a frozen constraint and naming ``.mutable()`` as the way out (for ``from_rule``: build with ``Constraint.from_rule`` and ``.freeze()`` the result; for ``soften``: add the constraint with ``freeze=False``), like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#963 <https://github.com/PyPSA/linopy/issues/963>`__)
 * ``linopy.merge`` with ``join="outer"``, ``"left"`` or ``"right"`` no longer raises an xarray ``MergeError`` when only one operand carries an auxiliary coordinate on a joined dimension. The constant was reindexed with a fill value of ``0`` that was also written into the auxiliary coordinate, so it no longer matched the coefficients' copy. The coordinate now follows its rows onto the joined grid and is NaN where no operand provides it, under both semantics. The same holds for ``.add`` / ``.sub`` / ``.mul`` / ``.div`` with a constant and a reindexing ``join=``.
 * A frozen constraint caches the label-to-position mapping of its matrix columns by weak reference and only rebuilds it when the constraint or the set of variables changes. While a persistent snapshot holds the arrays, repeated matrix assembly on an unchanged model returns the same objects, so the snapshot diff can again skip the comparison of untouched frozen constraints by object identity; one-off exports such as ``to_file`` retain no extra memory. (`#933 <https://github.com/PyPSA/linopy/issues/933>`__)
 * ``Solver.close()`` no longer leaves dangling native handles behind. The solver model is now dropped before the environment that owns it, instead of after. And the COPT and MindOpt file interfaces no longer hand back a model they already disposed: after a file-based COPT or MindOpt solve, ``model.solver_model`` is ``None`` rather than a handle into freed memory. (`#899 <https://github.com/PyPSA/linopy/pull/899>`__)

linopy/common.py

@@ -7,6 +7,7 @@
 
 from __future__ import annotations
 
+import json
 import operator
 from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence
 from functools import cached_property, reduce, wraps
@@ -1318,54 +1319,70 @@ def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]:
 
     Suitable for embedding coordinate metadata as plain data variables in a
     Dataset that has its own unrelated dimensions (e.g. CSR netcdf format).
+    The variables are named by position, never by the (possibly dashed)
+    dimension or level names, since the netcdf reader splits variable names
+    on ``-``; the level names of a MultiIndex are stored as a JSON attribute.
     Reconstruct with :func:`coords_from_dataset`.
     """
     data_vars: dict[str, DataArray] = {}
-    for c in coords:
+    for i, c in enumerate(coords):
         if isinstance(c, pd.MultiIndex):
-            for level_name, level_values in zip(c.names, c.levels):
-                data_vars[f"_coord_{c.name}_level_{level_name}"] = DataArray(
-                    np.array(level_values),
-                    dims=[f"_coorddim_{c.name}_level_{level_name}"],
+            for j, level_values in enumerate(c.levels):
+                data_vars[f"_index{i}_level{j}"] = DataArray(
+                    np.array(level_values), dims=[f"_indexdim{i}_level{j}"]
                 )
-            data_vars[f"_coord_{c.name}_codes"] = DataArray(
+            data_vars[f"_index{i}_codes"] = DataArray(
                 np.array(c.codes).T,
-                dims=[f"_coorddim_{c.name}", f"_coorddim_{c.name}_nlevels"],
+                dims=[f"_indexdim{i}", f"_indexdim{i}_nlevels"],
+                attrs={"level_names": json.dumps([str(n) for n in c.names])},
             )
         else:
-            data_vars[f"_coord_{c.name}"] = DataArray(
-                np.array(c), dims=[f"_coorddim_{c.name}"]
-            )
+            data_vars[f"_index{i}"] = DataArray(np.array(c), dims=[f"_indexdim{i}"])
     return data_vars
 
 
 def coords_from_dataset(ds: Dataset, coord_dims: list[str]) -> list[pd.Index]:
     """
     Deserialize a list of pd.Index (including MultiIndex) from a Dataset.
 
-    Reconstructs coordinates previously serialized by :func:`coords_to_dataset_vars`.
+    Reconstructs coordinates previously serialized by :func:`coords_to_dataset_vars`,
+    or by its earlier name-keyed format (``_coord_<dim>``).
     """
     coords = []
-    for d in coord_dims:
-        if f"_coord_{d}_codes" in ds:
-            codes_2d = ds[f"_coord_{d}_codes"].values.T
+    for i, d in enumerate(coord_dims):
+        if f"_index{i}_codes" in ds:
+            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))
+        elif f"_index{i}" in ds:
+            coords.append(pd.Index(ds[f"_index{i}"].values, name=d))
+        elif f"_coord_{d}_codes" in ds:
+            prefix = f"_coord_{d}_level_"
             level_names = [
-                str(k)[len(f"_coord_{d}_level_") :]
-                for k in ds
-                if str(k).startswith(f"_coord_{d}_level_")
+                str(k)[len(prefix) :] for k in ds if str(k).startswith(prefix)
             ]
-            arrays = [
-                ds[f"_coord_{d}_level_{ln}"].values[codes_2d[i]]
-                for i, ln in enumerate(level_names)
-            ]
-            mi = pd.MultiIndex.from_arrays(arrays, names=level_names)
-            mi.name = d
-            coords.append(mi)
+            level_keys = [prefix + ln for ln in level_names]
+            codes_2d = ds[f"_coord_{d}_codes"].values.T
+            coords.append(_multiindex(ds, codes_2d, level_keys, level_names, d))
         else:
             coords.append(pd.Index(ds[f"_coord_{d}"].values, name=d))
     return coords
 
 
+def _multiindex(
+    ds: Dataset,
+    codes_2d: np.ndarray,
+    level_keys: list[str],
+    level_names: list[str],
+    name: str,
+) -> pd.MultiIndex:
+    arrays = [ds[k].values[codes_2d[j]] for j, k in enumerate(level_keys)]
+    mi = pd.MultiIndex.from_arrays(arrays, names=level_names)
+    mi.name = name
+    return mi
+
+
 def is_constant(x: SideLike) -> bool:
     """
     Check if the given object is a constant type or an expression type without