ENH: Check that initial conditions for simulations are feasible

[hmgaudecker]

Is your feature request related to a problem?

Inspired by a test case created by @buddejul:

from copy import deepcopy

import numpy as np
import jax.numpy as jnp
from lcm.entry_point import get_lcm_function
from tests.test_models.deterministic import ISKHAKOV_ET_AL_2017

PARAMS = {
    "beta": 0.95,
    "utility": {"disutility_of_work": 1},
    "next_wealth": {"interest_rate": 0.02},
    "next_lagged_retirement": {},
    "consumption_constraint": {},
    "absorbing_retirement_filter": {},
    "labor_income": {"wage": 1},
    "working": {},
}

INITIAL_STATES = {
    "wealth": jnp.array([0.25]),
    "lagged_retirement": jnp.array([0]),
}


def solve_and_simulate(model, params=PARAMS, initial_states=INITIAL_STATES):
    solve, params_dict = get_lcm_function(model, targets="solve")
    simulate, _ = get_lcm_function(model, targets="simulate")

    vf_arr_list = solve(params)

    return simulate(params, vf_arr_list=vf_arr_list, initial_states=initial_states)

def test_choice_constraint():
    result = solve_and_simulate(ISKHAKOV_ET_AL_2017)
    assert np.isfinite(result["value"]).all()

The test will fail because:

>>> result["value"]
period  initial_state_id
0       0                  -inf
1       0                  -inf
2       0                  -inf
Name: value, dtype: float32

Note that in the initial period, the consumption constraint

def consumption_constraint(consumption, wealth):
    return consumption <= wealth

cannot ever be fulfilled, because the minimum gridpoint of consumption is 1 and wealth is 0.25.

Describe the solution you'd like

Check that the initial states are valid in the sense that all of them admit at least one valid choice.

Maybe we can just apply filters and constraints to the initial states. Not fully obvious to me yet how to do that as part of an estimation when the outcome of this check may depend on parameters picked by the optimizer (i.e., what precisely we should do upon failure), but just sticking our head in the sand seems worse (will need to figure out much later what causes all those -infs)

Another thing is that we might have to allow extrapolating beyond the grid boundaries.