perf(int32): auto-widen labels to int64 on overflow instead of raising

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`` (configurable via ``linopy.options["label_dtype"]``, set to ``np.int64`` for the old behavior), cutting memory ~25% and speeding up model build 10-35%. Labels auto-widen to ``int64`` once a model exceeds the int32 maximum (~2.1 billion labels), emitting a ``UserWarning`` that recommends setting ``linopy.options["label_dtype"] = np.int64`` explicitly for large models; the widening is permanent for the process and can be reverted with ``linopy.options["label_dtype"] = np.int32``.
 * ``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/config.py

@@ -7,6 +7,7 @@
 
 from __future__ import annotations
 
+import warnings
 from typing import Any
 
 import numpy as np
@@ -46,6 +47,22 @@ def get_value(self, name: str) -> Any:
         else:
             raise KeyError(f"{name} is not a valid setting.")
 
+    def widen_label_dtype(self) -> None:
+        """Permanently widen ``label_dtype`` to ``int64`` (monotonic, survives ``reset``)."""
+        transitioned = self._current_values["label_dtype"] != np.int64
+        self._defaults["label_dtype"] = np.int64
+        self._current_values["label_dtype"] = np.int64
+        if transitioned:
+            warnings.warn(
+                "Auto-widened the internal label dtype to int64: the model "
+                "exceeded the int32 label limit (~2.1 billion). Set "
+                "linopy.options['label_dtype'] = np.int64 explicitly when large "
+                "models are expected -- it avoids the mid-build upcast and is "
+                "faster and lighter on memory.",
+                UserWarning,
+                stacklevel=2,
+            )
+
     def reset(self) -> None:
         self._current_values = self._defaults.copy()
 

linopy/io.py

@@ -26,6 +26,7 @@
 
 from linopy import solvers
 from linopy.common import to_polars
+from linopy.config import options
 from linopy.constants import CONCAT_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR
 from linopy.objective import Objective
 
@@ -1119,6 +1120,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:
+        options.widen_label_dtype()
+
     if "_relaxed_registry" in ds.attrs:
         m._relaxed_registry = json.loads(ds.attrs["_relaxed_registry"])
 

linopy/model.py

@@ -825,12 +825,8 @@ 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})."
-            )
+        if end > np.iinfo(options["label_dtype"]).max:
+            options.widen_label_dtype()
         data.labels.values = np.arange(
             start, end, dtype=options["label_dtype"]
         ).reshape(data.labels.shape)
@@ -978,15 +974,11 @@ 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.shape
-        )
+        if end > np.iinfo(options["label_dtype"]).max:
+            options.widen_label_dtype()
+        data.labels.values = np.arange(
+            start, end, dtype=options["label_dtype"]
+        ).reshape(data.labels.shape)
         self._cCounter += data.labels.size
         if mask is not None:
             data.labels.values = np.where(mask.values, data.labels.values, -1)

test/test_dtypes.py

@@ -1,5 +1,8 @@
 """Tests for int32 default label dtype."""
 
+from collections.abc import Generator
+from pathlib import Path
+
 import numpy as np
 import pytest
 
@@ -45,19 +48,55 @@ def test_solve_with_int32_labels() -> None:
     assert m.objective.value == pytest.approx(25.0)
 
 
-def test_overflow_guard_variables() -> None:
+@pytest.fixture
+def restore_label_dtype() -> Generator[None, None, None]:
+    yield
+    options._defaults["label_dtype"] = np.int32
+    options["label_dtype"] = np.int32
+
+
+def test_auto_widen_variables(restore_label_dtype: None) -> 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="Auto-widened"):
+        x = m.add_variables(lower=0, upper=1, coords=[range(5)], name="x")
+    assert x.labels.dtype == np.int64
+    assert options["label_dtype"] == np.int64
 
 
-def test_overflow_guard_constraints() -> None:
+def test_auto_widen_constraints(restore_label_dtype: None) -> 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"):
-        m.add_constraints(x >= 0, name="c")
+    m.add_constraints(x >= 0, name="c")
+    assert m.constraints["c"].labels.dtype == np.int64
+    assert options["label_dtype"] == np.int64
+
+
+def test_auto_widen_is_sticky(restore_label_dtype: None) -> None:
+    options.widen_label_dtype()
+    options.reset()
+    assert options["label_dtype"] == np.int64
+
+
+def test_auto_widen_survives_netcdf(restore_label_dtype: None, tmp_path: Path) -> None:
+    from linopy import read_netcdf
+
+    m = Model()
+    m._xCounter = np.iinfo(np.int32).max - 1
+    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)
+
+    options._defaults["label_dtype"] = np.int32
+    options["label_dtype"] = np.int32
+    with pytest.warns(UserWarning, match="Auto-widened"):
+        loaded = read_netcdf(path)
+
+    assert options["label_dtype"] == np.int64
+    assert loaded.variables["x"].labels.dtype == np.int64
+    assert (2 * loaded.variables["x"]).vars.dtype == np.int64
 
 
 def test_label_dtype_option_int64() -> None: