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 :).

[VincentJilesen]

Hi @dbrakenhoff ,

the current implementation does suffer a bit from readability en implicit abstract python behaviour.

It makes use of that scripts you import during runtime are objects. And uses this to determine what objects the user setup and witch ones are setup by other objects. Using a decorator also allows for this selectin if all objects are exclusively user-setup or automated-setup. If there exist object that can be done in both ways the decorator become very complicated again. I do not currently know is this case exist but you can proberly tell me.

Finding the correct class to load is done via de init_subclass method. This uses that the classes themself are object during runtime to collect all possible object that can be created by the factory-method. This also results in that you can setup any object from the JSON from any random object. Th is means that adding a separate function for load() under timflow requires a bit of manual setup which will need to be updated any time a class is renamed or added.

Adding explicit save and load methods for each class will most likely lead to many identical methods which violates the idea of never repeating yourself (DRY). This does help with readability to know that these methods exist but makes maintaining this function much more difficult.

Using a decorator instead of putting the new logic in the baseclass itself is fine but is doesn't guarantee that the code itself becomes simpeler. The same thing appleis to the init_subclass logic.

I hope helps clarify some of your questions and provides a startpoint for the next iteration of the implentation.