Implement Model Registrator

[cpaniaguam]

Implementing a model_registrator function for the HSSM constructor will provide a streamlined way to create derived models from a base model without needing to pass the entire set of metadata each time.

Advantages

• Simplified Model Creation: Users can create derived models using only partial metadata, reducing the complexity and effort required to specify all parameters repeatedly.
• Consistency: Ensures that derived models maintain consistency with the base model's configuration, minimizing the risk of errors due to missing or incorrect metadata.
• Efficiency: Saves time and reduces boilerplate code by allowing users to reuse the base model's metadata, focusing only on the changes needed for the derived model.
• Flexibility: Provides a flexible approach to model creation, enabling users to experiment with different configurations and parameters easily.
• Enhanced Usability: Improves the overall user experience by making the process of creating and managing hierarchical models more intuitive and less error-prone.

[cpaniaguam]

Here is an initial implementation using a closure to keep the state of the base model's metadata. It seems like a good and simple way to do this.

from typing import Callable

import bambi as bmb
import pandas as pd

from hssm import HSSM, load_data
from hssm.defaults import DefaultConfig


def register_model(base_model: HSSM) -> Callable[[DefaultConfig], HSSM]:
    # Extract the base model metadata
    base_metadata = base_model._init_args

    def model_generator(**partial_metadata) -> HSSM:
        updated_metadata = {**base_metadata, **partial_metadata}
        return HSSM(**updated_metadata)

    return model_generator


# Example usage
cav_data = load_data("cavanagh_theta")
base_model = HSSM(data=cav_data)
mymodel_generator = register_model(base_model)

base_model.lapse  # Uniform(lower: 0.0, upper: 20.0)
base_model.initval_jitter  # 0.01

new_model_metadata = {"initval_jitter": 0.05, "lapse": bmb.Prior("HalfNormal", sigma=1)}

my_derived_model = mymodel_generator(**new_model_metadata)

# Tests
assert isinstance(my_derived_model, HSSM)
assert my_derived_model.initval_jitter == 0.05
assert my_derived_model.lapse == bmb.Prior("HalfNormal", sigma=1)

# generate a dict from base_model._init_args without the keys in new_model_metadata
common_values = {
    key: value
    for key, value in base_model._init_args.items()
    if key not in new_model_metadata
}

# compare common_values with my_derived_model._init_args
assert all(
    [
        my_derived_model._init_args[key] == value
        for key, value in common_values.items()
        if not isinstance(value, pd.DataFrame)
    ]
)

Thoguhts? @AlexanderFengler @digicosmos86 @krishnbera

[digicosmos86]

Hmmmm it seems that you want to distill model metadata from a model that's already built, which is different from what I intend this to be.

hssm.register_model should only deal with model metadata. Basically the purpose is to allow people to register models that are not currently supported in HSSM and then use them as if they were default models. For example, if I have a new model called my_cool_model, the way to do this right now in HSSM is:

my_cool_model = hssm.HSSM(
    model="my_cool_model",
    model_config=custom_config_that_I need to define,
    loglik=my_likelihood_function,
    ...,
)

However, if you want to use this model again with different prior configurations, you have to feed the configs to hssm.HSSM again.

With hssm.register_model(), you can supply your model configs to it like this:

hssm.register_model(
    name="my_cool_model,
    model_config=custom_config,
    loglik=my_likelihood_function,
    ...,
)

Then when you build the model you can simply do:

hssm.HSSM("my_cool_model")

As if it was ddm or other models that are supported by hssm out-of-box.

Essentially hssm.register_model should only interact with model metadata, which it makes available for HSSM. I think what you are proposing has value. Maybe we can implement it in a method called derive_model for HSSM, which essentially makes a copy of the HSSM object with slightly different configurations.

Let me know if this makes sense. Happy to chat more if needed.

[AlexanderFengler]

Agree with @digicosmos86.

I think the crux here is how we think about model, and to be fair that was probably confusing.

We have "models" which are basically the likelihood functions / simulators corresponding to particular observation processes.

If a model is completely supported out of the box, a string description (e.g. 'ddm') is enough, but then we allow various extra things to facilitate passing custom models (e.g. you pass a config argument actively, and you pass the likelihood actively --> we construct a model out of those components for you).

The core of what we want to facilitate (at least how I understood it), is to clean up / encapsulate how custom models are created for two purposes:

1. Cleaner interface, and consequently a somewhat simplified hssm class.
2. Cleaner route to contributing new models. If the model is ready for 'registration' and everything checks out, it is ready for a clean PR. (more work needs to be done here ofc).