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,154 @@
+import inspect
+import json
+
+from numpy import array, ndarray
+from typing_extensions import Self
+
+
+class BaseIO:
+    # Registry for all subclasses.
+    _class_registry = {}
+    # Registry for all created objects with their kwargs for storing.
+    _obj_registry = {}
+    # Registry for model instance for storing.
+    _model_registry = {}
+
+    def __init_subclass__(cls) -> None:
+        """Add the subclass to the registry on inheritance."""
+        cls._class_registry[cls.__name__] = cls
+
+    def __new__(cls, *args, **kwargs) -> Self:
+        """Add all newly created object to a registry if they are created directly.
+
+        :return: instance of the (sub)class
+        """
+        instance = super().__new__(cls)
+        frame = inspect.currentframe()
+        caller = frame.f_back
+        if caller.f_code.co_name == "<module>":
+            # If a new Model object create a new list before adding it.
+            if "Model" in str(cls.__name__):
+                m = f"model{len(cls._obj_registry)}"
+                cls._model_registry.update({instance: m})
+                cls._obj_registry.update({m: []})
+                cls._obj_registry[m].append((instance, args, kwargs))
+            # Other objects are added to the list of the model they have been
+            # added to.
+            else:
+                if args != ():
+                    m_inst = args[0]
+                else:
+                    m_inst = kwargs.get("model", None)
+                    if m_inst is None:
+                        m_inst = kwargs.get("ml")
+                cls._obj_registry[cls._model_registry[m_inst]].append(
+                    (instance, args, kwargs)
+                )
+        return instance
+
+    def to_json(self, filepath) -> None:
+        """
+        Write the constructor arguments to a JSON-file.
+
+        :param filepath: Filepath for the to be created JSON-file.
+        """
+        data = {}
+        i = 0
+        for item in self._obj_registry[self._model_registry[self]]:
+            obj, args, kwargs = item
+            data.update({f"object{i}": obj.to_dict(args, kwargs)})
+            i += 1
+        with open(filepath, "w") as f:
+            f.write(json.dumps(data, indent=4))
+
+    def to_dict(self, args, kwargs):
+        """
+        Collect the constructor arguments into a dict.
+
+        :return: Dict with the arguments.
+        """
+        sig = inspect.signature(self.__init__)
+        bound = sig.bind(*args, **kwargs)
+        # Reference to class for recreation
+        data = {"_type": self.__class__.__name__}
+        data.update(
+            {
+                k: self._serialize(v)
+                for k, v in bound.arguments.items()
+                if k not in ("model", "ml")
+            }
+        )
+        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, ndarray):
+            return {"ndarray": value.tolist()}
+        return value
+
+    @classmethod
+    def from_json(cls, filepath):
+        """
+        Read the constructor arguments and potential addition attributes from a JSON-file.
+
+        :param filepath: Filepath to the to be created JSON-file.
+        """
+        cls._setup_model = None
+        with open(filepath, "r") as f:
+            data: dict = json.load(f)
+        for k, v in data.items():
+            if k == "object0":  # Model object is always first created.
+                obj = cls.from_dict(v)
+                continue
+            if "obj" not in locals():  # No model in json
+                raise ImportError("No main model found in the JSON-file.")
+            cls.from_dict(v)
+        return obj
+
+    @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"]
+        subclass = cls._class_registry[type_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, 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/constant.py

@@ -33,6 +33,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

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/inhomogeneity1d.py

@@ -379,6 +379,7 @@ def __init__(
         N=None,
         name=None,
     ):
+        self.topboundary = topboundary
         if c is None:
             c = []
         if z is None:
@@ -459,6 +460,7 @@ def __init__(
         N=None,
         name=None,
     ):
+        self.topboundary = topboundary
         if z is None:
             z = [1, 0]
         (

timflow/steady/model.py

@@ -19,6 +19,7 @@
 
 from timflow.steady.aquifer import Aquifer, SimpleAquifer
 from timflow.steady.aquifer_parameters import param_3d, param_maq
+from timflow.steady.base_io import BaseIO
 from timflow.steady.constant import ConstantStar
 from timflow.steady.plots import PlotSteady
 from timflow.version import check_tqdm_parallel
@@ -42,7 +43,7 @@ def _compute_velocity_mp(args):
     return i, vv
 
 
-class Model:
+class Model(BaseIO):
     """Create a model consisting of an arbitrary sequence of aquifers and leaky layers.
 
     Notes
@@ -982,6 +983,7 @@ class ModelMaq(Model):
     """
 
     def __init__(self, kaq=1, z=None, c=None, npor=0.3, topboundary="conf", hstar=None):
+        self.topboundary = topboundary
         if c is None:
             c = []
         if z is None:
@@ -1097,6 +1099,7 @@ class ModelXsection(Model):
     """
 
     def __init__(self, naq=1):
+        self.naq = naq
         self.elementlist = []
         self.elementdict = {}  # only elements that have a label
         self.aq = SimpleAquifer(self, naq)