Saving timflow.steady models as json files

[VincentJilesen]

First iteration of allowing a timflow.steady model to be written to a json-file.

Models can be saved with the to_json() method. And loaded with the from_json() method. The from_json() can currently be called from every class.

[dbrakenhoff]

Hi @VincentJilesen,

Just a minor random comment. I saw you mention that stored attributes do not necessarily match input args/kwargs. I think would be good to fix that so that input is consistent with the attributes, which I think is what you would expect as a user. But I was also curious to see whether I could just capture class constructor inputs (and not worry about attributes) and store those internally, so I asked AI and then tweaked the answer a bit myself, and came up with this bit of code. Just sharing it here in case you feel it might be useful :). Looks pretty cool, since it requires so little code... but there might be some downsides I haven't thought of yet.

import functools
import inspect


def auto_capture_init(init_func):
    """Wraps __init__ to capture ONLY top-most user inputs."""

    @functools.wraps(init_func)
    def wrapper(self, *args, **kwargs):
        # Guard: Only capture on the outermost __init__ call
        if not hasattr(self, "_init_args"):
            sig = inspect.signature(init_func)
            bound = sig.bind(self, *args, **kwargs)
            bound.apply_defaults()
            self._init_args = {
                k: v
                for k, v in bound.arguments.items()
                if k not in ["self", "model", "ml"]
            }

        return init_func(self, *args, **kwargs)

    return wrapper


class StorageMixin:
    """Base class that automatically applies auto_capture_init to all subclasses."""

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        if "__init__" in cls.__dict__:
            cls.__init__ = auto_capture_init(cls.__init__)

Anyway, curious to see what you come up with :). And do feel free to suggest some cleaning up of the Timflow code by aligning inputs to attributes. That would be really welcome as well.

[VincentJilesen]

I have resolved the input arguments issues. This should also prevent de attribute problem. I have also resolved the potential mixups with storing/loading multiple models in a single script.

[dbrakenhoff]

Hi @VincentJilesen,

Thanks for working on this! I took a look at your code, and I like how concise it is, especially the to/from methods and it nicely avoids touching all other code (but I'm rethinking the benefits of that last one a little bit). But as for how it actually works, I have to admit I have no idea what it's really doing 🙈. The __new__ stuff confuses me, e.g. checking on the caller, putting args/kwargs in different class level dictionaries. So I'm wondering if we cannot modify that __new__ stuff into to something a bit simpler and easier to understand?

I earlier mentioned autocapturing the init args/kwargs with some __init_subclass__ hackery. I chatted a bit more with AI about this, and I now think we're better off wrapping classes with a decorator that does this for us. That is much more explicit (and having to edit the other files with one little line around each class is perfectly fine) and easy to understand (and according to AI also plays much nicer with more complex inheritance trees).

So having said that, I'm now thinking it could look something like this:

# in base_io.py

def decorator_for_capturing_init_args():
    # this decorator has to make sure it collects the args from the user-facing class,
    # not the parent class if both are decorated. E.g. we want to collect the args from ModelMaq, not Model
    # if the user creates a ModelMaq.
    
    # stores the captured args in the instance's _init_kwargs attribute
    # (I also think almost all arguments to user-facing elements are defined as kwargs, 
    # so we can store everything as kwargs?)
    # cls._init_kwargs = {}

class BaseIO:

    # your existing code with maybe some adjustments:
    def to_dict():
         # exclude the parent model for objects that are added to a model
         ...

    def from_dict():
         # injects the model if it is passed and object expects it
         ...

    def save():  # or to_json()
        # maybe not even necessary, since we don't really want to save components?
        ...

    def load():  # or from_json()
        # maybe not even necessary, since we don't really want to save components?
        ...

    # + your existing serialization/deserialization logic


# in model.py

@decorator_for_capturing_init_args()
class Model(BaseIO):
    
    def __init__(self, ...):
        ...

    # etc.

    def to_dict(self):
        # overrides BaseIO.to_dict to include elements and inhomogeneities
        # uses to_dict() to add individual components
        ...

    @classmethod
    def from_dict():
        # overrides BaseIO.from_dict to handle elements and inhomogeneities
        # uses from_dict() to rebuild individual components
        ...

    def save():
        # simple json dump of to_dict()
        ...

    @classmethod
    def load():
        # load json, call from_dict()
        ...


# in well.py

@decorator_for_capturing_init_args()
class Well(BaseIO):
    def __init__(self, ...):
        ...

Then all we need to do is add the decorator to user-facing classes we want to save. All classes get a few extra methods, and store their input arguments in a private instance level attribute.

Saving models becomes ml.save() and loading models would be tfs.Model.load("my_model.json"). For this latter case we need to ensure the load function identifies the correct class from the json file to actually use , so that it correctly loads e.g. ModelXsection as well. (Maybe that is ugly, and it would be nicer to have a timflow.load() function, but I would be fine with a generic load on the Model base class.)

Anyway, I'm curious to hear your ideas on this, and whether you think this design could work, or if I'm missing something :).