diff --git a/benchmark/benchmark_sparse_export.py b/benchmark/benchmark_sparse_export.py new file mode 100644 index 000000000..9b0590f39 --- /dev/null +++ b/benchmark/benchmark_sparse_export.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +""" +Benchmark the build and export of a PyPSA-like dispatch model. + +Run as ``python benchmark/benchmark_sparse_export.py {sparse,dense,frozen}``. +``sparse`` uses ``Model(sparse=True)``, ``dense`` keeps mutable constraints and +``frozen`` freezes each constraint of a dense model. Every phase reports its +wall time, its peak traced memory and the operations that densified. The +``matrices``, ``to_highspy`` and ``to_file(lp)`` phases each build the matrices. +""" + +from __future__ import annotations + +import argparse +import gc +import tempfile +import time +import tracemalloc +import warnings +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +import numpy as np +import pandas as pd +import xarray as xr + +import linopy +from linopy.constants import PerformanceWarning +from linopy.io import to_highspy + + +class Phases: + def __init__(self) -> None: + self.rows: list[tuple[str, float, float, list[str]]] = [] + + @contextmanager + def __call__(self, name: str) -> Iterator[None]: + gc.collect() + tracemalloc.start() + start = time.perf_counter() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + yield + elapsed = time.perf_counter() - start + peak = tracemalloc.get_traced_memory()[1] + tracemalloc.stop() + densified = [ + str(w.message)[:90] + for w in caught + if issubclass(w.category, PerformanceWarning) + ] + self.rows.append((name, elapsed, peak / 1e6, densified)) + + def report(self) -> None: + for name, elapsed, peak, densified in self.rows: + note = "DENSIFY: " + "; ".join(densified) if densified else "" + print(f"{name:28s} {elapsed * 1e3:9.0f} ms peak {peak:8.1f} MB {note}") + total = sum(row[1] for row in self.rows) + print(f"{'TOTAL':28s} {total * 1e3:9.0f} ms") + + +def run(mode: str, n_bus: int, n_t: int, n_gen: int, n_line: int, n_sto: int) -> None: + linopy.options["semantics"] = "v1" + linopy.options["warn_on_densify"] = True + rng = np.random.default_rng(0) + + bus = pd.RangeIndex(n_bus, name="bus") + t = pd.RangeIndex(n_t, name="t") + gen = pd.RangeIndex(n_gen, name="gen") + line = pd.RangeIndex(n_line, name="line") + sto = pd.RangeIndex(n_sto, name="sto") + + gen_bus = rng.integers(0, n_bus, n_gen) + gen_bus[rng.random(n_gen) < 0.4] = 0 + gen_bus = pd.Series(gen_bus, index=gen, name="bus") + sto_bus = pd.Series(np.arange(n_sto) % n_bus, index=sto, name="bus") + bus0 = np.arange(n_line) % n_bus + bus1 = (bus0 + 1 + rng.integers(0, 3, n_line)) % n_bus + incidence = np.zeros((n_line, n_bus)) + incidence[np.arange(n_line), bus0] = -1 + incidence[np.arange(n_line), bus1] = 1 + incidence = xr.DataArray(incidence, coords=[line, bus]) + demand = xr.DataArray(rng.uniform(10, 100, (n_bus, n_t)), coords=[bus, t]) + pmax = xr.DataArray(rng.uniform(0, 1, (n_gen, n_t)), coords=[gen, t]) + + freeze = {"sparse": None, "dense": False, "frozen": True}[mode] + m = linopy.Model(sparse=mode == "sparse") + phase = Phases() + + with phase("vars"): + p = m.add_variables(0, coords=[gen, t], name="p") + flow = m.add_variables(-100, 100, coords=[line, t], name="flow") + soc = m.add_variables(0, 100, coords=[sto, t], name="soc") + ch = m.add_variables(0, 50, coords=[sto, t], name="ch") + dis = m.add_variables(0, 50, coords=[sto, t], name="dis") + with phase("expr: groupby nodal supply"): + supply = p.groupby(gen_bus).sum() + with phase("expr: sto groupby"): + sto_net = (dis - ch).groupby(sto_bus).sum() + with phase("expr: flow @ incidence"): + net_flow = flow @ incidence + with phase("expr: balance sum"): + balance = supply + sto_net + net_flow + with phase("cons: balance"): + m.add_constraints(balance == demand, name="balance", freeze=freeze) + with phase("expr: soc shift"): + soc_expr = soc - 0.99 * soc.shift(t=1) - 0.95 * ch + dis / 0.95 + with phase("cons: soc"): + m.add_constraints(soc_expr == 0, name="soc", freeze=freeze) + with phase("expr+cons: p <= pmax"): + m.add_constraints(p <= pmax * 100, name="p_max", freeze=freeze) + with phase("objective"): + m.add_objective((p * 2).sum() + (ch + dis).sum()) + with phase("matrices"): + m.matrices.c + with phase("to_highspy"): + to_highspy(m) + with phase("to_file(lp)"), tempfile.TemporaryDirectory() as tmp: + m.to_file(Path(tmp) / "model.lp", progress=False) + + print( + f"=== {mode} buses={n_bus} snapshots={n_t} generators={n_gen} " + f"lines={n_line} storage={n_sto} nvars={m.nvars} ncons={m.ncons}" + ) + phase.report() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("mode", choices=["sparse", "dense", "frozen"]) + parser.add_argument("--buses", type=int, default=200) + parser.add_argument("--snapshots", type=int, default=720) + parser.add_argument("--generators", type=int, default=2000) + parser.add_argument("--lines", type=int, default=300) + parser.add_argument("--storage", type=int, default=200) + args = parser.parse_args() + run( + args.mode, args.buses, args.snapshots, args.generators, args.lines, args.storage + ) + + +if __name__ == "__main__": + main() diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 42e8bee49..45661be79 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -41,6 +41,7 @@ Upcoming Version * ``@``/``dot`` against a sparse constant no longer scales with the total number of variables in the model: each chunk of the sparse product now runs on only the variables it uses, which removes several seconds of allocation overhead on models with millions of variables. (`#990 `__) * Deprecated in favour of ``Model(sparse=True)``, each with a ``FutureWarning`` and to be removed with the legacy semantics: ``Model(freeze_constraints=...)`` and the ``Model.freeze_constraints`` setter, ``groupby(...).sum(sparse=...)`` and ``linopy.options["sparse_groupby"]``. They keep their current behaviour until then, except that ``@`` ignores ``sparse_groupby``, and netcdf files that store ``freeze_constraints`` still load. (`#976 `__) * Adding a frozen constraint from a sparse expression no longer copies the lhs matrix when every row stays active, and picks the mask and the row scaling at the active rows without expanding them over the full coordinate grid. This roughly halves the peak memory of ``add_constraints`` on a sparse model. (`#977 `__) +* Building ``Model.matrices`` no longer allocates label-sized scaling lookups and skips ``eliminate_zeros`` for frozen constraints. On a model with 2M variables this makes each build of a sparse or frozen model about 20 to 35 ms faster. (`#1008 `__) *Other* diff --git a/linopy/dualization.py b/linopy/dualization.py index edcfffdf0..27bf68e72 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 083e673f2..d2f084635 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1743,11 +1743,15 @@ def _map_solution(self) -> DataArray: Replace variable labels by solution values. """ m = self.model - M = m.matrices - sol = pd.Series(M.sol, M.vlabels) + labels = self.vars.values + known = np.append(m.variables.label_index.label_to_pos != -1, True) + if not known[labels].all(): + raise KeyError("Expression references variables missing from the model.") + sol = np.full(m._xCounter + 1, np.nan) + for var in m.variables.data.values(): + sol[var.labels.values] = var.solution.values sol[-1] = np.nan - idx = np.ravel(self.vars) - values = np.asarray(sol[idx]).reshape(self.vars.shape) + values = sol[labels] return xr.DataArray(values, dims=self.vars.dims, coords=self.vars.coords) @property diff --git a/linopy/matrices.py b/linopy/matrices.py index bca66d70f..08e328e41 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -15,28 +15,17 @@ from numpy import ndarray from linopy import expressions -from linopy.constraints import CSRConstraint -from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup +from linopy.constraints import ConstraintBase, CSRConstraint if TYPE_CHECKING: from linopy.model import Model -def _stack(csrs: list) -> scipy.sparse.csr_array | None: - """ - Vertically stack CSR blocks, or None when there are none. - - Explicit zeros are dropped: expressions that broadcast against a dense - coordinate store one coefficient per pair, most of them zero, and a zero - coefficient never changes a constraint. Keeping them only inflates the - stored nnz handed to the solvers/writers (e.g. ``highspy.addRows`` scales - with stored nnz), so we prune them once, centrally, for every backend. - """ +def _stack(csrs: list[scipy.sparse.csr_array]) -> scipy.sparse.csr_array | None: + """Vertically stack CSR blocks, or None when there are none.""" if not csrs: return None - stacked = cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr")) - stacked.eliminate_zeros() - return stacked + return cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr")) def _concat(arrays: list, dtype: type | None = None) -> ndarray: @@ -44,6 +33,13 @@ def _concat(arrays: list, dtype: type | None = None) -> ndarray: return np.concatenate(arrays) if arrays else np.array([], dtype=dtype) +def _row_scaling(c: ConstraintBase) -> ndarray: + """Scaling of the active rows of a constraint, in row order.""" + if isinstance(c, CSRConstraint): + return c._scaling + return c.scaling.values.ravel()[c.active_row_mask()] + + def _binval_per_row(binval: int | np.ndarray, n: int) -> ndarray: """Broadcast an indicator triggering value to one entry per active row.""" if np.ndim(binval) == 0: @@ -72,18 +68,12 @@ 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) - ) + self.vlabels: ndarray = m.variables.label_index.vlabels lb_list = [] ub_list = [] vtypes_list = [] + scaling_list = [] for name, var in m.variables.items(): labels = var.labels.values.ravel() @@ -101,7 +91,9 @@ def _build_vars(self) -> None: lb_list.append(var.lower.values.ravel()[mask]) ub_list.append(var.upper.values.ravel()[mask]) vtypes_list.append(np.full(mask.sum(), vtype)) + scaling_list.append(var.solver_scaling.values.ravel()[mask]) + self.var_scaling: ndarray = _concat(scaling_list, dtype=float) if lb_list: self.lb: ndarray = np.concatenate(lb_list) * self.var_scaling self.ub: ndarray = np.concatenate(ub_list) * self.var_scaling @@ -115,15 +107,13 @@ 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 + csr: scipy.sparse.csr_array, row_scaling: 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 @@ -144,8 +134,8 @@ def scale_rows_and_cols( 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) + csr, _, b, sense = cc.to_matrix_with_rhs(label_index) + csr, b = scale_rows_and_cols(csr, cc._scaling, b) ind_csrs.append(csr) ind_b.append(b) ind_sense.append(sense) @@ -153,8 +143,10 @@ def scale_rows_and_cols( 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) + csr, _, b, sense = c.to_matrix_with_rhs(label_index) + if not isinstance(c, CSRConstraint): + csr.eliminate_zeros() + csr, b = scale_rows_and_cols(csr, _row_scaling(c), b) reg_csrs.append(csr) reg_b.append(b) reg_sense.append(sense) diff --git a/linopy/solvers.py b/linopy/solvers.py index 6e1ec74d0..600e3d6d6 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -2799,10 +2799,8 @@ def _build_solver_model( rowind = np.empty(0, dtype=np.int64) rowcoef = np.empty(0, dtype=float) - lb = np.asarray(M.lb, dtype=float) - ub = np.asarray(M.ub, dtype=float) - np.place(lb, np.isneginf(lb), -xpress.infinity) - np.place(ub, np.isposinf(ub), xpress.infinity) + lb = np.where(np.isneginf(M.lb), -xpress.infinity, M.lb) + ub = np.where(np.isposinf(M.ub), xpress.infinity, M.ub) rowtype: np.ndarray rhs: np.ndarray diff --git a/test/test_expressions.py b/test/test_expressions.py index 53c5a57ce..90cc15386 100644 --- a/test/test_expressions.py +++ b/test/test_expressions.py @@ -141,3 +141,19 @@ def test_expressions_solution() -> None: assert isinstance(sol, xr.Dataset) assert "double_x" in sol assert (sol["double_x"] == 4).all() + + +def test_expression_solution_maps_variable_values() -> None: + m = Model() + y = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="y") + z = m.add_variables(coords=[pd.RangeIndex(2, name="j")], name="z") + m._mock_solve() + y.solution = xr.DataArray([10.0, 20.0, 30.0], coords=y.coords) + z.solution = xr.DataArray([1.0, 2.0], coords=z.coords) + expected = 2 * y.solution + z.solution + xr.testing.assert_equal((2 * y + z).solution, expected.rename("solution")) + expected = y.solution * z.solution + xr.testing.assert_equal((y * z).solution, expected.rename("solution")) + m.remove_variables("z") + with pytest.raises(KeyError): + (2 * y + z).solution diff --git a/test/test_matrices.py b/test/test_matrices.py index 6da3eaf1d..04317f8aa 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 @@ -68,7 +69,8 @@ def test_matrices_duplicated_variables() -> None: assert np.isin(np.unique(np.array(A)), [0.0, 2.0]).all() -def test_matrices_drops_explicit_zeros() -> None: +@pytest.mark.parametrize("freeze", [False, True]) +def test_matrices_drops_explicit_zeros(freeze: bool) -> None: # https://github.com/PyPSA/linopy/issues/814 # Expressions that broadcast against a dense coordinate store one coefficient # per pair, most of them structurally zero. Those must not reach A, whose @@ -79,7 +81,9 @@ def test_matrices_drops_explicit_zeros() -> None: coeff = xr.DataArray( np.eye(4), dims=["j", "i"], coords={"j": range(4), "i": range(4)} ) - m.add_constraints((coeff * x.rename(dim_0="i")).sum("i") <= 1, name="c") + m.add_constraints( + (coeff * x.rename(dim_0="i")).sum("i") <= 1, name="c", freeze=freeze + ) A = m.matrices.A assert A is not None @@ -98,3 +102,19 @@ def test_matrices_float_c() -> None: c = m.matrices.c assert np.all(c == np.array([1.5, 1.5])) + + +def test_matrices_sol_aligned_with_vlabels() -> None: + m = Model() + i = pd.RangeIndex(3, name="i") + x = m.add_variables(coords=[i], name="x", mask=pd.Series([True, False, True], i)) + y = m.add_variables(coords=[i], name="y") + m.add_constraints(x + y >= 0, name="c") + with pytest.raises(ValueError, match="not optimized"): + m.matrices.sol + m._mock_solve() + x.solution = xr.DataArray([1.0, np.nan, 3.0], coords=[i]) + y.solution = xr.DataArray([4.0, 5.0, 6.0], coords=[i]) + M = m.matrices + np.testing.assert_array_equal(M.vlabels, [0, 2, 3, 4, 5]) + np.testing.assert_array_equal(M.sol, [1.0, 3.0, 4.0, 5.0, 6.0]) diff --git a/test/test_scaling.py b/test/test_scaling.py index 16c318421..709172c99 100644 --- a/test/test_scaling.py +++ b/test/test_scaling.py @@ -19,6 +19,38 @@ def _dense(matrix: Any) -> np.ndarray: return matrix.toarray() +@pytest.mark.parametrize("freeze", [False, True]) +def test_masked_scaling_in_matrices(freeze: bool) -> None: + m = Model() + i = pd.RangeIndex(3, name="i") + mask = xr.DataArray([True, False, True], coords=[i]) + x = m.add_variables( + 0, + 1, + coords=[i], + name="x", + mask=mask, + scaling=xr.DataArray([2.0, 3.0, 4.0], coords=[i]), + ) + y = m.add_variables(0, 1, coords=[i], name="y") + m.add_constraints( + x.where(mask) + y >= 1, + name="c", + mask=mask, + scaling=xr.DataArray([5.0, 6.0, 7.0], coords=[i]), + freeze=freeze, + ) + + matrices = m.matrices + + np.testing.assert_allclose(matrices.var_scaling, [2.0, 4.0, 1.0, 1.0, 1.0]) + np.testing.assert_allclose(matrices.b, [5.0, 7.0]) + np.testing.assert_allclose( + _dense(matrices.A), + [[5 / 2, 0.0, 5.0, 0.0, 0.0], [0.0, 7 / 4, 0.0, 0.0, 7.0]], + ) + + def test_variable_constraint_and_objective_scaling_in_matrices() -> None: m = Model() i = pd.Index(["a", "b"], name="i")