diff --git a/timflow/steady/aquifer.py b/timflow/steady/aquifer.py index 5f884ca1..8f5a8781 100644 --- a/timflow/steady/aquifer.py +++ b/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. diff --git a/timflow/steady/base_io.py b/timflow/steady/base_io.py new file mode 100644 index 00000000..faff6e22 --- /dev/null +++ b/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 == "": + # 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 diff --git a/timflow/steady/constant.py b/timflow/steady/constant.py index dde0baba..38fd6951 100644 --- a/timflow/steady/constant.py +++ b/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 diff --git a/timflow/steady/element.py b/timflow/steady/element.py index 8095f98e..09e58f4e 100644 --- a/timflow/steady/element.py +++ b/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, diff --git a/timflow/steady/inhomogeneity1d.py b/timflow/steady/inhomogeneity1d.py index 6baa2e0c..41f1f315 100644 --- a/timflow/steady/inhomogeneity1d.py +++ b/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] ( diff --git a/timflow/steady/model.py b/timflow/steady/model.py index 30caf3f4..34c13ea4 100644 --- a/timflow/steady/model.py +++ b/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)