perf: configurable float storage dtype (default float32, ~½ coeffs memory)

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)
 

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:

linopy/expressions.py

@@ -601,18 +601,19 @@ 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)
 
         if TERM_DIM not in data.dims:
             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)
@@ -805,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:
@@ -820,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(
@@ -851,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)

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,

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:
@@ -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})
+
+
+@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