From 416d9c637946bf65c354166eec79aee0a233a466 Mon Sep 17 00:00:00 2001 From: Fabian Date: Wed, 23 Sep 2026 17:09:36 +0200 Subject: [PATCH 01/15] feat(csr): emit flat/to_polars directly from the CSR store (#968) --- doc/release_notes.rst | 1 + linopy/expressions.py | 37 ++++++++++++++++++++++++++++--------- test/test_csr.py | 15 +++++++++++++++ 3 files changed, 44 insertions(+), 9 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 53541eff..d4276b30 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -61,6 +61,7 @@ Upcoming Version * ``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 `__) * ``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 `__) * ``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 `__) +* ``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 `__) * 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 `__) **Bug fixes** diff --git a/linopy/expressions.py b/linopy/expressions.py index 0726afc2..61e45428 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -2894,13 +2894,15 @@ def flat(self) -> pd.DataFrame: ------- df : pandas.DataFrame """ - ds = self.data + if self._csr is not None: + df = self._csr_terms(self._csr).to_pandas() + else: - def mask_func(data: dict) -> pd.Series: - mask = (data["vars"] != -1) & (data["coeffs"] != 0) - return mask + def mask_func(data: dict) -> pd.Series: + mask = (data["vars"] != -1) & (data["coeffs"] != 0) + return mask - df = to_dataframe(ds, mask_func=mask_func) + df = to_dataframe(self.data, mask_func=mask_func) df = df.groupby("vars", as_index=False).sum() check_has_nulls(df, name=self.type) return df @@ -2975,18 +2977,35 @@ def to_polars(self) -> pl.DataFrame: ------- df : polars.DataFrame """ - if self.is_constant: + if self._csr is not None: + df = self._csr_terms(self._csr) + elif self.is_constant: df = pl.DataFrame( {"const": self.data["const"].values.reshape(-1)} ).with_columns(pl.lit(None).alias("coeffs"), pl.lit(None).alias("vars")) return df.select(["vars", "coeffs", "const"]) - - df = to_polars(self.data) - df = filter_nulls_polars(df) + else: + df = filter_nulls_polars(to_polars(self.data)) df = maybe_group_terms_polars(df) check_has_nulls_polars(df, name=self.type) return df + def _csr_terms(self, csr: CSRLinearExpression) -> pl.DataFrame: + """ + Stored terms of a CSR backing as ``coeffs``, ``vars`` and ``const`` + columns, dropping absent cells and zero coefficients like the dense + long format. + """ + rows = np.repeat(np.arange(csr.n_cells), np.diff(csr.csr.indptr)) + df = pl.DataFrame( + { + "const": csr.const[rows], + "coeffs": csr.csr.data.astype(float), + "vars": csr.csr.indices.astype(self.model._dtypes["labels"]), + } + ) + return filter_nulls_polars(df.filter(pl.col("const").is_not_nan())) + def simplify(self) -> LinearExpression: """ Simplify the linear expression by combining terms with the same variable. diff --git a/test/test_csr.py b/test/test_csr.py index bb30e86a..295f3ff1 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1397,6 +1397,21 @@ def test_warn_on_densify_names_the_reason(op: str, enabled: bool) -> None: assert notices[0].filename == __file__ +@pytest.mark.parametrize("scale", [1.0, 0.0], ids=["plain", "zeros"]) +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_flat_and_to_polars_served_without_densifying(build: str, scale: float) -> None: + require_v1() + sparse, _ = sparse_and_dense(build) + sparse = scale * sparse + assert sparse._csr is not None + dense = sparse._csr.to_dense() + with no_densify(): + got_pl, got_flat = sparse.to_polars(), sparse.flat + assert sparse.is_sparse + assert got_pl.sort("vars").equals(dense.to_polars().sort("vars")) + pd.testing.assert_frame_equal(got_flat, dense.flat) + + def grid_operand(e: LinearExpression) -> xr.DataArray: """Positive values over the expression's full grid.""" values = np.random.default_rng(1).uniform(1, 2, e.shape[:-1]) From 080fb8f2181f535ac681ed17aafaeee174e89ca9 Mon Sep 17 00:00:00 2001 From: Fabian Date: Wed, 23 Sep 2026 17:17:36 +0200 Subject: [PATCH 02/15] feat(csr): keep a CSR-backed objective sparse through export and solve; Objective owns its name (#967) --- doc/release_notes.rst | 1 + linopy/io.py | 23 ++++----- linopy/matrices.py | 19 ++------ linopy/objective.py | 48 ++++++++++++++++++- linopy/persistent/snapshot.py | 16 +------ test/test_csr.py | 87 +++++++++++++++++++++++++++++------ 6 files changed, 136 insertions(+), 58 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index d4276b30..158f86ae 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -62,6 +62,7 @@ Upcoming Version * ``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 `__) * ``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 `__) * ``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 `__) +* 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 `__) * 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 `__) **Bug fixes** diff --git a/linopy/io.py b/linopy/io.py index 08701ed0..718757e5 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -12,6 +12,7 @@ import time import warnings from collections.abc import Callable, Iterable +from dataclasses import replace from importlib.metadata import version from io import BufferedWriter from pathlib import Path @@ -30,7 +31,7 @@ to_polars, ) from linopy.constants import CONCAT_DIM, FACTOR_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR -from linopy.objective import Objective +from linopy.objective import Objective, linear_part from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup if TYPE_CHECKING: @@ -301,14 +302,7 @@ def objective_to_file( elif m.is_quadratic: df = _scale_objective_dataframe(df, variable_scaling, m.objective.scaling) - linear_terms = df.filter(pl.col("vars1").eq(-1) | pl.col("vars2").eq(-1)) - linear_terms = linear_terms.with_columns( - pl.when(pl.col("vars1").eq(-1)) - .then(pl.col("vars2")) - .otherwise(pl.col("vars1")) - .alias("vars") - ) - objective_write_linear_terms(f, linear_terms, print_variable) + objective_write_linear_terms(f, linear_part(df), print_variable) quads = df.filter(pl.col("vars1").ne(-1) & pl.col("vars2").ne(-1)) objective_write_quadratic_terms(f, quads, print_variable) @@ -1091,8 +1085,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ) for name, expr in m.expressions.items() ] - objective = m.objective.data - objective = objective.assign_attrs( + objective = m.objective.to_netcdf_ds().assign_attrs( sense=m.objective.sense, scaling=m.objective.scaling, **{EXPR_TYPE_ATTR: m.objective.expression.type}, @@ -1399,8 +1392,12 @@ def _copy_con_data(con: ConstraintBase) -> xr.Dataset: new_model, ) - obj_expr = type(m.objective.expression)( - m.objective.expression.data.copy(deep=deep), new_model + expr = m.objective.expression + csr = expr._csr if isinstance(expr, LinearExpression) else None + obj_expr = ( + type(expr)(expr.data.copy(deep=deep), new_model) + if csr is None + else LinearExpression._from_csr(replace(csr, model=new_model), new_model) ) new_model._objective = Objective( obj_expr, new_model, m.objective.sense, m.objective.scaling diff --git a/linopy/matrices.py b/linopy/matrices.py index 9aec46fa..82da4a88 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -174,22 +174,9 @@ def c(self) -> ndarray: label_index = m.variables.label_index label_to_pos = label_index.label_to_pos - expr = m.objective.expression - if isinstance(expr, expressions.QuadraticExpression): - # vars has shape (_factor=2, _term); linear terms have one factor == -1 - vars_2d = expr.data.vars.values # shape (2, n_term) - coeffs_all = expr.data.coeffs.values.ravel() - vars1, vars2 = vars_2d[0], vars_2d[1] - linear = (vars1 == -1) | (vars2 == -1) - var_labels = np.where(vars1[linear] != -1, vars1[linear], vars2[linear]) - coeffs = coeffs_all[linear] - else: - var_labels = expr.data.vars.values.ravel() - coeffs = expr.data.coeffs.values.ravel() - - mask = var_labels != -1 - positions = label_to_pos[var_labels[mask]] - scaled_coeffs = coeffs[mask] / self.var_scaling[positions] + var_labels, coeffs = m.objective.linear_terms() + positions = label_to_pos[var_labels] + scaled_coeffs = coeffs / self.var_scaling[positions] scaled_coeffs = scaled_coeffs * m.objective.scaling np.add.at(result, positions, scaled_coeffs) return result diff --git a/linopy/objective.py b/linopy/objective.py index 6cc2cd6f..2dc6af4d 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -53,6 +53,19 @@ def _objwrap(obj: Objective, *args: Any, **kwargs: Any) -> Objective: return _objwrap +def linear_part(df: pl.DataFrame) -> pl.DataFrame: + """ + Linear terms of a quadratic ``to_polars`` frame, their variable in ``vars``. + """ + linear = df.filter(pl.col("vars1").eq(-1) | pl.col("vars2").eq(-1)) + return linear.with_columns( + pl.when(pl.col("vars1").eq(-1)) + .then(pl.col("vars2")) + .otherwise(pl.col("vars1")) + .alias("vars") + ) + + class Objective: """ An objective expression containing all relevant information. @@ -92,12 +105,19 @@ def __repr__(self) -> str: return f"Objective:\n----------\n{expr_string}\n{sense_string}\n{value_string}" + @property + def name(self) -> str: + """ + Returns the name of the objective, owned by the objective itself. + """ + return "objective" + @property def attrs(self) -> dict[str, Any]: """ Returns the attributes of the objective. """ - return self.expression.attrs + return {"name": self.name} @property def coords(self) -> DatasetCoordinates: @@ -133,6 +153,29 @@ def to_polars(self, **kwargs: Any) -> pl.DataFrame: """ return self.expression.to_polars(**kwargs) + def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: + """ + Returns the variable labels and coefficients of the linear objective + terms, zero coefficients dropped, read from ``to_polars``. + """ + df = self.to_polars() + if self.is_quadratic: + df = linear_part(df) + df = df.drop_nulls("vars").cast({"vars": pl.Int64, "coeffs": pl.Float64}) + return df["vars"].to_numpy(), df["coeffs"].to_numpy() + + def to_netcdf_ds(self) -> Dataset: + """ + Returns the objective in the dense layout for netcdf serialization. + + A sparse objective is expanded without dropping its sparse backing; + after the setter's full sum it holds a single cell. + """ + expr = self.expression + csr = expr._csr if isinstance(expr, expressions.LinearExpression) else None + ds = expr.data if csr is None else csr.to_dense().data + return ds.assign_attrs(name=self.name) + @property def coeffs(self) -> DataArray: """ @@ -197,7 +240,8 @@ def expression( if (expr.const != 0.0) and not np.isnan(expr.const): raise ValueError("Constant values in objective function not supported.") - expr.attrs["name"] = "objective" + if not expr.is_sparse: + expr.attrs["name"] = self.name self._expression = expr @property diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py index 76396d5a..ee09390c 100644 --- a/linopy/persistent/snapshot.py +++ b/linopy/persistent/snapshot.py @@ -7,7 +7,6 @@ import numpy as np import pandas as pd -from linopy import expressions from linopy.constraints import Constraint if TYPE_CHECKING: @@ -38,19 +37,8 @@ def _objective_linear_vector(model: Model) -> np.ndarray: vlabels = model.variables.label_index.vlabels label_to_pos = model.variables.label_index.label_to_pos result = np.zeros(len(vlabels), dtype=np.float64) - expr = model.objective.expression - if isinstance(expr, expressions.QuadraticExpression): - vars_2d = expr.data.vars.values - coeffs_all = expr.data.coeffs.values.ravel() - vars1, vars2 = vars_2d[0], vars_2d[1] - linear = (vars1 == -1) | (vars2 == -1) - var_labels = np.where(vars1[linear] != -1, vars1[linear], vars2[linear]) - coeffs = coeffs_all[linear] - else: - var_labels = expr.data.vars.values.ravel() - coeffs = expr.data.coeffs.values.ravel() - mask = var_labels != -1 - np.add.at(result, label_to_pos[var_labels[mask]], coeffs[mask]) + var_labels, coeffs = model.objective.linear_terms() + np.add.at(result, label_to_pos[var_labels], coeffs) return result diff --git a/test/test_csr.py b/test/test_csr.py index 295f3ff1..acda8774 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -499,6 +499,22 @@ def test_nan_grouper_raises_eagerly() -> None: (1.0 * c.gen_p).groupby(gbus).sum(sparse=True) +TERM_LINE = re.compile(r"^[+-][0-9.e+-]+ x[0-9]+$") + + +def canon_lp(text: str) -> list[str]: + """LP file lines with each block of term lines sorted.""" + out: list[str] = [] + buf: list[str] = [] + for line in text.splitlines(): + if TERM_LINE.match(line): + buf.append(line) + else: + out += sorted(buf) + [line] + buf = [] + return out + sorted(buf) + + def test_lp_files_identical(tmp_path: Path) -> None: require_v1() sizes = (7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1) @@ -510,19 +526,6 @@ def test_lp_files_identical(tmp_path: Path) -> None: ) c2.m.add_objective((1.0 * c2.gen_p).sum()) - term_line = re.compile(r"^[+-][0-9.e+-]+ x[0-9]+$") - - def canon_lp(text: str) -> list[str]: - out: list[str] = [] - buf: list[str] = [] - for line in text.splitlines(): - if term_line.match(line): - buf.append(line) - else: - out += sorted(buf) + [line] - buf = [] - return out + sorted(buf) - f1, f2 = tmp_path / "eager.lp", tmp_path / "sparse.lp" c1.m.to_file(f1) c2.m.to_file(f2) @@ -1412,6 +1415,64 @@ def test_flat_and_to_polars_served_without_densifying(build: str, scale: float) pd.testing.assert_frame_equal(got_flat, dense.flat) +def objective_twins(build: str, scale: float) -> tuple[Model, Model]: + """Twin models, the first with the sparse build as objective, the second dense.""" + sparse, _ = sparse_and_dense(build) + _, dense = sparse_and_dense(build) + with no_densify(): + sparse.model.add_objective(scale * (sparse - sparse.const.fillna(0))) + dense.model.add_objective(scale * (dense - dense.const.fillna(0))) + return sparse.model, dense.model + + +def test_sparse_objective_rejects_constant_without_densifying() -> None: + require_v1() + sparse, _ = sparse_and_dense("absent") + with no_densify(), pytest.raises(ValueError, match="Constant values"): + sparse.model.add_objective(sparse) + + +@pytest.mark.parametrize("scale", [1.0, 0.0], ids=["plain", "zeros"]) +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_objective_stays_csr_and_exports_like_dense( + build: str, scale: float, tmp_path: Path +) -> None: + require_v1() + ms, md = objective_twins(build, scale) + with no_densify(): + c = ms.matrices.c + ms.to_file(tmp_path / "sparse.lp") + ms.to_netcdf(tmp_path / "sparse.nc") + copied = ms.copy() + assert ms.objective.attrs == {"name": "objective"} + repr(ms.objective) + assert ms.objective.expression.is_sparse and copied.objective.expression.is_sparse + md.to_file(tmp_path / "dense.lp") + md.to_netcdf(tmp_path / "dense.nc") + assert np.array_equal(c, md.matrices.c) + lp_sparse, lp_dense = (tmp_path / f"{k}.lp" for k in ("sparse", "dense")) + assert canon_lp(lp_sparse.read_text()) == canon_lp(lp_dense.read_text()) + rs, rd = (linopy.read_netcdf(tmp_path / f"{k}.nc") for k in ("sparse", "dense")) + assert_cells_equal(rs.objective.expression, rd.objective.expression, ()) + assert rs.objective.expression.attrs["name"] == "objective" + + +@pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs") +@pytest.mark.parametrize("io_api", ["lp", "direct"]) +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_sparse_objective_solves_like_dense(build: str, io_api: str) -> None: + require_v1() + values = [] + for m in objective_twins(build, 1.0): + for var in m.variables.data.values(): + var.update(lower=0, upper=1) + m.objective.sense = "max" + with no_densify(): + m.solve("highs", io_api=io_api) + values.append(m.objective.value) + assert values[0] == pytest.approx(values[1]) + + def grid_operand(e: LinearExpression) -> xr.DataArray: """Positive values over the expression's full grid.""" values = np.random.default_rng(1).uniform(1, 2, e.shape[:-1]) From 5faae66500fb92a451c8c6f8d37958901ed09574 Mon Sep 17 00:00:00 2001 From: Fabian Date: Wed, 23 Sep 2026 17:37:32 +0200 Subject: [PATCH 03/15] fix(csr): build sparse flat/to_polars from numpy arrays, no pyarrow needed --- linopy/expressions.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index 61e45428..de8b5f9e 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -2895,7 +2895,7 @@ def flat(self) -> pd.DataFrame: df : pandas.DataFrame """ if self._csr is not None: - df = self._csr_terms(self._csr).to_pandas() + df = pd.DataFrame(self._csr_terms(self._csr)) else: def mask_func(data: dict) -> pd.Series: @@ -2978,7 +2978,7 @@ def to_polars(self) -> pl.DataFrame: df : polars.DataFrame """ if self._csr is not None: - df = self._csr_terms(self._csr) + df = pl.DataFrame(self._csr_terms(self._csr)) elif self.is_constant: df = pl.DataFrame( {"const": self.data["const"].values.reshape(-1)} @@ -2990,21 +2990,21 @@ def to_polars(self) -> pl.DataFrame: check_has_nulls_polars(df, name=self.type) return df - def _csr_terms(self, csr: CSRLinearExpression) -> pl.DataFrame: + def _csr_terms(self, csr: CSRLinearExpression) -> dict[str, np.ndarray]: """ - Stored terms of a CSR backing as ``coeffs``, ``vars`` and ``const`` + Stored terms of a CSR backing as ``const``, ``coeffs`` and ``vars`` columns, dropping absent cells and zero coefficients like the dense long format. """ rows = np.repeat(np.arange(csr.n_cells), np.diff(csr.csr.indptr)) - df = pl.DataFrame( - { - "const": csr.const[rows], - "coeffs": csr.csr.data.astype(float), - "vars": csr.csr.indices.astype(self.model._dtypes["labels"]), - } - ) - return filter_nulls_polars(df.filter(pl.col("const").is_not_nan())) + const = csr.const[rows] + coeffs = csr.csr.data.astype(float) + keep = ~np.isnan(const) & (coeffs != 0) + return { + "const": const[keep], + "coeffs": coeffs[keep], + "vars": csr.csr.indices[keep].astype(self.model._dtypes["labels"]), + } def simplify(self) -> LinearExpression: """ From 8ffdcf18dc22a25284e1d26c7fcfba72e6cce4fe Mon Sep 17 00:00:00 2001 From: Fabian Date: Wed, 23 Sep 2026 17:37:32 +0200 Subject: [PATCH 04/15] feat(csr): aux coords on CSRConstraint, explicit frozen errors, expression rhs and mask on frozen constraints (#941, #963, #970) --- doc/release_notes.rst | 5 +- linopy/constraints.py | 121 ++++++++++++++++++++++++++++-------------- linopy/csr.py | 28 +++++++++- linopy/expressions.py | 3 ++ linopy/model.py | 16 ++---- test/test_csr.py | 116 ++++++++++++++++++++++++++++++++++++++-- 6 files changed, 232 insertions(+), 57 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 158f86ae..b6ebb1c5 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -60,7 +60,7 @@ 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 `__) * ``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 `__) * ``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 `__) -* ``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 `__) +* ``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 `__) * ``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 `__) * 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 `__) * 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 `__) @@ -75,7 +75,8 @@ 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 `__) * ``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 `__) * 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 `__) - +* 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 `__) +* ``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, like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#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 `__) * ``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 `__) diff --git a/linopy/constraints.py b/linopy/constraints.py index fd88ab57..abfbb122 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -17,6 +17,7 @@ TYPE_CHECKING, Any, NamedTuple, + NoReturn, overload, ) from warnings import warn @@ -46,8 +47,6 @@ check_has_nulls, check_has_nulls_polars, contains_labels, - coords_from_dataset, - coords_to_dataset_vars, filter_nulls_polars, format_coord, format_single_constraint, @@ -58,7 +57,6 @@ get_label_position, get_printout_labels, has_optimized_model, - is_constant, iterate_slices, maybe_group_terms_polars, maybe_replace_sign, @@ -592,6 +590,13 @@ def _positional_csr( return csr +_UNSUPPORTED = "is not supported on a frozen constraint" + + +def _frozen_error(attr: str, what: str = "is read-only") -> AttributeError: + return AttributeError(f"CSRConstraint.{attr} {what}; call .mutable() to modify.") + + _PositionalCache = tuple[scipy.sparse.csr_array, np.ndarray, "weakref.ref[np.ndarray]"] @@ -715,7 +720,7 @@ def attrs(self) -> dict[str, Any]: @property def coords(self) -> DatasetCoordinates: - return Dataset(coords=self._grid.indexes).coords + return self._grid.to_dataset().coords @property def dims(self) -> Frozen[Hashable, int]: @@ -785,21 +790,38 @@ def assign_labels( a zero coefficient counts as a term. ``scaling`` is a row scaling over the full flat grid. """ - keep = np.diff(self._csr.indptr) > 0 - positions = self._active_positions[keep] - csr = self._csr[keep] - csr.eliminate_zeros() - changes: dict[str, Any] = dict( - csr=csr, - active_positions=positions, + kept = self._kept(np.diff(self._csr.indptr) > 0) + kept._csr.eliminate_zeros() + changes: dict[str, Any] = dict(cindex=cindex, name=name) + if scaling is not None: + changes["scaling"] = scaling[kept._active_positions] + return kept._replace(**changes) + + def _kept(self, keep: np.ndarray) -> CSRConstraint: + """Copy holding only the active rows where ``keep`` is True.""" + + def rows(values: Any) -> Any: + is_rows = isinstance(values, np.ndarray) and values.ndim + return values[keep] if is_rows else values + + return self._replace( + csr=self._csr[keep], + active_positions=self._active_positions[keep], rhs=self._rhs[keep], - sign=self._sign if isinstance(self._sign, str) else self._sign[keep], - cindex=cindex, - name=name, + sign=rows(self._sign), + scaling=self._scaling[keep], + dual=rows(self._dual), + binvar_labels=rows(self._binvar_labels), + binval=rows(self._binval), ) - if scaling is not None: - changes["scaling"] = scaling[positions] - return self._replace(**changes) + + def masked(self, mask: DataArray) -> CSRConstraint: + """ + Copy with the cells where the boolean ``mask`` is False made inactive, + without the dense rectangle. ``mask`` must lie on the constraint grid. + """ + flat = mask.transpose(*self._grid.dims).to_numpy().reshape(-1) + return self._kept(flat[self._active_positions]) def _assign_coords(self, **coords: Any) -> CSRConstraint: """ @@ -823,7 +845,7 @@ def _active_to_dataarray( ) -> DataArray: full = np.full(self.full_size, fill, dtype=active_values.dtype) full[self.active_positions] = active_values - return DataArray(full.reshape(self.shape), coords=self._grid.coords) + return self._grid.dataarray(full) @property def labels(self) -> DataArray: @@ -843,6 +865,10 @@ def coeffs(self) -> DataArray: ) return self.data.coeffs + @coeffs.setter + def coeffs(self, value: ConstantLike) -> None: + raise _frozen_error("coeffs") + @property def vars(self) -> DataArray: """Get variable labels DataArray, shape (*coord_dims, _term).""" @@ -854,13 +880,21 @@ def vars(self) -> DataArray: ) return self.data.vars + @vars.setter + def vars(self, value: variables.Variable | DataArray) -> None: + raise _frozen_error("vars") + @property def sign(self) -> DataArray: """Get sign DataArray.""" if isinstance(self._sign, str): - return DataArray(np.full(self.shape, self._sign), coords=self._grid.coords) + return self._grid.dataarray(np.full(self.full_size, self._sign)) return self._active_to_dataarray(self._sign, fill="") + @sign.setter + def sign(self, value: SignLike) -> None: + raise _frozen_error("sign") + @property def rhs(self) -> DataArray: """Get RHS DataArray, shape (*coord_dims).""" @@ -868,9 +902,7 @@ def rhs(self) -> DataArray: @rhs.setter def rhs(self, value: ConstantLike) -> None: - raise AttributeError( - "CSRConstraint.rhs is read-only; call .mutable() to modify." - ) + raise _frozen_error("rhs") @property def scaling(self) -> DataArray: @@ -879,9 +911,7 @@ def scaling(self) -> DataArray: @scaling.setter def scaling(self, value: ConstantLike) -> None: - raise AttributeError( - "CSRConstraint.scaling is read-only; call .mutable() to modify." - ) + raise _frozen_error("scaling") @property def lhs(self) -> expressions.LinearExpression: @@ -891,9 +921,21 @@ def lhs(self) -> expressions.LinearExpression: @lhs.setter def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: - raise AttributeError( - "CSRConstraint.lhs is read-only; call .mutable() to modify term structure." - ) + raise _frozen_error("lhs") + + @property + def loc(self) -> LocIndexer: + raise _frozen_error("loc", _UNSUPPORTED) + + def update(self, *args: Any, **kwargs: Any) -> NoReturn: + raise _frozen_error("update", _UNSUPPORTED) + + def soften(self, *args: Any, **kwargs: Any) -> NoReturn: + raise _frozen_error("soften", _UNSUPPORTED) + + @classmethod + def from_rule(cls, *args: Any, **kwargs: Any) -> NoReturn: + raise _frozen_error("from_rule", _UNSUPPORTED) @property @has_optimized_model @@ -951,7 +993,7 @@ def _to_dataset(self, nterm: int) -> Dataset: ) dim_names = self.coord_names - xr_coords = self._grid.indexes + xr_coords = self._grid.indexes | self._grid.aux dims_with_term = dim_names + [TERM_DIM] coeffs_da = DataArray( coeffs_2d.reshape(shape + (nterm,)), @@ -969,7 +1011,7 @@ def _to_dataset(self, nterm: int) -> Dataset: labels_flat[active_positions] = self.active_labels() ds = assign_multiindex_safe( ds, - labels=DataArray(labels_flat.reshape(shape), coords=self._grid.coords), + labels=self._grid.dataarray(labels_flat), ) return ds @@ -1100,7 +1142,7 @@ def to_netcdf_ds(self) -> Dataset: } if isinstance(self._sign, np.ndarray): data_vars["_sign"] = DataArray(self._sign, dims=["_flat"]) - data_vars.update(coords_to_dataset_vars(self._grid.coords)) + data_vars.update(self._grid.to_netcdf_vars()) if self._dual is not None: data_vars["dual"] = DataArray(self._dual, dims=["_flat"]) dim_names = list(self._grid.dims) @@ -1156,7 +1198,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: coord_dims = attrs["coord_dims"] if isinstance(coord_dims, str): coord_dims = [coord_dims] - grid = Grid.from_coords(coords_from_dataset(ds, coord_dims)) + grid = Grid.from_netcdf_vars(ds, coord_dims) dual = ds["dual"].values if "dual" in ds else None if "_active_positions" in ds: active_positions = ds["_active_positions"].values @@ -1358,7 +1400,7 @@ def from_dense( ) csr.sum_duplicates() csr.eliminate_zeros() - grid = Grid.from_coords(con.indexes[d] for d in con.coord_dims) + grid = Grid.from_dataset(con.data, map(str, con.coord_dims)) rhs = con.rhs.values.ravel()[active_mask] scaling = con.scaling.values.ravel()[active_mask] sign_vals = con.sign.values.ravel() @@ -1420,20 +1462,19 @@ def from_csr( active, rhs_flat[active], sign, - grid=Grid(expr.grid.indexes), + grid=expr.grid, model=expr.model, ) def csr_rhs(expr: CSRLinearExpression, rhs: Any) -> DataArray | str: """ - Return ``rhs`` as a DataArray on the expression grid, or the reason the - sparse path cannot take it: a non-constant rhs, one that is no - DataArray-like, or one with helper dims or dims outside the grid falls - back to the dense path. + Return the constant ``rhs`` as a DataArray on the expression grid, or the + reason the sparse path cannot take it: an rhs that is no DataArray-like, + or one with helper dims or dims outside the grid falls back to the dense + path. An expression rhs is moved to the lhs before, see + :meth:`LinearExpression.to_constraint`. """ - if not is_constant(rhs): - return "constraint with a non-constant rhs" try: da = as_dataarray(rhs) except (TypeError, ValueError): diff --git a/linopy/csr.py b/linopy/csr.py index 1a24b2da..72fb8e5c 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -38,6 +38,7 @@ import scipy.sparse from xarray import DataArray, Dataset +from linopy.common import coords_from_dataset, coords_to_dataset_vars from linopy.config import options from linopy.constants import HELPER_DIMS, TERM_DIM, PerformanceWarning from linopy.semantics import ( @@ -56,6 +57,9 @@ AuxCoords: TypeAlias = dict[str, tuple[str | tuple[()], np.ndarray]] """Auxiliary coordinates as ``name -> (grid dim, values)``, dim ``()`` for a scalar.""" +_AUX_PREFIX = "_aux_" +_COORDDIM_PREFIX = "_coorddim_" + @dataclass(frozen=True, eq=False) class Grid: @@ -85,6 +89,28 @@ def from_dataset(cls, ds: Dataset | DataArray, dims: Iterable[str]) -> Grid: indexes = {d: ds.get_index(d).rename(d) for d in dims} return cls(indexes, _aux_coords(ds, set(dims))) + @classmethod + def from_netcdf_vars(cls, ds: Dataset, dims: Iterable[str]) -> Grid: + """Read back a grid written by :meth:`to_netcdf_vars`.""" + aux: AuxCoords = {} + for k in ds: + name = str(k) + if name.startswith(_AUX_PREFIX): + da = ds[k] + d = str(da.dims[0]).removeprefix(_COORDDIM_PREFIX) if da.ndim else () + aux[name.removeprefix(_AUX_PREFIX)] = (d, da.to_numpy()) + return cls(cls.from_coords(coords_from_dataset(ds, list(dims))).indexes, aux) + + def to_netcdf_vars(self) -> dict[str, DataArray]: + """The indexes and auxiliary coordinates as plain data variables for netcdf.""" + aux = { + f"{_AUX_PREFIX}{n}": DataArray( + v, dims=[f"{_COORDDIM_PREFIX}{d}"] if isinstance(d, str) else [] + ) + for n, (d, v) in self.aux.items() + } + return coords_to_dataset_vars(self.coords) | aux + @property def dims(self) -> tuple[str, ...]: return tuple(self.indexes) @@ -687,7 +713,7 @@ def _aux_coords(ds: Dataset | DataArray, dims: set[str]) -> AuxCoords: """Scalar and one-dimensional auxiliary coordinates of ``ds`` lying on ``dims``.""" aux: AuxCoords = {} for n, c in ds.coords.items(): - if n in ds.dims: + if n in ds.xindexes: continue if c.ndim == 0: aux[str(n)] = ((), c.to_numpy()) diff --git a/linopy/expressions.py b/linopy/expressions.py index de8b5f9e..21474b3b 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -2669,6 +2669,9 @@ def to_constraint( self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None ) -> ConstraintBase: if self._csr is not None and isinstance(sign, str): + rhs = as_constant(rhs) + if not isinstance(rhs, CONSTANT_TYPES): + return self.sub(rhs, join=join).to_constraint(sign, 0) rhs_da = constraints.csr_rhs(self._csr, rhs) if isinstance(rhs_da, DataArray): return constraints.CSRConstraint.from_csr(self._csr, sign, rhs_da) diff --git a/linopy/model.py b/linopy/model.py index fd4066c9..a309f58c 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1312,7 +1312,7 @@ def add_constraints( if penalty is not None and freeze: raise ValueError( - "`penalty` cannot be combined with `freeze=True` (or a model default of `freeze_constraints=True`)," + "`penalty` cannot be combined with `freeze=True` (or a model default of `freeze_constraints=True`), " "since `soften` is not supported on frozen constraints." ) @@ -1329,17 +1329,11 @@ def add_constraints( rhs_da = as_dataarray(rhs) original_rhs_mask = (rhs_da.coords, rhs_da.dims, ~np.isnan(rhs_da.values)) - if ( - isinstance(lhs, LinearExpression) - and lhs.is_sparse - and freeze - and mask is not None - ): - mask = broadcast_to_coords(mask, lhs.coords, label="mask").astype(bool) - lhs = lhs.where(mask) - mask = None con = self._constraint_from_lhs(lhs, sign, rhs, coords) - if isinstance(con, CSRConstraint) and freeze and mask is None: + if isinstance(con, CSRConstraint) and freeze: + if mask is not None: + mask = broadcast_to_coords(mask, con.coords, label="mask") + con = con.masked(mask.astype(bool)) _check_infinities(con._sign, con._rhs, name) self.check_force_dim_names(con.coords.to_dataset()) enforce_no_multiindex(con, context=f"constraint {name!r}") diff --git a/test/test_csr.py b/test/test_csr.py index acda8774..b7d6f831 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1351,7 +1351,11 @@ def test_is_sparse_tracks_backing_and_repr_marks_it() -> None: "data": (lambda e, c: e.data, "`.data` read"), "merge": (lambda e, c: e + 1.0 * c.flow, "over different dimensions"), "matmul": (lambda e, c: e @ loc_operand(c.gbus.index), "sharing no dimension"), - "rhs": (lambda e, c: e <= 1.0 * c.gen_p.sum("gen"), "non-constant rhs"), + "rhs": ( + lambda e, c: e <= xr.DataArray([1.0, 2.0], coords=[LOC]), + "rhs over dimensions outside the grid", + ), + "rhs-expr": (lambda e, c: e <= 1.0 * c.gen_p.sum("gen"), "over different dim"), "mutable": (lambda e, c: (e == c.load).mutable(), "`mutable\\(\\)`"), "new-dim": (lambda e, c: e * xr.DataArray([1.0, 2.0], coords=[LOC]), "new dim"), "sum-kwargs": (lambda e, c: e.sum(dims="bus"), "`.data` read"), @@ -1818,14 +1822,120 @@ def test_selection_fallbacks_match_dense(op: str) -> None: } +@pytest.mark.parametrize("prebuilt", [False, True], ids=["lhs", "constraint"]) @pytest.mark.parametrize("mask", list(MASKS)) -def test_add_constraints_mask_freezes_sparse_and_matches_dense(mask: str) -> None: +def test_add_constraints_mask_freezes_sparse_and_matches_dense( + mask: str, prebuilt: bool +) -> None: require_v1() c1, c2 = base_model(), base_model() lhs = c2.balance_lhs(sparse=True) m = MASKS[mask](lhs) con1 = c1.m.add_constraints(c1.balance_lhs(False), ">=", c1.load, "bal", mask=m) with no_densify(): - con2 = c2.m.add_constraints(lhs, ">=", c2.load, "bal", mask=m, freeze=True) + if prebuilt: + con2 = c2.m.add_constraints(lhs >= c2.load, name="bal", mask=m, freeze=True) + else: + con2 = c2.m.add_constraints(lhs, ">=", c2.load, "bal", mask=m, freeze=True) + assert isinstance(con2, CSRConstraint) + assert_frozen_equal(con1, con2) + + +def observed_keys_group() -> LinearExpression: + """The #941 reproducer: a multi-key observed grouping with aux coords on ``group``.""" + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="s")], name="x") + expr = x.to_linexpr().assign_coords( + period=("s", [1, 1, 2, 2]), region=("s", ["n", "s", "n", "s"]) + ) + return expr.groupby(["period", "region"]).sum(observed=True, sparse=True) + + +AUX_BUILDS: dict[str, Callable[[], LinearExpression]] = { + "observed-keys": observed_keys_group, + "aux": lambda: SPARSE_BUILDS["aux"](base_model()), + "scalar": lambda: SPARSE_BUILDS["grouped"](base_model()).sel(snapshot=1), +} + + +@pytest.mark.parametrize("build", list(AUX_BUILDS)) +def test_frozen_constraint_keeps_aux_coords_like_dense( + build: str, tmp_path: Path +) -> None: + require_v1() + sparse = AUX_BUILDS[build]() + m = sparse.model + assert sparse._csr is not None + ref = m.add_constraints(sparse._csr.to_dense() >= 1, name="dense") + with no_densify(): + con = m.add_constraints(sparse >= 1, name="sparse", freeze=True) + assert isinstance(con, CSRConstraint) + want = xr.Dataset(coords=ref.coords) + assert set(want.coords) > set(con.coord_names) + for got in (con, con.mutable(), ref.freeze()): + xr.testing.assert_identical(xr.Dataset(coords=got.coords), want) + m.to_netcdf(tmp_path / "m.nc") + read = linopy.read_netcdf(tmp_path / "m.nc").constraints["sparse"] + assert isinstance(read, CSRConstraint) + xr.testing.assert_identical(xr.Dataset(coords=read.coords), want) + + +FROZEN_MUTATIONS: dict[str, Callable[[CSRConstraint], Any]] = { + "loc": lambda con: con.loc[{"bus": "bus0"}], + "update": lambda con: con.update(rhs=2.0), + "from_rule": lambda con: type(con).from_rule(con.model, lambda m, i: None, [[0]]), + "soften": lambda con: con.soften(penalty=1.0), + **{ + attr: lambda con, attr=attr: setattr(con, attr, 1.0) + for attr in ["coeffs", "vars", "sign", "rhs", "lhs", "scaling"] + }, +} + + +@pytest.mark.parametrize("op", list(FROZEN_MUTATIONS)) +def test_frozen_constraint_mutation_names_mutable(op: str) -> None: + require_v1() + c = base_model() + con = c.m.add_constraints(c.balance_lhs(sparse=True) >= c.load, freeze=True) + assert isinstance(con, CSRConstraint) + with pytest.raises(AttributeError, match=rf"CSRConstraint\.{op} .*\.mutable\(\)"): + FROZEN_MUTATIONS[op](con) + + +EXPR_RHS: dict[str, Callable[[Case, bool], LinearExpression]] = { + "expr": lambda c, sparse: (1.0 * c.flow).groupby(c.bus1).sum(sparse=sparse), + "expr-const": lambda c, sparse: ( + (1.0 * c.flow).groupby(c.bus1).sum(sparse=sparse) + c.load + ), + "dense-expr-const": lambda c, sparse: (1.0 * c.flow).groupby(c.bus1).sum() - 2.0, +} + + +@pytest.mark.parametrize("form", ["operator", "add_constraints"]) +@pytest.mark.parametrize("rhs", list(EXPR_RHS)) +def test_expression_rhs_freezes_sparse_and_matches_dense(rhs: str, form: str) -> None: + require_v1() + c1, c2 = base_model(), base_model() + lhs1 = (c1.eff * c1.gen_p).groupby(c1.gbus).sum() + con1 = c1.m.add_constraints(lhs1 <= EXPR_RHS[rhs](c1, False), name="c", freeze=True) + lhs2 = (c2.eff * c2.gen_p).groupby(c2.gbus).sum(sparse=True) + rhs2 = EXPR_RHS[rhs](c2, True) + with no_densify(): + if form == "operator": + con2 = c2.m.add_constraints(lhs2 <= rhs2, name="c", freeze=True) + else: + con2 = c2.m.add_constraints(lhs2, "<=", rhs2, name="c", freeze=True) assert isinstance(con2, CSRConstraint) assert_frozen_equal(con1, con2) + + +@pytest.mark.parametrize("default", [False, True], ids=["argument", "model-default"]) +def test_penalty_with_freeze_raises(default: bool) -> None: + require_v1() + c = base_model() + c.m.freeze_constraints = default + lhs = c.balance_lhs(sparse=True) + with pytest.raises(ValueError, match="`penalty` cannot be combined with `freeze"): + c.m.add_constraints( + lhs >= c.load, penalty=1.0, freeze=None if default else True + ) From 037f91cbfc966b88865010c88f4610556304f493 Mon Sep 17 00:00:00 2001 From: Fabian Date: Wed, 23 Sep 2026 17:59:35 +0200 Subject: [PATCH 05/15] fix(csr): dash-safe netcdf coord names, densify notice on unfrozen rebuild, join for constant rhs, 0-d merge --- doc/release_notes.rst | 3 +- linopy/common.py | 65 +++++++++++++++++++++++-------------- linopy/constraints.py | 14 ++++++-- linopy/csr.py | 34 ++++++++++++-------- linopy/expressions.py | 2 +- linopy/model.py | 10 ++++++ test/test_common.py | 33 +++++++++++++++---- test/test_csr.py | 75 ++++++++++++++++++++++++++++++++++++++++++- 8 files changed, 185 insertions(+), 51 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index b6ebb1c5..fa19ade5 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -76,7 +76,8 @@ Upcoming Version * ``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 `__) * 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 `__) * 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 `__) -* ``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, like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#963 `__) +* ``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), like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#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 `__) * ``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 `__) diff --git a/linopy/common.py b/linopy/common.py index ba055859..79291050 100644 --- a/linopy/common.py +++ b/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,24 +1319,25 @@ 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 @@ -1343,29 +1345,44 @@ 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_``). """ 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 diff --git a/linopy/constraints.py b/linopy/constraints.py index abfbb122..66439408 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -593,8 +593,12 @@ def _positional_csr( _UNSUPPORTED = "is not supported on a frozen constraint" -def _frozen_error(attr: str, what: str = "is read-only") -> AttributeError: - return AttributeError(f"CSRConstraint.{attr} {what}; call .mutable() to modify.") +def _frozen_error( + attr: str, + what: str = "is read-only", + remedy: str = "call .mutable() to modify", +) -> AttributeError: + return AttributeError(f"CSRConstraint.{attr} {what}; {remedy}.") _PositionalCache = tuple[scipy.sparse.csr_array, np.ndarray, "weakref.ref[np.ndarray]"] @@ -935,7 +939,11 @@ def soften(self, *args: Any, **kwargs: Any) -> NoReturn: @classmethod def from_rule(cls, *args: Any, **kwargs: Any) -> NoReturn: - raise _frozen_error("from_rule", _UNSUPPORTED) + raise _frozen_error( + "from_rule", + "is not supported", + "build with Constraint.from_rule and call .freeze() on the result", + ) @property @has_optimized_model diff --git a/linopy/csr.py b/linopy/csr.py index 72fb8e5c..22d6e694 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -57,8 +57,8 @@ AuxCoords: TypeAlias = dict[str, tuple[str | tuple[()], np.ndarray]] """Auxiliary coordinates as ``name -> (grid dim, values)``, dim ``()`` for a scalar.""" -_AUX_PREFIX = "_aux_" -_COORDDIM_PREFIX = "_coorddim_" +_AUX_PREFIX = "_aux" +_SCALAR_DIM = "_scalar" @dataclass(frozen=True, eq=False) @@ -92,22 +92,25 @@ def from_dataset(cls, ds: Dataset | DataArray, dims: Iterable[str]) -> Grid: @classmethod def from_netcdf_vars(cls, ds: Dataset, dims: Iterable[str]) -> Grid: """Read back a grid written by :meth:`to_netcdf_vars`.""" - aux: AuxCoords = {} - for k in ds: - name = str(k) - if name.startswith(_AUX_PREFIX): - da = ds[k] - d = str(da.dims[0]).removeprefix(_COORDDIM_PREFIX) if da.ndim else () - aux[name.removeprefix(_AUX_PREFIX)] = (d, da.to_numpy()) + aux: AuxCoords = { + da.attrs["name"]: (da.attrs.get("dim", ()), da.to_numpy()) + for k, da in ds.data_vars.items() + if str(k).startswith(_AUX_PREFIX) + } return cls(cls.from_coords(coords_from_dataset(ds, list(dims))).indexes, aux) def to_netcdf_vars(self) -> dict[str, DataArray]: - """The indexes and auxiliary coordinates as plain data variables for netcdf.""" + """ + The indexes and auxiliary coordinates as plain data variables for + netcdf, named by position with the coordinate names as attributes. + """ aux = { - f"{_AUX_PREFIX}{n}": DataArray( - v, dims=[f"{_COORDDIM_PREFIX}{d}"] if isinstance(d, str) else [] + f"{_AUX_PREFIX}{j}": DataArray( + v, + dims=[f"{_AUX_PREFIX}dim{j}"] if isinstance(d, str) else [], + attrs={"name": n} | ({"dim": d} if isinstance(d, str) else {}), ) - for n, (d, v) in self.aux.items() + for j, (n, (d, v)) in enumerate(self.aux.items()) } return coords_to_dataset_vars(self.coords) | aux @@ -369,6 +372,9 @@ def from_dense(cls, ds: Dataset, model: Model) -> CSRLinearExpression: """Convert a dense expression to CSR form on its own coordinate grid.""" grid_dims = tuple(str(d) for d in ds.coeffs.dims if d not in HELPER_DIMS) grid = Grid.from_dataset(ds, grid_dims) + if not grid_dims: + scalar = ds.expand_dims(_SCALAR_DIM) + return cls._from_scatter(scalar, model, grid, _SCALAR_DIM, {}, False) first = grid_dims[0] codes = {first: np.arange(len(grid.indexes[first]))} return cls._from_scatter(ds, model, grid, first, codes, False) @@ -394,7 +400,7 @@ def _from_scatter( stays absent, as on the dense v1 merge path. """ grid_dims = grid.dims - slot = min(grid_dims.index(d) for d in scatter_codes) + slot = min((grid_dims.index(d) for d in scatter_codes), default=0) transposed = [d for d in grid_dims if d not in scatter_codes] transposed.insert(slot, member_dim) member_rows = _member_rows(grid, scatter_codes, ds.sizes[member_dim]) diff --git a/linopy/expressions.py b/linopy/expressions.py index 21474b3b..a415151c 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -2670,7 +2670,7 @@ def to_constraint( ) -> ConstraintBase: if self._csr is not None and isinstance(sign, str): rhs = as_constant(rhs) - if not isinstance(rhs, CONSTANT_TYPES): + if join is not None or not isinstance(rhs, CONSTANT_TYPES): return self.sub(rhs, join=join).to_constraint(sign, 0) rhs_da = constraints.csr_rhs(self._csr, rhs) if isinstance(rhs_da, DataArray): diff --git a/linopy/model.py b/linopy/model.py index a309f58c..93d80b19 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -61,6 +61,7 @@ Constraints, CSRConstraint, ) +from linopy.csr import _densify_notice from linopy.dualization import dualize from linopy.expressions import ( Expressions, @@ -1308,6 +1309,7 @@ def add_constraints( name = self._resolve_constraint_name(name) if freeze is None: freeze = self.freeze_constraints + chunked = bool(freeze and self.chunk) freeze = freeze and not self.chunk if penalty is not None and freeze: @@ -1345,6 +1347,14 @@ def add_constraints( self._cCounter += con.full_size con = con.assign_labels(cindex, name, scaling_grid.values.ravel()) return self.constraints.add(con) + if isinstance(con, CSRConstraint): + if penalty is not None: + reason = "`penalty` given, softening needs a mutable constraint" + elif chunked: + reason = "chunked model, `Model.chunk` adds constraints unfrozen" + else: + reason = "constraint added unfrozen, `freeze=False`" + _densify_notice(reason) data = con.data _check_infinities(data.sign, data.rhs, name) diff --git a/test/test_common.py b/test/test_common.py index 7109314e..ead6a541 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -79,20 +79,39 @@ def test_assign_multiindex_safe() -> None: def test_coords_dataset_vars_roundtrip_multiindex() -> None: """MultiIndex and plain coords survive serialization to Dataset vars and back.""" mi = pd.MultiIndex.from_product( - [[2020, 2030], ["t1", "t2"]], names=("period", "timestep") + [[2020, 2030], ["t1", "t2"]], names=("period", "time-step") ) - mi.name = "snapshot" - plain = pd.Index([1, 2, 3], name="simple") + mi.name = "snap-shot" + plain = pd.Index([1, 2, 3], name="my-simple") ds = xr.Dataset(coords_to_dataset_vars([mi, plain])) - restored = coords_from_dataset(ds, ["snapshot", "simple"]) + assert not any("-" in str(k) for k in ds) + restored = coords_from_dataset(ds, ["snap-shot", "my-simple"]) assert isinstance(restored[0], pd.MultiIndex) assert restored[0].equals(mi) - assert list(restored[0].names) == ["period", "timestep"] - assert restored[0].name == "snapshot" + assert list(restored[0].names) == ["period", "time-step"] + assert restored[0].name == "snap-shot" assert restored[1].equals(plain) - assert restored[1].name == "simple" + assert restored[1].name == "my-simple" + + +def test_coords_from_dataset_reads_name_keyed_format() -> None: + """Files written before the positional format keyed the variables by name.""" + ds = xr.Dataset( + { + "_coord_snapshot_level_period": ("a", [2020, 2030]), + "_coord_snapshot_level_timestep": ("b", ["t1", "t2"]), + "_coord_snapshot_codes": (("c", "d"), [[0, 0], [0, 1], [1, 0], [1, 1]]), + "_coord_simple": ("e", [1, 2, 3]), + } + ) + mi, plain = coords_from_dataset(ds, ["snapshot", "simple"]) + want = pd.MultiIndex.from_product( + [[2020, 2030], ["t1", "t2"]], names=("period", "timestep") + ) + assert mi.equals(want) and mi.name == "snapshot" + assert plain.equals(pd.Index([1, 2, 3])) and plain.name == "simple" def test_iterate_slices_basic() -> None: diff --git a/test/test_csr.py b/test/test_csr.py index b7d6f831..65c2e66a 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1347,6 +1347,16 @@ def test_is_sparse_tracks_backing_and_repr_marks_it() -> None: assert not (1.0 * c.gen_p).is_sparse +def add_chunked(e: LinearExpression, c: Case) -> Any: + c.m.chunk = {"bus": 2} + return c.m.add_constraints(e >= 1, freeze=True) + + +def add_with_penalty(e: LinearExpression, c: Case) -> Any: + c.m.add_objective(1.0 * c.gen_p.sum()) + return c.m.add_constraints(e >= 1, penalty=1.0) + + DENSIFY_OPS: dict[str, tuple[Callable[[LinearExpression, Case], Any], str]] = { "data": (lambda e, c: e.data, "`.data` read"), "merge": (lambda e, c: e + 1.0 * c.flow, "over different dimensions"), @@ -1357,6 +1367,12 @@ def test_is_sparse_tracks_backing_and_repr_marks_it() -> None: ), "rhs-expr": (lambda e, c: e <= 1.0 * c.gen_p.sum("gen"), "over different dim"), "mutable": (lambda e, c: (e == c.load).mutable(), "`mutable\\(\\)`"), + "add-unfrozen": ( + lambda e, c: c.m.add_constraints(e >= 1), + "constraint added unfrozen, `freeze=False`", + ), + "add-chunked": (add_chunked, "chunked model"), + "add-penalty": (add_with_penalty, "`penalty` given"), "new-dim": (lambda e, c: e * xr.DataArray([1.0, 2.0], coords=[LOC]), "new dim"), "sum-kwargs": (lambda e, c: e.sum(dims="bus"), "`.data` read"), "groupby-fallback": ( @@ -1534,6 +1550,28 @@ def test_elementwise_constant_ops_stay_csr_and_match_dense( assert_linequal(res, func(dense, x)) +DENSE_ADDENDS: dict[str, Callable[[Case, LinearExpression], LinearExpression]] = { + "same-grid": lambda c, dense: 2 * dense, + "scalar-var": lambda c, dense: 2 * c.gen_p.isel(gen=0, snapshot=0), +} + + +@pytest.mark.parametrize( + ("build", "addend"), + [(b, "same-grid") for b in SPARSE_BUILDS] + [("zero-dim", "scalar-var")], +) +def test_merge_with_dense_expression_stays_csr_and_matches_dense( + build: str, addend: str +) -> None: + require_v1() + c = base_model() + sparse, dense = sparse_and_dense(build, c) + other = DENSE_ADDENDS[addend](c, dense) + with no_densify(): + res = sparse + other + assert_sparse_matches(res, dense + other) + + @pytest.mark.parametrize("join", ["inner", "outer", "left", "right"]) @pytest.mark.parametrize("fill_value", [None, linopy.ABSENT], ids=["fill", "absent"]) @pytest.mark.parametrize("method", ["add", "sub", "mul", "div"]) @@ -1604,6 +1642,24 @@ def test_join_with_constant_keeps_aux_coords_like_dense( assert_sparse_matches(res, getattr(dense, op)(x, join=join)) +@pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) +def test_to_constraint_join_with_constant_matches_dense(join: JoinOptions) -> None: + require_v1() + cons = [] + for sparse in (False, True): + c = base_model() + lhs = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=sparse) + with no_densify(): + con = lhs.to_constraint("<=", c.load.isel(bus=[0, 2]), join=join) + cons.append(c.m.add_constraints(con, name="c", freeze=sparse)) + dense, frozen = cons + assert isinstance(frozen, CSRConstraint) + xr.testing.assert_identical( + xr.Dataset(coords=frozen.coords), xr.Dataset(coords=dense.coords) + ) + assert_frozen_equal(dense, frozen) + + def assert_sparse_matches(res: LinearExpression, want: LinearExpression) -> None: """CSR backing kept, and equal to the dense reference up to term layout.""" assert res.is_sparse @@ -1851,8 +1907,21 @@ def observed_keys_group() -> LinearExpression: return expr.groupby(["period", "region"]).sum(observed=True, sparse=True) +def dashed_group() -> LinearExpression: + """An observed grouping whose dims and aux coords carry dashes, the netcdf name separator.""" + m = Model() + s = pd.RangeIndex(4, name="my-s") + x = m.add_variables(coords=[s], name="x") + grouper = pd.DataFrame( + {"my-bus": ["a", "a", "b", "b"], "tag": [1, 1, 2, 2]}, index=s + ) + grouped = x.to_linexpr().groupby(grouper).sum(sparse=True, observed=True) + return grouped.rename({"group": "my-group"}) + + AUX_BUILDS: dict[str, Callable[[], LinearExpression]] = { "observed-keys": observed_keys_group, + "dashed": dashed_group, "aux": lambda: SPARSE_BUILDS["aux"](base_model()), "scalar": lambda: SPARSE_BUILDS["grouped"](base_model()).sel(snapshot=1), } @@ -1883,7 +1952,6 @@ def test_frozen_constraint_keeps_aux_coords_like_dense( FROZEN_MUTATIONS: dict[str, Callable[[CSRConstraint], Any]] = { "loc": lambda con: con.loc[{"bus": "bus0"}], "update": lambda con: con.update(rhs=2.0), - "from_rule": lambda con: type(con).from_rule(con.model, lambda m, i: None, [[0]]), "soften": lambda con: con.soften(penalty=1.0), **{ attr: lambda con, attr=attr: setattr(con, attr, 1.0) @@ -1902,6 +1970,11 @@ def test_frozen_constraint_mutation_names_mutable(op: str) -> None: FROZEN_MUTATIONS[op](con) +def test_frozen_from_rule_names_the_mutable_route() -> None: + with pytest.raises(AttributeError, match=r"Constraint\.from_rule .*\.freeze\(\)"): + CSRConstraint.from_rule(Model(), lambda m, i: None, [[0]]) + + EXPR_RHS: dict[str, Callable[[Case, bool], LinearExpression]] = { "expr": lambda c, sparse: (1.0 * c.flow).groupby(c.bus1).sum(sparse=sparse), "expr-const": lambda c, sparse: ( From aa441bc6b9470f10b8f826af0c0c0577b16299d0 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 09:26:00 +0200 Subject: [PATCH 06/15] fix(csr): soften on a frozen constraint points to freeze=False, not the detached mutable() copy --- doc/release_notes.rst | 2 +- linopy/constraints.py | 6 +++++- test/test_csr.py | 22 ++++++++++++++++++---- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index fa19ade5..67399d44 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -77,7 +77,7 @@ Upcoming Version * 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 `__) * 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 `__) * ``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), like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#963 `__) +* ``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 `__) * ``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 `__) * ``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 `__) diff --git a/linopy/constraints.py b/linopy/constraints.py index 66439408..09089b2b 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -935,7 +935,11 @@ def update(self, *args: Any, **kwargs: Any) -> NoReturn: raise _frozen_error("update", _UNSUPPORTED) def soften(self, *args: Any, **kwargs: Any) -> NoReturn: - raise _frozen_error("soften", _UNSUPPORTED) + raise _frozen_error( + "soften", + _UNSUPPORTED, + "add the constraint with freeze=False to soften it", + ) @classmethod def from_rule(cls, *args: Any, **kwargs: Any) -> NoReturn: diff --git a/test/test_csr.py b/test/test_csr.py index 65c2e66a..5e9f4121 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1952,7 +1952,6 @@ def test_frozen_constraint_keeps_aux_coords_like_dense( FROZEN_MUTATIONS: dict[str, Callable[[CSRConstraint], Any]] = { "loc": lambda con: con.loc[{"bus": "bus0"}], "update": lambda con: con.update(rhs=2.0), - "soften": lambda con: con.soften(penalty=1.0), **{ attr: lambda con, attr=attr: setattr(con, attr, 1.0) for attr in ["coeffs", "vars", "sign", "rhs", "lhs", "scaling"] @@ -1970,9 +1969,24 @@ def test_frozen_constraint_mutation_names_mutable(op: str) -> None: FROZEN_MUTATIONS[op](con) -def test_frozen_from_rule_names_the_mutable_route() -> None: - with pytest.raises(AttributeError, match=r"Constraint\.from_rule .*\.freeze\(\)"): - CSRConstraint.from_rule(Model(), lambda m, i: None, [[0]]) +FROZEN_UNSUPPORTED: dict[str, tuple[Callable[[CSRConstraint], Any], str]] = { + "soften": (lambda con: con.soften(penalty=1.0), r"freeze=False"), + "from_rule": ( + lambda con: CSRConstraint.from_rule(con.model, lambda m, i: None, [[0]]), + r"Constraint\.from_rule .*\.freeze\(\)", + ), +} + + +@pytest.mark.parametrize("op", list(FROZEN_UNSUPPORTED)) +def test_frozen_unsupported_names_the_working_route(op: str) -> None: + require_v1() + c = base_model() + con = c.m.add_constraints(c.balance_lhs(sparse=True) >= c.load, freeze=True) + assert isinstance(con, CSRConstraint) + call, remedy = FROZEN_UNSUPPORTED[op] + with pytest.raises(AttributeError, match=rf"CSRConstraint\.{op} .*{remedy}"): + call(con) EXPR_RHS: dict[str, Callable[[Case, bool], LinearExpression]] = { From 8f279319b52954620dddb3b28cda04fe7111248d Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 12:57:32 +0200 Subject: [PATCH 07/15] fix(types): widen frozen-constraint test helper, narrow objective expression, named setter in FROZEN_MUTATIONS --- test/test_csr.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/test/test_csr.py b/test/test_csr.py index 5e9f4121..1e90c740 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -92,7 +92,7 @@ def canon(df: pl.DataFrame) -> pl.DataFrame: ) -def assert_frozen_equal(con1: Constraint, con2: CSRConstraint) -> None: +def assert_frozen_equal(con1: ConstraintBase, con2: ConstraintBase) -> None: d1, d2 = canon(con1.to_polars()), canon(con2.to_polars()) assert d1["labels"].equals(d2["labels"]) assert d1["vars"].equals(d2["vars"]) @@ -1473,7 +1473,9 @@ def test_objective_stays_csr_and_exports_like_dense( lp_sparse, lp_dense = (tmp_path / f"{k}.lp" for k in ("sparse", "dense")) assert canon_lp(lp_sparse.read_text()) == canon_lp(lp_dense.read_text()) rs, rd = (linopy.read_netcdf(tmp_path / f"{k}.nc") for k in ("sparse", "dense")) - assert_cells_equal(rs.objective.expression, rd.objective.expression, ()) + es, ed = rs.objective.expression, rd.objective.expression + assert isinstance(es, LinearExpression) and isinstance(ed, LinearExpression) + assert_cells_equal(es, ed, ()) assert rs.objective.expression.attrs["name"] == "objective" @@ -1949,11 +1951,15 @@ def test_frozen_constraint_keeps_aux_coords_like_dense( xr.testing.assert_identical(xr.Dataset(coords=read.coords), want) +def setter(attr: str) -> Callable[[CSRConstraint], None]: + return lambda con: setattr(con, attr, 1.0) + + FROZEN_MUTATIONS: dict[str, Callable[[CSRConstraint], Any]] = { "loc": lambda con: con.loc[{"bus": "bus0"}], "update": lambda con: con.update(rhs=2.0), **{ - attr: lambda con, attr=attr: setattr(con, attr, 1.0) + attr: setter(attr) for attr in ["coeffs", "vars", "sign", "rhs", "lhs", "scaling"] }, } From 0cf4a2006dd46c15eaaaab81ba5b11f09f76bdf1 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 13:41:52 +0200 Subject: [PATCH 08/15] perf(objective): read linear objective terms from numpy arrays, not a polars round trip --- linopy/objective.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/linopy/objective.py b/linopy/objective.py index 2dc6af4d..e0182b4c 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -23,6 +23,7 @@ from xarray.core.utils import Frozen from linopy import expressions +from linopy.constants import FACTOR_DIM from linopy.types import ConstantLike if TYPE_CHECKING: @@ -156,13 +157,24 @@ def to_polars(self, **kwargs: Any) -> pl.DataFrame: def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: """ Returns the variable labels and coefficients of the linear objective - terms, zero coefficients dropped, read from ``to_polars``. + terms, read from the stored arrays; duplicate labels are not summed. """ - df = self.to_polars() - if self.is_quadratic: - df = linear_part(df) - df = df.drop_nulls("vars").cast({"vars": pl.Int64, "coeffs": pl.Float64}) - return df["vars"].to_numpy(), df["coeffs"].to_numpy() + expr = self.expression + if isinstance(expr, expressions.LinearExpression) and expr._csr is not None: + terms = expr._csr_terms(expr._csr) + return terms["vars"], terms["coeffs"] + coeffs = expr.data.coeffs.values.ravel() + vars = expr.data.vars + if isinstance(expr, expressions.QuadraticExpression): + factors = np.moveaxis(vars.values, vars.get_axis_num(FACTOR_DIM), 0) + vars1, vars2 = factors[0].ravel(), factors[1].ravel() + linear = (vars1 == -1) | (vars2 == -1) + labels = np.where(vars1 != -1, vars1, vars2)[linear] + coeffs = coeffs[linear] + else: + labels = vars.values.ravel() + mask = labels != -1 + return labels[mask], coeffs[mask] def to_netcdf_ds(self) -> Dataset: """ From 59146b705704ed269a357e49d4ec8336c25c0a73 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 14:31:27 +0200 Subject: [PATCH 09/15] perf(csr): read objective terms and scale matrix blocks from raw arrays, skip unit scaling --- linopy/matrices.py | 21 ++++++++++++--------- linopy/objective.py | 13 +++++++++---- test/test_csr.py | 4 ++++ 3 files changed, 25 insertions(+), 13 deletions(-) diff --git a/linopy/matrices.py b/linopy/matrices.py index 82da4a88..bca66d70 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -116,6 +116,7 @@ def _build_cons(self) -> None: label_index = m.variables.label_index label_to_pos = label_index.label_to_pos con_scaling_by_label = constraint_scaling_lookup(m) + unit_cols = bool((self.var_scaling == 1).all()) def scale_rows_and_cols( csr: scipy.sparse.csr_array, con_labels: np.ndarray, b: np.ndarray @@ -123,18 +124,20 @@ def scale_rows_and_cols( if csr.shape[0] == 0: return csr, b row_scaling = con_scaling_by_label[con_labels] + unit_rows = bool((row_scaling == 1).all()) + if unit_rows and unit_cols: + return csr, b # With solver variables y = Scol * x, constraints A x = b become # Srow * A * Scol^-1 * y = Srow * b. - csr = cast( - scipy.sparse.csr_array, - csr.multiply(row_scaling[:, np.newaxis]).tocsr(), + data = csr.data + if not unit_rows: + data = data * np.repeat(row_scaling, np.diff(csr.indptr)) + if not unit_cols: + data = data / self.var_scaling[csr.indices] + scaled = scipy.sparse.csr_array( + (data, csr.indices, csr.indptr), shape=csr.shape ) - if csr.shape[1] and len(self.var_scaling): - csr = cast( - scipy.sparse.csr_array, - csr.multiply(1 / self.var_scaling[np.newaxis, :]).tocsr(), - ) - return csr, b * row_scaling + return scaled, b * row_scaling reg_csrs, reg_b, reg_sense = [], [], [] ind_csrs, ind_b, ind_sense, ind_binvar, ind_binval = [], [], [], [], [] diff --git a/linopy/objective.py b/linopy/objective.py index e0182b4c..72afdb33 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -157,12 +157,17 @@ def to_polars(self, **kwargs: Any) -> pl.DataFrame: def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: """ Returns the variable labels and coefficients of the linear objective - terms, read from the stored arrays; duplicate labels are not summed. + terms, read from the stored arrays; absent terms and zero coefficients + are dropped, duplicate labels are not summed. """ expr = self.expression if isinstance(expr, expressions.LinearExpression) and expr._csr is not None: - terms = expr._csr_terms(expr._csr) - return terms["vars"], terms["coeffs"] + csr = expr._csr.csr + present = ~np.isnan(expr._csr.const) + keep = np.repeat(present, np.diff(csr.indptr)) & (csr.data != 0) + if keep.all(): + return csr.indices, csr.data + return csr.indices[keep], csr.data[keep] coeffs = expr.data.coeffs.values.ravel() vars = expr.data.vars if isinstance(expr, expressions.QuadraticExpression): @@ -173,7 +178,7 @@ def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: coeffs = coeffs[linear] else: labels = vars.values.ravel() - mask = labels != -1 + mask = (labels != -1) & (coeffs != 0) return labels[mask], coeffs[mask] def to_netcdf_ds(self) -> Dataset: diff --git a/test/test_csr.py b/test/test_csr.py index 1e90c740..9c1f6116 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1461,6 +1461,7 @@ def test_objective_stays_csr_and_exports_like_dense( ms, md = objective_twins(build, scale) with no_densify(): c = ms.matrices.c + terms = ms.objective.linear_terms() ms.to_file(tmp_path / "sparse.lp") ms.to_netcdf(tmp_path / "sparse.nc") copied = ms.copy() @@ -1470,6 +1471,9 @@ def test_objective_stays_csr_and_exports_like_dense( md.to_file(tmp_path / "dense.lp") md.to_netcdf(tmp_path / "dense.nc") assert np.array_equal(c, md.matrices.c) + dense_terms = md.objective.linear_terms() + assert sorted(zip(*map(list, terms))) == sorted(zip(*map(list, dense_terms))) + assert (dense_terms[1] != 0).all() lp_sparse, lp_dense = (tmp_path / f"{k}.lp" for k in ("sparse", "dense")) assert canon_lp(lp_sparse.read_text()) == canon_lp(lp_dense.read_text()) rs, rd = (linopy.read_netcdf(tmp_path / f"{k}.nc") for k in ("sparse", "dense")) From 1d5b58091d90763de321ed2e3a283c2303823696 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 14:41:50 +0200 Subject: [PATCH 10/15] perf(csr): CSR store indices follow the model label dtype (#828) --- linopy/constraints.py | 11 ++++++-- linopy/csr.py | 59 ++++++++++++++++++++++++++++++------------- test/test_csr.py | 22 ++++++++++++++++ 3 files changed, 73 insertions(+), 19 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 09089b2b..44999ea4 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -77,7 +77,13 @@ PerformanceWarning, SIGNS_pretty, ) -from linopy.csr import Grid, _densify_notice, csr_nterm, csr_to_term_arrays +from linopy.csr import ( + Grid, + _densify_notice, + csr_nterm, + csr_to_term_arrays, + index_dtype, +) from linopy.scaling import ensure_scaling, validate_scaling from linopy.semantics import check_user_nan from linopy.types import ( @@ -1978,7 +1984,8 @@ def _matrix_export_data( data = coeffs_final[valid_final] counts = valid_final.sum(axis=1) - indptr = np.empty(len(con_labels) + 1, dtype=np.int32) + dtype = index_dtype(len(data), (len(con_labels),), self.model) + indptr = np.empty(len(con_labels) + 1, dtype=dtype) indptr[0] = 0 np.cumsum(counts, out=indptr[1:]) return con_labels, row_mask, vlabel_cols, data, indptr diff --git a/linopy/csr.py b/linopy/csr.py index 22d6e694..d410ffd3 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -269,6 +269,14 @@ class CSRLinearExpression: grid: Grid model: Model + def __post_init__(self) -> None: + csr = self.csr + dtype = index_dtype(csr.nnz, csr.shape, self.model) + if csr.indices.dtype != dtype or csr.indptr.dtype != dtype: + indices, indptr = csr.indices.astype(dtype), csr.indptr.astype(dtype) + csr = scipy.sparse.csr_array((csr.data, indices, indptr), shape=csr.shape) + object.__setattr__(self, "csr", csr) + @property def shape(self) -> tuple[int, ...]: return self.grid.shape @@ -412,9 +420,8 @@ def _from_scatter( keep = (vars_ != -1) & ~np.isnan(coeffs) full_size = grid.size - coo = scipy.sparse.coo_array( - (coeffs[keep], (rows[keep], vars_[keep])), - shape=(full_size, model._xCounter), + csr = coo_to_csr( + coeffs[keep], rows[keep], vars_[keep], (full_size, model._xCounter), model ) const_vals = ds.const.transpose(*transposed).to_numpy().reshape(-1) @@ -424,7 +431,7 @@ def _from_scatter( const[cell_rows] = 0.0 np.add.at(const, cell_rows, const_vals) - return cls(scipy.sparse.csr_array(coo), const, grid, model) + return cls(csr, const, grid, model) def aggregated(self, grid: Grid, rows: np.ndarray) -> CSRLinearExpression: """ @@ -435,14 +442,13 @@ def aggregated(self, grid: Grid, rows: np.ndarray) -> CSRLinearExpression: cells no row lands in are absent. Auxiliary coordinates are ``grid``'s. """ coo = self.csr.tocoo() - coo = scipy.sparse.coo_array( - (coo.data, (rows[coo.coords[0]], coo.coords[1])), - shape=(grid.size, self.csr.shape[1]), - ) + shape = (grid.size, self.csr.shape[1]) + rows_ = rows[coo.coords[0]] + csr = coo_to_csr(coo.data, rows_, coo.coords[1], shape, self.model) weights = np.nan_to_num(self.const) const = np.bincount(rows, weights=weights, minlength=grid.size).astype(float) const[np.bincount(rows, minlength=grid.size) == 0] = np.nan - return replace(self, csr=scipy.sparse.csr_array(coo), const=const, grid=grid) + return replace(self, csr=csr, const=const, grid=grid) def summed(self, dims: Iterable[str]) -> CSRLinearExpression: """ @@ -527,14 +533,13 @@ def reindexed(self, grid: Grid, fill: float = np.nan) -> CSRLinearExpression: rows = row_map[coo.coords[0][keep]] cols = coo.coords[1][keep] n_cells = grid.size - coo = scipy.sparse.coo_array( - (coo.data[keep], (rows, cols)), shape=(n_cells, self.csr.shape[1]) - ) + shape = (n_cells, self.csr.shape[1]) + csr = coo_to_csr(coo.data[keep], rows, cols, shape, self.model) const = np.full(n_cells, fill) const[row_map[valid]] = self.const[valid] return replace( self, - csr=scipy.sparse.csr_array(coo), + csr=csr, const=const, grid=self.grid.conformed(grid), ) @@ -569,12 +574,10 @@ def added(self, other: CSRLinearExpression) -> CSRLinearExpression: cols = np.concatenate([a.coords[1], b.coords[1]]) data = np.concatenate([a.data, b.data]) present = ~np.isnan(const)[rows] - coo = scipy.sparse.coo_array( - (data[present], (rows[present], cols[present])), shape=shape - ) + csr = coo_to_csr(data[present], rows[present], cols[present], shape, self.model) enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in (self, other)]) grid = replace(self.grid, aux=other.grid.aux | self.grid.aux) - return replace(self, csr=scipy.sparse.csr_array(coo), const=const, grid=grid) + return replace(self, csr=csr, const=const, grid=grid) def contracted( self, @@ -670,6 +673,28 @@ def to_dense(self) -> LinearExpression: return LinearExpression(absorb_absence(ds), self.model) +def index_dtype(nnz: int, shape: tuple[int, ...], model: Model) -> np.dtype: + """ + Index dtype of a CSR store: the model's label dtype, widened to int64 + only when the nonzeros or the shape outgrow it. + """ + dtype = np.dtype(model._dtypes["labels"]) + return dtype if max(nnz, *shape) <= np.iinfo(dtype).max else np.dtype(np.int64) + + +def coo_to_csr( + data: np.ndarray, + rows: np.ndarray, + cols: np.ndarray, + shape: tuple[int, int], + model: Model, +) -> scipy.sparse.csr_array: + """Build a CSR store from COO triplets in the model's index dtype.""" + dtype = index_dtype(len(data), shape, model) + coords = (rows.astype(dtype, copy=False), cols.astype(dtype, copy=False)) + return scipy.sparse.csr_array(scipy.sparse.coo_array((data, coords), shape=shape)) + + def csr_nterm(csr: scipy.sparse.csr_array) -> int: """Widest CSR row, floored at one term.""" return max(int(np.diff(csr.indptr).max(initial=0)), 1) diff --git a/test/test_csr.py b/test/test_csr.py index 9c1f6116..cc6f2a85 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1435,6 +1435,28 @@ def test_flat_and_to_polars_served_without_densifying(build: str, scale: float) pd.testing.assert_frame_equal(got_flat, dense.flat) +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_sparse_store_indices_follow_model_label_dtype(build: str) -> None: + require_v1() + sparse, _ = sparse_and_dense(build) + m = sparse.model + con = m.add_constraints(sparse >= 1, name="con", freeze=True) + assert isinstance(con, CSRConstraint) and sparse._csr is not None + dtypes = {sparse._csr.csr.indices.dtype, con._csr.indices.dtype} + assert dtypes == {np.dtype(m.dtypes["labels"]), m.matrices.A.indices.dtype} + + +@pytest.mark.parametrize("label_dtype", [np.int32, np.int64]) +def test_sparse_group_sum_indices_widen_with_model(label_dtype: type) -> None: + require_v1() + m = Model(dtypes={"labels": label_dtype}) + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + group = xr.DataArray([0, 0, 1, 1], coords=[x.indexes["i"]], name="g") + expr = x.groupby(group).sum(sparse=True) + assert expr._csr is not None + assert expr._csr.csr.indices.dtype == expr._csr.csr.indptr.dtype == label_dtype + + def objective_twins(build: str, scale: float) -> tuple[Model, Model]: """Twin models, the first with the sparse build as objective, the second dense.""" sparse, _ = sparse_and_dense(build) From 3db05c289b7f08fb74db716e817367afbcceee2b Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 14:45:09 +0200 Subject: [PATCH 11/15] test(csr): narrow frozen store in label dtype test --- test/test_csr.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_csr.py b/test/test_csr.py index cc6f2a85..24e039b0 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1441,7 +1441,8 @@ def test_sparse_store_indices_follow_model_label_dtype(build: str) -> None: sparse, _ = sparse_and_dense(build) m = sparse.model con = m.add_constraints(sparse >= 1, name="con", freeze=True) - assert isinstance(con, CSRConstraint) and sparse._csr is not None + assert isinstance(con, CSRConstraint) + assert sparse._csr is not None and con._csr is not None dtypes = {sparse._csr.csr.indices.dtype, con._csr.indices.dtype} assert dtypes == {np.dtype(m.dtypes["labels"]), m.matrices.A.indices.dtype} From 24283a038638af830f51d3c8fd1604e9b963d891 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 14:46:13 +0200 Subject: [PATCH 12/15] test(csr): narrow optional matrix A in label dtype test --- test/test_csr.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/test/test_csr.py b/test/test_csr.py index 24e039b0..b75f49de 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1441,10 +1441,11 @@ def test_sparse_store_indices_follow_model_label_dtype(build: str) -> None: sparse, _ = sparse_and_dense(build) m = sparse.model con = m.add_constraints(sparse >= 1, name="con", freeze=True) - assert isinstance(con, CSRConstraint) - assert sparse._csr is not None and con._csr is not None + A = m.matrices.A + assert isinstance(con, CSRConstraint) and sparse._csr is not None + assert A is not None dtypes = {sparse._csr.csr.indices.dtype, con._csr.indices.dtype} - assert dtypes == {np.dtype(m.dtypes["labels"]), m.matrices.A.indices.dtype} + assert dtypes == {np.dtype(m.dtypes["labels"]), A.indices.dtype} @pytest.mark.parametrize("label_dtype", [np.int32, np.int64]) From cc738344e14b07bfb44a16c7693aab062cecc2d5 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 15:59:46 +0200 Subject: [PATCH 13/15] refactor(expressions): linear_terms lives on the expression classes CSRLinearExpression.live_terms is the single sparse term filter, shared by _csr_terms and LinearExpression.linear_terms. QuadraticExpression aligns factors with coeffs by dim name. Objective.linear_terms delegates. --- linopy/csr.py | 8 ++++++ linopy/expressions.py | 41 +++++++++++++++++++++++++++---- linopy/objective.py | 25 ++----------------- test/test_quadratic_expression.py | 11 +++++++++ 4 files changed, 57 insertions(+), 28 deletions(-) diff --git a/linopy/csr.py b/linopy/csr.py index d410ffd3..c9a148da 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -467,6 +467,14 @@ def summed(self, dims: Iterable[str]) -> CSRLinearExpression: ) return self.aggregated(grid, rows).filled(0.0) + def live_terms(self) -> np.ndarray: + """ + Mask over the stored terms: the cell is present and the coefficient + is nonzero. + """ + present = np.repeat(~np.isnan(self.const), np.diff(self.csr.indptr)) + return present & (self.csr.data != 0) + def pruned(self) -> CSRLinearExpression: """Drop explicit zero coefficients; cell activeness stays with ``const``.""" csr = self.csr.copy() diff --git a/linopy/expressions.py b/linopy/expressions.py index a415151c..9a4b148d 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -867,6 +867,9 @@ def flat(self) -> pd.DataFrame: ... @abstractmethod def to_polars(self) -> pl.DataFrame: ... + @abstractmethod + def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: ... + def __init__(self, data: Dataset | Any | None, model: Model) -> None: from linopy.model import Model @@ -2999,16 +3002,31 @@ def _csr_terms(self, csr: CSRLinearExpression) -> dict[str, np.ndarray]: columns, dropping absent cells and zero coefficients like the dense long format. """ - rows = np.repeat(np.arange(csr.n_cells), np.diff(csr.csr.indptr)) - const = csr.const[rows] - coeffs = csr.csr.data.astype(float) - keep = ~np.isnan(const) & (coeffs != 0) + keep = csr.live_terms() + const = np.repeat(csr.const, np.diff(csr.csr.indptr)) return { "const": const[keep], - "coeffs": coeffs[keep], + "coeffs": csr.csr.data[keep].astype(float), "vars": csr.csr.indices[keep].astype(self.model._dtypes["labels"]), } + def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: + """ + Variable labels and coefficients of the stored terms, without a dense + round trip for a CSR backing; absent terms and zero coefficients are + dropped, duplicate labels are not summed. + """ + if self._csr is not None: + csr = self._csr.csr + keep = self._csr.live_terms() + if keep.all(): + return csr.indices, csr.data + return csr.indices[keep], csr.data[keep] + labels = self.data.vars.values.ravel() + coeffs = self.data.coeffs.values.ravel() + keep = (labels != -1) & (coeffs != 0) + return labels[keep], coeffs[keep] + def simplify(self) -> LinearExpression: """ Simplify the linear expression by combining terms with the same variable. @@ -3483,6 +3501,19 @@ def mask_func(data: dict) -> pd.Series: check_has_nulls(df, name=self.type) return df + def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: + """ + Variable labels and coefficients of the terms with a single variable; + zero coefficients are dropped, duplicate labels are not summed. + """ + coeffs = self.data.coeffs + factors = self.data.vars.transpose(FACTOR_DIM, *coeffs.dims).values + vars1, vars2 = factors[0].ravel(), factors[1].ravel() + labels = np.where(vars1 != -1, vars1, vars2) + values = coeffs.values.ravel() + keep = ((vars1 == -1) | (vars2 == -1)) & (labels != -1) & (values != 0) + return labels[keep], values[keep] + def to_polars(self, **kwargs: Any) -> pl.DataFrame: """ Convert the expression to a polars DataFrame. diff --git a/linopy/objective.py b/linopy/objective.py index 72afdb33..ad7b6bdd 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -23,7 +23,6 @@ from xarray.core.utils import Frozen from linopy import expressions -from linopy.constants import FACTOR_DIM from linopy.types import ConstantLike if TYPE_CHECKING: @@ -157,29 +156,9 @@ def to_polars(self, **kwargs: Any) -> pl.DataFrame: def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: """ Returns the variable labels and coefficients of the linear objective - terms, read from the stored arrays; absent terms and zero coefficients - are dropped, duplicate labels are not summed. + terms, see ``LinearExpression.linear_terms``. """ - expr = self.expression - if isinstance(expr, expressions.LinearExpression) and expr._csr is not None: - csr = expr._csr.csr - present = ~np.isnan(expr._csr.const) - keep = np.repeat(present, np.diff(csr.indptr)) & (csr.data != 0) - if keep.all(): - return csr.indices, csr.data - return csr.indices[keep], csr.data[keep] - coeffs = expr.data.coeffs.values.ravel() - vars = expr.data.vars - if isinstance(expr, expressions.QuadraticExpression): - factors = np.moveaxis(vars.values, vars.get_axis_num(FACTOR_DIM), 0) - vars1, vars2 = factors[0].ravel(), factors[1].ravel() - linear = (vars1 == -1) | (vars2 == -1) - labels = np.where(vars1 != -1, vars1, vars2)[linear] - coeffs = coeffs[linear] - else: - labels = vars.values.ravel() - mask = (labels != -1) & (coeffs != 0) - return labels[mask], coeffs[mask] + return self.expression.linear_terms() def to_netcdf_ds(self) -> Dataset: """ diff --git a/test/test_quadratic_expression.py b/test/test_quadratic_expression.py index 40ab8c41..fa9381b5 100644 --- a/test/test_quadratic_expression.py +++ b/test/test_quadratic_expression.py @@ -296,6 +296,17 @@ def test_quadratic_expression_to_polars(x: Variable, y: Variable) -> None: assert len(df) == expr.nterm * 2 +@pytest.mark.parametrize("factor_last", [False, True]) +def test_quadratic_expression_linear_terms( + x: Variable, y: Variable, factor_last: bool +) -> None: + expr = x * y + 3 * x + 0 * y + if factor_last: + expr = QuadraticExpression(expr.data.transpose(..., FACTOR_DIM), expr.model) + labels, coeffs = expr.linear_terms() + assert sorted(zip(labels.tolist(), coeffs.tolist())) == [(0, 3.0), (1, 3.0)] + + def test_quadratic_expression_constant_to_polars() -> None: m = Model() arr = pd.Series(index=pd.Index([0, 1], name="t"), data=[10, 20]) From a817a5efaeeb0c2dfae9a230a5ab6d32a964d88f Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 16:12:51 +0200 Subject: [PATCH 14/15] docs(csr): docstrings state current behavior, not development history --- linopy/constraints.py | 14 +++--- linopy/csr.py | 105 +++++++++++++++++++++--------------------- linopy/expressions.py | 53 ++++++++++++--------- linopy/objective.py | 5 +- 4 files changed, 93 insertions(+), 84 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 44999ea4..1a61d8bb 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -1189,8 +1189,8 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: """ Reconstruct a Constraint from a netcdf Dataset (CSR format). - Files without the ``_csr_columns`` attribute were written before #926 - and hold dense variable positions instead of labels. + Files without the ``_csr_columns`` attribute hold dense variable + positions instead of labels. """ attrs = ds.attrs shape = tuple(attrs["shape"]) @@ -1465,12 +1465,10 @@ def from_csr( """ Staple sign and rhs onto a CSR-backed lhs to form an unassigned CSRConstraint. - The sparse counterpart of :meth:`from_dense`: instead of converting a - dense :class:`Constraint`, it realizes a - :class:`~linopy.csr.CSRLinearExpression` directly. The expression's - label columns are kept as they are, its constant moves to the rhs, and - rows with a NaN rhs are inactive, as on the dense path. ``rhs`` must - come from :func:`csr_rhs`. + Builds directly from a :class:`~linopy.csr.CSRLinearExpression`. The + expression's label columns are kept as they are, its constant moves to + the rhs, and rows with a NaN rhs are inactive. ``rhs`` must come from + :func:`csr_rhs`. """ sign = maybe_replace_sign(sign) rhs_flat = _rhs_grid_values(expr, rhs) - expr.const diff --git a/linopy/csr.py b/linopy/csr.py index c9a148da..ea5d0da7 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -1,22 +1,22 @@ """ The sparse backing of a LinearExpression: ``A @ x + c`` in CSR form. -``expr.groupby(g).sum(sparse=True)`` (or ``linopy.options["sparse_groupby"]`` -under v1) returns an ordinary :class:`~linopy.expressions.LinearExpression` -backed by a :class:`CSRLinearExpression` instead of the dense dataset — same -public type, different backing, akin to dask-backed xarray objects. The CSR -form is canonical (duplicate variables summed, terms label-ordered) and ragged -along ``_term``, so the group-size padding of issue #745 has no analog; -grouping, ``sum``, ``merge``/``+``/``-``, scaling and ``@``/``dot`` -(:meth:`contracted`) become sparse linear algebra. Zero policy: the structural -operations (grouping and ``sum`` via :meth:`aggregated`, merge via +Under v1 semantics, ``expr.groupby(g).sum(sparse=True)`` (or +``linopy.options["sparse_groupby"]``) returns an ordinary +:class:`~linopy.expressions.LinearExpression` backed by a +:class:`CSRLinearExpression`: same public type, different backing, akin to +dask-backed xarray objects. The CSR form is canonical (duplicate variables +summed, terms label-ordered) and ragged along ``_term``, with no fixed term +count per row; grouping, ``sum``, ``merge``/``+``/``-``, scaling and +``@``/``dot`` (:meth:`contracted`) are sparse linear algebra. Zero policy: the +structural operations (grouping and ``sum`` via :meth:`aggregated`, merge via :meth:`added`, scaling, reindexing) go through COO and keep explicit zero -coefficients, like the dense path; only the product with a constant matrix, -``@``/``dot``, prunes them. Either way cell activeness is carried by ``const`` -alone (issue #925), so the two differ in term layout only. -Anything without a sparse branch expands through ``.data`` to the -mathematically identical dense rectangle in canonical term layout — the reason -the feature is v1-gated, where term layout is non-contractual. +coefficients; only the product with a constant matrix, ``@``/``dot``, prunes +them. Either way cell activeness is carried by ``const`` alone, independent of +term layout. +Any operation without a sparse branch expands the expression through +``.data`` to the mathematically identical dense rectangle in canonical term +layout; this is valid because v1 semantics do not fix the term layout. This module documents the CSR structure only, working on plain datasets. The one bridge back to a dense type is :meth:`CSRLinearExpression.to_dense`, which @@ -292,7 +292,8 @@ def nterm(self) -> int: def cell(self, indices: tuple[Any, ...]) -> tuple[np.ndarray, np.ndarray, float]: """ Coefficients, label-ordered variable labels and constant of the grid - cell at ``indices``, as in the dense form: an absent cell has no terms. + cell at ``indices``. An absent cell returns empty coefficient and + label arrays. """ row = int(np.ravel_multi_index(indices, self.grid.shape)) if indices else 0 const = float(self.const[row]) @@ -314,18 +315,18 @@ def from_grouper( coord_dims: tuple[str, ...], ) -> CSRLinearExpression: """ - Build the grouped sum directly in CSR form (no padded rectangle). + Build the grouped sum in CSR form. The grouper is conformed to the expression's member index by label (upstream alignment checks guarantee equal label sets) and group - labels are sorted, matching the dense kernel's output grid. A - DataFrame grouper (one column per key) yields one grid dim per key - -- the cartesian grid, absent combinations being empty cells -- or, - ``stacked``, a single ``group_dim`` over the observed key combinations - only, the key values attached as auxiliary coordinates. The new dims - take the member dim's slot in ``coord_dims``, as on the dense path. - ``source`` is a dense expression dataset or an already CSR-backed - expression, which is regrouped through :meth:`aggregated`. + labels are sorted. A DataFrame grouper (one column per key) yields + one grid dim per key -- the cartesian grid, absent combinations being + empty cells -- or, ``stacked``, a single ``group_dim`` over the + observed key combinations only, the key values attached as auxiliary + coordinates. The new dims take the member dim's slot in + ``coord_dims``. ``source`` is a dense expression dataset or an + already CSR-backed expression, which is regrouped through + :meth:`aggregated`. """ ds = source if isinstance(source, Dataset) else source.grid.to_dataset() member_dim = str(grouper.index.name) @@ -401,11 +402,11 @@ def _from_scatter( Scatter an expression's terms into grid rows (conceptually ``G @ A``): ``member_dim`` lands in the contiguous block of grid dims named by ``scatter_codes`` (one row-position array per dim), every other grid - dim maps one-to-one, and the COO→CSR conversion sums duplicates -- - which is the group sum. Cells no member lands in stay absent (NaN - const). With ``skipna`` the constant is reduced as by the dense group - kernel (NaN members count as 0); without it an absent cell (NaN const) - stays absent, as on the dense v1 merge path. + dim maps one-to-one, and the COO to CSR conversion sums duplicate + variables, giving the group sum. Cells no member lands in stay absent + (NaN const). With ``skipna``, NaN member constants count as 0 in the + sum; without it, a NaN member constant propagates and leaves the cell + absent (NaN const). """ grid_dims = grid.dims slot = min((grid_dims.index(d) for d in scatter_codes), default=0) @@ -438,8 +439,9 @@ def aggregated(self, grid: Grid, rows: np.ndarray) -> CSRLinearExpression: Sum source rows into the cells of ``grid``, row ``i`` landing in cell ``rows[i]`` (conceptually ``G @ A``). Goes through COO, so duplicate variables are summed and explicit zeros kept, as by :meth:`added`. The - constant is reduced as by the dense group kernel (NaN counts as 0); - cells no row lands in are absent. Auxiliary coordinates are ``grid``'s. + constant is the NaN-skipping sum of its rows' constants (NaN counts as + 0); cells no row lands in are absent. Auxiliary coordinates are + ``grid``'s. """ coo = self.csr.tocoo() shape = (grid.size, self.csr.shape[1]) @@ -452,9 +454,9 @@ def aggregated(self, grid: Grid, rows: np.ndarray) -> CSRLinearExpression: def summed(self, dims: Iterable[str]) -> CSRLinearExpression: """ - Sum over grid dimensions, as the dense ``sum``: the kept dims stay in - grid order with their auxiliary coordinates, and every kept cell is - present, its constant the NaN-skipping sum of its members. + Sum over grid dimensions. The kept dims stay in grid order with their + auxiliary coordinates, and every kept cell is present, its constant + the NaN-skipping sum of its members. """ dims = set(dims) grid = self.grid.reordered(d for d in self.grid.dims if d not in dims) @@ -501,8 +503,8 @@ def scaled( def with_const(self, const: np.ndarray) -> CSRLinearExpression: """ - Replace the per-cell constant, leaving the terms alone. A cell it makes - absent (NaN) drops its terms (v1 dead-term invariant). + Replace the per-cell constant, leaving the terms alone. A cell made + absent (NaN) has its terms dropped: an absent cell carries no terms. """ absent = np.isnan(const) counts = np.diff(self.csr.indptr) @@ -529,10 +531,10 @@ def taken(self, rows: np.ndarray, grid: Grid) -> CSRLinearExpression: def reindexed(self, grid: Grid, fill: float = np.nan) -> CSRLinearExpression: """ - Remap rows onto a new grid, possibly in a new dim order, without the - dense rectangle: dropped labels vanish, new labels get ``fill`` as - their constant (NaN: absent cells). Auxiliary coordinates follow the - rows; those of ``grid`` are ignored. + Remap rows onto a new grid, possibly in a new dim order: dropped + labels vanish, new labels get ``fill`` as their constant (NaN: absent + cells). Auxiliary coordinates follow the rows; those of ``grid`` are + ignored. """ row_map, valid = grid.indexer(self.grid) @@ -570,10 +572,9 @@ def added(self, other: CSRLinearExpression) -> CSRLinearExpression: Sparse matrix addition == merge along the term dimension. Goes through COO so explicit zero coefficients survive (scipy's ``+`` drops them), keeping a cell with only zero-coefficient terms distinguishable from - an empty cell, as on the dense path. A cell absent in either operand - is absent in the sum and carries no terms (v1 dead-term invariant). - Auxiliary coordinates propagate and conflicting ones raise (§11), as - on the dense path. + an empty cell. A cell absent in either operand is absent in the sum + and carries no terms. Auxiliary coordinates propagate and conflicting + ones raise (§11). """ const = self.const + other.const a, b = self.csr.tocoo(), other.csr.tocoo() @@ -605,11 +606,11 @@ def contracted( ``kron(I_kept, matrix.T) @ csr``, evaluated in chunks of the kept axis so the operator never grows with the kept size. - The result is compact canonical form: duplicate variables summed, terms - label-ordered and explicit zeros pruned -- unlike :meth:`added`, the - sparse product drops them, so cell activeness is carried by ``const`` - alone (see issue #925). Auxiliary coordinates on kept dims propagate, - those on contracted dims drop. + The result is in compact canonical form: duplicate variables summed, + terms label-ordered and explicit zeros pruned -- unlike :meth:`added`, + the sparse product drops them, so cell activeness is carried by + ``const`` alone. Auxiliary coordinates on kept dims propagate, those + on contracted dims drop. """ contracted_dims = tuple(contracted_dims) kept = tuple(d for d in self.grid.dims if d not in contracted_dims) @@ -656,8 +657,8 @@ def to_dense(self) -> LinearExpression: """ Expand to the dense equivalent in canonical form: terms label-ordered, duplicates summed, padded to the widest cell with the usual fill. - Absent cells (NaN const) carry no terms, per the v1 dead-term invariant. - The expanded dataset is wrapped in a :class:`LinearExpression`. + Absent cells (NaN const) carry no terms. The expanded dataset is + wrapped in a :class:`LinearExpression`. """ from linopy.expressions import LinearExpression diff --git a/linopy/expressions.py b/linopy/expressions.py index 9a4b148d..b4b19bc2 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -567,9 +567,9 @@ def sum( Parameters ---------- use_fallback : bool - Fall back to the previous, slower groupby-sum implementation, kept - as an escape hatch. Leave at False unless the default misbehaves. - Defaults to False. + Use the slower fallback groupby-sum implementation instead of the + default one. Kept as an escape hatch. Leave at False unless the + default misbehaves. Defaults to False. sparse : bool, optional Build the grouped sum in CSR form behind the ordinary LinearExpression type — no group-size padding; a still-sparse @@ -2556,14 +2556,20 @@ def _selected( return csr.taken(flat, grid) def sel(self, *args: Any, **kwargs: Any) -> LinearExpression: - """Select by label as ``Dataset.sel``; a CSR-backed expression stays sparse.""" + """ + Select by label as ``Dataset.sel``. For a CSR-backed expression, + returns a CSR-backed result when the selection stays on the grid. + """ csr = self._selected(lambda rows: rows.sel(*args, **kwargs), "sel") if csr is None: return super().sel(*args, **kwargs) return type(self)._from_csr(csr, self._model) def isel(self, *args: Any, **kwargs: Any) -> LinearExpression: - """Select by position as ``Dataset.isel``; a CSR-backed expression stays sparse.""" + """ + Select by position as ``Dataset.isel``. For a CSR-backed expression, + returns a CSR-backed result when the selection stays on the grid. + """ csr = self._selected(lambda rows: rows.isel(*args, **kwargs), "isel") if csr is None: return super().isel(*args, **kwargs) @@ -2822,7 +2828,8 @@ def __matmul__( ``(self * other).sum(dim)``. The result is then the compact canonical form -- duplicate variables summed, terms label-ordered, explicit zeros pruned -- so its term count may differ from the dense path's - while the values agree. A CSR-backed expression stays CSR-backed. + while the values agree. Returns a CSR-backed result when ``self`` is + CSR-backed. """ other = as_constant(other) other_is_const = not isinstance(other, LinearExpression | variables.Variable) @@ -2849,8 +2856,8 @@ def _sparse_matmul(self, other: ConstantLike) -> LinearExpression | None: labels, a zero-size grid, an operand sharing no dimension with the grid, and an unlabelled output dimension. The result is the compact canonical form of :meth:`CSRLinearExpression.contracted`, so its term - count may differ from the dense path's while the values agree; it - stays CSR-backed when the input was. + count may differ from the dense path's while the values agree; the + result is CSR-backed when the input was. """ if is_nan_scalar(other): check_user_nan(op_kind="mul") @@ -2924,8 +2931,9 @@ def reindex( **indexers_kwargs: Any, ) -> LinearExpression: """ - Conform to new coordinates as ``Dataset.reindex``; a CSR-backed - expression stays sparse when only labels change. + Conform to new coordinates as ``Dataset.reindex``. For a CSR-backed + expression, returns a CSR-backed result when only grid labels change + and no other keyword argument is given. """ indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex") csr = self._csr @@ -2953,8 +2961,8 @@ def rename( **names: Any, ) -> LinearExpression: """ - Rename dimensions as ``Dataset.rename``; a CSR-backed expression - stays sparse when only grid dims are relabelled. + Rename dimensions as ``Dataset.rename``. For a CSR-backed expression, + returns a CSR-backed result when only grid dims are relabelled. """ name_dict = either_dict_or_kwargs(name_dict, names, "rename") csr = self._csr @@ -3012,9 +3020,10 @@ def _csr_terms(self, csr: CSRLinearExpression) -> dict[str, np.ndarray]: def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: """ - Variable labels and coefficients of the stored terms, without a dense - round trip for a CSR backing; absent terms and zero coefficients are - dropped, duplicate labels are not summed. + Variable labels and coefficients of the stored terms. For a CSR-backed + expression, reads the CSR arrays directly; otherwise reads the dense + term arrays. Absent terms and zero coefficients are dropped, duplicate + labels are not summed. """ if self._csr is not None: csr = self._csr.csr @@ -3633,10 +3642,10 @@ def _aligned( ) -> list[CSRLinearExpression] | str: """ Conform the CSR expressions to the grid an explicit join produces, the - cells the join creates carrying ``fill`` as constant. None where the dense path - owns the semantics: ``exact`` and the auto-detected join raise there on - differing grids, ``override`` on differing shapes, any join on - non-unique labels; the reason is returned instead. + cells the join creates carrying ``fill`` as constant. Returns the reason + as a string where the dense path owns the semantics instead: ``exact`` + and the auto-detected join raise there on differing grids, ``override`` + on differing shapes, any join on non-unique labels. """ template = csrs[0].grid dims = template.dims @@ -3670,9 +3679,9 @@ def _try_csr_merge( the template order first. Grids that differ in their labels are aligned row-wise onto the joined grid, the cells the join creates carrying the fill of the dense path (zero, or NaN for ``fill_value=ABSENT``). Auxiliary - coordinates are checked for conflicts on the operands as given (§11, as - on the dense path) and follow their rows onto the joined grid. Returns - None to fall through to the dense path. + coordinates are checked for conflicts on the operands as given and follow + their rows onto the joined grid. Returns None to fall through to the + dense path. """ if not any(type(e) is LinearExpression and e._csr is not None for e in exprs): return None diff --git a/linopy/objective.py b/linopy/objective.py index ad7b6bdd..f3801ce9 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -164,8 +164,9 @@ def to_netcdf_ds(self) -> Dataset: """ Returns the objective in the dense layout for netcdf serialization. - A sparse objective is expanded without dropping its sparse backing; - after the setter's full sum it holds a single cell. + The expression setter reduces the objective to a single cell, so a + CSR-backed expression is densified into that one cell here, without + mutating the stored CSR data. """ expr = self.expression csr = expr._csr if isinstance(expr, expressions.LinearExpression) else None From 6785da3da36fe3c70819ead3c4e91fe8010fb2d7 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 24 Sep 2026 16:53:23 +0200 Subject: [PATCH 15/15] test(quadratic): narrow expression type in linear_terms test --- test/test_quadratic_expression.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/test_quadratic_expression.py b/test/test_quadratic_expression.py index fa9381b5..8355c5b3 100644 --- a/test/test_quadratic_expression.py +++ b/test/test_quadratic_expression.py @@ -301,6 +301,7 @@ def test_quadratic_expression_linear_terms( x: Variable, y: Variable, factor_last: bool ) -> None: expr = x * y + 3 * x + 0 * y + assert isinstance(expr, QuadraticExpression) if factor_last: expr = QuadraticExpression(expr.data.transpose(..., FACTOR_DIM), expr.model) labels, coeffs = expr.linear_terms()