refactor(int32): per-model label dtype instead of sticky process-global

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(label_dtype=np.int64)`` upfront to avoid the mid-build upcast.
 * ``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/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})
 
 

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,
 )

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._label_dtype)
         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._label_dtype)
             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._label_dtype,
         )
 
     @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._label_dtype,
         )
 
     @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._label_dtype),
             index=unique_labels,
         )
         df["key"] = df.labels.map(map_labels)

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._label_dtype)
             )
         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._label_dtype)
+            )
 
         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._label_dtype)
         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._label_dtype)
         coeffs = combined.isel({CV_DIM: 1})
 
         # Create new dataset with simplified data

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._label_dtype = np.int64
+
     if "_relaxed_registry" in ds.attrs:
         m._relaxed_registry = json.loads(ds.attrs["_relaxed_registry"])
 

linopy/model.py

@@ -35,7 +35,6 @@
     replace_by_map,
     to_path,
 )
-from linopy.config import options
 from linopy.constants import (
     GREATER_EQUAL,
     HELPER_DIMS,
@@ -140,6 +139,7 @@ class Model:
     _termination_condition: str
     _xCounter: int
     _cCounter: int
+    _label_dtype: type[np.signedinteger]
     _varnameCounter: int
     _connameCounter: int
     _pwlCounter: int
@@ -163,6 +163,7 @@ class Model:
         # TODO: move counters to Variables and Constraints class
         "_xCounter",
         "_cCounter",
+        "_label_dtype",
         "_varnameCounter",
         "_connameCounter",
         "_pwlCounter",
@@ -189,6 +190,7 @@ def __init__(
         auto_mask: bool = False,
         freeze_constraints: bool = False,
         set_names_in_solver_io: bool = True,
+        label_dtype: type[np.signedinteger] = np.int32,
     ) -> None:
         """
         Initialize the linopy model.
@@ -218,11 +220,22 @@ def __init__(
         set_names_in_solver_io : bool
             Whether direct solver exports should include variable and
             constraint names by default. The default is True.
+        label_dtype : np.int32 or np.int64, optional
+            Integer dtype used for the model's variable and constraint labels.
+            The default ``np.int32`` halves label memory but caps a model at
+            ~2.1 billion labels; the model widens itself to ``np.int64``
+            automatically when that limit is hit. Pass ``np.int64`` upfront
+            for very large models to avoid the mid-build upcast.
 
         Returns
         -------
         linopy.Model
         """
+        if label_dtype not in (np.int32, np.int64):
+            raise ValueError(
+                f"label_dtype must be np.int32 or np.int64, got {label_dtype}"
+            )
+        self._label_dtype: type[np.signedinteger] = label_dtype
         self._variables: Variables = Variables({}, model=self)
         self._constraints: Constraints = Constraints({}, model=self)
         self._objective: Objective = Objective(LinearExpression(None, self), self)
@@ -495,6 +508,36 @@ def solver_dir(self, value: str | Path) -> None:
             raise TypeError("'solver_dir' must path-like.")
         self._solver_dir = Path(value)
 
+    @property
+    def label_dtype(self) -> type[np.signedinteger]:
+        """
+        Integer dtype used for this model's variable and constraint labels.
+
+        Set via the ``label_dtype`` argument of ``Model()`` and automatically
+        widened to ``int64`` once the labels outgrow int32.
+        """
+        return self._label_dtype
+
+    def _widen_label_dtype(self) -> None:
+        """Widen this model's label dtype to ``int64`` (monotonic, never narrows)."""
+        if self._label_dtype == np.int64:
+            return
+        self._label_dtype = np.int64
+        warnings.warn(
+            "The model exceeded the int32 label limit (~2.1 billion labels); "
+            "its label dtype was widened to int64. Pass label_dtype=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._label_dtype).max:
+            self._widen_label_dtype()
+        return np.arange(start, end, dtype=self._label_dtype)
+
     @property
     def dataset_attrs(self) -> list[str]:
         return ["parameters"]
@@ -825,15 +868,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 +1015,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

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._label_dtype)
         )
 
     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._label_dtype))
 
     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._label_dtype)
             )
         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._label_dtype,
         )
 
     @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._label_dtype),
             index=unique_labels,
         )
         df["key"] = df.labels.map(map_labels)