Saving timflow.steady models as json files

timflow/steady/aquifer.py

@@ -12,12 +12,13 @@
 import numpy as np
 import pandas as pd
 
+from timflow.steady.base_io import BaseIO
 from timflow.steady.constant import ConstantStar
 
 __all__ = ["Aquifer", "SimpleAquifer"]
 
 
-class AquiferData:
+class AquiferData(BaseIO):
     def __init__(self, model, kaq, c, z, npor, ltype, model3d=False):
         """Initialize aquifer data.
 

timflow/steady/base_io.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import inspect
+from functools import wraps
+from importlib import import_module
+from typing import TYPE_CHECKING, TypeVar
+
+from numpy import array, ndarray
+
+if TYPE_CHECKING:
+    from timflow.steady import Model
+
+T = TypeVar("T")
+
+
+def store_input(cls: type[T]) -> type[T]:
+
+    original_init = cls.__init__
+
+    @wraps(original_init)
+    def new_init(self, *args, **kwargs) -> None:
+        original_init(self, *args, **kwargs)
+
+        model_instance: Model | None
+        if "Model" in self.__class__.__name__:
+            model_instance = self
+        else:
+            if args != ():
+                model_instance = args[0]
+            else:
+                model_instance = kwargs.pop("model", None)  # remove model ref
+                if model_instance is None:
+                    model_instance = kwargs.pop("ml", None)  # remove model ref
+        if model_instance is not None:
+            # Prevent the reference to the model object from being stored
+            # this is unused and might complicate pickling.
+            if len(args) != 0:
+                args = args[1:] # model ref always first posarg
+
+            model_instance._obj_registry.append(
+                {
+                    "class": f"{cls.__module__}.{cls.__qualname__}",
+                    "args": args,
+                    "kwargs": kwargs,
+                }
+            )
+
+    cls.__init__ = new_init
+
+    return cls
+
+
+class BaseIO:
+    @classmethod
+    def to_dict(cls, args: tuple, kwargs: dict):
+        """
+        Collect the constructor arguments into a dict.
+
+        :return: Dict with the arguments.
+        """
+        sig = inspect.signature(cls.__init__)
+        if "Model" not in cls.__name__:
+            if "model" not in kwargs or "ml" not in kwargs:
+                args = args + ("model dummy",)  # add dummy for sig.bind
+        bound = sig.bind(cls, *args, **kwargs)
+        # Reference to class for recreation
+        data = {"_type": f"{cls.__module__}.{cls.__qualname__}"}
+        data.update(
+            {
+                k: cls._serialize(v)
+                for k, v in bound.arguments.items()
+                if k not in ("model", "ml", "self")
+            }
+        )
+        return data
+
+    @classmethod
+    def _serialize(cls, value):
+        """Convert python objects to exportable types.
+
+        :param value: Object for export.
+        :return: Object in exportable form.
+        """
+        if isinstance(value, list):
+            return [cls._serialize(v) for v in value]
+        if isinstance(value, dict):
+            return {k: cls._serialize(v) for k, v in value.items()}
+        if isinstance(value, tuple):
+            return {"tuple": [cls._serialize(v) for v in value]}
+        if isinstance(value, ndarray):
+            return {"ndarray": value.tolist()}
+        return value
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        """Factory method to create an instance of this (sub)class.
+
+        :param data: Dict with parameters
+        :return: Instance of this (sub)class.
+        """
+        type_name: str = data["_type"]
+        module_name = ".".join(type_name.split(".")[:-1])
+        class_name = type_name.split(".")[-1]
+        module = import_module(module_name)
+        subclass = getattr(module, class_name)
+        sig = inspect.signature(subclass.__init__)
+        constructor_args = {}
+
+        for name in sig.parameters:
+            if name in ("model", "ml"):
+                constructor_args[name] = cls._setup_model
+            if name != "self" and name in data:
+                constructor_args[name] = cls._deserialize(data.pop(name))
+        obj = subclass(**constructor_args)
+        if cls._setup_model is None:
+            cls._setup_model = obj
+        return obj
+
+    @classmethod
+    def _deserialize(cls, value):
+        """Convert a dict of values to the right python objects.
+
+        :param value: Imported object
+        :return: Object as correct python-type.
+        """
+        if isinstance(value, dict) and "_type" in value:
+            return cls.from_dict(value)
+        if isinstance(value, dict) and "ndarray" in value:
+            return array(value["ndarray"])
+        if isinstance(value, dict) and "tuple" in value:
+            return tuple(cls._deserialize(v) for v in value["tuple"])
+        if isinstance(value, list):
+            return [cls._deserialize(v) for v in value]
+        if isinstance(value, dict):
+            return {k: cls._deserialize(v) for k, v in value.items()}
+        return value

timflow/steady/circareasink.py

@@ -10,11 +10,13 @@
 import numpy as np
 from scipy.special import i0, i1, k0, k1
 
+from timflow.steady.base_io import store_input
 from timflow.steady.element import Element
 
 __all__ = ["CircAreaSink"]
 
 
+@store_input
 class CircAreaSink(Element):
     """Class to create a circular area-sink.
 

timflow/steady/constant.py

@@ -10,6 +10,7 @@
 
 import numpy as np
 
+from timflow.steady.base_io import store_input
 from timflow.steady.element import Element
 from timflow.steady.equation import PotentialEquation
 
@@ -33,6 +34,7 @@ def __init__(
         )
         # Defined here and not in Element as other elements can have multiple parameters
         # per layers:
+        self.layer = layer
         self.nparam = 1
         self.nunknowns = 0
         self.xr = xr
@@ -71,6 +73,7 @@ def disvecinf(self, x, y, aq=None):
         return rv
 
 
+@store_input
 class Constant(ConstantBase, PotentialEquation):
     """Specify the head at one point in the model in one layer.
 
@@ -178,6 +181,7 @@ def setparams(self, sol):
 
 # class ConstantStar(Element, PotentialEquation):
 # I don't think we need the equation
+# @store_input
 class ConstantStar(Element):
     """Constant representing the particular solution inside a semi-confined aquifer.
 

timflow/steady/element.py

@@ -11,10 +11,12 @@ def initialize(self):
 
 import numpy as np
 
+from timflow.steady.base_io import BaseIO
+
 __all__ = ["Element"]
 
 
-class Element:
+class Element(BaseIO):
     """Base class for all timflow.steady elements.
 
     Elements represent physical features in the aquifer system such as wells,

timflow/steady/inhomogeneity.py

@@ -19,6 +19,7 @@
 
 from timflow.steady.aquifer import AquiferData
 from timflow.steady.aquifer_parameters import param_3d, param_maq
+from timflow.steady.base_io import store_input
 from timflow.steady.constant import ConstantInside, ConstantStar
 from timflow.steady.element import Element
 from timflow.steady.intlinesink import (
@@ -145,6 +146,7 @@ def create_elements(self):
             c.inhomelement = True
 
 
+@store_input
 class PolygonInhomMaq(PolygonInhom):
     """Create a polygonal inhomogeneity.
 
@@ -240,6 +242,7 @@ def __init__(
         )
 
 
+@store_input
 class PolygonInhom3D(PolygonInhom):
     """Create a multi-layer model object consisting of many aquifer layers.
 
@@ -545,6 +548,7 @@ def create_elements(self):
             c.inhomelement = True
 
 
+@store_input
 class BuildingPitMaq(BuildingPit):
     """Element to simulate a building pit with an impermeable wall in ModelMaq.
 
@@ -627,6 +631,7 @@ def __init__(
         )
 
 
+@store_input
 class BuildingPit3D(BuildingPit):
     """Element to simulate a building pit with an impermeable wall in Model3D.
 
@@ -917,6 +922,7 @@ def create_elements(self):
             c.inhomelement = True
 
 
+@store_input
 class LeakyBuildingPitMaq(LeakyBuildingPit):
     """Element to simulate a building pit with a leaky wall in ModelMaq.
 
@@ -1005,6 +1011,7 @@ def __init__(
         )
 
 
+@store_input
 class LeakyBuildingPit3D(LeakyBuildingPit):
     """Element to simulate a building pit with a leaky wall in Model3D.
 

timflow/steady/inhomogeneity1d.py

@@ -17,6 +17,7 @@
 
 from timflow.steady.aquifer import AquiferData
 from timflow.steady.aquifer_parameters import param_3d, param_maq
+from timflow.steady.base_io import store_input
 from timflow.steady.constant import ConstantStar
 from timflow.steady.linesink1d import FluxDiffLineSink1D, HeadDiffLineSink1D
 from timflow.steady.stripareasink import XsectionAreaSinkInhom
@@ -326,6 +327,7 @@ def plot(
         return ax
 
 
+@store_input
 class XsectionMaq(Xsection):
     """Cross-section inhomogeneity for a multi-aquifer sequence.
 
@@ -379,6 +381,7 @@ def __init__(
         N=None,
         name=None,
     ):
+        self.topboundary = topboundary
         if c is None:
             c = []
         if z is None:
@@ -395,6 +398,7 @@ def __init__(
         )
 
 
+@store_input
 class Xsection3D(Xsection):
     """Cross-section inhomogeneity consisting of stacked aquifer layers.
 
@@ -459,6 +463,7 @@ def __init__(
         N=None,
         name=None,
     ):
+        self.topboundary = topboundary
         if z is None:
             z = [1, 0]
         (
@@ -487,6 +492,7 @@ def __init__(self, model, x1, x2, kaq, c, z, npor, ltype, hstar, N, name=None):
         super().__init__(model, x1, x2, kaq, c, z, npor, ltype, hstar, N, name=name)
 
 
+@store_input
 class StripInhomMaq(XsectionMaq):
     def __init__(
         self,
@@ -510,6 +516,7 @@ def __init__(
         super().__init__(model, x1, x2, kaq, z, c, npor, topboundary, hstar, N, name)
 
 
+@store_input
 class StripInhom3D(Xsection3D):
     def __init__(
         self,

timflow/steady/linedoublet.py

@@ -13,6 +13,7 @@
 import numpy as np
 
 from timflow.bessel.besselnumba import disbesldv, potbesldv
+from timflow.steady.base_io import store_input
 from timflow.steady.controlpoints import controlpoints
 from timflow.steady.element import Element
 from timflow.steady.equation import DisvecEquation, LeakyWallEquation
@@ -185,6 +186,7 @@ def plot(self, ax=None, layer=None):
             ax.plot([self.x1, self.x2], [self.y1, self.y2], "k")
 
 
+@store_input
 class ImpermeableWall(LineDoubletHoBase, DisvecEquation):
     """Create a segment of an impermeable wall, which is simulated with a line-doublet.
 
@@ -251,6 +253,7 @@ def setparams(self, sol):
         self.parameters[:, 0] = sol
 
 
+@store_input
 class LeakyWall(LineDoubletHoBase, LeakyWallEquation):
     """Create a segment of a leaky wall, which is simulated with a line-doublet.
 
@@ -423,6 +426,7 @@ def plot(self, ax=None, layer=None):
             ax.plot(self.x, self.y, "k")
 
 
+@store_input
 class ImpermeableWallString(LineDoubletStringBase, DisvecEquation):
     """Create a string of impermeable wall segments consisting of line-doublets.
 
@@ -473,6 +477,7 @@ def setparams(self, sol):
         self.parameters[:, 0] = sol
 
 
+@store_input
 class LeakyWallString(LineDoubletStringBase, LeakyWallEquation):
     """Create a string of leaky wall segments consisting of line-doublets.
 
@@ -525,6 +530,7 @@ def setparams(self, sol):
         self.parameters[:, 0] = sol
 
 
+@store_input
 class ImpLineDoublet(ImpermeableWall):
     """Deprecated alias for :class:`.ImpermeableWall`.
 
@@ -542,6 +548,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
 
+@store_input
 class ImpLineDoubletString(ImpermeableWallString):
     """Deprecated alias for :class:`.ImpermeableWallString`.
 
@@ -559,6 +566,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
 
+@store_input
 class LeakyLineDoublet(LeakyWall):
     """Deprecated alias for :class:`.LeakyWall`.
 
@@ -576,6 +584,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
 
+@store_input
 class LeakyLineDoubletString(LeakyWallString):
     """Deprecated alias for :class:`.LeakyWallString`.