perf: matrix accessor rewrite

linopy/__init__.py

@@ -14,8 +14,13 @@
 import linopy.monkey_patch_xarray  # noqa: F401
 from linopy.common import align
 from linopy.config import options
-from linopy.constants import EQUAL, GREATER_EQUAL, LESS_EQUAL
-from linopy.constraints import Constraint, Constraints
+from linopy.constants import EQUAL, GREATER_EQUAL, LESS_EQUAL, PerformanceWarning
+from linopy.constraints import (
+    Constraint,
+    ConstraintBase,
+    Constraints,
+    MutableConstraint,
+)
 from linopy.expressions import LinearExpression, QuadraticExpression, merge
 from linopy.io import read_netcdf
 from linopy.model import Model, Variable, Variables, available_solvers
@@ -30,8 +35,11 @@
 
 __all__ = (
     "Constraint",
+    "ConstraintBase",
     "Constraints",
+    "MutableConstraint",
     "EQUAL",
+    "PerformanceWarning",
     "GREATER_EQUAL",
     "LESS_EQUAL",
     "LinearExpression",

linopy/common.py

@@ -10,7 +10,7 @@
 import operator
 import os
 from collections.abc import Callable, Generator, Hashable, Iterable, Sequence
-from functools import partial, reduce, wraps
+from functools import cached_property, partial, reduce, wraps
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload
 from warnings import warn
@@ -40,7 +40,7 @@
 )
 
 if TYPE_CHECKING:
-    from linopy.constraints import Constraint
+    from linopy.constraints import ConstraintBase
     from linopy.expressions import LinearExpression, QuadraticExpression
     from linopy.variables import Variable
 
@@ -554,7 +554,7 @@ def fill_missing_coords(
     return ds
 
 
-T = TypeVar("T", Dataset, "Variable", "LinearExpression", "Constraint")
+T = TypeVar("T", Dataset, "Variable", "LinearExpression", "ConstraintBase")
 
 
 @overload
@@ -583,10 +583,10 @@ def iterate_slices(
 
 @overload
 def iterate_slices(
-    ds: Constraint,
+    ds: ConstraintBase,
     slice_size: int | None = 10_000,
     slice_dims: list | None = None,
-) -> Generator[Constraint, None, None]: ...
+) -> Generator[ConstraintBase, None, None]: ...
 
 
 def iterate_slices(
@@ -939,6 +939,49 @@ def find_single(value: int) -> tuple[str, dict] | tuple[None, None]:
         raise ValueError("Array's with more than two dimensions is not supported")
 
 
+class VariableLabelIndex:
+    """
+    Index for O(1) mapping between variable labels and dense positions.
+
+    Both arrays are computed lazily and cached:
+    - ``vlabels``: active variable labels in encounter order, shape (n_active_vars,)
+    - ``label_to_pos``: derived from vlabels; size _xCounter, maps label -> position (-1 if masked)
+
+    Invalidated by clearing the instance ``__dict__`` when variables are added or removed.
+    """
+
+    def __init__(self, variables: Any) -> None:
+        self._variables = variables
+
+    @cached_property
+    def vlabels(self) -> np.ndarray:
+        label_lists = []
+        for _, var in self._variables.items():
+            labels = var.labels.values.ravel()
+            mask = labels != -1
+            label_lists.append(labels[mask])
+        return (
+            np.concatenate(label_lists) if label_lists else np.array([], dtype=np.intp)
+        )
+
+    @cached_property
+    def label_to_pos(self) -> np.ndarray:
+        vlabels = self.vlabels
+        n = self._variables.model._xCounter
+        label_to_pos = np.full(n, -1, dtype=np.intp)
+        label_to_pos[vlabels] = np.arange(len(vlabels), dtype=np.intp)
+        return label_to_pos
+
+    @property
+    def n_active_vars(self) -> int:
+        return len(self.vlabels)
+
+    def invalidate(self) -> None:
+        """Clear cached arrays so they are recomputed on next access."""
+        self.__dict__.pop("vlabels", None)
+        self.__dict__.pop("label_to_pos", None)
+
+
 def get_label_position(
     obj: Any,
     values: int | np.ndarray,
@@ -1306,7 +1349,7 @@ def align(
     "Variable",
     "LinearExpression",
     "QuadraticExpression",
-    "Constraint",
+    "ConstraintBase",
 )
 
 
@@ -1358,6 +1401,60 @@ def __call__(self) -> bool:
         return self.value
 
 
+def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]:
+    """
+    Serialize a list of pd.Index (including MultiIndex) to a DataArray dict.
+
+    Suitable for embedding coordinate metadata as plain data variables in a
+    Dataset that has its own unrelated dimensions (e.g. CSR netcdf format).
+    Reconstruct with :func:`coords_from_dataset`.
+    """
+    data_vars: dict[str, DataArray] = {}
+    for c in 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}"],
+                )
+            data_vars[f"_coord_{c.name}_codes"] = DataArray(
+                np.array(c.codes).T,
+                dims=[f"_coorddim_{c.name}", f"_coorddim_{c.name}_nlevels"],
+            )
+        else:
+            data_vars[f"_coord_{c.name}"] = DataArray(
+                np.array(c), dims=[f"_coorddim_{c.name}"]
+            )
+    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`.
+    """
+    coords = []
+    for d in coord_dims:
+        if f"_coord_{d}_codes" in ds:
+            codes_2d = ds[f"_coord_{d}_codes"].values.T
+            level_names = [
+                k[len(f"_coord_{d}_level_") :]
+                for k in ds
+                if k.startswith(f"_coord_{d}_level_")
+            ]
+            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)
+        else:
+            coords.append(pd.Index(ds[f"_coord_{d}"].values, name=d))
+    return coords
+
+
 def is_constant(x: SideLike) -> bool:
     """
     Check if the given object is a constant type or an expression type without

linopy/constants.py

@@ -18,6 +18,11 @@
 GREATER_EQUAL = ">="
 LESS_EQUAL = "<="
 
+
+class PerformanceWarning(UserWarning):
+    """Warning raised when an operation triggers expensive Dataset reconstruction."""
+
+
 long_EQUAL = "=="
 short_GREATER_EQUAL = ">"
 short_LESS_EQUAL = "<"