diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 42e8bee4..92e4ae28 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -73,6 +73,12 @@ Upcoming Version * ``LinearExpression.flat`` and ``LinearExpression.to_polars`` on a CSR-backed expression are emitted directly from the sparse backing instead of expanding the dense rectangle; the rows equal those of the dense path, with absent cells and zero coefficients dropped. (`#968 `__) * A CSR-backed linear objective stays sparse when set with ``Model.add_objective`` and when the model is exported or solved: the objective vector ``matrices.c``, LP/MPS files, netcdf output, ``Model.copy``, persistent snapshots and direct solver APIs all read it without expanding the dense rectangle, with results equal to the dense objective. The objective name is now owned by the ``Objective`` itself (``Objective.name`` and ``Objective.attrs``) instead of being written into the expression's attributes. (`#967 `__) * Persistent snapshots of tz-aware ``DatetimeIndex`` coordinates no longer materialise an object array of ``Timestamp`` per container per capture and diff. Coordinates are stored as UTC-ns arrays with the timezone identity carried alongside, making snapshot capture ~24x and warm-start diffs ~33x faster on tz-aware models, while naive and tz-aware coordinates — and differing timezones — stay correctly unequal. (`#960 `__) +* The MOSEK direct API builds bound keys, bounds and constraint names vectorised instead of looping over every row and column in Python. +* The LP writer computes the scaling lookups once per file, skips scaling when every factor is 1, does not re-sort rows that are already grouped (as frozen constraints deliver them) and builds each constraint line from fewer string branches. Constraint writing is about 2x faster; the output is byte-identical. +* ``linopy.merge`` of CSR-backed expressions concatenates all operands once instead of folding them pairwise, so an N-way merge scales with the total number of nonzeros instead of quadratically with the operand count. +* ``LinearExpression.solution`` and ``QuadraticExpression.solution`` no longer rebuild the full constraint matrix; the solution is mapped by variable label, and a CSR-backed expression is evaluated on its sparse backing without densifying. Roughly 10x faster and 6x less memory on a large model. ``Model.dualize`` reads ``model.matrices`` once instead of four times. +* ``model.matrices`` assembles the constraint blocks into preallocated CSR buffers in one pass, without ``vstack``, ``eliminate_zeros`` or a full-grid scaling round trip; scaling is skipped when it is identity. Label lookups use the model's label dtype, so ``matrices.clabels`` and ``matrices.indicator_binvar`` are now ``int32`` by default. On an 8M-row model this is 45% faster with 200 MB less peak memory, with bit-identical output. +* Duals of frozen constraints are read back on the active rows directly instead of through three full-grid arrays, and freezing a masked constraint gathers rows once instead of up to three times. Mixed-sign constraints convert their sense vectorised instead of in a Python loop. **Bug fixes** diff --git a/linopy/common.py b/linopy/common.py index 79291050..16b1da0c 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -930,13 +930,16 @@ def label_to_pos(self) -> np.ndarray: """ Mapping from variable label to dense position, shape (_xCounter,). + Positions share the model's label dtype, since they never exceed a label. + 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) - label_to_pos[vlabels] = np.arange(len(vlabels), dtype=np.intp) + model = self._variables.model + dtype = model._dtypes["labels"] + label_to_pos = np.full(model._xCounter, -1, dtype=dtype) + label_to_pos[vlabels] = np.arange(len(vlabels), dtype=dtype) return label_to_pos @property @@ -969,17 +972,17 @@ def clabels(self) -> np.ndarray: for c in self._constraints.data.values() if not c.is_indicator ] - return ( - np.concatenate(label_lists) if label_lists else np.array([], dtype=np.intp) - ) + dtype = self._constraints.model._dtypes["labels"] + return np.concatenate([np.array([], dtype=dtype), *label_lists], dtype=dtype) @cached_property def label_to_pos(self) -> np.ndarray: """Mapping from constraint label to dense position, shape (_cCounter,).""" clabels = self.clabels - n = self._constraints.model._cCounter - label_to_pos = np.full(n, -1, dtype=np.intp) - label_to_pos[clabels] = np.arange(len(clabels), dtype=np.intp) + model = self._constraints.model + dtype = model._dtypes["labels"] + label_to_pos = np.full(model._cCounter, -1, dtype=dtype) + label_to_pos[clabels] = np.arange(len(clabels), dtype=dtype) return label_to_pos @property diff --git a/linopy/constraints.py b/linopy/constraints.py index 320c863d..ef8c99b4 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -939,31 +939,39 @@ def _replace(self, **changes: Any) -> CSRConstraint: return new def assign_labels( - self, cindex: int, name: str, scaling: float | DataArray = 1.0 + self, + cindex: int, + name: str, + scaling: float | DataArray = 1.0, + mask: np.ndarray | None = None, ) -> CSRConstraint: """ Return a copy labelled from ``cindex`` and named ``name``. Rows without terms are dropped, as when freezing a dense constraint; - a zero coefficient counts as a term. ``scaling`` is a scalar or a row - scaling broadcast on the grid; its distinct values are validated - without expanding a broadcast view. + a zero coefficient counts as a term. Active rows where the boolean + ``mask`` is False are dropped in the same gather. ``scaling`` is a + scalar or a row scaling broadcast on the grid; its distinct values are + validated without expanding a broadcast view. """ values = np.asarray(scaling) distinct = values[tuple(slice(None) if s else 0 for s in values.strides)] validate_scaling(distinct, "constraint scaling") - kept = self._kept(np.diff(self._csr.indptr) > 0) + keep = np.diff(self._csr.indptr) > 0 + if mask is not None: + keep &= mask + kept = self._kept(keep) csr = kept._csr if not csr.data.all(): csr = csr.copy() if csr is self._csr else csr csr.eliminate_zeros() if isinstance(scaling, DataArray): - row_scaling = kept._active_values(scaling) + row_scaling = kept.active_values(scaling) else: row_scaling = np.full(csr.shape[0], float(scaling)) return kept._replace(csr=csr, cindex=cindex, name=name, scaling=row_scaling) - def _active_values(self, values: DataArray) -> np.ndarray: + def active_values(self, values: DataArray) -> np.ndarray: """ Values of ``values``, broadcast on the grid, at the active rows. @@ -1002,13 +1010,6 @@ def rows(values: Any) -> Any: binval=rows(self._binval), ) - def masked(self, mask: DataArray) -> CSRConstraint: - """ - Copy with the cells where the boolean ``mask`` is False made inactive, - without the dense rectangle. ``mask`` must lie on the constraint grid. - """ - return self._kept(self._active_values(mask).astype(bool)) - def _assign_coords(self, **coords: Any) -> CSRConstraint: """ Reassign coordinate values on the constraint, keeping the shape. @@ -1448,7 +1449,7 @@ def to_matrix_with_rhs( 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]) + sense = self._sign.astype("U1") return ( self._to_positional_csr(label_index), self.active_labels(), @@ -1472,7 +1473,9 @@ def sanitize_zeros(self) -> CSRConstraint: external holders of the previous arrays (e.g. a ModelSnapshot sharing them) keep a valid baseline. """ - zeros = np.abs(self._csr.data) <= 1e-10 + data = self._csr.data + zeros = data <= 1e-10 + zeros &= data >= -1e-10 if zeros.any(): csr = self._csr.copy() csr.data[zeros] = 0 @@ -1617,14 +1620,12 @@ def from_dense( scaling = con.scaling.values.ravel()[active_mask] sign_vals = con.sign.values.ravel() active_signs = sign_vals[active_mask] - unique_signs = np.unique(active_signs) - if len(unique_signs) == 0: + sign: str | np.ndarray + if not len(active_signs): full_unique_signs = np.unique(sign_vals) - sign: str | np.ndarray = ( - str(full_unique_signs.item()) if len(full_unique_signs) == 1 else "=" - ) - elif len(unique_signs) == 1: - sign = str(unique_signs[0]) + sign = str(full_unique_signs.item()) if len(full_unique_signs) == 1 else "=" + elif (active_signs == active_signs[0]).all(): + sign = str(active_signs[0]) else: sign = active_signs dual = ( @@ -2221,12 +2222,7 @@ def to_matrix_with_rhs( csr.sum_duplicates() b = self.rhs.values.ravel()[row_mask] - sign_flat = self.sign.values.ravel()[row_mask] - unique_signs = np.unique(sign_flat) - if len(unique_signs) == 1: - sense = np.full(len(con_labels), str(unique_signs[0])[0], dtype="U1") - else: - sense = sign_flat.astype("U1") + sense = self.sign.values.ravel()[row_mask].astype("U1") return csr, con_labels, b, sense def sanitize_zeros(self) -> Constraint: diff --git a/linopy/csr.py b/linopy/csr.py index 2222bd50..b7bad36a 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -28,6 +28,7 @@ from __future__ import annotations +import functools import operator from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, field, replace @@ -572,25 +573,30 @@ def same_grid(self, other: CSRLinearExpression) -> bool: """Whether both live on the same cells, auxiliary coordinates aside.""" return self.grid.same_layout(other.grid) - def added(self, other: CSRLinearExpression) -> CSRLinearExpression: + def added(self, *others: CSRLinearExpression) -> CSRLinearExpression: """ - Sparse matrix addition == merge along the term dimension. Goes through - COO so explicit zero coefficients survive (scipy's ``+`` drops them), + Sparse matrix addition == merge along the term dimension, over any + number of operands on the same grid in one COO pass. Goes through COO + so explicit zero coefficients survive (scipy's ``+`` drops them), keeping a cell with only zero-coefficient terms distinguishable from - an empty cell. A cell absent in either operand is absent in the sum - and carries no terms. Auxiliary coordinates propagate and conflicting - ones raise (§11). - """ - const = self.const + other.const - a, b = self.csr.tocoo(), other.csr.tocoo() - shape = (self.n_cells, max(a.shape[1], b.shape[1])) - rows = np.concatenate([a.coords[0], b.coords[0]]) - cols = np.concatenate([a.coords[1], b.coords[1]]) - data = np.concatenate([a.data, b.data]) + an empty cell. A cell absent in any operand is absent in the sum and + carries no terms. Auxiliary coordinates propagate, earlier operands + taking precedence, and conflicting ones raise (§11). + """ + parts = (self, *others) + const = functools.reduce(np.add, (p.const for p in parts)) + coos = [p.csr.tocoo() for p in parts] + shape = (self.n_cells, max(c.shape[1] for c in coos)) + rows = np.concatenate([c.coords[0] for c in coos]) + cols = np.concatenate([c.coords[1] for c in coos]) + data = np.concatenate([c.data for c in coos]) present = ~np.isnan(const)[rows] csr = coo_to_csr(data[present], rows[present], cols[present], shape, self.model) - enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in (self, other)]) - grid = replace(self.grid, aux=other.grid.aux | self.grid.aux) + enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in parts]) + aux: AuxCoords = {} + for p in reversed(parts): + aux |= p.grid.aux + grid = replace(self.grid, aux=aux) return replace(self, csr=csr, const=const, grid=grid) def contracted( diff --git a/linopy/dualization.py b/linopy/dualization.py index edcfffdf..27bf68e7 100644 --- a/linopy/dualization.py +++ b/linopy/dualization.py @@ -479,12 +479,13 @@ def _add_dual_feasibility_constraints( dual_vars : dict ``{constraint_name: dual_variable}`` as returned by ``_add_dual_variables()``. """ - A = m.matrices.A + M = m.matrices + A = M.A if A is None: raise ValueError("Constraint matrix is None, model has no constraints.") - vlabels = np.asarray(m.matrices.vlabels, dtype=np.int64) - clabels = np.asarray(m.matrices.clabels, dtype=np.int64) + vlabels = np.asarray(M.vlabels, dtype=np.int64) + clabels = np.asarray(M.clabels, dtype=np.int64) flat_con_to_dual = _build_flat_con_to_dual_label_lookup(m, dual_vars) if not len(flat_con_to_dual): @@ -496,7 +497,7 @@ def _add_dual_feasibility_constraints( flat_v, flat_d, nnz_data = _extract_dual_feas_entries( A, vlabels, clabels, flat_con_to_dual ) - c_lookup = _build_obj_coeff_lookup(vlabels, m.matrices.c) + c_lookup = _build_obj_coeff_lookup(vlabels, M.c) logger.debug("Building dual feasibility constraints for each primal variable.") for var_name, var in m.variables.items(): diff --git a/linopy/expressions.py b/linopy/expressions.py index 827a319a..8aad59de 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1737,17 +1737,25 @@ def mask(self) -> None: return None @has_optimized_model + def _label_solution(self) -> np.ndarray: + """ + Solution values indexed by variable label, with a trailing NaN that + label ``-1`` reads. + """ + sol = np.full(self.model._xCounter + 1, np.nan) + for _, var in self.model.variables.items(): + labels = var.labels.values.ravel() + mask = labels != -1 + sol[labels[mask]] = var.solution.values.ravel()[mask] + return sol + def _map_solution(self) -> DataArray: """ Replace variable labels by solution values. """ - m = self.model - M = m.matrices - sol = pd.Series(M.sol, M.vlabels) - sol[-1] = np.nan - idx = np.ravel(self.vars) - values = np.asarray(sol[idx]).reshape(self.vars.shape) - return xr.DataArray(values, dims=self.vars.dims, coords=self.vars.coords) + labels = self.vars + values = self._label_solution()[labels.values] + return xr.DataArray(values, dims=labels.dims, coords=labels.coords) @property def solution(self) -> DataArray: @@ -2496,6 +2504,20 @@ def const(self) -> DataArray: def const(self, value: DataArray) -> None: self._data = assign_multiindex_safe(self.data, const=value) + @property + def solution(self) -> DataArray: + """ + Get the optimal values of the expression. + + The function raises an error in case no model is set as a + reference or the model is not optimized. + """ + csr = self._csr + if csr is None: + return super().solution + sol = np.nan_to_num(self._label_solution()[: csr.csr.shape[1]]) + return csr.grid.dataarray(csr.csr @ sol + csr.const, name="solution") + def _combined_with_constant( self, self_const: DataArray, @@ -3736,9 +3758,7 @@ def _try_csr_merge( return None csrs = aligned - combined = csrs[0] - for csr in csrs[1:]: - combined = combined.added(csr) + combined = csrs[0].added(*csrs[1:]) return LinearExpression._from_csr(combined, exprs[0].model) diff --git a/linopy/io.py b/linopy/io.py index 3f2bb8f4..81fff260 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -119,23 +119,30 @@ def _lookup_positive_labels(lookup: np.ndarray, labels: np.ndarray) -> np.ndarra return values +def _non_identity(lookup: np.ndarray) -> np.ndarray | None: + """Return the scaling lookup, or None if all factors are 1.""" + return None if (lookup == 1).all() else lookup + + def _scale_objective_dataframe( - df: pl.DataFrame, variable_scaling: np.ndarray, objective_scaling: float + df: pl.DataFrame, variable_scaling: np.ndarray | None, objective_scaling: float ) -> pl.DataFrame: """Apply column scaling and row-like objective scaling to objective terms.""" if df.is_empty(): return df - if "vars" in df.columns: - scales = _lookup_positive_labels(variable_scaling, df["vars"].to_numpy()) - else: - scale1 = _lookup_positive_labels(variable_scaling, df["vars1"].to_numpy()) - scale2 = _lookup_positive_labels(variable_scaling, df["vars2"].to_numpy()) - scales = scale1 * scale2 + if variable_scaling is not None: + if "vars" in df.columns: + scales = _lookup_positive_labels(variable_scaling, df["vars"].to_numpy()) + else: + scale1 = _lookup_positive_labels(variable_scaling, df["vars1"].to_numpy()) + scale2 = _lookup_positive_labels(variable_scaling, df["vars2"].to_numpy()) + scales = scale1 * scale2 + df = df.with_columns(pl.col("coeffs") / pl.Series(scales)) - return df.with_columns( - (pl.col("coeffs") / pl.Series(scales) * objective_scaling).alias("coeffs") - ) + if objective_scaling != 1: + df = df.with_columns(pl.col("coeffs") * objective_scaling) + return df def _scale_bounds_dataframe( @@ -154,19 +161,22 @@ def _scale_bounds_dataframe( def _scale_constraint_dataframe( df: pl.DataFrame, - variable_scaling: np.ndarray, - constraint_scaling: np.ndarray, + variable_scaling: np.ndarray | None, + constraint_scaling: np.ndarray | None, ) -> pl.DataFrame: """Divide by column scaling and multiply by row scaling.""" - if df.is_empty(): - return df - row_scales = _lookup_positive_labels(constraint_scaling, df["labels"].to_numpy()) - var_scales = _lookup_positive_labels(variable_scaling, df["vars"].to_numpy()) - row_scale_series = pl.Series(row_scales) - return df.with_columns( - (pl.col("coeffs") / pl.Series(var_scales) * row_scale_series).alias("coeffs"), - (pl.col("rhs") * row_scale_series).alias("rhs"), - ) + if variable_scaling is not None: + var_scales = _lookup_positive_labels(variable_scaling, df["vars"].to_numpy()) + df = df.with_columns(pl.col("coeffs") / pl.Series(var_scales)) + if constraint_scaling is not None: + row_scales = _lookup_positive_labels( + constraint_scaling, df["labels"].to_numpy() + ) + row_scale_series = pl.Series(row_scales) + df = df.with_columns( + pl.col("coeffs") * row_scale_series, pl.col("rhs") * row_scale_series + ) + return df def format_coord(coord: str) -> str: @@ -281,6 +291,7 @@ def objective_write_quadratic_terms( def objective_to_file( m: Model, f: BufferedWriter, + variable_scaling: np.ndarray | None, progress: bool = False, explicit_coordinate_names: bool = False, ) -> None: @@ -293,7 +304,6 @@ def objective_to_file( print_variable, _ = get_printers( m, explicit_coordinate_names=explicit_coordinate_names ) - variable_scaling = variable_scaling_lookup(m) sense = m.objective.sense f.write(f"{sense}\n\nobj:\n\n".encode()) @@ -325,6 +335,7 @@ def _binary_has_nondefault_bounds(var: Variable) -> bool: def bounds_to_file( m: Model, f: BufferedWriter, + variable_scaling: np.ndarray | None, progress: bool = False, slice_size: int = 2_000_000, explicit_coordinate_names: bool = False, @@ -348,7 +359,6 @@ def bounds_to_file( print_variable, _ = get_printers( m, explicit_coordinate_names=explicit_coordinate_names ) - variable_scaling = variable_scaling_lookup(m) f.write(b"\n\nbounds\n\n") if progress: @@ -362,7 +372,8 @@ def bounds_to_file( var = m.variables[name] for var_slice in var.iterate_slices(slice_size): df = var_slice.to_polars() - df = _scale_bounds_dataframe(df, variable_scaling) + if variable_scaling is not None: + df = _scale_bounds_dataframe(df, variable_scaling) columns = [ *signed_number(pl.col("lower")), @@ -557,6 +568,7 @@ def sos_to_file( def indicator_constraints_to_file( m: Model, f: BufferedWriter, + variable_scaling: np.ndarray, explicit_coordinate_names: bool = False, ) -> None: """ @@ -574,7 +586,6 @@ def indicator_constraints_to_file( print_variable_scalar, _ = get_printers_scalar( m, explicit_coordinate_names=explicit_coordinate_names ) - variable_scaling = variable_scaling_lookup(m) for con in m.constraints.indicator.data.values(): ic_data = con.data @@ -619,6 +630,8 @@ def indicator_constraints_to_file( def constraints_to_file( m: Model, f: BufferedWriter, + variable_scaling: np.ndarray | None, + constraint_scaling: np.ndarray | None, progress: bool = False, lazy: bool = False, slice_size: int = 2_000_000, @@ -631,8 +644,7 @@ def constraints_to_file( print_variable, print_constraint = get_printers( m, explicit_coordinate_names=explicit_coordinate_names ) - variable_scaling = variable_scaling_lookup(m) - constraint_scaling = constraint_scaling_lookup(m) + scaled = variable_scaling is not None or constraint_scaling is not None f.write(b"\n\ns.t.\n\n") names = list(regular) @@ -643,48 +655,53 @@ def constraints_to_file( colour=TQDM_COLOR, ) - # to make this even faster, we can use polars expression - # https://docs.pola.rs/user-guide/expressions/plugins/#output-data-types for name in names: con = regular[name] for con_slice in con.iterate_slices(slice_size): df = con_slice.to_polars() - df = _scale_constraint_dataframe(df, variable_scaling, constraint_scaling) - if df.height == 0: continue + if scaled: + df = _scale_constraint_dataframe( + df, variable_scaling, constraint_scaling + ) + if not df["labels"].is_sorted(): + df = df.sort("labels", maintain_order=True) + _write_constraint_rows(df, f, print_constraint, print_variable) + - # Sort by labels and mark first/last occurrences - df = df.sort("labels").with_columns( +def _write_constraint_rows( + df: pl.DataFrame, + f: BufferedWriter, + print_constraint: Callable, + print_variable: Callable, +) -> None: + labels = df["labels"].to_numpy() + first = np.empty(len(labels), dtype=bool) + first[0] = True + np.not_equal(labels[1:], labels[:-1], out=first[1:]) + last = np.append(first[1:], True) + columns = [ + pl.when(pl.Series(first)).then( + pl.concat_str( + [*print_constraint(pl.col("labels")), pl.lit(":\n")], + ignore_nulls=True, + ) + ), + *signed_number(pl.col("coeffs")), + *print_variable(pl.col("vars")), + pl.when(pl.Series(last)).then( + pl.concat_str( [ - pl.col("labels").is_first_distinct().alias("is_first_in_group"), - (pl.col("labels") != pl.col("labels").shift(-1)) - .fill_null(True) - .alias("is_last_in_group"), + pl.lit("\n"), + pl.col("sign"), + pl.lit(" "), + pl.col("rhs").cast(pl.String), ] ) - - row_labels = print_constraint(pl.col("labels")) - col_labels = print_variable(pl.col("vars")) - columns = [ - pl.when(pl.col("is_first_in_group")).then(row_labels[0]), - pl.when(pl.col("is_first_in_group")).then(row_labels[1]), - pl.when(pl.col("is_first_in_group")).then(pl.lit(":\n")).alias(":"), - *signed_number(pl.col("coeffs")), - col_labels[0], - col_labels[1], - pl.when(pl.col("is_last_in_group")).then(pl.lit("\n")), - pl.when(pl.col("is_last_in_group")).then(pl.col("sign")), - pl.when(pl.col("is_last_in_group")).then(pl.lit(" ")), - pl.when(pl.col("is_last_in_group")).then(pl.col("rhs").cast(pl.String)), - ] - - _format_and_write(df, columns, f) - - # in the future, we could use lazy dataframes when they support appending - # tp existent files - # formatted = df.lazy().select(pl.concat_str(columns, ignore_nulls=True)) - # formatted.sink_csv(f, **kwargs) + ), + ] + _format_and_write(df, columns, f) def to_lp_file( @@ -697,13 +714,21 @@ def to_lp_file( ) -> None: with open(fn, mode="wb") as f: start = time.time() + variable_scaling = variable_scaling_lookup(m) + active_variable_scaling = _non_identity(variable_scaling) objective_to_file( - m, f, progress=progress, explicit_coordinate_names=explicit_coordinate_names + m, + f, + active_variable_scaling, + progress=progress, + explicit_coordinate_names=explicit_coordinate_names, ) constraints_to_file( m, f=f, + variable_scaling=active_variable_scaling, + constraint_scaling=_non_identity(constraint_scaling_lookup(m)), progress=progress, slice_size=slice_size, explicit_coordinate_names=explicit_coordinate_names, @@ -711,11 +736,13 @@ def to_lp_file( indicator_constraints_to_file( m, f=f, + variable_scaling=variable_scaling, explicit_coordinate_names=explicit_coordinate_names, ) bounds_to_file( m, f=f, + variable_scaling=active_variable_scaling, progress=progress, slice_size=slice_size, explicit_coordinate_names=explicit_coordinate_names, diff --git a/linopy/matrices.py b/linopy/matrices.py index bca66d70..35dfa682 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -7,8 +7,9 @@ from __future__ import annotations -from functools import cached_property -from typing import TYPE_CHECKING, cast +from collections.abc import Callable +from functools import cached_property, partial +from typing import TYPE_CHECKING, NamedTuple, cast import numpy as np import scipy.sparse @@ -16,15 +17,66 @@ from linopy import expressions from linopy.constraints import CSRConstraint -from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup +from linopy.csr import index_dtype if TYPE_CHECKING: + from linopy.common import VariableLabelIndex + from linopy.constraints import ConstraintBase from linopy.model import Model -def _stack(csrs: list) -> scipy.sparse.csr_array | None: +class _RowBlock(NamedTuple): """ - Vertically stack CSR blocks, or None when there are none. + Rows of one constraint. ``positional`` returns its positional CSR; for a + frozen constraint it is built only while the rows are gathered. + """ + + positional: Callable[[], scipy.sparse.csr_array] + n_rows: int + nnz: int + rhs: ndarray + sign: str | ndarray + row_scaling: ndarray | None + + +def _unit_or(scaling: ndarray) -> ndarray | None: + return None if (scaling == 1).all() else scaling + + +def _row_block(con: ConstraintBase, label_index: VariableLabelIndex) -> _RowBlock: + if isinstance(con, CSRConstraint): + stored = con._csr + return _RowBlock( + partial(con._to_positional_csr, label_index), + stored.shape[0], + int(np.count_nonzero(stored.data)), + con._rhs, + con._sign, + _unit_or(con._scaling), + ) + csr, _, rhs, sense = con.to_matrix_with_rhs(label_index) + scaling = _unit_or(con.scaling.values.ravel()) + return _RowBlock( + lambda: csr, + csr.shape[0], + int(np.count_nonzero(csr.data)), + rhs, + sense, + None if scaling is None else scaling[con.active_row_mask()], + ) + + +def _stack( + blocks: list[_RowBlock], + label_index: VariableLabelIndex, + col_scaling: ndarray | None, + model: Model, +) -> tuple[scipy.sparse.csr_array | None, ndarray, ndarray]: + """ + Gather the scaled row blocks into one preallocated CSR, rhs and sense. + + With solver variables y = Scol * x, constraints A x = b become + Srow * A * Scol^-1 * y = Srow * b. Explicit zeros are dropped: expressions that broadcast against a dense coordinate store one coefficient per pair, most of them zero, and a zero @@ -32,16 +84,52 @@ def _stack(csrs: list) -> scipy.sparse.csr_array | None: stored nnz handed to the solvers/writers (e.g. ``highspy.addRows`` scales with stored nnz), so we prune them once, centrally, for every backend. """ - if not csrs: - return None - stacked = cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr")) - stacked.eliminate_zeros() - return stacked - - -def _concat(arrays: list, dtype: type | None = None) -> ndarray: - """Concatenate arrays, or an empty array when there are none.""" - return np.concatenate(arrays) if arrays else np.array([], dtype=dtype) + if not blocks: + return None, np.array([]), np.array([], dtype=object) + nnz = sum(block.nnz for block in blocks) + n_rows = sum(block.n_rows for block in blocks) + n_cols = label_index.n_active_vars + dtype = index_dtype(nnz, (n_rows, n_cols), model) + data = np.empty(nnz, dtype=float) + indices = np.empty(nnz, dtype=dtype) + indptr = np.empty(n_rows + 1, dtype=dtype) + b = np.empty(n_rows, dtype=float) + sense = np.empty(n_rows, dtype="U1") + indptr[0] = 0 + pos = row = 0 + for block in blocks: + source = block.positional() + count = block.nnz + block_data, block_indices, block_indptr = ( + source.data, + source.indices, + source.indptr, + ) + if count < source.nnz: + keep = block_data != 0 + block_data, block_indices = block_data[keep], block_indices[keep] + block_indptr = np.concatenate([[0], np.cumsum(keep)])[block_indptr] + rows = slice(row, row + source.shape[0]) + entries = slice(pos, pos + count) + if block.row_scaling is None: + data[entries] = block_data + b[rows] = block.rhs + else: + row_counts = np.diff(block_indptr) + np.multiply( + block_data, np.repeat(block.row_scaling, row_counts), out=data[entries] + ) + np.multiply(block.rhs, block.row_scaling, out=b[rows]) + if col_scaling is not None: + data[entries] /= col_scaling[block_indices] + indices[entries] = block_indices + block_offsets = indptr[rows.start + 1 : rows.stop + 1] + block_offsets[:] = block_indptr[1:] + block_offsets += pos + sense[rows] = block.sign + pos, row = entries.stop, rows.stop + A = scipy.sparse.csr_array((data, indices, indptr), shape=(n_rows, n_cols)) + return A, b, sense def _binval_per_row(binval: int | np.ndarray, n: int) -> ndarray: @@ -72,23 +160,17 @@ 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 - var_scaling_by_label = variable_scaling_lookup(m) - self.var_scaling: ndarray = ( - var_scaling_by_label[self.vlabels] - if len(self.vlabels) - else np.array([], dtype=float) - ) - - lb_list = [] - ub_list = [] - vtypes_list = [] - + self.vlabels: ndarray = m.variables.label_index.vlabels + n = len(self.vlabels) + self.var_scaling: ndarray = np.empty(n, dtype=float) + self.lb: ndarray = np.empty(n, dtype=float) + self.ub: ndarray = np.empty(n, dtype=float) + self.vtypes: ndarray = np.empty(n, dtype="U1") + + pos = 0 for name, var in m.variables.items(): - labels = var.labels.values.ravel() - mask = labels != -1 - + mask = var.labels.values.ravel() != -1 + cols = slice(pos, pos + int(np.count_nonzero(mask))) if name in m.binaries: vtype = "B" elif name in m.integers: @@ -97,77 +179,48 @@ def _build_vars(self) -> None: vtype = "S" else: vtype = "C" + self.vtypes[cols] = vtype + self.var_scaling[cols] = var.solver_scaling.values.ravel()[mask] + self.lb[cols] = var.lower.values.ravel()[mask] + self.ub[cols] = var.upper.values.ravel()[mask] + pos = cols.stop - 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 lb_list: - self.lb: ndarray = np.concatenate(lb_list) * self.var_scaling - self.ub: ndarray = np.concatenate(ub_list) * self.var_scaling - self.vtypes: ndarray = np.concatenate(vtypes_list) - else: - self.lb = np.array([]) - self.ub = np.array([]) - self.vtypes = np.array([], dtype=object) + if not (self.var_scaling == 1).all(): + self.lb *= self.var_scaling + self.ub *= self.var_scaling def _build_cons(self) -> None: m = self._parent label_index = m.variables.label_index label_to_pos = label_index.label_to_pos - con_scaling_by_label = constraint_scaling_lookup(m) - unit_cols = bool((self.var_scaling == 1).all()) - - def scale_rows_and_cols( - csr: scipy.sparse.csr_array, con_labels: np.ndarray, b: np.ndarray - ) -> tuple[scipy.sparse.csr_array, np.ndarray]: - if csr.shape[0] == 0: - return csr, b - row_scaling = con_scaling_by_label[con_labels] - unit_rows = bool((row_scaling == 1).all()) - if unit_rows and unit_cols: - return csr, b - # With solver variables y = Scol * x, constraints A x = b become - # Srow * A * Scol^-1 * y = Srow * b. - data = csr.data - if not unit_rows: - data = data * np.repeat(row_scaling, np.diff(csr.indptr)) - if not unit_cols: - data = data / self.var_scaling[csr.indices] - scaled = scipy.sparse.csr_array( - (data, csr.indices, csr.indptr), shape=csr.shape - ) - return scaled, b * row_scaling + col_scaling = None if (self.var_scaling == 1).all() else self.var_scaling - reg_csrs, reg_b, reg_sense = [], [], [] - ind_csrs, ind_b, ind_sense, ind_binvar, ind_binval = [], [], [], [], [] + regular, indicator = [], [] + binvar, binval = [], [] for c in m.constraints.data.values(): - if c.is_indicator: - cc = c if isinstance(c, CSRConstraint) else c.freeze() - csr, con_labels, b, sense = cc.to_matrix_with_rhs(label_index) - csr, b = scale_rows_and_cols(csr, con_labels, b) - ind_csrs.append(csr) - ind_b.append(b) - ind_sense.append(sense) - ind_binvar.append(label_to_pos[cc._binvar_labels]) - binval = cast("int | np.ndarray", cc._binval) - ind_binval.append(_binval_per_row(binval, len(b))) - else: - csr, con_labels, b, sense = c.to_matrix_with_rhs(label_index) - csr, b = scale_rows_and_cols(csr, con_labels, b) - reg_csrs.append(csr) - reg_b.append(b) - reg_sense.append(sense) + if not c.is_indicator: + regular.append(_row_block(c, label_index)) + continue + cc = c if isinstance(c, CSRConstraint) else c.freeze() + indicator.append(_row_block(cc, label_index)) + binvar.append(label_to_pos[cc._binvar_labels]) + cc_binval = cast("int | np.ndarray", cc._binval) + binval.append(_binval_per_row(cc_binval, len(cc._rhs))) self.clabels: ndarray = m.constraints.label_index.clabels - self.A: scipy.sparse.csr_array | None = _stack(reg_csrs) - self.b: ndarray = _concat(reg_b) - self.sense: ndarray = _concat(reg_sense, dtype=object) - self.indicator_A: scipy.sparse.csr_array | None = _stack(ind_csrs) - self.indicator_b: ndarray = _concat(ind_b) - self.indicator_sense: ndarray = _concat(ind_sense, dtype=object) - self.indicator_binvar: ndarray = _concat(ind_binvar, dtype=np.intp) - self.indicator_binval: ndarray = _concat(ind_binval, dtype=np.intp) + self.A: scipy.sparse.csr_array | None + self.A, self.b, self.sense = _stack(regular, label_index, col_scaling, m) + self.indicator_A: scipy.sparse.csr_array | None + self.indicator_A, self.indicator_b, self.indicator_sense = _stack( + indicator, label_index, col_scaling, m + ) + label_dtype = m._dtypes["labels"] + self.indicator_binvar: ndarray = np.concatenate( + [np.array([], dtype=label_dtype), *binvar] + ) + self.indicator_binval: ndarray = ( + np.concatenate(binval) if binval else np.array([], dtype=np.intp) + ) @cached_property def c(self) -> ndarray: diff --git a/linopy/model.py b/linopy/model.py index c873f684..8c041441 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -127,11 +127,13 @@ DtypeKey = Literal["labels"] -def _check_infinities(sign: Any, rhs: Any, name: str) -> None: +def _check_infinities( + sign: Any, rhs: Any, name: str, where: np.ndarray | None = None +) -> None: invalid = ((sign == LESS_EQUAL) & (rhs == -np.inf)) | ( (sign == GREATER_EQUAL) & (rhs == np.inf) ) - if np.any(invalid): + if np.any(invalid if where is None else invalid & where): raise ValueError(f"Constraint {name} contains incorrect infinite values.") @@ -1371,9 +1373,14 @@ def add_constraints( con = self._constraint_from_lhs(lhs, sign, rhs, coords) if isinstance(con, CSRConstraint) and freeze: - if mask is not None: - con = con.masked(broadcast_to_coords(mask, con.coords, label="mask")) - _check_infinities(con._sign, con._rhs, name) + row_mask = ( + None + if mask is None + else con.active_values( + broadcast_to_coords(mask, con.coords, label="mask") + ).astype(bool) + ) + _check_infinities(con._sign, con._rhs, name, row_mask) self.check_force_dim_names(con.coords.to_dataset()) enforce_no_multiindex(con, context=f"constraint {name!r}") row_scaling = ( @@ -1385,7 +1392,7 @@ def add_constraints( ) cindex = self._cCounter self._cCounter += con.full_size - con = con.assign_labels(cindex, name, row_scaling) + con = con.assign_labels(cindex, name, row_scaling, row_mask) return self._soften_added(self.constraints.add(con), penalty) if isinstance(con, CSRConstraint): if chunked: @@ -2523,6 +2530,10 @@ def assign_result( if con.is_indicator: continue start, end = con.range + if isinstance(con, CSRConstraint): + active = dual[start:end][con.active_positions] + con._dual = active * con._scaling / self.objective.scaling + continue coords = {dim: con.coords[dim] for dim in con.coord_dims} values = ( dual[start:end].reshape(con.shape) diff --git a/linopy/solvers.py b/linopy/solvers.py index 6e1ec74d..473868da 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -1726,11 +1726,11 @@ def _build_solver_model( print_variables, print_constraints = linopy.io.get_printers_scalar( model, explicit_coordinate_names=explicit_coordinate_names ) - lp = h.getLp() - lp.col_names_ = print_variables(M.vlabels) + mdl = h.getModel() + mdl.lp_.col_names_ = print_variables(M.vlabels) if len(M.clabels): - lp.row_names_ = print_constraints(M.clabels) - h.passModel(lp) + mdl.lp_.row_names_ = print_constraints(M.clabels) + h.passModel(mdl) Q = M.Q if Q is not None: @@ -3526,56 +3526,49 @@ def _build_solver_model( np.arange(0, len(labels)), "%0", [len(labels)], None, [0], labels ) - bkx = [ - ( - ( - (mosek.boundkey.ra if lb < ub else mosek.boundkey.fx) - if ub < np.inf - else mosek.boundkey.lo - ) - if (lb > -np.inf) - else (mosek.boundkey.up if (ub < np.inf) else mosek.boundkey.fr) - ) - for (lb, ub) in zip(M.lb, M.ub) - ] - blx = [b if b > -np.inf else 0.0 for b in M.lb] - bux = [b if b < np.inf else 0.0 for b in M.ub] - task.putvarboundslice(0, model.nvars, bkx, blx, bux) + bk = mosek.boundkey + keys = np.array([bk.fr, bk.lo, bk.up, bk.fx, bk.ra], dtype=object) + lb_fin = M.lb > -np.inf + ub_fin = M.ub < np.inf + bkx_code = np.select( + [lb_fin & ub_fin & (M.lb < M.ub), lb_fin & ub_fin, lb_fin, ub_fin], + [4, 3, 1, 2], + default=0, + ) + blx = np.where(lb_fin, M.lb, 0.0) + bux = np.where(ub_fin, M.ub, 0.0) + task.putvarboundslice(0, model.nvars, keys[bkx_code].tolist(), blx, bux) if len(model.binaries.labels) + len(model.integers.labels) > 0: - idx = [i for (i, v) in enumerate(M.vtypes) if v in ["B", "I"]] + idx = np.flatnonzero(np.isin(M.vtypes, ["B", "I"])).astype(np.int32) task.putvartypelist(idx, [mosek.variabletype.type_int] * len(idx)) if len(model.constraints) > 0: if set_names: names = print_constraints(M.clabels) - for i, n in enumerate(names): - task.putconname(i, n) - bkc = [ - ( - (mosek.boundkey.up if b < np.inf else mosek.boundkey.fr) - if s == "<" - else ( - (mosek.boundkey.lo if b > -np.inf else mosek.boundkey.up) - if s == ">" - else mosek.boundkey.fx - ) + task.generateconnames( + np.arange(0, len(names)), "%0", [len(names)], None, [0], names ) - for s, b in zip(M.sense, M.b) - ] - blc = [b if b > -np.inf else 0.0 for b in M.b] - buc = [b if b < np.inf else 0.0 for b in M.b] + b_lo = M.b > -np.inf + b_up = M.b < np.inf + leq = M.sense == "<" + geq = M.sense == ">" + bkc_code = np.select( + [leq & b_up, leq, geq & b_lo, geq], [2, 0, 1, 2], default=3 + ) + blc = np.where(b_lo, M.b, 0.0) + buc = np.where(b_up, M.b, 0.0) if M.A is not None: A = M.A.tocsr() task.putarowslice( 0, model.ncons, A.indptr[:-1], A.indptr[1:], A.indices, A.data ) - task.putconboundslice(0, model.ncons, bkc, blc, buc) + task.putconboundslice(0, model.ncons, keys[bkc_code].tolist(), blc, buc) if M.Q is not None: Q = (0.5 * tril(M.Q + M.Q.transpose())).tocoo() task.putqobj(Q.row, Q.col, Q.data) - task.putclist(list(np.arange(model.nvars)), M.c) + task.putclist(np.arange(model.nvars, dtype=np.int32), M.c) if model.objective.sense == "max": task.putobjsense(mosek.objsense.maximize) diff --git a/test/test_csr.py b/test/test_csr.py index a3349105..e89d0c7b 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -11,6 +11,7 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager from dataclasses import dataclass, replace +from functools import partial from pathlib import Path from typing import Any @@ -456,12 +457,21 @@ def build(sparse: bool) -> ConstraintBase: np.testing.assert_array_equal(dense.labels.values, sparse.labels.values) +@pytest.mark.parametrize("masked", [False, True], ids=["unmasked", "masked"]) @pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) -def test_frozen_invalid_infinite_rhs_raises(sparse: bool) -> None: +def test_frozen_invalid_infinite_rhs_raises(sparse: bool, masked: bool) -> None: require_v1() c = base_model(sparse=sparse) + valid = c.load.bus != "bus0" + rhs = c.load.where(valid, -np.inf) + mask = valid if masked else None + add = partial(c.m.add_constraints, name="bal", freeze=True, mask=mask) + if masked and sparse: + con = add(c.balance_lhs() <= rhs) + assert con.ncons == int(valid.sum()) * c.load.sizes["snapshot"] + return with pytest.raises(ValueError, match="incorrect infinite values"): - c.m.add_constraints(c.balance_lhs() <= -np.inf, name="bal", freeze=True) + add(c.balance_lhs() <= rhs) ROW_SCALINGS: dict[str, Callable[[xr.DataArray], Any]] = { @@ -748,6 +758,62 @@ def test_reindex_stays_csr_and_matches_dense(indexers: dict) -> None: assert_linequal(sparse, dense) +def case_twins( + build: Callable[[Case], LinearExpression], +) -> Callable[[], tuple[LinearExpression, LinearExpression]]: + def twins() -> tuple[LinearExpression, LinearExpression]: + c1, c2 = twin_models() + return build(c1), build(c2) + + return twins + + +def keyed_observed_twins() -> tuple[LinearExpression, LinearExpression]: + dense, sparse = (keyed_model(sparse=s)[1] for s in (False, True)) + keys = ["period", "season"] + return dense.groupby(keys).sum(observed=True), sparse.groupby(keys).sum( + observed=True + ) + + +@pytest.mark.parametrize("nan_every_other", [False, True], ids=["full", "nan"]) +@pytest.mark.parametrize( + "twins", + [ + case_twins(lambda c: c.balance_lhs() + 2.0), + case_twins(lambda c: c.gen_sum().reindex(bus=["bus3", "bus9", "bus0"])), + case_twins( + lambda c: linopy.merge( + cross_grid_parts(c) + cross_grid_parts(c, ("line3", "line4"))[1:], + join="outer", + fill_value=linopy.ABSENT, + cls=LinearExpression, + ) + ), + keyed_observed_twins, + ], + ids=["composed", "absent_cell", "absent_merge", "aux_coords"], +) +def test_csr_solution_matches_dense( + twins: Callable[[], tuple[LinearExpression, LinearExpression]], + nan_every_other: bool, +) -> None: + require_v1() + dense, sparse = twins() + assert sparse._csr is not None + rng = np.random.default_rng(0) + for name, var in dense.model.variables.items(): + values = rng.uniform(-1, 1, var.shape) + if nan_every_other: + values.ravel()[::2] = np.nan + for m in (dense.model, sparse.model): + m.variables[name].solution = xr.DataArray(values, coords=var.labels.coords) + m._status = "ok" + sol = sparse.solution + assert sparse._csr is not None + xr.testing.assert_allclose(sol, dense.solution) + + def test_reindex_falls_back_to_dense_for_unsupported_kwargs() -> None: require_v1() c1, c2 = twin_models() @@ -894,16 +960,20 @@ def test_transposed_grid_exact_merge_stays_csr_and_matches_dense() -> None: assert_terms_equal(res, linopy.merge(dense, cls=LinearExpression)) +@pytest.mark.parametrize("fill_value", [None, linopy.ABSENT], ids=["fill", "absent"]) @pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) -def test_three_operand_cross_grid_merge_matches_dense(join: JoinOptions) -> None: +def test_three_operand_cross_grid_merge_matches_dense( + join: JoinOptions, fill_value: Any +) -> None: require_v1() c1, c2 = twin_models() third_lines = ("line3", "line4") sparse = cross_grid_parts(c2) + cross_grid_parts(c2, third_lines)[1:] dense = cross_grid_parts(c1) + cross_grid_parts(c1, third_lines)[1:] - res = linopy.merge(sparse, join=join, cls=LinearExpression) + kwargs: dict[str, Any] = {"join": join, "fill_value": fill_value} + res = linopy.merge(sparse, **kwargs, cls=LinearExpression) assert res._csr is not None - assert_terms_equal(res, linopy.merge(dense, join=join, cls=LinearExpression)) + assert_terms_equal(res, linopy.merge(dense, **kwargs, cls=LinearExpression)) def test_cross_grid_merge_absent_fill_matches_dense() -> None: diff --git a/test/test_io.py b/test/test_io.py index 67919367..0277e325 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -20,6 +20,7 @@ from linopy import LESS_EQUAL, Model, available_solvers, read_netcdf from linopy.constants import FACTOR_DIM +from linopy.constraints import Constraint from linopy.expressions import LinearExpression, QuadraticExpression from linopy.io import CONTAINER_ORDER_ATTR, signed_number from linopy.testing import assert_exprequal, assert_model_equal @@ -610,9 +611,12 @@ def test_to_gurobipy(model: Model) -> None: @pytest.mark.skipif("highs" not in available_solvers, reason="Highspy not installed") -def test_to_highspy(model: Model) -> None: - h = model.to_highspy() - assert h.getLp().num_col_ > 0 +@pytest.mark.parametrize("set_names", [True, False]) +def test_to_highspy(model: Model, set_names: bool) -> None: + lp = model.to_highspy(set_names=set_names).getLp() + assert lp.num_col_ > 0 + assert len(lp.col_names_) == (lp.num_col_ if set_names else 0) + assert len(lp.row_names_) == (lp.num_row_ if set_names else 0) @pytest.mark.skipif("mosek" not in available_solvers, reason="Mosek not installed") @@ -715,6 +719,16 @@ def test_model_set_names_in_solver_io(model: Model) -> None: assert model.objective.value == pytest.approx(expected_obj) +@pytest.mark.skipif("highs" not in available_solvers, reason="Highspy not installed") +def test_highs_direct_qp_keeps_hessian_with_names() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(2, name="i")], name="x") + m.add_constraints(x.sum() >= 1) + m.add_objective((x * x).sum() + x.sum()) + m.solve(solver_name="highs", io_api="direct", set_names=True) + assert m.objective.value == pytest.approx(1.5) + + def test_to_blocks(tmp_path: Path) -> None: m: Model = Model() @@ -1000,3 +1014,70 @@ def bound2(m: Model, i: int) -> object: m_mutable.to_file(fn_mutable) assert fn_frozen.read_text() == fn_mutable.read_text() + + +def _lp_constraint_section(m: Model, path: Path, slice_size: int = 2_000_000) -> str: + m.to_file(path, progress=False, slice_size=slice_size) + return path.read_text().split("s.t.\n\n")[1].split("\n\nbounds")[0] + + +@pytest.mark.parametrize("freeze", [True, False]) +@pytest.mark.parametrize( + "slice_size", + [ + pytest.param(2_000_000, id="default-slices"), + pytest.param(1, id="slice-1"), + ], +) +@pytest.mark.parametrize( + ("scaled", "expected"), + [ + pytest.param( + True, + "c0:\n+2.0 x0\n-4.0 x2\n<= 3.0\n" + "c1:\n+1.0 x1\n-4.0 x3\n<= 3.0\n" + "c2:\n+2.0 x0\n+1.0 x1\n>= -0.0\n", + id="scaled", + ), + pytest.param( + False, + "c0:\n+1.0 x0\n-2.0 x2\n<= 1.5\n" + "c1:\n+1.0 x1\n-2.0 x3\n<= 1.5\n" + "c2:\n+1.0 x0\n+1.0 x1\n>= -0.0\n", + id="unscaled", + ), + ], +) +def test_to_file_lp_constraint_section( + tmp_path: Path, + freeze: bool, + slice_size: int, + scaled: bool, + expected: str, +) -> None: + m = Model() + i = pd.RangeIndex(2, name="i") + row_scaling = 2.0 if scaled else 1.0 + x = m.add_variables(coords=[i], name="x", scaling=[1.0, 2.0] if scaled else 1.0) + y = m.add_variables(coords=[i], name="y") + m.add_constraints(x - 2 * y <= 1.5, name="a", freeze=freeze, scaling=row_scaling) + m.add_constraints(x.sum() >= -0.0, name="b", freeze=freeze, scaling=row_scaling) + m.add_objective(x.sum()) + + fn = tmp_path / "constraints.lp" + assert _lp_constraint_section(m, fn, slice_size) == expected + + +def test_to_file_lp_unsorted_constraint_labels(tmp_path: Path) -> None: + m = Model() + i = pd.RangeIndex(3, name="i") + x = m.add_variables(coords=[i], name="x") + y = m.add_variables(coords=[i], name="y") + m.add_constraints(x + 2 * y <= 1, name="a", freeze=False) + m.add_objective(x.sum()) + expected = _lp_constraint_section(m, tmp_path / "sorted.lp") + + con = m.constraints["a"] + m.constraints.data["a"] = Constraint(con.data.isel(i=slice(None, None, -1)), m, "a") + assert not m.constraints["a"].to_polars()["labels"].is_sorted() + assert _lp_constraint_section(m, tmp_path / "unsorted.lp") == expected diff --git a/test/test_matrices.py b/test/test_matrices.py index 6da3eaf1..da6743d1 100644 --- a/test/test_matrices.py +++ b/test/test_matrices.py @@ -7,6 +7,7 @@ import numpy as np import pandas as pd +import pytest import xarray as xr from linopy import EQUAL, GREATER_EQUAL, Model @@ -98,3 +99,44 @@ def test_matrices_float_c() -> None: c = m.matrices.c assert np.all(c == np.array([1.5, 1.5])) + + +@pytest.mark.parametrize("freeze", [False, True]) +def test_matrices_scaled_masked_mixed_signs(freeze: bool) -> None: + m = Model() + i = pd.RangeIndex(3, name="i") + x = m.add_variables(0, 4, coords=[i], name="x", scaling=[1.0, 2.0, 4.0]) + y = m.add_variables(coords=[i], name="y") + z = m.add_variables(coords=[i], name="z", binary=True) + sign = xr.DataArray(["<=", ">=", "="], coords=[i]) + m.add_constraints( + 2 * x + y - y, + sign, + xr.DataArray([1.0, 2.0, 3.0], coords=[i]), + name="c", + scaling=xr.DataArray([10.0, 1.0, 5.0], coords=[i]), + mask=xr.DataArray([True, False, True], coords=[i]), + freeze=freeze, + ) + m.add_constraints(x + 3 * y >= 1, name="d", freeze=not freeze) + m.add_indicator_constraints(z, 1, x <= 2, name="ind") + M = m.matrices + + assert M.A is not None and M.indicator_A is not None + col_scaling = np.array([1.0, 0.5, 0.25]) + expected = np.zeros((5, 9)) + expected[0, 0], expected[1, 2] = 20.0, 2.5 + expected[2:, :3] = np.diag(col_scaling) + expected[2:, 3:6] = 3 * np.eye(3) + np.testing.assert_array_equal(M.A.toarray(), expected) + assert M.A.nnz == 8 + np.testing.assert_array_equal(M.b, [10.0, 15.0, 1.0, 1.0, 1.0]) + np.testing.assert_array_equal(M.sense, ["<", "=", ">", ">", ">"]) + np.testing.assert_array_equal(M.clabels, [0, 2, 3, 4, 5]) + expected_ind = np.zeros((3, 9)) + expected_ind[:, :3] = np.diag(col_scaling) + np.testing.assert_array_equal(M.indicator_A.toarray(), expected_ind) + np.testing.assert_array_equal(M.indicator_b, [2.0, 2.0, 2.0]) + np.testing.assert_array_equal(M.indicator_binvar, [6, 7, 8]) + np.testing.assert_array_equal(M.lb[:3], [0.0, 0.0, 0.0]) + np.testing.assert_array_equal(M.ub[:3], [4.0, 8.0, 16.0]) diff --git a/test/test_scaling.py b/test/test_scaling.py index 16c31842..5685121d 100644 --- a/test/test_scaling.py +++ b/test/test_scaling.py @@ -180,13 +180,14 @@ def test_indicator_constraint_lp_export_uses_scaled_values(tmp_path: Path) -> No assert f"<= {40.0 * 4}" in text -def test_assign_result_unscales_solution_objective_and_dual() -> None: +@pytest.mark.parametrize("freeze", [False, True]) +def test_assign_result_unscales_solution_objective_and_dual(freeze: bool) -> None: m = Model() i = pd.Index(["a", "b"], name="i") x = m.add_variables(coords=[i], name="x", scaling=[10.0, 100.0]) b = m.add_variables(binary=True, name="b", scaling=50.0) - m.add_constraints(x + b >= 1, name="c", scaling=[2.0, 4.0]) + m.add_constraints(x + b >= 1, name="c", scaling=[2.0, 4.0], freeze=freeze) m.add_objective(x.sum() + b, scaling=10.0) primal = np.full(m._xCounter, np.nan)