diff --git a/doc/release_notes.rst b/doc/release_notes.rst index f8464a7d..665a5023 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -19,7 +19,7 @@ Upcoming Version *Other* -* Default internal integer labels to ``int32`` (configurable via ``linopy.options["label_dtype"]``, set to ``np.int64`` for the old behavior), cutting memory ~25% and speeding up model build 10-35%. Raises ``ValueError`` if labels exceed the int32 maximum. +* 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``). * ``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 2c45c999..49cbb4a9 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -18,6 +18,7 @@ import pandas as pd import polars as pl from numpy import nan, signedinteger +from numpy.typing import DTypeLike from polars.datatypes import DataTypeClass from xarray import DataArray, Dataset, apply_ufunc, broadcast from xarray import align as xr_align @@ -292,9 +293,12 @@ def maybe_group_terms_polars(df: pl.DataFrame) -> pl.DataFrame: return df.select(keys + ["coeffs"] + rest) -def save_join(*dataarrays: DataArray, integer_dtype: bool = False) -> Dataset: +def save_join(*dataarrays: DataArray, fill_dtype: DTypeLike | None = None) -> Dataset: """ Join multiple xarray Dataarray's to a Dataset and warn if coordinates are not equal. + + If ``fill_dtype`` is given, values filled in by an outer join are set to -1 + and the arrays are cast to that dtype (used for integer label arrays). """ try: arrs = xr_align(*dataarrays, join="exact") @@ -304,8 +308,8 @@ def save_join(*dataarrays: DataArray, integer_dtype: bool = False) -> Dataset: UserWarning, ) arrs = xr_align(*dataarrays, join="outer") - if integer_dtype: - arrs = tuple([ds.fillna(-1).astype(options["label_dtype"]) for ds in arrs]) + if fill_dtype is not None: + arrs = tuple([ds.fillna(-1).astype(fill_dtype) for ds in arrs]) return Dataset({ds.name: ds for ds in arrs}) diff --git a/linopy/config.py b/linopy/config.py index 6a28d43f..5d269c4e 100644 --- a/linopy/config.py +++ b/linopy/config.py @@ -9,10 +9,6 @@ from typing import Any -import numpy as np - -_VALID_LABEL_DTYPES = {np.int32, np.int64} - class OptionSettings: """Runtime configuration knobs (e.g. display widths). Use as a context manager or set values directly via ``options(key=value)``.""" @@ -34,10 +30,6 @@ def set_value(self, **kwargs: Any) -> None: for k, v in kwargs.items(): if k not in self._defaults: raise KeyError(f"{k} is not a valid setting.") - if k == "label_dtype" and v not in _VALID_LABEL_DTYPES: - raise ValueError( - f"label_dtype must be one of {_VALID_LABEL_DTYPES}, got {v}" - ) self._current_values[k] = v def get_value(self, name: str) -> Any: @@ -70,5 +62,4 @@ def __repr__(self) -> str: options = OptionSettings( display_max_rows=14, display_max_terms=6, - label_dtype=np.int32, ) diff --git a/linopy/constraints.py b/linopy/constraints.py index ceef45c9..dbd2d2ee 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -752,7 +752,7 @@ def _to_dataset(self, nterm: int) -> Dataset: # Map active row i -> flat position in full shape via con_labels active_positions = self.active_positions coeffs_2d = np.zeros((full_size, nterm), dtype=csr.dtype) - vars_2d = np.full((full_size, nterm), -1, dtype=options["label_dtype"]) + vars_2d = np.full((full_size, nterm), -1, dtype=self._model._dtypes["labels"]) if csr.nnz > 0: row_indices = np.repeat(active_positions, counts) term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) @@ -776,7 +776,7 @@ def _to_dataset(self, nterm: int) -> Dataset: ) ds = Dataset({"coeffs": coeffs_da, "vars": vars_da}) if self._cindex is not None: - labels_flat = np.full(full_size, -1, dtype=options["label_dtype"]) + labels_flat = np.full(full_size, -1, dtype=self._model._dtypes["labels"]) labels_flat[active_positions] = self._con_labels ds = assign_multiindex_safe( ds, @@ -1933,7 +1933,7 @@ def labels(self) -> Dataset: """ return save_join( *[v.labels.rename(k) for k, v in self.items()], - integer_dtype=True, + fill_dtype=self.model._dtypes["labels"], ) @property @@ -1954,7 +1954,7 @@ def rename_term_dim(ds: DataArray) -> DataArray: return save_join( *[rename_term_dim(v.vars.rename(k)) for k, v in self.items()], - integer_dtype=True, + fill_dtype=self.model._dtypes["labels"], ) @property @@ -2191,7 +2191,7 @@ def flat(self) -> pd.DataFrame: df = pd.concat(dfs, ignore_index=True) unique_labels = df.labels.unique() map_labels = pd.Series( - np.arange(len(unique_labels), dtype=options["label_dtype"]), + np.arange(len(unique_labels), dtype=self.model._dtypes["labels"]), index=unique_labels, ) df["key"] = df.labels.map(map_labels) diff --git a/linopy/expressions.py b/linopy/expressions.py index dd8c5c6c..63ac295f 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -510,6 +510,9 @@ def to_polars(self) -> pl.DataFrame: ... def __init__(self, data: Dataset | Any | None, model: Model) -> None: from linopy.model import Model + if not isinstance(model, Model): + raise ValueError("model must be an instance of linopy.Model") + if data is None: da = xr.DataArray([], dims=[TERM_DIM]) data = Dataset({"coeffs": da, "vars": da, "const": 0.0}) @@ -533,7 +536,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: if np.issubdtype(data.vars, np.floating): data = assign_multiindex_safe( - data, vars=data.vars.fillna(-1).astype(options["label_dtype"]) + 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) @@ -561,9 +564,6 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: # TODO: add a warning here, routines should be safe against this data = data.drop_vars(drop_dims) - if not isinstance(model, Model): - raise ValueError("model must be an instance of linopy.Model") - self._model = model self._data = cast(Dataset, data) @@ -1618,7 +1618,9 @@ def sanitize(self) -> Self: linopy.LinearExpression """ if not np.issubdtype(self.vars.dtype, np.integer): - return self.assign(vars=self.vars.fillna(-1).astype(options["label_dtype"])) + return self.assign( + vars=self.vars.fillna(-1).astype(self.model._dtypes["labels"]) + ) return self @@ -2022,12 +2024,12 @@ def _simplify_row(vars_row: np.ndarray, coeffs_row: np.ndarray) -> np.ndarray: # Combined has dimensions (.., CV_DIM, TERM_DIM) # Drop terms where all vars are -1 (i.e., empty terms across all coordinates) - vars = combined.isel({CV_DIM: 0}).astype(options["label_dtype"]) + vars = combined.isel({CV_DIM: 0}).astype(self.model._dtypes["labels"]) non_empty_terms = (vars != -1).any(dim=[d for d in vars.dims if d != TERM_DIM]) combined = combined.isel({TERM_DIM: non_empty_terms}) # Extract vars and coeffs from the combined result - vars = combined.isel({CV_DIM: 0}).astype(options["label_dtype"]) + vars = combined.isel({CV_DIM: 0}).astype(self.model._dtypes["labels"]) coeffs = combined.isel({CV_DIM: 1}) # Create new dataset with simplified data diff --git a/linopy/io.py b/linopy/io.py index bccbdecf..e4859f69 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1119,6 +1119,9 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: if k in ds.attrs: setattr(m, k, ds.attrs[k]) + if max(m._xCounter, m._cCounter) > np.iinfo(np.int32).max: + m._dtypes["labels"] = np.int64 + if "_relaxed_registry" in ds.attrs: m._relaxed_registry = json.loads(ds.attrs["_relaxed_registry"]) diff --git a/linopy/model.py b/linopy/model.py index b1477b27..f814dbf1 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -13,7 +13,8 @@ from collections.abc import Callable, Mapping, Sequence from pathlib import Path from tempfile import NamedTemporaryFile, gettempdir -from typing import TYPE_CHECKING, Any, Literal, overload +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Literal, get_args, overload from warnings import warn import numpy as np @@ -35,7 +36,6 @@ replace_by_map, to_path, ) -from linopy.config import options from linopy.constants import ( GREATER_EQUAL, HELPER_DIMS, @@ -112,6 +112,8 @@ logger = logging.getLogger(__name__) +DtypeKey = Literal["labels"] + class Model: """ @@ -140,6 +142,7 @@ class Model: _termination_condition: str _xCounter: int _cCounter: int + _dtypes: dict[DtypeKey, type[np.signedinteger]] _varnameCounter: int _connameCounter: int _pwlCounter: int @@ -163,6 +166,7 @@ class Model: # TODO: move counters to Variables and Constraints class "_xCounter", "_cCounter", + "_dtypes", "_varnameCounter", "_connameCounter", "_pwlCounter", @@ -181,6 +185,25 @@ class Model: "__weakref__", ) + @staticmethod + def _resolve_dtypes( + dtypes: Mapping[DtypeKey, type[np.signedinteger]] | None, + ) -> dict[DtypeKey, type[np.signedinteger]]: + """Validate the ``dtypes`` argument and merge it onto the defaults.""" + resolved: dict[DtypeKey, type[np.signedinteger]] = {"labels": np.int32} + 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}" + ) + resolved[key] = dtype + return resolved + def __init__( self, solver_dir: str | None = None, @@ -189,6 +212,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, ) -> None: """ Initialize the linopy model. @@ -218,11 +242,21 @@ def __init__( set_names_in_solver_io : bool 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. Returns ------- linopy.Model """ + self._dtypes: dict[DtypeKey, type[np.signedinteger]] = self._resolve_dtypes( + dtypes + ) self._variables: Variables = Variables({}, model=self) self._constraints: Constraints = Constraints({}, model=self) self._objective: Objective = Objective(LinearExpression(None, self), self) @@ -495,6 +529,37 @@ def solver_dir(self, value: str | Path) -> None: raise TypeError("'solver_dir' must path-like.") self._solver_dir = Path(value) + @property + def dtypes(self) -> Mapping[DtypeKey, type[np.signedinteger]]: + """ + Read-only mapping of the model's integer dtypes. + + Currently holds only ``"labels"``, the dtype of the variable and + constraint labels, which widens to ``int64`` automatically once the + labels outgrow int32. + """ + return MappingProxyType(self._dtypes) + + def _widen_label_dtype(self) -> None: + """Widen this model's label dtype to ``int64`` (monotonic, never narrows).""" + if self._dtypes["labels"] == np.int64: + return + self._dtypes["labels"] = np.int64 + warnings.warn( + "The model exceeded the int32 label limit (~2.1 billion labels); " + "its label dtype was widened to int64. Pass " + 'dtypes={"labels": np.int64} to Model() when building large models ' + "to avoid the mid-build upcast.", + UserWarning, + stacklevel=4, + ) + + 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: + self._widen_label_dtype() + return np.arange(start, end, dtype=self._dtypes["labels"]) + @property def dataset_attrs(self) -> list[str]: return ["parameters"] @@ -825,15 +890,9 @@ def add_variables( start = self._xCounter end = start + data.labels.size - label_dtype = options["label_dtype"] - if end > np.iinfo(label_dtype).max: - raise ValueError( - f"Number of labels ({end}) exceeds the maximum value for " - f"{label_dtype.__name__} ({np.iinfo(label_dtype).max})." - ) - data.labels.values = np.arange( - start, end, dtype=options["label_dtype"] - ).reshape(data.labels.shape) + data.labels.values = self._allocate_labels(start, end).reshape( + data.labels.shape + ) self._xCounter += data.labels.size if mask is not None: @@ -978,13 +1037,7 @@ def _allocate_constraint_labels( """Assign label ranges from the constraint counter and apply an optional mask.""" start = self._cCounter end = start + data.labels.size - label_dtype = options["label_dtype"] - if end > np.iinfo(label_dtype).max: - raise ValueError( - f"Number of labels ({end}) exceeds the maximum value for " - f"{label_dtype.__name__} ({np.iinfo(label_dtype).max})." - ) - data.labels.values = np.arange(start, end, dtype=label_dtype).reshape( + data.labels.values = self._allocate_labels(start, end).reshape( data.labels.shape ) self._cCounter += data.labels.size diff --git a/linopy/variables.py b/linopy/variables.py index 54679524..c2e247bb 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -1273,7 +1273,7 @@ def ffill(self, dim: str, limit: None = None) -> Variable: .fillna(self._fill_value) ) return self.assign_multiindex_safe( - labels=data.labels.astype(options["label_dtype"]) + labels=data.labels.astype(self.model._dtypes["labels"]) ) def bfill(self, dim: str, limit: None = None) -> Variable: @@ -1301,7 +1301,7 @@ def bfill(self, dim: str, limit: None = None) -> Variable: .map(DataArray.bfill, dim=dim, limit=limit) .fillna(self._fill_value) ) - return self.assign(labels=data.labels.astype(options["label_dtype"])) + return self.assign(labels=data.labels.astype(self.model._dtypes["labels"])) def sanitize(self) -> Variable: """ @@ -1313,7 +1313,7 @@ def sanitize(self) -> Variable: """ if issubdtype(self.labels.dtype, floating): return self.assign( - labels=self.labels.fillna(-1).astype(options["label_dtype"]) + labels=self.labels.fillna(-1).astype(self.model._dtypes["labels"]) ) return self @@ -1726,7 +1726,8 @@ def labels(self) -> Dataset: Get the labels of all variables. """ return save_join( - *[v.labels.rename(k) for k, v in self.items()], integer_dtype=True + *[v.labels.rename(k) for k, v in self.items()], + fill_dtype=self.model._dtypes["labels"], ) @property @@ -2037,7 +2038,7 @@ def flat(self) -> pd.DataFrame: df = pd.concat([self[k].flat for k in self], ignore_index=True) unique_labels = df.labels.unique() map_labels = pd.Series( - np.arange(len(unique_labels), dtype=options["label_dtype"]), + np.arange(len(unique_labels), dtype=self.model._dtypes["labels"]), index=unique_labels, ) df["key"] = df.labels.map(map_labels) diff --git a/test/test_dtypes.py b/test/test_dtypes.py index b30c7eac..1284d654 100644 --- a/test/test_dtypes.py +++ b/test/test_dtypes.py @@ -1,14 +1,16 @@ """Tests for int32 default label dtype.""" +import pickle +from pathlib import Path + import numpy as np import pytest from linopy import Model -from linopy.config import options def test_default_label_dtype_is_int32() -> None: - assert options["label_dtype"] == np.int32 + assert Model().dtypes["labels"] == np.int32 def test_variable_labels_are_int32() -> None: @@ -45,31 +47,88 @@ def test_solve_with_int32_labels() -> None: assert m.objective.value == pytest.approx(25.0) -def test_overflow_guard_variables() -> None: +def test_auto_widen_variables() -> None: m = Model() m._xCounter = np.iinfo(np.int32).max - 1 - with pytest.raises(ValueError, match="exceeds the maximum"): - m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") + with pytest.warns(UserWarning, match="widened to int64"): + x = m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") + assert x.labels.dtype == np.int64 + assert m.dtypes["labels"] == np.int64 + # the widening is per-model: other models are untouched + other = Model() + y = other.add_variables(lower=0, upper=1, coords=[range(5)], name="y") + assert y.labels.dtype == np.int32 -def test_overflow_guard_constraints() -> None: +def test_auto_widen_constraints() -> None: m = Model() x = m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") m._cCounter = np.iinfo(np.int32).max - 1 - with pytest.raises(ValueError, match="exceeds the maximum"): + with pytest.warns(UserWarning, match="widened to int64"): m.add_constraints(x >= 0, name="c") + assert m.constraints["c"].labels.dtype == np.int64 + assert m.dtypes["labels"] == np.int64 + + +def test_widen_applies_to_expressions() -> None: + m = Model() + m._xCounter = np.iinfo(np.int32).max - 1 + with pytest.warns(UserWarning, match="widened to int64"): + x = m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") + assert (2 * x + 1).vars.dtype == np.int64 + + +def test_label_dtype_init_arg() -> None: + m = Model(dtypes={"labels": np.int64}) + assert m.dtypes["labels"] == np.int64 + x = m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") + assert x.labels.dtype == np.int64 + assert (2 * x + 1).vars.dtype == np.int64 + + +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] + + +def test_dtypes_init_arg_rejects_unknown_key() -> None: + with pytest.raises(ValueError, match="only supports the keys"): + Model(dtypes={"coeffs": np.int64}) # type: ignore[dict-item] -def test_label_dtype_option_int64() -> None: - with options: - options["label_dtype"] = np.int64 - m = Model() - x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") - assert x.labels.dtype == np.int64 - expr = 2 * x + 1 - assert expr.vars.dtype == np.int64 +def test_dtypes_is_read_only_mapping() -> None: + dtypes = Model().dtypes + assert set(dtypes) == {"labels"} + with pytest.raises(TypeError): + dtypes["labels"] = np.int64 # type: ignore[index] + + +def test_auto_widen_survives_netcdf(tmp_path: Path) -> None: + from linopy import read_netcdf + + m = Model() + m._xCounter = np.iinfo(np.int32).max - 1 + with pytest.warns(UserWarning, match="widened to int64"): + x = m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") + m.add_constraints(x >= 0, name="c") + path = tmp_path / "widened.nc" + m.to_netcdf(path) + + loaded = read_netcdf(path) + + assert loaded.dtypes["labels"] == np.int64 + assert loaded.variables["x"].labels.dtype == np.int64 + assert (2 * loaded.variables["x"]).vars.dtype == np.int64 + + +def test_auto_widen_survives_pickle() -> None: + m = Model() + m._xCounter = np.iinfo(np.int32).max - 1 + with pytest.warns(UserWarning, match="widened to int64"): + x = m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") + m.add_constraints(x >= 0, name="c") + loaded = pickle.loads(pickle.dumps(m)) -def test_label_dtype_rejects_invalid() -> None: - with pytest.raises(ValueError, match="label_dtype must be one of"): - options["label_dtype"] = np.float64 + assert loaded.dtypes["labels"] == np.int64 + assert (2 * loaded.variables["x"]).vars.dtype == np.int64