opensourceeconomics / lcm Goto Github PK

A bug corresponding to this was first discussed in issue #40. The change was implemented in #42, but no tests were added to the test suite. We should implement a simple test that checks this.

LCM Version: af6661e

Example

A state variable that only influences the next period. In this case, the last period utility compute_ccv function does not depend on this state variable, since the next_-functions are not included. But, since it is a state variable, LCM tries to vmap over it, resulting in an error.

Potential Solution

For each period compare variables over which we want to vmap, and variables that are in the signature of the functions we wish to vmap over. Then update the state space and space info accordingly.

Problem description

On this branch, I test the results of lcm against the analytical solution by Iskhakov et al. (2017). In this model, agents face a discrete retirement decision (absorbing state) and a continuous consumption decision.

When testing the value function of lcm, I set up the model in the following way. Note that the wealth and consumption grid starts at 0 and there are T=3 periods.

I included 0 in the consumption and wealth grid because I think this is necessary in order to correctly test the numerical results against an analytical solution. Else, retired agents being in the n-1 lowest possible states (in an n period problem with equal consumption and wealth grid) would necessarily violate their budget constraint, which produces incorrect results that may spill over to the other states.

from lcm.example_models import PHELPS_DEATON_NO_BORROWING
from lcm.entry_point import get_lcm_function

model = PHELPS_DEATON_NO_BORROWING
model["n_periods"] = 3
model["choices"]["consumption"]["start"] = 0
model["choices"]["consumption"]["stop"] = 100
model["choices"]["consumption"]["n_points"] = 10_000
model["states"]["wealth"]["start"] = 0
model["states"]["wealth"]["stop"] = 100
model["states"]["wealth"]["n_points"] = 10_000
solve_model, params_template = get_lcm_function(model=model)

params_template["beta"] = 0.98
params_template["labor_income"]["wage"] = 20.0
params_template["next_wealth"]["interest_rate"] = 0.0
params_template["utility"]["delta"] = 1.0

numerical_solution = solve_model(params=params_template)

The value functions of the first two periods (for the retired agent) seem to be wrong. The last period value function is correct.

Value function in period 2

For the initially retired, I get the following result for the value function on the wealth grid:

[ -inf , -inf, -9.118039 , ..., 7.7455106, 7.7457085, 7.745907 ]

The correct/analytical solution would be:

[-inf, -10.49036936,  -9.11793795, ..., 7.74551444, 7.74571051, 7.74590656]

The -inf in the second entry is due to the value function being computed on a grid, this doesn't show up in the analytical solution (only the first -inf entry is correct). What surprises me is that the following values on the grid are imprecise as well. When working on a grid with 10_000 grid points, only after the 800th or so grid points, results are always close up to 5 or 6 decimals. Maybe the issues at the first few grid points induce numerical imprecision somehow?

The solution for the workers coincides with the analytical solution.

Value function in period 1

I get a list of nan, both for the initially retired and for the agents still working. This problem does not occur if the minimum grid value is larger than 0.

We should think about naming conventions and abbreviations. Examples:

• scs: state choice space
• ss: state space
• ccv: conditional continuation values

https://github.com/iza-institute-of-labor-economics/lcm/blob/b187c937394c6dcfa7efe083d0ec4e1fe7e208dd/src/lcm/create_state_space.py#L204

When we use jax logical_and we cannot use more than 2 masks. At least it fails when I try to. Can we just use functools reduce or is there a better way?

The simulate function only returns values and choices, but not states. In the future, we want all three of these objects to be returned.

There are two modules example_models.py and example_models_stochastic.py containing multiple models. Then there is also the get_models.py module, defining some other models. This is cluttered and leads to bugs.

I propose we create a _models folder in which there are many more modules to reduce clutter. On top we should have a models.py module that imports all finished models from _models and puts them in __all__.

Currently, randomness is handled via numpy.random. This should be replaced by a JAX random number generator.

Code Sample, a copy-pastable example

The following model specification leads to an error when solving the model:

from lcm.entry_point import get_lcm_function
from pybaum import tree_update

def utility(consumption, working):
    return consumption - working

def next_wealth(wealth, consumption, working):
    return wealth - consumption + working

def next_wealth_constraint(next_wealth):
    return next_wealth >= 0

USER_MODEL = {
    "functions": {
        "utility": utility,
        "next_wealth": next_wealth,
        "next_wealth_constraint": next_wealth_constraint,
    },
    "choices": {
        "working": {"options": [0, 1]},
        "consumption": {"options": [0, 1]},
    },
    "states": {
        "wealth": {"options": [0, 1]},
    },
    "n_periods": 2,
}

solve_model, params_template = get_lcm_function(model=USER_MODEL)
params = tree_update(params_template, {"beta": 0.9})
solve_model(params)

Additional Information

It seems like having a discrete state causes the problems because there is no error when "wealth": {"grid_type": "linspace", "start": 0, "stop": 10, "n_points": 3}.

Error

>       solve_model(params)

tests/test_illustrative_model.py:13:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/lcm/solve_brute.py:60: in solve
    conditional_continuation_values = solve_continuous_problem(
src/lcm/solve_brute.py:139: in solve_continuous_problem
    utilities, feasibilities = gridmapped(
src/lcm/dispatchers.py:162: in allow_kwargs_wrapper
    positional += convert_kwargs_to_args(kwargs, parameters)
src/lcm/dispatchers.py:180: in convert_kwargs_to_args
    sorted_kwargs = dict(sorted(kwargs.items(), key=lambda kw: parameters.index(kw[0])))
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

kw = ('vf_arr', Array([0, 1], dtype=int64))

>   sorted_kwargs = dict(sorted(kwargs.items(), key=lambda kw: parameters.index(kw[0])))
E   ValueError: 'vf_arr' is not in list

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.

Is your feature request related to a problem?

Currently, we ignore many of the D ruff rules related to missing docstrings.

Describe the solution you'd like

Add more docstrings.

Currently, dependencies of stochastic next functions can only be discrete states. In principle, we want to allow for arbitrary arguments. There are a few things that need to be tackled. The current implementation (create_params.py) does:

dependencies = list(
    inspect.signature(model["functions"][f"next_{var}"]).parameters,
)

_check_variables_are_all_discrete_states(
    variables=dependencies,
    variable_info=variable_info,
    msg_suffix=(
        f"The function next_{var} can only depend on discrete state variables."
    ),
)

What we probably should do:

We don't need the list of arguments here, but the root nodes in a dag that contains the next function and all auxiliary functions the user might have provided. Another workaround would be to build a next function that only depends on states and choices via dags and inspects its signature.

What remains to be done:

Extend the logic to continuous variables.

If in the model-specification dictionary a continuous state variable is followed by a discrete state variable (see below), the following error is shown:

ValueError: Interpolation axes need to be the last entries in axis_order.

Code Sample, a copy-pastable example

    "states": {
        "wealth": {
            "grid_type": "linspace",
            "start": 100,
            "stop": 10_000,
            "n_points": 2000,
        },
        "age_youngest_child": {"options": list(range(19))},
    },

If the order is reversed, it works fine.

Expected Output

• At some point lcm should probably reorder the state variable definitions internally such that no error is thrown
• Alternatively, at least the error message should be improved

Bug does not have a high priority (should we add a label to indicate that?)