From b1dd4de0bbaba045c64c5304d7b49210ffb08e60 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:46:07 +0200 Subject: [PATCH 1/3] perf(sign): store constraint sign as int8 category codes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The constraint sign was stored as a numpy `=`). Encode it internally as int8 category codes (1 byte) — an 8x reduction on an array as large as the variable-labels array, mirroring the int32-labels win. The mapping is exact/categorical, so there is no precision risk. `constraint.sign` still returns the canonical string DataArrays: the codes are decoded at the public `.sign` accessor. The sign is encoded once, idempotently, in `Constraint.__init__` and `Constraint._update_data`; all raw readers of the stored sign array were audited and routed through decoding (LP indicator writer in io.py, `Constraint.to_polars`/`.flat`, the infinite-RHS check in `add_constraints`, and the dtype-aware `reindex`/`reindex_like` fill). Default is ON (`Model.dtypes["sign"] == np.int8`); opt out with `Model(dtypes={"sign": np.str_})` to restore the legacy ` --- doc/release_notes.rst | 1 + linopy/common.py | 65 +++++++++++++++++++++++++++++++ linopy/constants.py | 9 +++++ linopy/constraints.py | 44 +++++++++++++++++++-- linopy/io.py | 2 +- linopy/model.py | 77 +++++++++++++++++++++++-------------- test/test_dtypes.py | 89 ++++++++++++++++++++++++++++++++++++++++++- 7 files changed, 252 insertions(+), 35 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 8ba3a0ee..a7933a76 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 constraint sign internally as ``int8`` category codes instead of an 8-byte ``="`` strings (decoded at the public boundary), so consumers are unaffected. Opt out with ``Model(dtypes={"sign": np.str_})`` to restore the legacy string storage (mainly to A/B measure the saving). * ``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/common.py b/linopy/common.py index 49cbb4a9..c4bf5489 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -27,6 +27,8 @@ from linopy.config import options from linopy.constants import ( + CODE_TO_SIGN, + SIGN_TO_CODE, SIGNS, SIGNS_alternative, SIGNS_pretty, @@ -97,6 +99,69 @@ def maybe_replace_signs(sign: DataArray) -> DataArray: return apply_ufunc(func, sign, dask="parallelized", output_dtypes=[sign.dtype]) +def _encode_sign_codes(values: np.ndarray) -> np.ndarray: + """Map a numpy array of sign strings to int8 category codes.""" + codes = np.full(values.shape, -1, dtype=np.int8) + for sign, code in SIGN_TO_CODE.items(): + codes[values == sign] = code + for alt, canonical in sign_replace_dict.items(): + codes[values == alt] = SIGN_TO_CODE[canonical] + if (codes == -1).any(): + invalid = np.unique(values[codes == -1]).tolist() + raise ValueError( + f"Invalid constraint sign(s) {invalid}; " + f"expected one of {sorted(SIGN_TO_CODE)}." + ) + return codes + + +def _decode_sign_codes(codes: np.ndarray) -> np.ndarray: + """Map int8 category codes back to canonical sign strings.""" + return CODE_TO_SIGN[codes] + + +def encode_signs(sign: DataArray, dtype: DTypeLike = np.int8) -> DataArray: + """ + Encode canonical sign strings as compact int8 category codes. + + ``dtype`` selects the storage format: ``np.int8`` (default) stores 1-byte + category codes, ``np.str_`` keeps the legacy `` DataArray: + """ + Decode int8 sign category codes to canonical string DataArrays. + + String input (legacy `` str: """ Format a string to a valid python variable name. diff --git a/linopy/constants.py b/linopy/constants.py index 7936ef1c..d974be1a 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -37,6 +37,15 @@ class PerformanceWarning(UserWarning): short_LESS_EQUAL: LESS_EQUAL, } +# Compact int8 category codes for the constraint sign. +# +# The sign field takes only three canonical values, so it is stored internally +# as 1-byte category codes instead of an 8-byte `` pd.DataFrame: stacklevel=2, ) ds = self.data + if "sign" in ds: + ds = ds.assign(sign=decode_signs(ds["sign"])) def mask_func(data: dict) -> pd.Series: mask = (data["vars"] != -1) & (data["coeffs"] != 0) @@ -1171,6 +1176,14 @@ def __init__( if not skip_broadcast: (data,) = xr.broadcast(data, exclude=[TERM_DIM]) + # Store the sign as compact int8 category codes (decoded back to + # strings by the ``.sign`` accessor). Idempotent: a no-op when the + # sign is already in the model's configured storage dtype. + sign_da = data["sign"] + encoded_sign = encode_signs(sign_da, model._dtypes["sign"]) + if encoded_sign is not sign_da: + data = assign_multiindex_safe(data, sign=encoded_sign) + self._assigned = "labels" in data self._data = data self._model = model @@ -1248,7 +1261,7 @@ def vars(self, value: variables.Variable | DataArray) -> None: @property def sign(self) -> DataArray: - return self.data.sign + return decode_signs(self.data.sign) @sign.setter def sign(self, value: SignLike) -> None: @@ -1328,6 +1341,10 @@ def _update_data(self, **fields: Any) -> None: Writes that touch the lhs structure (``coeffs``, ``vars``) flip ``_coef_dirty``. Other fields (``rhs``, ``sign``, …) leave it alone. """ + if "sign" in fields: + fields["sign"] = encode_signs( + DataArray(fields["sign"]), self._model._dtypes["sign"] + ) self._data = assign_multiindex_safe(self.data, **fields) if "coeffs" in fields or "vars" in fields: self._coef_dirty = True @@ -1738,7 +1755,7 @@ def to_polars(self) -> pl.DataFrame: labels_masked = labels_flat[mask] rhs_flat = np.broadcast_to(ds["rhs"].values, ds["labels"].shape).reshape(-1) - sign_values = ds["sign"].values + sign_values = self.sign.values sign_flat = np.broadcast_to(sign_values, ds["labels"].shape).reshape(-1) all_same_sign = len(sign_flat) > 0 and ( sign_flat[0] == sign_flat[-1] and (sign_flat[0] == sign_flat).all() @@ -1776,14 +1793,33 @@ def to_polars(self) -> pl.DataFrame: shift = conwrap(Dataset.shift) swap_dims = conwrap(Dataset.swap_dims) set_index = conwrap(Dataset.set_index) - reindex = conwrap(Dataset.reindex, fill_value=FILL_VALUE) - reindex_like = conwrap(Dataset.reindex_like, fill_value=FILL_VALUE) rename = conwrap(Dataset.rename) rename_dims = conwrap(Dataset.rename_dims) roll = conwrap(Dataset.roll) stack = conwrap(Dataset.stack) unstack = conwrap(Dataset.unstack) + def _reindex_fill_value(self) -> dict[str, Any]: + """``FILL_VALUE`` with the sign fill matched to the current storage dtype.""" + sign_fill: str | int = ( + EQUAL + if np.issubdtype(self.data["sign"].dtype, np.str_) + else SIGN_TO_CODE[EQUAL] + ) + return {**FILL_VALUE, "sign": sign_fill} + + def reindex(self, *args: Any, **kwargs: Any) -> Constraint: + """Wrapper for xarray ``Dataset.reindex`` for linopy.Constraint.""" + kwargs.setdefault("fill_value", self._reindex_fill_value()) + return self.__class__(self.data.reindex(*args, **kwargs), self.model, self.name) + + def reindex_like(self, *args: Any, **kwargs: Any) -> Constraint: + """Wrapper for xarray ``Dataset.reindex_like`` for linopy.Constraint.""" + kwargs.setdefault("fill_value", self._reindex_fill_value()) + return self.__class__( + self.data.reindex_like(*args, **kwargs), self.model, self.name + ) + @dataclass(repr=False) class Constraints: diff --git a/linopy/io.py b/linopy/io.py index 462fa5b8..8fe9255f 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -509,7 +509,7 @@ def indicator_constraints_to_file( ).flatten() coeffs_flat = ic_data.coeffs.values.reshape(len(labels_flat), -1) vars_flat = ic_data.vars.values.reshape(len(labels_flat), -1) - sign_flat = np.broadcast_to(ic_data.sign.values, labels_flat.shape).flatten() + sign_flat = np.broadcast_to(con.sign.values, labels_flat.shape).flatten() rhs_flat = np.broadcast_to(ic_data.rhs.values, labels_flat.shape).flatten() for i in range(len(labels_flat)): diff --git a/linopy/model.py b/linopy/model.py index 24594c9c..2b229b1d 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 @@ -32,6 +32,7 @@ from linopy.common import ( assign_multiindex_safe, best_int, + decode_signs, maybe_replace_signs, replace_by_map, to_path, @@ -112,7 +113,8 @@ logger = logging.getLogger(__name__) -DtypeKey = Literal["labels"] +DtypeKey = Literal["labels", "sign"] +DtypeValue = type[np.signedinteger] | type[np.str_] class Model: @@ -142,7 +144,7 @@ class Model: _termination_condition: str _xCounter: int _cCounter: int - _dtypes: dict[DtypeKey, type[np.signedinteger]] + _dtypes: dict[DtypeKey, DtypeValue] _varnameCounter: int _connameCounter: int _pwlCounter: int @@ -187,20 +189,23 @@ class Model: @staticmethod def _resolve_dtypes( - dtypes: Mapping[DtypeKey, type[np.signedinteger]] | None, - ) -> dict[DtypeKey, type[np.signedinteger]]: + dtypes: Mapping[DtypeKey, DtypeValue] | None, + ) -> dict[DtypeKey, DtypeValue]: """Validate the ``dtypes`` argument and merge it onto the defaults.""" - resolved: dict[DtypeKey, type[np.signedinteger]] = {"labels": np.int32} + resolved: dict[DtypeKey, DtypeValue] = {"labels": np.int32, "sign": np.int8} + allowed: dict[DtypeKey, tuple[DtypeValue, ...]] = { + "labels": (np.int32, np.int64), + "sign": (np.int8, np.str_), + } 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): - raise ValueError( - f"dtypes[{key!r}] must be np.int32 or np.int64, got {dtype}" - ) + if dtype not in allowed[key]: + names = " or ".join(f"np.{d.__name__}" for d in allowed[key]) + raise ValueError(f"dtypes[{key!r}] must be {names}, got {dtype}") resolved[key] = dtype return resolved @@ -212,7 +217,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, DtypeValue] | None = None, ) -> None: """ Initialize the linopy model. @@ -243,20 +248,28 @@ 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. + Storage dtypes for the model's data, exposed read-only as + ``Model.dtypes``. Two keys are supported: + + * ``"labels"`` — the variable and constraint label dtype, + ``np.int32`` (default) or ``np.int64``. 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. + * ``"sign"`` — how the constraint sign is stored, ``np.int8`` + (default, compact 1-byte category codes) or ``np.str_`` (legacy + `` None: self._solver_dir = Path(value) @property - def dtypes(self) -> Mapping[DtypeKey, type[np.signedinteger]]: + def dtypes(self) -> Mapping[DtypeKey, DtypeValue]: """ - Read-only mapping of the model's integer dtypes. + Read-only mapping of the model's storage 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 ``"sign"`` — how the constraint sign is stored (``np.int8`` + category codes by default, ``np.str_`` for legacy `` 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"]) + label_dtype = cast(type[np.signedinteger], self._dtypes["labels"]) + return np.arange(start, end, dtype=label_dtype) @property def dataset_attrs(self) -> list[str]: @@ -1149,9 +1165,12 @@ def add_constraints( data = self._constraint_data_from_lhs(lhs, sign, rhs, coords) + # ``data.sign`` is stored as int8 category codes; decode before the + # string comparison below. + data_sign = decode_signs(data.sign) invalid_infinity_values = ( - (data.sign == LESS_EQUAL) & (data.rhs == -np.inf) - ) | ((data.sign == GREATER_EQUAL) & (data.rhs == np.inf)) # noqa: F821 + (data_sign == LESS_EQUAL) & (data.rhs == -np.inf) + ) | ((data_sign == GREATER_EQUAL) & (data.rhs == np.inf)) # noqa: F821 if invalid_infinity_values.any(): raise ValueError(f"Constraint {name} contains incorrect infinite values.") diff --git a/test/test_dtypes.py b/test/test_dtypes.py index 1284d654..adf5f190 100644 --- a/test/test_dtypes.py +++ b/test/test_dtypes.py @@ -98,11 +98,98 @@ 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", "sign"} with pytest.raises(TypeError): dtypes["labels"] = np.int64 # type: ignore[index] +def test_default_sign_dtype_is_int8() -> None: + assert Model().dtypes["sign"] == np.int8 + + +def test_sign_stored_as_int8_but_read_as_strings() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + m.add_constraints(x <= 5, name="c") + con = m.constraints["c"] + # Stored compactly as int8 category codes ... + assert con.data["sign"].dtype == np.int8 + # ... but decoded to canonical strings at the public boundary. + assert con.sign.dtype.kind == "U" + assert set(np.unique(con.sign.values)) == {"<="} + + +@pytest.mark.parametrize( + ("build", "expected"), + [ + (lambda x: x <= 5, "<="), + (lambda x: x >= 1, ">="), + (lambda x: x == 3, "="), + ], +) +def test_all_three_senses_round_trip(build, expected) -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + m.add_constraints(build(x), name="c") + con = m.constraints["c"] + assert con.data["sign"].dtype == np.int8 + assert set(np.unique(con.sign.values)) == {expected} + + +def test_sign_int8_is_eightfold_smaller() -> None: + m = Model() + x = m.add_variables(lower=0, upper=1, coords=[range(100_000)], name="x") + m.add_constraints(x <= 1, name="c") + compact = m.constraints["c"].data["sign"] + assert compact.dtype == np.int8 + legacy_nbytes = m.constraints["c"].sign.astype(" None: + m = Model(dtypes={"sign": np.str_}) + assert m.dtypes["sign"] == np.str_ + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + m.add_constraints(x <= 5, name="c") + con = m.constraints["c"] + assert con.data["sign"].dtype.kind == "U" + assert set(np.unique(con.sign.values)) == {"<="} + + +def test_sign_dtype_init_arg_rejects_invalid() -> None: + with pytest.raises(ValueError, match="dtypes\\['sign'\\] must be"): + Model(dtypes={"sign": np.float64}) # type: ignore[dict-item] + + +def test_sign_update_round_trips() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + m.add_constraints(x <= 5, name="c") + con = m.constraints["c"] + con.update(sign=">=") + assert con.data["sign"].dtype == np.int8 + assert set(np.unique(con.sign.values)) == {">="} + + +@pytest.mark.skipif( + not pytest.importorskip("highspy", reason="highspy not installed"), + reason="highspy not installed", +) +def test_mixed_senses_solve_correctly() -> 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_constraints(x >= 2, name="c2") + m.add_constraints(y == 4, name="c3") + m.add_objective(x + 2 * y, sense="max") + m.solve("highs") + # x = 10, y = 4 -> 10 + 8 = 18 + assert m.objective.value == pytest.approx(18.0) + assert m.solution["x"].item() == pytest.approx(10.0) + assert m.solution["y"].item() == pytest.approx(4.0) + + def test_auto_widen_survives_netcdf(tmp_path: Path) -> None: from linopy import read_netcdf From f96cd4408fa4a46364af2cde64545e4b9ca3205f Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:26:47 +0200 Subject: [PATCH 2/3] perf(sign): export int8 codes without materialising a pl.Enum via replace_strict in the mixed case. The full --- linopy/constraints.py | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 8c56aa05..dc611bb4 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -63,6 +63,7 @@ ) from linopy.config import options from linopy.constants import ( + CODE_TO_SIGN, EQUAL, GREATER_EQUAL, HELPER_DIMS, @@ -1755,8 +1756,14 @@ def to_polars(self) -> pl.DataFrame: labels_masked = labels_flat[mask] rhs_flat = np.broadcast_to(ds["rhs"].values, ds["labels"].shape).reshape(-1) - sign_values = self.sign.values - sign_flat = np.broadcast_to(sign_values, ds["labels"].shape).reshape(-1) + # Keep the sign as its compact storage codes here and only turn them + # into the "=" / "<=" / ">=" strings at the last moment, so the export + # never materialises a full-size ``="]) all_same_sign = len(sign_flat) > 0 and ( sign_flat[0] == sign_flat[-1] and (sign_flat[0] == sign_flat).all() ) @@ -1766,14 +1773,27 @@ def to_polars(self) -> pl.DataFrame: "rhs": rhs_flat[mask], } if all_same_sign: + first = sign_flat[0] + first_str = ( + CODE_TO_SIGN[first] + if np.issubdtype(sign_flat.dtype, np.integer) + else first + ) short = pl.DataFrame(short_data).with_columns( - pl.lit(sign_flat[0]).cast(pl.Enum(["=", "<=", ">="])).alias("sign") + pl.lit(first_str).cast(sign_enum).alias("sign") ) - else: + elif np.issubdtype(sign_flat.dtype, np.integer): short_data["sign"] = pl.Series( - "sign", sign_flat[mask], dtype=pl.Enum(["=", "<=", ">="]) + "sign", sign_flat[mask], dtype=pl.Int8 + ).replace_strict( + list(range(len(CODE_TO_SIGN))), + CODE_TO_SIGN.tolist(), + return_dtype=sign_enum, ) short = pl.DataFrame(short_data) + else: + short_data["sign"] = pl.Series("sign", sign_flat[mask], dtype=sign_enum) + short = pl.DataFrame(short_data) df = long.join(short, on="labels", how="inner") return df[["labels", "coeffs", "vars", "sign", "rhs"]] From ba58b0dbbd0ea716b087dc184b964c5efe7cf804 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:47:54 +0200 Subject: [PATCH 3/3] test(sign): annotate parametrized test params for mypy CI runs mypy over test/ with disallow_untyped_defs; annotate the build/expected params of test_all_three_senses_round_trip. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/test_dtypes.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/test/test_dtypes.py b/test/test_dtypes.py index adf5f190..3fb71cd9 100644 --- a/test/test_dtypes.py +++ b/test/test_dtypes.py @@ -1,7 +1,9 @@ """Tests for int32 default label dtype.""" import pickle +from collections.abc import Callable from pathlib import Path +from typing import Any import numpy as np import pytest @@ -127,7 +129,9 @@ def test_sign_stored_as_int8_but_read_as_strings() -> None: (lambda x: x == 3, "="), ], ) -def test_all_three_senses_round_trip(build, expected) -> None: +def test_all_three_senses_round_trip( + build: Callable[[Any], Any], expected: str +) -> None: m = Model() x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") m.add_constraints(build(x), name="c")