refac: store names in PiecewiseFormulation, add IO persistence

PiecewiseFormulation now stores variable/constraint names as strings
with a model reference. Properties return live Views on access.
This makes serialization trivial — persist as JSON in netcdf attrs,
reconstruct on load.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: FBumann

linopy/io.py

@@ -1147,6 +1147,11 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset:
     ds = ds.assign_attrs(scalars)
     if m._relaxed_registry:
         ds.attrs["_relaxed_registry"] = json.dumps(m._relaxed_registry)
+    if m._piecewise_formulations:
+        ds.attrs["_piecewise_formulations"] = json.dumps(
+            {name: {"method": pwl.method, "variables": pwl.variable_names, "constraints": pwl.constraint_names}
+             for name, pwl in m._piecewise_formulations.items()}
+        )
     ds.attrs = non_bool_dict(ds.attrs)
 
     for k in ds:
@@ -1244,6 +1249,18 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset:
     if "_relaxed_registry" in ds.attrs:
         m._relaxed_registry = json.loads(ds.attrs["_relaxed_registry"])
 
+    if "_piecewise_formulations" in ds.attrs:
+        from linopy.piecewise import PiecewiseFormulation
+
+        for name, d in json.loads(ds.attrs["_piecewise_formulations"]).items():
+            m._piecewise_formulations[name] = PiecewiseFormulation(
+                name=name,
+                method=d["method"],
+                variable_names=d["variables"],
+                constraint_names=d["constraints"],
+                model=m,
+            )
+
     return m
 
 

linopy/model.py

@@ -533,8 +533,8 @@ def _piecewise_names(self) -> set[str]:
         """Return all variable/constraint names belonging to piecewise formulations."""
         names: set[str] = set()
         for pwl in self._piecewise_formulations.values():
-            names.update(pwl.variables)
-            names.update(pwl.constraints)
+            names.update(pwl.variable_names)
+            names.update(pwl.constraint_names)
         return names
 
     def __getitem__(self, key: str) -> Variable:

linopy/piecewise.py

@@ -8,7 +8,6 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from dataclasses import dataclass
 from numbers import Real
 from typing import TYPE_CHECKING, Literal, TypeAlias
 
@@ -62,19 +61,40 @@
 # ---------------------------------------------------------------------------
 
 
-@dataclass(repr=False)
 class PiecewiseFormulation:
     """
     Result of ``add_piecewise_formulation``.
 
     Groups all auxiliary variables and constraints created by a single
-    piecewise formulation.
+    piecewise formulation. Stores only names internally; ``variables``
+    and ``constraints`` properties return live views from the model.
     """
 
-    name: str
-    method: str
-    variables: Variables
-    constraints: Constraints
+    __slots__ = ("name", "method", "variable_names", "constraint_names", "_model")
+
+    def __init__(
+        self,
+        name: str,
+        method: str,
+        variable_names: list[str],
+        constraint_names: list[str],
+        model: Model,
+    ) -> None:
+        self.name = name
+        self.method = method
+        self.variable_names = variable_names
+        self.constraint_names = constraint_names
+        self._model = model
+
+    @property
+    def variables(self) -> Variables:
+        """View of the auxiliary variables in this formulation."""
+        return self._model.variables[self.variable_names]
+
+    @property
+    def constraints(self) -> Constraints:
+        """View of the auxiliary constraints in this formulation."""
+        return self._model.constraints[self.constraint_names]
 
     def __repr__(self) -> str:
         # Collect user-facing dims with sizes (skip internal _ prefixed dims)
@@ -769,8 +789,9 @@ def add_piecewise_formulation(
     result = PiecewiseFormulation(
         name=name,
         method=resolved_method,
-        variables=model.variables[new_vars],
-        constraints=model.constraints[new_cons],
+        variable_names=new_vars,
+        constraint_names=new_cons,
+        model=model,
     )
     model._piecewise_formulations[name] = result
     return result