From 59bff2768ca1d5df81480bb49184badc8291fc83 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:16:49 +0200 Subject: [PATCH 1/3] perf: store float arrays in configurable per-model "values" dtype (default float32) Extend the per-model dtypes mechanism (Model._dtypes) with a new "values" key controlling the storage dtype of the model's float arrays: coefficients, constants, right-hand sides and variable lower/upper bounds. coeffs is the single largest array in a built model, so defaulting to float32 roughly halves peak build memory. Solvers upcast to double internally, so solutions are unaffected; the LP/MPS writers and direct backends handle float32 fine. - model.py: DtypeKey gains "values"; annotations widened to type[np.number]; _resolve_dtypes validates "values" as np.float32/np.float64 (default np.float32) and "labels" as before; add_variables casts bounds to the values dtype; docstrings updated. - expressions.py: LinearExpression.__init__ casts coeffs and const to the model's values dtype (all expression construction funnels through here, including simplify() results). - constraints.py: Constraint.__init__ casts rhs to the model's values dtype. Pass Model(dtypes={"values": np.float64}) to restore full float64 storage. The float32 default is TEMPORARY, to surface the reduction in CI memory benchmarks, and is intended to flip to opt-in (float64 default) later. Co-Authored-By: Claude Opus 4.8 (1M context) --- doc/release_notes.rst | 1 + linopy/constraints.py | 4 +++ linopy/expressions.py | 11 ++++--- linopy/model.py | 71 ++++++++++++++++++++++++++++--------------- test/test_dtypes.py | 63 +++++++++++++++++++++++++++++++++++++- 5 files changed, 119 insertions(+), 31 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 8ba3a0ee..b0124747 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -20,6 +20,7 @@ Upcoming Version *Other* * Default internal integer labels to ``int32``, cutting memory ~25% and speeding up model build 10-35%. Models exceeding the int32 maximum (~2.1 billion labels) widen to ``int64`` automatically with a ``UserWarning``; pass ``Model(dtypes={"labels": np.int64})`` upfront to avoid the mid-build upcast (exposed read-only via ``Model.dtypes``). +* Store the model's float arrays (coefficients, constants, right-hand sides and variable bounds) in a per-model ``"values"`` dtype, configurable via ``Model(dtypes={"values": ...})`` and exposed read-only through ``Model.dtypes``. ``coeffs`` is the single largest array in a built model, so ``float32`` roughly halves peak build memory; pass ``Model(dtypes={"values": np.float64})`` to keep full double-precision storage. Solvers upcast to double internally, so solutions are unaffected. **Note:** the default is *temporarily* ``np.float32`` so the reduction shows up in the CI memory benchmarks; it is intended to flip to opt-in (``np.float64`` default) in a later release. * ``add_variables(binary=True, ...)`` now accepts ``lower``/``upper`` bounds, as long as they are 0 or 1. Previously binary bounds could only be set via the ``.lower``/``.upper`` setters after creation. (https://github.com/PyPSA/linopy/issues/776) * ``add_piecewise_formulation`` gained an ``active_fill`` parameter that gates a partial ``active`` (defined over a subset of the indexed dimension, or masked) as always-active (``1``) or always-off (``0``); without it, a partial ``active`` — which was previously zeroed silently — now raises. Useful when one formulation mixes gated and ungated entities (e.g. committable and non-committable units sharing a ``status``). ``active_fill`` is transitional and will be removed once v1 semantics make ``active.reindex(coords).fillna(value)`` sufficient. (https://github.com/PyPSA/linopy/issues/796) diff --git a/linopy/constraints.py b/linopy/constraints.py index dbd2d2ee..6f1810de 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -1166,6 +1166,10 @@ def __init__( if attr not in data: raise ValueError(f"missing '{attr}' in data") + values_dtype = model._dtypes["values"] + if data["rhs"].dtype != values_dtype: + data = assign_multiindex_safe(data, rhs=data["rhs"].astype(values_dtype)) + data = data.assign_attrs(name=name) if not skip_broadcast: diff --git a/linopy/expressions.py b/linopy/expressions.py index f13fc6f2..7c340d43 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -601,8 +601,9 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: data = assign_multiindex_safe( data, vars=data.vars.fillna(-1).astype(model._dtypes["labels"]) ) - if not np.issubdtype(data.coeffs, np.floating): - data["coeffs"].values = data.coeffs.values.astype(float) + values_dtype = model._dtypes["values"] + if data.coeffs.dtype != values_dtype: + data["coeffs"].values = data.coeffs.values.astype(values_dtype) data = fill_missing_coords(data) @@ -610,9 +611,9 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: raise ValueError("data must contain one dimension ending with '_term'") if "const" not in data: - data = data.assign(const=0.0) - elif not np.issubdtype(data.const, np.floating): - data = assign_multiindex_safe(data, const=data.const.astype(float)) + data = data.assign(const=values_dtype(0)) + elif data.const.dtype != values_dtype: + data = assign_multiindex_safe(data, const=data.const.astype(values_dtype)) (data,) = xr.broadcast(data, exclude=HELPER_DIMS) data = cast(Dataset, data) diff --git a/linopy/model.py b/linopy/model.py index 24594c9c..fb9c3996 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -14,7 +14,7 @@ from pathlib import Path from tempfile import NamedTemporaryFile, gettempdir from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Literal, get_args, overload +from typing import TYPE_CHECKING, Any, Literal, cast, get_args, overload from warnings import warn import numpy as np @@ -112,7 +112,7 @@ logger = logging.getLogger(__name__) -DtypeKey = Literal["labels"] +DtypeKey = Literal["labels", "values"] class Model: @@ -142,7 +142,7 @@ class Model: _termination_condition: str _xCounter: int _cCounter: int - _dtypes: dict[DtypeKey, type[np.signedinteger]] + _dtypes: dict[DtypeKey, type[np.number]] _varnameCounter: int _connameCounter: int _pwlCounter: int @@ -187,20 +187,27 @@ class Model: @staticmethod def _resolve_dtypes( - dtypes: Mapping[DtypeKey, type[np.signedinteger]] | None, - ) -> dict[DtypeKey, type[np.signedinteger]]: + dtypes: Mapping[DtypeKey, type[np.number]] | None, + ) -> dict[DtypeKey, type[np.number]]: """Validate the ``dtypes`` argument and merge it onto the defaults.""" - resolved: dict[DtypeKey, type[np.signedinteger]] = {"labels": np.int32} + resolved: dict[DtypeKey, type[np.number]] = { + "labels": np.int32, + "values": np.float32, + } for key, dtype in (dtypes or {}).items(): if key not in get_args(DtypeKey): raise ValueError( f"dtypes only supports the keys {list(get_args(DtypeKey))}, " f"got unknown key {key!r}" ) - if dtype not in (np.int32, np.int64): + if key == "labels" and dtype not in (np.int32, np.int64): raise ValueError( f"dtypes[{key!r}] must be np.int32 or np.int64, got {dtype}" ) + if key == "values" and dtype not in (np.float32, np.float64): + raise ValueError( + f"dtypes[{key!r}] must be np.float32 or np.float64, got {dtype}" + ) resolved[key] = dtype return resolved @@ -212,7 +219,7 @@ def __init__( auto_mask: bool = False, freeze_constraints: bool = False, set_names_in_solver_io: bool = True, - dtypes: Mapping[DtypeKey, type[np.signedinteger]] | None = None, + dtypes: Mapping[DtypeKey, type[np.number]] | None = None, ) -> None: """ Initialize the linopy model. @@ -243,20 +250,27 @@ def __init__( Whether direct solver exports should include variable and constraint names by default. The default is True. dtypes : mapping, optional - Integer dtypes for the model's data, exposed read-only as - ``Model.dtypes``. Only ``"labels"`` is supported, e.g. - ``Model(dtypes={"labels": np.int64})``. The default ``np.int32`` - halves label memory but caps the model at ~2.1 billion labels, - after which it widens to ``np.int64`` automatically; pass - ``np.int64`` upfront to avoid that mid-build upcast. + Per-model dtypes for the model's data, exposed read-only as + ``Model.dtypes``. Two keys are supported: + + * ``"labels"`` (``np.int32`` or ``np.int64``, default ``np.int32``): + the dtype of the variable and constraint labels. The default + halves label memory but caps the model at ~2.1 billion labels, + after which it widens to ``np.int64`` automatically; pass + ``np.int64`` upfront to avoid that mid-build upcast. + * ``"values"`` (``np.float32`` or ``np.float64``, default + ``np.float32``): the dtype of the stored float arrays + (coefficients, constants, right-hand sides and variable bounds). + ``np.float32`` roughly halves the memory of the largest arrays; + pass ``np.float64`` to keep full double-precision storage. + + e.g. ``Model(dtypes={"labels": np.int64, "values": np.float64})``. Returns ------- linopy.Model """ - self._dtypes: dict[DtypeKey, type[np.signedinteger]] = self._resolve_dtypes( - dtypes - ) + self._dtypes: dict[DtypeKey, type[np.number]] = self._resolve_dtypes(dtypes) self._variables: Variables = Variables({}, model=self) self._constraints: Constraints = Constraints({}, model=self) self._objective: Objective = Objective(LinearExpression(None, self), self) @@ -530,13 +544,14 @@ def solver_dir(self, value: str | Path) -> None: self._solver_dir = Path(value) @property - def dtypes(self) -> Mapping[DtypeKey, type[np.signedinteger]]: + def dtypes(self) -> Mapping[DtypeKey, type[np.number]]: """ - Read-only mapping of the model's integer dtypes. + Read-only mapping of the model's dtypes. - Currently holds only ``"labels"``, the dtype of the variable and - constraint labels, which widens to ``int64`` automatically once the - labels outgrow int32. + Holds ``"labels"``, the dtype of the variable and constraint labels, + which widens to ``int64`` automatically once the labels outgrow int32, + and ``"values"``, the dtype of the stored float arrays (coefficients, + constants, right-hand sides and variable bounds). """ return MappingProxyType(self._dtypes) @@ -556,7 +571,8 @@ def _widen_label_dtype(self) -> None: def _allocate_labels(self, start: int, end: int) -> np.ndarray: """Return the label range ``[start, end)``, widening the dtype on overflow.""" - if end > np.iinfo(self._dtypes["labels"]).max: + label_dtype = cast(type[np.signedinteger], self._dtypes["labels"]) + if end > np.iinfo(label_dtype).max: self._widen_label_dtype() return np.arange(start, end, dtype=self._dtypes["labels"]) @@ -854,8 +870,13 @@ def add_variables( "Semi-continuous variables require a positive scalar lower bound." ) - lower_da = broadcast_to_coords(lower, coords, label="lower bound", **kwargs) - upper_da = broadcast_to_coords(upper, coords, label="upper bound", **kwargs) + values_dtype = self._dtypes["values"] + lower_da = broadcast_to_coords( + lower, coords, label="lower bound", **kwargs + ).astype(values_dtype) + upper_da = broadcast_to_coords( + upper, coords, label="upper bound", **kwargs + ).astype(values_dtype) data = Dataset( { "lower": lower_da, diff --git a/test/test_dtypes.py b/test/test_dtypes.py index 1284d654..734497fc 100644 --- a/test/test_dtypes.py +++ b/test/test_dtypes.py @@ -98,11 +98,72 @@ def test_dtypes_init_arg_rejects_unknown_key() -> None: def test_dtypes_is_read_only_mapping() -> None: dtypes = Model().dtypes - assert set(dtypes) == {"labels"} + assert set(dtypes) == {"labels", "values"} with pytest.raises(TypeError): dtypes["labels"] = np.int64 # type: ignore[index] +def test_default_values_dtype_is_float32() -> None: + assert Model().dtypes["values"] == np.float32 + + +def test_default_bounds_are_float32() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + assert x.lower.dtype == np.float32 + assert x.upper.dtype == np.float32 + + +def test_default_coeffs_and_const_are_float32() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + expr = 2 * x + 1 + assert expr.coeffs.dtype == np.float32 + assert expr.const.dtype == np.float32 + + +def test_default_rhs_and_constraint_coeffs_are_float32() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + m.add_constraints(2 * x + 1 <= 5, name="c") + con = m.constraints["c"] + assert con.rhs.dtype == np.float32 + assert con.coeffs.dtype == np.float32 + + +def test_values_dtype_init_arg_float64() -> None: + m = Model(dtypes={"values": np.float64}) + assert m.dtypes["values"] == np.float64 + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + assert x.lower.dtype == np.float64 + assert x.upper.dtype == np.float64 + expr = 2 * x + 1 + assert expr.coeffs.dtype == np.float64 + assert expr.const.dtype == np.float64 + m.add_constraints(2 * x + 1 <= 5, name="c") + assert m.constraints["c"].rhs.dtype == np.float64 + + +def test_values_dtype_init_arg_rejects_invalid() -> None: + with pytest.raises(ValueError, match="dtypes\\['values'\\] must be"): + Model(dtypes={"values": np.float16}) # type: ignore[dict-item] + + +@pytest.mark.skipif( + not pytest.importorskip("highspy", reason="highspy not installed"), + reason="highspy not installed", +) +def test_solve_with_float32_values() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, name="x") + y = m.add_variables(lower=0, upper=10, name="y") + m.add_constraints(x + y <= 15, name="c1") + m.add_objective(x + 2 * y, sense="max") + assert m.constraints["c1"].coeffs.dtype == np.float32 + m.solve("highs") + assert m.objective.value == pytest.approx(25.0, abs=1e-4) + + def test_auto_widen_survives_netcdf(tmp_path: Path) -> None: from linopy import read_netcdf From 5ddb06e921a646ebf7f414877a26925078a8f40c Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:48:48 +0200 Subject: [PATCH 2/3] test(values): drop unused type:ignore comments for mypy CI's warn_unused_ignores flags the two # type: ignore[dict-item] on the invalid-dtype rejection tests as unnecessary; np.float64/np.float16 are valid type[np.number] values (the rejection is a runtime check). Co-Authored-By: Claude Opus 4.8 (1M context) --- test/test_dtypes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_dtypes.py b/test/test_dtypes.py index 734497fc..f0496a69 100644 --- a/test/test_dtypes.py +++ b/test/test_dtypes.py @@ -88,7 +88,7 @@ def test_label_dtype_init_arg() -> None: def test_label_dtype_init_arg_rejects_invalid() -> None: with pytest.raises(ValueError, match="dtypes\\['labels'\\] must be"): - Model(dtypes={"labels": np.float64}) # type: ignore[dict-item] + Model(dtypes={"labels": np.float64}) def test_dtypes_init_arg_rejects_unknown_key() -> None: @@ -146,7 +146,7 @@ def test_values_dtype_init_arg_float64() -> None: def test_values_dtype_init_arg_rejects_invalid() -> None: with pytest.raises(ValueError, match="dtypes\\['values'\\] must be"): - Model(dtypes={"values": np.float16}) # type: ignore[dict-item] + Model(dtypes={"values": np.float16}) @pytest.mark.skipif( From d23e0d773e67610bc2f9bc3185349b1d5eef2151 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 16 Jul 2026 14:04:58 +0200 Subject: [PATCH 3/3] perf(values): keep float32 constant ops in float32 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multiplying/adding a float64 array against a float32 expression let numpy upcast to a float64 broadcast intermediate that __init__ then downcast back to float32 — a large transient that raised op peak memory ~17-27% over master (CodSpeed test_op[expr_{mul,add}_array_bcast]). Cast the other operand to the model value dtype first (extracted as _as_value_dtype), so the elementwise op stays float32 end to end. Peak now drops below master (mul_array_bcast 1.04MB -> 0.69MB). Co-Authored-By: Claude Opus 4.8 (1M context) --- linopy/expressions.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/linopy/expressions.py b/linopy/expressions.py index 7c340d43..d8714f22 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -806,6 +806,18 @@ def _align_constant( ) return self_const, aligned, True + def _as_value_dtype(self, da: DataArray) -> DataArray: + """ + Cast ``da`` to the model's configured value dtype. + + Applied to the other operand of an elementwise op so that a float32 + expression stays float32 throughout, rather than letting numpy upcast to + a large (broadcast) float64 intermediate that ``__init__`` only downcasts + again — a transient that would double the operation's peak memory. + """ + values_dtype = self.model._dtypes["values"] + return da.astype(values_dtype) if da.dtype != values_dtype else da + def _add_constant( self, other: ConstantLike, join: JoinOptions | None = None ) -> Self: @@ -821,6 +833,7 @@ def _add_constant( ) da = da.fillna(0) self_const = self_const.fillna(0) + da = self._as_value_dtype(da) if needs_data_reindex: return self.__class__( self.data.reindex_like(self_const, fill_value=self._fill_value).assign( @@ -852,6 +865,7 @@ def _apply_constant_op( ) factor = factor.fillna(fill_value) self_const = self_const.fillna(0) + factor = self._as_value_dtype(factor) if needs_data_reindex: data = self.data.reindex_like(self_const, fill_value=self._fill_value) coeffs = data.coeffs.fillna(0)