From fb9f6abce7294b4ee2423d647ce21929339fa6fc Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Thu, 19 Mar 2026 00:43:10 +0100 Subject: [PATCH 01/20] perf: add to_matrix_via_csr --- linopy/constraints.py | 73 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/linopy/constraints.py b/linopy/constraints.py index d3ebef19..f52c4f5d 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -603,6 +603,46 @@ def mask_func(data: pd.DataFrame) -> pd.Series: check_has_nulls(df, name=f"{self.type} {self.name}") return df + def to_matrix(self) -> tuple[scipy.sparse.csr_matrix, np.ndarray]: + """ + Construct a CSR matrix representation of this constraint. + + Returns + ------- + matrix : scipy.sparse.csr_matrix + Shape (n_labels, model._xCounter). Rows correspond to individual + constraint labels, columns to variable labels. Missing entries + (labels or vars == -1) are excluded. + labels : np.ndarray + 1D array of shape (n_labels,) mapping each row back to the + original constraint label. + + Notes + ----- + Assumes that constraint labels are monotonuously increasing! + """ + coeffs_shape = self.coeffs.values.shape + broadcast_labels = np.ravel( + np.broadcast_to(self.labels.values[..., np.newaxis], coeffs_shape) + ) + vars_flat = self.vars.values.ravel() + coeffs_flat = self.coeffs.values.ravel() + + valid = (broadcast_labels != -1) & (vars_flat != -1) + broadcast_labels = broadcast_labels[valid] + vars_flat = vars_flat[valid] + coeffs_flat = coeffs_flat[valid] + + changes = np.r_[True, broadcast_labels[1:] != broadcast_labels[:-1]] + con_labels = broadcast_labels[changes] + indptr = np.r_[np.nonzero(changes)[0], len(broadcast_labels)] + + shape = (len(con_labels), self.model._xCounter) + # Note: duplicate (row, col) entries are not summed in CSR format. + # They will be summed automatically upon conversion to CSC or dense. + matrix = scipy.sparse.csr_matrix((coeffs_flat, vars_flat, indptr), shape=shape) + return matrix, con_labels + def to_polars(self) -> pl.DataFrame: """ Convert the constraint to a polars DataFrame. @@ -1117,6 +1157,39 @@ def to_matrix(self, filter_missings: bool = True) -> scipy.sparse.csc_matrix: (cons.coeffs, (cons.labels, cons.vars)), shape=shape ) + def to_matrix_via_csr( + self, + ) -> tuple[scipy.sparse.csc_matrix, np.ndarray, np.ndarray]: + """ + Construct a constraint matrix in sparse format by stacking per-constraint CSR matrices. + + Returns + ------- + matrix : scipy.sparse.csc_matrix + Shape (n_con_labels, n_var_labels), containing only non-empty rows and columns. + con_labels : np.ndarray + Shape (n_con_labels,), maps each matrix row to the original constraint label. + var_labels : np.ndarray + Shape (n_var_labels,), maps each matrix column to the original variable label. + """ + if not len(self): + raise ValueError("No constraints available to convert to matrix.") + + matrices, con_labels_list = zip(*(c.to_matrix() for c in self.data.values())) + csc: scipy.sparse.csc_matrix = scipy.sparse.vstack(matrices).tocsc() + con_labels = np.concatenate(con_labels_list) + + indptr = csc.indptr + nonempty_cols = indptr[1:] != indptr[:-1] + new_indptr = np.r_[0, indptr[1:][nonempty_cols]] + (var_labels,) = np.nonzero(nonempty_cols) + + matrix = scipy.sparse.csc_matrix( + (csc.data, csc.indices, new_indptr), + shape=(csc.shape[0], len(var_labels)), + ) + return matrix, con_labels, var_labels + def reset_dual(self) -> None: """ Reset the stored solution of variables. From 80a8b70438700a750cc66f57cfcd054c62b76bd4 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Thu, 19 Mar 2026 01:23:20 +0100 Subject: [PATCH 02/20] perf: improve per-constraint csr matrix construction --- linopy/constraints.py | 30 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index f52c4f5d..daecad39 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -616,26 +616,24 @@ def to_matrix(self) -> tuple[scipy.sparse.csr_matrix, np.ndarray]: labels : np.ndarray 1D array of shape (n_labels,) mapping each row back to the original constraint label. - - Notes - ----- - Assumes that constraint labels are monotonuously increasing! """ - coeffs_shape = self.coeffs.values.shape - broadcast_labels = np.ravel( - np.broadcast_to(self.labels.values[..., np.newaxis], coeffs_shape) - ) - vars_flat = self.vars.values.ravel() + vars = self.vars.values + labels_flat = self.labels.values.ravel() + vars_2d = vars.reshape(len(labels_flat), -1) coeffs_flat = self.coeffs.values.ravel() - valid = (broadcast_labels != -1) & (vars_flat != -1) - broadcast_labels = broadcast_labels[valid] - vars_flat = vars_flat[valid] - coeffs_flat = coeffs_flat[valid] + valid_vars = vars_2d != -1 + labels_valid = (labels_flat != -1) & valid_vars.any(axis=1) + valid = (labels_valid[:, np.newaxis] & valid_vars).ravel() - changes = np.r_[True, broadcast_labels[1:] != broadcast_labels[:-1]] - con_labels = broadcast_labels[changes] - indptr = np.r_[np.nonzero(changes)[0], len(broadcast_labels)] + con_labels = labels_flat[labels_valid] + counts = valid_vars.sum(axis=1)[labels_valid] + indptr = np.empty(len(con_labels) + 1, dtype=np.int32) + indptr[0] = 0 + np.cumsum(counts, out=indptr[1:]) + + vars_flat = vars.ravel()[valid] + coeffs_flat = coeffs_flat[valid] shape = (len(con_labels), self.model._xCounter) # Note: duplicate (row, col) entries are not summed in CSR format. From 89115c2a2cea76911018d5c8d1f232d704c40bb7 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Thu, 19 Mar 2026 11:33:18 +0100 Subject: [PATCH 03/20] Add conversion functions --- linopy/constraints.py | 164 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 157 insertions(+), 7 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index daecad39..13696102 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -73,6 +73,125 @@ if TYPE_CHECKING: from linopy.model import Model + +def csr_to_constraint_dataset( + csr: scipy.sparse.csr_array, + con_labels: np.ndarray, + coords: list[pd.Index], + nterm: int, +) -> Dataset: + """ + Reconstruct a constraint Dataset (labels, coeffs, vars) from a CSR matrix. + + Parameters + ---------- + csr : scipy.sparse.csr_array + Shape (n_con_labels, n_vars). Rows correspond to non-masked constraints. + con_labels : np.ndarray + 0-based flat indices into the constraint grid (*coord_dims) for each row. + coords : list of pd.Index + One pd.Index per coordinate dimension, defining the constraint grid shape. + nterm : int + Size of the term dimension in the reconstructed Dataset. + + Returns + ------- + Dataset with variables 'labels', 'coeffs', 'vars' and dimensions + (*coord_dims, _term). + """ + shape = tuple(len(c) for c in coords) + n_total = int(np.prod(shape)) if shape else 1 + dim_names = [c.name for c in coords] + + labels_flat = np.full(n_total, -1, dtype=np.int64) + labels_flat[con_labels] = con_labels + + coeffs_2d = np.zeros((n_total, nterm), dtype=csr.dtype) + vars_2d = np.full((n_total, nterm), -1, dtype=np.int64) + + # For each entry in csr.data/indices, compute its (row_idx, term_col) position. + # term_col is the within-row offset (0, 1, 2, ...) for each entry. + counts = np.diff(csr.indptr) # number of entries per CSR row + term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) + row_indices = con_labels[np.repeat(np.arange(len(con_labels)), counts)] + + vars_2d[row_indices, term_cols] = csr.indices + coeffs_2d[row_indices, term_cols] = csr.data + + xr_coords = {c.name: c for c in coords} + labels_da = DataArray(labels_flat.reshape(shape), coords=xr_coords, dims=dim_names) + coeffs_da = DataArray( + coeffs_2d.reshape(shape + (nterm,)), + coords=xr_coords, + dims=dim_names + [TERM_DIM], + ) + vars_da = DataArray( + vars_2d.reshape(shape + (nterm,)), + coords=xr_coords, + dims=dim_names + [TERM_DIM], + ) + + return Dataset({"labels": labels_da, "coeffs": coeffs_da, "vars": vars_da}) + + +def csr2_to_constraint_dataset( + csr: scipy.sparse.csr_array, + coords: list[pd.Index], + nterm: int, +) -> Dataset: + """ + Reconstruct a constraint Dataset (labels, coeffs, vars) from a CSR matrix + where all flat positions are rows (inverse of Constraint.to_matrix2). + + Parameters + ---------- + csr : scipy.sparse.csr_array + Shape (n_total, n_vars). Each row corresponds to a flat position in the + constraint grid, including masked (empty) positions. + coords : list of pd.Index + One pd.Index per coordinate dimension, defining the constraint grid shape. + nterm : int + Size of the term dimension in the reconstructed Dataset. + + Returns + ------- + Dataset with variables 'labels', 'coeffs', 'vars' and dimensions + (*coord_dims, _term). + """ + shape = tuple(len(c) for c in coords) + n_total = csr.shape[0] + dim_names = [c.name for c in coords] + + counts = np.diff(csr.indptr) + nonempty = counts > 0 + + labels_flat = np.where(nonempty, np.arange(n_total), -1) + + coeffs_2d = np.zeros((n_total, nterm), dtype=csr.dtype) + vars_2d = np.full((n_total, nterm), -1, dtype=np.int64) + + term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) + row_indices = np.repeat(np.arange(n_total), counts) + + vars_2d[row_indices, term_cols] = csr.indices + coeffs_2d[row_indices, term_cols] = csr.data + + xr_coords = {c.name: c for c in coords} + labels_da = DataArray(labels_flat.reshape(shape), coords=xr_coords, dims=dim_names) + coeffs_da = DataArray( + coeffs_2d.reshape(shape + (nterm,)), + coords=xr_coords, + dims=dim_names + [TERM_DIM], + ) + vars_da = DataArray( + vars_2d.reshape(shape + (nterm,)), + coords=xr_coords, + dims=dim_names + [TERM_DIM], + ) + + return Dataset({"labels": labels_da, "coeffs": coeffs_da, "vars": vars_da}) + + FILL_VALUE = {"labels": -1, "rhs": np.nan, "coeffs": 0, "vars": -1, "sign": "="} @@ -603,13 +722,13 @@ def mask_func(data: pd.DataFrame) -> pd.Series: check_has_nulls(df, name=f"{self.type} {self.name}") return df - def to_matrix(self) -> tuple[scipy.sparse.csr_matrix, np.ndarray]: + def to_matrix(self) -> tuple[scipy.sparse.csr_array, np.ndarray]: """ Construct a CSR matrix representation of this constraint. Returns ------- - matrix : scipy.sparse.csr_matrix + matrix : scipy.sparse.csr_array Shape (n_labels, model._xCounter). Rows correspond to individual constraint labels, columns to variable labels. Missing entries (labels or vars == -1) are excluded. @@ -638,9 +757,40 @@ def to_matrix(self) -> tuple[scipy.sparse.csr_matrix, np.ndarray]: shape = (len(con_labels), self.model._xCounter) # Note: duplicate (row, col) entries are not summed in CSR format. # They will be summed automatically upon conversion to CSC or dense. - matrix = scipy.sparse.csr_matrix((coeffs_flat, vars_flat, indptr), shape=shape) + matrix = scipy.sparse.csr_array((coeffs_flat, vars_flat, indptr), shape=shape) return matrix, con_labels + def to_matrix2(self) -> scipy.sparse.csr_array: + """ + Construct a CSR matrix representation of this constraint including masked rows. + + Unlike to_matrix(), all flat positions in the constraint grid are included as + rows (masked entries become empty rows). Shape is + (len(labels_flat), model._xCounter). + + Returns + ------- + matrix : scipy.sparse.csr_array + """ + vars = self.vars.values + labels_flat = self.labels.values.ravel() + vars_2d = vars.reshape(len(labels_flat), -1) + coeffs_flat = self.coeffs.values.ravel() + + valid_2d = (labels_flat != -1)[:, np.newaxis] & (vars_2d != -1) + cols = vars_2d[valid_2d] + data = coeffs_flat.reshape(vars_2d.shape)[valid_2d] + + counts = valid_2d.sum(axis=1) + indptr = np.empty(len(labels_flat) + 1, dtype=np.int32) + indptr[0] = 0 + np.cumsum(counts, out=indptr[1:]) + + shape = (len(labels_flat), self.model._xCounter) + # Note: duplicate (row, col) entries are not summed in CSR format. + # They will be summed automatically upon conversion to CSC or dense. + return scipy.sparse.csr_array((data, cols, indptr), shape=shape) + def to_polars(self) -> pl.DataFrame: """ Convert the constraint to a polars DataFrame. @@ -1157,13 +1307,13 @@ def to_matrix(self, filter_missings: bool = True) -> scipy.sparse.csc_matrix: def to_matrix_via_csr( self, - ) -> tuple[scipy.sparse.csc_matrix, np.ndarray, np.ndarray]: + ) -> tuple[scipy.sparse.csc_array, np.ndarray, np.ndarray]: """ Construct a constraint matrix in sparse format by stacking per-constraint CSR matrices. Returns ------- - matrix : scipy.sparse.csc_matrix + matrix : scipy.sparse.csc_array Shape (n_con_labels, n_var_labels), containing only non-empty rows and columns. con_labels : np.ndarray Shape (n_con_labels,), maps each matrix row to the original constraint label. @@ -1174,7 +1324,7 @@ def to_matrix_via_csr( raise ValueError("No constraints available to convert to matrix.") matrices, con_labels_list = zip(*(c.to_matrix() for c in self.data.values())) - csc: scipy.sparse.csc_matrix = scipy.sparse.vstack(matrices).tocsc() + csc: scipy.sparse.csc_array = scipy.sparse.vstack(matrices).tocsc() con_labels = np.concatenate(con_labels_list) indptr = csc.indptr @@ -1182,7 +1332,7 @@ def to_matrix_via_csr( new_indptr = np.r_[0, indptr[1:][nonempty_cols]] (var_labels,) = np.nonzero(nonempty_cols) - matrix = scipy.sparse.csc_matrix( + matrix = scipy.sparse.csc_array( (csc.data, csc.indices, new_indptr), shape=(csc.shape[0], len(var_labels)), ) From bcd0228796114b2242dc3de61de99d2370868d21 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Thu, 19 Mar 2026 15:35:42 +0100 Subject: [PATCH 04/20] feat: add ability to freeze constraints into csr --- linopy/__init__.py | 12 +- linopy/common.py | 10 +- linopy/constants.py | 5 + linopy/constraints.py | 1190 +++++++++++++++++++------------- linopy/expressions.py | 33 +- linopy/io.py | 4 +- linopy/model.py | 29 +- linopy/piecewise.py | 22 +- linopy/variables.py | 26 +- test/test_constraint.py | 49 +- test/test_scalar_constraint.py | 18 +- 11 files changed, 824 insertions(+), 574 deletions(-) diff --git a/linopy/__init__.py b/linopy/__init__.py index b1dc33b9..12a1cb9e 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -14,8 +14,13 @@ import linopy.monkey_patch_xarray # noqa: F401 from linopy.common import align from linopy.config import options -from linopy.constants import EQUAL, GREATER_EQUAL, LESS_EQUAL -from linopy.constraints import Constraint, Constraints +from linopy.constants import EQUAL, GREATER_EQUAL, LESS_EQUAL, PerformanceWarning +from linopy.constraints import ( + Constraint, + ConstraintBase, + Constraints, + MutableConstraint, +) from linopy.expressions import LinearExpression, QuadraticExpression, merge from linopy.io import read_netcdf from linopy.model import Model, Variable, Variables, available_solvers @@ -30,8 +35,11 @@ __all__ = ( "Constraint", + "ConstraintBase", "Constraints", + "MutableConstraint", "EQUAL", + "PerformanceWarning", "GREATER_EQUAL", "LESS_EQUAL", "LinearExpression", diff --git a/linopy/common.py b/linopy/common.py index 09f67355..ac0b3af3 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -40,7 +40,7 @@ ) if TYPE_CHECKING: - from linopy.constraints import Constraint + from linopy.constraints import ConstraintBase from linopy.expressions import LinearExpression, QuadraticExpression from linopy.variables import Variable @@ -554,7 +554,7 @@ def fill_missing_coords( return ds -T = TypeVar("T", Dataset, "Variable", "LinearExpression", "Constraint") +T = TypeVar("T", Dataset, "Variable", "LinearExpression", "ConstraintBase") @overload @@ -583,10 +583,10 @@ def iterate_slices( @overload def iterate_slices( - ds: Constraint, + ds: ConstraintBase, slice_size: int | None = 10_000, slice_dims: list | None = None, -) -> Generator[Constraint, None, None]: ... +) -> Generator[ConstraintBase, None, None]: ... def iterate_slices( @@ -1306,7 +1306,7 @@ def align( "Variable", "LinearExpression", "QuadraticExpression", - "Constraint", + "ConstraintBase", ) diff --git a/linopy/constants.py b/linopy/constants.py index 00bbd705..96957be1 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -18,6 +18,11 @@ GREATER_EQUAL = ">=" LESS_EQUAL = "<=" + +class PerformanceWarning(UserWarning): + """Warning raised when an operation triggers expensive Dataset reconstruction.""" + + long_EQUAL = "==" short_GREATER_EQUAL = ">" short_LESS_EQUAL = "<" diff --git a/linopy/constraints.py b/linopy/constraints.py index 13696102..8d0fd783 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -7,6 +7,8 @@ from __future__ import annotations import functools +import warnings +from abc import ABC, abstractmethod from collections.abc import Callable, Hashable, ItemsView, Iterator, Sequence from dataclasses import dataclass from itertools import product @@ -60,6 +62,7 @@ HELPER_DIMS, LESS_EQUAL, TERM_DIM, + PerformanceWarning, SIGNS_pretty, ) from linopy.types import ( @@ -74,124 +77,6 @@ from linopy.model import Model -def csr_to_constraint_dataset( - csr: scipy.sparse.csr_array, - con_labels: np.ndarray, - coords: list[pd.Index], - nterm: int, -) -> Dataset: - """ - Reconstruct a constraint Dataset (labels, coeffs, vars) from a CSR matrix. - - Parameters - ---------- - csr : scipy.sparse.csr_array - Shape (n_con_labels, n_vars). Rows correspond to non-masked constraints. - con_labels : np.ndarray - 0-based flat indices into the constraint grid (*coord_dims) for each row. - coords : list of pd.Index - One pd.Index per coordinate dimension, defining the constraint grid shape. - nterm : int - Size of the term dimension in the reconstructed Dataset. - - Returns - ------- - Dataset with variables 'labels', 'coeffs', 'vars' and dimensions - (*coord_dims, _term). - """ - shape = tuple(len(c) for c in coords) - n_total = int(np.prod(shape)) if shape else 1 - dim_names = [c.name for c in coords] - - labels_flat = np.full(n_total, -1, dtype=np.int64) - labels_flat[con_labels] = con_labels - - coeffs_2d = np.zeros((n_total, nterm), dtype=csr.dtype) - vars_2d = np.full((n_total, nterm), -1, dtype=np.int64) - - # For each entry in csr.data/indices, compute its (row_idx, term_col) position. - # term_col is the within-row offset (0, 1, 2, ...) for each entry. - counts = np.diff(csr.indptr) # number of entries per CSR row - term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) - row_indices = con_labels[np.repeat(np.arange(len(con_labels)), counts)] - - vars_2d[row_indices, term_cols] = csr.indices - coeffs_2d[row_indices, term_cols] = csr.data - - xr_coords = {c.name: c for c in coords} - labels_da = DataArray(labels_flat.reshape(shape), coords=xr_coords, dims=dim_names) - coeffs_da = DataArray( - coeffs_2d.reshape(shape + (nterm,)), - coords=xr_coords, - dims=dim_names + [TERM_DIM], - ) - vars_da = DataArray( - vars_2d.reshape(shape + (nterm,)), - coords=xr_coords, - dims=dim_names + [TERM_DIM], - ) - - return Dataset({"labels": labels_da, "coeffs": coeffs_da, "vars": vars_da}) - - -def csr2_to_constraint_dataset( - csr: scipy.sparse.csr_array, - coords: list[pd.Index], - nterm: int, -) -> Dataset: - """ - Reconstruct a constraint Dataset (labels, coeffs, vars) from a CSR matrix - where all flat positions are rows (inverse of Constraint.to_matrix2). - - Parameters - ---------- - csr : scipy.sparse.csr_array - Shape (n_total, n_vars). Each row corresponds to a flat position in the - constraint grid, including masked (empty) positions. - coords : list of pd.Index - One pd.Index per coordinate dimension, defining the constraint grid shape. - nterm : int - Size of the term dimension in the reconstructed Dataset. - - Returns - ------- - Dataset with variables 'labels', 'coeffs', 'vars' and dimensions - (*coord_dims, _term). - """ - shape = tuple(len(c) for c in coords) - n_total = csr.shape[0] - dim_names = [c.name for c in coords] - - counts = np.diff(csr.indptr) - nonempty = counts > 0 - - labels_flat = np.where(nonempty, np.arange(n_total), -1) - - coeffs_2d = np.zeros((n_total, nterm), dtype=csr.dtype) - vars_2d = np.full((n_total, nterm), -1, dtype=np.int64) - - term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) - row_indices = np.repeat(np.arange(n_total), counts) - - vars_2d[row_indices, term_cols] = csr.indices - coeffs_2d[row_indices, term_cols] = csr.data - - xr_coords = {c.name: c for c in coords} - labels_da = DataArray(labels_flat.reshape(shape), coords=xr_coords, dims=dim_names) - coeffs_da = DataArray( - coeffs_2d.reshape(shape + (nterm,)), - coords=xr_coords, - dims=dim_names + [TERM_DIM], - ) - vars_da = DataArray( - vars_2d.reshape(shape + (nterm,)), - coords=xr_coords, - dims=dim_names + [TERM_DIM], - ) - - return Dataset({"labels": labels_da, "coeffs": coeffs_da, "vars": vars_da}) - - FILL_VALUE = {"labels": -1, "rhs": np.nan, "coeffs": 0, "vars": -1, "sign": "="} @@ -199,7 +84,9 @@ def conwrap( method: Callable, *default_args: Any, **new_default_kwargs: Any ) -> Callable: @functools.wraps(method) - def _conwrap(con: Constraint, *args: Any, **kwargs: Any) -> Constraint: + def _conwrap( + con: MutableConstraint, *args: Any, **kwargs: Any + ) -> MutableConstraint: for k, v in new_default_kwargs.items(): kwargs.setdefault(k, v) return con.__class__( @@ -215,199 +102,192 @@ def _conwrap(con: Constraint, *args: Any, **kwargs: Any) -> Constraint: return _conwrap -def _con_unwrap(con: Constraint | Dataset) -> Dataset: - return con.data if isinstance(con, Constraint) else con +def _con_unwrap(con: ConstraintBase | Dataset) -> Dataset: + return con.data if isinstance(con, ConstraintBase) else con -class Constraint: +class ConstraintBase(ABC): """ - Projection to a single constraint in a model. + Abstract base class for Constraint and MutableConstraint. - The Constraint class is a subclass of xr.DataArray hence most xarray - functions can be applied to it. + Provides all read-only properties and methods shared by both the immutable + Constraint (CSR-backed) and the mutable MutableConstraint (Dataset-backed). """ - __slots__ = ("_data", "_model", "_assigned") - _fill_value = FILL_VALUE - def __init__( - self, - data: Dataset, - model: Model, - name: str = "", - skip_broadcast: bool = False, - ) -> None: - """ - Initialize the Constraint. + @property + @abstractmethod + def data(self) -> Dataset: + """Get the underlying xarray Dataset representation.""" - Parameters - ---------- - labels : xarray.DataArray - labels of the constraint. - model : linopy.Model - Underlying model. - name : str - Name of the constraint. - """ + @property + @abstractmethod + def model(self) -> Model: + """Get the model reference.""" - from linopy.model import Model + @property + @abstractmethod + def name(self) -> str: + """Get the constraint name.""" - if not isinstance(data, Dataset): - raise ValueError(f"data must be a Dataset, got {type(data)}") + @property + @abstractmethod + def is_assigned(self) -> bool: + """Whether the constraint has been assigned labels by the model.""" - if not isinstance(model, Model): - raise ValueError(f"model must be a Model, got {type(model)}") + @property + @abstractmethod + def labels(self) -> DataArray: + """Get the labels DataArray.""" - # check that `labels`, `lower` and `upper`, `sign` and `mask` are in data - for attr in ("coeffs", "vars", "sign", "rhs"): - if attr not in data: - raise ValueError(f"missing '{attr}' in data") + @property + @abstractmethod + def coeffs(self) -> DataArray: + """Get the LHS coefficients DataArray.""" - data = data.assign_attrs(name=name) + @property + @abstractmethod + def vars(self) -> DataArray: + """Get the LHS variable labels DataArray.""" - if not skip_broadcast: - (data,) = xr.broadcast(data, exclude=[TERM_DIM]) + @property + @abstractmethod + def sign(self) -> DataArray: + """Get the constraint sign DataArray.""" - self._assigned = "labels" in data - self._data = data - self._model = model + @property + @abstractmethod + def rhs(self) -> DataArray: + """Get the RHS DataArray.""" + + @property + @abstractmethod + def dual(self) -> DataArray: + """Get the dual values DataArray.""" def __getitem__( self, selector: str | int | slice | list | tuple | dict - ) -> Constraint: + ) -> MutableConstraint: """ Get selection from the constraint. - This is a wrapper around the xarray __getitem__ method. It returns a - new object with the selected data. + Returns a MutableConstraint with the selected data. """ - data = Dataset({k: self.data[k][selector] for k in self.data}, attrs=self.attrs) - return self.__class__(data, self.model, self.name) + data = Dataset( + {k: self.data[k][selector] for k in self.data}, attrs=self.data.attrs + ) + return MutableConstraint(data, self.model, self.name) @property def attrs(self) -> dict[str, Any]: - """ - Get the attributes of the constraint. - """ + """Get the attributes of the constraint.""" return self.data.attrs @property def coords(self) -> DatasetCoordinates: - """ - Get the coordinates of the constraint. - """ + """Get the coordinates of the constraint.""" return self.data.coords @property def indexes(self) -> Indexes: - """ - Get the indexes of the constraint. - """ + """Get the indexes of the constraint.""" return self.data.indexes @property def dims(self) -> Frozen[Hashable, int]: - """ - Get the dimensions of the constraint. - """ + """Get the dimensions of the constraint.""" return self.data.dims @property def sizes(self) -> Frozen[Hashable, int]: - """ - Get the sizes of the constraint. - """ + """Get the sizes of the constraint.""" return self.data.sizes @property def nterm(self) -> int: - """ - Get the number of terms in the constraint. - """ - return self.lhs.nterm + """Get the number of terms in the constraint.""" + return self.data.sizes.get(TERM_DIM, 1) @property def ndim(self) -> int: - """ - Get the number of dimensions of the constraint. - """ + """Get the number of dimensions of the constraint.""" return self.rhs.ndim @property def shape(self) -> tuple[int, ...]: - """ - Get the shape of the constraint. - """ + """Get the shape of the constraint.""" return self.rhs.shape @property def size(self) -> int: - """ - Get the size of the constraint. - """ + """Get the size of the constraint.""" return self.rhs.size @property - def loc(self) -> LocIndexer: - return LocIndexer(self) - - @property - def data(self) -> Dataset: + def ncons(self) -> int: """ - Get the underlying DataArray. + Get the number of active constraints (non-masked, with at least one valid variable). """ - return self._data + labels = self.labels.values + vars_arr = self.vars.values + if labels.ndim == 0: + return int(labels != FILL_VALUE["labels"] and (vars_arr != -1).any()) + return int( + ((labels != FILL_VALUE["labels"]) & (vars_arr != -1).any(axis=-1)).sum() + ) @property - def labels(self) -> DataArray: - """ - Get the labels of the constraint. - """ - return self.data.get("labels", DataArray([])) + def coord_dims(self) -> tuple[Hashable, ...]: + return tuple(k for k in self.dims if k not in HELPER_DIMS) @property - def model(self) -> Model: - """ - Get the model of the constraint. - """ - return self._model + def coord_sizes(self) -> dict[Hashable, int]: + return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS} @property - def name(self) -> str: - """ - Return the name of the constraint. - """ - return self.attrs["name"] + def coord_names(self) -> list[str]: + """Get the names of the coordinates.""" + return get_dims_with_index_levels(self.data, self.coord_dims) @property - def coord_dims(self) -> tuple[Hashable, ...]: - return tuple(k for k in self.dims if k not in HELPER_DIMS) + def type(self) -> str: + """Get the type string of the constraint.""" + return "Constraint" if self.is_assigned else "Constraint (unassigned)" @property - def coord_sizes(self) -> dict[Hashable, int]: - return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS} + def term_dim(self) -> str: + """Return the term dimension of the constraint.""" + return TERM_DIM @property - def coord_names(self) -> list[str]: + def mask(self) -> DataArray | None: """ - Get the names of the coordinates. + Get the mask of the constraint. + + The mask indicates on which coordinates the constraint is enabled + (True) and disabled (False). """ - return get_dims_with_index_levels(self.data, self.coord_dims) + if self.is_assigned: + return (self.labels != FILL_VALUE["labels"]).astype(bool) + return None @property - def is_assigned(self) -> bool: - return self._assigned + def lhs(self) -> expressions.LinearExpression: + """Get the left-hand-side linear expression of the constraint.""" + data = self.data[["coeffs", "vars"]].rename({self.term_dim: TERM_DIM}) + return expressions.LinearExpression(data, self.model) + + def __contains__(self, value: Any) -> bool: + return self.data.__contains__(value) def __repr__(self) -> str: - """ - Print the constraint arrays. - """ + """Print the constraint arrays.""" max_lines = options["display_max_rows"] dims = list(self.coord_sizes.keys()) ndim = len(dims) dim_names = self.coord_names dim_sizes = list(self.coord_sizes.values()) - size = np.prod(dim_sizes) # that the number of theoretical printouts + size = np.prod(dim_sizes) masked_entries = (~self.mask).sum().values if self.mask is not None else 0 lines = [] @@ -443,81 +323,556 @@ def __repr__(self) -> str: lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") elif size == 1: expr = print_single_expression( - self.coeffs.values, self.vars.values, 0, self.model + self.coeffs.values, self.vars.values, 0, self.model + ) + lines.append( + f"{header_string}\n{'-' * len(header_string)}\n{expr} {SIGNS_pretty[self.sign.item()]} {self.rhs.item()}" + ) + else: + lines.append(f"{header_string}\n{'-' * len(header_string)}\n") + + return "\n".join(lines) + + def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None: + """ + Print the linear expression. + + Parameters + ---------- + display_max_rows : int + Maximum number of rows to be displayed. + display_max_terms : int + Maximum number of terms to be displayed. + """ + with options as opts: + opts.set_value( + display_max_rows=display_max_rows, display_max_terms=display_max_terms + ) + print(self) + + @property + def flat(self) -> pd.DataFrame: + """ + Convert the constraint to a pandas DataFrame. + + The resulting DataFrame represents a long table format of the all + non-masked constraints with non-zero coefficients. It contains the + columns `labels`, `coeffs`, `vars`, `rhs`, `sign`. + """ + ds = self.data + + def mask_func(data: pd.DataFrame) -> pd.Series: + mask = (data["vars"] != -1) & (data["coeffs"] != 0) + if "labels" in data: + mask &= data["labels"] != -1 + return mask + + df = to_dataframe(ds, mask_func=mask_func) + + # Group repeated variables in the same constraint + agg_custom = {k: "first" for k in list(df.columns)} + agg_standards = dict(coeffs="sum", rhs="first", sign="first") + agg = {**agg_custom, **agg_standards} + df = df.groupby(["labels", "vars"], as_index=False).aggregate(agg) + check_has_nulls(df, name=f"{self.type} {self.name}") + return df + + def to_matrix(self) -> scipy.sparse.csr_array: + """ + Construct a CSR matrix representation of this constraint. + + All flat positions in the constraint grid are included as rows; + masked entries (labels == -1) become empty rows. Shape is + (len(labels_flat), model._xCounter). + + Returns + ------- + matrix : scipy.sparse.csr_array + """ + vars = self.vars.values + labels_flat = self.labels.values.ravel() + vars_2d = vars.reshape(len(labels_flat), -1) + coeffs_flat = self.coeffs.values.ravel() + + valid_2d = (labels_flat != -1)[:, np.newaxis] & (vars_2d != -1) + cols = vars_2d[valid_2d] + data = coeffs_flat.reshape(vars_2d.shape)[valid_2d] + + counts = valid_2d.sum(axis=1) + indptr = np.empty(len(labels_flat) + 1, dtype=np.int32) + indptr[0] = 0 + np.cumsum(counts, out=indptr[1:]) + + shape = (len(labels_flat), self.model._xCounter) + # Note: duplicate (row, col) entries are not summed in CSR format. + # They will be summed automatically upon conversion to CSC or dense. + return scipy.sparse.csr_array((data, cols, indptr), shape=shape) + + iterate_slices = iterate_slices + + +class Constraint(ConstraintBase): + """ + Immutable constraint backed by a CSR sparse matrix. + + Parameters + ---------- + csr : scipy.sparse.csr_array + Shape (n_flat, model._xCounter). Each row is a flat position in the + constraint grid (including masked/empty rows). + rhs : np.ndarray + Shape (n_flat,). Right-hand-side values. + sign : str + Constraint sign: one of '=', '<=', '>='. + Note: per-element signs are not supported (documented regression vs MutableConstraint). + coords : list of pd.Index + One index per coordinate dimension defining the constraint grid. + model : Model + The linopy model this constraint belongs to. + name : str + Name of the constraint. + cindex : int or None + Starting label assigned by the model. None if not yet assigned. + dual : np.ndarray or None + Shape (n_flat,). Dual values after solving, or None. + """ + + __slots__ = ( + "_csr", + "_rhs", + "_sign", + "_coords", + "_model", + "_name", + "_cindex", + "_dual", + ) + + def __init__( + self, + csr: scipy.sparse.csr_array, + rhs: np.ndarray, + sign: str, + coords: list[pd.Index], + model: Model, + name: str = "", + cindex: int | None = None, + dual: np.ndarray | None = None, + ) -> None: + self._csr = csr + self._rhs = rhs + self._sign = sign + self._coords = coords + self._model = model + self._name = name + self._cindex = cindex + self._dual = dual + + @property + def model(self) -> Model: + return self._model + + @property + def name(self) -> str: + return self._name + + @property + def is_assigned(self) -> bool: + return self._cindex is not None + + @property + def range(self) -> tuple[int, int]: + """Return the (start, end) label range of the constraint.""" + if self._cindex is None: + raise AttributeError("Constraint has not been assigned labels yet.") + return (self._cindex, self._cindex + self._csr.shape[0]) + + @property + def _nonempty(self) -> np.ndarray: + """Boolean array of shape (n_flat,) — True where row is non-masked.""" + return np.diff(self._csr.indptr).astype(bool) + + @property + def ncons(self) -> int: + return int(np.diff(self._csr.indptr).astype(bool).sum()) + + @property + def attrs(self) -> dict[str, Any]: + d: dict[str, Any] = {"name": self._name} + if self._cindex is not None: + d["label_range"] = (self._cindex, self._cindex + self._csr.shape[0]) + return d + + @property + def dims(self) -> Frozen[Hashable, int]: + d: dict[Hashable, int] = {c.name: len(c) for c in self._coords} + nterm = int(self._csr.indptr.max()) if self._csr.nnz > 0 else 1 + d[TERM_DIM] = nterm + return Frozen(d) + + @property + def sizes(self) -> Frozen[Hashable, int]: + return self.dims + + @property + def indexes(self) -> Indexes: + return Indexes({c.name: c for c in self._coords}) + + @property + def nterm(self) -> int: + return int(self._csr.indptr.max()) if self._csr.nnz > 0 else 1 + + @property + def coord_names(self) -> list[str]: + return [c.name for c in self._coords] + + @property + def labels(self) -> DataArray: + """Get labels DataArray, shape (*coord_dims).""" + if self._cindex is None: + return DataArray([]) + n_flat = self._csr.shape[0] + labels_flat = np.where( + self._nonempty, + np.arange(self._cindex, self._cindex + n_flat), + -1, + ) + shape = tuple(len(c) for c in self._coords) + dim_names = [c.name for c in self._coords] + xr_coords = {c.name: c for c in self._coords} + return DataArray( + labels_flat.reshape(shape) if shape else labels_flat, + coords=xr_coords, + dims=dim_names, + ) + + @property + def coeffs(self) -> DataArray: + """Get coefficients DataArray, shape (*coord_dims, _term).""" + warnings.warn( + "Accessing .coeffs on a Constraint triggers full Dataset reconstruction. " + "Use .to_matrix() for efficient access.", + PerformanceWarning, + stacklevel=2, + ) + return self.data.coeffs + + @property + def vars(self) -> DataArray: + """Get variable labels DataArray, shape (*coord_dims, _term).""" + warnings.warn( + "Accessing .vars on a Constraint triggers full Dataset reconstruction. " + "Use .to_matrix() for efficient access.", + PerformanceWarning, + stacklevel=2, + ) + return self.data.vars + + @property + def sign(self) -> DataArray: + """Get sign DataArray (scalar, same sign for all entries).""" + shape = tuple(len(c) for c in self._coords) + dim_names = [c.name for c in self._coords] + xr_coords = {c.name: c for c in self._coords} + return DataArray( + np.full(shape, self._sign) if shape else np.array(self._sign), + coords=xr_coords, + dims=dim_names, + ) + + @property + def rhs(self) -> DataArray: + """Get RHS DataArray, shape (*coord_dims).""" + shape = tuple(len(c) for c in self._coords) + dim_names = [c.name for c in self._coords] + xr_coords = {c.name: c for c in self._coords} + return DataArray( + self._rhs.reshape(shape) if shape else self._rhs, + coords=xr_coords, + dims=dim_names, + ) + + @property + @has_optimized_model + def dual(self) -> DataArray: + """Get dual values DataArray, shape (*coord_dims).""" + if self._dual is None: + raise AttributeError( + "Underlying is optimized but does not have dual values stored." + ) + shape = tuple(len(c) for c in self._coords) + dim_names = [c.name for c in self._coords] + xr_coords = {c.name: c for c in self._coords} + return DataArray( + self._dual.reshape(shape) if shape else self._dual, + coords=xr_coords, + dims=dim_names, + ) + + def _to_dataset(self, nterm: int) -> Dataset: + """ + Reconstruct labels/coeffs/vars Dataset from the CSR matrix. + + Parameters + ---------- + nterm : int + Number of terms per row (width of the dense term block). + + Returns + ------- + Dataset with variables ``labels``, ``coeffs``, ``vars``. + """ + csr = self._csr + n_total = csr.shape[0] + counts = np.diff(csr.indptr) + nonempty = counts > 0 + coeffs_2d = np.zeros((n_total, nterm), dtype=csr.dtype) + vars_2d = np.full((n_total, nterm), -1, dtype=np.int64) + if csr.nnz > 0: + row_indices = np.repeat(np.arange(n_total, dtype=np.int32), counts) + term_cols = np.arange(csr.nnz, dtype=np.int32) - np.repeat( + csr.indptr[:-1].astype(np.int32), counts + ) + vars_2d[row_indices, term_cols] = csr.indices + coeffs_2d[row_indices, term_cols] = csr.data + shape = tuple(len(c) for c in self._coords) + dim_names = [c.name for c in self._coords] + xr_coords = {c.name: c for c in self._coords} + term_coords = {TERM_DIM: np.arange(nterm)} + dims_with_term = dim_names + [TERM_DIM] + coords_with_term = {**xr_coords, **term_coords} + coeffs_da = DataArray( + coeffs_2d.reshape(shape + (nterm,)) if shape else coeffs_2d, + coords=coords_with_term, + dims=dims_with_term, + ) + vars_da = DataArray( + vars_2d.reshape(shape + (nterm,)) if shape else vars_2d, + coords=coords_with_term, + dims=dims_with_term, + ) + ds = Dataset({"coeffs": coeffs_da, "vars": vars_da}) + if self._cindex is not None: + labels_flat = np.where( + nonempty, np.arange(self._cindex, self._cindex + n_total), -1 + ) + ds["labels"] = DataArray( + labels_flat.reshape(shape) if shape else labels_flat, + coords=xr_coords, + dims=dim_names, + ) + return ds + + @property + def data(self) -> Dataset: + """Reconstruct the xarray Dataset from the CSR representation.""" + nterm = int(self._csr.indptr.max()) if self._csr.nnz > 0 else 1 + ds = self._to_dataset(nterm) + sign_arr = np.full(tuple(len(c) for c in self._coords) or (1,), self._sign) + rhs_arr = self._rhs.reshape(tuple(len(c) for c in self._coords) or (1,)) + shape = tuple(len(c) for c in self._coords) + dim_names = [c.name for c in self._coords] + xr_coords = {c.name: c for c in self._coords} + ds = ds.assign( + sign=DataArray( + sign_arr if shape else np.array(self._sign), + coords=xr_coords, + dims=dim_names, + ), + rhs=DataArray( + rhs_arr if shape else self._rhs, coords=xr_coords, dims=dim_names + ), + ) + if self._dual is not None: + ds = ds.assign( + dual=DataArray( + self._dual.reshape(shape) if shape else self._dual, + coords=xr_coords, + dims=dim_names, + ) + ) + return ds.assign_attrs(name=self._name) + + def __repr__(self) -> str: + """Print the constraint without reconstructing the full Dataset.""" + max_lines = options["display_max_rows"] + coords = self._coords + shape = tuple(len(c) for c in coords) + dim_names = [c.name for c in coords] + dim_sizes = list(shape) + size = int(np.prod(shape)) if shape else 1 + nonempty = self._nonempty # shape (size,) + nterm = int(self._csr.indptr.max()) if self._csr.nnz > 0 else 1 + + # Dense arrays for coeffs/vars, built without going through data property + csr = self._csr + counts = np.diff(csr.indptr) + coeffs_2d = np.zeros((size, nterm), dtype=csr.dtype) + vars_2d = np.full((size, nterm), -1, dtype=np.int64) + if csr.nnz > 0: + row_idx = np.repeat(np.arange(size, dtype=np.int32), counts) + term_cols = np.arange(csr.nnz, dtype=np.int32) - np.repeat( + csr.indptr[:-1].astype(np.int32), counts + ) + vars_2d[row_idx, term_cols] = csr.indices + coeffs_2d[row_idx, term_cols] = csr.data + + coeffs_nd = coeffs_2d.reshape(shape + (nterm,)) if shape else coeffs_2d + vars_nd = vars_2d.reshape(shape + (nterm,)) if shape else vars_2d + rhs_nd = self._rhs.reshape(shape) if shape else self._rhs + masked_entries = int((~nonempty).sum()) + + header_string = f"{self.type} `{self._name}`" if self._name else f"{self.type}" + lines = [] + + if size > 1 or len(dim_sizes) > 0: + for indices in generate_indices_for_printout(dim_sizes, max_lines): + if indices is None: + lines.append("\t\t...") + else: + coord = [coords[i][int(ind)] for i, ind in enumerate(indices)] + flat_idx = int(np.ravel_multi_index(indices, shape)) if shape else 0 + if nonempty[flat_idx]: + expr = print_single_expression( + coeffs_nd[indices], vars_nd[indices], 0, self._model + ) + sign = SIGNS_pretty[self._sign] + rhs = rhs_nd[indices] + line = print_coord(coord) + f": {expr} {sign} {rhs}" + else: + line = print_coord(coord) + ": None" + lines.append(line) + lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) + + shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes)) + mask_str = f" - {masked_entries} masked entries" if masked_entries else "" + underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) + lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") + elif size == 1: + expr = print_single_expression( + coeffs_nd.ravel(), vars_nd.ravel(), 0, self._model ) lines.append( - f"{header_string}\n{'-' * len(header_string)}\n{expr} {SIGNS_pretty[self.sign.item()]} {self.rhs.item()}" + f"{header_string}\n{'-' * len(header_string)}\n{expr} {SIGNS_pretty[self._sign]} {self._rhs.item()}" ) else: lines.append(f"{header_string}\n{'-' * len(header_string)}\n") return "\n".join(lines) - def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None: + def to_matrix(self) -> scipy.sparse.csr_array: + """Return the stored CSR matrix directly (no reconstruction needed).""" + return self._csr + + def freeze(self) -> Constraint: + """Return self (already immutable).""" + return self + + def mutable(self) -> MutableConstraint: + """Convert to a MutableConstraint.""" + return MutableConstraint(self.data, self._model, self._name) + + @classmethod + def from_mutable( + cls, + con: MutableConstraint, + cindex: int | None = None, + ) -> Constraint: """ - Print the linear expression. + Create a Constraint from a MutableConstraint. Parameters ---------- - display_max_rows : int - Maximum number of rows to be displayed. - display_max_terms : int - Maximum number of terms to be displayed. - """ - with options as opts: - opts.set_value( - display_max_rows=display_max_rows, display_max_terms=display_max_terms + con : MutableConstraint + cindex : int or None + Starting label index, if assigned. + """ + csr = con.to_matrix() + coords = [con.indexes[d] for d in con.coord_dims] + rhs = con.rhs.values.ravel() + sign_vals = con.sign.values.ravel() + unique_signs = np.unique(sign_vals) + if len(unique_signs) > 1: + raise ValueError( + "Constraint has per-element signs; cannot freeze to immutable Constraint. " + "This is a known limitation — use MutableConstraint instead." ) - print(self) + sign = str(unique_signs[0]) if len(unique_signs) == 1 else "=" + dual = con.data["dual"].values.ravel() if "dual" in con.data else None + return cls( + csr, rhs, sign, coords, con.model, con.name, cindex=cindex, dual=dual + ) - def __contains__(self, value: Any) -> bool: - return self.data.__contains__(value) + +class MutableConstraint(ConstraintBase): + """ + Mutable constraint backed by an xarray Dataset. + + This is the original Constraint implementation, renamed to MutableConstraint. + Supports setters, xarray operations via conwrap, and from_rule construction. + """ + + __slots__ = ("_data", "_model", "_assigned") + + def __init__( + self, + data: Dataset, + model: Model, + name: str = "", + skip_broadcast: bool = False, + ) -> None: + from linopy.model import Model + + if not isinstance(data, Dataset): + raise ValueError(f"data must be a Dataset, got {type(data)}") + + if not isinstance(model, Model): + raise ValueError(f"model must be a Model, got {type(model)}") + + for attr in ("coeffs", "vars", "sign", "rhs"): + if attr not in data: + raise ValueError(f"missing '{attr}' in data") + + data = data.assign_attrs(name=name) + + if not skip_broadcast: + (data,) = xr.broadcast(data, exclude=[TERM_DIM]) + + self._assigned = "labels" in data + self._data = data + self._model = model @property - def type(self) -> str: - """ - Get the type of the constraint. - """ - return "Constraint" if self.is_assigned else "Constraint (unassigned)" + def data(self) -> Dataset: + return self._data @property - def range(self) -> tuple[int, int]: - """ - Return the range of the constraint. - """ - return self.data.attrs["label_range"] + def model(self) -> Model: + return self._model @property - def term_dim(self) -> str: - """ - Return the term dimension of the constraint. - """ - return TERM_DIM + def name(self) -> str: + return self.attrs["name"] @property - def mask(self) -> DataArray | None: - """ - Get the mask of the constraint. + def is_assigned(self) -> bool: + return self._assigned - The mask indicates on which coordinates the constraint is enabled - (True) and disabled (False). + @property + def range(self) -> tuple[int, int]: + """Return the range of the constraint.""" + return self.data.attrs["label_range"] - Returns - ------- - xr.DataArray - """ - if self.is_assigned: - return (self.data.labels != FILL_VALUE["labels"]).astype(bool) - return None + @property + def loc(self) -> LocIndexer: + return LocIndexer(self) @property - def coeffs(self) -> DataArray: - """ - Get the left-hand-side coefficients of the constraint. + def labels(self) -> DataArray: + return self.data.get("labels", DataArray([])) - The function raises an error in case no model is set as a - reference. - """ + @property + def coeffs(self) -> DataArray: return self.data.coeffs @coeffs.setter @@ -527,12 +882,6 @@ def coeffs(self, value: ConstantLike) -> None: @property def vars(self) -> DataArray: - """ - Get the left-hand-side variables of the constraint. - - The function raises an error in case no model is set as a - reference. - """ return self.data.vars @vars.setter @@ -544,34 +893,8 @@ def vars(self, value: variables.Variable | DataArray) -> None: value = value.broadcast_like(self.coeffs, exclude=[self.term_dim]) self._data = assign_multiindex_safe(self.data, vars=value) - @property - def lhs(self) -> expressions.LinearExpression: - """ - Get the left-hand-side linear expression of the constraint. - - The function raises an error in case no model is set as a - reference. - """ - data = self.data[["coeffs", "vars"]].rename({self.term_dim: TERM_DIM}) - return expressions.LinearExpression(data, self.model) - - @lhs.setter - def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: - value = expressions.as_expression( - value, self.model, coords=self.coords, dims=self.coord_dims - ) - self._data = self.data.drop_vars(["coeffs", "vars"]).assign( - coeffs=value.coeffs, vars=value.vars, rhs=self.rhs - value.const - ) - @property def sign(self) -> DataArray: - """ - Get the signs of the constraint. - - The function raises an error in case no model is set as a - reference. - """ return self.data.sign @sign.setter @@ -582,12 +905,6 @@ def sign(self, value: SignLike) -> None: @property def rhs(self) -> DataArray: - """ - Get the right hand side constants of the constraint. - - The function raises an error in case no model is set as a - reference. - """ return self.data.rhs @rhs.setter @@ -598,15 +915,23 @@ def rhs(self, value: ExpressionLike) -> None: self.lhs = self.lhs - value.reset_const() self._data = assign_multiindex_safe(self.data, rhs=value.const) + @property + def lhs(self) -> expressions.LinearExpression: + data = self.data[["coeffs", "vars"]].rename({self.term_dim: TERM_DIM}) + return expressions.LinearExpression(data, self.model) + + @lhs.setter + def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: + value = expressions.as_expression( + value, self.model, coords=self.coords, dims=self.coord_dims + ) + self._data = self.data.drop_vars(["coeffs", "vars"]).assign( + coeffs=value.coeffs, vars=value.vars, rhs=self.rhs - value.const + ) + @property @has_optimized_model def dual(self) -> DataArray: - """ - Get the dual values of the constraint. - - The function raises an error in case no model is set as a - reference or the model status is not okay. - """ if "dual" not in self.data: raise AttributeError( "Underlying is optimized but does not have dual values stored." @@ -615,21 +940,32 @@ def dual(self) -> DataArray: @dual.setter def dual(self, value: ConstantLike) -> None: - """ - Get the dual values of the constraint. - """ value = DataArray(value).broadcast_like(self.labels) self._data = assign_multiindex_safe(self.data, dual=value) + def freeze(self) -> Constraint: + """Convert to an immutable Constraint.""" + cindex = ( + int(self.data.attrs["label_range"][0]) + if "label_range" in self.data.attrs + else None + ) + return Constraint.from_mutable(self, cindex=cindex) + + def mutable(self) -> MutableConstraint: + """Return self (already mutable).""" + return self + @classmethod - def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constraint: + def from_rule( + cls, model: Model, rule: Callable, coords: CoordsLike + ) -> MutableConstraint: """ Create a constraint from a rule and a set of coordinates. This functionality mirrors the assignment of constraints as done by Pyomo. - Parameters ---------- model : linopy.Model @@ -647,14 +983,13 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai The order and size of coords has to be same as the argument list followed by `model` in function `rule`. - Returns ------- - linopy.Constraint + linopy.MutableConstraint Examples -------- - >>> from linopy import Model, LinearExpression, Constraint + >>> from linopy import Model, LinearExpression, MutableConstraint >>> m = Model() >>> coords = pd.RangeIndex(10), ["a", "b"] >>> x = m.add_variables(0, 100, coords) @@ -664,14 +999,13 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai ... else: ... return i * x.at[i, j] >= 0 ... - >>> con = Constraint.from_rule(m, bound, coords) + >>> con = MutableConstraint.from_rule(m, bound, coords) >>> con = m.add_constraints(con) """ if not isinstance(coords, DataArrayCoordinates): coords = DataArray(coords=coords).coords shape = list(map(len, coords.values())) - # test output type output = rule(model, *[c.values[0] for c in coords.values()]) if not isinstance(output, AnonymousScalarConstraint) and output is not None: msg = f"`rule` has to return AnonymousScalarConstraint not {type(output)}." @@ -691,106 +1025,6 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai data = lhs.data.assign(sign=sign, rhs=rhs) return cls(data, model=model) - @property - def flat(self) -> pd.DataFrame: - """ - Convert the constraint to a pandas DataFrame. - - The resulting DataFrame represents a long table format of the all - non-masked constraints with non-zero coefficients. It contains the - columns `labels`, `coeffs`, `vars`, `rhs`, `sign`. - - Returns - ------- - df : pandas.DataFrame - """ - ds = self.data - - def mask_func(data: pd.DataFrame) -> pd.Series: - mask = (data["vars"] != -1) & (data["coeffs"] != 0) - if "labels" in data: - mask &= data["labels"] != -1 - return mask - - df = to_dataframe(ds, mask_func=mask_func) - - # Group repeated variables in the same constraint - agg_custom = {k: "first" for k in list(df.columns)} - agg_standards = dict(coeffs="sum", rhs="first", sign="first") - agg = {**agg_custom, **agg_standards} - df = df.groupby(["labels", "vars"], as_index=False).aggregate(agg) - check_has_nulls(df, name=f"{self.type} {self.name}") - return df - - def to_matrix(self) -> tuple[scipy.sparse.csr_array, np.ndarray]: - """ - Construct a CSR matrix representation of this constraint. - - Returns - ------- - matrix : scipy.sparse.csr_array - Shape (n_labels, model._xCounter). Rows correspond to individual - constraint labels, columns to variable labels. Missing entries - (labels or vars == -1) are excluded. - labels : np.ndarray - 1D array of shape (n_labels,) mapping each row back to the - original constraint label. - """ - vars = self.vars.values - labels_flat = self.labels.values.ravel() - vars_2d = vars.reshape(len(labels_flat), -1) - coeffs_flat = self.coeffs.values.ravel() - - valid_vars = vars_2d != -1 - labels_valid = (labels_flat != -1) & valid_vars.any(axis=1) - valid = (labels_valid[:, np.newaxis] & valid_vars).ravel() - - con_labels = labels_flat[labels_valid] - counts = valid_vars.sum(axis=1)[labels_valid] - indptr = np.empty(len(con_labels) + 1, dtype=np.int32) - indptr[0] = 0 - np.cumsum(counts, out=indptr[1:]) - - vars_flat = vars.ravel()[valid] - coeffs_flat = coeffs_flat[valid] - - shape = (len(con_labels), self.model._xCounter) - # Note: duplicate (row, col) entries are not summed in CSR format. - # They will be summed automatically upon conversion to CSC or dense. - matrix = scipy.sparse.csr_array((coeffs_flat, vars_flat, indptr), shape=shape) - return matrix, con_labels - - def to_matrix2(self) -> scipy.sparse.csr_array: - """ - Construct a CSR matrix representation of this constraint including masked rows. - - Unlike to_matrix(), all flat positions in the constraint grid are included as - rows (masked entries become empty rows). Shape is - (len(labels_flat), model._xCounter). - - Returns - ------- - matrix : scipy.sparse.csr_array - """ - vars = self.vars.values - labels_flat = self.labels.values.ravel() - vars_2d = vars.reshape(len(labels_flat), -1) - coeffs_flat = self.coeffs.values.ravel() - - valid_2d = (labels_flat != -1)[:, np.newaxis] & (vars_2d != -1) - cols = vars_2d[valid_2d] - data = coeffs_flat.reshape(vars_2d.shape)[valid_2d] - - counts = valid_2d.sum(axis=1) - indptr = np.empty(len(labels_flat) + 1, dtype=np.int32) - indptr[0] = 0 - np.cumsum(counts, out=indptr[1:]) - - shape = (len(labels_flat), self.model._xCounter) - # Note: duplicate (row, col) entries are not summed in CSR format. - # They will be summed automatically upon conversion to CSC or dense. - return scipy.sparse.csr_array((data, cols, indptr), shape=shape) - def to_polars(self) -> pl.DataFrame: """ Convert the constraint to a polars DataFrame. @@ -813,8 +1047,6 @@ def to_polars(self) -> pl.DataFrame: long = maybe_group_terms_polars(long) check_has_nulls_polars(long, name=f"{self.type} {self.name}") - # Build short DataFrame (labels, rhs, sign) without xarray broadcast. - # Apply labels mask directly instead of filter_nulls_polars. labels_flat = ds["labels"].values.reshape(-1) mask = labels_flat != -1 labels_masked = labels_flat[mask] @@ -843,55 +1075,29 @@ def to_polars(self) -> pl.DataFrame: df = long.join(short, on="labels", how="inner") return df[["labels", "coeffs", "vars", "sign", "rhs"]] - # Wrapped function which would convert variable to dataarray + # Wrapped xarray methods — only available on MutableConstraint assign = conwrap(Dataset.assign) - assign_multiindex_safe = conwrap(assign_multiindex_safe) - assign_attrs = conwrap(Dataset.assign_attrs) - assign_coords = conwrap(Dataset.assign_coords) - - # bfill = conwrap(Dataset.bfill) - broadcast_like = conwrap(Dataset.broadcast_like) - chunk = conwrap(Dataset.chunk) - drop_sel = conwrap(Dataset.drop_sel) - drop_isel = conwrap(Dataset.drop_isel) - expand_dims = conwrap(Dataset.expand_dims) - - # ffill = conwrap(Dataset.ffill) - sel = conwrap(Dataset.sel) - isel = conwrap(Dataset.isel) - shift = conwrap(Dataset.shift) - swap_dims = conwrap(Dataset.swap_dims) - set_index = conwrap(Dataset.set_index) - - reindex = conwrap(Dataset.reindex, fill_value=_fill_value) - - reindex_like = conwrap(Dataset.reindex_like, fill_value=_fill_value) - + reindex = conwrap(Dataset.reindex, fill_value=FILL_VALUE) + reindex_like = conwrap(Dataset.reindex_like, fill_value=FILL_VALUE) rename = conwrap(Dataset.rename) - rename_dims = conwrap(Dataset.rename_dims) - roll = conwrap(Dataset.roll) - stack = conwrap(Dataset.stack) - unstack = conwrap(Dataset.unstack) - iterate_slices = iterate_slices - @dataclass(repr=False) class Constraints: @@ -899,7 +1105,7 @@ class Constraints: A constraint container used for storing multiple constraint arrays. """ - data: dict[str, Constraint] + data: dict[str, ConstraintBase] model: Model _label_position_index: LabelPositionIndex | None = None @@ -940,17 +1146,17 @@ def __repr__(self) -> str: return r @overload - def __getitem__(self, names: str) -> Constraint: ... + def __getitem__(self, names: str) -> ConstraintBase: ... @overload def __getitem__(self, names: list[str]) -> Constraints: ... - def __getitem__(self, names: str | list[str]) -> Constraint | Constraints: + def __getitem__(self, names: str | list[str]) -> ConstraintBase | Constraints: if isinstance(names, str): return self.data[names] return Constraints({name: self.data[name] for name in names}, self.model) - def __getattr__(self, name: str) -> Constraint: + def __getattr__(self, name: str) -> ConstraintBase: # If name is an attribute of self (including methods and properties), return that if name in self.data: return self.data[name] @@ -980,7 +1186,7 @@ def __len__(self) -> int: def __iter__(self) -> Iterator[str]: return self.data.__iter__() - def items(self) -> ItemsView[str, Constraint]: + def items(self) -> ItemsView[str, ConstraintBase]: return self.data.items() def _ipython_key_completions_(self) -> list[str]: @@ -993,12 +1199,15 @@ def _ipython_key_completions_(self) -> list[str]: """ return list(self) - def add(self, constraint: Constraint) -> None: + def add(self, constraint: ConstraintBase, freeze: bool = False) -> ConstraintBase: """ - Add a constraint to the constraints constrainer. + Add a constraint to the constraints container. """ + if freeze and isinstance(constraint, MutableConstraint): + constraint = constraint.freeze() self.data[constraint.name] = constraint self._invalidate_label_position_index() + return constraint def remove(self, name: str) -> None: """ @@ -1085,33 +1294,7 @@ def ncons(self) -> int: This excludes constraints with missing labels or where all variables are masked (vars == -1). """ - total = 0 - for con in self.data.values(): - labels = con.labels.values - vars_arr = con.vars.values - - # Handle scalar constraint (single constraint, labels is 0-d) - if labels.ndim == 0: - # Scalar: valid if label != -1 and any var != -1 - if labels != -1 and (vars_arr != -1).any(): - total += 1 - continue - - # Array constraint: labels has constraint dimensions, vars has - # constraint dimensions + _term dimension - valid_labels = labels != -1 - - # Check if any variable in each constraint is valid (not -1) - # vars has shape (..., n_terms) where ... matches labels shape - has_valid_var = (vars_arr != -1).any(axis=-1) - - active = valid_labels & has_valid_var - - if con.mask is not None: - active = active & con.mask.values - - total += int(active.sum()) - return total + return sum(con.ncons for con in self.data.values()) @property def inequalities(self) -> Constraints: @@ -1323,7 +1506,18 @@ def to_matrix_via_csr( if not len(self): raise ValueError("No constraints available to convert to matrix.") - matrices, con_labels_list = zip(*(c.to_matrix() for c in self.data.values())) + matrices = [] + con_labels_list = [] + for c in self.data.values(): + csr = c.to_matrix() + matrices.append(csr) + nonempty = np.diff(csr.indptr).astype(bool) + start = ( + c._cindex + if isinstance(c, Constraint) + else c.data.attrs["label_range"][0] + ) + con_labels_list.append(np.flatnonzero(nonempty) + start) csc: scipy.sparse.csc_array = scipy.sparse.vstack(matrices).tocsc() con_labels = np.concatenate(con_labels_list) @@ -1404,6 +1598,6 @@ def rhs(self) -> int | float | np.floating | np.integer: """ return self._rhs - def to_constraint(self) -> Constraint: + def to_constraint(self) -> MutableConstraint: data = self.lhs.to_linexpr().data.assign(sign=self.sign, rhs=self.rhs) - return Constraint(data=data, model=self.lhs.model) + return MutableConstraint(data=data, model=self.lhs.model) diff --git a/linopy/expressions.py b/linopy/expressions.py index d2ae9022..80eb15c9 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -90,7 +90,10 @@ ) if TYPE_CHECKING: - from linopy.constraints import AnonymousScalarConstraint, Constraint + from linopy.constraints import ( + AnonymousScalarConstraint, + MutableConstraint, + ) from linopy.model import Model from linopy.piecewise import PiecewiseConstraintDescriptor, PiecewiseExpression from linopy.variables import ScalarVariable, Variable @@ -673,9 +676,11 @@ def __truediv__(self: GenericExpression, other: SideLike) -> GenericExpression: def __le__(self, rhs: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __le__(self, rhs: SideLike) -> Constraint: ... + def __le__(self, rhs: SideLike) -> MutableConstraint: ... - def __le__(self, rhs: SideLike) -> Constraint | PiecewiseConstraintDescriptor: + def __le__( + self, rhs: SideLike + ) -> MutableConstraint | PiecewiseConstraintDescriptor: descriptor = _to_piecewise_constraint_descriptor(self, rhs, "<=") if descriptor is not None: return descriptor @@ -685,9 +690,11 @@ def __le__(self, rhs: SideLike) -> Constraint | PiecewiseConstraintDescriptor: def __ge__(self, rhs: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __ge__(self, rhs: SideLike) -> Constraint: ... + def __ge__(self, rhs: SideLike) -> MutableConstraint: ... - def __ge__(self, rhs: SideLike) -> Constraint | PiecewiseConstraintDescriptor: + def __ge__( + self, rhs: SideLike + ) -> MutableConstraint | PiecewiseConstraintDescriptor: descriptor = _to_piecewise_constraint_descriptor(self, rhs, ">=") if descriptor is not None: return descriptor @@ -697,9 +704,11 @@ def __ge__(self, rhs: SideLike) -> Constraint | PiecewiseConstraintDescriptor: def __eq__(self, rhs: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __eq__(self, rhs: SideLike) -> Constraint: ... + def __eq__(self, rhs: SideLike) -> MutableConstraint: ... - def __eq__(self, rhs: SideLike) -> Constraint | PiecewiseConstraintDescriptor: + def __eq__( + self, rhs: SideLike + ) -> MutableConstraint | PiecewiseConstraintDescriptor: descriptor = _to_piecewise_constraint_descriptor(self, rhs, "==") if descriptor is not None: return descriptor @@ -818,7 +827,7 @@ def le( self: GenericExpression, rhs: SideLike, join: str | None = None, - ) -> Constraint: + ) -> MutableConstraint: """ Less than or equal constraint. @@ -837,7 +846,7 @@ def ge( self: GenericExpression, rhs: SideLike, join: str | None = None, - ) -> Constraint: + ) -> MutableConstraint: """ Greater than or equal constraint. @@ -856,7 +865,7 @@ def eq( self: GenericExpression, rhs: SideLike, join: str | None = None, - ) -> Constraint: + ) -> MutableConstraint: """ Equality constraint. @@ -1112,7 +1121,7 @@ def cumsum( def to_constraint( self, sign: SignLike, rhs: SideLike, join: str | None = None - ) -> Constraint: + ) -> MutableConstraint: """ Convert a linear expression to a constraint. @@ -1172,7 +1181,7 @@ def to_constraint( data = assign_multiindex_safe( all_to_lhs[["coeffs", "vars"]], sign=sign, rhs=computed_rhs ) - return constraints.Constraint(data, model=self.model) + return constraints.MutableConstraint(data, model=self.model) def reset_const(self: GenericExpression) -> GenericExpression: """ diff --git a/linopy/io.py b/linopy/io.py index 2213cbb5..39b56783 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1168,10 +1168,10 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: m : linopy.Model """ from linopy.model import ( - Constraint, Constraints, LinearExpression, Model, + MutableConstraint, Variable, Variables, ) @@ -1224,7 +1224,7 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: constraints = {} for k in sorted(con_names): name = remove_prefix(k, "constraints") - constraints[name] = Constraint(get_prefix(ds, k), m, name) + constraints[name] = MutableConstraint(get_prefix(ds, k), m, name) m._constraints = Constraints(constraints, m) objective = get_prefix(ds, "objective") diff --git a/linopy/model.py b/linopy/model.py index 06e814c6..b1310831 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -46,7 +46,12 @@ ModelStatus, TerminationCondition, ) -from linopy.constraints import AnonymousScalarConstraint, Constraint, Constraints +from linopy.constraints import ( + AnonymousScalarConstraint, + ConstraintBase, + Constraints, + MutableConstraint, +) from linopy.expressions import ( LinearExpression, QuadraticExpression, @@ -726,7 +731,8 @@ def add_constraints( name: str | None = None, coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = None, mask: MaskLike | None = None, - ) -> Constraint: + freeze: bool = False, + ) -> ConstraintBase: """ Assign a new, possibly multi-dimensional array of constraints to the model. @@ -738,7 +744,7 @@ def add_constraints( Parameters ---------- - lhs : linopy.LinearExpression/linopy.Constraint/callable + lhs : linopy.LinearExpression/linopy.ConstraintBase/callable Left hand side of the constraint(s) or optionally full constraint. In case a linear expression is passed, `sign` and `rhs` must not be None. @@ -760,12 +766,14 @@ def add_constraints( Boolean mask with False values for constraints which are skipped. The shape of the mask has to match the shape the added constraints. Default is None. - + freeze : bool, optional + If True, convert the constraint to an immutable CSR-backed Constraint + before returning. Default is False. Returns ------- - labels : linopy.model.Constraint - Array containing the labels of the added constraints. + constraint : linopy.ConstraintBase + The added constraint (MutableConstraint by default, or Constraint if freeze=True). """ msg_sign_rhs_none = f"Arguments `sign` and `rhs` cannot be None when passing along with a {type(lhs)}." @@ -801,12 +809,12 @@ def add_constraints( rule = lhs if sign is not None or rhs is not None: raise ValueError(msg_sign_rhs_none) - data = Constraint.from_rule(self, rule, coords).data + data = MutableConstraint.from_rule(self, rule, coords).data elif isinstance(lhs, AnonymousScalarConstraint): if sign is not None or rhs is not None: raise ValueError(msg_sign_rhs_none) data = lhs.to_constraint().data - elif isinstance(lhs, Constraint): + elif isinstance(lhs, ConstraintBase): if sign is not None or rhs is not None: raise ValueError(msg_sign_rhs_none) data = lhs.data @@ -883,9 +891,8 @@ def add_constraints( if self.chunk: data = data.chunk(self.chunk) - constraint = Constraint(data, name=name, model=self, skip_broadcast=True) - self.constraints.add(constraint) - return constraint + constraint = MutableConstraint(data, name=name, model=self, skip_broadcast=True) + return self.constraints.add(constraint, freeze=freeze) def add_objective( self, diff --git a/linopy/piecewise.py b/linopy/piecewise.py index 78f7be65..deb0ef3c 100644 --- a/linopy/piecewise.py +++ b/linopy/piecewise.py @@ -40,7 +40,7 @@ ) if TYPE_CHECKING: - from linopy.constraints import Constraint + from linopy.constraints import ConstraintBase from linopy.expressions import LinearExpression from linopy.model import Model from linopy.types import LinExprLike @@ -686,7 +686,7 @@ def _add_pwl_lp( sign: str, x_points: DataArray, y_points: DataArray, -) -> Constraint: +) -> ConstraintBase: """Add pure LP tangent-line constraints.""" dx = x_points.diff(BREAKPOINT_DIM) dy = y_points.diff(BREAKPOINT_DIM) @@ -729,7 +729,7 @@ def _add_pwl_sos2_core( y_points: DataArray, lambda_mask: DataArray | None, active: LinearExpression | None = None, -) -> Constraint: +) -> ConstraintBase: """ Core SOS2 formulation linking x_expr and target_expr via breakpoints. @@ -780,7 +780,7 @@ def _add_pwl_incremental_core( y_points: DataArray, bp_mask: DataArray | None, active: LinearExpression | None = None, -) -> Constraint: +) -> ConstraintBase: """ Core incremental formulation linking x_expr and target_expr. @@ -846,7 +846,7 @@ def _add_pwl_incremental_core( model.add_constraints(delta_var <= binary_var, name=inc_link_name) # Order constraints: y_{i+1} ≤ δ_i for i = 0..n-2 - fill_con: Constraint | None = None + fill_con: ConstraintBase | None = None if n_segments >= 2: delta_lo = delta_var.isel({LP_SEG_DIM: slice(None, -1)}, drop=True) delta_hi = delta_var.isel({LP_SEG_DIM: slice(1, None)}, drop=True) @@ -884,7 +884,7 @@ def _add_dpwl_sos2_core( y_points: DataArray, lambda_mask: DataArray | None, active: LinearExpression | None = None, -) -> Constraint: +) -> ConstraintBase: """ Core disjunctive SOS2 formulation with separate x/y points. @@ -948,11 +948,11 @@ def _add_dpwl_sos2_core( def add_piecewise_constraints( model: Model, - descriptor: PiecewiseConstraintDescriptor | Constraint, + descriptor: PiecewiseConstraintDescriptor | ConstraintBase, method: Literal["sos2", "incremental", "auto", "lp"] = "auto", name: str | None = None, skip_nan_check: bool = False, -) -> Constraint: +) -> ConstraintBase: """ Add a piecewise linear constraint from a :class:`PiecewiseConstraintDescriptor`. @@ -975,7 +975,7 @@ def add_piecewise_constraints( Returns ------- - Constraint + ConstraintBase """ if not isinstance(descriptor, PiecewiseConstraintDescriptor): raise TypeError( @@ -1064,7 +1064,7 @@ def _add_continuous( method: str, skip_nan_check: bool, active: LinearExpression | None = None, -) -> Constraint: +) -> ConstraintBase: """Handle continuous (non-disjunctive) piecewise constraints.""" convexity: Literal["convex", "concave", "linear", "mixed"] | None = None @@ -1182,7 +1182,7 @@ def _add_disjunctive( mask: DataArray | None, method: str, active: LinearExpression | None = None, -) -> Constraint: +) -> ConstraintBase: """Handle disjunctive piecewise constraints.""" if method == "lp": raise ValueError("Pure LP method is not supported for disjunctive constraints") diff --git a/linopy/variables.py b/linopy/variables.py index 4332a037..f5e52a88 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -64,7 +64,7 @@ ) if TYPE_CHECKING: - from linopy.constraints import AnonymousScalarConstraint, Constraint + from linopy.constraints import AnonymousScalarConstraint, MutableConstraint from linopy.expressions import ( GenericExpression, LinearExpression, @@ -535,27 +535,33 @@ def __rsub__(self, other: ConstantLike) -> LinearExpression: def __le__(self, other: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __le__(self, other: SideLike) -> Constraint: ... + def __le__(self, other: SideLike) -> MutableConstraint: ... - def __le__(self, other: SideLike) -> Constraint | PiecewiseConstraintDescriptor: + def __le__( + self, other: SideLike + ) -> MutableConstraint | PiecewiseConstraintDescriptor: return self.to_linexpr().__le__(other) @overload def __ge__(self, other: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __ge__(self, other: SideLike) -> Constraint: ... + def __ge__(self, other: SideLike) -> MutableConstraint: ... - def __ge__(self, other: SideLike) -> Constraint | PiecewiseConstraintDescriptor: + def __ge__( + self, other: SideLike + ) -> MutableConstraint | PiecewiseConstraintDescriptor: return self.to_linexpr().__ge__(other) @overload # type: ignore[override] def __eq__(self, other: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __eq__(self, other: SideLike) -> Constraint: ... + def __eq__(self, other: SideLike) -> MutableConstraint: ... - def __eq__(self, other: SideLike) -> Constraint | PiecewiseConstraintDescriptor: + def __eq__( + self, other: SideLike + ) -> MutableConstraint | PiecewiseConstraintDescriptor: return self.to_linexpr().__eq__(other) def __gt__(self, other: Any) -> NotImplementedType: @@ -639,7 +645,7 @@ def div( """ return self.to_linexpr().div(other, join=join) - def le(self, rhs: SideLike, join: str | None = None) -> Constraint: + def le(self, rhs: SideLike, join: str | None = None) -> MutableConstraint: """ Less than or equal constraint. @@ -654,7 +660,7 @@ def le(self, rhs: SideLike, join: str | None = None) -> Constraint: """ return self.to_linexpr().le(rhs, join=join) - def ge(self, rhs: SideLike, join: str | None = None) -> Constraint: + def ge(self, rhs: SideLike, join: str | None = None) -> MutableConstraint: """ Greater than or equal constraint. @@ -669,7 +675,7 @@ def ge(self, rhs: SideLike, join: str | None = None) -> Constraint: """ return self.to_linexpr().ge(rhs, join=join) - def eq(self, rhs: SideLike, join: str | None = None) -> Constraint: + def eq(self, rhs: SideLike, join: str | None = None) -> MutableConstraint: """ Equality constraint. diff --git a/test/test_constraint.py b/test/test_constraint.py index bfd29a6e..483d0835 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -24,7 +24,13 @@ short_LESS_EQUAL, sign_replace_dict, ) -from linopy.constraints import AnonymousScalarConstraint, Constraint, Constraints +from linopy.constraints import ( + AnonymousScalarConstraint, + Constraint, + ConstraintBase, + Constraints, + MutableConstraint, +) @pytest.fixture @@ -56,10 +62,25 @@ def test_constraint_repr(c: linopy.constraints.Constraint) -> None: c.__repr__() +def test_constraint_repr_equivalent_to_mutable( + c: linopy.constraints.Constraint, +) -> None: + """Constraint (CSR-backed) and MutableConstraint repr must be identical.""" + frozen = c.freeze() + assert repr(frozen) == repr(c) + + def test_constraints_repr(m: Model) -> None: m.constraints.__repr__() +def test_add_constraints_freeze(m: Model, x: linopy.Variable) -> None: + c = m.add_constraints(x >= 1, name="frozen_c", freeze=True) + assert isinstance(c, linopy.constraints.Constraint) + assert isinstance(m.constraints["frozen_c"], linopy.constraints.Constraint) + assert c.ncons == 10 + + def test_constraint_name(c: linopy.constraints.Constraint) -> None: assert c.name == "c" @@ -228,7 +249,7 @@ def test_constraint_wrapped_methods(x: linopy.Variable, y: linopy.Variable) -> N def test_anonymous_constraint_sel(x: linopy.Variable, y: linopy.Variable) -> None: expr = 10 * x + y con = expr <= 10 - assert isinstance(con.sel(first=[1, 2]), Constraint) + assert isinstance(con.sel(first=[1, 2]), ConstraintBase) def test_anonymous_constraint_swap_dims(x: linopy.Variable, y: linopy.Variable) -> None: @@ -236,7 +257,7 @@ def test_anonymous_constraint_swap_dims(x: linopy.Variable, y: linopy.Variable) con = expr <= 10 con = con.assign_coords({"third": ("second", con.indexes["second"] + 100)}) con = con.swap_dims({"second": "third"}) - assert isinstance(con, Constraint) + assert isinstance(con, ConstraintBase) assert con.coord_dims == ("first", "third") @@ -245,7 +266,7 @@ def test_anonymous_constraint_set_index(x: linopy.Variable, y: linopy.Variable) con = expr <= 10 con = con.assign_coords({"third": ("second", con.indexes["second"] + 100)}) con = con.set_index({"multi": ["second", "third"]}) - assert isinstance(con, Constraint) + assert isinstance(con, ConstraintBase) assert con.coord_dims == ( "first", "multi", @@ -256,13 +277,13 @@ def test_anonymous_constraint_set_index(x: linopy.Variable, y: linopy.Variable) def test_anonymous_constraint_loc(x: linopy.Variable, y: linopy.Variable) -> None: expr = 10 * x + y con = expr <= 10 - assert isinstance(con.loc[[1, 2]], Constraint) + assert isinstance(con.loc[[1, 2]], ConstraintBase) def test_anonymous_constraint_getitem(x: linopy.Variable, y: linopy.Variable) -> None: expr = 10 * x + y con = expr <= 10 - assert isinstance(con[1], Constraint) + assert isinstance(con[1], ConstraintBase) def test_constraint_from_rule(m: Model, x: linopy.Variable, y: linopy.Variable) -> None: @@ -270,8 +291,8 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint: return (i - 1) * x.at[i - 1] + y.at[j] >= 0 if i % 2 else i * x.at[i] >= 0 coords = [x.coords["first"], y.coords["second"]] - con = Constraint.from_rule(m, bound, coords) - assert isinstance(con, Constraint) + con = MutableConstraint.from_rule(m, bound, coords) + assert isinstance(con, ConstraintBase) assert con.lhs.nterm == 2 repr(con) # test repr @@ -285,8 +306,8 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint | None: return None coords = [x.coords["first"], y.coords["second"]] - con = Constraint.from_rule(m, bound, coords) - assert isinstance(con, Constraint) + con = MutableConstraint.from_rule(m, bound, coords) + assert isinstance(con, ConstraintBase) assert isinstance(con.lhs.vars, xr.DataArray) assert con.lhs.nterm == 2 assert (con.lhs.vars.loc[0, :] == -1).all() @@ -419,8 +440,8 @@ def test_constraint_labels_setter_invalid(c: linopy.constraints.Constraint) -> N def test_constraint_sel(c: linopy.constraints.Constraint) -> None: - assert isinstance(c.sel(first=[1, 2]), Constraint) - assert isinstance(c.isel(first=[1, 2]), Constraint) + assert isinstance(c.sel(first=[1, 2]), ConstraintBase) + assert isinstance(c.isel(first=[1, 2]), ConstraintBase) def test_constraint_flat(c: linopy.constraints.Constraint) -> None: @@ -429,7 +450,7 @@ def test_constraint_flat(c: linopy.constraints.Constraint) -> None: def test_iterate_slices(c: linopy.constraints.Constraint) -> None: for i in c.iterate_slices(slice_size=2): - assert isinstance(i, Constraint) + assert isinstance(i, ConstraintBase) assert c.coord_dims == i.coord_dims @@ -581,7 +602,7 @@ def test_constraint_with_helper_dims_as_coords(m: Model) -> None: data = xr.Dataset({"coeffs": coeffs, "vars": vars, "sign": sign, "rhs": rhs}) assert set(HELPER_DIMS).intersection(set(data.coords)) - con = Constraint(data, m, "c") + con = MutableConstraint(data, m, "c") expr = m.add_constraints(con) assert not set(HELPER_DIMS).intersection(set(expr.data.coords)) diff --git a/test/test_scalar_constraint.py b/test/test_scalar_constraint.py index cf5b3724..4872fada 100644 --- a/test/test_scalar_constraint.py +++ b/test/test_scalar_constraint.py @@ -6,7 +6,7 @@ import linopy from linopy import GREATER_EQUAL, Model, Variable -from linopy.constraints import AnonymousScalarConstraint, Constraint +from linopy.constraints import AnonymousScalarConstraint, ConstraintBase @pytest.fixture @@ -32,20 +32,20 @@ def test_anonymous_scalar_constraint_type(x: Variable) -> None: def test_simple_constraint_type(m: Model, x: Variable) -> None: - c: Constraint = m.add_constraints(x.at[0] >= 0) - assert isinstance(c, linopy.constraints.Constraint) + c: ConstraintBase = m.add_constraints(x.at[0] >= 0) + assert isinstance(c, linopy.constraints.ConstraintBase) def test_compound_constraint_type(m: Model, x: Variable) -> None: - c: Constraint = m.add_constraints(x.at[0] + x.at[1] >= 0) - assert isinstance(c, linopy.constraints.Constraint) + c: ConstraintBase = m.add_constraints(x.at[0] + x.at[1] >= 0) + assert isinstance(c, linopy.constraints.ConstraintBase) def test_explicit_simple_constraint_type(m: Model, x: Variable) -> None: - c: Constraint = m.add_constraints(x.at[0], GREATER_EQUAL, 0) - assert isinstance(c, linopy.constraints.Constraint) + c: ConstraintBase = m.add_constraints(x.at[0], GREATER_EQUAL, 0) + assert isinstance(c, linopy.constraints.ConstraintBase) def test_explicit_compound_constraint_type(m: Model, x: Variable) -> None: - c: Constraint = m.add_constraints(x.at[0] + x.at[1], GREATER_EQUAL, 0) - assert isinstance(c, linopy.constraints.Constraint) + c: ConstraintBase = m.add_constraints(x.at[0] + x.at[1], GREATER_EQUAL, 0) + assert isinstance(c, linopy.constraints.ConstraintBase) From b4dc1eacc9203e49e14a1c7111dda57a4f89a7b5 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Thu, 19 Mar 2026 15:57:00 +0100 Subject: [PATCH 05/20] Add io.to_netcdf support for frozen Constraint --- linopy/constraints.py | 54 +++++++++++++++++++++++++++++++++++++++++++ linopy/io.py | 9 ++++++-- test/test_io.py | 17 ++++++++++++++ 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 8d0fd783..617464a7 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -408,6 +408,10 @@ def to_matrix(self) -> scipy.sparse.csr_array: # They will be summed automatically upon conversion to CSC or dense. return scipy.sparse.csr_array((data, cols, indptr), shape=shape) + def to_netcdf_ds(self) -> Dataset: + """Return a Dataset representation suitable for netcdf serialization.""" + return self.data + iterate_slices = iterate_slices @@ -764,6 +768,56 @@ def to_matrix(self) -> scipy.sparse.csr_array: """Return the stored CSR matrix directly (no reconstruction needed).""" return self._csr + def to_netcdf_ds(self) -> Dataset: + """Return a Dataset with raw CSR components for netcdf serialization.""" + from xarray import DataArray + + csr = self._csr + data_vars: dict[str, DataArray] = { + "indptr": DataArray(csr.indptr, dims=["_indptr"]), + "indices": DataArray(csr.indices, dims=["_nnz"]), + "data": DataArray(csr.data, dims=["_nnz"]), + "rhs": DataArray(self._rhs, dims=["_flat"]), + } + for c in self._coords: + data_vars[f"_coord_{c.name}"] = DataArray( + np.array(c), dims=[f"_coorddim_{c.name}"] + ) + if self._dual is not None: + data_vars["dual"] = DataArray(self._dual, dims=["_flat"]) + dim_names = [c.name for c in self._coords] + return Dataset( + data_vars, + attrs={ + "_linopy_format": "csr", + "sign": self._sign, + "cindex": self._cindex if self._cindex is not None else -1, + "shape": list(csr.shape), + "coord_dims": dim_names, + "name": self._name, + }, + ) + + @classmethod + def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint: + """Reconstruct a Constraint from a netcdf Dataset (CSR format).""" + attrs = ds.attrs + shape = tuple(attrs["shape"]) + csr = scipy.sparse.csr_array( + (ds["data"].values, ds["indices"].values, ds["indptr"].values), + shape=shape, + ) + rhs = ds["rhs"].values + sign = attrs["sign"] + cindex = int(attrs["cindex"]) + cindex = cindex if cindex >= 0 else None + coord_dims = attrs["coord_dims"] + if isinstance(coord_dims, str): + coord_dims = [coord_dims] + coords = [pd.Index(ds[f"_coord_{d}"].values, name=d) for d in coord_dims] + dual = ds["dual"].values if "dual" in ds else None + return cls(csr, rhs, sign, coords, model, name, cindex=cindex, dual=dual) + def freeze(self) -> Constraint: """Return self (already immutable).""" return self diff --git a/linopy/io.py b/linopy/io.py index 39b56783..8c0d8541 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1131,7 +1131,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: with_prefix(var.data, f"variables-{name}") for name, var in m.variables.items() ] cons = [ - with_prefix(con.data, f"constraints-{name}") + with_prefix(con.to_netcdf_ds(), f"constraints-{name}") for name, con in m.constraints.items() ] objective = m.objective.data @@ -1167,6 +1167,7 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: ------- m : linopy.Model """ + from linopy.constraints import Constraint from linopy.model import ( Constraints, LinearExpression, @@ -1224,7 +1225,11 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: constraints = {} for k in sorted(con_names): name = remove_prefix(k, "constraints") - constraints[name] = MutableConstraint(get_prefix(ds, k), m, name) + con_ds = get_prefix(ds, k) + if con_ds.attrs.get("_linopy_format") == "csr": + constraints[name] = Constraint.from_netcdf_ds(con_ds, m, name) + else: + constraints[name] = MutableConstraint(con_ds, m, name) m._constraints = Constraints(constraints, m) objective = get_prefix(ds, "objective") diff --git a/test/test_io.py b/test/test_io.py index e8ded144..3d269636 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -78,6 +78,23 @@ def test_model_to_netcdf(model: Model, tmp_path: Path) -> None: assert_model_equal(m, p) +def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None: + from linopy.constraints import Constraint + + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") + m.add_constraints(x >= 1, name="c", freeze=True) + + assert isinstance(m.constraints["c"], Constraint) + + fn = tmp_path / "test_frozen.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert isinstance(p.constraints["c"], Constraint) + assert_model_equal(m, p) + + def test_model_to_netcdf_with_sense(model: Model, tmp_path: Path) -> None: m = model m.objective.sense = "max" From 304a2e7ac895fc514b96bfd87bba6d437561a827 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Thu, 19 Mar 2026 17:15:07 +0100 Subject: [PATCH 06/20] fix: re-implement matrices --- linopy/constraints.py | 85 ++++++-------- linopy/matrices.py | 254 ++++++++++++++++++++-------------------- linopy/model.py | 14 +-- test/test_constraint.py | 62 ++++++---- 4 files changed, 211 insertions(+), 204 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 617464a7..99e82114 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -404,8 +404,6 @@ def to_matrix(self) -> scipy.sparse.csr_array: np.cumsum(counts, out=indptr[1:]) shape = (len(labels_flat), self.model._xCounter) - # Note: duplicate (row, col) entries are not summed in CSR format. - # They will be summed automatically upon conversion to CSC or dense. return scipy.sparse.csr_array((data, cols, indptr), shape=shape) def to_netcdf_ds(self) -> Dataset: @@ -1516,75 +1514,66 @@ def flat(self) -> pd.DataFrame: df["key"] = df.labels.map(map_labels) return df - def to_matrix(self, filter_missings: bool = True) -> scipy.sparse.csc_matrix: - """ - Construct a constraint matrix in sparse format. - - Missing values, i.e. -1 in labels and vars, are ignored filtered - out. - """ - # TODO: rename "filter_missings" to "~labels_as_coordinates" - cons = self.flat - - if not len(self): - raise ValueError("No constraints available to convert to matrix.") - - if filter_missings: - vars = self.model.variables.flat - shape = (cons.key.max() + 1, vars.key.max() + 1) - cons["vars"] = cons.vars.map(vars.set_index("labels").key) - return scipy.sparse.csc_matrix( - (cons.coeffs, (cons.key, cons.vars)), shape=shape - ) - else: - shape = self.model.shape - return scipy.sparse.csc_matrix( - (cons.coeffs, (cons.labels, cons.vars)), shape=shape - ) - - def to_matrix_via_csr( - self, - ) -> tuple[scipy.sparse.csc_array, np.ndarray, np.ndarray]: + def to_matrix( + self, filter_missings: bool = True + ) -> tuple[scipy.sparse.csc_array, np.ndarray, np.ndarray | None]: """ Construct a constraint matrix in sparse format by stacking per-constraint CSR matrices. + Parameters + ---------- + filter_missings : bool, default True + If True, also strip empty columns and return ``var_labels`` for + remapping columns back to original variable labels. + If False, return full-width CSC with shape + ``(n_active_cons, model._xCounter)`` and ``var_labels=None``. + ``con_labels`` is always returned. + Returns ------- matrix : scipy.sparse.csc_array - Shape (n_con_labels, n_var_labels), containing only non-empty rows and columns. + Shape ``(n_active_cons, n_active_vars)`` when + ``filter_missings=True``, or ``(n_active_cons, + model._xCounter)`` when ``filter_missings=False``. con_labels : np.ndarray - Shape (n_con_labels,), maps each matrix row to the original constraint label. - var_labels : np.ndarray - Shape (n_var_labels,), maps each matrix column to the original variable label. + Shape ``(n_active_cons,)``, maps each matrix row to the + original constraint label. + var_labels : np.ndarray or None + Shape ``(n_active_vars,)``, maps each matrix column to the + original variable label. ``None`` when + ``filter_missings=False``. """ if not len(self): raise ValueError("No constraints available to convert to matrix.") - matrices = [] + active_csrs = [] con_labels_list = [] for c in self.data.values(): csr = c.to_matrix() - matrices.append(csr) nonempty = np.diff(csr.indptr).astype(bool) + active_csrs.append(csr[nonempty]) start = ( c._cindex if isinstance(c, Constraint) else c.data.attrs["label_range"][0] ) con_labels_list.append(np.flatnonzero(nonempty) + start) - csc: scipy.sparse.csc_array = scipy.sparse.vstack(matrices).tocsc() + csc: scipy.sparse.csc_array = scipy.sparse.vstack(active_csrs).tocsc() + csc.sum_duplicates() con_labels = np.concatenate(con_labels_list) - indptr = csc.indptr - nonempty_cols = indptr[1:] != indptr[:-1] - new_indptr = np.r_[0, indptr[1:][nonempty_cols]] - (var_labels,) = np.nonzero(nonempty_cols) - - matrix = scipy.sparse.csc_array( - (csc.data, csc.indices, new_indptr), - shape=(csc.shape[0], len(var_labels)), - ) - return matrix, con_labels, var_labels + if filter_missings: + indptr = csc.indptr + nonempty_cols = indptr[1:] != indptr[:-1] + new_indptr = np.r_[0, indptr[1:][nonempty_cols]] + (var_labels,) = np.nonzero(nonempty_cols) + matrix = scipy.sparse.csc_array( + (csc.data, csc.indices, new_indptr), + shape=(csc.shape[0], len(var_labels)), + ) + return matrix, con_labels, var_labels + else: + return csc, con_labels, None def reset_dual(self) -> None: """ diff --git a/linopy/matrices.py b/linopy/matrices.py index e1489e76..b81d50db 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -7,175 +7,179 @@ from __future__ import annotations -from functools import cached_property from typing import TYPE_CHECKING import numpy as np -import pandas as pd from numpy import ndarray -from pandas.core.indexes.base import Index -from pandas.core.series import Series from scipy.sparse._csc import csc_matrix from linopy import expressions +from linopy.constraints import Constraint if TYPE_CHECKING: from linopy.model import Model -def create_vector( - indices: Series | Index, - values: Series | ndarray, - fill_value: str | float | int = np.nan, - shape: int | None = None, -) -> ndarray: - """Create a vector of a size equal to the maximum index plus one.""" - if shape is None: - max_value = indices.max() - if not isinstance(max_value, np.integer | int): - raise ValueError("Indices must be integers.") - shape = max_value + 1 - vector = np.full(shape, fill_value) - vector[indices] = values - return vector - - class MatrixAccessor: """ Helper class to quickly access model related vectors and matrices. + + All arrays are compact — only active (non-masked) entries are included. + Position i in variable-side arrays corresponds to vlabels[i]. + Position i in constraint-side arrays corresponds to clabels[i]. """ def __init__(self, model: Model) -> None: self._parent = model + self._build_vars() + self._build_cons() - def clean_cached_properties(self) -> None: - """Clear the cache for all cached properties of an object""" - - for cached_prop in ["flat_vars", "flat_cons", "sol", "dual"]: - # check existence of cached_prop without creating it - if cached_prop in self.__dict__: - delattr(self, cached_prop) - - @cached_property - def flat_vars(self) -> pd.DataFrame: + def _build_vars(self) -> None: m = self._parent - return m.variables.flat - @cached_property - def flat_cons(self) -> pd.DataFrame: - m = self._parent - return m.constraints.flat + labels_list = [] + lb_list = [] + ub_list = [] + vtypes_list = [] - @property - def vlabels(self) -> ndarray: - """Vector of labels of all non-missing variables.""" - df: pd.DataFrame = self.flat_vars - return create_vector(df.key, df.labels, -1) + for name, var in m.variables.items(): + labels = var.labels.values.ravel() + mask = labels != -1 + active_labels = labels[mask] - @property - def vtypes(self) -> ndarray: - """Vector of types of all non-missing variables.""" - m = self._parent - df: pd.DataFrame = self.flat_vars - specs = [] - for name in m.variables: if name in m.binaries: - val = "B" + vtype = "B" elif name in m.integers: - val = "I" + vtype = "I" elif name in m.semi_continuous: - val = "S" + vtype = "S" else: - val = "C" - specs.append(pd.Series(val, index=m.variables[name].flat.labels)) - - ds = pd.concat(specs) - ds = df.set_index("key").labels.map(ds) - return create_vector(ds.index, ds.to_numpy(), fill_value="") - - @property - def lb(self) -> ndarray: - """Vector of lower bounds of all non-missing variables.""" - df: pd.DataFrame = self.flat_vars - return create_vector(df.key, df.lower) - - @cached_property - def sol(self) -> ndarray: - """Vector of solution values of all non-missing variables.""" - if not self._parent.status == "ok": - raise ValueError("Model is not optimized.") - if "solution" not in self.flat_vars: - del self.flat_vars # clear cache - df: pd.DataFrame = self.flat_vars - return create_vector(df.key, df.solution, fill_value=np.nan) - - @cached_property - def dual(self) -> ndarray: - """Vector of dual values of all non-missing constraints.""" - if not self._parent.status == "ok": - raise ValueError("Model is not optimized.") - if "dual" not in self.flat_cons: - del self.flat_cons # clear cache - df: pd.DataFrame = self.flat_cons - if "dual" not in df: - raise AttributeError( - "Underlying is optimized but does not have dual values stored." - ) - return create_vector(df.key, df.dual, fill_value=np.nan) - - @property - def ub(self) -> ndarray: - """Vector of upper bounds of all non-missing variables.""" - df: pd.DataFrame = self.flat_vars - return create_vector(df.key, df.upper) - - @property - def clabels(self) -> ndarray: - """Vector of labels of all non-missing constraints.""" - df: pd.DataFrame = self.flat_cons - if df.empty: - return np.array([], dtype=int) - return create_vector(df.key, df.labels, fill_value=-1) - - @property - def A(self) -> csc_matrix | None: - """Constraint matrix of all non-missing constraints and variables.""" + vtype = "C" + + labels_list.append(active_labels) + lb_list.append(var.lower.values.ravel()[mask]) + ub_list.append(var.upper.values.ravel()[mask]) + vtypes_list.append(np.full(mask.sum(), vtype)) + + if labels_list: + vlabels = np.concatenate(labels_list) + order = np.argsort(vlabels) + self.vlabels: ndarray = vlabels[order] + self.lb: ndarray = np.concatenate(lb_list)[order] + self.ub: ndarray = np.concatenate(ub_list)[order] + self.vtypes: ndarray = np.concatenate(vtypes_list)[order] + else: + self.vlabels = np.array([], dtype=np.intp) + self.lb = np.array([]) + self.ub = np.array([]) + self.vtypes = np.array([], dtype=object) + + def _build_cons(self) -> None: m = self._parent - if not len(m.constraints): - return None - A: csc_matrix = m.constraints.to_matrix(filter_missings=False) - return A[self.clabels][:, self.vlabels] - @property - def sense(self) -> ndarray: - """Vector of senses of all non-missing constraints.""" - df: pd.DataFrame = self.flat_cons - return create_vector(df.key, df.sign.astype(np.dtype(" ndarray: - """Vector of right-hand-sides of all non-missing constraints.""" - df: pd.DataFrame = self.flat_cons - return create_vector(df.key, df.rhs) + self.b = np.concatenate(b_list) if b_list else np.array([]) + self.sense = ( + np.concatenate(sense_list) if sense_list else np.array([], dtype=object) + ) @property def c(self) -> ndarray: - """Vector of objective coefficients of all non-missing variables.""" + """Objective coefficients aligned with vlabels.""" m = self._parent + result = np.zeros(len(self.vlabels)) + ds = m.objective.flat if isinstance(m.objective.expression, expressions.QuadraticExpression): - ds = ds[(ds.vars1 == -1) | (ds.vars2 == -1)] + ds = ds[(ds.vars1 == -1) | (ds.vars2 == -1)].copy() ds["vars"] = ds.vars1.where(ds.vars1 != -1, ds.vars2) - vars: pd.Series = ds.vars.map(self.flat_vars.set_index("labels").key) - shape: int = self.flat_vars.key.max() + 1 - return create_vector(vars, ds.coeffs, fill_value=0.0, shape=shape) + var_labels = ds.vars.values + coeffs = ds.coeffs.values + mask = var_labels != -1 + active_labels = var_labels[mask] + positions = np.searchsorted(self.vlabels, active_labels) + valid = (positions < len(self.vlabels)) & ( + self.vlabels[positions] == active_labels + ) + np.add.at(result, positions[valid], coeffs[mask][valid]) + return result @property def Q(self) -> csc_matrix | None: - """Matrix objective coefficients of quadratic terms of all non-missing variables.""" + """Quadratic objective matrix, shape (n_active_vars, n_active_vars).""" m = self._parent expr = m.objective.expression if not isinstance(expr, expressions.QuadraticExpression): return None return expr.to_matrix()[self.vlabels][:, self.vlabels] + + @property + def sol(self) -> ndarray: + """Solution values aligned with vlabels.""" + if not self._parent.status == "ok": + raise ValueError("Model is not optimized.") + m = self._parent + result = np.full(len(self.vlabels), np.nan) + for _, var in m.variables.items(): + labels = var.labels.values.ravel() + mask = labels != -1 + active_labels = labels[mask] + positions = np.searchsorted(self.vlabels, active_labels) + result[positions] = var.solution.values.ravel()[mask] + return result + + @property + def dual(self) -> ndarray: + """Dual values aligned with clabels.""" + if not self._parent.status == "ok": + raise ValueError("Model is not optimized.") + m = self._parent + dual_list = [] + has_dual = False + for c in m.constraints.data.values(): + csr = c.to_matrix() + nonempty = np.diff(csr.indptr).astype(bool) + active_rows = np.flatnonzero(nonempty) + if isinstance(c, Constraint): + if c._dual is not None: + dual_list.append(c._dual[active_rows]) + has_dual = True + else: + dual_list.append(np.full(len(active_rows), np.nan)) + else: + if "dual" in c.data: + dual_list.append(c.dual.values.ravel()[active_rows]) + has_dual = True + else: + dual_list.append(np.full(len(active_rows), np.nan)) + if not has_dual: + raise AttributeError( + "Underlying is optimized but does not have dual values stored." + ) + return np.concatenate(dual_list) if dual_list else np.array([]) diff --git a/linopy/model.py b/linopy/model.py index b1310831..20d47e40 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -134,8 +134,6 @@ class Model: _chunk: T_Chunks _force_dim_names: bool _solver_dir: Path - matrices: MatrixAccessor - __slots__ = ( # containers "_variables", @@ -161,7 +159,6 @@ class Model: "_solver_dir", "solver_model", "solver_name", - "matrices", ) def __init__( @@ -219,7 +216,10 @@ def __init__( gettempdir() if solver_dir is None else solver_dir ) - self.matrices: MatrixAccessor = MatrixAccessor(self) + @property + def matrices(self) -> MatrixAccessor: + """Matrix representation of the model, computed fresh on each access.""" + return MatrixAccessor(self) @property def variables(self) -> Variables: @@ -1397,9 +1397,6 @@ def solve( sanitize_zeros=sanitize_zeros, sanitize_infinities=sanitize_infinities ) - # clear cached matrix properties potentially present from previous solve commands - self.matrices.clean_cached_properties() - # check io_api if io_api is not None and io_api not in IO_APIS: raise ValueError( @@ -1615,9 +1612,6 @@ def _mock_solve( ) -> tuple[str, str]: solver_name = "mock" - # clear cached matrix properties potentially present from previous solve commands - self.matrices.clean_cached_properties() - logger.info(f" Solve problem using {solver_name.title()} solver") # reset result self.reset_solution() diff --git a/test/test_constraint.py b/test/test_constraint.py index 483d0835..098b91b8 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -609,8 +609,17 @@ def test_constraint_with_helper_dims_as_coords(m: Model) -> None: def test_constraint_matrix(m: Model) -> None: - A = m.constraints.to_matrix() - assert A.shape == (10, 14) + # filter_missings=True: strips empty columns, returns labels for remapping + A, con_labels, var_labels = m.constraints.to_matrix() + assert A.shape == (10, 10) # only x (10 vars) appears in constraints + assert con_labels is not None + assert var_labels is not None + + # filter_missings=False: full width, con_labels always returned, var_labels=None + A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=False) + assert A.shape == (10, m.shape[1]) # all 14 vars as columns + assert len(con_labels) == 10 + assert var_labels is None def test_constraint_matrix_masked_variables() -> None: @@ -627,48 +636,59 @@ def test_constraint_matrix_masked_variables() -> None: x = m.add_variables(coords=[range(10)], mask=mask) m.add_variables() m.add_constraints(x, EQUAL, 0) - A = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (5, 6) - assert A.shape == (m.ncons, m.nvars) + # filter_missings=True: only vars/cons that appear in constraints + A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=True) + assert A.shape == (m.ncons, 5) # 5 active x vars; scalar var not in any constraint + assert len(con_labels) == m.ncons + assert len(var_labels) == 5 - A = m.constraints.to_matrix(filter_missings=False) - assert A.shape == m.shape + # filter_missings=False: full width, no remapping labels + A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=False) + assert A.shape == (m.ncons, m.shape[1]) + assert con_labels is not None + assert var_labels is None def test_constraint_matrix_masked_constraints() -> None: """ Test constraint matrix with missing constraints. """ - # now with missing variables m = Model() mask = pd.Series([False] * 5 + [True] * 5) x = m.add_variables(coords=[range(10)]) m.add_variables() m.add_constraints(x, EQUAL, 0, mask=mask) - A = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (5, 11) - assert A.shape == (m.ncons, m.nvars) + # filter_missings=True: only active cons and vars they reference + # active cons are indices 5-9, which reference vars 5-9 only + A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=True) + assert A.shape == (m.ncons, m.ncons) # 5 active cons; 5 referenced vars + assert len(con_labels) == m.ncons - A = m.constraints.to_matrix(filter_missings=False) - assert A.shape == m.shape + A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=False) + assert A.shape == (m.ncons, m.shape[1]) + assert con_labels is not None + assert var_labels is None def test_constraint_matrix_masked_constraints_and_variables() -> None: """ - Test constraint matrix with missing constraints. + Test constraint matrix with missing constraints and variables. """ - # now with missing variables m = Model() mask = pd.Series([False] * 5 + [True] * 5) x = m.add_variables(coords=[range(10)], mask=mask) m.add_variables() m.add_constraints(x, EQUAL, 0, mask=mask) - A = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (5, 6) - assert A.shape == (m.ncons, m.nvars) - - A = m.constraints.to_matrix(filter_missings=False) - assert A.shape == m.shape + # filter_missings=True: 5 active cons x 5 active vars + A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=True) + assert A.shape == (m.ncons, m.ncons) # both masks align: 5x5 + assert len(con_labels) == m.ncons + assert len(var_labels) == m.ncons + + A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=False) + assert A.shape == (m.ncons, m.shape[1]) + assert con_labels is not None + assert var_labels is None def test_get_name_by_label() -> None: From 0fc2673673f12909528e12e055ff1a777e1cd630 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Fri, 20 Mar 2026 08:50:35 +0100 Subject: [PATCH 07/20] Move sum_duplicates --- linopy/constraints.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 99e82114..5ceb6e9d 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -404,7 +404,9 @@ def to_matrix(self) -> scipy.sparse.csr_array: np.cumsum(counts, out=indptr[1:]) shape = (len(labels_flat), self.model._xCounter) - return scipy.sparse.csr_array((data, cols, indptr), shape=shape) + csr = scipy.sparse.csr_array((data, cols, indptr), shape=shape) + csr.sum_duplicates() + return csr def to_netcdf_ds(self) -> Dataset: """Return a Dataset representation suitable for netcdf serialization.""" @@ -1559,7 +1561,6 @@ def to_matrix( ) con_labels_list.append(np.flatnonzero(nonempty) + start) csc: scipy.sparse.csc_array = scipy.sparse.vstack(active_csrs).tocsc() - csc.sum_duplicates() con_labels = np.concatenate(con_labels_list) if filter_missings: From 3c8c5d69d469316ef5e3e52ed7c6b6afeb0fa600 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Sat, 21 Mar 2026 17:56:55 +0100 Subject: [PATCH 08/20] feat: VariableLabelIndex --- linopy/common.py | 45 ++++++++++++++++++++++++++++++++++++++++++++- linopy/variables.py | 11 +++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index ac0b3af3..2e0d9af1 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -10,7 +10,7 @@ import operator import os from collections.abc import Callable, Generator, Hashable, Iterable, Sequence -from functools import partial, reduce, wraps +from functools import cached_property, partial, reduce, wraps from pathlib import Path from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload from warnings import warn @@ -939,6 +939,49 @@ def find_single(value: int) -> tuple[str, dict] | tuple[None, None]: raise ValueError("Array's with more than two dimensions is not supported") +class VariableLabelIndex: + """ + Index for O(1) mapping between variable labels and dense positions. + + Both arrays are computed lazily and cached: + - ``vlabels``: active variable labels in encounter order, shape (n_active_vars,) + - ``label_to_pos``: derived from vlabels; size _xCounter, maps label -> position (-1 if masked) + + Invalidated by clearing the instance ``__dict__`` when variables are added or removed. + """ + + def __init__(self, variables: Any) -> None: + self._variables = variables + + @cached_property + def vlabels(self) -> np.ndarray: + label_lists = [] + for _, var in self._variables.items(): + labels = var.labels.values.ravel() + mask = labels != -1 + label_lists.append(labels[mask]) + return ( + np.concatenate(label_lists) if label_lists else np.array([], dtype=np.intp) + ) + + @cached_property + def label_to_pos(self) -> np.ndarray: + vlabels = self.vlabels + n = self._variables.model._xCounter + label_to_pos = np.full(n, -1, dtype=np.intp) + label_to_pos[vlabels] = np.arange(len(vlabels), dtype=np.intp) + return label_to_pos + + @property + def n_active_vars(self) -> int: + return len(self.vlabels) + + def invalidate(self) -> None: + """Clear cached arrays so they are recomputed on next access.""" + self.__dict__.pop("vlabels", None) + self.__dict__.pop("label_to_pos", None) + + def get_label_position( obj: Any, values: int | np.ndarray, diff --git a/linopy/variables.py b/linopy/variables.py index f5e52a88..8a465a39 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -33,6 +33,7 @@ from linopy.common import ( LabelPositionIndex, LocIndexer, + VariableLabelIndex, as_dataarray, assign_multiindex_safe, check_has_nulls, @@ -1327,6 +1328,7 @@ class Variables: data: dict[str, Variable] model: Model _label_position_index: LabelPositionIndex | None = None + _variable_label_index: VariableLabelIndex | None = None dataset_attrs = ["labels", "lower", "upper"] dataset_names = ["Labels", "Lower bounds", "Upper bounds"] @@ -1436,6 +1438,15 @@ def _invalidate_label_position_index(self) -> None: """Invalidate the label position index cache.""" if self._label_position_index is not None: self._label_position_index.invalidate() + if self._variable_label_index is not None: + self._variable_label_index.invalidate() + + @property + def label_index(self) -> VariableLabelIndex: + """Index for O(1) label->position mapping and compact vlabels array.""" + if self._variable_label_index is None: + self._variable_label_index = VariableLabelIndex(self) + return self._variable_label_index @property def attrs(self) -> dict[Any, Any]: From 0b9de003028131ff0c511373cd04172898093e4c Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Sun, 22 Mar 2026 02:00:20 +0100 Subject: [PATCH 09/20] fix: until solve --- linopy/common.py | 54 +++++ linopy/constraints.py | 495 ++++++++++++++++++++++++-------------- linopy/matrices.py | 94 ++++---- linopy/model.py | 31 ++- test/test_constraint.py | 205 ++++++++-------- test/test_constraints.py | 4 +- test/test_model.py | 5 +- test/test_optimization.py | 2 +- 8 files changed, 543 insertions(+), 347 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 2e0d9af1..c47e9775 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -1401,6 +1401,60 @@ def __call__(self) -> bool: return self.value +def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]: + """ + Serialize a list of pd.Index (including MultiIndex) to a DataArray dict. + + Suitable for embedding coordinate metadata as plain data variables in a + Dataset that has its own unrelated dimensions (e.g. CSR netcdf format). + Reconstruct with :func:`coords_from_dataset`. + """ + data_vars: dict[str, DataArray] = {} + for c in coords: + if isinstance(c, pd.MultiIndex): + for level_name, level_values in zip(c.names, c.levels): + data_vars[f"_coord_{c.name}_level_{level_name}"] = DataArray( + np.array(level_values), + dims=[f"_coorddim_{c.name}_level_{level_name}"], + ) + data_vars[f"_coord_{c.name}_codes"] = DataArray( + np.array(c.codes).T, + dims=[f"_coorddim_{c.name}", f"_coorddim_{c.name}_nlevels"], + ) + else: + data_vars[f"_coord_{c.name}"] = DataArray( + np.array(c), dims=[f"_coorddim_{c.name}"] + ) + return data_vars + + +def coords_from_dataset(ds: Dataset, coord_dims: list[str]) -> list[pd.Index]: + """ + Deserialize a list of pd.Index (including MultiIndex) from a Dataset. + + Reconstructs coordinates previously serialized by :func:`coords_to_dataset_vars`. + """ + coords = [] + for d in coord_dims: + if f"_coord_{d}_codes" in ds: + codes_2d = ds[f"_coord_{d}_codes"].values.T + level_names = [ + k[len(f"_coord_{d}_level_") :] + for k in ds + if k.startswith(f"_coord_{d}_level_") + ] + arrays = [ + ds[f"_coord_{d}_level_{ln}"].values[codes_2d[i]] + for i, ln in enumerate(level_names) + ] + mi = pd.MultiIndex.from_arrays(arrays, names=level_names) + mi.name = d + coords.append(mi) + else: + coords.append(pd.Index(ds[f"_coord_{d}"].values, name=d)) + return coords + + def is_constant(x: SideLike) -> bool: """ Check if the given object is a constant type or an expression type without diff --git a/linopy/constraints.py b/linopy/constraints.py index 5ceb6e9d..096632d3 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -33,10 +33,13 @@ from linopy.common import ( LabelPositionIndex, LocIndexer, + VariableLabelIndex, align_lines_by_delimiter, assign_multiindex_safe, check_has_nulls, check_has_nulls_polars, + coords_from_dataset, + coords_to_dataset_vars, filter_nulls_polars, format_string_as_variable_name, generate_indices_for_printout, @@ -166,6 +169,33 @@ def rhs(self) -> DataArray: def dual(self) -> DataArray: """Get the dual values DataArray.""" + @abstractmethod + def has_variable(self, variable: variables.Variable) -> bool: + """Check if the constraint references any of the given variable labels.""" + + @abstractmethod + def sanitize_zeros(self) -> ConstraintBase: + """Remove terms with zero or near-zero coefficients.""" + + @abstractmethod + def sanitize_missings(self) -> ConstraintBase: + """Mask out rows where all variables are missing (-1).""" + + @abstractmethod + def sanitize_infinities(self) -> ConstraintBase: + """Mask out rows with invalid infinite RHS values.""" + + @abstractmethod + def to_matrix_with_rhs( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: + """ + Return (csr, con_labels, b, sense) in one pass. + + Avoids computing the CSR matrix twice when both the matrix and + the RHS/sense vectors are needed. + """ + def __getitem__( self, selector: str | int | slice | list | tuple | dict ) -> MutableConstraint: @@ -377,36 +407,59 @@ def mask_func(data: pd.DataFrame) -> pd.Series: check_has_nulls(df, name=f"{self.type} {self.name}") return df - def to_matrix(self) -> scipy.sparse.csr_array: + def to_matrix( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray]: """ - Construct a CSR matrix representation of this constraint. + Construct a dense CSR matrix for this constraint. + + Only active (non-masked) rows are included. Column indices are dense + positions in the active variable array, as given by ``label_index``. - All flat positions in the constraint grid are included as rows; - masked entries (labels == -1) become empty rows. Shape is - (len(labels_flat), model._xCounter). + Parameters + ---------- + label_index : VariableLabelIndex + Variable label index providing ``label_to_pos`` and ``n_active_vars``. Returns ------- - matrix : scipy.sparse.csr_array + csr : scipy.sparse.csr_array + Shape (n_active_cons, n_active_vars). + con_labels : np.ndarray + Active constraint labels in row order. """ - vars = self.vars.values + label_to_pos = label_index.label_to_pos labels_flat = self.labels.values.ravel() - vars_2d = vars.reshape(len(labels_flat), -1) - coeffs_flat = self.coeffs.values.ravel() + vars_vals = self.vars.values + n_rows = len(labels_flat) + vars_2d = ( + vars_vals.reshape(n_rows, -1) + if n_rows > 0 + else vars_vals.reshape(0, max(1, vars_vals.size)) + ) + + # Single boolean mask combining both filter conditions + row_mask = (labels_flat != -1) & (vars_2d != -1).any(axis=1) + con_labels = labels_flat[row_mask] + vars_final = vars_2d[row_mask] + valid_final = vars_final != -1 - valid_2d = (labels_flat != -1)[:, np.newaxis] & (vars_2d != -1) - cols = vars_2d[valid_2d] - data = coeffs_flat.reshape(vars_2d.shape)[valid_2d] + coeffs_final = self.coeffs.values.ravel().reshape(vars_2d.shape)[row_mask] - counts = valid_2d.sum(axis=1) - indptr = np.empty(len(labels_flat) + 1, dtype=np.int32) + cols = label_to_pos[vars_final[valid_final]] + data = coeffs_final[valid_final] + + counts = valid_final.sum(axis=1) + n_active_cons = len(con_labels) + indptr = np.empty(n_active_cons + 1, dtype=np.int32) indptr[0] = 0 np.cumsum(counts, out=indptr[1:]) - shape = (len(labels_flat), self.model._xCounter) - csr = scipy.sparse.csr_array((data, cols, indptr), shape=shape) + csr = scipy.sparse.csr_array( + (data, cols, indptr), shape=(n_active_cons, label_index.n_active_vars) + ) csr.sum_duplicates() - return csr + return csr, con_labels def to_netcdf_ds(self) -> Dataset: """Return a Dataset representation suitable for netcdf serialization.""" @@ -443,6 +496,7 @@ class Constraint(ConstraintBase): __slots__ = ( "_csr", + "_con_labels", "_rhs", "_sign", "_coords", @@ -455,6 +509,7 @@ class Constraint(ConstraintBase): def __init__( self, csr: scipy.sparse.csr_array, + con_labels: np.ndarray, rhs: np.ndarray, sign: str, coords: list[pd.Index], @@ -464,6 +519,7 @@ def __init__( dual: np.ndarray | None = None, ) -> None: self._csr = csr + self._con_labels = con_labels self._rhs = rhs self._sign = sign self._coords = coords @@ -491,14 +547,9 @@ def range(self) -> tuple[int, int]: raise AttributeError("Constraint has not been assigned labels yet.") return (self._cindex, self._cindex + self._csr.shape[0]) - @property - def _nonempty(self) -> np.ndarray: - """Boolean array of shape (n_flat,) — True where row is non-masked.""" - return np.diff(self._csr.indptr).astype(bool) - @property def ncons(self) -> int: - return int(np.diff(self._csr.indptr).astype(bool).sum()) + return self._csr.shape[0] @property def attrs(self) -> dict[str, Any]: @@ -510,10 +561,16 @@ def attrs(self) -> dict[str, Any]: @property def dims(self) -> Frozen[Hashable, int]: d: dict[Hashable, int] = {c.name: len(c) for c in self._coords} - nterm = int(self._csr.indptr.max()) if self._csr.nnz > 0 else 1 - d[TERM_DIM] = nterm + d[TERM_DIM] = self.nterm return Frozen(d) + @property + def active_positions(self) -> np.ndarray: + """Flat positions of active (non-masked) rows in the original coord shape.""" + if self._cindex is None: + return np.arange(self._csr.shape[0]) + return self._con_labels - self._cindex + @property def sizes(self) -> Frozen[Hashable, int]: return self.dims @@ -524,7 +581,7 @@ def indexes(self) -> Indexes: @property def nterm(self) -> int: - return int(self._csr.indptr.max()) if self._csr.nnz > 0 else 1 + return int(np.diff(self._csr.indptr).max()) if self._csr.nnz > 0 else 1 @property def coord_names(self) -> list[str]: @@ -535,17 +592,14 @@ def labels(self) -> DataArray: """Get labels DataArray, shape (*coord_dims).""" if self._cindex is None: return DataArray([]) - n_flat = self._csr.shape[0] - labels_flat = np.where( - self._nonempty, - np.arange(self._cindex, self._cindex + n_flat), - -1, - ) shape = tuple(len(c) for c in self._coords) + full_size = int(np.prod(shape)) if shape else 1 + labels_flat = np.full(full_size, -1, dtype=np.int64) + labels_flat[self.active_positions] = self._con_labels dim_names = [c.name for c in self._coords] xr_coords = {c.name: c for c in self._coords} return DataArray( - labels_flat.reshape(shape) if shape else labels_flat, + labels_flat.reshape(shape), coords=xr_coords, dims=dim_names, ) @@ -590,8 +644,11 @@ def rhs(self) -> DataArray: shape = tuple(len(c) for c in self._coords) dim_names = [c.name for c in self._coords] xr_coords = {c.name: c for c in self._coords} + full_size = int(np.prod(shape)) if shape else 1 + rhs_full = np.full(full_size, np.nan) + rhs_full[self.active_positions] = self._rhs return DataArray( - self._rhs.reshape(shape) if shape else self._rhs, + rhs_full.reshape(shape) if shape else rhs_full.reshape(()), coords=xr_coords, dims=dim_names, ) @@ -607,12 +664,21 @@ def dual(self) -> DataArray: shape = tuple(len(c) for c in self._coords) dim_names = [c.name for c in self._coords] xr_coords = {c.name: c for c in self._coords} + full_size = int(np.prod(shape)) if shape else 1 + dual_full = np.full(full_size, np.nan) + dual_full[self.active_positions] = self._dual return DataArray( - self._dual.reshape(shape) if shape else self._dual, + dual_full.reshape(shape) if shape else dual_full.reshape(()), coords=xr_coords, dims=dim_names, ) + @dual.setter + def dual(self, value: DataArray) -> None: + """Set dual values from a DataArray aligned with the full coord shape.""" + vals = np.asarray(value).ravel() + self._dual = vals[self.active_positions] + def _to_dataset(self, nterm: int) -> Dataset: """ Reconstruct labels/coeffs/vars Dataset from the CSR matrix. @@ -627,41 +693,42 @@ def _to_dataset(self, nterm: int) -> Dataset: Dataset with variables ``labels``, ``coeffs``, ``vars``. """ csr = self._csr - n_total = csr.shape[0] + n_active = csr.shape[0] # active rows only (no masked rows in CSR) counts = np.diff(csr.indptr) - nonempty = counts > 0 - coeffs_2d = np.zeros((n_total, nterm), dtype=csr.dtype) - vars_2d = np.full((n_total, nterm), -1, dtype=np.int64) + shape = tuple(len(c) for c in self._coords) + full_size = int(np.prod(shape)) if shape else n_active + + # Map active row i -> flat position in full shape via con_labels + active_positions = self.active_positions + coeffs_2d = np.zeros((full_size, nterm), dtype=csr.dtype) + vars_2d = np.full((full_size, nterm), -1, dtype=np.int64) if csr.nnz > 0: - row_indices = np.repeat(np.arange(n_total, dtype=np.int32), counts) - term_cols = np.arange(csr.nnz, dtype=np.int32) - np.repeat( - csr.indptr[:-1].astype(np.int32), counts - ) - vars_2d[row_indices, term_cols] = csr.indices + row_indices = np.repeat(active_positions, counts) + term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) + # csr.indices are column positions into vlabels; map back to variable labels + vlabels = self._model.variables.label_index.vlabels + vars_2d[row_indices, term_cols] = vlabels[csr.indices] coeffs_2d[row_indices, term_cols] = csr.data - shape = tuple(len(c) for c in self._coords) + dim_names = [c.name for c in self._coords] xr_coords = {c.name: c for c in self._coords} - term_coords = {TERM_DIM: np.arange(nterm)} dims_with_term = dim_names + [TERM_DIM] - coords_with_term = {**xr_coords, **term_coords} coeffs_da = DataArray( - coeffs_2d.reshape(shape + (nterm,)) if shape else coeffs_2d, - coords=coords_with_term, + coeffs_2d.reshape(shape + (nterm,)), + coords=xr_coords, dims=dims_with_term, ) vars_da = DataArray( - vars_2d.reshape(shape + (nterm,)) if shape else vars_2d, - coords=coords_with_term, + vars_2d.reshape(shape + (nterm,)), + coords=xr_coords, dims=dims_with_term, ) ds = Dataset({"coeffs": coeffs_da, "vars": vars_da}) if self._cindex is not None: - labels_flat = np.where( - nonempty, np.arange(self._cindex, self._cindex + n_total), -1 - ) + labels_flat = np.full(full_size, -1, dtype=np.int64) + labels_flat[active_positions] = self._con_labels ds["labels"] = DataArray( - labels_flat.reshape(shape) if shape else labels_flat, + labels_flat.reshape(shape), coords=xr_coords, dims=dim_names, ) @@ -670,32 +737,41 @@ def _to_dataset(self, nterm: int) -> Dataset: @property def data(self) -> Dataset: """Reconstruct the xarray Dataset from the CSR representation.""" - nterm = int(self._csr.indptr.max()) if self._csr.nnz > 0 else 1 - ds = self._to_dataset(nterm) - sign_arr = np.full(tuple(len(c) for c in self._coords) or (1,), self._sign) - rhs_arr = self._rhs.reshape(tuple(len(c) for c in self._coords) or (1,)) + ds = self._to_dataset(self.nterm) shape = tuple(len(c) for c in self._coords) dim_names = [c.name for c in self._coords] xr_coords = {c.name: c for c in self._coords} + full_size = int(np.prod(shape)) if shape else 1 + active_pos = self.active_positions + rhs_full = np.full(full_size, np.nan) + rhs_full[active_pos] = self._rhs ds = ds.assign( sign=DataArray( - sign_arr if shape else np.array(self._sign), + np.full(shape, self._sign), coords=xr_coords, dims=dim_names, ), rhs=DataArray( - rhs_arr if shape else self._rhs, coords=xr_coords, dims=dim_names + rhs_full.reshape(shape) if shape else rhs_full.reshape(()), + coords=xr_coords, + dims=dim_names, ), ) if self._dual is not None: + dual_full = np.full(full_size, np.nan) + dual_full[active_pos] = self._dual ds = ds.assign( dual=DataArray( - self._dual.reshape(shape) if shape else self._dual, + dual_full.reshape(shape) if shape else dual_full.reshape(()), coords=xr_coords, dims=dim_names, ) ) - return ds.assign_attrs(name=self._name) + attrs: dict[str, Any] = {"name": self._name} + if self._cindex is not None: + n_flat = int(np.prod(shape)) if shape else 1 + attrs["label_range"] = (self._cindex, self._cindex + n_flat) + return ds.assign_attrs(attrs) def __repr__(self) -> str: """Print the constraint without reconstructing the full Dataset.""" @@ -705,30 +781,28 @@ def __repr__(self) -> str: dim_names = [c.name for c in coords] dim_sizes = list(shape) size = int(np.prod(shape)) if shape else 1 - nonempty = self._nonempty # shape (size,) - nterm = int(self._csr.indptr.max()) if self._csr.nnz > 0 else 1 - - # Dense arrays for coeffs/vars, built without going through data property + # Map active rows (CSR is active-only) back to full flat positions csr = self._csr - counts = np.diff(csr.indptr) - coeffs_2d = np.zeros((size, nterm), dtype=csr.dtype) - vars_2d = np.full((size, nterm), -1, dtype=np.int64) - if csr.nnz > 0: - row_idx = np.repeat(np.arange(size, dtype=np.int32), counts) - term_cols = np.arange(csr.nnz, dtype=np.int32) - np.repeat( - csr.indptr[:-1].astype(np.int32), counts - ) - vars_2d[row_idx, term_cols] = csr.indices - coeffs_2d[row_idx, term_cols] = csr.data + nterm = self.nterm + active_positions = self.active_positions + masked_entries = size - len(active_positions) - coeffs_nd = coeffs_2d.reshape(shape + (nterm,)) if shape else coeffs_2d - vars_nd = vars_2d.reshape(shape + (nterm,)) if shape else vars_2d - rhs_nd = self._rhs.reshape(shape) if shape else self._rhs - masked_entries = int((~nonempty).sum()) + # pos_to_row[flat_idx] -> CSR row index, or -1 if masked + # active_positions is sorted (labels assigned in order) + pos_to_row = np.full(size, -1, dtype=np.intp) + pos_to_row[active_positions] = np.arange(len(active_positions), dtype=np.intp) header_string = f"{self.type} `{self._name}`" if self._name else f"{self.type}" lines = [] + def row_expr(row: int) -> str: + start, end = int(csr.indptr[row]), int(csr.indptr[row + 1]) + vars_row = np.full(nterm, -1, dtype=np.int64) + coeffs_row = np.zeros(nterm, dtype=csr.dtype) + vars_row[: end - start] = csr.indices[start:end] + coeffs_row[: end - start] = csr.data[start:end] + return f"{print_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[self._sign]} {self._rhs[row]}" + if size > 1 or len(dim_sizes) > 0: for indices in generate_indices_for_printout(dim_sizes, max_lines): if indices is None: @@ -736,16 +810,9 @@ def __repr__(self) -> str: else: coord = [coords[i][int(ind)] for i, ind in enumerate(indices)] flat_idx = int(np.ravel_multi_index(indices, shape)) if shape else 0 - if nonempty[flat_idx]: - expr = print_single_expression( - coeffs_nd[indices], vars_nd[indices], 0, self._model - ) - sign = SIGNS_pretty[self._sign] - rhs = rhs_nd[indices] - line = print_coord(coord) + f": {expr} {sign} {rhs}" - else: - line = print_coord(coord) + ": None" - lines.append(line) + row = pos_to_row[flat_idx] + body = row_expr(row) if row >= 0 else "None" + lines.append(print_coord(coord) + f": {body}") lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes)) @@ -753,36 +820,30 @@ def __repr__(self) -> str: underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") elif size == 1: - expr = print_single_expression( - coeffs_nd.ravel(), vars_nd.ravel(), 0, self._model - ) - lines.append( - f"{header_string}\n{'-' * len(header_string)}\n{expr} {SIGNS_pretty[self._sign]} {self._rhs.item()}" - ) + body = row_expr(0) if len(active_positions) > 0 else "None" + lines.append(f"{header_string}\n{'-' * len(header_string)}\n{body}") else: lines.append(f"{header_string}\n{'-' * len(header_string)}\n") return "\n".join(lines) - def to_matrix(self) -> scipy.sparse.csr_array: - """Return the stored CSR matrix directly (no reconstruction needed).""" - return self._csr + def to_matrix( + self, label_index: VariableLabelIndex | None = None + ) -> tuple[scipy.sparse.csr_array, np.ndarray]: + """Return the stored CSR matrix and con_labels.""" + return self._csr, self._con_labels def to_netcdf_ds(self) -> Dataset: """Return a Dataset with raw CSR components for netcdf serialization.""" - from xarray import DataArray - csr = self._csr data_vars: dict[str, DataArray] = { "indptr": DataArray(csr.indptr, dims=["_indptr"]), "indices": DataArray(csr.indices, dims=["_nnz"]), "data": DataArray(csr.data, dims=["_nnz"]), "rhs": DataArray(self._rhs, dims=["_flat"]), + "_con_labels": DataArray(self._con_labels, dims=["_flat"]), } - for c in self._coords: - data_vars[f"_coord_{c.name}"] = DataArray( - np.array(c), dims=[f"_coorddim_{c.name}"] - ) + data_vars.update(coords_to_dataset_vars(self._coords)) if self._dual is not None: data_vars["dual"] = DataArray(self._dual, dims=["_flat"]) dim_names = [c.name for c in self._coords] @@ -814,9 +875,56 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint: coord_dims = attrs["coord_dims"] if isinstance(coord_dims, str): coord_dims = [coord_dims] - coords = [pd.Index(ds[f"_coord_{d}"].values, name=d) for d in coord_dims] + coords = coords_from_dataset(ds, coord_dims) dual = ds["dual"].values if "dual" in ds else None - return cls(csr, rhs, sign, coords, model, name, cindex=cindex, dual=dual) + if "_con_labels" in ds: + con_labels = ds["_con_labels"].values + elif cindex is not None: + con_labels = np.arange(cindex, cindex + len(rhs), dtype=np.intp) + else: + con_labels = np.arange(len(rhs), dtype=np.intp) + return cls( + csr, con_labels, rhs, sign, coords, model, name, cindex=cindex, dual=dual + ) + + def has_variable(self, variable: variables.Variable) -> bool: + vlabels = self._model.variables.label_index.vlabels + return bool( + np.isin(vlabels[self._csr.indices], variable.labels.values.ravel()).any() + ) + + def to_matrix_with_rhs( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: + """Return (csr, con_labels, b, sense) — all pre-stored, no recomputation.""" + sense = np.full(len(self._rhs), self._sign[0]) + return self._csr, self._con_labels, self._rhs, sense + + def sanitize_zeros(self) -> Constraint: + """Remove terms with zero or near-zero coefficients (mutates in-place).""" + self._csr.data[np.abs(self._csr.data) <= 1e-10] = 0 + self._csr.eliminate_zeros() + return self + + def sanitize_missings(self) -> Constraint: + """No-op: missing rows are already excluded during freezing.""" + return self + + def sanitize_infinities(self) -> Constraint: + """Mask out rows with invalid infinite RHS values (mutates in-place).""" + if self._sign == LESS_EQUAL: + invalid = self._rhs == np.inf + elif self._sign == GREATER_EQUAL: + invalid = self._rhs == -np.inf + else: + return self + if not invalid.any(): + return self + keep = ~invalid + self._csr = self._csr[keep] + self._con_labels = self._con_labels[keep] + self._rhs = self._rhs[keep] + return self def freeze(self) -> Constraint: """Return self (already immutable).""" @@ -826,6 +934,10 @@ def mutable(self) -> MutableConstraint: """Convert to a MutableConstraint.""" return MutableConstraint(self.data, self._model, self._name) + def to_polars(self) -> Any: + """Convert to polars DataFrame — delegates to mutable().""" + return self.mutable().to_polars() + @classmethod def from_mutable( cls, @@ -841,20 +953,36 @@ def from_mutable( cindex : int or None Starting label index, if assigned. """ - csr = con.to_matrix() + label_index = con.model.variables.label_index + csr, con_labels = con.to_matrix(label_index) coords = [con.indexes[d] for d in con.coord_dims] - rhs = con.rhs.values.ravel() + # Build active_mask aligned with con_labels (rows in csr) + # Use same filter as to_matrix: label != -1 AND at least one var != -1 + labels_flat = con.labels.values.ravel() + vars_flat = con.vars.values.reshape(len(labels_flat), -1) + active_mask = (labels_flat != -1) & (vars_flat != -1).any(axis=1) + rhs = con.rhs.values.ravel()[active_mask] sign_vals = con.sign.values.ravel() - unique_signs = np.unique(sign_vals) + unique_signs = np.unique(sign_vals[active_mask]) if len(unique_signs) > 1: raise ValueError( "Constraint has per-element signs; cannot freeze to immutable Constraint. " "This is a known limitation — use MutableConstraint instead." ) sign = str(unique_signs[0]) if len(unique_signs) == 1 else "=" - dual = con.data["dual"].values.ravel() if "dual" in con.data else None + dual = ( + con.data["dual"].values.ravel()[active_mask] if "dual" in con.data else None + ) return cls( - csr, rhs, sign, coords, con.model, con.name, cindex=cindex, dual=dual + csr, + con_labels, + rhs, + sign, + coords, + con.model, + con.name, + cindex=cindex, + dual=dual, ) @@ -997,6 +1125,44 @@ def dual(self, value: ConstantLike) -> None: value = DataArray(value).broadcast_like(self.labels) self._data = assign_multiindex_safe(self.data, dual=value) + def has_variable(self, variable: variables.Variable) -> bool: + return bool(self.data["vars"].isin(variable.labels.values.ravel()).any()) + + def to_matrix_with_rhs( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: + """Return (csr, con_labels, b, sense) in one pass.""" + csr, con_labels = self.to_matrix(label_index) + nonempty = np.diff(csr.indptr).astype(bool) + active_rows = np.flatnonzero(nonempty) + b = self.rhs.values.ravel()[active_rows] + sign_flat = self.sign.values.ravel()[active_rows] + sense = np.array([s[0] for s in sign_flat]) + return csr, con_labels, b, sense + + def sanitize_zeros(self) -> MutableConstraint: + """Remove terms with zero or near-zero coefficients.""" + not_zero = abs(self.coeffs) > 1e-10 + self.vars = self.vars.where(not_zero, -1) + self.coeffs = self.coeffs.where(not_zero) + return self + + def sanitize_missings(self) -> MutableConstraint: + """Mask out rows where all variables are missing (-1).""" + contains_non_missing = (self.vars != -1).any(self.term_dim) + labels = self.labels.where(contains_non_missing, -1) + self._data = assign_multiindex_safe(self.data, labels=labels) + return self + + def sanitize_infinities(self) -> MutableConstraint: + """Mask out rows with invalid infinite RHS values.""" + valid_infinity_values = ((self.sign == LESS_EQUAL) & (self.rhs == np.inf)) | ( + (self.sign == GREATER_EQUAL) & (self.rhs == -np.inf) + ) + labels = self.labels.where(~valid_infinity_values, -1) + self._data = assign_multiindex_safe(self.data, labels=labels) + return self + def freeze(self) -> Constraint: """Convert to an immutable Constraint.""" cindex = ( @@ -1368,34 +1534,27 @@ def sanitize_zeros(self) -> None: """ Filter out terms with zero and close-to-zero coefficient. """ - for name in self: - not_zero = abs(self[name].coeffs) > 1e-10 - con = self[name] - con.vars = self[name].vars.where(not_zero, -1) - con.coeffs = self[name].coeffs.where(not_zero) + for con in self.data.values(): + con.sanitize_zeros() def sanitize_missings(self) -> None: """ Set constraints labels to -1 where all variables in the lhs are missing. """ - for name in self: - con = self[name] - contains_non_missing = (con.vars != -1).any(con.term_dim) - labels = self[name].labels.where(contains_non_missing, -1) - con._data = assign_multiindex_safe(con.data, labels=labels) + for con in self.data.values(): + con.sanitize_missings() def sanitize_infinities(self) -> None: """ - Replace infinite values in the constraints with a large value. + Remove constraints whose RHS is an invalid infinity. + + Constraints with ``rhs == inf`` and sign ``<=``, or ``rhs == -inf`` + and sign ``>=``, are trivially satisfied and are masked out (label set + to -1). """ - for name in self: - con = self[name] - valid_infinity_values = ((con.sign == LESS_EQUAL) & (con.rhs == np.inf)) | ( - (con.sign == GREATER_EQUAL) & (con.rhs == -np.inf) - ) - labels = con.labels.where(~valid_infinity_values, -1) - con._data = assign_multiindex_safe(con.data, labels=labels) + for con in self.data.values(): + con.sanitize_infinities() def get_name_by_label(self, label: int | float) -> str: """ @@ -1516,73 +1675,55 @@ def flat(self) -> pd.DataFrame: df["key"] = df.labels.map(map_labels) return df - def to_matrix( - self, filter_missings: bool = True - ) -> tuple[scipy.sparse.csc_array, np.ndarray, np.ndarray | None]: + def to_matrix(self) -> tuple[scipy.sparse.csr_array, np.ndarray]: """ Construct a constraint matrix in sparse format by stacking per-constraint CSR matrices. - Parameters - ---------- - filter_missings : bool, default True - If True, also strip empty columns and return ``var_labels`` for - remapping columns back to original variable labels. - If False, return full-width CSC with shape - ``(n_active_cons, model._xCounter)`` and ``var_labels=None``. - ``con_labels`` is always returned. + Each per-constraint CSR is already dense: rows are active constraints + only, column indices are dense variable positions (not raw labels). + Shape is ``(n_active_cons, n_active_vars)``. Returns ------- - matrix : scipy.sparse.csc_array - Shape ``(n_active_cons, n_active_vars)`` when - ``filter_missings=True``, or ``(n_active_cons, - model._xCounter)`` when ``filter_missings=False``. + matrix : scipy.sparse.csr_array + Shape ``(n_active_cons, n_active_vars)``. con_labels : np.ndarray Shape ``(n_active_cons,)``, maps each matrix row to the original constraint label. - var_labels : np.ndarray or None - Shape ``(n_active_vars,)``, maps each matrix column to the - original variable label. ``None`` when - ``filter_missings=False``. """ if not len(self): raise ValueError("No constraints available to convert to matrix.") - active_csrs = [] + label_index = self.model.variables.label_index + csrs = [] con_labels_list = [] for c in self.data.values(): - csr = c.to_matrix() - nonempty = np.diff(csr.indptr).astype(bool) - active_csrs.append(csr[nonempty]) - start = ( - c._cindex - if isinstance(c, Constraint) - else c.data.attrs["label_range"][0] - ) - con_labels_list.append(np.flatnonzero(nonempty) + start) - csc: scipy.sparse.csc_array = scipy.sparse.vstack(active_csrs).tocsc() - con_labels = np.concatenate(con_labels_list) - - if filter_missings: - indptr = csc.indptr - nonempty_cols = indptr[1:] != indptr[:-1] - new_indptr = np.r_[0, indptr[1:][nonempty_cols]] - (var_labels,) = np.nonzero(nonempty_cols) - matrix = scipy.sparse.csc_array( - (csc.data, csc.indices, new_indptr), - shape=(csc.shape[0], len(var_labels)), - ) - return matrix, con_labels, var_labels - else: - return csc, con_labels, None + csr, con_labels = c.to_matrix(label_index) + csrs.append(csr) + con_labels_list.append(con_labels) + return scipy.sparse.vstack(csrs, format="csr"), np.concatenate(con_labels_list) def reset_dual(self) -> None: """ Reset the stored solution of variables. """ for k, c in self.items(): - if "dual" in c: - c._data = c.data.drop_vars("dual") + if isinstance(c, Constraint): + if c._dual is not None: + self.data[k] = Constraint( + c._csr, + c._con_labels, + c._rhs, + c._sign, + c._coords, + c._model, + c._name, + cindex=c._cindex, + dual=None, + ) + else: + if "dual" in c.data: + c._data = c.data.drop_vars("dual") class AnonymousScalarConstraint: diff --git a/linopy/matrices.py b/linopy/matrices.py index b81d50db..93f3e981 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -10,8 +10,8 @@ from typing import TYPE_CHECKING import numpy as np +import scipy.sparse from numpy import ndarray -from scipy.sparse._csc import csc_matrix from linopy import expressions from linopy.constraints import Constraint @@ -36,8 +36,9 @@ def __init__(self, model: Model) -> None: def _build_vars(self) -> None: m = self._parent + label_index = m.variables.label_index + self.vlabels: ndarray = label_index.vlabels - labels_list = [] lb_list = [] ub_list = [] vtypes_list = [] @@ -45,7 +46,6 @@ def _build_vars(self) -> None: for name, var in m.variables.items(): labels = var.labels.values.ravel() mask = labels != -1 - active_labels = labels[mask] if name in m.binaries: vtype = "B" @@ -56,20 +56,15 @@ def _build_vars(self) -> None: else: vtype = "C" - labels_list.append(active_labels) lb_list.append(var.lower.values.ravel()[mask]) ub_list.append(var.upper.values.ravel()[mask]) vtypes_list.append(np.full(mask.sum(), vtype)) - if labels_list: - vlabels = np.concatenate(labels_list) - order = np.argsort(vlabels) - self.vlabels: ndarray = vlabels[order] - self.lb: ndarray = np.concatenate(lb_list)[order] - self.ub: ndarray = np.concatenate(ub_list)[order] - self.vtypes: ndarray = np.concatenate(vtypes_list)[order] + if lb_list: + self.lb: ndarray = np.concatenate(lb_list) + self.ub: ndarray = np.concatenate(ub_list) + self.vtypes: ndarray = np.concatenate(vtypes_list) else: - self.vlabels = np.array([], dtype=np.intp) self.lb = np.array([]) self.ub = np.array([]) self.vtypes = np.array([], dtype=object) @@ -81,28 +76,23 @@ def _build_cons(self) -> None: self.clabels: ndarray = np.array([], dtype=np.intp) self.b: ndarray = np.array([]) self.sense: ndarray = np.array([], dtype=object) - self.A: csc_matrix | None = None + self.A: scipy.sparse.csr_array | None = None return - A_full, clabels, _ = m.constraints.to_matrix(filter_missings=False) - self.A = A_full[:, self.vlabels] - self.clabels = clabels - + label_index = m.variables.label_index + csrs = [] + clabels_list = [] b_list = [] sense_list = [] for c in m.constraints.data.values(): - csr = c.to_matrix() - nonempty = np.diff(csr.indptr).astype(bool) - active_rows = np.flatnonzero(nonempty) - - if isinstance(c, Constraint): - b_list.append(c._rhs[active_rows]) - sense_list.append(np.full(len(active_rows), c._sign[0])) - else: - b_list.append(c.rhs.values.ravel()[active_rows]) - sign_flat = c.sign.values.ravel()[active_rows] - sense_list.append(np.array([s[0] for s in sign_flat])) - + csr, con_labels, b, sense = c.to_matrix_with_rhs(label_index) + csrs.append(csr) + clabels_list.append(con_labels) + b_list.append(b) + sense_list.append(sense) + + self.A = scipy.sparse.vstack(csrs, format="csr") + self.clabels = np.concatenate(clabels_list) self.b = np.concatenate(b_list) if b_list else np.array([]) self.sense = ( np.concatenate(sense_list) if sense_list else np.array([], dtype=object) @@ -114,24 +104,27 @@ def c(self) -> ndarray: m = self._parent result = np.zeros(len(self.vlabels)) - ds = m.objective.flat - if isinstance(m.objective.expression, expressions.QuadraticExpression): - ds = ds[(ds.vars1 == -1) | (ds.vars2 == -1)].copy() - ds["vars"] = ds.vars1.where(ds.vars1 != -1, ds.vars2) + 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() - var_labels = ds.vars.values - coeffs = ds.coeffs.values mask = var_labels != -1 - active_labels = var_labels[mask] - positions = np.searchsorted(self.vlabels, active_labels) - valid = (positions < len(self.vlabels)) & ( - self.vlabels[positions] == active_labels - ) - np.add.at(result, positions[valid], coeffs[mask][valid]) + np.add.at(result, label_to_pos[var_labels[mask]], coeffs[mask]) return result @property - def Q(self) -> csc_matrix | None: + def Q(self) -> scipy.sparse.csc_matrix | None: """Quadratic objective matrix, shape (n_active_vars, n_active_vars).""" m = self._parent expr = m.objective.expression @@ -146,11 +139,12 @@ def sol(self) -> ndarray: raise ValueError("Model is not optimized.") m = self._parent result = np.full(len(self.vlabels), np.nan) + label_index = m.variables.label_index + label_to_pos = label_index.label_to_pos for _, var in m.variables.items(): labels = var.labels.values.ravel() mask = labels != -1 - active_labels = labels[mask] - positions = np.searchsorted(self.vlabels, active_labels) + positions = label_to_pos[labels[mask]] result[positions] = var.solution.values.ravel()[mask] return result @@ -160,19 +154,21 @@ def dual(self) -> ndarray: if not self._parent.status == "ok": raise ValueError("Model is not optimized.") m = self._parent + label_index = m.variables.label_index dual_list = [] has_dual = False for c in m.constraints.data.values(): - csr = c.to_matrix() - nonempty = np.diff(csr.indptr).astype(bool) - active_rows = np.flatnonzero(nonempty) if isinstance(c, Constraint): + # _dual is active-only if c._dual is not None: - dual_list.append(c._dual[active_rows]) + dual_list.append(c._dual) has_dual = True else: - dual_list.append(np.full(len(active_rows), np.nan)) + dual_list.append(np.full(len(c._con_labels), np.nan)) else: + csr, _ = c.to_matrix(label_index) + nonempty = np.diff(csr.indptr).astype(bool) + active_rows = np.flatnonzero(nonempty) if "dual" in c.data: dual_list.append(c.dual.values.ravel()[active_rows]) has_dual = True diff --git a/linopy/model.py b/linopy/model.py index 20d47e40..2c0c1467 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -9,6 +9,7 @@ import logging import os import re +import warnings from collections.abc import Callable, Mapping, Sequence from pathlib import Path from tempfile import NamedTemporaryFile, gettempdir @@ -731,7 +732,7 @@ def add_constraints( name: str | None = None, coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = None, mask: MaskLike | None = None, - freeze: bool = False, + freeze: bool = True, ) -> ConstraintBase: """ Assign a new, possibly multi-dimensional array of constraints to the @@ -768,7 +769,7 @@ def add_constraints( Default is None. freeze : bool, optional If True, convert the constraint to an immutable CSR-backed Constraint - before returning. Default is False. + before returning. Default is True. Returns ------- @@ -889,6 +890,8 @@ def add_constraints( data = data.assign_attrs(label_range=(start, end), name=name) if self.chunk: + if freeze: + raise ValueError("Chunked constraints cannot be frozen") data = data.chunk(self.chunk) constraint = MutableConstraint(data, name=name, model=self, skip_broadcast=True) @@ -944,18 +947,26 @@ def remove_variables(self, name: str) -> None: ------- None. """ - labels = self.variables[name].labels - self.variables.remove(name) + variable = self.variables[name] + + to_remove = [ + k for k, con in self.constraints.items() if con.has_variable(variable) + ] - for k in list(self.constraints): - vars = self.constraints[k].data["vars"] - vars = vars.where(~vars.isin(labels), -1) - self.constraints[k]._data = assign_multiindex_safe( - self.constraints[k].data, vars=vars + if to_remove: + warnings.warn( + f"Removing variable '{name}' also removes constraints {to_remove} " + "because they reference this variable.", + UserWarning, + stacklevel=2, ) + for k in to_remove: + self.constraints.remove(k) + + self.variables.remove(name) self.objective = self.objective.sel( - {TERM_DIM: ~self.objective.vars.isin(labels)} + {TERM_DIM: ~self.objective.vars.isin(variable.labels)} ) def remove_constraints(self, name: str | list[str]) -> None: diff --git a/test/test_constraint.py b/test/test_constraint.py index 098b91b8..9fbb042e 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -58,6 +58,11 @@ def c(m: Model) -> linopy.constraints.Constraint: return m.constraints["c"] +@pytest.fixture +def mc(m: Model) -> linopy.constraints.MutableConstraint: + return m.constraints["c"].mutable() + + def test_constraint_repr(c: linopy.constraints.Constraint) -> None: c.__repr__() @@ -316,13 +321,13 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint | None: def test_constraint_vars_getter( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable ) -> None: - assert_equal(c.vars.squeeze(), x.labels) + assert_equal(mc.vars.squeeze(), x.labels) -def test_constraint_coeffs_getter(c: linopy.constraints.Constraint) -> None: - assert (c.coeffs == 1).all() +def test_constraint_coeffs_getter(mc: linopy.constraints.MutableConstraint) -> None: + assert (mc.coeffs == 1).all() def test_constraint_sign_getter(c: linopy.constraints.Constraint) -> None: @@ -334,124 +339,130 @@ def test_constraint_rhs_getter(c: linopy.constraints.Constraint) -> None: def test_constraint_vars_setter( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable ) -> None: - c.vars = x - assert_equal(c.vars, x.labels) + mc.vars = x + assert_equal(mc.vars, x.labels) def test_constraint_vars_setter_with_array( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable ) -> None: - c.vars = x.labels - assert_equal(c.vars, x.labels) + mc.vars = x.labels + assert_equal(mc.vars, x.labels) def test_constraint_vars_setter_invalid( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable ) -> None: with pytest.raises(TypeError): - c.vars = pd.DataFrame(x.labels) + mc.vars = pd.DataFrame(x.labels) -def test_constraint_coeffs_setter(c: linopy.constraints.Constraint) -> None: - c.coeffs = 3 - assert (c.coeffs == 3).all() +def test_constraint_coeffs_setter(mc: linopy.constraints.MutableConstraint) -> None: + mc.coeffs = 3 + assert (mc.coeffs == 3).all() def test_constraint_lhs_setter( - c: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable, y: linopy.Variable ) -> None: - c.lhs = x + y - assert c.lhs.nterm == 2 - assert c.vars.notnull().all().item() - assert c.coeffs.notnull().all().item() + mc.lhs = x + y + assert mc.lhs.nterm == 2 + assert mc.vars.notnull().all().item() + assert mc.coeffs.notnull().all().item() def test_constraint_lhs_setter_with_variable( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable ) -> None: - c.lhs = x - assert c.lhs.nterm == 1 + mc.lhs = x + assert mc.lhs.nterm == 1 -def test_constraint_lhs_setter_with_constant(c: linopy.constraints.Constraint) -> None: - sizes = c.sizes - c.lhs = 10 - assert (c.rhs == -10).all() - assert c.lhs.nterm == 0 - assert c.sizes["first"] == sizes["first"] +def test_constraint_lhs_setter_with_constant( + mc: linopy.constraints.MutableConstraint, +) -> None: + sizes = mc.sizes + mc.lhs = 10 + assert (mc.rhs == -10).all() + assert mc.lhs.nterm == 0 + assert mc.sizes["first"] == sizes["first"] -def test_constraint_sign_setter(c: linopy.constraints.Constraint) -> None: - c.sign = EQUAL - assert (c.sign == EQUAL).all() +def test_constraint_sign_setter(mc: linopy.constraints.MutableConstraint) -> None: + mc.sign = EQUAL + assert (mc.sign == EQUAL).all() -def test_constraint_sign_setter_alternative(c: linopy.constraints.Constraint) -> None: - c.sign = long_EQUAL - assert (c.sign == EQUAL).all() +def test_constraint_sign_setter_alternative( + mc: linopy.constraints.MutableConstraint, +) -> None: + mc.sign = long_EQUAL + assert (mc.sign == EQUAL).all() -def test_constraint_sign_setter_invalid(c: linopy.constraints.Constraint) -> None: +def test_constraint_sign_setter_invalid( + mc: linopy.constraints.MutableConstraint, +) -> None: # Test that assigning lhs with other type that LinearExpression raises TypeError with pytest.raises(ValueError): - c.sign = "asd" + mc.sign = "asd" -def test_constraint_rhs_setter(c: linopy.constraints.Constraint) -> None: - sizes = c.sizes - c.rhs = 2 # type: ignore - assert (c.rhs == 2).all() - assert c.sizes == sizes +def test_constraint_rhs_setter(mc: linopy.constraints.MutableConstraint) -> None: + sizes = mc.sizes + mc.rhs = 2 # type: ignore + assert (mc.rhs == 2).all() + assert mc.sizes == sizes def test_constraint_rhs_setter_with_variable( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable ) -> None: - c.rhs = x # type: ignore - assert (c.rhs == 0).all() - assert (c.coeffs.isel({c.term_dim: -1}) == -1).all() - assert c.lhs.nterm == 2 + mc.rhs = x # type: ignore + assert (mc.rhs == 0).all() + assert (mc.coeffs.isel({mc.term_dim: -1}) == -1).all() + assert mc.lhs.nterm == 2 def test_constraint_rhs_setter_with_expression( - c: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable, y: linopy.Variable ) -> None: - c.rhs = x + y - assert (c.rhs == 0).all() - assert (c.coeffs.isel({c.term_dim: -1}) == -1).all() - assert c.lhs.nterm == 3 + mc.rhs = x + y + assert (mc.rhs == 0).all() + assert (mc.coeffs.isel({mc.term_dim: -1}) == -1).all() + assert mc.lhs.nterm == 3 def test_constraint_rhs_setter_with_expression_and_constant( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.MutableConstraint, x: linopy.Variable ) -> None: - c.rhs = x + 1 - assert (c.rhs == 1).all() - assert (c.coeffs.sum(c.term_dim) == 0).all() - assert c.lhs.nterm == 2 + mc.rhs = x + 1 + assert (mc.rhs == 1).all() + assert (mc.coeffs.sum(mc.term_dim) == 0).all() + assert mc.lhs.nterm == 2 def test_constraint_labels_setter_invalid(c: linopy.constraints.Constraint) -> None: - # Test that assigning labels raises FrozenInstanceError + # Test that assigning labels raises AttributeError (Constraint is frozen) with pytest.raises(AttributeError): c.labels = c.labels # type: ignore def test_constraint_sel(c: linopy.constraints.Constraint) -> None: - assert isinstance(c.sel(first=[1, 2]), ConstraintBase) - assert isinstance(c.isel(first=[1, 2]), ConstraintBase) + assert isinstance(c.mutable().sel(first=[1, 2]), ConstraintBase) + assert isinstance(c.mutable().isel(first=[1, 2]), ConstraintBase) def test_constraint_flat(c: linopy.constraints.Constraint) -> None: assert isinstance(c.flat, pd.DataFrame) -def test_iterate_slices(c: linopy.constraints.Constraint) -> None: - for i in c.iterate_slices(slice_size=2): +def test_iterate_slices(mc: linopy.constraints.MutableConstraint) -> None: + for i in mc.iterate_slices(slice_size=2): assert isinstance(i, ConstraintBase) - assert c.coord_dims == i.coord_dims + assert mc.coord_dims == i.coord_dims def test_constraint_to_polars(c: linopy.constraints.Constraint) -> None: @@ -460,9 +471,8 @@ def test_constraint_to_polars(c: linopy.constraints.Constraint) -> None: def test_constraint_to_polars_mixed_signs(m: Model, x: linopy.Variable) -> None: """Test to_polars when a constraint has mixed sign values across dims.""" - # Create a constraint, then manually patch the sign to have mixed values - m.add_constraints(x >= 0, name="mixed") - con = m.constraints["mixed"] + # Use MutableConstraint so sign data can be patched + con = m.add_constraints(x >= 0, name="mixed", freeze=False) # Replace sign data with mixed signs across the first dimension n = con.data.sizes["first"] signs = np.array(["<=" if i % 2 == 0 else ">=" for i in range(n)]) @@ -485,8 +495,12 @@ def test_constraint_assignment_sanitize_zeros( ) -> None: m.add_constraints(0 * x + y == 0, name="c2") m.constraints.sanitize_zeros() - assert m.constraints["c2"].vars[0, 0, 0].item() == -1 - assert np.isnan(m.constraints["c2"].coeffs[0, 0, 0].item()) + c2 = m.constraints["c2"] + assert c2.nterm == 1 + assert c2.has_variable(y) + assert not c2.has_variable(x) + csr, _ = c2.to_matrix() + assert (csr.data == 1).all() def test_constraint_assignment_with_args( @@ -609,17 +623,11 @@ def test_constraint_with_helper_dims_as_coords(m: Model) -> None: def test_constraint_matrix(m: Model) -> None: - # filter_missings=True: strips empty columns, returns labels for remapping - A, con_labels, var_labels = m.constraints.to_matrix() - assert A.shape == (10, 10) # only x (10 vars) appears in constraints - assert con_labels is not None - assert var_labels is not None - - # filter_missings=False: full width, con_labels always returned, var_labels=None - A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=False) - assert A.shape == (10, m.shape[1]) # all 14 vars as columns + # Returns (csr_array, con_labels) — dense: active rows and active-var columns + A, con_labels = m.constraints.to_matrix() + n_active_vars = len(m.variables.label_index.vlabels) + assert A.shape == (10, n_active_vars) assert len(con_labels) == 10 - assert var_labels is None def test_constraint_matrix_masked_variables() -> None: @@ -630,23 +638,16 @@ def test_constraint_matrix_masked_variables() -> None: missing. The matrix shoud not be built for constraints which have variables which are missing. """ - # now with missing variables m = Model() mask = pd.Series([False] * 5 + [True] * 5) x = m.add_variables(coords=[range(10)], mask=mask) m.add_variables() m.add_constraints(x, EQUAL, 0) - # filter_missings=True: only vars/cons that appear in constraints - A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (m.ncons, 5) # 5 active x vars; scalar var not in any constraint + # Returns dense matrix: active rows only, all active-var columns + A, con_labels = m.constraints.to_matrix() + n_active_vars = len(m.variables.label_index.vlabels) + assert A.shape == (m.ncons, n_active_vars) assert len(con_labels) == m.ncons - assert len(var_labels) == 5 - - # filter_missings=False: full width, no remapping labels - A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=False) - assert A.shape == (m.ncons, m.shape[1]) - assert con_labels is not None - assert var_labels is None def test_constraint_matrix_masked_constraints() -> None: @@ -658,17 +659,12 @@ def test_constraint_matrix_masked_constraints() -> None: x = m.add_variables(coords=[range(10)]) m.add_variables() m.add_constraints(x, EQUAL, 0, mask=mask) - # filter_missings=True: only active cons and vars they reference - # active cons are indices 5-9, which reference vars 5-9 only - A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (m.ncons, m.ncons) # 5 active cons; 5 referenced vars + # active cons are indices 5-9, which reference vars 5-9 only (all active) + A, con_labels = m.constraints.to_matrix() + n_active_vars = len(m.variables.label_index.vlabels) + assert A.shape == (m.ncons, n_active_vars) assert len(con_labels) == m.ncons - A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=False) - assert A.shape == (m.ncons, m.shape[1]) - assert con_labels is not None - assert var_labels is None - def test_constraint_matrix_masked_constraints_and_variables() -> None: """ @@ -679,16 +675,11 @@ def test_constraint_matrix_masked_constraints_and_variables() -> None: x = m.add_variables(coords=[range(10)], mask=mask) m.add_variables() m.add_constraints(x, EQUAL, 0, mask=mask) - # filter_missings=True: 5 active cons x 5 active vars - A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (m.ncons, m.ncons) # both masks align: 5x5 + # both masks align: 5 active cons x all active vars (5 x + 1 scalar) + A, con_labels = m.constraints.to_matrix() + n_active_vars = len(m.variables.label_index.vlabels) + assert A.shape == (m.ncons, n_active_vars) assert len(con_labels) == m.ncons - assert len(var_labels) == m.ncons - - A, con_labels, var_labels = m.constraints.to_matrix(filter_missings=False) - assert A.shape == (m.ncons, m.shape[1]) - assert con_labels is not None - assert var_labels is None def test_get_name_by_label() -> None: diff --git a/test/test_constraints.py b/test/test_constraints.py index 9a467c8c..1d88d7e4 100644 --- a/test/test_constraints.py +++ b/test/test_constraints.py @@ -116,7 +116,9 @@ def test_constraint_assignment_chunked() -> None: lower = pd.DataFrame(np.zeros((10, 10))) upper = pd.Series(np.ones(10)) x = m.add_variables(lower, upper) - m.add_constraints(x, GREATER_EQUAL, 0, name="c") + with pytest.raises(ValueError, match="Chunked constraints cannot be frozen"): + m.add_constraints(x, GREATER_EQUAL, 0, name="c") + m.add_constraints(x, GREATER_EQUAL, 0, name="c", freeze=False) assert m.constraints.coeffs.c.data.shape == ( 10, 10, diff --git a/test/test_model.py b/test/test_model.py index c363fe4c..53e62b26 100644 --- a/test/test_model.py +++ b/test/test_model.py @@ -104,10 +104,11 @@ def test_remove_variable() -> None: assert "x" in m.variables - m.remove_variables("x") + with pytest.warns(UserWarning, match="con0"): + m.remove_variables("x") assert "x" not in m.variables - assert not m.constraints.con0.vars.isin(x.labels).any() + assert "con0" not in m.constraints assert not m.objective.vars.isin(x.labels).any() diff --git a/test/test_optimization.py b/test/test_optimization.py index cdac8e61..870cc154 100644 --- a/test/test_optimization.py +++ b/test/test_optimization.py @@ -298,7 +298,7 @@ def modified_model() -> Model: x = m.add_variables(coords=[lower.index], name="x", binary=True) y = m.add_variables(lower, name="y") - c = m.add_constraints(x + y, GREATER_EQUAL, 10) + c = m.add_constraints(x + y, GREATER_EQUAL, 10, freeze=False) y.lower = 9 c.lhs = 2 * x + y From 1122b163ad2048f16b385d0535f984c6e5d7fcec Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Mon, 23 Mar 2026 09:38:42 +0100 Subject: [PATCH 10/20] fix: disentangle range and ncons --- linopy/constraints.py | 116 ++++++++++++++---------------------------- 1 file changed, 37 insertions(+), 79 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 096632d3..fae5a03a 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -540,12 +540,20 @@ def name(self) -> str: def is_assigned(self) -> bool: return self._cindex is not None + @property + def shape(self) -> tuple[int, ...]: + return tuple(len(c) for c in self._coords) + + @property + def full_size(self) -> int: + return int(np.prod(shape)) if (shape := self.shape) else 1 + @property def range(self) -> tuple[int, int]: """Return the (start, end) label range of the constraint.""" if self._cindex is None: raise AttributeError("Constraint has not been assigned labels yet.") - return (self._cindex, self._cindex + self._csr.shape[0]) + return (self._cindex, self._cindex + self.full_size) @property def ncons(self) -> int: @@ -555,7 +563,7 @@ def ncons(self) -> int: def attrs(self) -> dict[str, Any]: d: dict[str, Any] = {"name": self._name} if self._cindex is not None: - d["label_range"] = (self._cindex, self._cindex + self._csr.shape[0]) + d["label_range"] = (self._cindex, self._cindex + self.full_size) return d @property @@ -592,17 +600,11 @@ def labels(self) -> DataArray: """Get labels DataArray, shape (*coord_dims).""" if self._cindex is None: return DataArray([]) - shape = tuple(len(c) for c in self._coords) - full_size = int(np.prod(shape)) if shape else 1 + shape = self.shape + full_size = self.full_size labels_flat = np.full(full_size, -1, dtype=np.int64) labels_flat[self.active_positions] = self._con_labels - dim_names = [c.name for c in self._coords] - xr_coords = {c.name: c for c in self._coords} - return DataArray( - labels_flat.reshape(shape), - coords=xr_coords, - dims=dim_names, - ) + return DataArray(labels_flat.reshape(shape), coords=self._coords) @property def coeffs(self) -> DataArray: @@ -629,29 +631,15 @@ def vars(self) -> DataArray: @property def sign(self) -> DataArray: """Get sign DataArray (scalar, same sign for all entries).""" - shape = tuple(len(c) for c in self._coords) - dim_names = [c.name for c in self._coords] - xr_coords = {c.name: c for c in self._coords} - return DataArray( - np.full(shape, self._sign) if shape else np.array(self._sign), - coords=xr_coords, - dims=dim_names, - ) + return DataArray(np.full(self.shape, self._sign), coords=self._coords) @property def rhs(self) -> DataArray: """Get RHS DataArray, shape (*coord_dims).""" - shape = tuple(len(c) for c in self._coords) - dim_names = [c.name for c in self._coords] - xr_coords = {c.name: c for c in self._coords} - full_size = int(np.prod(shape)) if shape else 1 - rhs_full = np.full(full_size, np.nan) + shape = self.shape + rhs_full = np.full(self.full_size, np.nan) rhs_full[self.active_positions] = self._rhs - return DataArray( - rhs_full.reshape(shape) if shape else rhs_full.reshape(()), - coords=xr_coords, - dims=dim_names, - ) + return DataArray(rhs_full.reshape(shape), coords=self._coords) @property @has_optimized_model @@ -661,17 +649,9 @@ def dual(self) -> DataArray: raise AttributeError( "Underlying is optimized but does not have dual values stored." ) - shape = tuple(len(c) for c in self._coords) - dim_names = [c.name for c in self._coords] - xr_coords = {c.name: c for c in self._coords} - full_size = int(np.prod(shape)) if shape else 1 - dual_full = np.full(full_size, np.nan) + dual_full = np.full(self.full_size, np.nan) dual_full[self.active_positions] = self._dual - return DataArray( - dual_full.reshape(shape) if shape else dual_full.reshape(()), - coords=xr_coords, - dims=dim_names, - ) + return DataArray(dual_full.reshape(self.shape), coords=self._coords) @dual.setter def dual(self, value: DataArray) -> None: @@ -693,10 +673,9 @@ def _to_dataset(self, nterm: int) -> Dataset: Dataset with variables ``labels``, ``coeffs``, ``vars``. """ csr = self._csr - n_active = csr.shape[0] # active rows only (no masked rows in CSR) counts = np.diff(csr.indptr) - shape = tuple(len(c) for c in self._coords) - full_size = int(np.prod(shape)) if shape else n_active + shape = self.shape + full_size = self.full_size # Map active row i -> flat position in full shape via con_labels active_positions = self.active_positions @@ -710,7 +689,7 @@ def _to_dataset(self, nterm: int) -> Dataset: vars_2d[row_indices, term_cols] = vlabels[csr.indices] coeffs_2d[row_indices, term_cols] = csr.data - dim_names = [c.name for c in self._coords] + dim_names = self.coord_names xr_coords = {c.name: c for c in self._coords} dims_with_term = dim_names + [TERM_DIM] coeffs_da = DataArray( @@ -727,60 +706,39 @@ def _to_dataset(self, nterm: int) -> Dataset: if self._cindex is not None: labels_flat = np.full(full_size, -1, dtype=np.int64) labels_flat[active_positions] = self._con_labels - ds["labels"] = DataArray( - labels_flat.reshape(shape), - coords=xr_coords, - dims=dim_names, - ) + ds["labels"] = DataArray(labels_flat.reshape(shape), coords=self._coords) return ds @property def data(self) -> Dataset: """Reconstruct the xarray Dataset from the CSR representation.""" ds = self._to_dataset(self.nterm) - shape = tuple(len(c) for c in self._coords) - dim_names = [c.name for c in self._coords] - xr_coords = {c.name: c for c in self._coords} - full_size = int(np.prod(shape)) if shape else 1 + shape = self.shape active_pos = self.active_positions - rhs_full = np.full(full_size, np.nan) + rhs_full = np.full(self.full_size, np.nan) rhs_full[active_pos] = self._rhs ds = ds.assign( - sign=DataArray( - np.full(shape, self._sign), - coords=xr_coords, - dims=dim_names, - ), - rhs=DataArray( - rhs_full.reshape(shape) if shape else rhs_full.reshape(()), - coords=xr_coords, - dims=dim_names, - ), + sign=DataArray(np.full(shape, self._sign), coords=self._coords), + rhs=DataArray(rhs_full.reshape(shape), coords=self._coords), ) if self._dual is not None: - dual_full = np.full(full_size, np.nan) + dual_full = np.full(self.full_size, np.nan) dual_full[active_pos] = self._dual ds = ds.assign( - dual=DataArray( - dual_full.reshape(shape) if shape else dual_full.reshape(()), - coords=xr_coords, - dims=dim_names, - ) + dual=DataArray(dual_full.reshape(shape), coords=self._coords) ) attrs: dict[str, Any] = {"name": self._name} if self._cindex is not None: - n_flat = int(np.prod(shape)) if shape else 1 - attrs["label_range"] = (self._cindex, self._cindex + n_flat) + attrs["label_range"] = (self._cindex, self._cindex + self.full_size) return ds.assign_attrs(attrs) def __repr__(self) -> str: """Print the constraint without reconstructing the full Dataset.""" max_lines = options["display_max_rows"] coords = self._coords - shape = tuple(len(c) for c in coords) - dim_names = [c.name for c in coords] - dim_sizes = list(shape) - size = int(np.prod(shape)) if shape else 1 + shape = self.shape + dim_names = self.coord_names + size = self.full_size # Map active rows (CSR is active-only) back to full flat positions csr = self._csr nterm = self.nterm @@ -803,19 +761,19 @@ def row_expr(row: int) -> str: coeffs_row[: end - start] = csr.data[start:end] return f"{print_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[self._sign]} {self._rhs[row]}" - if size > 1 or len(dim_sizes) > 0: - for indices in generate_indices_for_printout(dim_sizes, max_lines): + if size > 1: + for indices in generate_indices_for_printout(shape, max_lines): if indices is None: lines.append("\t\t...") else: coord = [coords[i][int(ind)] for i, ind in enumerate(indices)] - flat_idx = int(np.ravel_multi_index(indices, shape)) if shape else 0 + flat_idx = int(np.ravel_multi_index(indices, shape)) row = pos_to_row[flat_idx] body = row_expr(row) if row >= 0 else "None" lines.append(print_coord(coord) + f": {body}") lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) - shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes)) + shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, shape)) mask_str = f" - {masked_entries} masked entries" if masked_entries else "" underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") From 19125ac2d59f2caf4afd420a7a3cfc675651465c Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Mon, 23 Mar 2026 10:47:09 +0100 Subject: [PATCH 11/20] fix: don't freeze if model is chunked --- linopy/model.py | 4 +--- test/test_constraints.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/linopy/model.py b/linopy/model.py index 2c0c1467..6f9ad444 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -890,12 +890,10 @@ def add_constraints( data = data.assign_attrs(label_range=(start, end), name=name) if self.chunk: - if freeze: - raise ValueError("Chunked constraints cannot be frozen") data = data.chunk(self.chunk) constraint = MutableConstraint(data, name=name, model=self, skip_broadcast=True) - return self.constraints.add(constraint, freeze=freeze) + return self.constraints.add(constraint, freeze=freeze and not self.chunk) def add_objective( self, diff --git a/test/test_constraints.py b/test/test_constraints.py index 1d88d7e4..9a467c8c 100644 --- a/test/test_constraints.py +++ b/test/test_constraints.py @@ -116,9 +116,7 @@ def test_constraint_assignment_chunked() -> None: lower = pd.DataFrame(np.zeros((10, 10))) upper = pd.Series(np.ones(10)) x = m.add_variables(lower, upper) - with pytest.raises(ValueError, match="Chunked constraints cannot be frozen"): - m.add_constraints(x, GREATER_EQUAL, 0, name="c") - m.add_constraints(x, GREATER_EQUAL, 0, name="c", freeze=False) + m.add_constraints(x, GREATER_EQUAL, 0, name="c") assert m.constraints.coeffs.c.data.shape == ( 10, 10, From 58879fc8aad4c002597aa2492c78bc795e213e36 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Tue, 24 Mar 2026 10:13:20 +0100 Subject: [PATCH 12/20] fix typing errors --- linopy/common.py | 26 ++++++++------- linopy/constants.py | 6 ++-- linopy/constraints.py | 37 +++++++++++++++++----- linopy/expressions.py | 51 ++++++++++++++++-------------- linopy/io.py | 14 ++++----- linopy/matrices.py | 4 +-- linopy/model.py | 2 +- linopy/solvers.py | 2 +- linopy/testing.py | 4 +-- linopy/types.py | 70 ++++++++++++++++++++--------------------- linopy/variables.py | 15 ++++----- test/test_constraint.py | 5 ++- 12 files changed, 133 insertions(+), 103 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index c47e9775..4ab94d02 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -19,6 +19,7 @@ import pandas as pd import polars as pl from numpy import arange, signedinteger +from polars.datatypes import DataTypeClass from xarray import DataArray, Dataset, apply_ufunc, broadcast from xarray import align as xr_align from xarray.core import dtypes, indexing @@ -327,7 +328,7 @@ def check_has_nulls(df: pd.DataFrame, name: str) -> None: raise ValueError(f"Fields {name} contains nan's in field(s) {fields}") -def infer_schema_polars(ds: Dataset) -> dict[Hashable, pl.DataType]: +def infer_schema_polars(ds: Dataset) -> dict[str, DataTypeClass]: """ Infer the polars data schema from a xarray dataset. @@ -339,21 +340,22 @@ def infer_schema_polars(ds: Dataset) -> dict[Hashable, pl.DataType]: ------- dict: A dictionary mapping column names to their corresponding Polars data types. """ - schema = {} + schema: dict[str, DataTypeClass] = {} np_major_version = int(np.__version__.split(".")[0]) use_int32 = os.name == "nt" and np_major_version < 2 for name, array in ds.items(): + name = str(name) if np.issubdtype(array.dtype, np.integer): schema[name] = pl.Int32 if use_int32 else pl.Int64 elif np.issubdtype(array.dtype, np.floating): - schema[name] = pl.Float64 # type: ignore + schema[name] = pl.Float64 elif np.issubdtype(array.dtype, np.bool_): - schema[name] = pl.Boolean # type: ignore + schema[name] = pl.Boolean elif np.issubdtype(array.dtype, np.object_): - schema[name] = pl.Object # type: ignore + schema[name] = pl.Object else: - schema[name] = pl.Utf8 # type: ignore - return schema # type: ignore + schema[name] = pl.Utf8 + return schema def to_polars(ds: Dataset, **kwargs: Any) -> pl.DataFrame: @@ -429,7 +431,7 @@ def filter_nulls_polars(df: pl.DataFrame) -> pl.DataFrame: if "labels" in df.columns: cond.append(pl.col("labels").ne(-1)) - cond = reduce(operator.and_, cond) # type: ignore + cond = reduce(operator.and_, cond) # type: ignore[arg-type] return df.filter(cond) @@ -655,7 +657,7 @@ def iterate_slices( start = i * chunk_size end = min(start + chunk_size, size_of_leading_dim) slice_dict = {leading_dim: slice(start, end)} - yield ds.isel(slice_dict) + yield ds.isel(slice_dict) # type: ignore[attr-defined] def _remap(array: np.ndarray, mapping: np.ndarray) -> np.ndarray: @@ -1367,7 +1369,7 @@ def __getitem__( # expand the indexer so we can handle Ellipsis labels = indexing.expanded_indexer(key, self.object.ndim) key = dict(zip(self.object.dims, labels)) - return self.object.sel(key) + return self.object.sel(key) # type: ignore[attr-defined] class EmptyDeprecationWrapper: @@ -1439,9 +1441,9 @@ def coords_from_dataset(ds: Dataset, coord_dims: list[str]) -> list[pd.Index]: if f"_coord_{d}_codes" in ds: codes_2d = ds[f"_coord_{d}_codes"].values.T level_names = [ - k[len(f"_coord_{d}_level_") :] + str(k)[len(f"_coord_{d}_level_") :] for k in ds - if k.startswith(f"_coord_{d}_level_") + if str(k).startswith(f"_coord_{d}_level_") ] arrays = [ ds[f"_coord_{d}_level_{ln}"].values[codes_2d[i]] diff --git a/linopy/constants.py b/linopy/constants.py index 96957be1..690f93c4 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -214,9 +214,11 @@ def process(cls, status: str, termination_condition: str) -> "Status": @classmethod def from_termination_condition( - cls, termination_condition: Union["TerminationCondition", str] + cls, termination_condition: Union["TerminationCondition", str, None] ) -> "Status": - termination_condition = TerminationCondition.process(termination_condition) + termination_condition = TerminationCondition.process( + termination_condition if termination_condition is not None else "unknown" + ) solver_status = SolverStatus.from_termination_condition(termination_condition) return cls(solver_status, termination_condition) diff --git a/linopy/constraints.py b/linopy/constraints.py index fae5a03a..a161aaec 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -169,6 +169,11 @@ def rhs(self) -> DataArray: def dual(self) -> DataArray: """Get the dual values DataArray.""" + @dual.setter + @abstractmethod + def dual(self, value: DataArray) -> None: + """Set the dual values DataArray.""" + @abstractmethod def has_variable(self, variable: variables.Variable) -> bool: """Check if the constraint references any of the given variable labels.""" @@ -185,6 +190,18 @@ def sanitize_missings(self) -> ConstraintBase: def sanitize_infinities(self) -> ConstraintBase: """Mask out rows with invalid infinite RHS values.""" + @abstractmethod + def to_polars(self) -> pl.DataFrame: + """Convert constraint to a polars DataFrame.""" + + @abstractmethod + def freeze(self) -> Constraint: + """Return an immutable Constraint (CSR-backed).""" + + @abstractmethod + def mutable(self) -> MutableConstraint: + """Return a mutable MutableConstraint.""" + @abstractmethod def to_matrix_with_rhs( self, label_index: VariableLabelIndex @@ -298,7 +315,8 @@ def mask(self) -> DataArray | None: (True) and disabled (False). """ if self.is_assigned: - return (self.labels != FILL_VALUE["labels"]).astype(bool) + result: DataArray = self.labels != FILL_VALUE["labels"] # type: ignore[assignment] + return result.astype(bool) return None @property @@ -391,7 +409,7 @@ def flat(self) -> pd.DataFrame: """ ds = self.data - def mask_func(data: pd.DataFrame) -> pd.Series: + def mask_func(data: dict) -> pd.Series: mask = (data["vars"] != -1) & (data["coeffs"] != 0) if "labels" in data: mask &= data["labels"] != -1 @@ -593,7 +611,7 @@ def nterm(self) -> int: @property def coord_names(self) -> list[str]: - return [c.name for c in self._coords] + return [str(c.name) for c in self._coords] @property def labels(self) -> DataArray: @@ -828,8 +846,8 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint: ) rhs = ds["rhs"].values sign = attrs["sign"] - cindex = int(attrs["cindex"]) - cindex = cindex if cindex >= 0 else None + _cindex_raw = int(attrs["cindex"]) + cindex: int | None = _cindex_raw if _cindex_raw >= 0 else None coord_dims = attrs["coord_dims"] if isinstance(coord_dims, str): coord_dims = [coord_dims] @@ -892,7 +910,7 @@ def mutable(self) -> MutableConstraint: """Convert to a MutableConstraint.""" return MutableConstraint(self.data, self._model, self._name) - def to_polars(self) -> Any: + def to_polars(self) -> pl.DataFrame: """Convert to polars DataFrame — delegates to mutable().""" return self.mutable().to_polars() @@ -1598,6 +1616,8 @@ def set_blocks(self, block_map: np.ndarray) -> None: N = block_map.max() for name, constraint in self.items(): + if not isinstance(constraint, MutableConstraint): + self.data[name] = constraint = constraint.mutable() res = xr.full_like(constraint.labels, N + 1, dtype=block_map.dtype) entries = replace_by_map(constraint.vars, block_map) @@ -1679,9 +1699,12 @@ def reset_dual(self) -> None: cindex=c._cindex, dual=None, ) - else: + elif isinstance(c, MutableConstraint): if "dual" in c.data: c._data = c.data.drop_vars("dual") + else: + msg = f"reset_dual encountered an unknown constraint type: {type(c)}" + raise NotImplementedError(msg) class AnonymousScalarConstraint: diff --git a/linopy/expressions.py b/linopy/expressions.py index 80eb15c9..58325829 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -30,6 +30,7 @@ from xarray import Coordinates, DataArray, Dataset, IndexVariable from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates from xarray.core.indexes import Indexes +from xarray.core.types import JoinOptions from xarray.core.utils import Frozen try: @@ -38,7 +39,7 @@ from xarray.computation.rolling import DatasetRolling except ImportError: import xarray.core.rolling - from xarray.core.rolling import DatasetRolling # type: ignore + from xarray.core.rolling import DatasetRolling # type: ignore[no-redef] from types import EllipsisType, NotImplementedType @@ -369,6 +370,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: f"data must be an instance of {supported_types}, got {type(data)}" ) + data = cast(Dataset, data) if not set(data).issuperset({"coeffs", "vars"}): raise ValueError( "data must contain the fields 'coeffs' and 'vars' or 'const'" @@ -390,6 +392,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: data = assign_multiindex_safe(data, const=data.const.astype(float)) (data,) = xr.broadcast(data, exclude=HELPER_DIMS) + data = cast(Dataset, data) (coeffs_vars,) = xr.broadcast(data[["coeffs", "vars"]], exclude=[FACTOR_DIM]) coeffs_vars_dict = {str(k): v for k, v in coeffs_vars.items()} data = assign_multiindex_safe(data, **coeffs_vars_dict) @@ -406,7 +409,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: raise ValueError("model must be an instance of linopy.Model") self._model = model - self._data = data + self._data = cast(Dataset, data) def __repr__(self) -> str: """ @@ -545,13 +548,13 @@ def _multiply_by_linear_expression( res = res + self.const * other.reset_const() if other.has_constant: res = res + self.reset_const() * other.const - return res + return cast(QuadraticExpression, res) def _align_constant( self: GenericExpression, other: DataArray, fill_value: float = 0, - join: str | None = None, + join: JoinOptions | None = None, ) -> tuple[DataArray, DataArray, bool]: """ Align a constant DataArray with self.const. @@ -589,12 +592,12 @@ def _align_constant( self.const, other, join=join, - fill_value=fill_value, # type: ignore[call-overload] + fill_value=fill_value, ) return self_const, aligned, True def _add_constant( - self: GenericExpression, other: ConstantLike, join: str | None = None + self: GenericExpression, other: ConstantLike, join: JoinOptions | None = None ) -> GenericExpression: # NaN values in self.const or other are filled with 0 (additive identity) # so that missing data does not silently propagate through arithmetic. @@ -620,7 +623,7 @@ def _apply_constant_op( other: ConstantLike, op: Callable[[DataArray, DataArray], DataArray], fill_value: float, - join: str | None = None, + join: JoinOptions | None = None, ) -> GenericExpression: """ Apply a constant operation (mul, div, etc.) to this expression with a scalar or array. @@ -648,12 +651,12 @@ def _apply_constant_op( return self.assign(coeffs=op(coeffs, factor), const=op(self_const, factor)) def _multiply_by_constant( - self: GenericExpression, other: ConstantLike, join: str | None = None + self: GenericExpression, other: ConstantLike, join: JoinOptions | None = None ) -> GenericExpression: return self._apply_constant_op(other, operator.mul, fill_value=0, join=join) def _divide_by_constant( - self: GenericExpression, other: ConstantLike, join: str | None = None + self: GenericExpression, other: ConstantLike, join: JoinOptions | None = None ) -> GenericExpression: return self._apply_constant_op(other, operator.truediv, fill_value=1, join=join) @@ -727,7 +730,7 @@ def __lt__(self, other: Any) -> NotImplementedType: def add( self: GenericExpression, other: SideLike, - join: str | None = None, + join: JoinOptions | None = None, ) -> GenericExpression | QuadraticExpression: """ Add an expression to others. @@ -755,7 +758,7 @@ def add( def sub( self: GenericExpression, other: SideLike, - join: str | None = None, + join: JoinOptions | None = None, ) -> GenericExpression | QuadraticExpression: """ Subtract others from expression. @@ -774,7 +777,7 @@ def sub( def mul( self: GenericExpression, other: SideLike, - join: str | None = None, + join: JoinOptions | None = None, ) -> GenericExpression | QuadraticExpression: """ Multiply the expr by a factor. @@ -799,7 +802,7 @@ def mul( def div( self: GenericExpression, other: VariableLike | ConstantLike, - join: str | None = None, + join: JoinOptions | None = None, ) -> GenericExpression | QuadraticExpression: """ Divide the expr by a factor. @@ -826,7 +829,7 @@ def div( def le( self: GenericExpression, rhs: SideLike, - join: str | None = None, + join: JoinOptions | None = None, ) -> MutableConstraint: """ Less than or equal constraint. @@ -843,9 +846,9 @@ def le( return self.to_constraint(LESS_EQUAL, rhs, join=join) def ge( - self: GenericExpression, + self, rhs: SideLike, - join: str | None = None, + join: JoinOptions | None = None, ) -> MutableConstraint: """ Greater than or equal constraint. @@ -862,9 +865,9 @@ def ge( return self.to_constraint(GREATER_EQUAL, rhs, join=join) def eq( - self: GenericExpression, + self, rhs: SideLike, - join: str | None = None, + join: JoinOptions | None = None, ) -> MutableConstraint: """ Equality constraint. @@ -1015,7 +1018,7 @@ def _map_solution(self) -> DataArray: sol = pd.Series(m.matrices.sol, m.matrices.vlabels) sol[-1] = np.nan idx = np.ravel(self.vars) - values = sol[idx].to_numpy().reshape(self.vars.shape) + values = np.asarray(sol[idx]).reshape(self.vars.shape) return xr.DataArray(values, dims=self.vars.dims, coords=self.vars.coords) @property @@ -1120,7 +1123,7 @@ def cumsum( return self.rolling(dim=dim_dict).sum(keep_attrs=keep_attrs, skipna=skipna) def to_constraint( - self, sign: SignLike, rhs: SideLike, join: str | None = None + self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None ) -> MutableConstraint: """ Convert a linear expression to a constraint. @@ -1734,7 +1737,7 @@ def flat(self) -> pd.DataFrame: """ ds = self.data - def mask_func(data: pd.DataFrame) -> pd.Series: + def mask_func(data: dict) -> pd.Series: mask = (data["vars"] != -1) & (data["coeffs"] != 0) return mask @@ -2217,7 +2220,7 @@ def solution(self) -> DataArray: return sol.rename("solution") def to_constraint( - self, sign: SignLike, rhs: SideLike, join: str | None = None + self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None ) -> NotImplementedType: raise NotImplementedError( "Quadratic expressions cannot be used in constraints." @@ -2233,7 +2236,7 @@ def flat(self) -> DataFrame: ).to_dataset(FACTOR_DIM) ds = self.data.drop_vars("vars").assign(vars) - def mask_func(data: pd.DataFrame) -> pd.Series: + def mask_func(data: dict) -> pd.Series: mask = ((data["vars1"] != -1) | (data["vars2"] != -1)) & ( data["coeffs"] != 0 ) @@ -2350,7 +2353,7 @@ def merge( ], dim: str = TERM_DIM, cls: type[GenericExpression] = None, # type: ignore - join: str | None = None, + join: JoinOptions | None = None, **kwargs: Any, ) -> GenericExpression: """ diff --git a/linopy/io.py b/linopy/io.py index 8c0d8541..fa340a5a 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1167,15 +1167,15 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: ------- m : linopy.Model """ - from linopy.constraints import Constraint - from linopy.model import ( + from linopy.constraints import ( + Constraint, + ConstraintBase, Constraints, - LinearExpression, - Model, MutableConstraint, - Variable, - Variables, ) + from linopy.expressions import LinearExpression + from linopy.model import Model + from linopy.variables import Variable, Variables if isinstance(path, str): path = Path(path) @@ -1222,7 +1222,7 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: cons = [str(k) for k in ds if str(k).startswith("constraints")] con_names = list({str(k).rsplit("-", 1)[0] for k in cons}) - constraints = {} + constraints: dict[str, ConstraintBase] = {} for k in sorted(con_names): name = remove_prefix(k, "constraints") con_ds = get_prefix(ds, k) diff --git a/linopy/matrices.py b/linopy/matrices.py index 93f3e981..9496c2a7 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numpy as np import scipy.sparse @@ -91,7 +91,7 @@ def _build_cons(self) -> None: b_list.append(b) sense_list.append(sense) - self.A = scipy.sparse.vstack(csrs, format="csr") + self.A = cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr")) self.clabels = np.concatenate(clabels_list) self.b = np.concatenate(b_list) if b_list else np.array([]) self.sense = ( diff --git a/linopy/model.py b/linopy/model.py index 6f9ad444..e0723982 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1311,7 +1311,7 @@ def solve( sanitize_zeros: bool = True, sanitize_infinities: bool = True, slice_size: int = 2_000_000, - remote: RemoteHandler | OetcHandler = None, # type: ignore + remote: RemoteHandler | OetcHandler | None = None, progress: bool | None = None, mock_solve: bool = False, reformulate_sos: bool | Literal["auto"] = False, diff --git a/linopy/solvers.py b/linopy/solvers.py index 10731547..a8e950ce 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -54,7 +54,7 @@ FILE_IO_APIS = ["lp", "lp-polars", "mps"] IO_APIS = FILE_IO_APIS + ["direct"] -available_solvers = [] +available_solvers: list[str] = [] which = "where" if os.name == "nt" else "which" diff --git a/linopy/testing.py b/linopy/testing.py index 0392064e..e6a58a0d 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -2,7 +2,7 @@ from xarray.testing import assert_equal -from linopy.constraints import Constraint, _con_unwrap +from linopy.constraints import ConstraintBase, _con_unwrap from linopy.expressions import LinearExpression, QuadraticExpression, _expr_unwrap from linopy.model import Model from linopy.variables import Variable, _var_unwrap @@ -29,7 +29,7 @@ def assert_quadequal( return assert_equal(_expr_unwrap(a), _expr_unwrap(b)) -def assert_conequal(a: Constraint, b: Constraint, strict: bool = True) -> None: +def assert_conequal(a: ConstraintBase, b: ConstraintBase, strict: bool = True) -> None: """ Assert that two constraints are equal. diff --git a/linopy/types.py b/linopy/types.py index 7238c552..d41b54ec 100644 --- a/linopy/types.py +++ b/linopy/types.py @@ -2,7 +2,7 @@ from collections.abc import Hashable, Iterable, Mapping, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, TypeAlias import numpy import polars as pl @@ -11,7 +11,10 @@ from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates if TYPE_CHECKING: - from linopy.constraints import AnonymousScalarConstraint, Constraint + from linopy.constraints import ( + AnonymousScalarConstraint, + ConstraintBase, + ) from linopy.expressions import ( LinearExpression, QuadraticExpression, @@ -20,37 +23,34 @@ from linopy.piecewise import PiecewiseConstraintDescriptor from linopy.variables import ScalarVariable, Variable -# Type aliases using Union for Python 3.9 compatibility -CoordsLike = Union[ # noqa: UP007 - Sequence[Sequence | Index | DataArray], - Mapping, - DataArrayCoordinates, - DatasetCoordinates, -] -DimsLike = Union[str, Iterable[Hashable]] # noqa: UP007 +CoordsLike: TypeAlias = ( + Sequence[Sequence | Index | DataArray] + | Mapping + | DataArrayCoordinates + | DatasetCoordinates +) +DimsLike: TypeAlias = str | Iterable[Hashable] -ConstantLike = Union[ # noqa: UP007 - int, - float, - numpy.floating, - numpy.integer, - numpy.ndarray, - DataArray, - Series, - DataFrame, - pl.Series, -] -SignLike = Union[str, numpy.ndarray, DataArray, Series, DataFrame] # noqa: UP007 -VariableLike = Union["ScalarVariable", "Variable"] -ExpressionLike = Union[ - "ScalarLinearExpression", - "LinearExpression", - "QuadraticExpression", -] -ConstraintLike = Union[ - "Constraint", "AnonymousScalarConstraint", "PiecewiseConstraintDescriptor" -] -LinExprLike = Union["Variable", "LinearExpression"] -MaskLike = Union[numpy.ndarray, DataArray, Series, DataFrame] # noqa: UP007 -SideLike = Union[ConstantLike, VariableLike, ExpressionLike] # noqa: UP007 -PathLike = Union[str, Path] # noqa: UP007 +ConstantLike: TypeAlias = ( + int + | float + | numpy.floating + | numpy.integer + | numpy.ndarray + | DataArray + | Series + | DataFrame + | pl.Series +) +SignLike: TypeAlias = str | numpy.ndarray | DataArray | Series | DataFrame +VariableLike: TypeAlias = ScalarVariable | Variable +ExpressionLike: TypeAlias = ( + ScalarLinearExpression | LinearExpression | QuadraticExpression +) +ConstraintLike: TypeAlias = ( + ConstraintBase | AnonymousScalarConstraint | PiecewiseConstraintDescriptor +) +LinExprLike: TypeAlias = Variable | LinearExpression +MaskLike: TypeAlias = numpy.ndarray | DataArray | Series | DataFrame +SideLike: TypeAlias = ConstantLike | VariableLike | ExpressionLike +PathLike: TypeAlias = str | Path diff --git a/linopy/variables.py b/linopy/variables.py index 8a465a39..0f5d05e4 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -27,6 +27,7 @@ from xarray import DataArray, Dataset, broadcast from xarray.core.coordinates import DatasetCoordinates from xarray.core.indexes import Indexes +from xarray.core.types import JoinOptions from xarray.core.utils import Frozen import linopy.expressions as expressions @@ -579,7 +580,7 @@ def __contains__(self, value: str) -> bool: return self.data.__contains__(value) def add( - self, other: SideLike, join: str | None = None + self, other: SideLike, join: JoinOptions | None = None ) -> LinearExpression | QuadraticExpression: """ Add variables to linear expressions or other variables. @@ -596,7 +597,7 @@ def add( return self.to_linexpr().add(other, join=join) def sub( - self, other: SideLike, join: str | None = None + self, other: SideLike, join: JoinOptions | None = None ) -> LinearExpression | QuadraticExpression: """ Subtract linear expressions or other variables from the variables. @@ -613,7 +614,7 @@ def sub( return self.to_linexpr().sub(other, join=join) def mul( - self, other: ConstantLike, join: str | None = None + self, other: ConstantLike, join: JoinOptions | None = None ) -> LinearExpression | QuadraticExpression: """ Multiply variables with a coefficient. @@ -630,7 +631,7 @@ def mul( return self.to_linexpr().mul(other, join=join) def div( - self, other: ConstantLike, join: str | None = None + self, other: ConstantLike, join: JoinOptions | None = None ) -> LinearExpression | QuadraticExpression: """ Divide variables with a coefficient. @@ -646,7 +647,7 @@ def div( """ return self.to_linexpr().div(other, join=join) - def le(self, rhs: SideLike, join: str | None = None) -> MutableConstraint: + def le(self, rhs: SideLike, join: JoinOptions | None = None) -> MutableConstraint: """ Less than or equal constraint. @@ -661,7 +662,7 @@ def le(self, rhs: SideLike, join: str | None = None) -> MutableConstraint: """ return self.to_linexpr().le(rhs, join=join) - def ge(self, rhs: SideLike, join: str | None = None) -> MutableConstraint: + def ge(self, rhs: SideLike, join: JoinOptions | None = None) -> MutableConstraint: """ Greater than or equal constraint. @@ -676,7 +677,7 @@ def ge(self, rhs: SideLike, join: str | None = None) -> MutableConstraint: """ return self.to_linexpr().ge(rhs, join=join) - def eq(self, rhs: SideLike, join: str | None = None) -> MutableConstraint: + def eq(self, rhs: SideLike, join: JoinOptions | None = None) -> MutableConstraint: """ Equality constraint. diff --git a/test/test_constraint.py b/test/test_constraint.py index 9fbb042e..74389e4c 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -26,7 +26,6 @@ ) from linopy.constraints import ( AnonymousScalarConstraint, - Constraint, ConstraintBase, Constraints, MutableConstraint, @@ -54,7 +53,7 @@ def y(m: Model) -> linopy.Variable: @pytest.fixture -def c(m: Model) -> linopy.constraints.Constraint: +def c(m: Model) -> linopy.constraints.ConstraintBase: return m.constraints["c"] @@ -227,7 +226,7 @@ def test_constraint_inherited_properties( def test_constraint_wrapped_methods(x: linopy.Variable, y: linopy.Variable) -> None: - con: Constraint = 10 * x + y <= 10 + con: MutableConstraint = 10 * x + y <= 10 # Test wrapped methods con.assign({"new_var": xr.DataArray(np.zeros((2, 2)), coords=[range(2), range(2)])}) From 7fe0392ead74467eb94cc9769d4c80a637e6ef1e Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Tue, 24 Mar 2026 10:31:31 +0100 Subject: [PATCH 13/20] fix: bring back forward-refs --- linopy/types.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/linopy/types.py b/linopy/types.py index d41b54ec..34f7a5a9 100644 --- a/linopy/types.py +++ b/linopy/types.py @@ -2,7 +2,7 @@ from collections.abc import Hashable, Iterable, Mapping, Sequence from pathlib import Path -from typing import TYPE_CHECKING, TypeAlias +from typing import TYPE_CHECKING, TypeAlias, Union import numpy import polars as pl @@ -43,14 +43,23 @@ | pl.Series ) SignLike: TypeAlias = str | numpy.ndarray | DataArray | Series | DataFrame -VariableLike: TypeAlias = ScalarVariable | Variable -ExpressionLike: TypeAlias = ( - ScalarLinearExpression | LinearExpression | QuadraticExpression -) -ConstraintLike: TypeAlias = ( - ConstraintBase | AnonymousScalarConstraint | PiecewiseConstraintDescriptor -) -LinExprLike: TypeAlias = Variable | LinearExpression MaskLike: TypeAlias = numpy.ndarray | DataArray | Series | DataFrame -SideLike: TypeAlias = ConstantLike | VariableLike | ExpressionLike PathLike: TypeAlias = str | Path + +# These reference types only available under TYPE_CHECKING, so use Union with strings +VariableLike: TypeAlias = Union["ScalarVariable", "Variable"] +ExpressionLike: TypeAlias = Union[ + "ScalarLinearExpression", "LinearExpression", "QuadraticExpression" +] +ConstraintLike: TypeAlias = Union[ + "ConstraintBase", "AnonymousScalarConstraint", "PiecewiseConstraintDescriptor" +] +LinExprLike: TypeAlias = Union["Variable", "LinearExpression"] +SideLike: TypeAlias = Union[ + ConstantLike, + "ScalarVariable", + "Variable", + "ScalarLinearExpression", + "LinearExpression", + "QuadraticExpression", +] From 92752e2416aec5071856d9ffb90f5cc451bd9a29 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Tue, 24 Mar 2026 10:39:17 +0100 Subject: [PATCH 14/20] fix issues in tests --- linopy/model.py | 33 +++++++++++++++++++++++++++++++++ test/test_constraint.py | 2 +- test/test_linear_expression.py | 4 +++- 3 files changed, 37 insertions(+), 2 deletions(-) diff --git a/linopy/model.py b/linopy/model.py index e0723982..e3663287 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -49,6 +49,7 @@ ) from linopy.constraints import ( AnonymousScalarConstraint, + Constraint, ConstraintBase, Constraints, MutableConstraint, @@ -720,6 +721,38 @@ def add_sos_constraints( add_piecewise_constraints = add_piecewise_constraints + @overload + def add_constraints( + self, + lhs: VariableLike + | ExpressionLike + | ConstraintLike + | Sequence[tuple[ConstantLike, VariableLike | str]] + | Callable, + sign: SignLike | None = ..., + rhs: ConstantLike | VariableLike | ExpressionLike | None = ..., + name: str | None = ..., + coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = ..., + mask: MaskLike | None = ..., + freeze: Literal[True] = ..., + ) -> Constraint: ... + + @overload + def add_constraints( + self, + lhs: VariableLike + | ExpressionLike + | ConstraintLike + | Sequence[tuple[ConstantLike, VariableLike | str]] + | Callable, + sign: SignLike | None = ..., + rhs: ConstantLike | VariableLike | ExpressionLike | None = ..., + name: str | None = ..., + coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = ..., + mask: MaskLike | None = ..., + freeze: Literal[False] = ..., + ) -> MutableConstraint: ... + def add_constraints( self, lhs: VariableLike diff --git a/test/test_constraint.py b/test/test_constraint.py index 74389e4c..da3bbf25 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -498,7 +498,7 @@ def test_constraint_assignment_sanitize_zeros( assert c2.nterm == 1 assert c2.has_variable(y) assert not c2.has_variable(x) - csr, _ = c2.to_matrix() + csr, _ = c2.to_matrix(m.variables.label_index) assert (csr.data == 1).all() diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index d3b8d426..27ad5c34 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -14,6 +14,7 @@ import polars as pl import pytest import xarray as xr +from xarray.core.types import JoinOptions from xarray.testing import assert_equal from linopy import LinearExpression, Model, QuadraticExpression, Variable, merge @@ -1920,7 +1921,8 @@ def test_add_constant_join_override(self, a: Variable, c: Variable) -> None: def test_add_same_coords_all_joins(self, a: Variable, c: Variable) -> None: expr_a = 1 * a + 5 const = xr.DataArray([1, 2, 3], dims=["i"], coords={"i": [0, 1, 2]}) - for join in ["override", "outer", "inner"]: + joins: list[JoinOptions] = ["override", "outer", "inner"] + for join in joins: result = expr_a.add(const, join=join) assert list(result.coords["i"].values) == [0, 1, 2] np.testing.assert_array_equal(result.const.values, [6, 7, 8]) From 83dd58ac6c965384e7e2291272733970f6490152 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Tue, 24 Mar 2026 10:44:05 +0100 Subject: [PATCH 15/20] fix: add doc strings to VariableLabelIndex --- linopy/common.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/linopy/common.py b/linopy/common.py index 4ab94d02..a90a5fab 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -957,6 +957,7 @@ def __init__(self, variables: Any) -> None: @cached_property def vlabels(self) -> np.ndarray: + """Active variable labels in encounter order, shape (n_active_vars,).""" label_lists = [] for _, var in self._variables.items(): labels = var.labels.values.ravel() @@ -968,6 +969,12 @@ def vlabels(self) -> np.ndarray: @cached_property def label_to_pos(self) -> np.ndarray: + """ + Mapping from variable label to dense position, shape (_xCounter,). + + Position i in the active variable array corresponds to label vlabels[i]. + Masked or unused labels map to -1. + """ vlabels = self.vlabels n = self._variables.model._xCounter label_to_pos = np.full(n, -1, dtype=np.intp) @@ -976,6 +983,7 @@ def label_to_pos(self) -> np.ndarray: @property def n_active_vars(self) -> int: + """Number of active (non-masked) variables.""" return len(self.vlabels) def invalidate(self) -> None: From 2bec451891897015167fd668969fb8d3ddfc4795 Mon Sep 17 00:00:00 2001 From: Jonas Hoersch Date: Tue, 24 Mar 2026 12:59:27 +0100 Subject: [PATCH 16/20] test: relax dtype assertions for Windows np.int32 compatibility Use np.issubdtype() instead of exact dtype comparisons to allow np.int32/np.float32 that Windows may produce. Co-Authored-By: Claude Sonnet 4.6 --- test/test_constraints.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/test/test_constraints.py b/test/test_constraints.py index 9a467c8c..1667bfec 100644 --- a/test/test_constraints.py +++ b/test/test_constraints.py @@ -36,10 +36,10 @@ def test_constraint_assignment() -> None: assert "con0" in getattr(m.constraints, attr) assert m.constraints.labels.con0.shape == (10, 10) - assert m.constraints.labels.con0.dtype == int - assert m.constraints.coeffs.con0.dtype in (int, float) - assert m.constraints.vars.con0.dtype in (int, float) - assert m.constraints.rhs.con0.dtype in (int, float) + assert np.issubdtype(m.constraints.labels.con0.dtype, np.integer) + assert np.issubdtype(m.constraints.coeffs.con0.dtype, np.number) + assert np.issubdtype(m.constraints.vars.con0.dtype, np.number) + assert np.issubdtype(m.constraints.rhs.con0.dtype, np.number) assert_conequal(m.constraints.con0, con0) @@ -90,10 +90,10 @@ def test_anonymous_constraint_assignment() -> None: assert "con0" in getattr(m.constraints, attr) assert m.constraints.labels.con0.shape == (10, 10) - assert m.constraints.labels.con0.dtype == int - assert m.constraints.coeffs.con0.dtype in (int, float) - assert m.constraints.vars.con0.dtype in (int, float) - assert m.constraints.rhs.con0.dtype in (int, float) + assert np.issubdtype(m.constraints.labels.con0.dtype, np.integer) + assert np.issubdtype(m.constraints.coeffs.con0.dtype, np.number) + assert np.issubdtype(m.constraints.vars.con0.dtype, np.number) + assert np.issubdtype(m.constraints.rhs.con0.dtype, np.number) def test_constraint_assignment_with_tuples() -> None: From 3b2a41530848e05f7834ac6beb9556a32f8d360a Mon Sep 17 00:00:00 2001 From: Fabian Hofmann Date: Wed, 25 Mar 2026 16:44:43 +0100 Subject: [PATCH 17/20] fix: review fixes for #630 (matrix accessor rewrite) (#632) * fix: review fixes for PR #630 (matrix accessor rewrite) - Fix __repr__ passing CSR positions instead of variable labels - Fix set_blocks failing on frozen Constraint - Extract _active_to_dataarray helper to reduce DRY violations - Simplify reset_dual to direct mutation instead of reconstruction - Add tests for freeze/mutable roundtrip, VariableLabelIndex, to_matrix_with_rhs, from_mutable mixed signs, repr correctness * feat: support mixed per-element signs in frozen Constraint, add rhs/lhs setters - Store _sign as str (uniform, fast) or np.ndarray (mixed, per-row) - Add rhs/lhs setters on Constraint via _refreeze_after pattern - Invalidate _dual on mutation; update netcdf serialization for array signs - Add tests for setters, mixed-sign freeze/roundtrip/sanitize/repr/netcdf * feat: mixed per-element signs in frozen Constraint, rhs/lhs setters, DRY helpers --- linopy/constraints.py | 145 ++++++++++++++++++-------------- test/test_constraint.py | 182 ++++++++++++++++++++++++++++++++++++++++ test/test_io.py | 25 ++++++ 3 files changed, 291 insertions(+), 61 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index a161aaec..a6d9b2c9 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -497,9 +497,9 @@ class Constraint(ConstraintBase): constraint grid (including masked/empty rows). rhs : np.ndarray Shape (n_flat,). Right-hand-side values. - sign : str - Constraint sign: one of '=', '<=', '>='. - Note: per-element signs are not supported (documented regression vs MutableConstraint). + sign : str or np.ndarray + Constraint sign. Either a single str ('=', '<=', '>=') for uniform + signs, or a per-row np.ndarray of sign strings for mixed signs. coords : list of pd.Index One index per coordinate dimension defining the constraint grid. model : Model @@ -529,7 +529,7 @@ def __init__( csr: scipy.sparse.csr_array, con_labels: np.ndarray, rhs: np.ndarray, - sign: str, + sign: str | np.ndarray, coords: list[pd.Index], model: Model, name: str = "", @@ -613,16 +613,19 @@ def nterm(self) -> int: def coord_names(self) -> list[str]: return [str(c.name) for c in self._coords] + def _active_to_dataarray( + self, active_values: np.ndarray, fill: float | int | str = -1 + ) -> 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._coords) + @property def labels(self) -> DataArray: """Get labels DataArray, shape (*coord_dims).""" if self._cindex is None: return DataArray([]) - shape = self.shape - full_size = self.full_size - labels_flat = np.full(full_size, -1, dtype=np.int64) - labels_flat[self.active_positions] = self._con_labels - return DataArray(labels_flat.reshape(shape), coords=self._coords) + return self._active_to_dataarray(self._con_labels, fill=-1) @property def coeffs(self) -> DataArray: @@ -648,16 +651,39 @@ def vars(self) -> DataArray: @property def sign(self) -> DataArray: - """Get sign DataArray (scalar, same sign for all entries).""" - return DataArray(np.full(self.shape, self._sign), coords=self._coords) + """Get sign DataArray.""" + if isinstance(self._sign, str): + return DataArray(np.full(self.shape, self._sign), coords=self._coords) + return self._active_to_dataarray(self._sign, fill="") @property def rhs(self) -> DataArray: """Get RHS DataArray, shape (*coord_dims).""" - shape = self.shape - rhs_full = np.full(self.full_size, np.nan) - rhs_full[self.active_positions] = self._rhs - return DataArray(rhs_full.reshape(shape), coords=self._coords) + return self._active_to_dataarray(self._rhs, fill=np.nan) + + @rhs.setter + def rhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: + self._refreeze_after(lambda mc: setattr(mc, "rhs", value)) + + @property + def lhs(self) -> expressions.LinearExpression: + """Get LHS as LinearExpression (triggers Dataset reconstruction).""" + return self.mutable().lhs + + @lhs.setter + def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: + self._refreeze_after(lambda mc: setattr(mc, "lhs", value)) + + def _refreeze_after(self, mutate: Callable[[MutableConstraint], None]) -> None: + mc = self.mutable() + mutate(mc) + refrozen = Constraint.from_mutable(mc, self._cindex) + self._csr = refrozen._csr + self._con_labels = refrozen._con_labels + self._rhs = refrozen._rhs + self._sign = refrozen._sign + self._coords = refrozen._coords + self._dual = None @property @has_optimized_model @@ -667,9 +693,7 @@ def dual(self) -> DataArray: raise AttributeError( "Underlying is optimized but does not have dual values stored." ) - dual_full = np.full(self.full_size, np.nan) - dual_full[self.active_positions] = self._dual - return DataArray(dual_full.reshape(self.shape), coords=self._coords) + return self._active_to_dataarray(self._dual, fill=np.nan) @dual.setter def dual(self, value: DataArray) -> None: @@ -731,24 +755,10 @@ def _to_dataset(self, nterm: int) -> Dataset: def data(self) -> Dataset: """Reconstruct the xarray Dataset from the CSR representation.""" ds = self._to_dataset(self.nterm) - shape = self.shape - active_pos = self.active_positions - rhs_full = np.full(self.full_size, np.nan) - rhs_full[active_pos] = self._rhs - ds = ds.assign( - sign=DataArray(np.full(shape, self._sign), coords=self._coords), - rhs=DataArray(rhs_full.reshape(shape), coords=self._coords), - ) + ds = ds.assign(sign=self.sign, rhs=self.rhs) if self._dual is not None: - dual_full = np.full(self.full_size, np.nan) - dual_full[active_pos] = self._dual - ds = ds.assign( - dual=DataArray(dual_full.reshape(shape), coords=self._coords) - ) - attrs: dict[str, Any] = {"name": self._name} - if self._cindex is not None: - attrs["label_range"] = (self._cindex, self._cindex + self.full_size) - return ds.assign_attrs(attrs) + ds = ds.assign(dual=self._active_to_dataarray(self._dual, fill=np.nan)) + return ds.assign_attrs(self.attrs) def __repr__(self) -> str: """Print the constraint without reconstructing the full Dataset.""" @@ -777,7 +787,8 @@ def row_expr(row: int) -> str: coeffs_row = np.zeros(nterm, dtype=csr.dtype) vars_row[: end - start] = csr.indices[start:end] coeffs_row[: end - start] = csr.data[start:end] - return f"{print_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[self._sign]} {self._rhs[row]}" + sign = self._sign if isinstance(self._sign, str) else self._sign[row] + return f"{print_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[sign]} {self._rhs[row]}" if size > 1: for indices in generate_indices_for_printout(shape, max_lines): @@ -819,21 +830,22 @@ def to_netcdf_ds(self) -> Dataset: "rhs": DataArray(self._rhs, dims=["_flat"]), "_con_labels": DataArray(self._con_labels, dims=["_flat"]), } + if isinstance(self._sign, np.ndarray): + data_vars["_sign"] = DataArray(self._sign, dims=["_flat"]) data_vars.update(coords_to_dataset_vars(self._coords)) if self._dual is not None: data_vars["dual"] = DataArray(self._dual, dims=["_flat"]) dim_names = [c.name for c in self._coords] - return Dataset( - data_vars, - attrs={ - "_linopy_format": "csr", - "sign": self._sign, - "cindex": self._cindex if self._cindex is not None else -1, - "shape": list(csr.shape), - "coord_dims": dim_names, - "name": self._name, - }, - ) + attrs: dict[str, Any] = { + "_linopy_format": "csr", + "cindex": self._cindex if self._cindex is not None else -1, + "shape": list(csr.shape), + "coord_dims": dim_names, + "name": self._name, + } + if isinstance(self._sign, str): + attrs["sign"] = self._sign + return Dataset(data_vars, attrs=attrs) @classmethod def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint: @@ -845,7 +857,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint: shape=shape, ) rhs = ds["rhs"].values - sign = attrs["sign"] + sign: str | np.ndarray = ds["_sign"].values if "_sign" in ds else attrs["sign"] _cindex_raw = int(attrs["cindex"]) cindex: int | None = _cindex_raw if _cindex_raw >= 0 else None coord_dims = attrs["coord_dims"] @@ -873,7 +885,10 @@ def to_matrix_with_rhs( self, label_index: VariableLabelIndex ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: """Return (csr, con_labels, b, sense) — all pre-stored, no recomputation.""" - sense = np.full(len(self._rhs), self._sign[0]) + if isinstance(self._sign, str): + sense = np.full(len(self._rhs), self._sign[0]) + else: + sense = np.array([s[0] for s in self._sign]) return self._csr, self._con_labels, self._rhs, sense def sanitize_zeros(self) -> Constraint: @@ -888,18 +903,25 @@ def sanitize_missings(self) -> Constraint: def sanitize_infinities(self) -> Constraint: """Mask out rows with invalid infinite RHS values (mutates in-place).""" - if self._sign == LESS_EQUAL: - invalid = self._rhs == np.inf - elif self._sign == GREATER_EQUAL: - invalid = self._rhs == -np.inf + if isinstance(self._sign, str): + if self._sign == LESS_EQUAL: + invalid = self._rhs == np.inf + elif self._sign == GREATER_EQUAL: + invalid = self._rhs == -np.inf + else: + return self else: - return self + invalid = ((self._sign == LESS_EQUAL) & (self._rhs == np.inf)) | ( + (self._sign == GREATER_EQUAL) & (self._rhs == -np.inf) + ) if not invalid.any(): return self keep = ~invalid self._csr = self._csr[keep] self._con_labels = self._con_labels[keep] self._rhs = self._rhs[keep] + if not isinstance(self._sign, str): + self._sign = self._sign[keep] return self def freeze(self) -> Constraint: @@ -939,13 +961,14 @@ def from_mutable( active_mask = (labels_flat != -1) & (vars_flat != -1).any(axis=1) rhs = con.rhs.values.ravel()[active_mask] sign_vals = con.sign.values.ravel() - unique_signs = np.unique(sign_vals[active_mask]) - if len(unique_signs) > 1: - raise ValueError( - "Constraint has per-element signs; cannot freeze to immutable Constraint. " - "This is a known limitation — use MutableConstraint instead." - ) - sign = str(unique_signs[0]) if len(unique_signs) == 1 else "=" + active_signs = sign_vals[active_mask] + unique_signs = np.unique(active_signs) + if len(unique_signs) == 0: + sign: str | np.ndarray = "=" + elif len(unique_signs) == 1: + sign = str(unique_signs[0]) + else: + sign = active_signs dual = ( con.data["dual"].values.ravel()[active_mask] if "dual" in con.data else None ) diff --git a/test/test_constraint.py b/test/test_constraint.py index da3bbf25..091e4e54 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -705,3 +705,185 @@ def test_constraints_inequalities(m: Model) -> None: def test_constraints_equalities(m: Model) -> None: assert isinstance(m.constraints.equalities, Constraints) + + +def test_freeze_mutable_roundtrip(m: Model) -> None: + frozen = m.constraints["c"] + assert isinstance(frozen, linopy.constraints.Constraint) + mc = frozen.mutable() + assert isinstance(mc, MutableConstraint) + refrozen = linopy.constraints.Constraint.from_mutable(mc, frozen._cindex) + assert_equal(frozen.labels, refrozen.labels) + assert_equal(frozen.rhs, refrozen.rhs) + assert_equal(frozen.sign, refrozen.sign) + np.testing.assert_array_equal(frozen._csr.toarray(), refrozen._csr.toarray()) + np.testing.assert_array_equal(frozen._con_labels, refrozen._con_labels) + + +def test_freeze_mutable_roundtrip_with_masking() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") + mask = xr.DataArray([True, False, True, False, True], dims=["i"]) + m.add_constraints(x.where(mask) >= 0, name="c") + frozen = m.constraints["c"] + mc = frozen.mutable() + refrozen = linopy.constraints.Constraint.from_mutable(mc, frozen._cindex) + assert_equal(frozen.labels, refrozen.labels) + assert_equal(frozen.rhs, refrozen.rhs) + assert frozen.ncons == refrozen.ncons == 3 + + +def test_from_mutable_mixed_signs() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") + m.add_constraints(x >= 0, name="mixed", freeze=False) + mc = m.constraints["mixed"] + assert isinstance(mc, MutableConstraint) + mc._data["sign"] = xr.DataArray(["<=", ">=", "<="], dims=["i"]) + frozen = linopy.constraints.Constraint.from_mutable(mc) + assert isinstance(frozen._sign, np.ndarray) + assert list(frozen._sign) == ["<=", ">=", "<="] + assert_equal(frozen.sign, mc.sign) + + +def test_variable_label_index(m: Model) -> None: + li = m.variables.label_index + assert li.n_active_vars > 0 + assert len(li.vlabels) == li.n_active_vars + assert li.label_to_pos.shape[0] == m._xCounter + for lbl in li.vlabels: + assert li.label_to_pos[lbl] >= 0 + assert (li.label_to_pos[li.vlabels] == np.arange(li.n_active_vars)).all() + + +def test_variable_label_index_invalidation(m: Model) -> None: + li = m.variables.label_index + old_vlabels = li.vlabels.copy() + m.add_variables(name="w") + li.invalidate() + assert len(li.vlabels) > len(old_vlabels) + + +def test_to_matrix_with_rhs(m: Model) -> None: + c = m.constraints["c"] + li = m.variables.label_index + csr, con_labels, b, sense = c.to_matrix_with_rhs(li) + assert csr.shape[0] == len(con_labels) + assert csr.shape[0] == len(b) + assert csr.shape[0] == len(sense) + assert all(s in ("<", ">", "=") for s in sense) + np.testing.assert_array_equal(b, c._rhs) + + +def test_to_matrix_with_rhs_mutable(m: Model) -> None: + mc = m.constraints["c"].mutable() + li = m.variables.label_index + csr, con_labels, b, sense = mc.to_matrix_with_rhs(li) + assert csr.shape[0] == len(con_labels) + assert csr.shape[0] == len(b) + assert csr.shape[0] == len(sense) + + +def test_constraint_repr_shows_variable_names(m: Model) -> None: + c = m.constraints["c"] + r = repr(c) + assert "x" in r + + +def test_freeze_mixed_signs_from_rule() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + coords = [pd.RangeIndex(4, name="i")] + + def bound(m, i): + if i % 2: + return x.at[i] >= i + return x.at[i] == 0.0 + + con = m.add_constraints(bound, coords=coords, name="mixed_rule") + assert isinstance(con, linopy.constraints.Constraint) + assert isinstance(con._sign, np.ndarray) + assert con.ncons == 4 + expected_signs = ["=", ">=", "=", ">="] + assert list(con._sign) == expected_signs + np.testing.assert_array_equal(con.sign.values, expected_signs) + + +def test_frozen_rhs_setter() -> None: + m = Model() + time = pd.RangeIndex(5, name="t") + x = m.add_variables(lower=0, coords=[time], name="x") + con = m.add_constraints(x >= 1, name="c") + assert isinstance(con, linopy.constraints.Constraint) + con.rhs = 10 + np.testing.assert_array_equal(con._rhs, np.full(5, 10.0)) + factor = pd.Series(range(5), index=time) + con.rhs = 2 * factor + np.testing.assert_array_equal(con._rhs, 2 * np.arange(5, dtype=float)) + + +def test_frozen_lhs_setter() -> None: + m = Model() + time = pd.RangeIndex(5, name="t") + x = m.add_variables(lower=0, coords=[time], name="x") + y = m.add_variables(lower=0, coords=[time], name="y") + con = m.add_constraints(x >= 0, name="c") + assert isinstance(con, linopy.constraints.Constraint) + con.lhs = 3 * x + 2 * y + lhs = con.mutable().lhs + assert lhs.nterm == 2 + + +def test_frozen_setter_invalidates_dual() -> None: + m = Model() + x = m.add_variables(lower=0, coords=[pd.RangeIndex(3, name="i")], name="x") + con = m.add_constraints(x >= 0, name="c") + con._dual = np.array([1.0, 2.0, 3.0]) + con.rhs = 10 + assert con._dual is None + + +def test_mixed_sign_to_matrix_with_rhs() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + coords = [pd.RangeIndex(4, name="i")] + + def bound(m, i): + if i % 2: + return x.at[i] >= i + return x.at[i] == 0.0 + + con = m.add_constraints(bound, coords=coords, name="c") + li = m.variables.label_index + csr, con_labels, b, sense = con.to_matrix_with_rhs(li) + assert len(sense) == 4 + assert list(sense) == ["=", ">", "=", ">"] + + +def test_mixed_sign_sanitize_infinities() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + m.add_constraints(x >= 0, name="c", freeze=False) + mc = m.constraints["c"] + mc._data["sign"] = xr.DataArray(["<=", ">=", "<=", ">="], dims=["i"]) + mc._data["rhs"] = xr.DataArray([np.inf, -np.inf, 1.0, 2.0], dims=["i"]) + frozen = mc.freeze() + frozen.sanitize_infinities() + assert frozen.ncons == 2 + np.testing.assert_array_equal(frozen._rhs, [1.0, 2.0]) + + +def test_mixed_sign_repr() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + coords = [pd.RangeIndex(4, name="i")] + + def bound(m, i): + if i % 2: + return x.at[i] >= i + return x.at[i] == 0.0 + + con = m.add_constraints(bound, coords=coords, name="c") + r = repr(con) + assert "≥" in r + assert "=" in r diff --git a/test/test_io.py b/test/test_io.py index 3d269636..1a97bd22 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -95,6 +95,31 @@ def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None: assert_model_equal(m, p) +def test_model_to_netcdf_mixed_sign_constraint(tmp_path: Path) -> None: + from linopy.constraints import Constraint + + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + + def bound(m, i): + if i % 2: + return x.at[i] >= i + return x.at[i] == 0.0 + + m.add_constraints(bound, coords=[pd.RangeIndex(4, name="i")], name="c") + assert isinstance(m.constraints["c"], Constraint) + + fn = tmp_path / "test_mixed_sign.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert isinstance(p.constraints["c"], Constraint) + import numpy as np + + np.testing.assert_array_equal(m.constraints["c"]._sign, p.constraints["c"]._sign) + assert_model_equal(m, p) + + def test_model_to_netcdf_with_sense(model: Model, tmp_path: Path) -> None: m = model m.objective.sense = "max" From 189f0c55f769c3aa49c47292e9854538d2b88f80 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 26 Mar 2026 09:24:33 +0100 Subject: [PATCH 18/20] rename Constraint to CSRConstraint --- linopy/__init__.py | 4 ++-- linopy/constraints.py | 28 ++++++++++++++-------------- linopy/io.py | 4 ++-- linopy/matrices.py | 4 ++-- linopy/model.py | 4 ++-- test/test_constraint.py | 38 +++++++++++++++++++------------------- test/test_io.py | 12 ++++++------ test/test_repr.py | 8 ++++---- 8 files changed, 51 insertions(+), 51 deletions(-) diff --git a/linopy/__init__.py b/linopy/__init__.py index 12a1cb9e..3236d64c 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -16,9 +16,9 @@ from linopy.config import options from linopy.constants import EQUAL, GREATER_EQUAL, LESS_EQUAL, PerformanceWarning from linopy.constraints import ( - Constraint, ConstraintBase, Constraints, + CSRConstraint, MutableConstraint, ) from linopy.expressions import LinearExpression, QuadraticExpression, merge @@ -34,7 +34,7 @@ pass __all__ = ( - "Constraint", + "CSRConstraint", "ConstraintBase", "Constraints", "MutableConstraint", diff --git a/linopy/constraints.py b/linopy/constraints.py index a6d9b2c9..c07c8db1 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -195,7 +195,7 @@ def to_polars(self) -> pl.DataFrame: """Convert constraint to a polars DataFrame.""" @abstractmethod - def freeze(self) -> Constraint: + def freeze(self) -> CSRConstraint: """Return an immutable Constraint (CSR-backed).""" @abstractmethod @@ -486,9 +486,9 @@ def to_netcdf_ds(self) -> Dataset: iterate_slices = iterate_slices -class Constraint(ConstraintBase): +class CSRConstraint(ConstraintBase): """ - Immutable constraint backed by a CSR sparse matrix. + Frozen constraint backed by a CSR sparse matrix. Parameters ---------- @@ -677,7 +677,7 @@ def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: def _refreeze_after(self, mutate: Callable[[MutableConstraint], None]) -> None: mc = self.mutable() mutate(mc) - refrozen = Constraint.from_mutable(mc, self._cindex) + refrozen = CSRConstraint.from_mutable(mc, self._cindex) self._csr = refrozen._csr self._con_labels = refrozen._con_labels self._rhs = refrozen._rhs @@ -848,7 +848,7 @@ def to_netcdf_ds(self) -> Dataset: return Dataset(data_vars, attrs=attrs) @classmethod - def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint: + def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: """Reconstruct a Constraint from a netcdf Dataset (CSR format).""" attrs = ds.attrs shape = tuple(attrs["shape"]) @@ -891,17 +891,17 @@ def to_matrix_with_rhs( sense = np.array([s[0] for s in self._sign]) return self._csr, self._con_labels, self._rhs, sense - def sanitize_zeros(self) -> Constraint: + def sanitize_zeros(self) -> CSRConstraint: """Remove terms with zero or near-zero coefficients (mutates in-place).""" self._csr.data[np.abs(self._csr.data) <= 1e-10] = 0 self._csr.eliminate_zeros() return self - def sanitize_missings(self) -> Constraint: + def sanitize_missings(self) -> CSRConstraint: """No-op: missing rows are already excluded during freezing.""" return self - def sanitize_infinities(self) -> Constraint: + def sanitize_infinities(self) -> CSRConstraint: """Mask out rows with invalid infinite RHS values (mutates in-place).""" if isinstance(self._sign, str): if self._sign == LESS_EQUAL: @@ -924,7 +924,7 @@ def sanitize_infinities(self) -> Constraint: self._sign = self._sign[keep] return self - def freeze(self) -> Constraint: + def freeze(self) -> CSRConstraint: """Return self (already immutable).""" return self @@ -941,7 +941,7 @@ def from_mutable( cls, con: MutableConstraint, cindex: int | None = None, - ) -> Constraint: + ) -> CSRConstraint: """ Create a Constraint from a MutableConstraint. @@ -1162,14 +1162,14 @@ def sanitize_infinities(self) -> MutableConstraint: self._data = assign_multiindex_safe(self.data, labels=labels) return self - def freeze(self) -> Constraint: + def freeze(self) -> CSRConstraint: """Convert to an immutable Constraint.""" cindex = ( int(self.data.attrs["label_range"][0]) if "label_range" in self.data.attrs else None ) - return Constraint.from_mutable(self, cindex=cindex) + return CSRConstraint.from_mutable(self, cindex=cindex) def mutable(self) -> MutableConstraint: """Return self (already mutable).""" @@ -1709,9 +1709,9 @@ def reset_dual(self) -> None: Reset the stored solution of variables. """ for k, c in self.items(): - if isinstance(c, Constraint): + if isinstance(c, CSRConstraint): if c._dual is not None: - self.data[k] = Constraint( + self.data[k] = CSRConstraint( c._csr, c._con_labels, c._rhs, diff --git a/linopy/io.py b/linopy/io.py index e2dfa9a2..610927c7 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1171,9 +1171,9 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: m : linopy.Model """ from linopy.constraints import ( - Constraint, ConstraintBase, Constraints, + CSRConstraint, MutableConstraint, ) from linopy.expressions import LinearExpression @@ -1230,7 +1230,7 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: name = remove_prefix(k, "constraints") con_ds = get_prefix(ds, k) if con_ds.attrs.get("_linopy_format") == "csr": - constraints[name] = Constraint.from_netcdf_ds(con_ds, m, name) + constraints[name] = CSRConstraint.from_netcdf_ds(con_ds, m, name) else: constraints[name] = MutableConstraint(con_ds, m, name) m._constraints = Constraints(constraints, m) diff --git a/linopy/matrices.py b/linopy/matrices.py index 9496c2a7..1fb59344 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -14,7 +14,7 @@ from numpy import ndarray from linopy import expressions -from linopy.constraints import Constraint +from linopy.constraints import CSRConstraint if TYPE_CHECKING: from linopy.model import Model @@ -158,7 +158,7 @@ def dual(self) -> ndarray: dual_list = [] has_dual = False for c in m.constraints.data.values(): - if isinstance(c, Constraint): + if isinstance(c, CSRConstraint): # _dual is active-only if c._dual is not None: dual_list.append(c._dual) diff --git a/linopy/model.py b/linopy/model.py index 6c2d639c..a381c56a 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -49,9 +49,9 @@ ) from linopy.constraints import ( AnonymousScalarConstraint, - Constraint, ConstraintBase, Constraints, + CSRConstraint, MutableConstraint, ) from linopy.expressions import ( @@ -740,7 +740,7 @@ def add_constraints( coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = ..., mask: MaskLike | None = ..., freeze: Literal[True] = ..., - ) -> Constraint: ... + ) -> CSRConstraint: ... @overload def add_constraints( diff --git a/test/test_constraint.py b/test/test_constraint.py index 091e4e54..bf73fae6 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -62,12 +62,12 @@ def mc(m: Model) -> linopy.constraints.MutableConstraint: return m.constraints["c"].mutable() -def test_constraint_repr(c: linopy.constraints.Constraint) -> None: +def test_constraint_repr(c: linopy.constraints.CSRConstraint) -> None: c.__repr__() def test_constraint_repr_equivalent_to_mutable( - c: linopy.constraints.Constraint, + c: linopy.constraints.CSRConstraint, ) -> None: """Constraint (CSR-backed) and MutableConstraint repr must be identical.""" frozen = c.freeze() @@ -80,12 +80,12 @@ def test_constraints_repr(m: Model) -> None: def test_add_constraints_freeze(m: Model, x: linopy.Variable) -> None: c = m.add_constraints(x >= 1, name="frozen_c", freeze=True) - assert isinstance(c, linopy.constraints.Constraint) - assert isinstance(m.constraints["frozen_c"], linopy.constraints.Constraint) + assert isinstance(c, linopy.constraints.CSRConstraint) + assert isinstance(m.constraints["frozen_c"], linopy.constraints.CSRConstraint) assert c.ncons == 10 -def test_constraint_name(c: linopy.constraints.Constraint) -> None: +def test_constraint_name(c: linopy.constraints.CSRConstraint) -> None: assert c.name == "c" @@ -100,7 +100,7 @@ def test_cannot_create_constraint_without_variable() -> None: _ = linopy.LinearExpression(12, model) == linopy.LinearExpression(13, model) -def test_constraints_getter(m: Model, c: linopy.constraints.Constraint) -> None: +def test_constraints_getter(m: Model, c: linopy.constraints.CSRConstraint) -> None: assert c.shape == (10,) assert isinstance(m.constraints[["c"]], Constraints) @@ -329,11 +329,11 @@ def test_constraint_coeffs_getter(mc: linopy.constraints.MutableConstraint) -> N assert (mc.coeffs == 1).all() -def test_constraint_sign_getter(c: linopy.constraints.Constraint) -> None: +def test_constraint_sign_getter(c: linopy.constraints.CSRConstraint) -> None: assert (c.sign == GREATER_EQUAL).all() -def test_constraint_rhs_getter(c: linopy.constraints.Constraint) -> None: +def test_constraint_rhs_getter(c: linopy.constraints.CSRConstraint) -> None: assert (c.rhs == 0).all() @@ -443,18 +443,18 @@ def test_constraint_rhs_setter_with_expression_and_constant( assert mc.lhs.nterm == 2 -def test_constraint_labels_setter_invalid(c: linopy.constraints.Constraint) -> None: +def test_constraint_labels_setter_invalid(c: linopy.constraints.CSRConstraint) -> None: # Test that assigning labels raises AttributeError (Constraint is frozen) with pytest.raises(AttributeError): c.labels = c.labels # type: ignore -def test_constraint_sel(c: linopy.constraints.Constraint) -> None: +def test_constraint_sel(c: linopy.constraints.CSRConstraint) -> None: assert isinstance(c.mutable().sel(first=[1, 2]), ConstraintBase) assert isinstance(c.mutable().isel(first=[1, 2]), ConstraintBase) -def test_constraint_flat(c: linopy.constraints.Constraint) -> None: +def test_constraint_flat(c: linopy.constraints.CSRConstraint) -> None: assert isinstance(c.flat, pd.DataFrame) @@ -464,7 +464,7 @@ def test_iterate_slices(mc: linopy.constraints.MutableConstraint) -> None: assert mc.coord_dims == i.coord_dims -def test_constraint_to_polars(c: linopy.constraints.Constraint) -> None: +def test_constraint_to_polars(c: linopy.constraints.CSRConstraint) -> None: assert isinstance(c.to_polars(), pl.DataFrame) @@ -709,10 +709,10 @@ def test_constraints_equalities(m: Model) -> None: def test_freeze_mutable_roundtrip(m: Model) -> None: frozen = m.constraints["c"] - assert isinstance(frozen, linopy.constraints.Constraint) + assert isinstance(frozen, linopy.constraints.CSRConstraint) mc = frozen.mutable() assert isinstance(mc, MutableConstraint) - refrozen = linopy.constraints.Constraint.from_mutable(mc, frozen._cindex) + refrozen = linopy.constraints.CSRConstraint.from_mutable(mc, frozen._cindex) assert_equal(frozen.labels, refrozen.labels) assert_equal(frozen.rhs, refrozen.rhs) assert_equal(frozen.sign, refrozen.sign) @@ -727,7 +727,7 @@ def test_freeze_mutable_roundtrip_with_masking() -> None: m.add_constraints(x.where(mask) >= 0, name="c") frozen = m.constraints["c"] mc = frozen.mutable() - refrozen = linopy.constraints.Constraint.from_mutable(mc, frozen._cindex) + refrozen = linopy.constraints.CSRConstraint.from_mutable(mc, frozen._cindex) assert_equal(frozen.labels, refrozen.labels) assert_equal(frozen.rhs, refrozen.rhs) assert frozen.ncons == refrozen.ncons == 3 @@ -740,7 +740,7 @@ def test_from_mutable_mixed_signs() -> None: mc = m.constraints["mixed"] assert isinstance(mc, MutableConstraint) mc._data["sign"] = xr.DataArray(["<=", ">=", "<="], dims=["i"]) - frozen = linopy.constraints.Constraint.from_mutable(mc) + frozen = linopy.constraints.CSRConstraint.from_mutable(mc) assert isinstance(frozen._sign, np.ndarray) assert list(frozen._sign) == ["<=", ">=", "<="] assert_equal(frozen.sign, mc.sign) @@ -801,7 +801,7 @@ def bound(m, i): return x.at[i] == 0.0 con = m.add_constraints(bound, coords=coords, name="mixed_rule") - assert isinstance(con, linopy.constraints.Constraint) + assert isinstance(con, linopy.constraints.CSRConstraint) assert isinstance(con._sign, np.ndarray) assert con.ncons == 4 expected_signs = ["=", ">=", "=", ">="] @@ -814,7 +814,7 @@ def test_frozen_rhs_setter() -> None: time = pd.RangeIndex(5, name="t") x = m.add_variables(lower=0, coords=[time], name="x") con = m.add_constraints(x >= 1, name="c") - assert isinstance(con, linopy.constraints.Constraint) + assert isinstance(con, linopy.constraints.CSRConstraint) con.rhs = 10 np.testing.assert_array_equal(con._rhs, np.full(5, 10.0)) factor = pd.Series(range(5), index=time) @@ -828,7 +828,7 @@ def test_frozen_lhs_setter() -> None: x = m.add_variables(lower=0, coords=[time], name="x") y = m.add_variables(lower=0, coords=[time], name="y") con = m.add_constraints(x >= 0, name="c") - assert isinstance(con, linopy.constraints.Constraint) + assert isinstance(con, linopy.constraints.CSRConstraint) con.lhs = 3 * x + 2 * y lhs = con.mutable().lhs assert lhs.nterm == 2 diff --git a/test/test_io.py b/test/test_io.py index 1a97bd22..09b19fdc 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -79,24 +79,24 @@ def test_model_to_netcdf(model: Model, tmp_path: Path) -> None: def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None: - from linopy.constraints import Constraint + from linopy.constraints import CSRConstraint m = Model() x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") m.add_constraints(x >= 1, name="c", freeze=True) - assert isinstance(m.constraints["c"], Constraint) + assert isinstance(m.constraints["c"], CSRConstraint) fn = tmp_path / "test_frozen.nc" m.to_netcdf(fn) p = read_netcdf(fn) - assert isinstance(p.constraints["c"], Constraint) + assert isinstance(p.constraints["c"], CSRConstraint) assert_model_equal(m, p) def test_model_to_netcdf_mixed_sign_constraint(tmp_path: Path) -> None: - from linopy.constraints import Constraint + from linopy.constraints import CSRConstraint m = Model() x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") @@ -107,13 +107,13 @@ def bound(m, i): return x.at[i] == 0.0 m.add_constraints(bound, coords=[pd.RangeIndex(4, name="i")], name="c") - assert isinstance(m.constraints["c"], Constraint) + assert isinstance(m.constraints["c"], CSRConstraint) fn = tmp_path / "test_mixed_sign.nc" m.to_netcdf(fn) p = read_netcdf(fn) - assert isinstance(p.constraints["c"], Constraint) + assert isinstance(p.constraints["c"], CSRConstraint) import numpy as np np.testing.assert_array_equal(m.constraints["c"]._sign, p.constraints["c"]._sign) diff --git a/test/test_repr.py b/test/test_repr.py index 9a7af893..0e854db0 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -6,7 +6,7 @@ import xarray as xr from linopy import Model, options -from linopy.constraints import Constraint +from linopy.constraints import MutableConstraint from linopy.expressions import LinearExpression from linopy.variables import Variable @@ -139,7 +139,7 @@ def test_single_array_linear_repr(var: Variable) -> None: @pytest.mark.parametrize("con", anonymous_constraints) -def test_anonymous_constraint_repr(con: Constraint) -> None: +def test_anonymous_constraint_repr(con: MutableConstraint) -> None: repr(con) @@ -162,7 +162,7 @@ def test_single_array_constraint_repr(var: Variable) -> None: @pytest.mark.parametrize("con", constraints) -def test_constraint_repr(con: Constraint) -> None: +def test_constraint_repr(con: MutableConstraint) -> None: repr(con) @@ -173,7 +173,7 @@ def test_empty_repr() -> None: @pytest.mark.parametrize("obj", [v, lv, cv_, cv]) -def test_print_options(obj: Variable | LinearExpression | Constraint) -> None: +def test_print_options(obj: Variable | LinearExpression | MutableConstraint) -> None: default_repr = repr(obj) with options as opts: opts.set_value(display_max_rows=20) From a256212457473b7bbf1416f3602ae28c4b37e8d4 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 26 Mar 2026 09:33:28 +0100 Subject: [PATCH 19/20] rename MutableConstraint to Constraint (original name) --- linopy/__init__.py | 4 ++-- linopy/constraints.py | 42 +++++++++++++++++++---------------------- linopy/expressions.py | 30 ++++++++++++----------------- linopy/io.py | 8 ++++---- linopy/model.py | 8 ++++---- linopy/variables.py | 26 ++++++++++--------------- test/test_constraint.py | 14 +++++++------- test/test_repr.py | 8 ++++---- 8 files changed, 62 insertions(+), 78 deletions(-) diff --git a/linopy/__init__.py b/linopy/__init__.py index 3236d64c..d9f4b64e 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -16,10 +16,10 @@ from linopy.config import options from linopy.constants import EQUAL, GREATER_EQUAL, LESS_EQUAL, PerformanceWarning from linopy.constraints import ( + Constraint, ConstraintBase, Constraints, CSRConstraint, - MutableConstraint, ) from linopy.expressions import LinearExpression, QuadraticExpression, merge from linopy.io import read_netcdf @@ -37,7 +37,7 @@ "CSRConstraint", "ConstraintBase", "Constraints", - "MutableConstraint", + "Constraint", "EQUAL", "PerformanceWarning", "GREATER_EQUAL", diff --git a/linopy/constraints.py b/linopy/constraints.py index c07c8db1..776d8b12 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -87,9 +87,7 @@ def conwrap( method: Callable, *default_args: Any, **new_default_kwargs: Any ) -> Callable: @functools.wraps(method) - def _conwrap( - con: MutableConstraint, *args: Any, **kwargs: Any - ) -> MutableConstraint: + def _conwrap(con: Constraint, *args: Any, **kwargs: Any) -> Constraint: for k, v in new_default_kwargs.items(): kwargs.setdefault(k, v) return con.__class__( @@ -199,7 +197,7 @@ def freeze(self) -> CSRConstraint: """Return an immutable Constraint (CSR-backed).""" @abstractmethod - def mutable(self) -> MutableConstraint: + def mutable(self) -> Constraint: """Return a mutable MutableConstraint.""" @abstractmethod @@ -215,7 +213,7 @@ def to_matrix_with_rhs( def __getitem__( self, selector: str | int | slice | list | tuple | dict - ) -> MutableConstraint: + ) -> Constraint: """ Get selection from the constraint. Returns a MutableConstraint with the selected data. @@ -223,7 +221,7 @@ def __getitem__( data = Dataset( {k: self.data[k][selector] for k in self.data}, attrs=self.data.attrs ) - return MutableConstraint(data, self.model, self.name) + return Constraint(data, self.model, self.name) @property def attrs(self) -> dict[str, Any]: @@ -674,7 +672,7 @@ def lhs(self) -> expressions.LinearExpression: def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: self._refreeze_after(lambda mc: setattr(mc, "lhs", value)) - def _refreeze_after(self, mutate: Callable[[MutableConstraint], None]) -> None: + def _refreeze_after(self, mutate: Callable[[Constraint], None]) -> None: mc = self.mutable() mutate(mc) refrozen = CSRConstraint.from_mutable(mc, self._cindex) @@ -928,9 +926,9 @@ def freeze(self) -> CSRConstraint: """Return self (already immutable).""" return self - def mutable(self) -> MutableConstraint: + def mutable(self) -> Constraint: """Convert to a MutableConstraint.""" - return MutableConstraint(self.data, self._model, self._name) + return Constraint(self.data, self._model, self._name) def to_polars(self) -> pl.DataFrame: """Convert to polars DataFrame — delegates to mutable().""" @@ -939,7 +937,7 @@ def to_polars(self) -> pl.DataFrame: @classmethod def from_mutable( cls, - con: MutableConstraint, + con: Constraint, cindex: int | None = None, ) -> CSRConstraint: """ @@ -985,7 +983,7 @@ def from_mutable( ) -class MutableConstraint(ConstraintBase): +class Constraint(ConstraintBase): """ Mutable constraint backed by an xarray Dataset. @@ -1139,21 +1137,21 @@ def to_matrix_with_rhs( sense = np.array([s[0] for s in sign_flat]) return csr, con_labels, b, sense - def sanitize_zeros(self) -> MutableConstraint: + def sanitize_zeros(self) -> Constraint: """Remove terms with zero or near-zero coefficients.""" not_zero = abs(self.coeffs) > 1e-10 self.vars = self.vars.where(not_zero, -1) self.coeffs = self.coeffs.where(not_zero) return self - def sanitize_missings(self) -> MutableConstraint: + def sanitize_missings(self) -> Constraint: """Mask out rows where all variables are missing (-1).""" contains_non_missing = (self.vars != -1).any(self.term_dim) labels = self.labels.where(contains_non_missing, -1) self._data = assign_multiindex_safe(self.data, labels=labels) return self - def sanitize_infinities(self) -> MutableConstraint: + def sanitize_infinities(self) -> Constraint: """Mask out rows with invalid infinite RHS values.""" valid_infinity_values = ((self.sign == LESS_EQUAL) & (self.rhs == np.inf)) | ( (self.sign == GREATER_EQUAL) & (self.rhs == -np.inf) @@ -1171,14 +1169,12 @@ def freeze(self) -> CSRConstraint: ) return CSRConstraint.from_mutable(self, cindex=cindex) - def mutable(self) -> MutableConstraint: + def mutable(self) -> Constraint: """Return self (already mutable).""" return self @classmethod - def from_rule( - cls, model: Model, rule: Callable, coords: CoordsLike - ) -> MutableConstraint: + def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constraint: """ Create a constraint from a rule and a set of coordinates. @@ -1422,7 +1418,7 @@ def add(self, constraint: ConstraintBase, freeze: bool = False) -> ConstraintBas """ Add a constraint to the constraints container. """ - if freeze and isinstance(constraint, MutableConstraint): + if freeze and isinstance(constraint, Constraint): constraint = constraint.freeze() self.data[constraint.name] = constraint self._invalidate_label_position_index() @@ -1639,7 +1635,7 @@ def set_blocks(self, block_map: np.ndarray) -> None: N = block_map.max() for name, constraint in self.items(): - if not isinstance(constraint, MutableConstraint): + if not isinstance(constraint, Constraint): self.data[name] = constraint = constraint.mutable() res = xr.full_like(constraint.labels, N + 1, dtype=block_map.dtype) entries = replace_by_map(constraint.vars, block_map) @@ -1722,7 +1718,7 @@ def reset_dual(self) -> None: cindex=c._cindex, dual=None, ) - elif isinstance(c, MutableConstraint): + elif isinstance(c, Constraint): if "dual" in c.data: c._data = c.data.drop_vars("dual") else: @@ -1787,6 +1783,6 @@ def rhs(self) -> int | float | np.floating | np.integer: """ return self._rhs - def to_constraint(self) -> MutableConstraint: + def to_constraint(self) -> Constraint: data = self.lhs.to_linexpr().data.assign(sign=self.sign, rhs=self.rhs) - return MutableConstraint(data=data, model=self.lhs.model) + return Constraint(data=data, model=self.lhs.model) diff --git a/linopy/expressions.py b/linopy/expressions.py index 58325829..6a46834b 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -93,7 +93,7 @@ if TYPE_CHECKING: from linopy.constraints import ( AnonymousScalarConstraint, - MutableConstraint, + Constraint, ) from linopy.model import Model from linopy.piecewise import PiecewiseConstraintDescriptor, PiecewiseExpression @@ -679,11 +679,9 @@ def __truediv__(self: GenericExpression, other: SideLike) -> GenericExpression: def __le__(self, rhs: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __le__(self, rhs: SideLike) -> MutableConstraint: ... + def __le__(self, rhs: SideLike) -> Constraint: ... - def __le__( - self, rhs: SideLike - ) -> MutableConstraint | PiecewiseConstraintDescriptor: + def __le__(self, rhs: SideLike) -> Constraint | PiecewiseConstraintDescriptor: descriptor = _to_piecewise_constraint_descriptor(self, rhs, "<=") if descriptor is not None: return descriptor @@ -693,11 +691,9 @@ def __le__( def __ge__(self, rhs: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __ge__(self, rhs: SideLike) -> MutableConstraint: ... + def __ge__(self, rhs: SideLike) -> Constraint: ... - def __ge__( - self, rhs: SideLike - ) -> MutableConstraint | PiecewiseConstraintDescriptor: + def __ge__(self, rhs: SideLike) -> Constraint | PiecewiseConstraintDescriptor: descriptor = _to_piecewise_constraint_descriptor(self, rhs, ">=") if descriptor is not None: return descriptor @@ -707,11 +703,9 @@ def __ge__( def __eq__(self, rhs: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __eq__(self, rhs: SideLike) -> MutableConstraint: ... + def __eq__(self, rhs: SideLike) -> Constraint: ... - def __eq__( - self, rhs: SideLike - ) -> MutableConstraint | PiecewiseConstraintDescriptor: + def __eq__(self, rhs: SideLike) -> Constraint | PiecewiseConstraintDescriptor: descriptor = _to_piecewise_constraint_descriptor(self, rhs, "==") if descriptor is not None: return descriptor @@ -830,7 +824,7 @@ def le( self: GenericExpression, rhs: SideLike, join: JoinOptions | None = None, - ) -> MutableConstraint: + ) -> Constraint: """ Less than or equal constraint. @@ -849,7 +843,7 @@ def ge( self, rhs: SideLike, join: JoinOptions | None = None, - ) -> MutableConstraint: + ) -> Constraint: """ Greater than or equal constraint. @@ -868,7 +862,7 @@ def eq( self, rhs: SideLike, join: JoinOptions | None = None, - ) -> MutableConstraint: + ) -> Constraint: """ Equality constraint. @@ -1124,7 +1118,7 @@ def cumsum( def to_constraint( self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None - ) -> MutableConstraint: + ) -> Constraint: """ Convert a linear expression to a constraint. @@ -1184,7 +1178,7 @@ def to_constraint( data = assign_multiindex_safe( all_to_lhs[["coeffs", "vars"]], sign=sign, rhs=computed_rhs ) - return constraints.MutableConstraint(data, model=self.model) + return constraints.Constraint(data, model=self.model) def reset_const(self: GenericExpression) -> GenericExpression: """ diff --git a/linopy/io.py b/linopy/io.py index 610927c7..715d1a3b 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1171,10 +1171,10 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: m : linopy.Model """ from linopy.constraints import ( + Constraint, ConstraintBase, Constraints, CSRConstraint, - MutableConstraint, ) from linopy.expressions import LinearExpression from linopy.model import Model @@ -1232,7 +1232,7 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: if con_ds.attrs.get("_linopy_format") == "csr": constraints[name] = CSRConstraint.from_netcdf_ds(con_ds, m, name) else: - constraints[name] = MutableConstraint(con_ds, m, name) + constraints[name] = Constraint(con_ds, m, name) m._constraints = Constraints(constraints, m) objective = get_prefix(ds, "objective") @@ -1288,7 +1288,7 @@ def copy(m: Model, include_solution: bool = False, deep: bool = True) -> Model: Model A deep or shallow copy of the model. """ - from linopy.constraints import Constraints, MutableConstraint + from linopy.constraints import Constraint, Constraints from linopy.expressions import LinearExpression from linopy.model import Model, Objective from linopy.variables import Variable, Variables @@ -1318,7 +1318,7 @@ def copy(m: Model, include_solution: bool = False, deep: bool = True) -> Model: new_model._constraints = Constraints( { - name: MutableConstraint( + name: Constraint( con.mutable().data.copy(deep=deep) if include_solution else con.mutable().data[m.constraints.dataset_attrs].copy(deep=deep), diff --git a/linopy/model.py b/linopy/model.py index a381c56a..748a7b60 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -49,10 +49,10 @@ ) from linopy.constraints import ( AnonymousScalarConstraint, + Constraint, ConstraintBase, Constraints, CSRConstraint, - MutableConstraint, ) from linopy.expressions import ( LinearExpression, @@ -756,7 +756,7 @@ def add_constraints( coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = ..., mask: MaskLike | None = ..., freeze: Literal[False] = ..., - ) -> MutableConstraint: ... + ) -> Constraint: ... def add_constraints( self, @@ -848,7 +848,7 @@ def add_constraints( rule = lhs if sign is not None or rhs is not None: raise ValueError(msg_sign_rhs_none) - data = MutableConstraint.from_rule(self, rule, coords).data + data = Constraint.from_rule(self, rule, coords).data elif isinstance(lhs, AnonymousScalarConstraint): if sign is not None or rhs is not None: raise ValueError(msg_sign_rhs_none) @@ -930,7 +930,7 @@ def add_constraints( if self.chunk: data = data.chunk(self.chunk) - constraint = MutableConstraint(data, name=name, model=self, skip_broadcast=True) + constraint = Constraint(data, name=name, model=self, skip_broadcast=True) return self.constraints.add(constraint, freeze=freeze and not self.chunk) def add_objective( diff --git a/linopy/variables.py b/linopy/variables.py index 341a0cb5..6cdaa83f 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -72,7 +72,7 @@ ) if TYPE_CHECKING: - from linopy.constraints import AnonymousScalarConstraint, MutableConstraint + from linopy.constraints import AnonymousScalarConstraint, Constraint from linopy.expressions import ( GenericExpression, LinearExpression, @@ -543,33 +543,27 @@ def __rsub__(self, other: ConstantLike) -> LinearExpression: def __le__(self, other: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __le__(self, other: SideLike) -> MutableConstraint: ... + def __le__(self, other: SideLike) -> Constraint: ... - def __le__( - self, other: SideLike - ) -> MutableConstraint | PiecewiseConstraintDescriptor: + def __le__(self, other: SideLike) -> Constraint | PiecewiseConstraintDescriptor: return self.to_linexpr().__le__(other) @overload def __ge__(self, other: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __ge__(self, other: SideLike) -> MutableConstraint: ... + def __ge__(self, other: SideLike) -> Constraint: ... - def __ge__( - self, other: SideLike - ) -> MutableConstraint | PiecewiseConstraintDescriptor: + def __ge__(self, other: SideLike) -> Constraint | PiecewiseConstraintDescriptor: return self.to_linexpr().__ge__(other) @overload # type: ignore[override] def __eq__(self, other: PiecewiseExpression) -> PiecewiseConstraintDescriptor: ... @overload - def __eq__(self, other: SideLike) -> MutableConstraint: ... + def __eq__(self, other: SideLike) -> Constraint: ... - def __eq__( - self, other: SideLike - ) -> MutableConstraint | PiecewiseConstraintDescriptor: + def __eq__(self, other: SideLike) -> Constraint | PiecewiseConstraintDescriptor: return self.to_linexpr().__eq__(other) def __gt__(self, other: Any) -> NotImplementedType: @@ -653,7 +647,7 @@ def div( """ return self.to_linexpr().div(other, join=join) - def le(self, rhs: SideLike, join: JoinOptions | None = None) -> MutableConstraint: + def le(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: """ Less than or equal constraint. @@ -668,7 +662,7 @@ def le(self, rhs: SideLike, join: JoinOptions | None = None) -> MutableConstrain """ return self.to_linexpr().le(rhs, join=join) - def ge(self, rhs: SideLike, join: JoinOptions | None = None) -> MutableConstraint: + def ge(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: """ Greater than or equal constraint. @@ -683,7 +677,7 @@ def ge(self, rhs: SideLike, join: JoinOptions | None = None) -> MutableConstrain """ return self.to_linexpr().ge(rhs, join=join) - def eq(self, rhs: SideLike, join: JoinOptions | None = None) -> MutableConstraint: + def eq(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: """ Equality constraint. diff --git a/test/test_constraint.py b/test/test_constraint.py index bf73fae6..e1b9f4b6 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -26,9 +26,9 @@ ) from linopy.constraints import ( AnonymousScalarConstraint, + Constraint, ConstraintBase, Constraints, - MutableConstraint, ) @@ -226,7 +226,7 @@ def test_constraint_inherited_properties( def test_constraint_wrapped_methods(x: linopy.Variable, y: linopy.Variable) -> None: - con: MutableConstraint = 10 * x + y <= 10 + con: Constraint = 10 * x + y <= 10 # Test wrapped methods con.assign({"new_var": xr.DataArray(np.zeros((2, 2)), coords=[range(2), range(2)])}) @@ -295,7 +295,7 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint: return (i - 1) * x.at[i - 1] + y.at[j] >= 0 if i % 2 else i * x.at[i] >= 0 coords = [x.coords["first"], y.coords["second"]] - con = MutableConstraint.from_rule(m, bound, coords) + con = Constraint.from_rule(m, bound, coords) assert isinstance(con, ConstraintBase) assert con.lhs.nterm == 2 repr(con) # test repr @@ -310,7 +310,7 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint | None: return None coords = [x.coords["first"], y.coords["second"]] - con = MutableConstraint.from_rule(m, bound, coords) + con = Constraint.from_rule(m, bound, coords) assert isinstance(con, ConstraintBase) assert isinstance(con.lhs.vars, xr.DataArray) assert con.lhs.nterm == 2 @@ -615,7 +615,7 @@ def test_constraint_with_helper_dims_as_coords(m: Model) -> None: data = xr.Dataset({"coeffs": coeffs, "vars": vars, "sign": sign, "rhs": rhs}) assert set(HELPER_DIMS).intersection(set(data.coords)) - con = MutableConstraint(data, m, "c") + con = Constraint(data, m, "c") expr = m.add_constraints(con) assert not set(HELPER_DIMS).intersection(set(expr.data.coords)) @@ -711,7 +711,7 @@ def test_freeze_mutable_roundtrip(m: Model) -> None: frozen = m.constraints["c"] assert isinstance(frozen, linopy.constraints.CSRConstraint) mc = frozen.mutable() - assert isinstance(mc, MutableConstraint) + assert isinstance(mc, Constraint) refrozen = linopy.constraints.CSRConstraint.from_mutable(mc, frozen._cindex) assert_equal(frozen.labels, refrozen.labels) assert_equal(frozen.rhs, refrozen.rhs) @@ -738,7 +738,7 @@ def test_from_mutable_mixed_signs() -> None: x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") m.add_constraints(x >= 0, name="mixed", freeze=False) mc = m.constraints["mixed"] - assert isinstance(mc, MutableConstraint) + assert isinstance(mc, Constraint) mc._data["sign"] = xr.DataArray(["<=", ">=", "<="], dims=["i"]) frozen = linopy.constraints.CSRConstraint.from_mutable(mc) assert isinstance(frozen._sign, np.ndarray) diff --git a/test/test_repr.py b/test/test_repr.py index 0e854db0..9a7af893 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -6,7 +6,7 @@ import xarray as xr from linopy import Model, options -from linopy.constraints import MutableConstraint +from linopy.constraints import Constraint from linopy.expressions import LinearExpression from linopy.variables import Variable @@ -139,7 +139,7 @@ def test_single_array_linear_repr(var: Variable) -> None: @pytest.mark.parametrize("con", anonymous_constraints) -def test_anonymous_constraint_repr(con: MutableConstraint) -> None: +def test_anonymous_constraint_repr(con: Constraint) -> None: repr(con) @@ -162,7 +162,7 @@ def test_single_array_constraint_repr(var: Variable) -> None: @pytest.mark.parametrize("con", constraints) -def test_constraint_repr(con: MutableConstraint) -> None: +def test_constraint_repr(con: Constraint) -> None: repr(con) @@ -173,7 +173,7 @@ def test_empty_repr() -> None: @pytest.mark.parametrize("obj", [v, lv, cv_, cv]) -def test_print_options(obj: Variable | LinearExpression | MutableConstraint) -> None: +def test_print_options(obj: Variable | LinearExpression | Constraint) -> None: default_repr = repr(obj) with options as opts: opts.set_value(display_max_rows=20) From 3ead682d5425fb2d98c00ea67fbf068b6d46311f Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 26 Mar 2026 09:49:32 +0100 Subject: [PATCH 20/20] fix: xpress crash with zero constraints and remove_variables not removing variable without referencing constraints --- linopy/model.py | 2 +- linopy/solvers.py | 31 +++++++++++++++++-------------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/linopy/model.py b/linopy/model.py index 748a7b60..040c6762 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1009,7 +1009,7 @@ def remove_variables(self, name: str) -> None: for k in to_remove: self.constraints.remove(k) - self.variables.remove(name) + self.variables.remove(name) self.objective = self.objective.sel( {TERM_DIM: ~self.objective.vars.isin(variable.labels)} diff --git a/linopy/solvers.py b/linopy/solvers.py index a8e950ce..8309113b 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -1721,20 +1721,23 @@ def get_solver_solution() -> Solution: sol = pd.Series(m.getSolution(), index=var, dtype=float) try: - try: # Try new API first - _dual = m.getDuals() - except AttributeError: # Fallback to old API - _dual = m.getDual() - - try: # Try new API first - constraints = m.getNameList( - xpress_Namespaces.ROW, 0, m.attributes.rows - 1 - ) - except AttributeError: # Fallback to old API - constraints = m.getnamelist( - xpress_Namespaces.ROW, 0, m.attributes.rows - 1 - ) - dual = pd.Series(_dual, index=constraints, dtype=float) + if m.attributes.rows == 0: + dual = pd.Series(dtype=float) + else: + try: # Try new API first + _dual = m.getDuals() + except AttributeError: # Fallback to old API + _dual = m.getDual() + + try: # Try new API first + constraints = m.getNameList( + xpress_Namespaces.ROW, 0, m.attributes.rows - 1 + ) + except AttributeError: # Fallback to old API + constraints = m.getnamelist( + xpress_Namespaces.ROW, 0, m.attributes.rows - 1 + ) + dual = pd.Series(_dual, index=constraints, dtype=float) except (xpress.SolverError, xpress.ModelError, SystemError): logger.warning("Dual values of MILP couldn't be parsed") dual = pd.Series(dtype=float)