perf: store constraint sign as int8 codes (8x smaller sign array, lossless)

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 ``<U2`` unicode array — an 8x reduction on an array as large as the variable-labels array. ``constraint.sign`` still returns the canonical ``"="`` / ``"<="`` / ``">="`` 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)
 

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 ``<U2`` string storage. The
+    call is idempotent — already-encoded int8 input is returned unchanged — so
+    it is cheap to re-run on every constraint construction. Signs are expected
+    to be canonical already (see :func:`maybe_replace_signs`); the alternative
+    spellings are still accepted defensively.
+    """
+    already_int = np.issubdtype(sign.dtype, np.integer)
+    if dtype is np.str_:
+        return decode_signs(sign) if already_int else sign
+    if already_int:
+        return sign
+    if sign.dtype.kind not in ("U", "S"):
+        # Degenerate dtype (e.g. the float/NaN fill produced by Dataset.shift on
+        # the encoded array); leave it untouched rather than reject it.
+        return sign
+    return apply_ufunc(
+        _encode_sign_codes, sign, dask="parallelized", output_dtypes=[np.int8]
+    )
+
+
+def decode_signs(sign: DataArray) -> DataArray:
+    """
+    Decode int8 sign category codes to canonical string DataArrays.
+
+    String input (legacy ``<U2`` storage) is returned unchanged, so callers can
+    decode without first checking how the sign happens to be stored.
+    """
+    if not np.issubdtype(sign.dtype, np.integer):
+        return sign
+    return apply_ufunc(
+        _decode_sign_codes,
+        sign,
+        dask="parallelized",
+        output_dtypes=[CODE_TO_SIGN.dtype],
+    )
+
+
 def format_string_as_variable_name(name: Hashable) -> str:
     """
     Format a string to a valid python variable name.

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 ``<U2`` unicode array — an 8x
+# reduction on an array that is as large as the variable-labels array. The
+# codes are a stable part of the netcdf/pickle format: never renumber them.
+SIGN_TO_CODE: dict[str, int] = {EQUAL: 0, LESS_EQUAL: 1, GREATER_EQUAL: 2}
+CODE_TO_SIGN: np.ndarray = np.array([EQUAL, LESS_EQUAL, GREATER_EQUAL], dtype="<U2")
+
 STASHED_LOWER = "_stashed_lower"
 STASHED_UPPER = "_stashed_upper"
 STASHED_ATTRS: list[str] = [STASHED_LOWER, STASHED_UPPER]

linopy/constraints.py

@@ -42,6 +42,8 @@
     check_has_nulls_polars,
     coords_from_dataset,
     coords_to_dataset_vars,
+    decode_signs,
+    encode_signs,
     filter_nulls_polars,
     format_coord,
     format_single_constraint,
@@ -61,10 +63,12 @@
 )
 from linopy.config import options
 from linopy.constants import (
+    CODE_TO_SIGN,
     EQUAL,
     GREATER_EQUAL,
     HELPER_DIMS,
     LESS_EQUAL,
+    SIGN_TO_CODE,
     TERM_DIM,
     PerformanceWarning,
     SIGNS_pretty,
@@ -441,6 +445,8 @@ def flat(self) -> 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 +1177,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 +1262,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 +1342,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,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 = ds["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 ``<U2`` array (which would wipe out the
+        # int8 storage saving; see #843). ``self.data["sign"]`` is int8 codes by
+        # default, or ``<U2`` strings under ``dtypes={"sign": np.str_}``.
+        sign_stored = self.data["sign"].values
+        sign_flat = np.broadcast_to(sign_stored, ds["labels"].shape).reshape(-1)
+        sign_enum = pl.Enum(["=", "<=", ">="])
         all_same_sign = len(sign_flat) > 0 and (
             sign_flat[0] == sign_flat[-1] and (sign_flat[0] == sign_flat).all()
         )
@@ -1749,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"]]
@@ -1776,14 +1813,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:

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