OpenSourceEconomics / lcm

Solution and simulation of life cycle models in Python with GPU acceleration.
Apache License 2.0
16 stars 1 forks source link

BUG: Stochastic next functions need to depend on at least one argument #39

Open timmens opened 1 year ago

timmens commented 1 year ago

This works

@lcm.mark.stochastic
def next_health(health):  # noqa: ARG001
    pass

This does not

@lcm.mark.stochastic
def next_partner():
    pass

The second version should also work. The problem currently is in the weighting function:

...
new_kwargs = [*function_parameters, "params"]

@with_signature(args=new_kwargs)
def weight_func(*args, **kwargs):
    args = all_as_args(args, kwargs, arg_names=new_kwargs)

    # by definition of new_kwargs the last argument is params
    params = args[-1]

    indices = [
        label_translators[arg_name](arg)
        for arg_name, arg in zip(function_parameters, args, strict=False)
    ]

    return params["shocks"][name][*indices]
 ...

In particular, the indices are empty in the second case because there is no argument, although we want an index for each initial state.