Repository: PyPSA/linopy (python)

Issue:
Serializing and deserializing linopy.Model
Hi,

I am currently exploring/trying to setup larger models in parallel (in individual processes) and pass them back to the main process. Because the individual models are fairly large but can be prepared individually and largely independend from each other. Later on linked specific instances are linked through a few additional constraints. 

However, although serializing into a pickle or dill works fine, when trying to serialize the pickle again, a recursion error is thrown and therefore, ProcessPoolExecutor cannot be used to prepare models in parallel. (I.e., ProcessPoolExecutor uses serialization to hand over data from one process to another) This can be easily checked with this example:

```
import dill
import pandas as pd

import linopy
import pickle

m = linopy.Model()
time = pd.Index(range(10), name="time")

x = m.add_variables(
    lower=0,
    coords=[time],
    name="x",
) # to be done in parallel process
y = m.add_variables(lower=0, coords=[time], name="y") # to be done in parallel process

factor = pd.Series(time, index=time) # to be done in parallel process

con1 = m.add_constraints(3 * x + 7 * y >= 10 * factor, name="con1") # to be done in parallel process
con2 = m.add_constraints(5 * x + 2 * y >= 3 * factor, name="con2") # to be done in parallel process

m.add_objective(x + 2 * y) # to be done in parallel process

with open("test.pkl", 'wb') as f:
    dill.dump(m, f)

with open("test.pkl", 'rb') as f:
    m2 = dill.load(f)

x.lower = 1 # or add whatever additional constraint
m.solve()
```

Which throws the following error:

```
Traceback (most recent call last):
  File "C:\github\test\linopy\test.py", line 29, in <module>
    m2 = dill.load(f)
         ^^^^^^^^^^^^
  File "C:\github\test\.venv\Lib\site-packages\dill\_dill.py", line 289, in load
    return Unpickler(file, ignore=ignore, **kwds).load()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\github\test\.venv\Lib\site-packages\dill\_dill.py", line 444, in load
    obj = StockUnpickler.load(self)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\github\test\.venv\Lib\site-packages\linopy\variables.py", line 1149, in __getattr__
    if name in self.data:
               ^^^^^^^^^
  File "C:\github\test\.venv\Lib\site-packages\linopy\variables.py", line 1149, in __getattr__
    if name in self.data:
               ^^^^^^^^^
  File "C:\github\test\.venv\Lib\site-packages\linopy\variables.py", line 1149, in __getattr__
    if name in self.data:
               ^^^^^^^^^
  [Previous line repeated 745 more times]
RecursionError: maximum recursion depth exceeded
```

Pull request:
Support pickling a linopy model
Closes #340 .

## Changes proposed in this Pull Request

Defining `__getstate__` and `__setstate__` helps avoid a potential recursion error during pickling, especially in classes that override `__getattr__`. Since `__getattr__` is called whenever an attribute is not found, an implicit or default pickling process might unintentionally trigger it, leading to infinite recursion. By explicitly specifying how the object's state should be serialized (`__getstate__`) and restored (`__setstate__`), the class ensures that only its actual attributes are processed, preventing unnecessary lookups and recursion errors during pickling and unpickling.


## Checklist

- [x] Code changes are sufficiently documented; i.e. new functions contain docstrings and further explanations may be given in `doc`.
- [x] Unit tests for new features were added (if applicable).
- [x] A note for the release notes `doc/release_notes.rst` of the upcoming release is included.
- [x] I consent to the release of this PR's code under the MIT license.


Patch:
diff --git a/doc/release_notes.rst b/doc/release_notes.rst
index ff3f74f..3eeaeb4 100644
--- a/doc/release_notes.rst
+++ b/doc/release_notes.rst
@@ -13,6 +13,10 @@ Version 0.5.2
 This is mainly affecting operations where single numerical items from  pandas objects
 are selected and used for multiplication.
 
+**Feature**
+
+* Support pickling models.
+
 Version 0.5.1
 --------------
 
diff --git a/linopy/constraints.py b/linopy/constraints.py
index 183c008..13c1087 100644
--- a/linopy/constraints.py
+++ b/linopy/constraints.py
@@ -756,6 +756,12 @@ class Constraints:
             f"Constraints has no attribute `{name}` or the attribute is not accessible, e.g. raises an error."
         )
 
+    def __getstate__(self) -> dict:
+        return self.__dict__
+
+    def __setstate__(self, d: dict) -> None:
+        self.__dict__.update(d)
+
     def __dir__(self) -> list[str]:
         base_attributes = list(super().__dir__())
         formatted_names = [
diff --git a/linopy/variables.py b/linopy/variables.py
index d2dd3cd..1c50441 100644
--- a/linopy/variables.py
+++ b/linopy/variables.py
@@ -1182,6 +1182,12 @@ class Variables:
             f"Variables has no attribute `{name}` or the attribute is not accessible / raises an error."
         )
 
+    def __getstate__(self) -> dict:
+        return self.__dict__
+
+    def __setstate__(self, d: dict) -> None:
+        self.__dict__.update(d)
+
     def __dir__(self) -> list[str]:
         base_attributes = list(super().__dir__())
         formatted_names = [


Test patch:
diff --git a/test/test_io.py b/test/test_io.py
index 25c3680..559f884 100644
--- a/test/test_io.py
+++ b/test/test_io.py
@@ -5,6 +5,7 @@ Created on Thu Mar 18 09:03:35 2021.
 @author: fabian
 """
 
+import pickle
 from pathlib import Path
 from typing import Union
 
@@ -109,6 +110,21 @@ def test_model_to_netcdf_with_status_and_condition(
     assert_model_equal(m, p)
 
 
+def test_pickle_model(model_with_dash_names: Model, tmp_path: Path) -> None:
+    m = model_with_dash_names
+    fn = tmp_path / "test.nc"
+    m._status = "ok"
+    m._termination_condition = "optimal"
+
+    with open(fn, "wb") as f:
+        pickle.dump(m, f)
+
+    with open(fn, "rb") as f:
+        p = pickle.load(f)
+
+    assert_model_equal(m, p)
+
+
 # skip it xarray version is 2024.01.0 due to issue https://github.com/pydata/xarray/issues/8628
 @pytest.mark.skipif(
     xr.__version__ in ["2024.1.0", "2024.1.1"],