diff --git a/docs/source/background/choice_dependent_sparsity.rst b/docs/source/background/choice_dependent_sparsity.rst new file mode 100644 index 00000000..02f80ba2 --- /dev/null +++ b/docs/source/background/choice_dependent_sparsity.rst @@ -0,0 +1,265 @@ +.. _choice_dependent_sparsity: + +Choice-Dependent Sparsity +========================== + +.. warning:: + + This page is AI-drafted (with Claude, based on reading the source and a design + discussion with a maintainer). It documents current behaviour as of writing; + verify against the source if something looks off. + +.. note:: + + This page assumes you have read :ref:`sparsity_conditions`, which covers the + basics of ``sparsity_condition``. This page answers a narrower question: how to + express restrictions that depend on a *choice*. + +The short answer +----------------- + +"Sparsity that depends on choice" is two different requirements wearing one name, +and ``dcegm`` already supports both -- through two different functions. Which one +you need depends on *which* choice you mean: + +.. list-table:: + :header-rows: 1 + :widths: 30 30 40 + + * - You want to restrict... + - Use + - Because + * - based on the choice made **last** period + - ``sparsity_condition``, via the ``lagged_choice`` argument + - ``lagged_choice`` is an ordinary discrete state variable + * - which choices are available **in** a state + - ``state_specific_choice_set`` + - this *is* the (state, choice)-level sparsity mechanism + +There is deliberately no ``choice`` argument to ``sparsity_condition``. The next +section explains why, and the two after that give the recipe for each case. + +Why ``sparsity_condition`` cannot take ``choice`` +-------------------------------------------------- + +Two reasons, one conceptual and one structural. + +**Conceptually**, a state's existence cannot depend on a choice made *from* that +state. The agent is in the state first and chooses second; asking "is this state +valid given the choice?" inverts the timing. What people usually mean when they +reach for this is one of the two supported cases above. + +**Structurally**, ``sparsity_condition`` is evaluated in +:func:`dcegm.pre_processing.model_structure.state_space.create_state_space`, which +runs *before* the state-choice space exists: + +.. code-block:: python + + # state_space.py, inside the period / endog-state / lagged-choice loop + state_dict = { + discrete_states_names[i]: state_value + for i, state_value in enumerate(state) + } + sparsity_output = sparsity_condition(**state_dict) + +``state_dict`` is built purely from discrete *state* variables. The state-choice +space is only assembled afterwards, in +``create_state_choice_space_and_child_state_mapping``, which loops over the +already-filtered valid states and calls ``state_specific_choice_set`` for each. +So the two functions run at different stages by construction, and each sees +exactly the information available at its stage. + +Case 1: restricting on last period's choice +-------------------------------------------- + +This needs nothing special -- ``lagged_choice`` is a state variable like any +other, so just declare it as an argument. The shipped ``with_exp`` toy model +(``toy_models/cons_ret_model_with_exp/state_space_objects.py``) does exactly +this: + +.. code-block:: python + + def sparsity_condition(period, experience, lagged_choice, model_specs): + max_exp_period = period + model_specs["max_init_experience"] + ... + # If experience is the maximum for this period, you must have worked + if (experience == max_exp_period) & (lagged_choice == 1): + return False + # Retirement is absorbing: if you worked last period, experience must + # be at least the period count + elif (lagged_choice == 0) & (experience < period): + return False + else: + return True + +Both branches are choice-dependent restrictions in the everyday sense -- they +just express it through the state variable that *records* the previous choice. + +Case 2: restricting which choices a state offers +------------------------------------------------- + +This is ``state_specific_choice_set``, which returns the feasible choice set for +a given state. It is called once per valid state during state-choice space +construction: + +.. code-block:: python + + # state_choice_space.py + feasible_choice_set = state_specific_choice_set(**this_period_state) + + for choice in feasible_choice_set: + state_choice_space_raw[idx, :-1] = state_vec + state_choice_space_raw[idx, -1] = choice + ... + +Any (state, choice) pair you omit from the returned set simply never enters the +state-choice space -- no solution container row, no EGM candidate, no upper +envelope call. That is precisely "sparsity at the choice level". + +The shipped ``dcegm_paper`` model +(``toy_models/cons_ret_model_dcegm_paper/state_space_objects.py``) uses it to make +retirement absorbing: + +.. code-block:: python + + def get_state_specific_feasible_choice_set(lagged_choice, model_specs): + n_choices = model_specs["n_choices"] + + # Once the agent chooses retirement, she can only choose retirement + # thereafter. Hence, retirement is an absorbing state. + if lagged_choice == 1: + feasible_choice_set = np.array([1]) + else: + feasible_choice_set = np.arange(n_choices) + + return feasible_choice_set + +.. warning:: + + Choice encodings are per-model, not a package-wide convention: in + ``dcegm_paper`` above ``lagged_choice == 1`` is retirement, while in the + ``with_exp`` sparsity condition quoted earlier ``lagged_choice == 0`` is + working. Read the model you are editing rather than carrying an encoding over + from an example. + +The two must agree +------------------- + +This is where most real bugs live. The two mechanisms are checked against each +other at model-build time, because they jointly have to produce a closed +transition system: every feasible (state, choice) must lead somewhere valid, and +every valid non-initial state should be reachable. + +**Check 1 -- every state-choice needs a valid child** (raises ``ValueError``, in +``test_child_state_mapping``): + +.. code-block:: text + + Some state-choice combinations have invalid child states. Please update + accordingly the deterministic law of motion or the proxy function. + +You get this when ``state_specific_choice_set`` offers a choice whose implied +child state (via ``next_period_deterministic_state``) was filtered out by +``sparsity_condition``. Fix by either removing that choice from the choice set, +loosening the sparsity condition, or proxying the child to a valid state (see +below). + +**Check 2 -- every valid state should be reachable** (raises ``UserWarning``): + +.. code-block:: text + + Some states are not child states of any state-choice combination or + stochastic transition. Please revisit the sparsity condition. + +You get this when ``sparsity_condition`` admits a state that no feasible +(state, choice) transitions into. It is a warning rather than an error because +an unreachable state is wasteful, not incorrect -- it is solved and stored, but +never read. Still worth fixing: it inflates every solution container. + +A useful way to hold the two checks in mind: Check 1 says your choice set is too +*permissive* relative to your sparsity condition; Check 2 says your sparsity +condition is too permissive relative to your choice set. + +The stochastic-state exception: proxying +------------------------------------------ + +There is one case where you cannot simply delete an invalid state: when it is +invalid because of a *stochastic* state component. Every state-choice is assumed +to have exactly ``n_stochastic_states`` children, so a missing one would break +that shape invariant. + +For these, ``sparsity_condition`` returns a **dictionary** instead of a boolean -- +the valid state to use in place of the invalid one: + +.. code-block:: python + + def sparsity_condition(period, lagged_choice, job_offer, survival, model_specs): + ... + # Dead agents receive no job offers; collapse the whole (dead, job_offer) + # family onto the job_offer = 0 representative. + job_offer_out = 0 if survival == 0 else job_offer + + return { + "period": period, + "lagged_choice": lagged_choice, + "survival": survival, + "job_offer": job_offer_out, + } + +Returning the *same* state you were handed means "this state is valid" -- which +is why a single ``return {...}`` can serve both roles, and why you rarely need to +mix ``return False`` and ``return {...}`` in one function. + +Two rules the build enforces: + +- The proxy target must itself be valid. If it is not, model setup fails with + ``The state ... is used as a proxy state for the state ... However, the proxy + state is also declared invalid by the sparsity condition.`` +- Proxy dictionaries must contain exactly the discrete state names, all + integer-valued -- checked in ``create_state_space``. + +Debugging +---------- + +When the checks above fire and the cause is not obvious, get the full picture of +what the sparsity condition did to every candidate state: + +.. code-block:: python + + from dcegm.pre_processing.setup_model import create_model_dict + + state_space_df = create_model_dict( + model_config=model_config, + model_specs=model_specs, + debug_info="state_space_df", + **model_funcs, + ) + +This returns a ``pandas.DataFrame`` over the *unfiltered* cross-product state +space, with one row per candidate: every discrete state variable as a column, +plus ``is_valid``, ``is_proxied``, and (when proxies exist) ``idxs_proxied_to``. +Filtering it to the state named in the error message usually localises the +problem in one step -- it is built by ``create_state_space(..., debugging=True)``, +so it reflects exactly the decisions the real build made. + +.. warning:: + + Use ``create_model_dict`` for this, not ``dcegm.setup_model``. ``setup_model`` + is a class whose ``__init__`` immediately indexes ``model_dict["model_config"]``, + so passing ``debug_info="state_space_df"`` there fails with + ``KeyError: 'model_config'`` -- the debug path returns a DataFrame, which has + no such key. + +Summary +-------- + +- Restricting on **last period's** choice: ``sparsity_condition``, via + ``lagged_choice``. Nothing special required. +- Restricting **this period's** available choices: ``state_specific_choice_set``. + This is the choice-level sparsity mechanism. +- ``sparsity_condition`` takes no ``choice`` argument, by design -- it runs before + the state-choice space exists, and state validity logically precedes choice. +- Keep the two consistent: every feasible (state, choice) needs a valid child + (``ValueError``), and every valid state should be reachable (``UserWarning``). +- Invalid-because-stochastic states get proxied to a valid representative rather + than dropped, to preserve the ``n_stochastic_states``-children invariant. diff --git a/docs/source/background/sparsity_conditions.rst b/docs/source/background/sparsity_conditions.rst index 224a1984..1ef7fecd 100644 --- a/docs/source/background/sparsity_conditions.rst +++ b/docs/source/background/sparsity_conditions.rst @@ -73,13 +73,52 @@ Explanation of Logic: - **c4/c5**: Forces retirement if agent is above maximum retirement age. - **Adjustments**: If the agent is *dead* (`survival == 0`), the job offer is set to 0, and the period is fixed at the final period to make the state absorbing. +Utility Functions at Proxied Terminal States +-------------------------------------------- + +Proxying an absorbing terminal state has a consequence for your **utility functions** that is easy to miss. In the example above, every dead state (``survival == 0``) is redirected to a single slot in the final period, where its value is stored once and reused. But that stored value is not the whole story: during the EGM step, whenever a *living* parent integrates over its (possibly dead) children, *dcegm* re-evaluates your periodic ``utility`` and ``marginal_utility`` at the child's own identity -- and for a dead child that identity still carries ``survival == 0``. + +So your periodic utility functions must branch on the terminal state and return the terminal (e.g. consume-all / bequest) value there, not the ordinary living value: + +.. code-block:: python + + def utility(consumption, survival, ..., params): + alive = utility_alive(consumption, ..., params) + dead = utility_final_consume_all(wealth=consumption, ..., params) + return jnp.where(survival == 0, dead, alive) + +Two things matter: + +- **It must match** ``utility_functions_final_period``. The dead child's value was stored with the final-period utility; the branch above is what the parent re-evaluates for that same child during EGM. If the two disagree, the two representations of the same terminal state are inconsistent and the Euler equation of every parent with a dead child is wrong. +- **The same applies to** ``marginal_utility``, and more strongly: the child's marginal utility is *always* recomputed at the proxied identity in the EGM step (it is never read back from storage), so a missing terminal branch there breaks the bequest term directly. + +When this branch is exercised +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The terminal branch of the *periodic* functions fires for every period solved by the main backward-induction scan -- that is, all periods **except the last two**: + +- the **last period** is solved entirely with ``utility_functions_final_period`` and never calls the periodic ``utility``; +- the **second-to-last period** is handed the final period's values and marginal utilities directly, so it too bypasses the periodic terminal branch; +- **every earlier period** reads its terminal children through the proxy and re-evaluates the periodic ``utility`` / ``marginal_utility`` at the ``survival == 0`` identity. + Best Practices -------------- - Make sparsity conditions **strict**: only allow logically valid state combinations. - Handle **absorbing states** like death or permanent retirement carefully. +- When a sparsity condition proxies an absorbing terminal state, make sure your periodic ``utility`` and ``marginal_utility`` branch on that state and agree with ``utility_functions_final_period`` (see `Utility Functions at Proxied Terminal States`_ above). - Ensure that any state created in the deterministic transition function also satisfies the sparsity condition. Failing to correctly define sparsity conditions will result in `ValueError` exceptions and warnings during model setup, as *dcegm* verifies that every `(state, choice)` pair leads to a valid next state. -To help you setup the correct sparsity conditions for your model, ``setup_model()`` provides a debug mode which can be triggered by passing the argument ``debug=state_space_df``. +If your restriction depends on a *choice* -- either the one made last period or which choices a state offers at all -- see :ref:`choice_dependent_sparsity`, which covers how the two mechanisms (``sparsity_condition`` and ``state_specific_choice_set``) divide that work between them. + +To help you setup the correct sparsity conditions for your model, a debug mode is available which returns a ``pandas.DataFrame`` over the unfiltered state space, showing ``is_valid``/``is_proxied`` per candidate state: + +.. code-block:: python + + from dcegm.pre_processing.setup_model import create_model_dict + + state_space_df = create_model_dict(..., debug_info="state_space_df") + +Note this must be called on ``create_model_dict``, not on ``setup_model``: the latter is a class that immediately indexes ``model_dict["model_config"]``, so the debug path (which returns a DataFrame) raises ``KeyError: 'model_config'`` there. diff --git a/docs/source/development/internals/law_of_motion.rst b/docs/source/development/internals/law_of_motion.rst new file mode 100644 index 00000000..6c681ab3 --- /dev/null +++ b/docs/source/development/internals/law_of_motion.rst @@ -0,0 +1,303 @@ +.. _law_of_motion_internals: + +The Law of Motion: Evaluation Granularity +========================================== + +.. warning:: + + This page is AI-drafted (with Claude, based on reading the source and a design + discussion with a maintainer). It documents current behaviour as of writing; + verify against the source if something looks off. The performance claims below + are reasoned, not measured -- see "What is not measured". + +What "the law of motion" means here +------------------------------------ + +Two user-supplied functions describe how a state transitions into the next period: + +- ``budget_constraint`` -- processed into + ``model_funcs["compute_assets_begin_of_period"]``. Maps the previous period's + end-of-period assets, this period's income shock, and the discrete state to + beginning-of-period wealth. This is the *only* law of motion that may declare + ``choice``, and it means the choice made in the period whose wealth is being + computed -- see :ref:`choice_dependent_budget` below. + + Its two non-state arguments are deliberately asymmetric in timing, and the names + say so: ``asset_end_of_previous_period`` really is carried over from ``t-1``, + while ``income_shock`` is realized in ``t``, the period whose wealth is being + computed. (It was called ``income_shock_previous_period`` until the shipped + example models' habit of earning income against ``lagged_choice`` was mistaken + for a framework property; ``_check_budget_does_not_declare_renamed_income_shock`` + in ``process_model_functions.py`` rejects the old name rather than letting the + signature filter raise a bare ``KeyError``.) +- ``next_period_continuous_state`` -- optional, only when the model has additional + continuous states (e.g. continuous experience). Maps this period's continuous + state to next period's. It may **not** declare ``choice``; doing so raises at + model build (``_check_choice_not_declared`` in ``process_model_functions.py``). + The previous period's choice, which is what accumulation processes actually + need, is the ordinary state variable ``lagged_choice``. + +Both are evaluated in +:func:`dcegm.law_of_motion.calc_law_of_motion_for_state_choices`, called from +``egm/interpolate_marginal_utility.py`` (the main solve path, once per batch) and +``final_periods.py`` (the terminal period). + +The nesting +------------ + +``calc_law_of_motion_for_state_choices`` evaluates the budget equation over a +four-dimensional product, one ``vmap`` level per axis: + +.. code-block:: text + + for each row (state-choice, or child state -- see below) + for each continuous-state combo point + for each end-of-period assets grid point + for each income shock draw + compute_assets_begin_of_period(...) + +Each level strips one axis, so the innermost call sees a single scalar tuple and +the user's function never has to know it is being batched. The grid that supplies +the assets axis is evaluated inside the outermost level, right next to where it +is consumed, rather than precomputed for the whole batch and passed in -- see +``_transitions_for_one_state``. + +Parent vs. child: two different state-choices +----------------------------------------------- + +This trips people up, so it is worth stating precisely. Two distinct state-choice +dicts flow into this function: + +``child_state_choices`` + The state-choice whose beginning-of-period wealth we are computing. In the + solve path this is the *child*. Its own identity is what the transition + function sees -- including ``lagged_choice``, which *is* the parent's choice. + +``representative_parent_state_choice_vec`` + A representative *parent* state-choice, used only to select whose continuous + grid supplies the values fed through the transition. Once grids are + state-choice-specific these are not interchangeable: the transition depends on + the child's identity, but the grid *values* must come from the parent that is + actually transitioning. + +Any one parent works as the representative, because +``check_continuous_grid_consistency_across_shared_children`` (run once at +model-build time) already guarantees every parent sharing a child agrees on its +own grid. The representative is picked in +``pre_processing/batches/child_state_dedup.py``. + +Two evaluation granularities +----------------------------- + +The transition into a child normally does not depend on the child's own choice, +only on its state. So unless the budget equation declares ``choice``, every +state-choice sharing a child state would compute a bit-identical transition, and +evaluating per state-choice does redundant work proportional to the number of +choices. + +``dcegm`` therefore has two implementations, dispatched between by +``calc_law_of_motion`` -- the single entry point every caller uses, so the +granularity decision lives in one place rather than being repeated per call site: + +.. list-table:: + :header-rows: 1 + :widths: 32 30 38 + + * - Function + - Evaluates once per + - Used when + * - ``calc_law_of_motion_for_state_choices`` + - child **state-choice** + - the budget equation declares ``choice`` + * - ``calc_law_of_motion_for_child_states`` + - unique child **state** + - it does not (the common case) + +.. note:: + + The granularity is chosen for both transitions at once, on the budget equation + alone. Since ``next_period_continuous_state`` can never declare ``choice``, a + model with a choice-dependent budget recomputes a bit-identical continuous + state for each of a child's choices. Splitting the two apart would mean + carrying both sets of dedup arrays and a gather in between, for a quantity of + shape ``(n_rows, n_combos)`` against wealth's + ``(n_rows, n_combos, n_savings, n_shocks)`` -- a small fraction of the work, so + the two are kept in one pass. + +Both call sites -- ``interpolate_value_and_marg_util`` for the main backward +induction, and ``solve_final_period`` for the terminal period -- go through +``calc_law_of_motion`` and so take the same shortcut. The final period needs its +own dedup arrays, since it is not built by ``compute_child_dedup_for_batch``; +those are constructed in ``last_two_periods.py`` +(``unique_final_period_states``, +``representative_second_last_period_parent_idx_per_final_state``, +``state_row_for_final_period_state_choice``) by deduplicating +``parent_states_final_period``. + +Crucially these are not two implementations of the transition math. +``calc_law_of_motion_for_child_states`` *calls* +``calc_law_of_motion_for_state_choices`` with deduplicated child states instead of +state-choices, then gathers the per-state result back out to state-choice +granularity: + +.. code-block:: python + + law_of_motion_per_state = calc_law_of_motion_for_state_choices( + child_state_choices=child_states, # deduplicated + representative_parent_state_choice_vec=representative_parent_state_choices, + ... + ) + + return { + "continuous_states": { + name: jnp.take(grid, state_row_for_state_choice, axis=0) + for name, grid in law_of_motion_per_state["continuous_states"].items() + }, + "assets_begin_of_period": jnp.take( + law_of_motion_per_state["assets_begin_of_period"], + state_row_for_state_choice, axis=0, + ), + } + +``state_row_for_state_choice`` maps each child state-choice back to its row in the +deduplicated child states. It is built in ``child_state_dedup.py`` alongside the +dedup arrays that already existed, and follows the same gather pattern +``calculate_candidate_solutions_from_euler_equation`` uses one stage later for +``marg_util_next``. + +How the choice is made +----------------------- + +By signature inspection, once at model-build time -- the same approach +``taste_shock_scale_is_scalar`` uses: + +.. code-block:: python + + # process_model_functions.py + def _declares_choice(func): + return "choice" in set(inspect.signature(func).parameters) + + budget_depends_on_choice = _declares_choice(budget_constraint) + +The result is stored as ``model_funcs["budget_depends_on_choice"]``. Inspection +happens on the *user's* function, before +``determine_function_arguments_and_partial_model_specs`` wraps it -- the wrapper's +``**kwargs`` signature would hide the real parameter names. + +The same inspection is used in the opposite direction for +``next_period_continuous_state``: ``_check_choice_not_declared`` raises a +``ValueError`` if it declares ``choice``, so the only function that can put the +solve on the per-state-choice path is the budget equation. + +All shipped toy models (``with_cont_exp``, ``with_exp``, ``dcegm_paper``) report +``False``, so the deduplicated path is what the test suite actually exercises; +``test_default_toy_models_take_the_state_level_fast_path`` asserts this so it +cannot silently regress. + +.. _choice_dependent_budget: + +Writing a choice-dependent budget equation +-------------------------------------------- + +Declare ``choice`` as an argument and it will be passed through: + +.. code-block:: python + + def budget_constraint( + period, lagged_choice, choice, asset_end_of_previous_period, + income_shock, params, model_specs, + ): + wealth = ... # usual computation + # e.g. a choice-specific fixed cost deducted at the start of the period + return wealth - model_specs["entry_cost"] * (choice == 1) + +Functions that do not declare ``choice`` are unaffected: +``determine_function_arguments_and_partial_model_specs`` filters kwargs down to +each function's own signature, so the extra key is simply ignored. + +.. important:: + + **A choice-dependent budget equation must return pre-consumption wealth.** + + The ``choice`` seen here is the agent's own choice in the period whose wealth + is being computed, so the timing is unusual: wealth is resolved *after* the + discrete choice but *before* the continuous consumption choice. ``dcegm`` + builds the agent's decision on exactly that reading -- it evaluates the budget + equation once per choice, and the agent then compares each choice's value at + that choice's own wealth, consuming out of it. Anything the budget equation + subtracts is therefore money the agent never gets to consume; anything it adds + is consumable in the same period. + + This is why only wealth works this way. A choice-specific cost or transfer is + naturally levied at the moment of choosing, and consumption adjusts to it + within the period. Every other state variable is fixed at the start of the + period, before the choice, so ``next_period_continuous_state`` is evaluated at + the *end* of the previous period and is not offered ``choice`` at all. The + parent's choice is available separately, and always has been, as the child's + ``lagged_choice``. + +.. note:: + + This did not work before the granularity split. Two places stripped ``choice`` + before calling user functions -- ``state_vec.pop("choice", None)`` in + ``law_of_motion.py`` and ``state_vec.pop("choice")`` in ``final_periods.py`` -- + so a budget equation declaring ``choice`` raised ``KeyError: 'choice'``. Both + pops were removable: the signature filter already does that job. + +Simulation follows the same split. ``assets_begin_of_period_for_each_choice`` +(``simulation/sim_utils.py``) evaluates the budget equation for every choice at +the start of a period, before the choice is drawn, and the realized wealth is the +column belonging to the choice actually made. The additional continuous states are +settled one step earlier, in ``transition_to_next_period`` at the end of the +previous period, and so carry no choice axis at all. + +Cost +----- + +Deduplication removes a factor of (choices per child state) from the *entire* +four-dimensional product above, not from one axis in isolation -- so it is a +straight ``1/n_choices`` reduction in how often the budget equation is evaluated, +independent of how large the combo/savings/shock axes are. What it pays instead is +a gather: pure data movement, no FLOPs beyond address computation. + +That trade should favour deduplication when the budget equation is +arithmetic-heavy, and matters little when it is trivial (``R * savings + income`` +fuses into the surrounding kernel at negligible cost, and the gather becomes pure +overhead). Since the fallback only triggers on an explicit ``choice`` argument, +models with a cheap choice-independent budget still get the dedup -- harmless, +just not a large win. + +What is not measured +--------------------- + +The paragraph above is reasoning, not measurement. Nothing in ``dcegm`` currently +profiles this, and the crossover point -- how many FLOPs the budget equation needs +before dedup beats the gather -- is unknown. Confirming it would need +``jax.block_until_ready()`` timing or the XLA profiler, comparing a model with a +genuinely heavy budget equation and a realistic choice count under both paths +(forceable by adding an unused ``choice`` argument, as +``test_state_specific_grids_law_of_motion.py`` does). The same caveat the +:ref:`batching_internals` page applies to its own performance claims applies here. + +Correctness testing +-------------------- + +The central guarantee is that the two granularities agree. +``tests/test_state_specific_grids_law_of_motion.py`` establishes this +(``test_state_level_and_state_choice_level_paths_agree_bit_for_bit``) by solving +the same economics twice -- once normally, once with a budget equation that declares an +otherwise-unused ``choice`` to force the per-state-choice path -- and requiring +bit-for-bit identical ``value``, ``policy`` and ``endog_grid``. A wrong +``state_row_for_state_choice`` mapping would surface there immediately, as +children would be fed another state's transition. + +That test is paired with a sensitivity check +(``test_genuinely_choice_dependent_budget_differs_from_choice_free_one``): a budget +equation that actually *uses* ``choice`` must produce a different solution. Without +it, the equivalence test would still pass if ``choice`` were being silently +dropped again -- the exact regression it is there to catch. + +The restriction on ``next_period_continuous_state`` has its own test, +``test_continuous_state_transition_declaring_choice_is_rejected`` in +``tests/test_state_specific_grids_reference.py``: a transition declaring +``choice`` must raise at model build, not be silently ignored. diff --git a/docs/source/development/internals/state_specific_continuous_grids_plan.md b/docs/source/development/internals/state_specific_continuous_grids_plan.md new file mode 100644 index 00000000..169c5127 --- /dev/null +++ b/docs/source/development/internals/state_specific_continuous_grids_plan.md @@ -0,0 +1,237 @@ +# State-choice-specific `assets_end_of_period` / `assets_begin_of_period` grids + +> **Note.** This page is AI-drafted (with Claude, based on reading the source and a +> design discussion with a maintainer) and is a *plan*, updated as phases land and +> design decisions change. It replaces an earlier plan doc at this same path for the +> now-complete "state-specific continuous state grids" feature (state-choice-specific +> grids for *additional* continuous states like "experience"), which shipped in full — +> see git history for that design if needed. This doc covers the follow-on work: making +> the two continuous states outside that feature's scope, `assets_end_of_period` and +> `assets_begin_of_period`, state-choice-specific too, plus a config ergonomics change +> (the `None`-grid convention) that applies to all state-choice-specific continuous +> states, not just these two. + +## Status + +**Done, tested, passing the full `dcegm` suite.** All phases below landed, with two +corrections to the original design (see "What changed from the original plan" +below): `assets_end_of_period` does *not* get the `None`-grid convention after all, +and a genuine pre-existing bug (unrelated to this work) had to be fixed as a +prerequisite for testing `assets_begin_of_period` state-specificity at all. + +### What changed from the original plan + +1. **`assets_end_of_period` does not support the `None`-grid convention, and is the one + name exempt from the (later-added) requirement that a declared array be `None` + whenever a `continuous_grid_functions` entry exists.** `check_model_config.py` uses + `len(model_config["continuous_states"]["assets_end_of_period"])` to compute + tuning-parameter defaults (`n_constrained_points_to_add`, `extra_wealth_grid_factor` + -derived `n_total_wealth_grid`) — this runs unconditionally, for every + `upper_envelope` method, before `continuous_grid_functions` is even known and before + any state-choice exists to pin a deferred size against; it isn't "for FUES" + specifically, even though only the `fues` branch actually ends up consuming the + result (`druedahl_jorgensen` overwrites it with its own `assets_begin_of_period` + -derived formula). Making its size deferrable would mean reordering + `check_model_config.py` so this `len()` call isn't unconditional, a materially + bigger and riskier change than this work needed. `assets_end_of_period` still fully + supports state-choice-specific *values* via `continuous_grid_functions` — a real + array must always be declared regardless, and is simply ignored once a grid_func is + given, the one case where that's allowed rather than rejected. +2. **A genuine, pre-existing bug turned up as a blocker**: `final_periods.py`'s + `solve_final_period` had a `has_additional_continuous_state` branch that, for + models *without* an additional continuous state, unconditionally read + `wealth_child_states_final_period` (the law-of-motion's `assets_end_of_period`-based + output) as the final-period storage grid — correct for FUES, but silently wrong + for Druedahl-Jorgensen (which should use `assets_begin_of_period` there instead). + This was apparently never exercised before: every existing DJ test in this repo + used a toy model with an additional continuous state, which goes through a + different (correct) branch. It surfaced immediately as a shape-mismatch crash + when writing the first `assets_begin_of_period`-state-specific end-to-end test + against a no-additional-continuous-state DJ model, and had to be fixed (a new + `_calc_dj_final_value_no_additional_state` branch) before that work could be + validated at all. Unrelated to state-specificity itself — the bug is present for + *any* DJ model without an additional continuous state, state-specific or not. +3. **`assets_begin_of_period` state-specificity together with an additional + continuous state is now supported, but only when `skip_endog_grid_storage` is + `True`** (`upper_envelope["method"] == "druedahl_jorgensen"` and at least two + choices), enforced by `process_continuous_grid_functions`. The n-D regular + (Druedahl-Jorgensen) interpolation path + (`interpnd_policy_and_value_for_child_states_on_own_regular_grids`) now threads + `wealth_grid` per child too (shape `(n_child_state_choices, n_wealth)`), + mirroring `additional_continuous_state_grids_per_child`; the matching per-child + computation (`vmap(compute_own_dj_wealth_grid, ...)`, the child's own identity, + same reasoning as the simple-1D-DJ branch) lives in + `_interpolate_value_and_marg_util_nd_regular` + (`interpolate_marginal_utility.py`). With a single choice the Druedahl-Jorgensen + upper envelope is skipped entirely (see `check_model_config.py`), so the stored + endogenous grid there is a genuinely different, non-fixed grid — a + state-specific `assets_begin_of_period` has nothing to plug into in that case, + and remains rejected. + +## Problem + +`continuous_grid_functions` (the top-level `create_model_dict` argument from the prior +feature) is already *accepted* for `"assets_end_of_period"` and `"assets_begin_of_period"` +at the config-validation layer: `process_continuous_grid_functions`'s `default_grids` +dict includes both, a user-supplied grid_func for either gets wrapped correctly, and +`state_specific_continuous_grid_names` (just `list(continuous_grid_functions.keys())`, +not filtered) would even run either through the Phase-1 child-sharing consistency check. + +But nothing downstream ever calls `continuous_grid_functions["assets_end_of_period"]` or +`["assets_begin_of_period"]`. Every consumer reads a single, model-wide global array +directly instead: + +- `assets_end_of_period`: always `continuous_states_info["assets_grid_end_of_period"]`, + read directly in `law_of_motion.py`, `solve_euler_equation.py`, + `interpolate_marginal_utility.py`, `final_periods.py`. +- `assets_begin_of_period` (Druedahl-Jorgensen's fixed common wealth grid, "m_grid"): + always `continuous_states_info["assets_begin_of_period"]` / its derived + `dj_wealth_grid`, read directly in `upper_evelope_wrapper.py`'s `drued_jorg_jax(...)` + call and broadcast everywhere via `broadcast_dj_wealth_grid` + (`sol_container.py`, `interface.py`, `simulation_interp.py`) whenever + `skip_endog_grid_storage` is `True`. + +So today a user can pass a `continuous_grid_functions` entry for either name, it +validates cleanly, and is then silently ignored at solve time. That's the gap. + +**Deliberately out of scope / unaffected:** the *additional* continuous states +mechanism (`additional_continuous_state_names`, `compute_own_continuous_grid_combos`, +the representative-parent threading through batches) stays exactly as it is, +scoped to continuous states other than these two — `check_model_config.py` already +excludes both names from `additional_continuous_states`, and that boundary is not +changing. `assets_end_of_period`/`assets_begin_of_period` get their own, separate +wiring, reusing the same underlying ideas (on-demand evaluation, representative +parent where needed) rather than being folded into the existing list. + +## Design + +### Two different mechanisms, not one + +The two names need different treatment because they play structurally different +roles in the solve: + +**`assets_end_of_period`** is an *input to the law of motion*: it's the exogenous +savings grid a parent's transition is evaluated over +(`calc_beginning_of_period_assets_for_single_state`'s `asset_end_of_previous_period` +comes from iterating over this grid, feeding the child's own +`compute_assets_begin_of_period` budget equation). This is exactly the same +parent→child relationship that made the *additional continuous states* feature need +representative-parent threading in the first place: the child's own identity is +correct for the transition call itself, but the grid *values* fed into it must come +from a representative parent's own grid. So `assets_end_of_period` needs the same +representative-parent mechanism, applied to this one additional name. + +**`assets_begin_of_period`** is not a law-of-motion input at all — it's the *output* +common grid that Druedahl-Jorgensen's upper envelope interpolates onto and stores +against (`m_grid`). Reading a child's own stored solution back doesn't need a +representative parent (same reasoning as reading a child's own *additional* +continuous-state combo grid back during interpolation — see +`interpolate_marginal_utility.py`'s "own grid" pattern) — it just needs to know the +right grid was used at solve time for *this* state-choice, indexed by its own +identity, no parent involved. And it only matters at all when +`skip_endog_grid_storage` is `True` (Druedahl-Jorgensen with ≥2 choices) — that's +precisely the case where the real per-state-choice `endog_grid` isn't stored, because +every state-choice's "endogenous" grid is by construction this fixed array. When +`skip_endog_grid_storage` is `False` (Druedahl-Jorgensen with 1 choice, or FUES), the +real stored `endog_grid` is read/used normally and `assets_begin_of_period` plays no +role in interpolation — so state-specificity for this name is meaningless outside +the `skip_endog_grid_storage=True` case, and should be gated on it explicitly (e.g. +rejected or ignored otherwise, mirroring the existing FUES/`assets_begin_of_period` +mutual-exclusion check in `check_model_config.py`). + +### The `None`-grid convention + +Separately, for *any* state-choice-specific continuous state (this includes the +existing "additional continuous states" mechanism too, retroactively, plus these two +new names): today, once a `continuous_grid_functions[name]` entry is given, the +*values* declared in `model_config["continuous_states"][name]` are entirely unused — +but the array must still be supplied, purely so its `len()` can be read off to fix +the grid size and validate every state-choice's own grid against it. That's a +redundant, awkward requirement: fabricating a same-length dummy array just to declare +a size. + +New convention, now enforced in both directions: `model_config["continuous_states"][name] += None` is required whenever a matching `continuous_grid_functions[name]` is given, and +required to be non-`None` otherwise. `None` with no function raises (nothing to build the +grid from); a real array *with* a function also raises (the array would be silently +unused, which is exactly the awkward, easy-to-misread state this convention exists to +rule out — see `process_continuous_grid_functions`). The key must still be *present* in +`model_config["continuous_states"]` (so a genuinely missing continuous state is still +caught as a missing key, not confused with this explicit opt-in). When `None` is +declared, the grid's size is pinned by evaluating the grid_func once on a representative +state-choice row instead of reading a precomputed length. `assets_end_of_period` is the +one name exempt from this: it must always keep a real array, function or not (see below). + +**Architectural consequence:** this moves *when* the expected length becomes known. +Today it's read from `model_config["continuous_states"][name]` in +`check_model_config.py`, before the state-choice space exists. Pinning it by +evaluation requires an actual state-choice row to call the function on, which only +exists after `create_state_choice_space_and_child_state_mapping` builds +`state_choice_space` — so "determine expected length for this name" needs to move to +alongside `evaluate_state_specific_continuous_grids` in `continuous_state_grids.py` +(which already evaluates every row's grid_func output for the consistency check; +deriving `expected_length` from row 0's own output shape there is a natural +extension, not a bolt-on). `n_continuous_state_combinations` (which sizes the global +solution containers, currently computed early in `check_model_config.py`) becomes a +two-stage value: known immediately for names with a real declared array, resolved +only after state-structure construction for names declared as `None`. + +## Phases + +1. **Config layer — done.** `None` allowed (and, once a `continuous_grid_functions[name]` + entry exists, *required*) in `model_config["continuous_states"][name]` for the + additional continuous states and `assets_begin_of_period`; `assets_end_of_period` is + exempt in both directions (see above) and must always keep a real array. + `process_continuous_grid_functions` raises if a `None`-eligible name is declared + `None` with no matching function, and equally raises if it's declared as a real array + *with* one. Also added: a validation rejecting `assets_begin_of_period` + state-specificity together with any additional continuous state. +2. **Grid-size pinning — done.** `evaluate_state_specific_continuous_grids` (in + `continuous_state_grids.py`) pins the length of a `None`-declared grid from the + first state-choice row it evaluates, alongside its existing per-row consistency + check. `_merge_resolved_state_specific_lengths` (in `state_choice_space.py`) + folds the resolved length back into `n_continuous_state_combinations` and, for + `assets_begin_of_period`, `n_total_wealth_grid`. +3. **`assets_end_of_period` representative-parent threading — done.** + `calc_law_of_motion_for_state_choices` evaluates `assets_end_of_period` from + `representative_parent_state_choice_vec` (the same representative-parent dict the + additional continuous states already use), replacing the old shared + `assets_grid_end_of_period` array parameter. Its "own grid" role (this + state-choice solving its own EGM problem in `solve_euler_equation.py`, or its own + terminal value in `final_periods.py`'s FUES branch) is self-referential, evaluated + from `state_choice_vec` directly, no representative parent needed there. +4. **`assets_begin_of_period` own-grid threading — done.** New + `compute_own_dj_wealth_grid` helper in `law_of_motion.py`, used wherever + `dj_wealth_grid`/`broadcast_dj_wealth_grid` used to be read as a shared array: + `upper_evelope_wrapper.py`'s DJ branch (solve), `interpolate_marginal_utility.py`'s + 1-D and n-D DJ interpolation branches (solve), `interp_interfaces.py`'s three + `interpolate_*_for_state_and_choice` functions, `simulation_interp.py`, and + `sol_interface.py`'s raw-grid accessor (readers). The shared array + (`continuous_states_info["dj_wealth_grid"]`) and the `broadcast_dj_wealth_grid` + helper have since been **removed entirely**: nothing reads a shared grid any more, + every reader recomputes each state-choice's own grid via + `compute_own_dj_wealth_grid`. Where the endogenous grid is not stored, callers + pass a shape-only placeholder (`None` in the solve carry, a zeros array of the + right shape in the readers) that is never read for its values. +5. **Readers — done**, folded into phase 4 above (the two mechanisms turned out to + share almost all their reader-side call sites, so weren't worth separating). + `simulate()`'s `interpolate_policy_and_value_for_all_agents` computes its own + per-agent Druedahl-Jorgensen wealth grid the same way. +6. **End-to-end validation — done.** Strong bit-for-bit "constant-but-different-grid + matches direct declaration" tests for both names, through the solve path and all + reader paths (`simulate()`, `choice_values_for_states`/`choice_policies_for_states`): + `tests/test_assets_end_of_period_state_specific.py`, + `tests/test_assets_begin_of_period_state_specific.py`. The latter also confirmed + the pre-existing `final_periods.py` bug fix and the multidim-rejection validation. + +## Open questions (resolved) + +- Which state-choice row is "representative" for pinning a `None`-declared grid's + size: row 0 of `state_choice_space`, unconditionally. Every other row's grid_func + output is validated against that pinned length in the same pass (mismatches raise, + same as the pre-existing declared-length validation) — a period-dependent grid_func + that changes size would already be caught this way, no separate check needed. +- The representative-parent machinery for `assets_end_of_period` reuses + `representative_parent_state_choice_vec` directly (no separate threading): it's evaluated + alongside the additional continuous states, from the same representative-parent + dict, inside `calc_law_of_motion_for_state_choices`. diff --git a/docs/source/guides/divorce_transition_without_lagged_state.md b/docs/source/guides/divorce_transition_without_lagged_state.md index bbfadfd1..ffbbe308 100644 --- a/docs/source/guides/divorce_transition_without_lagged_state.md +++ b/docs/source/guides/divorce_transition_without_lagged_state.md @@ -40,7 +40,7 @@ partner-status multiplier in `budget_constraint`, keyed off the *current* ```python def budget_constraint(period, lagged_choice, partner_state, - asset_end_of_previous_period, income_shock_previous_period, + asset_end_of_previous_period, income_shock, params, model_specs): multiplier = 2.0 if partner_state == 1 else 1.0 return multiplier * asset_end_of_previous_period * (1 + params["interest_rate"]) + income @@ -70,7 +70,7 @@ independent solve — the model still runs, converges, produces a ```python def budget_constraint(period, lagged_choice, partner_state, - asset_end_of_previous_period, income_shock_previous_period, + asset_end_of_previous_period, income_shock, params, model_specs): multiplier = jnp.where(partner_state == 1, 2.0, 1.0) own_income = params["y_work"] * (lagged_choice == 0) diff --git a/docs/source/guides/minimal_example.ipynb b/docs/source/guides/minimal_example.ipynb index 1775d616..a1b8e90e 100644 --- a/docs/source/guides/minimal_example.ipynb +++ b/docs/source/guides/minimal_example.ipynb @@ -125,7 +125,7 @@ " \"period\": jnp.zeros(n_agents),\n", " \"lagged_choice\": jnp.zeros(n_agents), # all agents start as workers\n", " \"experience\": jnp.ones(n_agents),\n", - " \"assets_begin_of_period\": jnp.ones(n_agents) * 10,\n", + " \"assets_end_of_previous_period\": jnp.ones(n_agents) * 10,\n", "}" ] }, diff --git a/docs/source/guides/practitioner_guide.rst b/docs/source/guides/practitioner_guide.rst index 95d16338..abd7168c 100644 --- a/docs/source/guides/practitioner_guide.rst +++ b/docs/source/guides/practitioner_guide.rst @@ -82,7 +82,7 @@ This guide explains how to specify, solve, simulate and potentially estimate str - `n_periods`: Number of period. Needs to be an integer larger than 1. - `choices`: Discrete choices of the model. Consecutive integers starting from 0. Either provided as a list or as a integer, which then is converted to a list with consecutive integers starting from 0 and to the integer minus 1. - `continuous_states`: Dictionary containing the grids for continuous variable. The dictionary requires - - `assets_end_of_period`: The grid for the end of period assets, which is required for the egm step. It is expected as a numpy array and with monotonic increasing values. + - `assets_end_of_period`: The grid for the end of period assets, which is required for the egm step. It is expected as a numpy array and with monotonic increasing values, or as `None` when the grid is supplied per state-choice through `continuous_grid_functions` instead. - `n_quad_points`: Number of quadrature points used for the integration over the income shock distribution. The quadrature points are used to approximate the integral of the value function over the income shock distribution. The number of quadrature points should be a positive integer. An example for a model configuration with the mandatory keys is: @@ -195,7 +195,24 @@ Thanks to it's efficient JAX implementation, `dcegm` is capable of solving and s **Initial Conditions** -Model simulation requires user provided *initial_conditions*. In contrast to the sparsity conditions, `dcegm` currently does not perform checks on initial conditions. Badly specified initial conditions may result in errors in the simulation. For instance, the user should make sure that `assets_begin_of_period` is not set to zero in the first period and that initial conditions do not violate the sparsity conditions (for example by providing more initial experience that the sparsity conditions allows for an agent to have in a given period). +Model simulation requires user provided *initial_conditions*. Among them, `dcegm` requires `assets_end_of_previous_period`: the assets agents **carry into** the first period, not the wealth they have available to spend in it. Every period, the first included, builds its beginning-of-period wealth by applying your budget constraint to the assets carried in plus that period's income shock, so the first period is not a special case. (This is also why supplying `assets_begin_of_period` raises a `ValueError` -- it named a different quantity.) + +The simulated panel is self-consistent row by row: each period draws its own income +shock, feeds it to your budget constraint, and reports it in that period's own +`income_shock` column. So for any period, applying your budget constraint to the +previous period's `savings` and this period's `income_shock` reproduces this period's +`assets_begin_of_period` -- with the assets you supplied standing in for `savings` in +the first period. + +The shock is realized in the period whose wealth your budget constraint computes, and +your budget constraint receives it under the name `income_shock`. Whether the *income* +it perturbs was earned earlier is a modelling decision that is entirely yours: the +example models shipped with `dcegm` make income depend on `lagged_choice`, so that a +retiring agent still collects the wage for the work done in the previous period, but +nothing in the framework requires that. You are free to make income depend on this +period's state alone. + +Apart from that key, `dcegm` currently does not perform checks on initial conditions. Badly specified initial conditions may result in errors in the simulation. For instance, the user should make sure that initial conditions do not violate the sparsity conditions (for example by providing more initial experience than the sparsity condition allows for an agent to have in a given period). **Model Timing** @@ -219,3 +236,40 @@ This implementation currently supports the following timing of events where each 5. State transitions that affect next-period states realize (e.g. health shocks and income shocks), determining :math:`s_{t+1}`. + +**Choice-dependent budget equations** + +By default every state variable, wealth included, is fixed at step 1 above -- +before the discrete choice. Wealth is the one exception `dcegm` allows: if your +`budget_constraint` declares `choice` in its signature, it is evaluated once per +choice, between steps 2 and 3, and the agent compares each choice's value at that +choice's own wealth. This is how you model a choice-specific cost or transfer, +for instance an entry cost paid on switching occupation: + +.. code-block:: python + + def budget_constraint( + period, lagged_choice, choice, asset_end_of_previous_period, + income_shock, params, model_specs, + ): + wealth = ... # usual computation + return wealth - model_specs["entry_cost"] * (choice == 1) + +.. important:: + + When `choice` is used, `dcegm` expects the budget equation to return + **pre-consumption wealth** -- the resources the agent then consumes out of in + the same period, at step 3. Whatever the budget equation subtracts is money + the agent never gets to consume, and whatever it adds is consumable + immediately. If you instead mean a cost that is *not* available for + consumption in that period at all, subtract it from end-of-period assets in + the following period's budget rather than here. + +No other user function may declare `choice`. In particular +`next_period_continuous_state` may not: additional continuous states (e.g. +experience) are settled at step 5, at the end of the previous period, so the +choice about to be made in the current period does not exist yet when they are +computed. `dcegm` raises a `ValueError` at model setup if you declare it there. +Use `lagged_choice` -- an ordinary state variable -- when you need the choice +made in the *previous* period, which is what accumulation processes such as +experience actually depend on. diff --git a/docs/source/guides/two_occupation_model.ipynb b/docs/source/guides/two_occupation_model.ipynb index 2a6a7e28..88c0a70d 100644 --- a/docs/source/guides/two_occupation_model.ipynb +++ b/docs/source/guides/two_occupation_model.ipynb @@ -290,7 +290,7 @@ " exp_green,\n", " exp_red,\n", " asset_end_of_previous_period,\n", - " income_shock_previous_period,\n", + " income_shock,\n", " params,\n", "):\n", " \"\"\"Budget constraint that determines the resource available to the agent in the next period given the state and choice in the previous period.\"\"\"\n", @@ -303,8 +303,8 @@ " )\n", " resource = (\n", " interest_factor * asset_end_of_previous_period\n", - " + (wage + income_shock_previous_period) * (lagged_choice != 2)\n", - " + (wage + income_shock_previous_period) * 0.5 * (lagged_choice == 2)\n", + " + (wage + income_shock) * (lagged_choice != 2)\n", + " + (wage + income_shock) * 0.5 * (lagged_choice == 2)\n", " )\n", " return jnp.maximum(resource, 0.5)" ] @@ -393,7 +393,7 @@ "n_agents = 1000\n", "states_initial = {\n", " \"n_agents\": n_agents,\n", - " \"assets_begin_of_period\": jnp.ones(n_agents),\n", + " \"assets_end_of_previous_period\": jnp.ones(n_agents),\n", " \"exp_green\": jnp.zeros(n_agents),\n", " \"exp_red\": jnp.zeros(n_agents),\n", " \"lagged_choice\": jnp.zeros(n_agents),\n", @@ -565,7 +565,7 @@ " exp_green,\n", " exp_red,\n", " asset_end_of_previous_period,\n", - " income_shock_previous_period,\n", + " income_shock,\n", " params,\n", "):\n", " \"\"\"Budget constraint for two continuous experience variables.\"\"\"\n", @@ -578,8 +578,8 @@ " )\n", " resource = (\n", " interest_factor * asset_end_of_previous_period\n", - " + (wage + income_shock_previous_period) * (lagged_choice != 2)\n", - " + (wage + income_shock_previous_period) * 0.5 * (lagged_choice == 2)\n", + " + (wage + income_shock) * (lagged_choice != 2)\n", + " + (wage + income_shock) * 0.5 * (lagged_choice == 2)\n", " )\n", " return jnp.maximum(resource, 0.5)" ] @@ -652,7 +652,7 @@ "n_agents = 1000\n", "states_initial = {\n", " \"n_agents\": n_agents,\n", - " \"assets_begin_of_period\": jnp.ones(n_agents),\n", + " \"assets_end_of_previous_period\": jnp.ones(n_agents),\n", " \"exp_green\": jnp.zeros(n_agents),\n", " \"exp_red\": jnp.zeros(n_agents),\n", " \"lagged_choice\": jnp.zeros(n_agents),\n", diff --git a/docs/source/guides/two_period_model_tutorial.ipynb b/docs/source/guides/two_period_model_tutorial.ipynb index 255f6bc1..0b28c938 100644 --- a/docs/source/guides/two_period_model_tutorial.ipynb +++ b/docs/source/guides/two_period_model_tutorial.ipynb @@ -397,7 +397,7 @@ " lagged_choice,\n", " ltc,\n", " asset_end_of_previous_period,\n", - " income_shock_previous_period,\n", + " income_shock,\n", " params,\n", "):\n", " interest_factor = 1 + params[\"interest_rate\"]\n", @@ -405,7 +405,7 @@ " wage = params[\"wage_avg\"]\n", " resource = (\n", " interest_factor * asset_end_of_previous_period\n", - " + (wage + income_shock_previous_period) * (1 - lagged_choice)\n", + " + (wage + income_shock) * (1 - lagged_choice)\n", " - ltc * health_costs\n", " )\n", " return jnp.maximum(resource, 0.5)" diff --git a/docs/source/index.rst b/docs/source/index.rst index 23253c5c..c1ba09ab 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -46,6 +46,7 @@ Check out our :ref:`guides` to find information on getting sta background/interface_plots.ipynb background/specify_exogenous_processes.md background/sparsity_conditions + background/choice_dependent_sparsity background/batching .. toctree:: @@ -57,6 +58,7 @@ Check out our :ref:`guides` to find information on getting sta development/changes development/roadmap development/internals/batching + development/internals/law_of_motion .. toctree:: diff --git a/src/dcegm/__init__.py b/src/dcegm/__init__.py index e31f6c02..dd7ebbdd 100644 --- a/src/dcegm/__init__.py +++ b/src/dcegm/__init__.py @@ -1 +1,2 @@ +from dcegm.interfaces.combined_solve import get_solve_from_small_models from dcegm.interfaces.model_class import setup_model diff --git a/src/dcegm/backward_induction.py b/src/dcegm/backward_induction.py index 426f9fd8..957c4797 100644 --- a/src/dcegm/backward_induction.py +++ b/src/dcegm/backward_induction.py @@ -22,6 +22,29 @@ def backward_induction( ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: """Do backward induction and solve for optimal policy and value function. + Structure of the solve. Backward induction runs from the final period back to + period 0, so the final period is filled first and period 0 last: + + 1. ``create_solution_container`` allocates the three arrays that are threaded + through everything below as the scan *carry*: + ``(value_solved, policy_solved, endog_grid_solved)``, each indexed by + state-choice with shape ``(n_state_choices, + n_continuous_state_combinations, n_total_wealth_grid)``. + ``endog_grid_solved`` is ``None`` when ``skip_endog_grid_storage`` is True + (the third carry element is then an empty pytree slot). + 2. ``solve_last_two_periods`` fills the final period (analytic -- everything is + consumed) and the second-to-last one, which is also special because policy and value + in child states are not interpolated and solved directly (last period - consume all). + 3. The remaining periods run through one ``jax.lax.scan`` per *segment*, with + ``solve_single_period`` as the body. Its ``xs`` is the per-batch slice of + ``batch_info``; see that function's docstring for what each entry is. Every + iteration reads its children out of the carry and writes its own rows back + in, which is why batches must be ordered so a state-choice's children are + always solved earlier. + + A segment whose batches do not evenly divide its state-choices has a leftover + batch, applied once after the scan (``batches_cover_all`` is False). + Args: params (dict): Dictionary containing the model parameters. income_shock_draws_unscaled (np.ndarray): 1d array of shape (n_quad_points,) @@ -51,10 +74,9 @@ def backward_induction( income_shock_draws_unscaled * income_shock_std + income_shock_mean ) - # Infer n_continuous_state_combinations from model structure - n_continuous_state_combinations = model_structure["continuous_state_space"][ - next(iter(model_structure["continuous_state_space"])) - ].shape[0] + n_continuous_state_combinations = continuous_states_info[ + "n_continuous_state_combinations" + ] ( value_solved, @@ -110,7 +132,7 @@ def backward_induction( xs=xs, params=params, continuous_grids_info=continuous_states_info, - continuous_state_space=model_structure["continuous_state_space"], + state_choice_space_dict=model_structure["state_choice_space_dict"], income_shocks_scaled=income_shocks_scaled, model_funcs=model_funcs, income_shock_weights=income_shock_weights, @@ -139,6 +161,7 @@ def backward_induction( segment_info["child_states_idxs"], segment_info["state_choices"], segment_info["state_choices_childs"], + segment_info["law_of_motion_arrays"], ), ) @@ -154,6 +177,7 @@ def backward_induction( last_batch_info["child_states_idxs"], last_batch_info["state_choices"], last_batch_info["state_choices_childs"], + last_batch_info["law_of_motion_arrays"], ), ) diff --git a/src/dcegm/check_func_outputs.py b/src/dcegm/check_func_outputs.py index 7649e986..ae6402b4 100644 --- a/src/dcegm/check_func_outputs.py +++ b/src/dcegm/check_func_outputs.py @@ -1,7 +1,7 @@ from jax import numpy as jnp -def check_budget_equation_and_return_wealth_plus_optional_aux(out, optional_aux=False): +def check_budget_equation_and_return_wealth_plus_optional_aux(out, optional_aux): if isinstance(out, tuple): if len(out) == 2: # Check validity of first element diff --git a/src/dcegm/egm/aggregate_marginal_utility.py b/src/dcegm/egm/aggregate_marginal_utility.py index a707ecda..f665f718 100644 --- a/src/dcegm/egm/aggregate_marginal_utility.py +++ b/src/dcegm/egm/aggregate_marginal_utility.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import Tuple, Union import jax.numpy as jnp @@ -7,11 +7,19 @@ def aggregate_marg_utils_and_exp_values( value_state_choice_specific: jnp.ndarray, marg_util_state_choice_specific: jnp.ndarray, reshape_state_choice_vec_to_mat: jnp.ndarray, - taste_shock_scale, - taste_shock_scale_is_scalar, + taste_shock_scale: Union[float, jnp.ndarray], + taste_shock_scale_is_scalar: bool, income_shock_weights: jnp.ndarray, ) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Compute the aggregate marginal utilities and expected values. + """EGM step 2: collapse children's values/marginal utilities to one per state. + + Two reductions, both taken over the child-choice axis with logit choice + probabilities and then over the income-shock axis with quadrature weights: + one on ``value`` (giving the expected value, i.e. the logsum), one on + ``marg_util`` (giving the aggregate marginal utility). Called once per + batch from ``solve_for_interpolated_values``, on the *interpolated* child + continuation values ``interpolate_value_and_marg_util`` (EGM step 1) just + computed. Args: value_state_choice_specific (jnp.ndarray): 3d array of shape @@ -25,7 +33,11 @@ def aggregate_marg_utils_and_exp_values( (n_states_current, n_choices_current) that reshapes the current period vector of feasible state-choice combinations to a matrix of shape (n_choices, n_choices). - taste_shock_scale (float): The taste shock scale. + taste_shock_scale: The taste shock scale -- a scalar, or one value per + state-choice; which of the two is given by + ``taste_shock_scale_is_scalar``. + taste_shock_scale_is_scalar: Whether ``taste_shock_scale`` is a single + scalar shared across all states, or one value per state-choice. income_shock_weights (jnp.ndarray): 1d array of shape (n_stochastic_quad_points,) containing the weights of the income shock quadrature. @@ -36,7 +48,7 @@ def aggregate_marg_utils_and_exp_values( - marg_util (np.ndarray): 2d array of shape (n_states, n_exog_savings) of the state-specific aggregate marginal utilities. - expected_value (np.ndarray): 2d array of shape (n_states, n_exog_savings) - of the state-specific aggregate expected values. + of the state-specific aggregate expected values (the logsum). """ choice_values_per_state = jnp.take( diff --git a/src/dcegm/egm/interpolate_marginal_utility.py b/src/dcegm/egm/interpolate_marginal_utility.py index d3c90f94..c79f3150 100644 --- a/src/dcegm/egm/interpolate_marginal_utility.py +++ b/src/dcegm/egm/interpolate_marginal_utility.py @@ -9,55 +9,88 @@ interp2d_policy_and_value_on_wealth_and_regular_grid, ) from dcegm.interpolation.interpnd_regular import ( - interpnd_policy_and_value_for_child_states_on_regular_grids, + interpnd_policy_and_value_for_child_states_on_own_regular_grids, +) +from dcegm.law_of_motion import ( + calc_law_of_motion, + compute_own_continuous_grid_combos, + compute_own_dj_wealth_grid, ) -from dcegm.law_of_motion import calc_law_of_motion_for_state_choices def interpolate_value_and_marg_util( - model_funcs, - state_choice_vec: Dict[str, int], + model_funcs: Dict[str, Any], + child_state_choices_with_proxy: Dict[str, int], continuous_grids_info: Dict[str, Any], income_shocks_scaled: jnp.ndarray, endog_grid_child_state_choice: jnp.ndarray, policy_child_state_choice: jnp.ndarray, value_child_state_choice: jnp.ndarray, - continuous_state_space, params: Dict[str, float], upper_envelope_method: str, - skip_endog_grid_storage: bool = False, + skip_endog_grid_storage: bool, + law_of_motion_arrays: Dict[str, Any], + state_choice_space_dict: Dict[str, int], ) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Interpolate value and policy for all child states and compute marginal utility. + """EGM step 1: the children's continuation values, seen from this period. + + Called once per batch from ``solve_single_period``. Does two things in order: + + 1. Applies the law of motion (``calc_law_of_motion``) to get, for every child + state-choice, the beginning-of-period wealth and additional continuous + states it would be reached at -- one value per end-of-period assets grid + point, continuous-state combination and income shock draw. + 2. Interpolates each child's *stored* policy and value at those points, and + evaluates the marginal utility there. + + Which interpolation routine runs depends on the model, not on any argument: + + ================================== ========================================== + additional continuous state? routine + ================================== ========================================== + no 1d on wealth alone + yes, ``upper_envelope="fues"`` 2d, irregular in wealth (FUES output is + not on a common grid) + yes, Druedahl-Jorgensen n-d regular, wealth axis shared per + state-choice + ================================== ========================================== Args: - compute_marginal_utility (callable): User-defined function to compute the - agent's marginal utility of consumption. - compute_utility (callable): Function for calculating the utility of consumption. - state_choice_vec (dict): Dictionary containing the state and choice of the agent. - assets_beginning_of_next_period (jnp.ndarray): 2d array of shape - (n_quad_stochastic, n_grid_wealth,) containing the agent's beginning of - period wealth. - endog_grid_child_state_choice (jnp.ndarray): 1d array containing the endogenous - wealth grid of the child state/choice pair. Shape (n_grid_wealth,). - policy_child_state_choice (jnp.ndarray): 1d array containing the - corresponding policy function values of the endogenous wealth grid of the - child state/choice pair. Shape (n_grid_wealth,). - value_child_state_choice (jnp.ndarray): 1d array containing the - corresponding value function values of the endogenous wealth grid of the - child state/choice pair. Shape (n_grid_wealth,). - has_second_continuous_state (bool): Boolean indicating whether the model - features a second continuous state variable. If False, the only - continuous state variable is consumption/savings. - params (dict): Dictionary containing the model parameters. + model_funcs: Processed model functions. + child_state_choices_with_proxy: State-choice dict for the (deduplicated) + children whose stored solution is being read. Carries proxy states + for children reused from another slot; the value/policy are stored on + the proxy's own grid, so its identity is the right one to read them. + continuous_grids_info: ``model_config["continuous_states_info"]``. + income_shocks_scaled: Quadrature points for the income shock, already + scaled. + endog_grid_child_state_choice: The children's stored endogenous wealth + grids, or ``None`` under ``skip_endog_grid_storage`` -- in which case + each child's grid is recomputed on demand from its own identity. + policy_child_state_choice: The children's stored policy. + value_child_state_choice: The children's stored value. + params: Model parameters. + upper_envelope_method: ``"fues"`` or ``"druedahl_jorgensen"``; selects the + interpolation routine together with the presence of a continuous state. + skip_endog_grid_storage: See ``endog_grid_child_state_choice``. + law_of_motion_arrays: The child index/dedup arrays ``calc_law_of_motion`` + reads, carrying only the branch this model takes -- assembled once at + model setup and threaded through the scan (see + ``bundle_law_of_motion_arrays`` in ``batch_creation.py``). Either a + per-child-state-choice representative parent, or the + per-unique-child-state dedup arrays; ``calc_law_of_motion`` decides + which from whether any transition function declares ``choice``. The + child state-choice dict itself is + ``child_state_choices_with_proxy`` above, passed on directly. + state_choice_space_dict: Full state-choice space; ``calc_law_of_motion`` + gathers the representative-parent indices in ``law_of_motion_arrays`` + out of this. Returns: - tuple: - - - value_interp (jnp.ndarray): 2d array of shape (n_wealth_grid, n_income_shocks) - containing the interpolated value function. - - marg_util_interp (jnp.ndarray): 2d array of shape (n_wealth_grid, n_income_shocks) - containing the interpolated marginal utilities for each wealth level and - income shock. + tuple of ``(value_interp, marg_util_interp)``, both shaped + ``(n_child_state_choices, n_continuous_combinations, n_wealth, + n_income_shocks)`` -- ready for ``aggregate_marg_utils_and_exp_values`` to + collapse the choice and income-shock axes. """ # Check if interpolation needs to be multidimensional and irregular @@ -65,16 +98,19 @@ def interpolate_value_and_marg_util( irregular = upper_envelope_method == "fues" # Compute the child continuous-state/wealth transitions on demand for exactly - # this batch's children, instead of reading from a precomputed whole-state-space - # structure (see law_of_motion.py). - law_of_motion = calc_law_of_motion_for_state_choices( - state_choice_vec=state_choice_vec, - continuous_state_space=continuous_state_space, - assets_grid_end_of_period=continuous_grids_info["assets_grid_end_of_period"], + # this batch's children. The continuous grids come from last period, via a + # representative parent state-choice. law_of_motion_arrays was assembled once at + # model setup and threaded in through the scan (see batch_creation.py). + law_of_motion = calc_law_of_motion( + law_of_motion_arrays=law_of_motion_arrays, + state_choice_space_dict=state_choice_space_dict, income_shocks_scaled=income_shocks_scaled, params=params, model_funcs=model_funcs, has_additional_continuous_states=multi_dim, + additional_continuous_state_names=continuous_grids_info[ + "additional_continuous_state_names" + ], ) wealth_child_states = law_of_motion["assets_begin_of_period"] continuous_states_next = law_of_motion["continuous_states"] @@ -87,7 +123,7 @@ def interpolate_value_and_marg_util( return _interpolate_value_and_marg_util_2d_irregular( compute_marginal_utility=compute_marginal_utility, compute_utility=compute_utility, - state_choice_vec=state_choice_vec, + state_choice_vec=child_state_choices_with_proxy, continuous_grids_info=continuous_grids_info, continuous_states_next=continuous_states_next, wealth_child_states=wealth_child_states, @@ -96,13 +132,14 @@ def interpolate_value_and_marg_util( value_child_state_choice=value_child_state_choice, params=params, discount_factor=discount_factor, + continuous_grid_functions=model_funcs["continuous_grid_functions"], ) elif multi_dim & (not irregular): return _interpolate_value_and_marg_util_nd_regular( compute_marginal_utility=compute_marginal_utility, compute_utility=compute_utility, - state_choice_vec=state_choice_vec, + state_choice_vec=child_state_choices_with_proxy, continuous_grids_info=continuous_grids_info, continuous_states_next=continuous_states_next, wealth_child_states=wealth_child_states, @@ -112,32 +149,45 @@ def interpolate_value_and_marg_util( params=params, discount_factor=discount_factor, skip_endog_grid_storage=skip_endog_grid_storage, + continuous_grid_functions=model_funcs["continuous_grid_functions"], ) else: # Selects inside if jorgensen_druedahl or fues (different treatment of budget constraint) - # Under DJ, the wealth grid is a fixed constant (not batched over children), so - # it is passed with in_axes=None, broadcast only across the (small) continuous - # combo axis rather than the (potentially large) children axis. - if skip_endog_grid_storage: - dj_wealth_grid = continuous_grids_info["dj_wealth_grid"] - endog_grid_arg = jnp.broadcast_to( - dj_wealth_grid, - (policy_child_state_choice.shape[1], dj_wealth_grid.shape[0]), - ) - endog_grid_in_axes = None - else: - endog_grid_arg = endog_grid_child_state_choice - endog_grid_in_axes = 0 + # Under DJ with skip_endog_grid_storage, each child's own wealth grid is + # evaluated inside interp1d_value_and_marg_util_for_state_choice, right + # next to where it's consumed, after vmapping down to a single (child) + # state-choice -- not precomputed for the whole batch upfront in a + # separate vmap and broadcast in as a paired array. state_choice_vec is + # the CHILD's own identity here (reading each child's own stored + # solution, interpolated on its own grid), not a parent, same reasoning + # as compute_own_continuous_grid_combos elsewhere in this file. + endog_grid_in_axes = None if skip_endog_grid_storage else 0 + endog_grid_arg = ( + None if skip_endog_grid_storage else endog_grid_child_state_choice + ) interp_for_single_state_choice = vmap( interp1d_value_and_marg_util_for_state_choice, - in_axes=(None, None, 0, 0, endog_grid_in_axes, 0, 0, None, None, None), + in_axes=( + None, + None, + 0, + 0, + endog_grid_in_axes, + 0, + 0, + None, + None, + None, + None, + None, + ), ) return interp_for_single_state_choice( compute_marginal_utility, compute_utility, - state_choice_vec, + child_state_choices_with_proxy, wealth_child_states, endog_grid_arg, policy_child_state_choice, @@ -145,6 +195,8 @@ def interpolate_value_and_marg_util( params, discount_factor, upper_envelope_method == "druedahl_jorgensen", + skip_endog_grid_storage, + model_funcs["continuous_grid_functions"], ) @@ -159,6 +211,8 @@ def interp1d_value_and_marg_util_for_state_choice( params: Dict[str, float], discount_factor: float, use_dj_interpolation: bool, + skip_endog_grid_storage: bool, + continuous_grid_functions: Dict[str, Callable], ) -> Tuple[jnp.ndarray, jnp.ndarray]: """Interpolate value and policy for given child state and compute marginal utility. @@ -193,6 +247,18 @@ def interp1d_value_and_marg_util_for_state_choice( containing the interpolated value function. """ + if skip_endog_grid_storage: + # Own DJ wealth grid, evaluated here for this one (already-sliced) + # state-choice, right next to where it's consumed below -- instead of + # precomputing the whole batch's grids upfront in a separate vmap and + # broadcasting the result in as a paired array. + own_dj_wealth_grid = compute_own_dj_wealth_grid( + state_choice_vec, continuous_grid_functions + ) + endog_grid_child_state_choice = jnp.broadcast_to( + own_dj_wealth_grid[None, :], + (policy_child_state_choice.shape[0], own_dj_wealth_grid.shape[0]), + ) endog_grid_child_state_choice = jnp.asarray(endog_grid_child_state_choice) policy_child_state_choice = jnp.asarray(policy_child_state_choice) value_child_state_choice = jnp.asarray(value_child_state_choice) @@ -236,8 +302,6 @@ def interp_on_single_wealth_point(wealth_point): value_interp, marg_util_interp = interp_over_single_wealth_and_income_shock_draw( assets_points ) - value_interp = jnp.asarray(value_interp) - marg_util_interp = jnp.asarray(marg_util_interp) # Add it back in the beginning return value_interp[None, :, :], marg_util_interp[None, :, :] @@ -255,6 +319,7 @@ def _interpolate_value_and_marg_util_2d_irregular( value_child_state_choice: jnp.ndarray, params: Dict[str, float], discount_factor: float, + continuous_grid_functions: Dict[str, Callable], ) -> Tuple[jnp.ndarray, jnp.ndarray]: """Interpolate value and marginal utility on the irregular FUES 2D grid. @@ -270,6 +335,9 @@ def _interpolate_value_and_marg_util_2d_irregular( ] continuous_state_child_states = continuous_states_next[continuous_state_name] + # Each child's own grid is now computed inside + # interp2d_value_and_marg_util_for_state_choice, on demand, after vmapping + # down to a single (child) state-choice -- see its docstring. interp_for_single_state_choice = vmap( interp2d_value_and_marg_util_for_state_choice, in_axes=( @@ -277,6 +345,7 @@ def _interpolate_value_and_marg_util_2d_irregular( None, 0, None, + None, 0, 0, 0, @@ -291,7 +360,8 @@ def _interpolate_value_and_marg_util_2d_irregular( compute_marginal_utility, compute_utility, state_choice_vec, - continuous_grids_info["additional_continuous_state_grids"], + continuous_grid_functions, + continuous_grids_info["additional_continuous_state_names"], wealth_child_states, continuous_state_child_states, endog_grid_child_state_choice, @@ -314,7 +384,8 @@ def _interpolate_value_and_marg_util_nd_regular( value_child_state_choice: jnp.ndarray, params: Dict[str, float], discount_factor: float, - skip_endog_grid_storage: bool = False, + skip_endog_grid_storage: bool, + continuous_grid_functions: Dict[str, Callable], ) -> Tuple[jnp.ndarray, jnp.ndarray]: """Interpolate value and marginal utility on the regular n-D grid. @@ -327,21 +398,39 @@ def _interpolate_value_and_marg_util_nd_regular( name: continuous_states_next[name] for name in continuous_state_names } - # interpnd_policy_and_value_for_child_states_on_regular_grids already treats - # wealth_grid as a single unbatched 1D array (it is not itself a mapped vmap - # argument), so under DJ we can pass the fixed grid directly with no broadcast. - wealth_grid = ( - continuous_grids_info["dj_wealth_grid"] - if skip_endog_grid_storage - else endog_grid_child_state_choice[0, 0] - ) + # Every child's own grid (for each additional continuous state, and for + # wealth) is now evaluated inside + # interpnd_policy_and_value_for_child_states_on_own_regular_grids's own + # per-child vmap, right next to the index/weight computation that consumes + # it -- not precomputed for the whole batch upfront in a separate vmap here + # and fed in as a paired array. state_choice_vec is the CHILD's own identity + # (we're reading each child's own stored solution, interpolated on its own + # grid), not a parent -- unlike law_of_motion.py's grid selection for the + # *transition into* a child, which needs a representative parent. + # + # wealth_grid_func is only reachable via the grid-function branch when + # skip_endog_grid_storage -- with a single choice the DJ upper envelope is + # skipped entirely (see check_model_config.py), so the stored endog_grid is a + # genuinely different (non-fixed) grid there instead, handled by the else + # branch below (see process_continuous_grid_functions, which forbids a + # state-specific assets_begin_of_period whenever skip_endog_grid_storage is + # False and additional continuous states exist). + if skip_endog_grid_storage: + + def wealth_grid_func(state_dict): + return compute_own_dj_wealth_grid(state_dict, continuous_grid_functions) + + else: + shared_wealth_grid = endog_grid_child_state_choice[0, 0] + + def wealth_grid_func(state_dict): + return shared_wealth_grid policy_interp, value_interp = ( - interpnd_policy_and_value_for_child_states_on_regular_grids( - additional_continuous_state_grids=continuous_grids_info[ - "additional_continuous_state_grids" - ], - wealth_grid=wealth_grid, + interpnd_policy_and_value_for_child_states_on_own_regular_grids( + continuous_grid_functions=continuous_grid_functions, + additional_continuous_state_names=continuous_state_names, + wealth_grid_func=wealth_grid_func, policy_grid_child_states=policy_child_state_choice, value_grid_child_states=value_child_state_choice, continuous_state_child_states=continuous_state_child_states, @@ -416,7 +505,8 @@ def interp2d_value_and_marg_util_for_state_choice( compute_marginal_utility: Callable, compute_utility: Callable, state_choice_vec: Dict[str, int], - continuous_state_space: Dict[str, jnp.ndarray], + continuous_grid_functions: Dict[str, Callable], + additional_continuous_state_names, assets_beginning_of_next_period: jnp.ndarray, continuous_state_beginning_of_next_period: jnp.ndarray, endog_grid_child_state_choice: jnp.ndarray, @@ -458,6 +548,20 @@ def interp2d_value_and_marg_util_for_state_choice( containing the interpolated value function. """ + # This state-choice's own grid, on demand -- computed here, after vmapping + # down to a single (child) state-choice, instead of precomputing the whole + # batch's grids upfront in a separate vmap and feeding the result in as a + # paired array. state_choice_vec is the CHILD's own identity here (we're + # reading this child's own stored solution, interpolated on its own grid), + # not a parent -- unlike law_of_motion.py's grid selection for the + # *transition into* a child, which needs a representative parent. Grids live + # on the state-choice space (that's where the solution itself lives), so + # "choice" is a legitimate part of that identity, same as for + # compute_utility/compute_marginal_utility below. + continuous_state_space = compute_own_continuous_grid_combos( + state_choice_vec, continuous_grid_functions, additional_continuous_state_names + ) + # FUES-only (irregular) branch: check_model_config.py enforces at most one # additional continuous state here, so this is the only continuous state besides # assets_begin_of_period. Pass it under its actual state name. diff --git a/src/dcegm/egm/solve_euler_equation.py b/src/dcegm/egm/solve_euler_equation.py index a3f859ab..bff06c3f 100644 --- a/src/dcegm/egm/solve_euler_equation.py +++ b/src/dcegm/egm/solve_euler_equation.py @@ -5,10 +5,11 @@ from jax import numpy as jnp from jax import vmap +from dcegm.law_of_motion import compute_own_continuous_grid_combos + def calculate_candidate_solutions_from_euler_equation( continuous_grids_info: Dict[str, Any], - continuous_state_space: Dict[str, jnp.ndarray], marg_util_next: jnp.ndarray, emax_next: jnp.ndarray, state_choice_mat: Dict[str, jnp.ndarray], @@ -16,7 +17,42 @@ def calculate_candidate_solutions_from_euler_equation( model_funcs: Dict[str, Any], params: Dict[str, float], ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Calculate candidates for the optimal policy and value function.""" + """EGM step 3: invert the Euler equation into (endog_grid, policy, value) + candidates. + + This solves/stores each state-choice in ``state_choice_mat``'s *own* problem (as + opposed to interpolating someone else's child, where the parent/child distinction in + law_of_motion.py matters) -- so the combo grid used here is just each row's own + grid, no representative-parent selection needed. Called once per batch from + ``solve_for_interpolated_values``, on the aggregate marginal utility/expected + value ``aggregate_marg_utils_and_exp_values`` (EGM step 2) just computed for + every *child* state. "Candidates" because these are not yet the final + solution: ``run_upper_envelope`` still has to discard the ones that are + not optimal (the DC-EGM refinement step). + + Args: + continuous_grids_info: ``model_config["continuous_states_info"]``. + marg_util_next: Aggregate marginal utility per child state and + exogenous-savings grid point, shape ``(n_child_states, + n_exog_savings)``. + emax_next: Aggregate expected value per child state and exogenous- + savings grid point, same shape as ``marg_util_next``. + state_choice_mat: State-choice dict for the rows being solved -- this + period's own state-choices, not the children. + idx_post_decision_child_states: For each (row, stochastic + realisation), the position of the resulting child state in + ``marg_util_next``/``emax_next``; out-of-bounds entries (states + not reachable from that row) are filled with ``nan`` and + integrated out via ``jnp.nansum`` in ``solve_euler_equation``. + model_funcs: Processed model functions. + params: Model parameters. + + Returns: + tuple ``(endog_grid, value, policy, expected_value)``, each shape + ``(n_state_choices, n_continuous_combinations, n_exog_savings)`` -- + the pre-upper-envelope candidate solution for every state-choice. + + """ feasible_marg_utils_child = jnp.take( marg_util_next, idx_post_decision_child_states, @@ -38,22 +74,15 @@ def calculate_candidate_solutions_from_euler_equation( value, expected_value, ) = vmap( - vmap( - vmap( - compute_optimal_policy_and_value, - in_axes=(1, 1, None, 0, None, None, None), - ), - in_axes=(1, 1, 0, None, None, None, None), - ), - in_axes=(0, 0, None, None, 0, None, None), + compute_optimal_policy_and_value_for_state_choice, + in_axes=(0, 0, 0, None, None, None), )( feasible_marg_utils_child, feasible_emax_child, - continuous_state_space, - continuous_grids_info["assets_grid_end_of_period"], state_choice_mat, model_funcs, params, + continuous_grids_info, ) return ( @@ -64,6 +93,57 @@ def calculate_candidate_solutions_from_euler_equation( ) +def compute_optimal_policy_and_value_for_state_choice( + feasible_marg_utils_child: jnp.ndarray, + feasible_emax_child: jnp.ndarray, + state_choice_vec: Any, + model_funcs: Dict[str, Any], + params: Dict[str, float], + continuous_grids_info: Dict[str, Any], +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Compute EGM candidates for one state-choice, across its own combo/wealth grid. + + Builds this state-choice's own combo grid *and* its own ``assets_end_of_period`` + grid on demand *after* vmapping down to a single state-choice, instead of + precomputing the whole batch's grids upfront in a separate vmap and feeding the + result in as a paired array. This state-choice is solving/storing its *own* + problem (not interpolating someone else's child), so its own grids are used + directly -- no representative-parent selection needed, unlike + law_of_motion.py's grid selection for a transition *into* a state. Grids live on the + state-choice space (that's where the solution itself lives), so ``state_choice_vec`` + -- including "choice" -- is exactly the identity a grid may depend on. + + """ + if continuous_grids_info["has_additional_continuous_state"]: + own_continuous_state_vec = compute_own_continuous_grid_combos( + state_choice_vec, + model_funcs["continuous_grid_functions"], + continuous_grids_info["additional_continuous_state_names"], + ) + else: + own_continuous_state_vec = {"dummy_cont": jnp.zeros(1)} + + assets_grid_end_of_period = model_funcs["continuous_grid_functions"][ + "assets_end_of_period" + ](**state_choice_vec) + + return vmap( + vmap( + compute_optimal_policy_and_value, + in_axes=(1, 1, None, 0, None, None, None), + ), + in_axes=(1, 1, 0, None, None, None, None), + )( + feasible_marg_utils_child, + feasible_emax_child, + own_continuous_state_vec, + assets_grid_end_of_period, + state_choice_vec, + model_funcs, + params, + ) + + def compute_optimal_policy_and_value( marg_util_next: jnp.ndarray, emax_next: jnp.ndarray, diff --git a/src/dcegm/final_periods.py b/src/dcegm/final_periods.py index dca3a2d1..92a60cbf 100644 --- a/src/dcegm/final_periods.py +++ b/src/dcegm/final_periods.py @@ -1,6 +1,6 @@ """Wrapper to solve the final period of the model.""" -from typing import Any, Callable, Dict, Tuple +from typing import Any, Dict, Optional, Tuple import jax.numpy as jnp from jax import vmap @@ -8,7 +8,10 @@ from dcegm.check_func_outputs import ( check_budget_equation_and_return_wealth_plus_optional_aux, ) -from dcegm.law_of_motion import calc_law_of_motion_for_state_choices +from dcegm.law_of_motion import ( + calc_law_of_motion, + compute_own_continuous_grid_combos, +) from dcegm.solve_single_period import solve_for_interpolated_values @@ -21,38 +24,64 @@ def solve_last_two_periods( model_funcs: Dict[str, Any], upper_envelope_method: str, skip_endog_grid_storage: bool, - last_two_period_batch_info, - value_solved, - policy_solved, - endog_grid_solved, - debug_info, -): - """Solves the last two periods of the model. - - The last two periods are solved using the EGM algorithm. The last period is - solved using the user-specified utility function and the second to last period - is solved using the user-specified utility function and the user-specified - bequest function. + last_two_period_batch_info: Dict[str, Any], + value_solved: jnp.ndarray, + policy_solved: jnp.ndarray, + endog_grid_solved: Optional[jnp.ndarray], + debug_info: Optional[Dict[str, bool]], +) -> Tuple[jnp.ndarray, ...]: + """Solve the final period and the second-to-last period. + + Called once, outside the ``jax.lax.scan`` over the remaining periods (see + ``backward_induction.py``), because both periods are special: + + 1. The final period is solved analytically via ``solve_final_period`` -- + everything is consumed, so there is no continuation value to + interpolate or Euler equation to invert. + 2. The second-to-last period reaches its own solution the normal EGM way + (``solve_for_interpolated_values``), but its children (the final + period) are the ``value``/``marginal_utility`` just computed in step 1 + directly, rather than the *stored*, then re-interpolated, solution + ``interpolate_value_and_marg_util`` reads for every other period. Args: - wealth_beginning_of_period (np.ndarray): 2d array of shape - (n_states, n_grid_wealth) of the wealth at the beginning of the - period. - params (dict): Dictionary of model parameters. - compute_utility (callable): User supplied utility function. - compute_marginal_utility (callable): User supplied marginal utility - function. - last_two_period_batch_info (dict): Dictionary containing information about the batch - size and the state space. - value_solved (np.ndarray): 3d array of shape - (n_states, n_grid_wealth, n_income_shocks) of the value function for - all states, end of period assets, and income shocks. - endog_grid_solved (np.ndarray): 3d array of shape - (n_states, n_grid_wealth, n_income_shocks) of the endogenous grid - for all states, end of period assets, and income shocks. + params: Model parameters. + continuous_states_info: ``model_config["continuous_states_info"]``. + model_structure: Model structure, in particular + ``state_space_dict``/``state_choice_space_dict``, used to gather + representative-parent and unique-child-state dicts for the law of + motion (see ``solve_final_period``). + income_shocks_scaled: Quadrature points for the income shock, already + scaled by its mean and standard deviation. + income_shock_weights: Quadrature weights matching + ``income_shocks_scaled``. + model_funcs: Processed model functions. + upper_envelope_method: ``"fues"`` or ``"druedahl_jorgensen"``. + skip_endog_grid_storage: Whether the endogenous grid is stored at all; + see ``endog_grid_solved``. + last_two_period_batch_info: Batch information for the final and + second-to-last period -- e.g. which final-period state-choices map + to which second-to-last-period parents; see + ``pre_processing/batches/last_two_periods.py``. + value_solved: The solution container, indexed by state-choice, filled + so far (empty at this point, since this runs before the main + scan). Shape ``(n_state_choices, n_continuous_state_combinations, + n_total_wealth_grid)``. + policy_solved: Same shape as ``value_solved``. + endog_grid_solved: Same shape as ``value_solved``, or ``None`` when + ``skip_endog_grid_storage`` is True. + debug_info: ``None`` in the normal solve. When given (with + ``"return_candidates"``), the pre-upper-envelope candidate + solutions for the second-to-last period are also returned. + + Returns: + ``(value_solved, policy_solved, endog_grid_solved)`` with the final + and second-to-last period filled in, plus the three candidate arrays + when ``debug_info["return_candidates"]`` is True. """ batch_info = last_two_period_batch_info + ( value_solved, policy_solved, @@ -60,13 +89,12 @@ def solve_last_two_periods( value_interp_final_period, marginal_utility_final_last_period, ) = solve_final_period( - idx_state_choices_final_period=batch_info["idx_state_choices_final_period"], - state_choice_mat_final_period=batch_info["state_choice_mat_final_period"], + batch_info=batch_info, + model_structure=model_structure, income_shocks_scaled=income_shocks_scaled, continuous_states_info=continuous_states_info, upper_envelope_method=upper_envelope_method, skip_endog_grid_storage=skip_endog_grid_storage, - model_structure=model_structure, params=params, model_funcs=model_funcs, value_solved=value_solved, @@ -75,14 +103,11 @@ def solve_last_two_periods( ) # Check if we have a scalar taste shock scale or state specific. Extract in each of the cases. - if model_funcs["taste_shock_function"]["taste_shock_scale_is_scalar"]: - taste_shock_scale = model_funcs["taste_shock_function"][ - "read_out_taste_shock_scale" - ](params) + ts_function = model_funcs["taste_shock_function"] + if ts_function["taste_shock_scale_is_scalar"]: + taste_shock_scale = ts_function["read_out_taste_shock_scale"](params) else: - taste_shock_scale_per_state_func = model_funcs["taste_shock_function"][ - "taste_shock_scale_per_state" - ] + taste_shock_scale_per_state_func = ts_function["taste_shock_scale_per_state"] taste_shock_scale = vmap(taste_shock_scale_per_state_func, in_axes=(0, None))( last_two_period_batch_info["state_choice_mat_final_period"], params ) @@ -98,13 +123,10 @@ def solve_last_two_periods( "state_to_choices_final_period" ], taste_shock_scale=taste_shock_scale, - taste_shock_scale_is_scalar=model_funcs["taste_shock_function"][ - "taste_shock_scale_is_scalar" - ], + taste_shock_scale_is_scalar=ts_function["taste_shock_scale_is_scalar"], params=params, income_shock_weights=income_shock_weights, continuous_grids_info=continuous_states_info, - continuous_state_space=model_structure["continuous_state_space"], model_funcs=model_funcs, debug_info=debug_info, ) @@ -122,7 +144,7 @@ def solve_last_two_periods( out_dict_second_last["endog_grid"] ) - # If we do not call the function in debug mode. Assign everything and return + # If we do not call the function in debug mode. Return arrays if debug_info is None: return ( value_solved, @@ -151,60 +173,105 @@ def solve_last_two_periods( def solve_final_period( - idx_state_choices_final_period, - state_choice_mat_final_period, + batch_info: Dict[str, Any], + model_structure: Dict[str, Any], income_shocks_scaled: jnp.ndarray, continuous_states_info: Dict[str, Any], upper_envelope_method: str, skip_endog_grid_storage: bool, - model_structure: Dict[str, Any], params: Dict[str, float], model_funcs: Dict[str, Any], - value_solved, - policy_solved, - endog_grid_solved, + value_solved: jnp.ndarray, + policy_solved: jnp.ndarray, + endog_grid_solved: Optional[jnp.ndarray], ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Compute solution to final period for policy and value function. + """Compute the final period's solution analytically (consumption = wealth). + + The final period's own state-choices act as *children* of the second-to- + last period's saving choices for the law of motion: ``calc_law_of_motion`` + is called exactly as it would be for any other period, with the final + period's own state-choices as ``child_state_choices`` and the + second-to-last period's representative parent supplying the grid the + transition is evaluated over (see ``rep_parent_idx_per_state``/ + ``rep_parent_idx_per_state_choice`` below). Since nothing is saved in the + final period, the resulting beginning-of-period wealth is both the + "child" wealth from that transition and this period's own wealth grid to + solve over -- there is no further recursion. - In the last period, everything is consumed, i.e. consumption = savings. Args: - + batch_info: Final-period batch information -- which state-choices are + solved, their representative second-to-last-period parents (at + both state-choice and unique-state granularity), and the + state-level dedup gather index; see + ``pre_processing/batches/last_two_periods.py``. + model_structure: Model structure, in particular + ``state_space_dict``/``state_choice_space_dict``, read via + ``batch_info``'s index arrays to build the law-of-motion inputs. + income_shocks_scaled: Quadrature points for the income shock, already + scaled by its mean and standard deviation. + continuous_states_info: ``model_config["continuous_states_info"]``. + upper_envelope_method: ``"fues"`` or ``"druedahl_jorgensen"``; selects + which wealth grid (``assets_begin_of_period`` vs. + ``assets_end_of_period``) the final period's solution is stored + on. + skip_endog_grid_storage: Whether the endogenous grid is stored at all; + see ``endog_grid_solved``. + params: Model parameters. + model_funcs: Processed model functions; the final-period-specific + ``compute_utility_final``/``compute_marginal_utility_final`` are + used here instead of the regular-period ones. + value_solved: The solution container, indexed by state-choice, filled + so far. Shape ``(n_state_choices, n_continuous_state_combinations, + n_total_wealth_grid)``. + policy_solved: Same shape as ``value_solved``. + endog_grid_solved: Same shape as ``value_solved``, or ``None`` when + ``skip_endog_grid_storage`` is True. Returns: tuple: - - marginal_utilities_choices (np.ndarray): 3d array of shape - (n_states, n_grid_wealth, n_income_shocks) of the marginal utility of - consumption for all final states, end of period assets, and - income shocks. - - final_value (np.ndarray): 3d array of shape - (n_states, n_grid_wealth, n_income_shocks) of the optimal - value function for all final states, end of period assets, and - income shocks. - - final_policy (np.ndarray): 3d array of shape - (n_states, n_grid_wealth, n_income_shocks) of the optimal - policy for all final states, end of period assets, and - income shocks. + + - value_solved: ``value_solved`` with the final period's rows filled + in. + - policy_solved: Likewise for policy. + - endog_grid_solved: Likewise for the endogenous grid (unchanged when + ``skip_endog_grid_storage``). + - value: The final period's own value, shape + ``(n_final_state_choices, n_continuous_combinations, n_exog_savings, + n_income_shocks)`` -- fed directly into + ``solve_for_interpolated_values`` for the second-to-last period, + bypassing the usual store-then-reinterpolate path. + - marg_util: The final period's own marginal utility, same shape. """ compute_utility = model_funcs["compute_utility_final"] compute_marginal_utility = model_funcs["compute_marginal_utility_final"] - law_of_motion_final_period = calc_law_of_motion_for_state_choices( - state_choice_vec=state_choice_mat_final_period, - continuous_state_space=model_structure["continuous_state_space"], - assets_grid_end_of_period=continuous_states_info["assets_grid_end_of_period"], + idx_state_choices_final_period = batch_info["idx_state_choices_final_period"] + state_choice_mat_final_period = batch_info["state_choice_mat_final_period"] + + has_additional_continuous_states = continuous_states_info[ + "has_additional_continuous_state" + ] + additional_continuous_state_names = continuous_states_info[ + "additional_continuous_state_names" + ] + + # Then call the law of motion to get the continuous states and wealth at the + # final period. law_of_motion_arrays was assembled once at model setup, keeping + # only the branch this model takes (see bundle_law_of_motion_arrays in + # batch_creation.py); the final period's own state-choices are the children. + final_period_cont_states = calc_law_of_motion( + law_of_motion_arrays=batch_info["law_of_motion_arrays"], + state_choice_space_dict=model_structure["state_choice_space_dict"], income_shocks_scaled=income_shocks_scaled, params=params, model_funcs=model_funcs, - has_additional_continuous_states=continuous_states_info[ - "has_additional_continuous_state" - ], + has_additional_continuous_states=has_additional_continuous_states, + additional_continuous_state_names=additional_continuous_state_names, ) - wealth_child_states_final_period = law_of_motion_final_period[ - "assets_begin_of_period" - ] - continuous_state_final = law_of_motion_final_period["continuous_states"] + wealth_final_period = final_period_cont_states["assets_begin_of_period"] + continuous_state_final = final_period_cont_states["continuous_states"] value, marg_util = vmap( vmap( @@ -221,34 +288,31 @@ def solve_final_period( )( state_choice_mat_final_period, continuous_state_final, - wealth_child_states_final_period, + wealth_final_period, params, compute_utility, compute_marginal_utility, ) - if continuous_states_info["has_additional_continuous_state"]: - # We also need to solve at the state space and the child states to store correctly - # For Druedahl Jorgensen wealth needs to be assets_begin_of_period + if ( + continuous_states_info["has_additional_continuous_state"] + or upper_envelope_method == "druedahl_jorgensen" + ): + # For Druedahl-Jorgensen the storage grid is assets_begin_of_period; for + # FUES with an additional continuous state it's assets_end_of_period. + # calc_value_and_budget_for_state_choice handles both, and also the + # no-additional-continuous-state Druedahl-Jorgensen case (a size-1 dummy + # combo axis, same convention as solve_euler_equation.py/ + # solve_single_period.py) -- so one branch now covers what used to be two. assets_begin = "assets_begin_of_period" in continuous_states_info.keys() - if upper_envelope_method == "druedahl_jorgensen": - asset_grid = continuous_states_info["assets_begin_of_period"] - else: - asset_grid = continuous_states_info["assets_grid_end_of_period"] values_regular, wealth_at_regular = vmap( - vmap( - vmap( - calc_value_and_budget_for_each_gridpoint, - in_axes=(None, None, 0, None, None, None, None), - ), - in_axes=(None, 0, None, None, None, None, None), - ), + calc_value_and_budget_for_state_choice, in_axes=(0, None, None, None, None, None, None), )( state_choice_mat_final_period, - model_structure["continuous_state_space"], - asset_grid, + model_funcs["continuous_grid_functions"], + continuous_states_info["additional_continuous_state_names"], params, compute_utility, model_funcs["compute_assets_begin_of_period"], @@ -262,7 +326,7 @@ def solve_final_period( middle_of_draws = int((value.shape[3] - 1) / 2) value_final = value[:, :, :, middle_of_draws] - wealth_to_save = wealth_child_states_final_period[:, :, :, middle_of_draws] + wealth_to_save = wealth_final_period[:, :, :, middle_of_draws] sort_idx = jnp.argsort(wealth_to_save, axis=2) wealth_sorted = jnp.take_along_axis(wealth_to_save, sort_idx, axis=2) values_sorted = jnp.take_along_axis(value_final, sort_idx, axis=2) @@ -326,6 +390,75 @@ def calc_value_and_marg_util_for_each_gridpoint( return value, marg_util +def calc_value_and_budget_for_state_choice( + state_choice_vec, + continuous_grid_functions, + additional_continuous_state_names, + params, + compute_utility, + compute_assets_begin_of_period, + assets_begin, +): + """Compute the final period's own value/budget for one state-choice. + + Builds this state-choice's own combo axis (what its own solve/storage is indexed + against) on demand *after* vmapping down to a single state-choice -- the final- + period analog of solve_euler_equation.py's job for every other period, on its own + separate code path since the final period has no continuation value and so doesn't + go through solve_euler_equation.py at all. ``state_choice_vec`` already is each + row's own identity here (no parent/child ambiguity, we're solving each row's own + terminal problem), so no representative-parent selection is needed, unlike + law_of_motion.py's grid selection for a transition *into* a state. Grids live on the + state-choice space (that's where the solution itself lives), so ``state_choice_vec`` + -- including "choice" -- is exactly the identity a grid may depend on. + + This state-choice's own asset grid (``assets_begin_of_period`` for + Druedahl-Jorgensen, ``assets_end_of_period`` for FUES) is evaluated on demand + right here, after vmapping down to a single state-choice, instead of being + precomputed for the whole batch upfront in a separate vmap and fed in as a + paired array -- same self-referential pattern as ``own_continuous_state_vec`` + below. The zero point that Druedahl-Jorgensen prepends elsewhere + (``compute_own_dj_wealth_grid``) is *not* added here: this is the raw + grid_func output, and the caller appends a single zero point uniformly for + every branch further down (see ``zeros_to_append`` in ``solve_final_period``). + + With no additional continuous state, ``additional_continuous_state_names`` is + empty and there is nothing to mesh -- a size-1 ``dummy_cont`` placeholder + stands in instead, the same convention ``solve_euler_equation.py`` and + ``solve_single_period.py`` already use, letting this one function (and the + vmap below) cover the Druedahl-Jorgensen no-additional-continuous-state case + too, rather than needing a dedicated combo-axis-free sibling. Downstream user + functions (``compute_utility``) silently ignore the extra ``dummy_cont`` + kwarg, same as they already do on those other paths. + + """ + if additional_continuous_state_names: + own_continuous_state_vec = compute_own_continuous_grid_combos( + state_choice_vec, + continuous_grid_functions, + additional_continuous_state_names, + ) + else: + own_continuous_state_vec = {"dummy_cont": jnp.zeros(1)} + grid_name = "assets_begin_of_period" if assets_begin else "assets_end_of_period" + asset_grid = continuous_grid_functions[grid_name](**state_choice_vec) + return vmap( + vmap( + calc_value_and_budget_for_each_gridpoint, + in_axes=(None, None, 0, None, None, None, None), + ), + in_axes=(None, 0, None, None, None, None, None), + )( + state_choice_vec, + own_continuous_state_vec, + asset_grid, + params, + compute_utility, + compute_assets_begin_of_period, + assets_begin, + ) + + def calc_value_and_budget_for_each_gridpoint( state_choice_vec, continuous_state_vec, @@ -335,18 +468,20 @@ def calc_value_and_budget_for_each_gridpoint( compute_assets_begin_of_period, assets_begin, ): - state_vec = state_choice_vec.copy() - state_vec.pop("choice") - if assets_begin: # If assets begin, the grid is directly the assets we start from wealth_final_period = asset_grid_point_end_of_previous_period else: + # "choice" is passed through (not stripped): a budget equation may declare + # it and return a different budget per choice, same as in law_of_motion.py. + # Functions that don't declare it are unaffected -- + # determine_function_arguments_and_partial_model_specs filters kwargs down + # to each function's own signature. out_budget = compute_assets_begin_of_period( - **state_vec, + **state_choice_vec, **continuous_state_vec, asset_end_of_previous_period=asset_grid_point_end_of_previous_period, - income_shock_previous_period=jnp.array(0.0), + income_shock=jnp.array(0.0), params=params, ) wealth_final_period = check_budget_equation_and_return_wealth_plus_optional_aux( diff --git a/src/dcegm/interfaces/combined_solve.py b/src/dcegm/interfaces/combined_solve.py new file mode 100644 index 00000000..52e6198d --- /dev/null +++ b/src/dcegm/interfaces/combined_solve.py @@ -0,0 +1,346 @@ +"""Assemble one big model's solution from several smaller sub-model solves. + +When the pooled solve of a model does not fit into (GPU) memory, the model can be split +by hand into several smaller ``setup_model`` instances whose discrete state-choice +spaces together tile the big model's state-choice space exactly -- e.g. one sub-model +per invariant "type" (education x sex): because such types never change along any +transition or proxy, the problem is block-diagonal in them and each block is an exact, +standalone sub-problem. + +``get_solve_from_small_models`` builds the big model, checks that the supplied sub- +models really are such a tiling, and returns a solve function that solves each sub-model +and scatters its solution into the big solution container. Sequentially, only one +block's device arrays are ever live, so the big solution never has to be solved (only +stored) in one piece. + +""" + +import warnings +from typing import Callable, Dict, List, Optional + +import jax +import jax.numpy as jnp +import numpy as np + +from dcegm.interfaces.index_functions import ( + get_state_choice_index_per_discrete_states_and_choices, +) +from dcegm.interfaces.model_class import setup_model +from dcegm.interfaces.sol_interface import model_solved + + +def get_solve_from_small_models( + small_models: List[setup_model], + parallel: bool, + model_config: Optional[Dict] = None, + model_specs: Optional[Dict] = None, + utility_functions: Optional[Dict[str, Callable]] = None, + utility_functions_final_period: Optional[Dict[str, Callable]] = None, + budget_constraint: Optional[Callable] = None, + state_space_functions: Optional[Dict[str, Callable]] = None, + stochastic_states_transitions: Optional[Dict[str, Callable]] = None, + shock_functions: Optional[Dict[str, Callable]] = None, + continuous_grid_functions: Optional[Dict[str, Callable]] = None, + alternative_sim_specifications: Optional[Dict[str, Callable]] = None, + debug_info: Optional[str] = None, + use_stochastic_sparsity: bool = False, + big_model: Optional[setup_model] = None, +) -> Callable[[Dict[str, jnp.ndarray]], model_solved]: + """Return a solve function that assembles the big solve from sub-model solves. + + The big model is either supplied via ``big_model`` (already built) or built here + from ``model_config`` plus the remaining ``setup_model`` build arguments. + + Args: + small_models: The sub-models to solve. Their discrete state-choice spaces must + tile the big model's state-choice space exactly (partition -- every big + state-choice covered by exactly one sub-model). Build them with the same + model functions as the big model, restricting only the invariant type + state(s) (see the module docstring). + parallel: If False, solve the sub-models one at a time in a Python loop, so + only one block's device arrays are live at once (the memory-saving path). + If True, dispatch the sub-model solves across ``jax.devices()`` (one block + per device, round-robin) and read them back together; this holds several + blocks' arrays at once and only speeds things up with more than one device. + big_model: An already-built big (full) model. If given, the remaining build + arguments are ignored. If omitted, the big model is built from + ``model_config`` and the build arguments below. + model_config: Config of the big (full) model (required when ``big_model`` is + omitted). + model_specs, utility_functions, utility_functions_final_period, + budget_constraint, state_space_functions, stochastic_states_transitions, + shock_functions, continuous_grid_functions, alternative_sim_specifications, + debug_info, use_stochastic_sparsity: The remaining ``setup_model`` build + arguments for the big model; forward the same objects used for the + sub-models. + + Returns: + A function ``solve(params) -> model_solved`` for the big model, whose + ``value``/``policy``/``endog_grid`` are the sub-model solves scattered into the + big solution container. The arrays are host (NumPy) arrays: the full solution + is stored on the host and never materialised on device in one piece. + + """ + if big_model is None: + if model_config is None: + raise ValueError( + "Provide either a pre-built big_model or model_config (plus the " + "remaining setup_model build arguments) to build the big model." + ) + big_model = setup_model( + model_config=model_config, + model_specs=model_specs, + utility_functions=utility_functions, + utility_functions_final_period=utility_functions_final_period, + budget_constraint=budget_constraint, + state_space_functions=state_space_functions, + stochastic_states_transitions=stochastic_states_transitions, + shock_functions=shock_functions, + continuous_grid_functions=continuous_grid_functions, + alternative_sim_specifications=alternative_sim_specifications, + debug_info=debug_info, + use_stochastic_sparsity=use_stochastic_sparsity, + ) + + _check_small_models_consistency(big_model, small_models) + big_rows_per_model = _compute_big_rows_per_model(big_model, small_models) + _check_partition(big_model, big_rows_per_model) + + # Compile each sub-model's solve once; reused across params. + small_solve_funcs = [m.get_solve_func() for m in small_models] + + store_endog = not big_model.model_config["upper_envelope"][ + "skip_endog_grid_storage" + ] + n_state_choices = big_model.model_structure["state_choice_space"].shape[0] + n_continuous = big_model.model_config["continuous_states_info"][ + "n_continuous_state_combinations" + ] + n_wealth = big_model.model_config["n_total_wealth_grid"] + container_shape = (n_state_choices, n_continuous, n_wealth) + + def solve_from_small(params: Dict[str, jnp.ndarray]) -> model_solved: + value = np.full(container_shape, np.nan) + policy = np.full(container_shape, np.nan) + endog_grid = np.full(container_shape, np.nan) if store_endog else None + + if parallel: + solutions = _dispatch_parallel(small_solve_funcs, params) + else: + solutions = None + + for i, rows in enumerate(big_rows_per_model): + sol = solutions[i] if parallel else small_solve_funcs[i](params) + # np.asarray pulls the block off device and scatters it into the host + # container; the block's device arrays are then free to be released. + value[rows] = np.asarray(sol.value) + policy[rows] = np.asarray(sol.policy) + if store_endog: + endog_grid[rows] = np.asarray(sol.endog_grid) + + return model_solved( + model=big_model, + params=params, + value=value, + policy=policy, + endog_grid=endog_grid, + ) + + return solve_from_small + + +def _dispatch_parallel(small_solve_funcs, params): + """Dispatch each sub-model solve on its own device (round-robin) without blocking. + + Returns the (still async) ``model_solved`` results; the caller blocks on them when + it reads the arrays out. With a single device this is effectively sequential. + + """ + devices = jax.devices() + if len(small_solve_funcs) > len(devices): + warnings.warn( + f"parallel=True dispatches {len(small_solve_funcs)} sub-model solves " + f"across {len(devices)} device(s); with fewer devices than sub-models the " + "solves cannot all run concurrently and several blocks are held in memory " + "at once. Use parallel=False to keep only one block live at a time.", + stacklevel=2, + ) + + results = [] + for i, solve_func in enumerate(small_solve_funcs): + with jax.default_device(devices[i % len(devices)]): + results.append(solve_func(params)) + return results + + +def _compute_big_rows_per_model(big_model, small_models): + """For each sub-model, the big state-choice row of each of its state-choices.""" + big_structure = big_model.model_structure + discrete_states_names = big_structure["discrete_states_names"] + + big_rows_per_model = [] + for model in small_models: + state_choice_space_dict = model.model_structure["state_choice_space_dict"] + states = { + name: jnp.asarray(state_choice_space_dict[name]) + for name in discrete_states_names + } + choices = jnp.asarray(state_choice_space_dict["choice"]) + big_rows = get_state_choice_index_per_discrete_states_and_choices( + model_structure=big_structure, + states=states, + choices=choices, + ) + big_rows_per_model.append(np.asarray(big_rows)) + return big_rows_per_model + + +def _check_partition(big_model, big_rows_per_model): + """Check the sub-models tile the big state-choice space exactly.""" + big_structure = big_model.model_structure + n_state_choices = big_structure["state_choice_space"].shape[0] + + indexer_dtype = np.asarray( + big_structure["map_state_choice_to_index_with_proxy"] + ).dtype + invalid_index = np.iinfo(indexer_dtype).max + + all_rows = np.concatenate([np.asarray(rows) for rows in big_rows_per_model]) + + if np.any(all_rows == invalid_index): + raise ValueError( + "Some sub-model state-choices do not exist in the big model (their " + "state-choice lookup returned the invalid index). The sub-models must be " + "restrictions of the big model built with the same state-space functions." + ) + + if all_rows.shape[0] != n_state_choices: + raise ValueError( + "The sub-models' state-choices do not sum to the big model's " + f"state-choices: got {all_rows.shape[0]} across sub-models, but the big " + f"model has {n_state_choices} state-choices." + ) + + if not np.array_equal(np.sort(all_rows), np.arange(n_state_choices)): + raise ValueError( + "The sub-models do not partition the big model's state-choice space " + "exactly: their mapped rows are not a permutation of all big rows (some " + "big state-choices are covered by more than one sub-model or by none)." + ) + + +def _check_small_models_consistency(big_model, small_models): + """Check the sub-models are shape- and structure-compatible with the big model.""" + if len(small_models) == 0: + raise ValueError("small_models must be a non-empty list of models.") + + big_config = big_model.model_config + big_structure = big_model.model_structure + big_continuous = big_config["continuous_states_info"] + + for i, model in enumerate(small_models): + config = model.model_config + structure = model.model_structure + continuous = config["continuous_states_info"] + prefix = f"Sub-model {i}:" + + if structure["discrete_states_names"] != big_structure["discrete_states_names"]: + raise ValueError( + f"{prefix} discrete_states_names differ from the big model: " + f"{structure['discrete_states_names']} vs " + f"{big_structure['discrete_states_names']}." + ) + + if not np.array_equal( + np.asarray(structure["choice_range"]), + np.asarray(big_structure["choice_range"]), + ): + raise ValueError(f"{prefix} choices differ from the big model.") + + if config["n_periods"] != big_config["n_periods"]: + raise ValueError(f"{prefix} n_periods differs from the big model.") + + if config["n_quad_points"] != big_config["n_quad_points"]: + raise ValueError(f"{prefix} n_quad_points differs from the big model.") + + if config["n_total_wealth_grid"] != big_config["n_total_wealth_grid"]: + raise ValueError( + f"{prefix} n_total_wealth_grid differs from the big model " + f"({config['n_total_wealth_grid']} vs " + f"{big_config['n_total_wealth_grid']})." + ) + + if ( + continuous["n_continuous_state_combinations"] + != big_continuous["n_continuous_state_combinations"] + ): + raise ValueError( + f"{prefix} n_continuous_state_combinations differs from the big model." + ) + + # `None` on both sides when the grid is state-choice-specific: there is no + # shared array to compare, and the per-state-choice grids come from the same + # continuous_grid_functions entry, already required to be the same object. + grid = continuous["assets_grid_end_of_period"] + big_grid = big_continuous["assets_grid_end_of_period"] + if (grid is None) != (big_grid is None): + raise ValueError( + f"{prefix} declares assets_end_of_period " + f"{'state-specifically' if grid is None else 'as an array'} while the " + "big model does not." + ) + if grid is not None and not np.allclose(np.asarray(grid), np.asarray(big_grid)): + raise ValueError( + f"{prefix} assets_end_of_period grid differs from the big model." + ) + + if ( + continuous["additional_continuous_state_names"] + != big_continuous["additional_continuous_state_names"] + ): + raise ValueError( + f"{prefix} additional continuous state names differ from the big model." + ) + for name, grid in big_continuous["additional_continuous_state_grids"].items(): + small_grid = continuous["additional_continuous_state_grids"][name] + if (grid is None) != (small_grid is None): + raise ValueError( + f"{prefix} additional continuous grid '{name}' differs from the " + "big model (one is state-specific and the other is not)." + ) + if grid is not None and not np.allclose( + np.asarray(small_grid), np.asarray(grid) + ): + raise ValueError( + f"{prefix} additional continuous grid '{name}' differs from the " + "big model." + ) + + if config["upper_envelope"]["method"] != big_config["upper_envelope"]["method"]: + raise ValueError( + f"{prefix} upper_envelope method differs from the big model." + ) + + if ( + config["upper_envelope"]["skip_endog_grid_storage"] + != big_config["upper_envelope"]["skip_endog_grid_storage"] + ): + raise ValueError( + f"{prefix} skip_endog_grid_storage differs from the big model." + ) + + if ( + structure["stochastic_states_names"] + != big_structure["stochastic_states_names"] + ): + raise ValueError( + f"{prefix} stochastic_states_names differ from the big model." + ) + + if not np.array_equal( + np.asarray(structure["stochastic_state_space"]), + np.asarray(big_structure["stochastic_state_space"]), + ): + raise ValueError( + f"{prefix} stochastic_state_space differs from the big model." + ) diff --git a/src/dcegm/interfaces/inspect_solution.py b/src/dcegm/interfaces/inspect_solution.py index 463d12e4..f323c996 100644 --- a/src/dcegm/interfaces/inspect_solution.py +++ b/src/dcegm/interfaces/inspect_solution.py @@ -1,5 +1,6 @@ import copy +import jax import jax.lax import jax.numpy as jnp import numpy as np @@ -18,7 +19,7 @@ def partially_solve( model_structure, params, n_periods, - return_candidates=False, + return_candidates, ): """Partially solve the model for the last n_periods. @@ -56,9 +57,9 @@ def partially_solve( relevant_state_choices_mask ] - n_continuous_state_combinations = model_structure["continuous_state_space"][ - next(iter(model_structure["continuous_state_space"])) - ].shape[0] + n_continuous_state_combinations = continuous_states_info[ + "n_continuous_state_combinations" + ] ( value_solved, policy_solved, @@ -73,13 +74,14 @@ def partially_solve( if return_candidates: n_assets_end_of_period = model_config["continuous_states_info"][ - "assets_grid_end_of_period" - ].shape[0] + "n_assets_end_of_period" + ] value_candidates, policy_candidates, endog_grid_candidates = ( create_solution_container( n_total_wealth_grid=n_assets_end_of_period, n_state_choices=relevant_state_choice_space.shape[0], n_continuous_state_combinations=n_continuous_state_combinations, + store_endog_grid=True, ) ) @@ -204,6 +206,10 @@ def partially_solve( - rescale_idx ) + law_of_motion_arrays_batch = jax.tree.map( + lambda leaf: leaf[id_batch], segment_info["law_of_motion_arrays"] + ) + xs = ( idx_to_solve, segment_info["child_state_choices_to_aggr_choice"][id_batch, :, :], @@ -212,6 +218,7 @@ def partially_solve( segment_info["child_states_idxs"][id_batch, :], state_choices_batch, state_choices_childs_batch, + law_of_motion_arrays_batch, ) carry = (value_solved, policy_solved, endog_grid_solved) single_period_out_dict = solve_single_period( @@ -219,7 +226,7 @@ def partially_solve( xs=xs, params=params, continuous_grids_info=continuous_states_info, - continuous_state_space=model_structure["continuous_state_space"], + state_choice_space_dict=model_structure["state_choice_space_dict"], income_shocks_scaled=income_shocks_scaled, model_funcs=model_funcs, income_shock_weights=income_shock_weights, diff --git a/src/dcegm/interfaces/interface.py b/src/dcegm/interfaces/interface.py index 97c3bf51..223ad01e 100644 --- a/src/dcegm/interfaces/interface.py +++ b/src/dcegm/interfaces/interface.py @@ -14,7 +14,6 @@ interpolate_policy_for_state_and_choice, interpolate_value_for_state_and_choice, ) -from dcegm.pre_processing.sol_container import broadcast_dj_wealth_grid def get_n_state_choice_period(model_structure): @@ -36,6 +35,24 @@ def get_n_state_choice_period(model_structure): ) +def _endog_grid_placeholder(shape): + """A stand-in for the endogenous grid when it is not stored. + + Under ``skip_endog_grid_storage`` the endogenous grid is never written: every + state-choice's "endogenous" grid is by construction its fixed + Druedahl-Jorgensen wealth grid, which each reader recomputes from that + state-choice's own identity (``_dj_wealth_grid_for_state_choice`` and + ``_interp_policy_and_value_multidim_dj_for_state_choice`` in + ``interpolation/interp_interfaces.py``). + + So nothing downstream reads these values -- only the shape matters, to keep the + ``in_axes=None`` contract of the vmaps below, which must not batch the wealth + grid over the (possibly large) query dimension. + + """ + return jnp.zeros(shape) + + def policy_and_value_for_states_and_choices( states, choices, @@ -93,10 +110,8 @@ def policy_and_value_for_states_and_choices( # DJ-constant: never batch the wealth grid over the (possibly large) query # dimension. Broadcast only across the (small) continuous combo axis and let # vmap treat it as invariant via in_axes=None. - endog_grid_state_choice = broadcast_dj_wealth_grid( - model_config["continuous_states_info"], - (value_grid_state_choice.shape[1],) - + model_config["continuous_states_info"]["dj_wealth_grid"].shape, + endog_grid_state_choice = _endog_grid_placeholder( + value_grid_state_choice.shape[1:] ) endog_grid_in_axes = None else: @@ -170,10 +185,8 @@ def value_for_state_and_choice( fill_value=jnp.nan, ) if model_config["upper_envelope"]["skip_endog_grid_storage"]: - endog_grid_state_choice = broadcast_dj_wealth_grid( - model_config["continuous_states_info"], - (value_grid_state_choice.shape[1],) - + model_config["continuous_states_info"]["dj_wealth_grid"].shape, + endog_grid_state_choice = _endog_grid_placeholder( + value_grid_state_choice.shape[1:] ) endog_grid_in_axes = None else: @@ -250,10 +263,8 @@ def policy_for_state_choice_vec( fill_value=jnp.nan, ) if model_config["upper_envelope"]["skip_endog_grid_storage"]: - endog_grid_state_choice = broadcast_dj_wealth_grid( - model_config["continuous_states_info"], - (value_grid_state_choice.shape[1],) - + model_config["continuous_states_info"]["dj_wealth_grid"].shape, + endog_grid_state_choice = _endog_grid_placeholder( + value_grid_state_choice.shape[1:] ) endog_grid_in_axes = None else: @@ -401,9 +412,7 @@ def choice_values_for_states( if model_config["upper_envelope"]["skip_endog_grid_storage"]: # Broadcast only across the combo axis, never across the (states, choices) # batch axes that the double vmap below maps over. - endog_grid_states = broadcast_dj_wealth_grid( - model_config["continuous_states_info"], value_grid_states.shape[2:] - ) + endog_grid_states = _endog_grid_placeholder(value_grid_states.shape[2:]) endog_grid_in_axes = None else: endog_grid_states = jnp.take( @@ -477,9 +486,7 @@ def choice_policies_for_states( fill_value=jnp.nan, ) if model_config["upper_envelope"]["skip_endog_grid_storage"]: - endog_grid_states = broadcast_dj_wealth_grid( - model_config["continuous_states_info"], value_grid_states.shape[2:] - ) + endog_grid_states = _endog_grid_placeholder(value_grid_states.shape[2:]) endog_grid_in_axes = None else: endog_grid_states = jnp.take( diff --git a/src/dcegm/interfaces/model_class.py b/src/dcegm/interfaces/model_class.py index d253c1c0..ffabef7c 100644 --- a/src/dcegm/interfaces/model_class.py +++ b/src/dcegm/interfaces/model_class.py @@ -1,10 +1,9 @@ import pickle as pkl -from functools import partial -from grp import struct_group -from typing import Callable, Dict +from typing import Any, Callable, Dict, Optional, Tuple, Union import jax import pandas as pd +from jax import numpy as jnp from dcegm.backward_induction import backward_induction from dcegm.interfaces.index_functions import ( @@ -40,26 +39,78 @@ class setup_model: + """User-facing entry point: builds a discrete-continuous life-cycle model. + + Constructing this class runs the full model-build pipeline once (state space, state- + choice space, batching for backward induction) and caches the result on the instance + as ``model_config``, ``model_funcs``, ``model_structure`` and ``batch_info``. Every + other method -- ``solve``, ``solve_and_simulate``, the ``get_*_func`` jit-compiling + factories, the likelihood-function factory, and the debug/inspection helpers -- + reuses that cached build rather than repeating it. + + """ + def __init__( self, - model_config: Dict, - model_specs: Dict, + model_config: Dict[str, Any], + model_specs: Dict[str, Any], utility_functions: Dict[str, Callable], utility_functions_final_period: Dict[str, Callable], budget_constraint: Callable, - state_space_functions: Dict[str, Callable] = None, - stochastic_states_transitions: Dict[str, Callable] = None, - shock_functions: Dict[str, Callable] = None, - alternative_sim_specifications: Dict[str, Callable] = None, - debug_info: str = None, - model_save_path: str = None, - model_load_path: str = None, + state_space_functions: Optional[Dict[str, Callable]] = None, + stochastic_states_transitions: Optional[Dict[str, Callable]] = None, + shock_functions: Optional[Dict[str, Callable]] = None, + continuous_grid_functions: Optional[Dict[str, Callable]] = None, + alternative_sim_specifications: Optional[Dict[str, Callable]] = None, + debug_info: Optional[str] = None, + model_save_path: Optional[str] = None, + model_load_path: Optional[str] = None, use_stochastic_sparsity: bool = False, - ): - """Setup the model and check if load or save is required. + ) -> None: + """Build (or load) the model and cache its structure on the instance. + + Exactly one of ``model_load_path``/``model_save_path`` may be given: if + ``model_load_path`` is set, the model is read back from disk instead of + rebuilt; if ``model_save_path`` is set, it is built as usual and then + also written to disk for a later ``model_load_path`` call. If neither + is given, the model is built fresh and kept only in memory. Args: - use_stochastic_sparsity (bool, optional): EXPERIMENTAL: Use stochastic transition sparsity. + model_config: Model configuration -- periods, choices, the + continuous grids, the upper-envelope method, and related + solver settings. + model_specs: User-supplied numerical model specifications (e.g. + economic parameters that are fixed across estimation, unlike + ``params``). Converted to a jax pytree via ``try_jax_array``. + utility_functions: User-supplied utility, marginal utility, and + inverse marginal utility functions for the regular periods. + utility_functions_final_period: The same, for the final period + (typically a simplified "consume everything" problem). + budget_constraint: User-supplied callable computing beginning-of- + period assets from the previous period's choice and shocks. + state_space_functions: Optional user-supplied functions for (i) + the state-specific feasible choice set and (ii) the + endogenous state update given a choice. + stochastic_states_transitions: Optional user-supplied transition + functions for stochastic (exogenous) discrete states. + shock_functions: Optional user-supplied taste-shock functions. + continuous_grid_functions: Optional user-supplied functions + returning (possibly state-specific) continuous grids, keyed by + grid name. + alternative_sim_specifications: Optional alternative simulation + functions, used to build ``self.alternative_sim_funcs`` via + ``generate_alternative_sim_functions``. If omitted, + ``self.alternative_sim_funcs`` is ``None`` and simulation + falls back to ``model_funcs``. + debug_info: Optional flag controlling how much extra (non-jit- + friendly) debug information the build keeps around, e.g. + ``"all"`` to keep ``map_state_choice_to_child_states`` for + ``get_child_states``. + model_save_path: Optional path to save the built model dict to. + model_load_path: Optional path to load a previously saved model + dict from, skipping the build. + use_stochastic_sparsity: EXPERIMENTAL: use stochastic transition + sparsity. """ if model_load_path is not None: @@ -72,6 +123,7 @@ def __init__( state_space_functions=state_space_functions, stochastic_states_transitions=stochastic_states_transitions, shock_functions=shock_functions, + continuous_grid_functions=continuous_grid_functions, path=model_load_path, use_stochastic_sparsity=use_stochastic_sparsity, ) @@ -85,6 +137,7 @@ def __init__( state_space_functions=state_space_functions, stochastic_states_transitions=stochastic_states_transitions, shock_functions=shock_functions, + continuous_grid_functions=continuous_grid_functions, path=model_save_path, debug_info=debug_info, use_stochastic_sparsity=use_stochastic_sparsity, @@ -99,6 +152,7 @@ def __init__( state_space_functions=state_space_functions, stochastic_states_transitions=stochastic_states_transitions, shock_functions=shock_functions, + continuous_grid_functions=continuous_grid_functions, debug_info=debug_info, use_stochastic_sparsity=use_stochastic_sparsity, ) @@ -129,26 +183,31 @@ def __init__( else: self.alternative_sim_funcs = None - def solve(self, params, load_sol_path=None, save_sol_path=None): - """Solve a discrete-continuous life-cycle model using the DC-EGM algorithm. + def solve( + self, + params: Dict[str, jnp.ndarray], + load_sol_path: Optional[str] = None, + save_sol_path: Optional[str] = None, + ) -> model_solved: + """Solve the model via backward induction for one set of parameters. + + Not jit-compiled -- each call retraces ``backward_induction`` from + scratch. For repeated solves at different parameters (e.g. inside an + optimizer or MCMC sampler), use ``get_solve_func`` instead, which + compiles once and reuses the compiled function. Args: - params (pd.DataFrame): Params DataFrame. - options (dict): Options dictionary. - utility_functions (Dict[str, callable]): Dictionary of three user-supplied - functions for computation of: - (i) utility - (ii) inverse marginal utility - (iii) next period marginal utility - budget_constraint (callable): Callable budget constraint. - state_space_functions (Dict[str, callable]): Dictionary of two user-supplied - functions to: - (i) get the state specific feasible choice set - (ii) update the endogenous part of the state by the choice - final_period_solution (callable): User-supplied function for solving the agent's - last period. - transition_function (callable): User-supplied function returning for each - state a transition matrix vector. + params: Model parameters. + load_sol_path: Optional path to a pickled solution dict + (``{"value", "policy", "endog_grid"}``); if given, the model is + not solved and this solution is loaded instead. + save_sol_path: Optional path to pickle the freshly computed + solution dict to. Ignored when ``load_sol_path`` is given. + + Returns: + The solved model, wrapping ``value``, ``policy`` and + ``endog_grid`` together with this model's structure for + downstream inspection/simulation. """ params_processed = process_params( @@ -185,24 +244,32 @@ def solve(self, params, load_sol_path=None, save_sol_path=None): def solve_and_simulate( self, - params, - states_initial, - seed, - load_sol_path=None, - save_sol_path=None, - ): - """Solve the model and simulate it. + params: Dict[str, jnp.ndarray], + states_initial: Dict[str, jnp.ndarray], + seed: int, + load_sol_path: Optional[str] = None, + save_sol_path: Optional[str] = None, + ) -> pd.DataFrame: + """Solve the model, then simulate a panel of agents forward through it. Args: - params: The parameters for the model. - states_initial: The initial states for the simulation. - wealth_initial: The initial wealth for the simulation. - n_periods: The number of periods to simulate. - seed: The random seed for the simulation. - alt_model_funcs_sim: Alternative model functions for simulation. + params: Model parameters, in any form ``process_params`` accepts. + states_initial: Initial discrete (and, if applicable, continuous) + states for the simulated agents, one array per state name, each of shape + ``(n_agents,)``. Must include ``assets_end_of_previous_period``: + the assets agents carry *into* the first period, not the wealth + they have to spend in it. The first period applies the budget + equation to it like every other period. + seed: Random seed for the simulation's taste shocks and income + draws. + load_sol_path: Optional path to a pickled solution dict to load + instead of solving. + save_sol_path: Optional path to pickle the freshly computed + solution dict to. Ignored when ``load_sol_path`` is given. Returns: - A dictionary containing the solution and simulation results. + A long-format panel with a ``(period, agent)`` MultiIndex, + containing simulated states, choices, consumption, and value. """ params_processed = process_params(params, self.params_check_info) @@ -245,8 +312,21 @@ def solve_and_simulate( sim_df = create_simulation_df(sim_dict) return sim_df - def get_solve_func(self): - """Create a fast function for solving that is jit compiled in the first call.""" + def get_solve_func(self) -> Callable: + """Build a solve function that is jit-compiled once, on its first call. + + The (large, array-valued) parts of ``model_structure``/``batch_info`` + are split off via ``split_structure_and_batch_info`` and threaded + through ``jax.jit`` as traced arguments; the remaining, non-jit- + friendly parts are closed over as static Python objects. Subsequent + calls at different ``params`` reuse the same compiled XLA executable. + + Returns: + A function ``solve_function(params) -> model_solved`` with the + same solving semantics as ``solve`` (without the load/save + options), but compiled. + + """ ( model_structure_for_jit, batch_info_for_jit, @@ -281,7 +361,7 @@ def solve_function_to_jit(params, model_structure_jit, batch_info_jit): solve_func = jax.jit(solve_function_to_jit) # Generate the function. The user only needs to provide params, but we call with the objects for jit. - def solve_function(params): + def solve_function(params: Dict[str, jnp.ndarray]) -> model_solved: """Solve the model for given params.""" value, policy, endog_grid = solve_func( params, model_structure_for_jit, batch_info_for_jit @@ -299,12 +379,36 @@ def solve_function(params): def get_solve_and_simulate_func( self, - states_initial, - seed, - slow_version=False, - ): - """Create a fast function for solving and simulation that is jit compiled in the - first call.""" + states_initial: Dict[str, jnp.ndarray], + seed: int, + slow_version: bool = False, + ) -> Callable[[Dict[str, jnp.ndarray]], pd.DataFrame]: + """Build a solve-and-simulate function, jit-compiled on its first call. + + Same jit strategy as ``get_solve_func`` (splitting ``model_structure``/ + ``batch_info`` into traced vs. static parts), but the returned + function also simulates a panel from the solution before returning, + in a single compiled call. + + Args: + states_initial: Initial states for the simulated agents, one array per state name, each of shape + ``(n_agents,)``. Must include ``assets_end_of_previous_period``: + the assets agents carry *into* the first period, not the wealth + they have to spend in it. The first period applies the budget + equation to it like every other period. + Fixed for every call of the returned function (only ``params`` + varies). + seed: Random seed for the simulation; likewise fixed. + slow_version: If True, skip ``jax.jit`` on the combined + solve-and-simulate step (useful for debugging with Python- + level tracebacks); the solve step itself is still whatever + ``backward_induction`` does internally. + + Returns: + A function ``solve_and_simulate_function(params) -> pd.DataFrame`` + with the same output as ``solve_and_simulate``. + + """ # Fix everything except params, solution of the model and model_structure which contains large arrays. sim_func = lambda params, value, policy, endog_gid, model_structure: simulate_all_periods( states_initial=states_initial, @@ -367,7 +471,7 @@ def solve_and_simulate_function_to_jit( solve_simulate_func = jax.jit(solve_and_simulate_function_to_jit) # Generate the function. The user only needs to provide params, but we call with the objects for jit. - def solve_and_simulate_function(params): + def solve_and_simulate_function(params: Dict[str, jnp.ndarray]) -> pd.DataFrame: sim_dict = solve_simulate_func( params, model_structure_for_jit, batch_info_for_jit ) @@ -378,15 +482,45 @@ def solve_and_simulate_function(params): def create_experimental_ll_func( self, - params_all, - observed_states, - observed_choices, - unobserved_state_specs=None, - return_model_solution=False, - use_probability_of_observed_states=True, - slow_version=False, - ): + params_all: Dict[str, float], + observed_states: Dict[str, Any], + observed_choices: jnp.ndarray, + unobserved_state_specs: Optional[Dict[str, Any]] = None, + return_model_solution: bool = False, + use_probability_of_observed_states: bool = True, + slow_version: bool = False, + ) -> Callable[[Dict[str, float]], jnp.ndarray]: + """Build a per-individual negative log-likelihood function for estimation. + + EXPERIMENTAL. Thin passthrough to + ``create_individual_likelihood_function``; see that function's + implementation for the exact choice-probability construction. + Args: + params_all: The full parameter dict; the function returned below + is called with only the subset being estimated, and this dict + supplies the rest. + observed_states: Observed discrete (and continuous, if any) states + from the data, one array per state name. + observed_choices: Observed choices from the data. + unobserved_state_specs: Optional specification of states that are + unobserved in the data and must be integrated/weighted over + when computing choice probabilities. + return_model_solution: If True, the returned function also + returns the ``(value, policy, endog_grid)`` solution dict + alongside the likelihood contributions. + use_probability_of_observed_states: Whether to weight by the + model-implied probability of the observed (partially latent) + states. + slow_version: If True, skip ``jax.jit`` on the likelihood function + (useful for debugging). + + Returns: + A function mapping the estimated-parameter subset to per- + individual negative log-likelihood contributions (and, if + ``return_model_solution``, also the solution dict). + + """ return create_individual_likelihood_function( income_shock_draws_unscaled=self.income_shock_draws_unscaled, income_shock_weights=self.income_shock_weights, @@ -404,8 +538,21 @@ def create_experimental_ll_func( slow_version=slow_version, ) - def validate_exogenous(self, params): + def validate_exogenous(self, params: Dict[str, jnp.ndarray]) -> bool: + """Check that every stochastic transition is a valid probability vector. + + Verifies, for each exogenous (stochastic) process and every state it + is evaluated at, that the transition probabilities are non-negative, + sum to one, and have the expected dimensionality. + + Args: + params: Model parameters, in any form ``process_params`` accepts. + + Returns: + True if every exogenous process validates; otherwise a + ``ValueError`` is raised by ``validate_stochastic_transition``. + """ return validate_stochastic_transition( params=params, model_structure=self.model_structure, @@ -413,15 +560,39 @@ def validate_exogenous(self, params): model_funcs=self.model_funcs, ) - def get_state_choices_idx(self, states): - """Get the indices of the state choices for given states.""" + def get_state_choices_idx(self, states: Dict[str, Any]) -> jnp.ndarray: + """Look up state-choice indices for given discrete states. + + Args: + states: Discrete state values, one array (or scalar) per state + name. + + Returns: + The index of each state's block in the state-choice space (i.e. + the row at which its first choice starts). + + """ return get_state_choice_index_per_discrete_states( states=states, map_state_choice_to_index=self.model_structure["map_state_choice_to_index"], discrete_states_names=self.model_structure["discrete_states_names"], ) - def get_child_states(self, state, choice): + def get_child_states(self, state: Dict[str, Any], choice: Any) -> pd.DataFrame: + """Look up the (deterministic) child states for a state-choice. + + Requires the model to have been built with ``debug_info="all"``, since + ``map_state_choice_to_child_states`` is only kept around in that mode. + + Args: + state: A discrete state (or a batch of them). + choice: The choice taken in ``state``. + + Returns: + A DataFrame with one row per child state, one column per discrete + state name. + + """ if "map_state_choice_to_child_states" not in self.model_structure: raise ValueError( "For this function the model needs to be created with debug_info='all'" @@ -437,9 +608,21 @@ def get_child_states(self, state, choice): } return pd.DataFrame(child_states) - def get_child_states_and_calc_trans_probs(self, state, choice, params): - """Get the child states for a given state and choice and calculate the - transition probabilities.""" + def get_child_states_and_calc_trans_probs( + self, state: Dict[str, Any], choice: Any, params: Dict[str, jnp.ndarray] + ) -> pd.DataFrame: + """``get_child_states`` plus each child's stochastic transition probability. + + Args: + state: A discrete state (or a batch of them). + choice: The choice taken in ``state``. + params: Model parameters, in any form ``process_params`` accepts. + + Returns: + The ``get_child_states`` DataFrame with an added ``trans_probs`` + column. + + """ child_states_df = self.get_child_states(state, choice) trans_probs = self.model_funcs["compute_stochastic_transition_vec"]( @@ -449,10 +632,34 @@ def get_child_states_and_calc_trans_probs(self, state, choice, params): return child_states_df def get_full_child_states_by_asset_id_and_probs( - self, state, choice, params, asset_id, second_continuous_id=None - ): - """Get the child states for a given state and choice and calculate the - transition probabilities.""" + self, + state: Dict[str, Any], + choice: Any, + params: Dict[str, jnp.ndarray], + asset_id: int, + second_continuous_id: Optional[int] = None, + ) -> pd.DataFrame: + """``get_child_states_and_calc_trans_probs`` plus next-period wealth by + quadrature point. + + Requires the model to have been built with ``debug_info="all"``. + + Args: + state: A discrete state (or a batch of them). + choice: The choice taken in ``state``. + params: Model parameters, in any form ``process_params`` accepts. + asset_id: Index into the end-of-period asset grid to evaluate the + law of motion at. + second_continuous_id: Index into the additional continuous-state + grid, required if and only if the model has one. + + Returns: + The ``get_child_states_and_calc_trans_probs`` DataFrame, with an + added ``assets_begin_of_period_quad_point_{i}`` column per income + quadrature point (and, with an additional continuous state, an + added column for its next-period value). + + """ if "map_state_choice_to_child_states" not in self.model_structure: raise ValueError( "For this function the model needs to be created with debug_info='all'" @@ -511,7 +718,22 @@ def get_full_child_states_by_asset_id_and_probs( child_states_df["trans_probs"] = trans_probs return child_states_df - def compute_law_of_motions(self, params): + def compute_law_of_motions(self, params: Dict[str, jnp.ndarray]) -> Dict[str, Any]: + """Compute the law of motion for the whole state(-choice) space at once. + + Debug/inspection entry point: unlike the main solve path (which + computes this on demand, per batch and period), this evaluates + ``calc_cont_grids_next_period`` once for every state-choice, which is + only tractable for small-to-medium models. + + Args: + params: Model parameters, in any form ``process_params`` accepts. + + Returns: + A dict with ``"assets_begin_of_period"`` and ``"continuous_states"`` + arrays, indexed by the full state(-choice) space. + + """ return calc_cont_grids_next_period( params=params, model_structure=self.model_structure, @@ -520,11 +742,36 @@ def compute_law_of_motions(self, params): income_shock_draws_unscaled=self.income_shock_draws_unscaled, ) - def get_n_state_choices_per_period(self): + def get_n_state_choices_per_period(self) -> pd.Series: + """Count state-choices per period. + + Returns: + A pandas Series indexed by period, with the number of + state-choices in each. + + """ return get_n_state_choice_period(self.model_structure) - def solve_partially(self, params, n_periods, return_candidates=False): + def solve_partially( + self, + params: Dict[str, jnp.ndarray], + n_periods: int, + return_candidates: bool = False, + ) -> Dict[str, Any]: + """Solve only the last ``n_periods`` periods, for debugging large models. + + Args: + params: Model parameters, in any form ``process_params`` accepts. + n_periods: Number of periods (counted from the end) to solve. + return_candidates: If True, also return the pre-upper-envelope + candidate solutions. + Returns: + A dict with the solved ``value``/``policy``/``endog_grid`` (and, + if ``return_candidates``, the candidate solutions) for the solved + periods. + + """ return partially_solve( income_shock_draws_unscaled=self.income_shock_draws_unscaled, income_shock_weights=self.income_shock_weights, @@ -538,8 +785,20 @@ def solve_partially(self, params, n_periods, return_candidates=False): ) def set_alternative_sim_funcs( - self, alternative_sim_specifications, alternative_specs=None - ): + self, + alternative_sim_specifications: Dict[str, Callable], + alternative_specs: Optional[Dict[str, Any]] = None, + ) -> None: + """(Re-)build ``self.alternative_sim_funcs`` for use in simulation. + + Args: + alternative_sim_specifications: Alternative simulation functions, + passed through to ``generate_alternative_sim_functions``. + alternative_specs: Optional alternative model specs to build them + against; defaults to this model's own ``model_specs`` if + omitted. + + """ if alternative_specs is None: self.alternative_sim_specs = self.model_specs alternative_specs_without_jax = self.specs_without_jax diff --git a/src/dcegm/interfaces/sol_interface.py b/src/dcegm/interfaces/sol_interface.py index 66d9415e..e6875b91 100644 --- a/src/dcegm/interfaces/sol_interface.py +++ b/src/dcegm/interfaces/sol_interface.py @@ -12,6 +12,7 @@ value_for_state_and_choice, ) from dcegm.interfaces.interface_checks import check_states_and_choices +from dcegm.law_of_motion import compute_own_dj_wealth_grid from dcegm.likelihood import ( calc_choice_probs_for_states, get_state_choice_index_per_discrete_states, @@ -20,7 +21,6 @@ generate_alternative_sim_functions, ) from dcegm.pre_processing.shared import try_jax_array -from dcegm.pre_processing.sol_container import broadcast_dj_wealth_grid from dcegm.simulation.sim_utils import create_simulation_df from dcegm.simulation.simulate import simulate_all_periods @@ -72,7 +72,20 @@ def set_alternative_sim_funcs( self.alternative_sim_funcs = alternative_sim_funcs def simulate(self, states_initial, seed): + """Simulate a panel of agents forward through the solved model. + Args: + states_initial: Initial states for the simulated agents, one array per state name, each of shape + ``(n_agents,)``. Must include ``assets_end_of_previous_period``: + the assets agents carry *into* the first period, not the wealth + they have to spend in it. The first period applies the budget + equation to it like every other period. + seed: Random seed for the simulation's taste shocks and income draws. + + Returns: + A long-format panel with a ``(period, agent)`` MultiIndex. + + """ sim_dict = simulate_all_periods( states_initial=states_initial, n_periods=self.model_config["n_periods"], @@ -196,8 +209,19 @@ def get_solution_for_discrete_state_choice(self, states, choices): fill_value=jnp.nan, ) if self.model_config["upper_envelope"]["skip_endog_grid_storage"]: - endog_grid = broadcast_dj_wealth_grid( - self.model_config["continuous_states_info"], value_grid.shape + # Each state-choice's own Druedahl-Jorgensen wealth grid, evaluated on + # demand -- self-referential, this is exactly the state-choice whose + # own stored solution is being read. + own_dj_wealth_grid = jax.vmap( + compute_own_dj_wealth_grid, in_axes=(0, None) + )(state_choices, self.model_funcs["continuous_grid_functions"]) + endog_grid = jnp.broadcast_to( + own_dj_wealth_grid[:, None, :], + ( + own_dj_wealth_grid.shape[0], + value_grid.shape[1], + own_dj_wealth_grid.shape[1], + ), ) else: endog_grid = jnp.take( @@ -238,6 +262,7 @@ def choice_probabilities_for_states(self, states): states=states, model_config=self.model_config, model_funcs=self.model_funcs, + model_structure=self.model_structure, ) def choice_values_for_states(self, states): diff --git a/src/dcegm/interpolation/interp2d_irregular.py b/src/dcegm/interpolation/interp2d_irregular.py index 66e580ed..fd532c3c 100644 --- a/src/dcegm/interpolation/interp2d_irregular.py +++ b/src/dcegm/interpolation/interp2d_irregular.py @@ -110,7 +110,7 @@ def interp2d_value_on_wealth_and_regular_grid( state_choice_vec: Dict[str, int], params: dict, discount_factor, - cont_state_name: str = "continuous_state", + cont_state_name: str, ): """Interpolate the value function on a 2D grid. diff --git a/src/dcegm/interpolation/interp_interfaces.py b/src/dcegm/interpolation/interp_interfaces.py index a9249b79..0e784840 100644 --- a/src/dcegm/interpolation/interp_interfaces.py +++ b/src/dcegm/interpolation/interp_interfaces.py @@ -15,6 +15,29 @@ from dcegm.interpolation.interpnd_regular import ( interpnd_policy_and_value_for_child_states_on_regular_grids, ) +from dcegm.law_of_motion import ( + compute_own_continuous_grid_combos, + compute_own_continuous_grids_raw, + compute_own_dj_wealth_grid, +) + + +def _dj_wealth_grid_for_state_choice( + state_choice_vec, endog_grid_state_choice, model_config, continuous_grid_functions +): + """This state-choice's own Druedahl-Jorgensen wealth grid (simple 1d case). + + Self-referential -- this is the exact state-choice whose own stored solution is + being read, so it's its own representative, same reasoning as + compute_own_continuous_grid_combos elsewhere in this file. Only used when + skip_endog_grid_storage is True (endog_grid isn't stored at all in that case, + since every state-choice's "endogenous" grid is by construction this fixed + array); otherwise the real stored grid is used unchanged. + + """ + if model_config["upper_envelope"]["skip_endog_grid_storage"]: + return compute_own_dj_wealth_grid(state_choice_vec, continuous_grid_functions) + return endog_grid_state_choice[0] def interpolate_value_for_state_and_choice( @@ -35,7 +58,7 @@ def interpolate_value_for_state_and_choice( multidim = continuous_states_info["has_additional_continuous_state"] - continuous_state_space = model_structure["continuous_state_space"] + continuous_grid_functions = model_funcs["continuous_grid_functions"] if (upper_envelope_method == "fues") & multidim: continuous_state_name = continuous_states_info[ @@ -44,7 +67,11 @@ def interpolate_value_for_state_and_choice( continuous_state = state_choice_vec[continuous_state_name] _, value = interp2d_policy_and_value_on_wealth_and_regular_grid( - continuous_state_space=continuous_state_space, + continuous_state_space={ + continuous_state_name: continuous_grid_functions[continuous_state_name]( + **state_choice_vec + ) + }, wealth_grid=endog_grid_state_choice, policy_grid=endog_grid_state_choice, value_grid=value_grid_state_choice, @@ -62,21 +89,26 @@ def interpolate_value_for_state_and_choice( value_grid_state_choice=value_grid_state_choice, endog_grid_state_choice=endog_grid_state_choice, state_choice_vec=state_choice_vec, - additional_continuous_state_grids=continuous_states_info[ - "additional_continuous_state_grids" - ], - continuous_state_space=continuous_state_space, + continuous_grid_functions=continuous_grid_functions, continuous_state_names=continuous_states_info[ "additional_continuous_state_names" ], compute_utility=compute_utility, params=params, discount_factor=discount_factor, + skip_endog_grid_storage=model_config["upper_envelope"][ + "skip_endog_grid_storage" + ], ) elif upper_envelope_method == "druedahl_jorgensen": value = interp1d_value_on_wealth_dj( wealth=state_choice_vec["assets_begin_of_period"], - wealth_grid=endog_grid_state_choice[0], + wealth_grid=_dj_wealth_grid_for_state_choice( + state_choice_vec, + endog_grid_state_choice, + model_config, + continuous_grid_functions, + ), value_grid=value_grid_state_choice[0], compute_utility=compute_utility, state_choice_vec=state_choice_vec, @@ -110,7 +142,7 @@ def interpolate_policy_for_state_and_choice( continuous_states_info = model_config["continuous_states_info"] upper_envelope_method = model_config["upper_envelope"]["method"] multidim = continuous_states_info["has_additional_continuous_state"] - continuous_state_space = model_structure["continuous_state_space"] + continuous_grid_functions = model_funcs["continuous_grid_functions"] if (upper_envelope_method == "fues") & multidim: continuous_state_name = continuous_states_info[ @@ -118,7 +150,11 @@ def interpolate_policy_for_state_and_choice( ][0] continuous_state = state_choice_vec[continuous_state_name] policy, _ = interp2d_policy_and_value_on_wealth_and_regular_grid( - continuous_state_space=continuous_state_space, + continuous_state_space={ + continuous_state_name: continuous_grid_functions[continuous_state_name]( + **state_choice_vec + ) + }, wealth_grid=endog_grid_state_choice, policy_grid=policy_grid_state_choice, value_grid=policy_grid_state_choice, @@ -143,7 +179,12 @@ def interpolate_policy_for_state_and_choice( elif upper_envelope_method == "druedahl_jorgensen": policy, _ = interp1d_policy_and_value_on_wealth_dj( wealth=state_choice_vec["assets_begin_of_period"], - wealth_grid=endog_grid_state_choice[0], + wealth_grid=_dj_wealth_grid_for_state_choice( + state_choice_vec, + endog_grid_state_choice, + model_config, + continuous_grid_functions, + ), policy_grid=policy_grid_state_choice[0], value_grid=value_grid_state_choice[0], compute_utility=model_funcs["compute_utility"], @@ -176,7 +217,7 @@ def interpolate_policy_and_value_for_state_and_choice( compute_utility = model_funcs["compute_utility"] discount_factor = model_funcs["read_funcs"]["discount_factor"](params) - continuous_state_space = model_structure["continuous_state_space"] + continuous_grid_functions = model_funcs["continuous_grid_functions"] multidim = continuous_states_info["has_additional_continuous_state"] @@ -186,7 +227,11 @@ def interpolate_policy_and_value_for_state_and_choice( ][0] continuous_state = state_choice_vec[continuous_state_name] policy, value = interp2d_policy_and_value_on_wealth_and_regular_grid( - continuous_state_space=continuous_state_space, + continuous_state_space={ + continuous_state_name: continuous_grid_functions[continuous_state_name]( + **state_choice_vec + ) + }, wealth_grid=endog_grid_state_choice, policy_grid=policy_grid_state_choice, value_grid=value_grid_state_choice, @@ -203,21 +248,26 @@ def interpolate_policy_and_value_for_state_and_choice( value_grid_state_choice=value_grid_state_choice, endog_grid_state_choice=endog_grid_state_choice, state_choice_vec=state_choice_vec, - additional_continuous_state_grids=continuous_states_info[ - "additional_continuous_state_grids" - ], - continuous_state_space=continuous_state_space, + continuous_grid_functions=continuous_grid_functions, continuous_state_names=continuous_states_info[ "additional_continuous_state_names" ], compute_utility=compute_utility, params=params, discount_factor=discount_factor, + skip_endog_grid_storage=model_config["upper_envelope"][ + "skip_endog_grid_storage" + ], ) elif upper_envelope_method == "druedahl_jorgensen": policy, value = interp1d_policy_and_value_on_wealth_dj( wealth=state_choice_vec["assets_begin_of_period"], - wealth_grid=endog_grid_state_choice[0], + wealth_grid=_dj_wealth_grid_for_state_choice( + state_choice_vec, + endog_grid_state_choice, + model_config, + continuous_grid_functions, + ), policy_grid=policy_grid_state_choice[0], value_grid=value_grid_state_choice[0], compute_utility=compute_utility, @@ -245,13 +295,37 @@ def _interp_policy_and_value_multidim_dj_for_state_choice( value_grid_state_choice, endog_grid_state_choice, state_choice_vec, - additional_continuous_state_grids, - continuous_state_space, + continuous_grid_functions, continuous_state_names, compute_utility, params, discount_factor, + skip_endog_grid_storage, ): + # Under skip_endog_grid_storage the endogenous grid is not stored at all, so + # endog_grid_state_choice carries no usable values (callers pass a placeholder + # of the right shape); this state-choice's own Druedahl-Jorgensen wealth grid is + # computed here instead, exactly as _dj_wealth_grid_for_state_choice does for + # the simple 1d case. It is the same grid for every continuous-state combo of a + # given state-choice -- that is what makes it "the fixed common wealth grid" -- + # so the one array serves both lookups below. + if skip_endog_grid_storage: + own_wealth_grid = compute_own_dj_wealth_grid( + state_choice_vec, continuous_grid_functions + ) + wealth_grid_first = own_wealth_grid + else: + wealth_grid_first = endog_grid_state_choice[0] + + # Self-referential: this state-choice's own solution is being queried, so its + # own grid (evaluated on itself) is the one it was solved on. + own_continuous_state_grids = compute_own_continuous_grids_raw( + state_choice_vec, continuous_grid_functions, continuous_state_names + ) + continuous_state_space = compute_own_continuous_grid_combos( + state_choice_vec, continuous_grid_functions, continuous_state_names + ) + continuous_state_child_states = { name: jnp.asarray(state_choice_vec[name])[None, None] for name in continuous_state_names @@ -262,8 +336,8 @@ def _interp_policy_and_value_multidim_dj_for_state_choice( if key not in {"assets_begin_of_period", *continuous_state_names} } policy_nd, value_nd = interpnd_policy_and_value_for_child_states_on_regular_grids( - additional_continuous_state_grids=additional_continuous_state_grids, - wealth_grid=endog_grid_state_choice[0], + additional_continuous_state_grids=own_continuous_state_grids, + wealth_grid=wealth_grid_first, policy_grid_child_states=policy_grid_state_choice[None, ...], value_grid_child_states=value_grid_state_choice[None, ...], continuous_state_child_states=continuous_state_child_states, @@ -289,7 +363,11 @@ def _interp_policy_and_value_multidim_dj_for_state_choice( policy_exact, value_exact = interp1d_policy_and_value_on_wealth_dj( wealth=state_choice_vec["assets_begin_of_period"], - wealth_grid=endog_grid_state_choice[combo_idx], + wealth_grid=( + own_wealth_grid + if skip_endog_grid_storage + else endog_grid_state_choice[combo_idx] + ), policy_grid=policy_grid_state_choice[combo_idx], value_grid=value_grid_state_choice[combo_idx], compute_utility=compute_utility, diff --git a/src/dcegm/interpolation/interpnd_regular.py b/src/dcegm/interpolation/interpnd_regular.py index f870a650..cfb2afeb 100644 --- a/src/dcegm/interpolation/interpnd_regular.py +++ b/src/dcegm/interpolation/interpnd_regular.py @@ -164,6 +164,135 @@ def _interp_one_child_state( return policy_final, value_final +def interpnd_policy_and_value_for_child_states_on_own_regular_grids( + continuous_grid_functions: Dict[str, Callable], + additional_continuous_state_names, + wealth_grid_func: Callable, + policy_grid_child_states: jnp.ndarray, + value_grid_child_states: jnp.ndarray, + continuous_state_child_states: Dict[str, jnp.ndarray], + wealth_child_states: jnp.ndarray, + state_choice_child_states: Dict[str, Any], + compute_utility: Callable, + params: Dict[str, Any], + discount_factor: float, +) -> tuple[jnp.ndarray, jnp.ndarray]: + """Interpolate policy/value on each child's own regular grid, with consume-all. + + Per-child-grid counterpart of + ``interpnd_policy_and_value_for_child_states_on_regular_grids`` above: each + child is interpolated against its *own* grid per continuous-state name and for + wealth, not one grid shared across all children, reusing every downstream + helper (corner weighting, wealth interpolation, consume-all overwrite) + unchanged -- only the index/weight precomputation differs (see + ``_precompute_interp_objects_own_grids``). Kept as a separate function rather + than changing the original's contract, since that's directly tested + (``tests/test_interpnd_regular.py``) with the shared-grid contract. + + Every child's own grid (for each additional continuous state, and for wealth + via ``wealth_grid_func``) is evaluated inside the same per-child vmap that + computes its interpolation indices/weights, not precomputed for the whole + batch upfront in a separate vmap and fed in as a paired array -- same pattern + as ``interp2d_value_and_marg_util_for_state_choice`` in + ``interpolate_marginal_utility.py``. + + """ + objs = _precompute_interp_objects_own_grids( + continuous_grid_functions=continuous_grid_functions, + additional_continuous_state_names=additional_continuous_state_names, + state_choice_child_states=state_choice_child_states, + continuous_state_child_states=continuous_state_child_states, + wealth_grid_func=wealth_grid_func, + wealth_child_states=wealth_child_states, + ) + + def _interp_one_child_state( + policy_grid_one_child, + value_grid_one_child, + regular_low_idxs_one_child, + regular_high_idxs_one_child, + regular_low_weights_one_child, + regular_high_weights_one_child, + wealth_points_one_child, + wealth_low_idxs_one_child, + wealth_high_idxs_one_child, + wealth_grid_one_child, + ): + return vmap( + _interp_policy_and_value_one_comb, + in_axes=( + None, + None, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + None, + None, + None, + ), + )( + policy_grid_one_child, + value_grid_one_child, + regular_low_idxs_one_child, + regular_high_idxs_one_child, + regular_low_weights_one_child, + regular_high_weights_one_child, + wealth_points_one_child, + wealth_low_idxs_one_child, + wealth_high_idxs_one_child, + objs["strides"], + objs["corner_table"], + wealth_grid_one_child, + ) + + policy_interp, value_interp = vmap( + _interp_one_child_state, + in_axes=(0, 0, 1, 1, 1, 1, 0, 0, 0, 0), + )( + policy_grid_child_states, + value_grid_child_states, + objs["regular_low_idxs"], + objs["regular_high_idxs"], + objs["regular_low_weights"], + objs["regular_high_weights"], + wealth_child_states, + objs["wealth_low_idxs"], + objs["wealth_high_idxs"], + objs["wealth_grid_per_child"], + ) + + # We need to interpolate the expected value at zero savings, because we only know it for the regular + # grid corners + expected_value_zero_savings = _interp_regular_only_all( + values_over_regular_grid_child_states=value_grid_child_states[..., 0], + regular_low_idxs=objs["regular_low_idxs"], + regular_high_idxs=objs["regular_high_idxs"], + regular_low_weights=objs["regular_low_weights"], + regular_high_weights=objs["regular_high_weights"], + strides=objs["strides"], + corner_table=objs["corner_table"], + ) + + consume_all_value = _compute_consume_all_value( + expected_value_zero_savings=expected_value_zero_savings, + wealth_child_states=wealth_child_states, + state_choice_child_states=state_choice_child_states, + continuous_state_child_states=continuous_state_child_states, + compute_utility=compute_utility, + params=params, + discount_factor=discount_factor, + ) + + overwrite_mask = consume_all_value > value_interp + policy_final = jnp.where(overwrite_mask, wealth_child_states, policy_interp) + value_final = jnp.where(overwrite_mask, consume_all_value, value_interp) + return policy_final, value_final + + def interpnd_value_for_child_states_on_regular_grids( additional_continuous_state_grids: Dict[str, jnp.ndarray], wealth_grid: jnp.ndarray, @@ -569,6 +698,126 @@ def _precompute_regular_indices_and_weights( ) +def _precompute_interp_objects_own_grids( + continuous_grid_functions: Dict[str, Callable], + additional_continuous_state_names, + state_choice_child_states: Dict[str, Any], + continuous_state_child_states: Dict[str, jnp.ndarray], + wealth_grid_func: Callable, + wealth_child_states: jnp.ndarray, +) -> Dict[str, jnp.ndarray]: + """Per-child-grid counterpart of ``_precompute_interp_objects`` above. + + Each child's own grid (for every additional continuous state, and for wealth) is + evaluated inside the same per-child ``vmap`` that computes its interpolation + indices/weights below -- not precomputed for the whole batch upfront in its own + separate ``vmap`` and fed in as a paired array. + + ``strides``/``corner_table`` only depend on the *number* of grid points per + dimension, which is the same for every child by construction (state-specific grids + may vary in value, not size -- enforced by + ``evaluate_state_specific_continuous_grids`` in ``continuous_state_grids.py``), so + that length is read once, from a single representative child's own grid (row 0), the + same pinning approach ``continuous_state_grids.py`` uses at model-structure-build + time. + + """ + state_names = list(additional_continuous_state_names) + representative_state = { + key: value[0] for key, value in state_choice_child_states.items() + } + regular_shape = [ + int(continuous_grid_functions[name](**representative_state).shape[0]) + for name in state_names + ] + strides = jnp.asarray(_compute_row_major_strides(regular_shape), dtype=jnp.int32) + regular_low_idxs, regular_high_idxs, regular_low_weights, regular_high_weights = ( + _precompute_regular_indices_and_weights_own_grids( + continuous_grid_functions=continuous_grid_functions, + state_choice_child_states=state_choice_child_states, + continuous_state_child_states=continuous_state_child_states, + state_names=state_names, + ) + ) + + def _wealth_grid_and_idx_for_one_child(state_dict, wealth_points_one_child): + wealth_grid_one_child = wealth_grid_func(state_dict) + wealth_high_idx, wealth_low_idx = get_index_high_and_low( + wealth_grid_one_child, wealth_points_one_child + ) + return wealth_grid_one_child, wealth_high_idx, wealth_low_idx + + wealth_grid_per_child, wealth_high_idxs, wealth_low_idxs = vmap( + _wealth_grid_and_idx_for_one_child, in_axes=(0, 0) + )(state_choice_child_states, wealth_child_states) + + corner_table = _corner_table(len(state_names)) + return { + "strides": strides, + "regular_low_idxs": regular_low_idxs, + "regular_high_idxs": regular_high_idxs, + "regular_low_weights": regular_low_weights, + "regular_high_weights": regular_high_weights, + "wealth_low_idxs": wealth_low_idxs, + "wealth_high_idxs": wealth_high_idxs, + "wealth_grid_per_child": wealth_grid_per_child, + "corner_table": corner_table, + } + + +def _precompute_regular_indices_and_weights_own_grids( + continuous_grid_functions: Dict[str, Callable], + state_choice_child_states: Dict[str, Any], + continuous_state_child_states: Dict[str, jnp.ndarray], + state_names, +): + """Per-child-grid counterpart of ``_precompute_regular_indices_and_weights`` above. + + Each child is interpolated against its *own* grid per continuous-state name, + evaluated here (one ``vmap`` per name, over the child axis) rather than read + from a precomputed ``(n_child_state_choices, n_points)`` array -- the grid + function and the index/weight computation that consumes it happen in the same + per-child call. + + Returns arrays of shape + ``(n_dims, n_child_state_choices, n_continuous_combinations)``, same as the + shared-grid version. + + """ + low_idxs = [] + high_idxs = [] + low_weights = [] + high_weights = [] + + for name in state_names: + points = continuous_state_child_states[name] + + def _grid_and_idx_for_one_child(state_dict, points_one_child, name=name): + grid_one_child = continuous_grid_functions[name](**state_dict) + high_idx, low_idx = get_index_high_and_low(grid_one_child, points_one_child) + return grid_one_child, high_idx, low_idx + + grid_per_child, high_idx, low_idx = vmap( + _grid_and_idx_for_one_child, in_axes=(0, 0) + )(state_choice_child_states, points) + x_low = jnp.take_along_axis(grid_per_child, low_idx, axis=1) + x_high = jnp.take_along_axis(grid_per_child, high_idx, axis=1) + high_w = (points - x_low) / (x_high - x_low) + high_w = jnp.where(jnp.isfinite(high_w), high_w, 0.0) + low_w = 1.0 - high_w + low_idxs.append(low_idx) + high_idxs.append(high_idx) + low_weights.append(low_w) + high_weights.append(high_w) + + return ( + jnp.stack(low_idxs, axis=0), + jnp.stack(high_idxs, axis=0), + jnp.stack(low_weights, axis=0), + jnp.stack(high_weights, axis=0), + ) + + def _take_last_axis(arr, idx): """Take along last axis with batched indices.""" if arr.ndim != idx.ndim + 1: diff --git a/src/dcegm/interpolation/simulation_interp.py b/src/dcegm/interpolation/simulation_interp.py index 2e0a9f85..d62dae85 100644 --- a/src/dcegm/interpolation/simulation_interp.py +++ b/src/dcegm/interpolation/simulation_interp.py @@ -10,6 +10,11 @@ from dcegm.interpolation.interpnd_regular import ( interpnd_policy_and_value_for_child_states_on_regular_grids, ) +from dcegm.law_of_motion import ( + compute_own_continuous_grid_combos, + compute_own_continuous_grids_raw, + compute_own_dj_wealth_grid, +) def interpolate_policy_and_value_for_all_agents( @@ -24,13 +29,11 @@ def interpolate_policy_and_value_for_all_agents( params, discrete_states_names, compute_utility, - continuous_state_space, - additional_continuous_state_grids, + continuous_grid_functions, upper_envelope_method, has_additional_continuous_state, discount_factor, - skip_endog_grid_storage=False, - dj_wealth_grid=None, + skip_endog_grid_storage, ): # 1D interpolation path is independent of upper-envelope method and only @@ -57,10 +60,13 @@ def interpolate_policy_and_value_for_all_agents( fill_value=jnp.nan, )[:, :, 0, :] if skip_endog_grid_storage: - # DJ-constant: never batch the wealth grid over agents/choices. Already - # matches the shape a single agent-choice item needs, so no broadcast at - # all is needed here. - endog_grid_agent = dj_wealth_grid + # DJ-constant: endog_grid isn't stored at all in this case (every + # state-choice's "endogenous" grid is by construction the fixed + # assets_begin_of_period grid), so this is just a placeholder -- the + # real (possibly state-choice-specific) wealth grid is computed on + # demand per agent-choice inside interp1d_policy_and_value_function + # instead, self-referential via that row's own identity. + endog_grid_agent = jnp.zeros(1) endog_grid_in_axes = None else: endog_grid_agent = jnp.take( @@ -76,7 +82,7 @@ def interpolate_policy_and_value_for_all_agents( vmap( interp1d_policy_and_value_function, in_axes=( - None, + 0, None, endog_grid_in_axes, 0, @@ -86,9 +92,24 @@ def interpolate_policy_and_value_for_all_agents( None, None, None, + None, + None, ), ), - in_axes=(0, 0, endog_grid_in_axes, 0, 0, None, None, None, None, None), + in_axes=( + 0, + 0, + endog_grid_in_axes, + 0, + 0, + None, + None, + None, + None, + None, + None, + None, + ), ) policy_agent, value_agent = vectorized_interp( @@ -102,6 +123,8 @@ def interpolate_policy_and_value_for_all_agents( compute_utility, discount_factor, upper_envelope_method == "druedahl_jorgensen", + skip_endog_grid_storage, + continuous_grid_functions, ) return policy_agent, value_agent @@ -142,7 +165,8 @@ def interpolate_policy_and_value_for_all_agents( vmap( interp2d_policy_and_value_function, in_axes=( - None, + 0, + # Wealth may differ by choice; the continuous state may not. None, None, 0, @@ -182,7 +206,7 @@ def interpolate_policy_and_value_for_all_agents( value_grid_agent, policy_grid_agent, choice_range, - continuous_state_space, + continuous_grid_functions, continuous_state_name, params, compute_utility, @@ -213,11 +237,11 @@ def interpolate_policy_and_value_for_all_agents( fill_value=jnp.nan, ) if skip_endog_grid_storage: - # DJ-constant: broadcast only across the combo axis, never across - # agents/choices. - endog_grid_agent = jnp.broadcast_to( - dj_wealth_grid, value_grid_agent.shape[2:] - ) + # endog_grid isn't stored in this case; interpnd_policy_and_value_function + # recomputes each state-choice's own Druedahl-Jorgensen wealth grid on + # demand. Only the shape matters here, to keep in_axes=None so the wealth + # grid is never batched over agents/choices. + endog_grid_agent = jnp.zeros(value_grid_agent.shape[2:]) endog_grid_in_axes = None else: endog_grid_agent = jnp.take( @@ -229,13 +253,16 @@ def interpolate_policy_and_value_for_all_agents( ) endog_grid_in_axes = 0 - additional_continuous_state_names = list(continuous_state_space.keys()) + additional_continuous_state_names = list( + continuous_state_beginning_of_period.keys() + ) vectorized_interp = vmap( vmap( interpnd_policy_and_value_function, in_axes=( - None, + 0, + # Wealth may differ by choice; the continuous state may not. None, None, endog_grid_in_axes, @@ -275,12 +302,12 @@ def interpolate_policy_and_value_for_all_agents( value_grid_agent, policy_grid_agent, choice_range, - continuous_state_space, - additional_continuous_state_grids, + continuous_grid_functions, additional_continuous_state_names, params, compute_utility, discount_factor, + skip_endog_grid_storage, ) return policy_agent, value_agent @@ -301,13 +328,23 @@ def interp1d_policy_and_value_function( compute_utility, discount_factor, use_dj_interpolation, + skip_endog_grid_storage, + continuous_grid_functions, ): state_choice_vec = {**state, "choice": choice} if use_dj_interpolation: + # This agent's own Druedahl-Jorgensen wealth grid, evaluated on demand when + # endog_grid isn't stored at all -- self-referential, this is exactly the + # state-choice whose own stored solution is being read. + wealth_grid = ( + compute_own_dj_wealth_grid(state_choice_vec, continuous_grid_functions) + if skip_endog_grid_storage + else endog_grid_agent + ) policy_interp, value_interp = interp1d_policy_and_value_on_wealth_dj( wealth=wealth_beginning_of_period, - wealth_grid=endog_grid_agent, + wealth_grid=wealth_grid, policy_grid=policy_agent, value_grid=value_agent, compute_utility=compute_utility, @@ -338,7 +375,7 @@ def interp2d_policy_and_value_function( value_agent, policy_agent, choice, - continuous_state_space, + continuous_grid_functions, continuous_state_name, params, compute_utility, @@ -346,8 +383,16 @@ def interp2d_policy_and_value_function( ): state_choice_vec = {**state, "choice": choice} + # Self-referential: the stored solution for this state-choice was solved on + # its own grid, so evaluating the grid function on itself reproduces it. + own_continuous_state_space = { + continuous_state_name: continuous_grid_functions[continuous_state_name]( + **state_choice_vec + ) + } + policy_interp, value_interp = interp2d_policy_and_value_on_wealth_and_regular_grid( - continuous_state_space=continuous_state_space, + continuous_state_space=own_continuous_state_space, wealth_grid=endog_grid_agent, policy_grid=policy_agent, value_grid=value_agent, @@ -372,12 +417,12 @@ def interpnd_policy_and_value_function( value_agent, policy_agent, choice, - continuous_state_space, - additional_continuous_state_grids, + continuous_grid_functions, additional_continuous_state_names, params, compute_utility, discount_factor, + skip_endog_grid_storage, ): state_choice_vec = {**state, "choice": choice} # The 1d interpolation below evaluates the utility function directly, so it needs @@ -390,6 +435,29 @@ def interpnd_policy_and_value_function( }, } + # Under skip_endog_grid_storage the endogenous grid is not stored, so + # endog_grid_agent carries no usable values (the caller passes a placeholder of + # the right shape); this state-choice's own Druedahl-Jorgensen wealth grid is + # computed here instead, same as interp1d_policy_and_value_function does for the + # simple 1d case. One array serves both lookups below: the fixed common wealth + # grid is by definition the same for every continuous-state combo. + if skip_endog_grid_storage: + own_wealth_grid = compute_own_dj_wealth_grid( + state_choice_vec, continuous_grid_functions + ) + wealth_grid_first = own_wealth_grid + else: + wealth_grid_first = endog_grid_agent[0] + + # Self-referential: the stored solution for this state-choice was solved on + # its own grid, so evaluating the grid function on itself reproduces it. + own_continuous_state_grids = compute_own_continuous_grids_raw( + state_choice_vec, continuous_grid_functions, additional_continuous_state_names + ) + own_continuous_state_space = compute_own_continuous_grid_combos( + state_choice_vec, continuous_grid_functions, additional_continuous_state_names + ) + continuous_state_child_states = { name: continuous_state_beginning_of_period[name][None, None] for name in additional_continuous_state_names @@ -400,8 +468,8 @@ def interpnd_policy_and_value_function( policy_interp, value_interp = ( interpnd_policy_and_value_for_child_states_on_regular_grids( - additional_continuous_state_grids=additional_continuous_state_grids, - wealth_grid=endog_grid_agent[0], + additional_continuous_state_grids=own_continuous_state_grids, + wealth_grid=wealth_grid_first, policy_grid_child_states=policy_agent[None, ...], value_grid_child_states=value_agent[None, ...], continuous_state_child_states=continuous_state_child_states, @@ -413,17 +481,22 @@ def interpnd_policy_and_value_function( ) ) - exact_mask = jnp.ones_like(next(iter(continuous_state_space.values())), dtype=bool) + exact_mask = jnp.ones_like( + next(iter(own_continuous_state_space.values())), dtype=bool + ) for name in additional_continuous_state_names: exact_mask = exact_mask & jnp.isclose( - continuous_state_space[name], continuous_state_beginning_of_period[name] + own_continuous_state_space[name], + continuous_state_beginning_of_period[name], ) has_exact_combo = jnp.any(exact_mask) combo_idx = jnp.argmax(exact_mask) policy_exact, value_exact = interp1d_policy_and_value_on_wealth_dj( wealth=wealth_beginning_of_period, - wealth_grid=endog_grid_agent[combo_idx], + wealth_grid=( + own_wealth_grid if skip_endog_grid_storage else endog_grid_agent[combo_idx] + ), policy_grid=policy_agent[combo_idx], value_grid=value_agent[combo_idx], compute_utility=compute_utility, diff --git a/src/dcegm/law_of_motion.py b/src/dcegm/law_of_motion.py index 50a76a3c..0e5943fe 100644 --- a/src/dcegm/law_of_motion.py +++ b/src/dcegm/law_of_motion.py @@ -1,3 +1,5 @@ +from typing import Any, Dict + import jax.numpy as jnp from jax import vmap @@ -6,71 +8,267 @@ ) +def calc_law_of_motion( + law_of_motion_arrays: Dict[str, jnp.ndarray], + state_choice_space_dict: Dict[str, jnp.ndarray], + income_shocks_scaled: jnp.ndarray, + params: Dict[str, float], + model_funcs: Dict[str, Any], + has_additional_continuous_states: bool, + additional_continuous_state_names: list, +) -> Dict[str, Any]: + """Compute the law of motion at whichever granularity is valid for this model. + + Single entry point for every caller (``solve_single_period.py`` via + ``interpolate_value_and_marg_util``, and ``final_periods.py``), so the + granularity decision lives here rather than being repeated at each call site. + Callers pass the raw ingredients -- index arrays plus + ``state_choice_space_dict`` -- rather than pre-gathered representative-parent + dicts, so only the branch actually taken pays for the gather. + + The budget equation is the only law of motion that may declare ``choice`` (the + child's own; every other state is fixed before the choice is made, see + ``_check_choice_not_declared`` in ``process_model_functions.py``). Unless it + does -- decided once at model-build time -- every state-choice sharing a child + state would compute a bit-identical transition, and the whole transition is + evaluated once per unique child *state* and gathered out instead. + + The two branches are alternatives, not nested: each calls the shared core + ``_calc_transitions_for_rows`` with different rows -- state-choices in one + case, deduplicated child states in the other -- so there is exactly one + implementation of the transition math. + + Args: + law_of_motion_arrays: The child arrays for the branch actually taken. + Assembled once at model setup (see ``bundle_law_of_motion_arrays`` in + ``pre_processing/batches/batch_creation.py``) and threaded through the + backward induction, so only the branch's arrays are carried. When the + budget depends on ``choice``: ``child_state_choices`` (the child's + own, *non-proxy* state-choice dict -- the transition into it uses its + real state, not the proxy value-reuse slot) and + ``rep_parent_state_choice_idx_per_child_state_choice``. Otherwise: + ``unique_child_states``, ``rep_parent_state_choice_idx_per_child_state`` + and ``state_row_for_state_choice``. Each representative-parent index + points into ``state_choice_space_dict``. + state_choice_space_dict: Full state-choice space; the representative- + parent index arrays are gathered out of this. + income_shocks_scaled: Quadrature points for the income shock, already + scaled by its mean and standard deviation. + params: Model parameters. + model_funcs: Processed model functions; in particular + ``budget_depends_on_choice`` selects the branch. + has_additional_continuous_states: Whether the model has an additional + continuous state besides wealth. + additional_continuous_state_names: Names of those additional + continuous states, if any. + + Returns: + dict with ``"assets_begin_of_period"`` (shape ``(n_children, + n_continuous_combinations, n_exog_savings, n_income_shocks)``) and + ``"continuous_states"`` (one array of that shape per additional + continuous-state name, or a size-1 ``"dummy_cont"`` placeholder when + there is none). + + """ + if model_funcs["budget_depends_on_choice"]: + # A representative parent's own state-choice, for each of this batch's + # deduplicated children -- used only to pick which state-choice's own + # continuous grid feeds the law of motion (see law_of_motion.py). Grids live on + # the state-choice space (that's where the solution itself lives), so this is a + # state-choice index, not a bare state. Any one parent works: + # check_continuous_grid_consistency_across_shared_children (run once at + # model-build time) guarantees every parent sharing a child agrees on its own + # grid. + rep_parent_state_choice_per_state_choice = { + key: var[ + law_of_motion_arrays[ + "rep_parent_state_choice_idx_per_child_state_choice" + ] + ] + for key, var in state_choice_space_dict.items() + } + return calc_law_of_motion_for_state_choices( + child_state_choices=law_of_motion_arrays["child_state_choices"], + representative_parent_state_choice_vec=rep_parent_state_choice_per_state_choice, + income_shocks_scaled=income_shocks_scaled, + params=params, + model_funcs=model_funcs, + has_additional_continuous_states=has_additional_continuous_states, + additional_continuous_state_names=additional_continuous_state_names, + ) + + else: + # Same, at the coarser unique-child-*state* granularity, for the law-of-motion + # fast path taken when the user's transition functions don't depend on "choice" + # (see interpolate_value_and_marg_util / law_of_motion.py). + rep_parent_state_choices_per_child_state = { + key: var[ + law_of_motion_arrays["rep_parent_state_choice_idx_per_child_state"] + ] + for key, var in state_choice_space_dict.items() + } + return calc_law_of_motion_for_child_states( + child_states=law_of_motion_arrays["unique_child_states"], + representative_parent_state_choices=( + rep_parent_state_choices_per_child_state + ), + state_row_for_state_choice=law_of_motion_arrays[ + "state_row_for_state_choice" + ], + income_shocks_scaled=income_shocks_scaled, + params=params, + model_funcs=model_funcs, + has_additional_continuous_states=has_additional_continuous_states, + additional_continuous_state_names=additional_continuous_state_names, + ) + + def calc_law_of_motion_for_state_choices( - state_choice_vec, - continuous_state_space, - assets_grid_end_of_period, + child_state_choices, income_shocks_scaled, params, model_funcs, has_additional_continuous_states, + additional_continuous_state_names, + representative_parent_state_choice_vec, ): - """Compute continuous-state and wealth transitions for a set of state-choices. + """Transitions for a set of child *state-choices*, one row each. + + Used when the budget equation declares ``choice``, so each of a state's choices + genuinely needs its own wealth. See ``calc_law_of_motion`` for how this is chosen, + and ``_calc_transitions_for_rows`` for the shared math. - ``state_choice_vec`` may or may not contain a ``"choice"`` key. It is dropped (via a - no-op-if-absent pop) before being passed to the user-supplied law-of-motion - functions, since the transition does not depend on it -- this is what lets - ``calc_cont_grids_next_period`` below reuse this function unchanged with the full - (choice-less) state space. + """ + return _calc_transitions_for_rows( + rows=child_state_choices, + grid_source_rows=representative_parent_state_choice_vec, + income_shocks_scaled=income_shocks_scaled, + params=params, + model_funcs=model_funcs, + has_additional_continuous_states=has_additional_continuous_states, + additional_continuous_state_names=additional_continuous_state_names, + ) + + +def _calc_transitions_for_rows( + rows, + grid_source_rows, + income_shocks_scaled, + params, + model_funcs, + has_additional_continuous_states, + additional_continuous_state_names, +): + """Continuous-state and wealth transitions for whatever rows are given. + + The single implementation of the transition math, shared by both public + entry points above. ``rows`` are child state-choices in one case and + deduplicated child states in the other; nothing here depends on which, because + a row is only ever used as "the state we are computing beginning-of-period + values for". + + Two *different* dicts flow in here, and keeping them apart is the whole point + of this function's signature: + + ``rows`` + Whose *beginning-of-period* continuous state/wealth we are computing -- + the child, in the main solve path (see ``solve_single_period.py``). May or + may not carry a ``"choice"`` key; it is passed through either way (see the + comment below). + + ``grid_source_rows`` + Whose own continuous grid supplies the *values* fed through the + law-of-motion function -- in the main solve path, a representative parent. + + These are not the same row in general once continuous grids are + state-choice-specific: the transition function itself correctly depends on the + child's own identity (e.g. its ``lagged_choice``, which is the parent's + choice), but the grid values fed into it must come from the parent's own grid. + This applies to ``assets_end_of_period`` just as much as to the additional + continuous states: it too is a law-of-motion input (the exogenous grid a + parent's transition is evaluated over, feeding the child's own budget + equation), so it is evaluated from the representative parent here too. + + Any one parent works as the representative, because + ``check_continuous_grid_consistency_across_shared_children`` (run once at + model-build time) guarantees every parent sharing a child agrees on its own + grid. + + Callers with no real parent/child relationship to trace -- the whole-state-space + debug entry point ``calc_cont_grids_next_period`` below -- pass ``rows`` + itself as the grid source, so each state is its own. That + is a degenerate but well-defined use of the same contract, and matches the + global-grid behavior exactly whenever grids are not state-choice-specific. """ - state_vec = dict(state_choice_vec) - state_vec.pop("choice", None) + # "choice" (the child's own) is passed straight through when present, so the + # budget equation may declare it and get a different wealth per choice -- e.g. a + # choice-specific cost deducted from beginning-of-period wealth. Functions that + # don't declare it are unaffected: determine_function_arguments_and_partial_model_specs + # filters kwargs down to each function's own signature. Callers passing a bare + # state space (no "choice" key at all) are likewise fine -- that is exactly the + # branch calc_law_of_motion routes to when the budget is choice-free. + state_vec = dict(rows) continuous_state_next_period = _get_continuous_state_next_period( has_additional_continuous_states=has_additional_continuous_states, - state_space_dict=state_vec, - continuous_state_space=continuous_state_space, + child_state_choices=rows, + representative_last_period_parent_states=grid_source_rows, + additional_continuous_state_names=additional_continuous_state_names, params=params, model_funcs=model_funcs, ) - def fix_assets_and_shocks_for_broadcast( - states, - continuous_state_vec, - asset_end_of_previous_period, - income_draw, + def _transitions_for_one_state( + states, continuous_state_vec, representative_parent_row ): - all_states = {**states, **continuous_state_vec} - assets_begin_of_period = calc_beginning_of_period_assets_for_single_state( - state_vec=all_states, - asset_end_of_previous_period=asset_end_of_previous_period, - income_shock_draw=income_draw, - params=params, - compute_assets_begin_of_period=model_funcs[ - "compute_assets_begin_of_period" - ], - aux_outs=False, + # Own assets_end_of_period grid, evaluated here for this one state's + # representative parent, right next to where it's consumed below -- + # instead of precomputing the whole batch's grids upfront in a separate + # vmap and feeding the result in as a paired array. + own_assets_grid_end_of_period = _own_assets_grid_end_of_period_for_one_state( + representative_parent_row, model_funcs["continuous_grid_functions"] ) - return assets_begin_of_period - assets_begin_of_next_period = vmap( - vmap( + def fix_assets_and_shocks_for_broadcast( + continuous_state_vec, + asset_end_of_previous_period, + income_draw, + ): + all_states = {**states, **continuous_state_vec} + return calc_beginning_of_period_assets_for_single_state( + state_vec=all_states, + asset_end_of_previous_period=asset_end_of_previous_period, + income_shock_draw=income_draw, + params=params, + compute_assets_begin_of_period=model_funcs[ + "compute_assets_begin_of_period" + ], + aux_outs=False, + ) + + return vmap( vmap( vmap( fix_assets_and_shocks_for_broadcast, - in_axes=(None, None, None, 0), + in_axes=(None, None, 0), ), - in_axes=(None, None, 0, None), + in_axes=(None, 0, None), ), - in_axes=(None, 0, None, None), - ), - in_axes=(0, 0, None, None), + in_axes=(0, None, None), + )( + continuous_state_vec, + own_assets_grid_end_of_period, + income_shocks_scaled, + ) + + assets_begin_of_next_period = vmap( + _transitions_for_one_state, + in_axes=(0, 0, 0), )( state_vec, continuous_state_next_period, - assets_grid_end_of_period, - income_shocks_scaled, + grid_source_rows, ) # Generate result dict @@ -80,6 +278,66 @@ def fix_assets_and_shocks_for_broadcast( } +def _own_assets_grid_end_of_period_for_one_state( + representative_parent_row, continuous_grid_functions +): + return continuous_grid_functions["assets_end_of_period"]( + **representative_parent_row + ) + + +def calc_law_of_motion_for_child_states( + child_states, + representative_parent_state_choices, + state_row_for_state_choice, + income_shocks_scaled, + params, + model_funcs, + has_additional_continuous_states, + additional_continuous_state_names, +): + """Law of motion once per unique child *state*, gathered out to state-choices. + + Sibling of ``calc_law_of_motion_for_state_choices`` above: both call the same + core (``_calc_transitions_for_rows``), this one with deduplicated child states + rather than child state-choices, and then gather the per-state result back out. + Only valid when the budget equation does not depend on ``choice`` (decided once + at model-build time, see ``budget_depends_on_choice`` in + ``process_model_functions.py``): the transition into a child is then a function + of the child's own state, not of the choice it goes on to make, so every + state-choice sharing a child state would otherwise recompute a bit-identical + result ``n_choices`` times. + + ``state_row_for_state_choice`` (built in ``child_state_dedup.py``) maps each + child state-choice back to its row in ``child_states``, so the per-state result + is expanded to the per-state-choice shape every downstream consumer expects -- + the same gather pattern ``calculate_candidate_solutions_from_euler_equation`` + already uses one stage later. + + """ + law_of_motion_per_state = _calc_transitions_for_rows( + rows=child_states, + grid_source_rows=representative_parent_state_choices, + income_shocks_scaled=income_shocks_scaled, + params=params, + model_funcs=model_funcs, + has_additional_continuous_states=has_additional_continuous_states, + additional_continuous_state_names=additional_continuous_state_names, + ) + + return { + "continuous_states": { + name: jnp.take(grid, state_row_for_state_choice, axis=0) + for name, grid in law_of_motion_per_state["continuous_states"].items() + }, + "assets_begin_of_period": jnp.take( + law_of_motion_per_state["assets_begin_of_period"], + state_row_for_state_choice, + axis=0, + ), + } + + def calc_cont_grids_next_period( params, income_shock_draws_unscaled, @@ -106,58 +364,172 @@ def calc_cont_grids_next_period( ) return calc_law_of_motion_for_state_choices( - state_choice_vec=state_space_dict, - continuous_state_space=model_structure["continuous_state_space"], - assets_grid_end_of_period=continuous_states_info["assets_grid_end_of_period"], + child_state_choices=state_space_dict, income_shocks_scaled=income_shocks_scaled, params=params, model_funcs=model_funcs, has_additional_continuous_states=continuous_states_info[ "has_additional_continuous_state" ], + additional_continuous_state_names=continuous_states_info[ + "additional_continuous_state_names" + ], + # No real parent/child relationship to trace here (whole-state-space + # debug entry point), so use each state's own identity. + representative_parent_state_choice_vec=state_space_dict, ) def _get_continuous_state_next_period( has_additional_continuous_states, - state_space_dict, - continuous_state_space, + child_state_choices, + representative_last_period_parent_states, + additional_continuous_state_names, params, model_funcs, ): if not has_additional_continuous_states: # Use an explicit zero-valued dummy continuous state with stable shape # (n_states, 1) to keep downstream shapes constant. - n_states = next(iter(state_space_dict.values())).shape[0] - dummy_name = "dummy_cont" - dummy_dtype = continuous_state_space[dummy_name].dtype + n_states = next(iter(child_state_choices.values())).shape[0] dummy_states = { - dummy_name: jnp.zeros((n_states, 1), dtype=dummy_dtype), + "dummy_cont": jnp.zeros((n_states, 1)), } return dummy_states - continuous_state_next_period = calculate_continuous_state( - discrete_states_beginning_of_period=state_space_dict, - continuous_states_end_of_last_period=continuous_state_space, - params=params, - compute_continuous_state=model_funcs["next_period_continuous_state"], + continuous_state_next_period = vmap( + _continuous_state_next_period_for_one_state, + in_axes=(0, 0, None, None, None, None), + )( + child_state_choices, + representative_last_period_parent_states, + model_funcs["continuous_grid_functions"], + additional_continuous_state_names, + params, + model_funcs["next_period_continuous_state"], ) _check_continuous_state_output_keys( continuous_state_output=continuous_state_next_period, - continuous_state_space=continuous_state_space, + expected_names=additional_continuous_state_names, ) return continuous_state_next_period +def _continuous_state_next_period_for_one_state( + state_dict, + representative_last_period_parent_state, + continuous_grid_functions, + additional_continuous_state_names, + params, + compute_continuous_state, +): + """Compute one state's beginning-of-period continuous state, across its own grid. + + Builds the grid on demand *after* vmapping down to a single state, instead of + precomputing the whole batch's grids upfront in a separate vmap and feeding the + result in as a paired array. Uses ``representative_parent_state_choice_vec`` (the + representative parent state-choice), not ``state_dict`` (the child) -- see the + docstring of ``calc_law_of_motion_for_state_choices`` above for why these differ; + for state-choices without a state-specific grid, ``continuous_grid_functions[name]`` + ignores its input and returns the same global grid for every row, so this is a no-op + relative to the old behavior in that case. + + """ + own_continuous_state_vec = compute_own_continuous_grid_combos( + representative_last_period_parent_state, + continuous_grid_functions, + additional_continuous_state_names, + ) + return vmap( + calc_continuous_state_for_each_grid_point, + in_axes=(None, 0, None, None), + )( + state_dict, + own_continuous_state_vec, + params, + compute_continuous_state, + ) + + +def compute_own_continuous_grids_raw( + state_dict, + continuous_grid_functions, + additional_continuous_state_names, +): + """Evaluate one state-choice's own grid for each additional continuous state. + + Unmeshed: each name's 1d grid on its own, not combined into combo points. + Companion to ``compute_own_continuous_grid_combos`` below (which meshes these), + used directly by n-D regular-grid interpolation + (``interpolation/interpnd_regular.py``), which needs each dimension's own grid + separately -- it combines them into combo points itself, via strides/corner + tables, rather than a dense meshgrid. Intended to be called via ``vmap`` over a + batch of state-choices -- one call here is one state-choice's own grids, not + the whole state-choice space. + + """ + return { + name: continuous_grid_functions[name](**state_dict) + for name in additional_continuous_state_names + } + + +def compute_own_continuous_grid_combos( + state_dict, + continuous_grid_functions, + additional_continuous_state_names, +): + """Evaluate one state's own grid for each additional continuous state and mesh them. + + Mirrors the meshgrid+ravel construction in + ``pre_processing/model_structure/model_structure.py``, but evaluates each name's + grid on demand via its (possibly state-specific) entry in + ``continuous_grid_functions`` instead of reading one precomputed global array. + Intended to be called via ``vmap`` over a batch of states -- one call here is one + state's own grid, not the whole state space. + + """ + own_grids_dict = compute_own_continuous_grids_raw( + state_dict, continuous_grid_functions, additional_continuous_state_names + ) + own_grids = [own_grids_dict[name] for name in additional_continuous_state_names] + meshed = jnp.meshgrid(*own_grids, indexing="ij") + return { + name: grid.ravel() + for name, grid in zip(additional_continuous_state_names, meshed) + } + + +def compute_own_dj_wealth_grid(state_dict, continuous_grid_functions): + """Evaluate one state-choice's own Druedahl-Jorgensen wealth grid on demand. + + The "m_grid" (``assets_begin_of_period``) with a zero-wealth point prepended -- the + grid every state-choice's policy/value is stored/interpolated against under + Druedahl-Jorgensen. Evaluated here on demand, per state-choice (so it may be state- + choice-specific), rather than read from any single stored array; this is the only + place that grid is materialised. Intended to be called via ``vmap`` over a batch of + state-choices -- one call here is one state-choice's own grid, not the whole state- + choice space. + + """ + return jnp.concatenate( + ( + jnp.zeros(1), + continuous_grid_functions["assets_begin_of_period"](**state_dict), + ) + ) + + def _check_continuous_state_output_keys( continuous_state_output, - continuous_state_space, + expected_names, ): - expected_keys = set(continuous_state_space.keys()) + expected_keys = set(expected_names) output_keys = set(continuous_state_output.keys()) if output_keys != expected_keys: raise ValueError( - "next_period_continuous_state output keys must match continuous_state_space keys. " + "next_period_continuous_state output keys must match the additional " + "continuous state names. " f"Expected {sorted(expected_keys)}, got {sorted(output_keys)}." ) @@ -173,7 +545,7 @@ def calc_beginning_of_period_assets_for_single_state( out_budget = compute_assets_begin_of_period( **state_vec, asset_end_of_previous_period=asset_end_of_previous_period, - income_shock_previous_period=income_shock_draw, + income_shock=income_shock_draw, params=params, ) checked_out = check_budget_equation_and_return_wealth_plus_optional_aux( @@ -217,6 +589,19 @@ def calculate_continuous_state( params, compute_continuous_state, ): + """Apply the law of motion for every state against one shared grid (outer product). + + ``continuous_states_end_of_last_period`` is a single grid (shape ``(n_combos,)`` + per continuous-state name), broadcast against every state in + ``discrete_states_beginning_of_period`` -- i.e. every state is evaluated against + every grid point. This is the original, still-current contract, used directly + by ``calc_cont_grids_next_period`` (the whole-state-space debug entry point) and + exercised directly by ``tests/test_law_of_motion.py``. The main solve path + instead pairs each state with its own (possibly state-specific) grid, computed + on demand inside ``_continuous_state_next_period_for_one_state`` above, rather + than broadcasting one grid shared across all of them. + + """ continuous_state_beginning_of_period = vmap( vmap( calc_continuous_state_for_each_grid_point, diff --git a/src/dcegm/likelihood.py b/src/dcegm/likelihood.py index 74a01c97..f47f310c 100644 --- a/src/dcegm/likelihood.py +++ b/src/dcegm/likelihood.py @@ -371,6 +371,7 @@ def partial_choice_prob_func( state_choice_indexes=data_from_observed[2], model_config=model_config, model_funcs=model_funcs, + model_structure=model_structure, ) return partial_choice_prob_func, data_from_observed_wrapped @@ -385,6 +386,7 @@ def calc_choice_prob_for_state_choices( state_choice_indexes, model_config, model_funcs, + model_structure, ): """This function interpolates the policy and value function for all agents. @@ -400,6 +402,7 @@ def calc_choice_prob_for_state_choices( states=states, model_config=model_config, model_funcs=model_funcs, + model_structure=model_structure, ) choice_probs = jnp.take_along_axis( choice_prob_across_choices, choices[:, None], axis=1 @@ -415,6 +418,7 @@ def calc_choice_probs_for_states( states, model_config, model_funcs, + model_structure, ): choice_values_per_state = choice_values_for_states( value_solved=value_solved, @@ -424,6 +428,7 @@ def calc_choice_probs_for_states( states=states, model_config=model_config, model_funcs=model_funcs, + model_structure=model_structure, ) if model_funcs["taste_shock_function"]["taste_shock_scale_is_scalar"]: diff --git a/src/dcegm/pre_processing/alternative_sim_functions.py b/src/dcegm/pre_processing/alternative_sim_functions.py index b2999344..4edbd05a 100644 --- a/src/dcegm/pre_processing/alternative_sim_functions.py +++ b/src/dcegm/pre_processing/alternative_sim_functions.py @@ -3,15 +3,14 @@ from dcegm.pre_processing.check_model_config import check_model_config_and_process from dcegm.pre_processing.model_functions.process_model_functions import ( + _check_budget_does_not_declare_renamed_income_shock, + _declares_choice, process_second_continuous_update_function, process_state_space_functions, ) from dcegm.pre_processing.model_functions.taste_shock_function import ( process_shock_functions, ) -from dcegm.pre_processing.model_functions.upper_evelope_wrapper import ( - create_upper_envelope_function, -) from dcegm.pre_processing.model_structure.stochastic_states import ( create_stochastic_state_mapping, create_stochastic_transition_function, @@ -74,7 +73,7 @@ def generate_alternative_sim_functions( stochastic_state_names, ) - print("Model setup complete.\n") + print("Alternative simulation functions ready.") return model_funcs @@ -85,7 +84,7 @@ def process_alternative_sim_functions( stochastic_states_transition, state_space_functions: Dict[str, Callable], budget_constraint: Callable, - shock_functions: Dict[str, Callable] = None, + shock_functions: Dict[str, Callable], ): """Create wrapped functions from user supplied functions. @@ -118,10 +117,6 @@ def process_alternative_sim_functions( agent's wealth matrices of the next period (t + 1). The inputs ```savings_grid```, ```income_shocks```, ```params``` and ```options``` are already partialled in. - - compute_upper_envelope (Callable): Function for calculating the upper envelope - of the policy and value function. If the number of discrete choices is 1, - this function is a dummy function that returns the policy and value - function as is, without performing a fast upper envelope scan. - transition_function (Callable): Partialled transition function that returns transition probabilities for each state. @@ -143,6 +138,9 @@ def process_alternative_sim_functions( state_space_functions, model_config=model_config, model_specs=model_specs, + additional_continuous_state_names=continuous_states_info[ + "additional_continuous_state_names" + ], ) ) @@ -155,18 +153,15 @@ def process_alternative_sim_functions( ) # Budget equation + _check_budget_does_not_declare_renamed_income_shock(budget_constraint) compute_assets_begin_of_period = ( determine_function_arguments_and_partial_model_specs( func=budget_constraint, model_specs=model_specs_jax, + not_allowed_state_choices=[], ) ) - # Upper envelope function - compute_upper_envelope = create_upper_envelope_function( - model_config=model_config, - ) - taste_shock_function_processed, taste_shock_scale_in_params = ( process_shock_functions( shock_functions=shock_functions, @@ -186,8 +181,11 @@ def process_alternative_sim_functions( "processed_stochastic_funcs": processed_stochastic_funcs_dict, "state_specific_choice_set": state_specific_choice_set, "next_period_deterministic_state": next_period_deterministic_state, - "compute_upper_envelope": compute_upper_envelope, "taste_shock_function": taste_shock_function_processed, + # The alternative budget equation has its own answer to "does wealth depend + # on the current choice?", and simulation reads it off these functions -- so + # it must be recomputed here rather than inherited. + "budget_depends_on_choice": _declares_choice(budget_constraint), } return alt_model_funcs, taste_shock_scale_in_params diff --git a/src/dcegm/pre_processing/batches/algo_batch_size.py b/src/dcegm/pre_processing/batches/algo_batch_size.py index d060b1ea..f7a0b63f 100644 --- a/src/dcegm/pre_processing/batches/algo_batch_size.py +++ b/src/dcegm/pre_processing/batches/algo_batch_size.py @@ -1,5 +1,7 @@ import numpy as np +from dcegm.pre_processing.batches.child_state_dedup import compute_child_dedup_for_batch + def determine_optimal_batch_size( bool_state_choices_to_batch, @@ -62,53 +64,42 @@ def determine_optimal_batch_size( child_states_to_integrate_exog = [] child_state_choices_to_aggr_choice = [] child_state_choice_idxs_to_interpolate = [] + representative_parent_state_choice_for_child = [] + unique_child_states_list = [] + representative_parent_state_choice_per_child_state_list = [] + state_row_for_state_choice_list = [] for i, batch in enumerate(batches_to_check): - # First get all child states and a mapping from the state-choice to the - # different child states due to exogenous change of states. - child_states_idxs = map_state_choice_to_child_states[batch] - unique_child_states, inverse_ids = np.unique( - child_states_idxs, return_index=False, return_inverse=True - ) - child_states_to_integrate_exog += [ - inverse_ids.reshape(child_states_idxs.shape) - ] - - # Next we use the child state indexes to get all unique child states and - # their corresponding state-choices. - child_states_batch = np.take(state_space, unique_child_states, axis=0) - child_states_tuple = tuple( - child_states_batch[:, i] for i in range(n_state_vars) - ) - unique_state_choice_idxs_childs = map_state_choice_to_index[ - child_states_tuple - ] - - # Now we create a mapping from the child-state choices back to the states - # with state-choices in columns for the choices ( + child_states_to_integrate_exog_batch, + child_state_choices_to_aggr_choice_batch, unique_child_state_choice_idxs, - inverse_child_state_choice_ids, - ) = np.unique( - unique_state_choice_idxs_childs, return_index=False, return_inverse=True + representative_parent_state_choice_batch, + unique_child_states_batch, + representative_parent_state_choice_per_child_state_batch, + state_row_for_state_choice_batch, + ) = compute_child_dedup_for_batch( + batch=batch, + map_state_choice_to_child_states=map_state_choice_to_child_states, + map_state_choice_to_index=map_state_choice_to_index, + state_space=state_space, + n_state_vars=n_state_vars, + invalid_state_idx=invalid_state_idx, + out_of_bounds_state_choice_idx=out_of_bounds_state_choice_idx, ) - - # Treat invalid choices: - if unique_child_state_choice_idxs[-1] == invalid_state_idx: - unique_child_state_choice_idxs = unique_child_state_choice_idxs[:-1] - inverse_child_state_choice_ids[ - inverse_child_state_choice_ids - >= np.max(inverse_child_state_choice_ids) - ] = out_of_bounds_state_choice_idx - - # Save the mapping from child-state-choices to child-states + child_states_to_integrate_exog += [child_states_to_integrate_exog_batch] child_state_choices_to_aggr_choice += [ - inverse_child_state_choice_ids.reshape( - unique_state_choice_idxs_childs.shape - ) + child_state_choices_to_aggr_choice_batch ] - # And the list of the unique child states. child_state_choice_idxs_to_interpolate += [unique_child_state_choice_idxs] + representative_parent_state_choice_for_child += [ + representative_parent_state_choice_batch + ] + unique_child_states_list += [unique_child_states_batch] + representative_parent_state_choice_per_child_state_list += [ + representative_parent_state_choice_per_child_state_batch + ] + state_row_for_state_choice_list += [state_row_for_state_choice_batch] # Now check if the smallest index of the child state choices is larger than # the maximum index of the batch, i.e. if all state choice relevant to @@ -119,14 +110,17 @@ def determine_optimal_batch_size( need_to_reduce_batchsize = True break - print("The batch size of the backwards induction is ", current_batch_size) - if not need_to_reduce_batchsize: batch_not_found = False return ( + size_last_period, batches_to_check, child_state_choice_idxs_to_interpolate, child_state_choices_to_aggr_choice, child_states_to_integrate_exog, + representative_parent_state_choice_for_child, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, ) diff --git a/src/dcegm/pre_processing/batches/batch_creation.py b/src/dcegm/pre_processing/batches/batch_creation.py index be5534b0..5da44f82 100644 --- a/src/dcegm/pre_processing/batches/batch_creation.py +++ b/src/dcegm/pre_processing/batches/batch_creation.py @@ -6,11 +6,180 @@ from dcegm.pre_processing.batches.single_segment import create_single_segment_of_batches +def _log_segment(id_segment, mode, bool_segment, segment_info, state_choice_space): + """Print one summary line per segment: period range, batch mode, batch shape.""" + periods = state_choice_space[bool_segment, 0] + p_lo, p_hi = int(periods.min()), int(periods.max()) + period_str = f"period {p_lo}" if p_lo == p_hi else f"periods {p_lo}-{p_hi}" + batches = np.asarray(segment_info["batches_state_choice_idx"]) + n_batches, width = batches.shape[0], batches.shape[1] + leftover = ( + "" if segment_info.get("batches_cover_all", True) else " (+1 leftover batch)" + ) + # period_max pads every batch to the widest period in the segment. + pad = ", padded to widest period" if mode == "period_max" else "" + batch_word = "batch" if n_batches == 1 else "batches" + print( + f" segment {id_segment}: {period_str} [{mode}] -> " + f"{n_batches} {batch_word} x {width} state-choices{pad}{leftover}" + ) + # largest_block shrinks the batch size from the segment's last-period size (the + # search seed) down to the final width; period_max does not search (start is None). + start = segment_info.get("search_start_size") + if start is not None: + print(f" start search: {start}") + print(f" end search: {width}") + + +def _log_total(segment_infos): + """Print the total number of sequential scan steps across all segments.""" + n_seg = segment_infos["n_segments"] + total = 0 + for s in range(n_seg): + info = segment_infos[f"batches_info_segment_{s}"] + n_batches = np.asarray(info["batches_state_choice_idx"]).shape[0] + total += n_batches + (0 if info.get("batches_cover_all", True) else 1) + print(f" {n_seg} segment(s), {total} scan steps total") + + +def _select_law_of_motion_arrays( + transition_depends_on_choice, + child_state_choices_no_proxy, + rep_parent_state_choice_idx_per_child_state_choice, + unique_child_states_state_dict, + rep_parent_state_choice_idx_per_child_state, + state_row_for_state_choice, +): + """Keep only the child arrays the taken law-of-motion branch reads. + + A choice-dependent budget equation evaluates the law of motion once per child + *state-choice*, so it needs the child's own (non-proxy) state-choice dict plus the + per-state-choice representative parent; otherwise the coarser per-unique- + child-*state* dedup suffices (see ``calc_law_of_motion`` in ``law_of_motion.py``). + Whichever branch is taken is a model-static property (``budget_depends_on_choice``), + so the selection is made once here at setup and the unused branch's arrays are never + threaded through the backward induction. + + ``child_state_choices_no_proxy`` is the transition child (its real state), not the + proxy value-reuse slot -- the proxy identity stays on the separate + ``state_choice_mat_child`` entry used to read the stored value/policy. + + """ + if transition_depends_on_choice: + return { + "child_state_choices": child_state_choices_no_proxy, + "rep_parent_state_choice_idx_per_child_state_choice": ( + rep_parent_state_choice_idx_per_child_state_choice + ), + } + return { + "unique_child_states": unique_child_states_state_dict, + "rep_parent_state_choice_idx_per_child_state": ( + rep_parent_state_choice_idx_per_child_state + ), + "state_row_for_state_choice": state_row_for_state_choice, + } + + +def _bundle_segment_law_of_motion_arrays(segment_info, transition_depends_on_choice): + """Fold a scan segment's per-branch law-of-motion arrays into one dict. + + Pops the four alternatives out of ``segment_info`` (and its leftover + ``last_batch_info``, if any) and replaces them with a single + ``law_of_motion_arrays`` entry carrying only the taken branch -- the dict + ``calc_law_of_motion`` reads, threaded through the backward-induction scan. + + """ + for info in (segment_info, segment_info.get("last_batch_info")): + if info is None: + continue + info["law_of_motion_arrays"] = _select_law_of_motion_arrays( + transition_depends_on_choice=transition_depends_on_choice, + child_state_choices_no_proxy=info.pop("state_choices_childs_no_proxy"), + rep_parent_state_choice_idx_per_child_state_choice=info.pop( + "rep_parent_state_choice_idx_per_child_state_choice" + ), + unique_child_states_state_dict=info.pop( + "state_choices_unique_child_states" + ), + rep_parent_state_choice_idx_per_child_state=info.pop( + "rep_parent_state_choice_idx_per_child_state" + ), + state_row_for_state_choice=info.pop("state_row_for_state_choice"), + ) + + +def _bundle_final_period_law_of_motion_arrays( + last_two_period_info, transition_depends_on_choice, state_space_dict +): + """Fold the final period's per-branch law-of-motion arrays into one dict. + + The final period reaches ``calc_law_of_motion`` outside the scan (see + ``final_periods.solve_final_period``), so it gets the same ``law_of_motion_arrays`` + bundle, just from the final-period-specific source keys. On the state-dedup branch + the unique child *states* are gathered into a state dict here (mirroring the per- + batch gather in ``single_segment.py``) rather than in the solve. + + """ + unique_final_period_states = last_two_period_info.pop("unique_final_period_states") + rep_parent_idx_per_state_choice = last_two_period_info.pop( + "rep_sec_last_period_parent_idx_per_final_state_choice" + ) + rep_parent_idx_per_state = last_two_period_info.pop( + "representative_second_last_period_parent_idx_per_final_state" + ) + state_row_for_state_choice = last_two_period_info.pop( + "state_row_for_final_period_state_choice" + ) + unique_child_states_state_dict = { + key: var[unique_final_period_states] for key, var in state_space_dict.items() + } + # The final period's children are its own (solved, real) state-choices, so the + # transition child is simply state_choice_mat_final_period -- no proxy is + # involved here (death proxies target the last period, which is solved). + last_two_period_info["law_of_motion_arrays"] = _select_law_of_motion_arrays( + transition_depends_on_choice=transition_depends_on_choice, + child_state_choices_no_proxy=last_two_period_info[ + "state_choice_mat_final_period" + ], + rep_parent_state_choice_idx_per_child_state_choice=rep_parent_idx_per_state_choice, + unique_child_states_state_dict=unique_child_states_state_dict, + rep_parent_state_choice_idx_per_child_state=rep_parent_idx_per_state, + state_row_for_state_choice=state_row_for_state_choice, + ) + + +def bundle_law_of_motion_arrays( + batch_info, transition_depends_on_choice, state_space_dict +): + """Assemble every ``law_of_motion_arrays`` bundle in a finished ``batch_info``. + + Run once at model setup: walks the final-period info and each scan segment + (including any leftover batch) and replaces the per-branch law-of-motion index + arrays with the single dict ``calc_law_of_motion`` reads, dropping the branch + that this model never takes. + + """ + _bundle_final_period_law_of_motion_arrays( + batch_info["last_two_period_info"], + transition_depends_on_choice=transition_depends_on_choice, + state_space_dict=state_space_dict, + ) + if batch_info["two_period_model"]: + return batch_info + for id_segment in range(batch_info["n_segments"]): + _bundle_segment_law_of_motion_arrays( + batch_info[f"batches_info_segment_{id_segment}"], + transition_depends_on_choice=transition_depends_on_choice, + ) + return batch_info + + def create_batches_and_information( model_structure, n_periods, - min_period_batch_segments=None, - batch_mode="largest_block", + min_period_batch_segments, + batch_mode, ): """Batches are used instead of periods to have chunks of equal sized state choices. The returned batch information dictionary contains the following arrays @@ -62,6 +231,7 @@ def create_batches_and_information( state_choice_space = model_structure["state_choice_space"] bool_state_choices_to_batch = state_choice_space[:, 0] < n_periods - 2 + print(" batches:") valid_batch_modes = {"largest_block", "period_max"} @@ -84,6 +254,13 @@ def create_batches_and_information( "n_segments": 1, "batches_info_segment_0": single_batch_segment_info, } + _log_segment( + 0, + batch_mode, + bool_state_choices_to_batch, + single_batch_segment_info, + state_choice_space, + ) else: @@ -145,6 +322,13 @@ def create_batches_and_information( batch_mode=batch_mode[id_segment], ) segment_infos[f"batches_info_segment_{id_segment}"] = segment_batch_info + _log_segment( + id_segment, + batch_mode[id_segment], + bool_state_choices_segment, + segment_batch_info, + state_choice_space, + ) # Set the bools to False which have been batched already bool_state_choices_to_batch = bool_state_choices_to_batch & split_cond @@ -159,6 +343,15 @@ def create_batches_and_information( segment_infos[f"batches_info_segment_{n_segments - 1}"] = ( last_segment_batch_info ) + _log_segment( + n_segments - 1, + batch_mode[n_segments - 1], + bool_state_choices_to_batch, + last_segment_batch_info, + state_choice_space, + ) + + _log_total(segment_infos) batch_info = { # First two bools determining the structure of solution functions we call diff --git a/src/dcegm/pre_processing/batches/child_state_dedup.py b/src/dcegm/pre_processing/batches/child_state_dedup.py new file mode 100644 index 00000000..78b99518 --- /dev/null +++ b/src/dcegm/pre_processing/batches/child_state_dedup.py @@ -0,0 +1,163 @@ +import numpy as np + + +def compute_child_dedup_for_batch( + batch, + map_state_choice_to_child_states, + map_state_choice_to_index, + state_space, + n_state_vars, + invalid_state_idx, + out_of_bounds_state_choice_idx, +): + """Deduplicate a batch's child states/state-choices; find a representative parent. + + Shared by both batch-size algorithms (``algo_batch_size.py``'s "largest_block" + and ``single_segment.py``'s "period_max") -- they deduplicated this identically + before this was factored out, so keeping one implementation avoids having to fix + the same bug twice. + + The batch code computes the continuation-value interpolation once per unique + child (deduplicated purely by discrete-state index, see the two ``np.unique`` + calls below) and reuses it for every parent state-choice that maps to it. That + computation needs the *parent's* own continuous grid as input to the child's + law-of-motion function (see ``law_of_motion.py``) -- not the child's. Since + the computation runs once per unique child rather than once per parent, it + needs *a* representative parent for each unique child; any + one of them works, because + ``check_continuous_grid_consistency_across_shared_children`` (run once at + model-build time, before batching) already guarantees every parent state-choice + sharing a child agrees on its own grid. + + The representative parent is a *state-choice* (an index into + ``batch``/``state_choice_space``), not a bare state: grids live on the + state-choice space (that's where the solution itself lives), so ``batch`` + already gives us exactly the identity we need, with no extra lookup required. + + The proxy (``map_state_choice_to_index`` here is the *with-proxy* indexer) is a + pure value-reuse pointer, kept out of the transition granularity: the rows are + the *non-proxy* child state-choices, so distinct children that reuse one solved + slot (death at different ages -> one last-period slot) stay separate for the law + of motion, and only the value/policy lookup follows the proxy (``value_slot``, + which then repeats across those rows). + + Also exposes the *state*-level dedup: ``unique_child_states``, + ``representative_parent_state_choice_per_child_state``, and + ``state_row_for_state_choice``. These let a caller evaluate something once per + unique child *state* (e.g. a law-of-motion function that doesn't depend on the + child's own choice) and then ``np.take``/``jnp.take`` the result out to the + state-choice granularity ``value_slot`` needs for reading stored policy/value -- + mirrors the existing ``child_states_to_integrate_exog`` gather used one stage + later, in ``calculate_candidate_solutions_from_euler_equation``. + + Returns: + tuple: + + - child_states_to_integrate_exog (np.ndarray): shape (len(batch), n_stochastic), + maps each (parent row, stochastic draw) to a position in the unique + child *state* space. + - child_state_choices_to_aggr_choice (np.ndarray): shape (n_unique_child_states, + n_choices), maps each (child state, choice) to its row in the non-proxy + child *state-choice* axis (or an out-of-bounds sentinel). + - value_slot (np.ndarray): one entry per non-proxy child state-choice; the + solved (proxy) state-choice index whose value/policy/endog that row + reads. Repeats across rows whose children share a proxy slot. + - representative_parent_state_choice_for_child (np.ndarray): same length as + value_slot; for each row, the state-choice index of one parent (from + this batch) that transitions to it. + - unique_child_states (np.ndarray): the deduplicated child *state* indices + (into ``state_space``) -- one entry per unique child state, collapsing + across that state's own choices. + - representative_parent_state_choice_per_child_state (np.ndarray): same + length as unique_child_states; for each unique child state, the + state-choice index of one parent that transitions to it (same value + used across all of that state's own choices in + representative_parent_state_choice_for_child). + - state_row_for_state_choice (np.ndarray): same length as value_slot; for + each non-proxy child state-choice, its row position in + unique_child_states -- the gather index needed to expand a per-state + result out to per-state-choice granularity. + + """ + child_states_idxs = map_state_choice_to_child_states[batch] + n_stochastic_states = child_states_idxs.shape[1] + + unique_child_states, first_occurrence_state, inverse_ids = np.unique( + child_states_idxs, return_index=True, return_inverse=True + ) + child_states_to_integrate_exog = inverse_ids.reshape(child_states_idxs.shape) + n_unique_child_states = unique_child_states.shape[0] + + # A representative parent state-choice (by local row within `batch`) for each + # unique child state -- the batch row/stochastic-draw where that child state + # first occurs. + representative_parent_row = first_occurrence_state // n_stochastic_states + representative_parent_state_choice_per_child_state = batch[ + representative_parent_row + ] + + child_states_batch = np.take(state_space, unique_child_states, axis=0) + child_states_tuple = tuple(child_states_batch[:, i] for i in range(n_state_vars)) + # Proxy solved state-choice index for each (unique child state, choice) cell. + # `map_state_choice_to_index` is the *with-proxy* indexer, so a child whose own + # solution is reused points at the reused (proxy) slot; cells the child does not + # admit map to `invalid_state_idx`. + proxy_state_choice_idxs_childs = map_state_choice_to_index[child_states_tuple] + + # The representative parent is a property of the child *state*, not of which of + # its choices we're looking at, so it is the same across all n_choices columns + # for a given child-state row. + representative_parent_per_cell = np.broadcast_to( + representative_parent_state_choice_per_child_state[:, None], + proxy_state_choice_idxs_childs.shape, + ) + state_row_per_cell = np.broadcast_to( + np.arange(n_unique_child_states)[:, None], + proxy_state_choice_idxs_childs.shape, + ) + + # One row per (child state, choice) the child admits -- the *non-proxy* + # granularity the law of motion needs. Distinct children that reuse the same + # solved slot (e.g. death at different ages, all proxied to one last-period + # slot) stay separate rows here, because the transition *into* them differs; + # only the value/policy lookup collapses, via the repeated ``value_slot`` below. + # Ordering the admitted cells by their proxy solved index keeps rows that share + # a proxy adjacent; when no proxy collapses two children (proxy == actual) each + # cell simply keeps its own row in solved-index order. + flat_proxy_idx = proxy_state_choice_idxs_childs.ravel() + admitted_cell = flat_proxy_idx != invalid_state_idx + ordered_cells = np.where(admitted_cell)[0][ + np.argsort(flat_proxy_idx[admitted_cell], kind="stable") + ] + + # value_slot: the solved (proxy) state-choice whose value/policy/endog this row + # reads -- repeated across rows that share a proxy. state_row_for_state_choice: + # the row's own (non-proxy) child *state* row, for expanding a per-state law of + # motion back out to per-state-choice granularity. + value_slot = flat_proxy_idx[ordered_cells] + state_row_for_state_choice = state_row_per_cell.ravel()[ordered_cells] + representative_parent_state_choice_for_child = ( + representative_parent_per_cell.ravel()[ordered_cells] + ) + + # For each (child state, choice) cell, its row position, or an out-of-bounds + # sentinel for choices the child does not admit -- the choice-aggregation map. + flat_aggr = np.full( + flat_proxy_idx.shape[0], + fill_value=out_of_bounds_state_choice_idx, + dtype=int, + ) + flat_aggr[ordered_cells] = np.arange(ordered_cells.shape[0]) + child_state_choices_to_aggr_choice = flat_aggr.reshape( + proxy_state_choice_idxs_childs.shape + ) + + return ( + child_states_to_integrate_exog, + child_state_choices_to_aggr_choice, + value_slot, + representative_parent_state_choice_for_child, + unique_child_states, + representative_parent_state_choice_per_child_state, + state_row_for_state_choice, + ) diff --git a/src/dcegm/pre_processing/batches/last_two_periods.py b/src/dcegm/pre_processing/batches/last_two_periods.py index 8638014d..c92c3819 100644 --- a/src/dcegm/pre_processing/batches/last_two_periods.py +++ b/src/dcegm/pre_processing/batches/last_two_periods.py @@ -66,12 +66,103 @@ def add_last_two_period_information( idx_state_choice_final_period ] + # A representative second-to-last-period *state-choice* for each final-period + # state-choice -- used only to pick which state-choice's own continuous grid + # feeds the law of motion in solve_final_period (see law_of_motion.py). Grids + # live on the state-choice space (that's where the solution itself lives), so + # the representative parent is a state-choice index, not a bare state. Mirrors + # the same computation the main backward induction loop does per batch (see + # child_state_dedup.py), specialized for the (undeduplicated) two-period case: + # solve_final_period computes the law of motion for every final-period + # state-choice individually, so we need a representative parent aligned with + # every entry of idx_state_choice_final_period, not just a deduplicated subset. + child_states_idxs_from_second_last = map_state_choice_to_child_states[ + idx_state_choice_second_last_period + ] + n_stochastic_states = child_states_idxs_from_second_last.shape[1] + unique_final_states, first_occurrence_flat = np.unique( + child_states_idxs_from_second_last, return_index=True + ) + representative_second_last_row = first_occurrence_flat // n_stochastic_states + # idx_state_choice_second_last_period[...] is already a state-choice index -- + # exactly the true (previous-period) parent state-choice we want, no need to + # collapse it down to a bare state. + representative_parent_state_choice_of_unique_final_state = ( + idx_state_choice_second_last_period[representative_second_last_row] + ) + + invalid_state_idx = np.iinfo(map_state_choice_to_child_states.dtype).max + valid_mask = unique_final_states != invalid_state_idx + + # parent_states_final_period (each final-period state-choice's own bare state) + # is used purely as a lookup key here, matching each final-period row to its + # corresponding entry in unique_final_states -- unrelated to grid selection. + # unique_final_states is sorted (np.unique guarantees this), so this is a + # vectorized lookup via searchsorted; any final-period state not actually + # reachable from the second-to-last period (e.g. an absorbing/initial-like + # state; already flagged elsewhere via a warning in test_child_state_mapping, + # not an error) falls back to its own state-choice identity, mirroring + # calc_law_of_motion_for_state_choices's own + # representative_parent_state_choice_vec=None -> self fallback. + insert_pos = np.clip( + np.searchsorted(unique_final_states, parent_states_final_period), + 0, + len(unique_final_states) - 1, + ) + found_mask = ( + unique_final_states[insert_pos] == parent_states_final_period + ) & valid_mask[insert_pos] + rep_sec_last_period_parent_idx_per_final_state_choice = np.where( + found_mask, + representative_parent_state_choice_of_unique_final_state[insert_pos], + idx_state_choice_final_period, + ) + + # State-level counterpart of the two arrays above, for the deduplicated law of + # motion (see calc_law_of_motion in law_of_motion.py). Mirrors what + # child_state_dedup.py builds per batch in the main backward-induction loop, so + # the final period takes the same shortcut instead of always evaluating the + # transition once per state-choice. + # + # parent_states_final_period is each final-period state-choice's own bare state, + # so np.unique over it gives exactly the dedup we need plus the gather index + # mapping every final-period state-choice back to its row. + ( + unique_final_period_states, + first_occurrence_of_final_state, + state_row_for_final_period_state_choice, + ) = np.unique(parent_states_final_period, return_index=True, return_inverse=True) + + # The representative parent is a property of the child *state*, so any of that + # state's own choices carries the same value and the first occurrence is + # representative. The one exception is the `found_mask=False` fallback above, + # which uses each state-choice's own identity and so does differ across a + # state's choices -- but that branch only fires for final-period states that no + # second-to-last-period state-choice transitions into (already warned about in + # test_child_state_mapping), whose final-period values are never aggregated + # back, so the choice of representative there is not observable. + representative_second_last_period_parent_idx_per_final_state = ( + rep_sec_last_period_parent_idx_per_final_state_choice[ + first_occurrence_of_final_state + ] + ) + last_two_period_info = { "idx_state_choices_final_period": idx_state_choice_final_period, "idx_state_choices_second_last_period": idx_state_choice_second_last_period, "idxs_parent_states_final_period": parent_states_final_period, "state_to_choices_final_period": state_to_choices_final_period, "child_states_second_last_period": child_states_second_last_period, + "rep_sec_last_period_parent_idx_per_final_state_choice": ( + rep_sec_last_period_parent_idx_per_final_state_choice + ), + "unique_final_period_states": unique_final_period_states, + "representative_second_last_period_parent_idx_per_final_state": ( + representative_second_last_period_parent_idx_per_final_state + ), + "state_row_for_final_period_state_choice": ( + state_row_for_final_period_state_choice + ), } state_choice_space_dict = model_structure["state_choice_space_dict"] diff --git a/src/dcegm/pre_processing/batches/single_segment.py b/src/dcegm/pre_processing/batches/single_segment.py index 2517e897..2081ff29 100644 --- a/src/dcegm/pre_processing/batches/single_segment.py +++ b/src/dcegm/pre_processing/batches/single_segment.py @@ -1,12 +1,39 @@ import numpy as np from dcegm.pre_processing.batches.algo_batch_size import determine_optimal_batch_size +from dcegm.pre_processing.batches.child_state_dedup import compute_child_dedup_for_batch + + +def build_no_proxy_child_state_choices( + state_space_dict, + unique_child_states, + state_row_for_state_choice, + proxy_child_state_choices, +): + """Reconstruct the *non-proxy* child state-choice dict for the law of motion. + + The proxy (value-reuse) child in ``proxy_child_state_choices`` may point at a + different discrete state than the real child (e.g. death at a different age all + reusing one last-period slot), so the transition *into* the child must use the + child's own state -- recovered from ``unique_child_states`` via + ``state_row_for_state_choice``. The choice is unaffected by the sparsity proxy (it + maps the state, keeping the choice axis), so it is taken straight from the proxy + child. Works for both the aligned per-segment arrays (2d, batch axis) and a single + leftover batch (1d). + + """ + child_state_index = np.take_along_axis( + unique_child_states, state_row_for_state_choice, axis=-1 + ) + no_proxy = {key: var[child_state_index] for key, var in state_space_dict.items()} + no_proxy["choice"] = proxy_child_state_choices["choice"] + return no_proxy def create_single_segment_of_batches( bool_state_choices_to_batch, model_structure, - batch_mode="largest_block", + batch_mode, ): """Create a single segment of evenly sized batches. @@ -16,23 +43,34 @@ def create_single_segment_of_batches( state_choice_space = model_structure["state_choice_space"] state_choice_space_dict = model_structure["state_choice_space_dict"] - state_space = model_structure["state_space"] discrete_states_names = model_structure["discrete_states_names"] + state_space = model_structure["state_space_incl_proxies"] + state_space_dict = { + name: state_space[:, i] for i, name in enumerate(discrete_states_names) + } map_state_choice_to_parent_state = model_structure[ "map_state_choice_to_parent_state" ] map_state_choice_to_child_states = model_structure[ - "map_state_choice_to_child_states" + "map_state_choice_to_child_states_actual" ] map_state_choice_to_index = model_structure["map_state_choice_to_index_with_proxy"] + # search_start_size is the seed the largest_block search shrinks from + # (size of the segment's last period); None for period_max, which does not search. + search_start_size = None if batch_mode == "largest_block": ( + search_start_size, batches_list, child_state_choice_idxs_to_interp_list, child_state_choices_to_aggr_choice_list, child_states_to_integrate_stochastic_list, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, ) = determine_optimal_batch_size( bool_state_choices_to_batch=bool_state_choices_to_batch, state_choice_space=state_choice_space, @@ -46,6 +84,10 @@ def create_single_segment_of_batches( child_states_to_integrate_stochastic_list, child_state_choices_to_aggr_choice_list, child_state_choice_idxs_to_interp_list, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, batches_cover_all, last_batch_info, ) = correct_for_uneven_last_batch( @@ -53,7 +95,12 @@ def create_single_segment_of_batches( child_states_to_integrate_stochastic_list, child_state_choices_to_aggr_choice_list, child_state_choice_idxs_to_interp_list, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, state_choice_space_dict, + state_space_dict, map_state_choice_to_parent_state, ) elif batch_mode == "period_max": @@ -62,6 +109,10 @@ def create_single_segment_of_batches( child_state_choice_idxs_to_interp_list, child_state_choices_to_aggr_choice_list, child_states_to_integrate_stochastic_list, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, ) = determine_period_max_batch_size( bool_state_choices_to_batch=bool_state_choices_to_batch, state_choice_space=state_choice_space, @@ -81,11 +132,17 @@ def create_single_segment_of_batches( child_states_to_integrate_stochastic_list, child_state_choices_to_aggr_choice_list, child_state_choice_idxs_to_interp_list, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, state_choice_space_dict, + state_space_dict, map_state_choice_to_parent_state, discrete_states_names, ) single_batch_segment_info["batches_cover_all"] = batches_cover_all + single_batch_segment_info["search_start_size"] = search_start_size if not batches_cover_all: single_batch_segment_info["last_batch_info"] = last_batch_info @@ -115,45 +172,43 @@ def determine_period_max_batch_size( child_states_to_integrate_exog = [] child_state_choices_to_aggr_choice = [] child_state_choice_idxs_to_interpolate = [] + representative_parent_state_choice_for_child = [] + unique_child_states_list = [] + representative_parent_state_choice_per_child_state_list = [] + state_row_for_state_choice_list = [] for period in periods_unique_desc: batch = idx_state_choice_raw[periods_to_batch == period] batches_to_check += [batch] - child_states_idxs = map_state_choice_to_child_states[batch] - unique_child_states, inverse_ids = np.unique( - child_states_idxs, return_index=False, return_inverse=True - ) - child_states_to_integrate_exog += [inverse_ids.reshape(child_states_idxs.shape)] - - child_states_batch = np.take(state_space, unique_child_states, axis=0) - child_states_tuple = tuple( - child_states_batch[:, i] for i in range(n_state_vars) - ) - unique_state_choice_idxs_childs = map_state_choice_to_index[child_states_tuple] - ( + child_states_to_integrate_exog_batch, + child_state_choices_to_aggr_choice_batch, unique_child_state_choice_idxs, - inverse_child_state_choice_ids, - ) = np.unique( - unique_state_choice_idxs_childs, return_index=False, return_inverse=True + representative_parent_state_choice_batch, + unique_child_states_batch, + representative_parent_state_choice_per_child_state_batch, + state_row_for_state_choice_batch, + ) = compute_child_dedup_for_batch( + batch=batch, + map_state_choice_to_child_states=map_state_choice_to_child_states, + map_state_choice_to_index=map_state_choice_to_index, + state_space=state_space, + n_state_vars=n_state_vars, + invalid_state_idx=invalid_state_idx, + out_of_bounds_state_choice_idx=out_of_bounds_state_choice_idx, ) - - if ( - len(unique_child_state_choice_idxs) > 0 - and unique_child_state_choice_idxs[-1] == invalid_state_idx - ): - unique_child_state_choice_idxs = unique_child_state_choice_idxs[:-1] - inverse_child_state_choice_ids[ - inverse_child_state_choice_ids >= np.max(inverse_child_state_choice_ids) - ] = out_of_bounds_state_choice_idx - - child_state_choices_to_aggr_choice += [ - inverse_child_state_choice_ids.reshape( - unique_state_choice_idxs_childs.shape - ) - ] + child_states_to_integrate_exog += [child_states_to_integrate_exog_batch] + child_state_choices_to_aggr_choice += [child_state_choices_to_aggr_choice_batch] child_state_choice_idxs_to_interpolate += [unique_child_state_choice_idxs] + representative_parent_state_choice_for_child += [ + representative_parent_state_choice_batch + ] + unique_child_states_list += [unique_child_states_batch] + representative_parent_state_choice_per_child_state_list += [ + representative_parent_state_choice_per_child_state_batch + ] + state_row_for_state_choice_list += [state_row_for_state_choice_batch] max_batch_size = max(len(batch) for batch in batches_to_check) @@ -177,6 +232,10 @@ def determine_period_max_batch_size( child_state_choice_idxs_to_interpolate, child_state_choices_to_aggr_choice, child_states_to_integrate_exog, + representative_parent_state_choice_for_child, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, ) @@ -185,7 +244,12 @@ def correct_for_uneven_last_batch( child_states_to_integrate_stochastic_list, child_state_choices_to_aggr_choice_list, child_state_choice_idxs_to_interp_list, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, state_choice_space_dict, + state_space_dict, map_state_choice_to_parent_state, ): """Check if the last batch has the same length as the others. @@ -216,6 +280,14 @@ def correct_for_uneven_last_batch( ] last_idx_to_aggregate_choice = child_state_choices_to_aggr_choice_list[-1] last_child_state_idx_interp = child_state_choice_idxs_to_interp_list[-1] + last_representative_parent_state_choice_for_child = ( + representative_parent_state_choice_for_child_list[-1] + ) + last_unique_child_states = unique_child_states_list[-1] + last_representative_parent_state_choice_per_child_state = ( + representative_parent_state_choice_per_child_state_list[-1] + ) + last_state_row_for_state_choice = state_row_for_state_choice_list[-1] last_state_choices = { key: var[last_batch] for key, var in state_choice_space_dict.items() @@ -227,6 +299,15 @@ def correct_for_uneven_last_batch( last_parent_state_idx_of_state_choice = map_state_choice_to_parent_state[ last_child_state_idx_interp ] + last_state_choices_unique_child_states = { + key: var[last_unique_child_states] for key, var in state_space_dict.items() + } + last_state_choices_childs_no_proxy = build_no_proxy_child_state_choices( + state_space_dict=state_space_dict, + unique_child_states=last_unique_child_states, + state_row_for_state_choice=last_state_row_for_state_choice, + proxy_child_state_choices=last_state_choices_childs, + ) last_batch_info = { "state_choice_idx": last_batch, @@ -237,6 +318,18 @@ def correct_for_uneven_last_batch( "child_state_choice_idxs_to_interp": last_child_state_idx_interp, "child_states_idxs": last_parent_state_idx_of_state_choice, "state_choices_childs": last_state_choices_childs, + "state_choices_childs_no_proxy": last_state_choices_childs_no_proxy, + "rep_parent_state_choice_idx_per_child_state_choice": ( + last_representative_parent_state_choice_for_child + ), + # State-level dedup (see child_state_dedup.py) for the cheap + # budget/law-of-motion path when a user's transition function doesn't + # depend on "choice". + "state_choices_unique_child_states": last_state_choices_unique_child_states, + "rep_parent_state_choice_idx_per_child_state": ( + last_representative_parent_state_choice_per_child_state + ), + "state_row_for_state_choice": last_state_row_for_state_choice, } batches_list = batches_list[:-1] child_states_to_integrate_stochastic_list = ( @@ -248,11 +341,23 @@ def correct_for_uneven_last_batch( child_state_choice_idxs_to_interp_list = child_state_choice_idxs_to_interp_list[ :-1 ] + representative_parent_state_choice_for_child_list = ( + representative_parent_state_choice_for_child_list[:-1] + ) + unique_child_states_list = unique_child_states_list[:-1] + representative_parent_state_choice_per_child_state_list = ( + representative_parent_state_choice_per_child_state_list[:-1] + ) + state_row_for_state_choice_list = state_row_for_state_choice_list[:-1] return ( batches_list, child_states_to_integrate_stochastic_list, child_state_choices_to_aggr_choice_list, child_state_choice_idxs_to_interp_list, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, batches_cover_all, last_batch_info, ) @@ -263,7 +368,12 @@ def prepare_and_align_batch_arrays( child_states_to_integrate_stochastic_list, child_state_choices_to_aggr_choice_list, child_state_choice_idxs_to_interp_list, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, state_choice_space_dict, + state_space_dict, map_state_choice_to_parent_state, discrete_states_names, ): @@ -290,10 +400,22 @@ def prepare_and_align_batch_arrays( ( child_state_choice_idxs_to_interp, child_state_choices_to_aggr_choice, + representative_parent_state_choice_idx, + unique_child_states, + rep_parent_state_choice_idx_per_child_state, + state_row_for_state_choice, ) = extend_child_state_choices_to_aggregate_choices( idx_to_aggregate_choice=child_state_choices_to_aggr_choice_list, max_child_state_index_batch=max_child_state_index_batch, idx_to_interpolate=child_state_choice_idxs_to_interp_list, + representative_parent_state_choice_for_child_list=( + representative_parent_state_choice_for_child_list + ), + unique_child_states_list=unique_child_states_list, + representative_parent_state_choice_per_child_state_list=( + representative_parent_state_choice_per_child_state_list + ), + state_row_for_state_choice_list=state_row_for_state_choice_list, out_of_bounds_state_choice_idx=out_of_bounds_state_choice_idx, ) parent_state_idx_of_state_choice = map_state_choice_to_parent_state[ @@ -303,6 +425,15 @@ def prepare_and_align_batch_arrays( key: var[child_state_choice_idxs_to_interp] for key, var in state_choice_space_dict.items() } + state_choices_unique_child_states = { + key: var[unique_child_states] for key, var in state_space_dict.items() + } + state_choices_childs_no_proxy = build_no_proxy_child_state_choices( + state_space_dict=state_space_dict, + unique_child_states=unique_child_states, + state_row_for_state_choice=state_row_for_state_choice, + proxy_child_state_choices=state_choices_childs, + ) batch_info = { # Now the batch array information. First the batch itself @@ -314,6 +445,29 @@ def prepare_and_align_batch_arrays( "child_state_choice_idxs_to_interp": child_state_choice_idxs_to_interp, "child_states_idxs": parent_state_idx_of_state_choice, "state_choices_childs": state_choices_childs, + # The non-proxy children, at the same per-row granularity as + # state_choices_childs: their real state (death at its actual age), for the + # law-of-motion transition. state_choices_childs above keeps the proxy + # identity, since that is where the reused value/policy is stored. + "state_choices_childs_no_proxy": state_choices_childs_no_proxy, + # State-choice index of a representative parent for each unique child -- + # used only to pick which state-choice's own continuous grid to feed the + # law of motion (see law_of_motion.py). Not to be confused with + # "child_states_idxs" above, which despite the name of + # the underlying map_state_choice_to_parent_state array means "the state + # this child state-choice belongs to" (its own state, dropping its + # choice) -- an unrelated, pre-existing field. + "rep_parent_state_choice_idx_per_child_state_choice": representative_parent_state_choice_idx, + # State-level dedup (see child_state_dedup.py): a law-of-motion/budget + # function that doesn't depend on "choice" can be evaluated once per + # unique child state here, then gathered out to state-choice granularity + # via state_row_for_state_choice -- mirrors child_states_to_integrate_exog + # one stage later. + "state_choices_unique_child_states": state_choices_unique_child_states, + "rep_parent_state_choice_idx_per_child_state": ( + rep_parent_state_choice_idx_per_child_state + ), + "state_row_for_state_choice": state_row_for_state_choice, } return batch_info @@ -322,6 +476,10 @@ def extend_child_state_choices_to_aggregate_choices( idx_to_aggregate_choice, max_child_state_index_batch, idx_to_interpolate, + representative_parent_state_choice_for_child_list, + unique_child_states_list, + representative_parent_state_choice_per_child_state_list, + state_row_for_state_choice_list, out_of_bounds_state_choice_idx, ): """In case of uneven batches, we need to extend the child state objects to cover the @@ -382,4 +540,67 @@ def extend_child_state_choices_to_aggregate_choices( id_batch, : len(idx_to_interpolate[id_batch]) ] = idx_to_interpolate[id_batch] - return child_state_choice_idxs_to_interp, child_state_choices_to_aggr_choice + # Representative parent state-choice (for continuous-grid selection, see + # law_of_motion.py) is aligned 1:1 with idx_to_interpolate, so it is padded the + # exact same way -- padded slots are never read downstream (the corresponding + # padded child_state_choice_idxs_to_interp entries just duplicate a valid state, + # so nothing indexes into these extra rows specifically). + dummy_parent_state_choice = representative_parent_state_choice_for_child_list[0][0] + representative_parent_state_choice_idx = np.full( + (n_batches, max_child_state_choices), + fill_value=dummy_parent_state_choice, + dtype=int, + ) + for id_batch in range(n_batches): + representative_parent_state_choice_idx[ + id_batch, : len(representative_parent_state_choice_for_child_list[id_batch]) + ] = representative_parent_state_choice_for_child_list[id_batch] + + # State-level dedup (see child_state_dedup.py): unique_child_states and + # representative_parent_state_choice_per_child_state have their own per-batch + # length (n_unique_child_states -- collapsed across each state's own choices), + # independent of max_child_state_choices above, so they get their own max and + # padding, the same recipe as child_state_choice_idxs_to_interp. + max_n_unique_child_states = max(map(len, unique_child_states_list)) + dummy_child_state = unique_child_states_list[0][0] + unique_child_states = np.full( + (n_batches, max_n_unique_child_states), + fill_value=dummy_child_state, + dtype=int, + ) + dummy_parent_state_choice_for_state = ( + representative_parent_state_choice_per_child_state_list[0][0] + ) + rep_parent_state_choice_idx_per_child_state = np.full( + (n_batches, max_n_unique_child_states), + fill_value=dummy_parent_state_choice_for_state, + dtype=int, + ) + for id_batch in range(n_batches): + n_unique = len(unique_child_states_list[id_batch]) + unique_child_states[id_batch, :n_unique] = unique_child_states_list[id_batch] + rep_parent_state_choice_idx_per_child_state[id_batch, :n_unique] = ( + representative_parent_state_choice_per_child_state_list[id_batch] + ) + + # state_row_for_state_choice is aligned 1:1 with idx_to_interpolate, so it is + # padded the exact same way, with the same "padded slots never read + # downstream" convention -- 0 is always a valid row of the (separately + # padded) unique_child_states array above, regardless of batch. + state_row_for_state_choice = np.zeros( + (n_batches, max_child_state_choices), + dtype=int, + ) + for id_batch in range(n_batches): + state_row_for_state_choice[ + id_batch, : len(state_row_for_state_choice_list[id_batch]) + ] = state_row_for_state_choice_list[id_batch] + + return ( + child_state_choice_idxs_to_interp, + child_state_choices_to_aggr_choice, + representative_parent_state_choice_idx, + unique_child_states, + rep_parent_state_choice_idx_per_child_state, + state_row_for_state_choice, + ) diff --git a/src/dcegm/pre_processing/check_model_config.py b/src/dcegm/pre_processing/check_model_config.py index 798a819e..e2e3c113 100644 --- a/src/dcegm/pre_processing/check_model_config.py +++ b/src/dcegm/pre_processing/check_model_config.py @@ -59,19 +59,23 @@ def check_model_config_and_process(model_config): raise ValueError( "model_config['assets_end_of_period'] must contain wealth as key." ) - # Check if it is an array + # Check if it is an array, or `None` -- see the note on `None` grids below. asset_grid = continuous_states_grids["assets_end_of_period"] - if not isinstance(asset_grid, (list, np.ndarray, jnp.ndarray)): + if asset_grid is not None and not isinstance( + asset_grid, (list, np.ndarray, jnp.ndarray) + ): raise ValueError( - "model_config['continuous_states']['assets_end_of_period'] must be a list or an array." + "model_config['continuous_states']['assets_end_of_period'] must be a list, " + "an array, or None (paired with a continuous_grid_functions entry)." ) # ToDo: Check if it is monotonic increasing continuous_states_info = {} - n_assets_end_of_period = len(asset_grid) - continuous_states_info["assets_grid_end_of_period"] = jnp.asarray( - continuous_states_grids["assets_end_of_period"], dtype=float + n_assets_end_of_period = None if asset_grid is None else len(asset_grid) + continuous_states_info["n_assets_end_of_period"] = n_assets_end_of_period + continuous_states_info["assets_grid_end_of_period"] = ( + None if asset_grid is None else jnp.asarray(asset_grid, dtype=float) ) additional_continuous_states = { @@ -83,8 +87,13 @@ def check_model_config_and_process(model_config): continuous_states_info["additional_continuous_state_names"] = list( additional_continuous_states.keys() ) + # A name left as `None` here has no default grid at all -- it must be paired + # with a continuous_grid_functions entry (validated in + # process_continuous_grid_functions, which is the first place both are known + # together) and is fully state-choice-specific, with no global fallback array. continuous_states_info["additional_continuous_state_grids"] = { - key: jnp.asarray(value) for key, value in additional_continuous_states.items() + key: (None if value is None else jnp.asarray(value)) + for key, value in additional_continuous_states.items() } continuous_states_info["n_additional_continuous_states"] = len( additional_continuous_states @@ -92,6 +101,39 @@ def check_model_config_and_process(model_config): continuous_states_info["has_additional_continuous_state"] = ( continuous_states_info["n_additional_continuous_states"] > 0 ) + # Names declared as `None` -- their size can only be pinned once a real + # state-choice exists to evaluate the grid function against (see + # continuous_state_grids.py's evaluate_state_specific_continuous_grids, run once + # the state-choice space is built), so n_continuous_state_combinations is left + # unresolved (None) here whenever any of these are present; resolved later and + # merged back into this dict. + continuous_states_info["state_specific_size_names"] = [ + key + for key, value in continuous_states_info[ + "additional_continuous_state_grids" + ].items() + if value is None + ] + if continuous_states_info["state_specific_size_names"]: + # Number of combo points spanned by the additional continuous states' grids + # can't be known yet -- at least one dimension's length is pending. + continuous_states_info["n_continuous_state_combinations"] = None + else: + # Number of combo points spanned by the additional continuous states' default + # grids (1 -- the dummy placeholder -- when there are none). Only the *count* + # is needed here, for sizing solution containers: every state-choice's own + # grid for a given name has this same length by construction (state-specific + # grids may vary in value, not size), regardless of which values it holds. + continuous_states_info["n_continuous_state_combinations"] = int( + np.prod( + [ + len(grid) + for grid in continuous_states_info[ + "additional_continuous_state_grids" + ].values() + ] + ) + ) processed_model_config["continuous_states_info"] = continuous_states_info @@ -129,25 +171,13 @@ def check_model_config_and_process(model_config): if "extra_wealth_grid_factor" in tuning_params else 0.2 ) - tuning_params["n_constrained_points_to_add"] = ( - tuning_params["n_constrained_points_to_add"] - if "n_constrained_points_to_add" in tuning_params - else n_assets_end_of_period // 10 - ) - - if ( - n_assets_end_of_period * (1 + tuning_params["extra_wealth_grid_factor"]) - < n_assets_end_of_period + tuning_params["n_constrained_points_to_add"] - ): - raise ValueError(f"""\n\n - When preparing the tuning parameters for the upper - envelope, we found the following contradicting parameters: \n - The extra wealth grid factor of {tuning_params["extra_wealth_grid_factor"]} is too small - to cover the {tuning_params["n_constrained_points_to_add"]} wealth points which are added in - the credit constrained part of the wealth grid. \n\n""") - tuning_params["n_total_wealth_grid"] = int( - n_assets_end_of_period * (1 + tuning_params["extra_wealth_grid_factor"]) - ) + # Left unresolved when assets_end_of_period is declared as `None`; filled in by + # resolve_tuning_params_from_assets_end_of_period once the grid length is pinned. + if n_assets_end_of_period is not None: + resolve_tuning_params_from_assets_end_of_period( + tuning_params=tuning_params, + n_assets_end_of_period=n_assets_end_of_period, + ) # Set jump threshold to default 2 if it is not given tuning_params["fues_jump_thresh"] = int( @@ -183,32 +213,51 @@ def check_model_config_and_process(model_config): "Specify 'assets_begin_of_period' in model_config['continuous_states'] when using " "the 'druedahl_jorgensen' upper envelope method." ) + assets_begin_of_period_grid = model_config["continuous_states"][ + "assets_begin_of_period" + ] + # `None` when the grid is fully state-choice-specific (paired with a + # continuous_grid_functions entry, validated in + # process_continuous_grid_functions); its length is then pinned later, once + # a real state-choice can be evaluated against. Otherwise store the declared + # array, used to pin the wealth-grid length (see continuous_state_grids.py). + # Either way, the actual per-state-choice wealth grid the Druedahl-Jorgensen + # upper envelope evaluates on -- and that every reader interpolates against + # -- is recomputed on demand from + # continuous_grid_functions["assets_begin_of_period"] (see + # compute_own_dj_wealth_grid), so no shared grid is stored here. processed_model_config["continuous_states_info"]["assets_begin_of_period"] = ( - jnp.asarray(model_config["continuous_states"]["assets_begin_of_period"]) - ) - # The Druedahl-Jorgensen upper envelope always evaluates on this fixed grid - # (see upper_envelope.jax.drued_jorg_jax), so the resulting "endogenous" grid - # is not actually endogenous. We precompute it once here so callers can reuse - # it instead of reading a stored (and redundant) endog_grid array. - processed_model_config["continuous_states_info"]["dj_wealth_grid"] = ( - jnp.concatenate( - ( - jnp.zeros(1), - processed_model_config["continuous_states_info"][ - "assets_begin_of_period" - ], - ) - ) + None + if assets_begin_of_period_grid is None + else jnp.asarray(assets_begin_of_period_grid) ) if upper_envelope["method"] == "fues": - processed_model_config["n_total_wealth_grid"] = tuning_params[ + if "assets_begin_of_period" in model_config["continuous_states"]: + raise ValueError( + "'assets_begin_of_period' is only used by the 'druedahl_jorgensen' " + "upper envelope method. It was found in " + "model_config['continuous_states'] together with " + "upper_envelope['method'] == 'fues', where it has no effect -- " + "either remove it or switch to " + "upper_envelope['method'] = 'druedahl_jorgensen'." + ) + # None while assets_end_of_period's length is still deferred; resolved + # alongside the tuning params it is derived from. + processed_model_config["n_total_wealth_grid"] = tuning_params.get( "n_total_wealth_grid" - ] + ) elif upper_envelope["method"] == "druedahl_jorgensen": - # Expected value at 0, so add 1 + # Expected value at 0, so add 1. None when assets_begin_of_period is + # state-specific (declared as `None`) -- resolved later, once its size can + # be pinned by evaluating the grid function against a real state-choice. + assets_begin_of_period_grid = model_config["continuous_states"][ + "assets_begin_of_period" + ] processed_model_config["n_total_wealth_grid"] = ( - len(model_config["continuous_states"]["assets_begin_of_period"]) + 1 + None + if assets_begin_of_period_grid is None + else len(assets_begin_of_period_grid) + 1 ) else: raise ValueError("Something wrong internally") @@ -220,11 +269,6 @@ def check_model_config_and_process(model_config): upper_envelope["method"] == "druedahl_jorgensen" and len(processed_model_config["choices"]) >= 2 ) - if upper_envelope["skip_endog_grid_storage"]: - assert ( - processed_model_config["continuous_states_info"]["dj_wealth_grid"].shape[0] - == processed_model_config["n_total_wealth_grid"] - ) if "min_period_batch_segments" in model_config.keys(): processed_model_config["min_period_batch_segments"] = model_config[ @@ -289,3 +333,32 @@ def check_model_config_and_process(model_config): processed_model_config["params_check_info"] = {} return processed_model_config + + +def resolve_tuning_params_from_assets_end_of_period( + tuning_params, n_assets_end_of_period +): + """Fill in the fues tuning params that scale with the end-of-period assets grid. + + Called eagerly from ``check_model_config_and_process`` for a declared array, and + from ``merge_resolved_state_specific_lengths`` when ``assets_end_of_period`` is + declared as ``None`` and its length is only pinned once a state-choice exists. + + """ + if "n_constrained_points_to_add" not in tuning_params: + tuning_params["n_constrained_points_to_add"] = n_assets_end_of_period // 10 + + if ( + n_assets_end_of_period * (1 + tuning_params["extra_wealth_grid_factor"]) + < n_assets_end_of_period + tuning_params["n_constrained_points_to_add"] + ): + raise ValueError(f"""\n\n + When preparing the tuning parameters for the upper + envelope, we found the following contradicting parameters: \n + The extra wealth grid factor of {tuning_params["extra_wealth_grid_factor"]} is too small + to cover the {tuning_params["n_constrained_points_to_add"]} wealth points which are added in + the credit constrained part of the wealth grid. \n\n""") + + tuning_params["n_total_wealth_grid"] = int( + n_assets_end_of_period * (1 + tuning_params["extra_wealth_grid_factor"]) + ) diff --git a/src/dcegm/pre_processing/check_params.py b/src/dcegm/pre_processing/check_params.py index 75759424..45d00b66 100644 --- a/src/dcegm/pre_processing/check_params.py +++ b/src/dcegm/pre_processing/check_params.py @@ -10,8 +10,7 @@ def process_params(params, params_check_info) -> Dict[str, float]: and discount factor. Args: - params (dict or tuple or pandas.Series or pandas.DataFrame): Model parameters - Support tuple and list as well? + params dict: Model parameters Returns: dict: Dictionary of model parameters. diff --git a/src/dcegm/pre_processing/model_functions/process_model_functions.py b/src/dcegm/pre_processing/model_functions/process_model_functions.py index 75e5d8d5..9bd67b0d 100644 --- a/src/dcegm/pre_processing/model_functions/process_model_functions.py +++ b/src/dcegm/pre_processing/model_functions/process_model_functions.py @@ -1,3 +1,4 @@ +import inspect from typing import Callable, Dict, Optional import jax @@ -6,9 +7,6 @@ from dcegm.pre_processing.model_functions.taste_shock_function import ( process_shock_functions, ) -from dcegm.pre_processing.model_functions.upper_evelope_wrapper import ( - create_upper_envelope_function, -) from dcegm.pre_processing.model_structure.stochastic_states import ( create_stochastic_transition_function, ) @@ -25,8 +23,9 @@ def process_model_functions_and_extract_info( utility_functions: Dict[str, Callable], utility_functions_final_period: Dict[str, Callable], budget_constraint: Callable, - stochastic_states_transitions: Optional[Dict[str, Callable]] = None, - shock_functions: Optional[Dict[str, Callable]] = None, + stochastic_states_transitions: Optional[Dict[str, Callable]], + shock_functions: Optional[Dict[str, Callable]], + continuous_grid_functions: Optional[Dict[str, Callable]], ): """Create wrapped functions from user supplied functions. @@ -59,10 +58,6 @@ def process_model_functions_and_extract_info( agent's wealth matrices of the next period (t + 1). The inputs ```savings_grid```, ```income_shocks```, ```params``` and ```options``` are already partialled in. - - compute_upper_envelope (Callable): Function for calculating the upper envelope - of the policy and value function. If the number of discrete choices is 1, - this function is a dummy function that returns the policy and value - function as is, without performing a fast upper envelope scan. - transition_function (Callable): Partialled transition function that returns transition probabilities for each state. @@ -80,17 +75,20 @@ def process_model_functions_and_extract_info( compute_utility = determine_function_arguments_and_partial_model_specs( func=utility_functions["utility"], model_specs=model_specs_jax, + not_allowed_state_choices=[], ) compute_marginal_utility = determine_function_arguments_and_partial_model_specs( func=utility_functions["marginal_utility"], model_specs=model_specs_jax, + not_allowed_state_choices=[], ) compute_inverse_marginal_utility = ( determine_function_arguments_and_partial_model_specs( func=utility_functions["inverse_marginal_utility"], model_specs=model_specs_jax, + not_allowed_state_choices=[], ) ) @@ -103,12 +101,14 @@ def process_model_functions_and_extract_info( compute_utility_final = determine_function_arguments_and_partial_model_specs( func=utility_functions_final_period["utility"], model_specs=model_specs_jax, + not_allowed_state_choices=[], ) compute_marginal_utility_final = ( determine_function_arguments_and_partial_model_specs( func=utility_functions_final_period["marginal_utility"], model_specs=model_specs_jax, + not_allowed_state_choices=[], ) ) @@ -142,19 +142,24 @@ def process_model_functions_and_extract_info( has_additional_continuous_states=has_additional_continuous_states, ) + continuous_grid_functions_processed, state_specific_continuous_grid_names = ( + process_continuous_grid_functions( + continuous_grid_functions=continuous_grid_functions, + model_config=model_config, + model_specs=model_specs_jax, + ) + ) + # Budget equation + _check_budget_does_not_declare_renamed_income_shock(budget_constraint) compute_assets_begin_of_period = ( determine_function_arguments_and_partial_model_specs( func=budget_constraint, model_specs=model_specs_jax, + not_allowed_state_choices=[], ) ) - # Upper envelope function - compute_upper_envelope = create_upper_envelope_function( - model_config=model_config, - ) - taste_shock_function_processed, taste_shock_scale_in_params = ( process_shock_functions( shock_functions=shock_functions, @@ -168,6 +173,8 @@ def process_model_functions_and_extract_info( "taste_shock_scale_in_params": taste_shock_scale_in_params } + budget_depends_on_choice = _declares_choice(budget_constraint) + model_funcs = { **utility_functions_processed, **utility_functions_final_period_processed, @@ -178,18 +185,76 @@ def process_model_functions_and_extract_info( "processed_stochastic_funcs": stochastic_transitions_dict, "state_specific_choice_set": state_specific_choice_set, "next_period_deterministic_state": next_period_deterministic_state, - "compute_upper_envelope": compute_upper_envelope, "taste_shock_function": taste_shock_function_processed, + "continuous_grid_functions": continuous_grid_functions_processed, + "state_specific_continuous_grid_names": state_specific_continuous_grid_names, + "budget_depends_on_choice": budget_depends_on_choice, } return model_funcs, model_config_processed +def _declares_choice(func): + """Does this user function take ``choice`` as an argument? + + Inspected on the *user's* function, before + ``determine_function_arguments_and_partial_model_specs`` wraps it (the + wrapper's ``**kwargs`` signature would hide the real parameter names). + + """ + return "choice" in set(inspect.signature(func).parameters) + + +def _check_budget_does_not_declare_renamed_income_shock(budget_constraint): + """Reject the old ``income_shock_previous_period`` argument name. + + The shock is realised in the period whose wealth the budget equation computes, + not in the one before it. The old name described how the shipped example models + happen to be written -- income depending on ``lagged_choice`` -- rather than + anything the framework requires, so it is now passed as ``income_shock``. + Without this check the rename would surface as a bare ``KeyError`` from the + signature filter in ``determine_function_arguments_and_partial_model_specs``. + + """ + if "income_shock_previous_period" in set( + inspect.signature(budget_constraint).parameters + ): + raise ValueError( + f"{budget_constraint.__name__}() declares " + "'income_shock_previous_period', which has been renamed to " + "'income_shock'. The shock is realised in the period whose " + "beginning-of-period wealth the budget equation computes, not in the " + "previous one; rename the argument to 'income_shock'." + ) + + +def _check_choice_not_declared(func, name): + """Reject ``choice`` in a law-of-motion function other than the budget equation. + + ``choice`` in a law of motion means the choice made in the period whose state + is being computed -- i.e. the state is resolved *after* the discrete choice. + Only wealth is modelled that way (see ``_declares_choice``'s callers and the + "pre-consumption wealth" contract in + ``docs/source/development/internals/law_of_motion.rst``); every other state is + fixed at the start of the period, before the choice. The previous period's + choice, which is what accumulation processes like experience actually need, is + available as the ordinary state variable ``lagged_choice``. + + """ + if _declares_choice(func): + raise ValueError( + f"{name}() declares 'choice', which is not supported. Only the budget " + "equation may depend on the current period's choice; every other state " + "is determined before the choice is made. Use 'lagged_choice' if you " + "need the choice made in the previous period." + ) + + def process_state_space_functions( state_space_functions, model_config, model_specs, - additional_continuous_state_names=None, + additional_continuous_state_names, ): state_space_functions = ( @@ -206,9 +271,9 @@ def state_specific_choice_set(**kwargs): return jnp.array(model_config["choices"]) else: - not_allowed_state_choices = ["assets_begin_of_period"] - if additional_continuous_state_names is not None: - not_allowed_state_choices += list(additional_continuous_state_names) + not_allowed_state_choices = ["assets_begin_of_period"] + list( + additional_continuous_state_names + ) state_specific_choice_set = ( determine_function_arguments_and_partial_model_specs( @@ -232,6 +297,7 @@ def next_period_deterministic_state(**kwargs): determine_function_arguments_and_partial_model_specs( func=state_space_functions["next_period_deterministic_state"], model_specs=model_specs, + not_allowed_state_choices=[], ) ) @@ -249,7 +315,9 @@ def next_period_deterministic_state(**kwargs): def process_sparsity_condition(state_space_functions, model_specs): if "sparsity_condition" in state_space_functions.keys(): sparsity_condition = determine_function_arguments_and_partial_model_specs( - func=state_space_functions["sparsity_condition"], model_specs=model_specs + func=state_space_functions["sparsity_condition"], + model_specs=model_specs, + not_allowed_state_choices=[], ) # ToDo: Error if sparsity condition takes second continuous state as input else: @@ -262,9 +330,9 @@ def sparsity_condition(**kwargs): def process_second_continuous_update_function( - state_space_functions=None, - model_specs=None, - has_additional_continuous_states=False, + state_space_functions, + model_specs, + has_additional_continuous_states, ): if has_additional_continuous_states: @@ -275,13 +343,149 @@ def process_second_continuous_update_function( "If additional continuous states are defined, provide " "'next_period_continuous_state' in state_space_functions." ) + _check_choice_not_declared( + state_space_functions["next_period_continuous_state"], + "next_period_continuous_state", + ) next_period_continuous_state = ( determine_function_arguments_and_partial_model_specs( func=state_space_functions["next_period_continuous_state"], model_specs=model_specs, + not_allowed_state_choices=[], ) ) else: next_period_continuous_state = None return next_period_continuous_state + + +def process_continuous_grid_functions( + continuous_grid_functions, model_config, model_specs +): + """Wrap every continuous state's grid into one uniform callable. + + ``continuous_grid_functions`` is an optional user-supplied + ``Dict[str, Callable]``, following the same convention as + ``shock_functions``/``stochastic_states_transitions``: a top-level argument to + ``create_model_dict``, not something nested inside ``model_config`` (which + holds data/config, never functions). Keys are continuous-state names (any of + the additional continuous states, ``assets_end_of_period``, or, if present, + ``assets_begin_of_period``); values are callables + ``(**discrete_state, choice) -> 1d array``, processed the same way as + ``sparsity_condition``/``next_period_deterministic_state`` -- grids live on + the state-choice space (that's where the solution itself lives, see + ``continuous_state_grids.py``), so a grid may depend on ``choice`` too, not + just the discrete state. + + Every continuous state of the model gets a callable in the returned dict, + state-choice-specified or not: names left unspecified keep today's behavior + via a callable that ignores its input and always returns the one global grid + declared in ``model_config["continuous_states"]``. + + These are evaluated on demand wherever a batch needs a state-choice's own + grid during solving, not precomputed for the whole state-choice space, to + avoid materializing an ``(n_state_choices, n_points)`` table. The one + exception is the one-time, NumPy-side consistency check in + ``continuous_state_grids.py``, run once during model-structure construction, + for which this function also returns the list of state-specific names (so + that check can skip names left at their default, which are trivially + consistent). + + """ + continuous_grid_functions = ( + {} if continuous_grid_functions is None else continuous_grid_functions + ) + if not isinstance(continuous_grid_functions, dict): + raise ValueError("continuous_grid_functions must be a dictionary.") + + continuous_states_info = model_config["continuous_states_info"] + default_grids = dict(continuous_states_info["additional_continuous_state_grids"]) + default_grids["assets_end_of_period"] = continuous_states_info[ + "assets_grid_end_of_period" + ] + if "assets_begin_of_period" in continuous_states_info: + default_grids["assets_begin_of_period"] = continuous_states_info[ + "assets_begin_of_period" + ] + + for name, grid_func in continuous_grid_functions.items(): + if name not in default_grids: + raise ValueError( + f"continuous_grid_functions contains the key '{name}', which is " + f"not a continuous state of this model. Valid keys are: " + f"{sorted(default_grids)}." + ) + if not callable(grid_func): + raise ValueError( + f"continuous_grid_functions['{name}'] must be a callable of the " + "form (**discrete_state) -> 1d array." + ) + if default_grids[name] is not None: + raise ValueError( + f"continuous_grid_functions['{name}'] is given, but " + f"model_config['continuous_states']['{name}'] is not None. A " + f"declared array is unused once a continuous_grid_functions entry " + f"takes over -- set model_config['continuous_states']['{name}'] to " + "None to make that explicit." + ) + + if ( + "assets_begin_of_period" in continuous_grid_functions + and continuous_states_info["n_additional_continuous_states"] > 0 + and not model_config["upper_envelope"]["skip_endog_grid_storage"] + ): + raise ValueError( + "continuous_grid_functions['assets_begin_of_period'] together with " + "additional continuous states requires skip_endog_grid_storage " + "(upper_envelope['method'] == 'druedahl_jorgensen' and at least two " + "choices). With a single choice the Druedahl-Jorgensen upper envelope " + "is skipped entirely (see check_model_config.py), so the stored " + "endogenous grid is not the fixed assets_begin_of_period grid there, " + "and a state-specific assets_begin_of_period has nothing to plug " + "into. Remove either the additional continuous state(s) or the " + "continuous_grid_functions['assets_begin_of_period'] entry, or add a " + "second choice." + ) + + continuous_grid_functions_processed = {} + for name, default_grid in default_grids.items(): + if name in continuous_grid_functions: + continuous_grid_functions_processed[name] = ( + determine_function_arguments_and_partial_model_specs( + func=continuous_grid_functions[name], + model_specs=model_specs, + not_allowed_state_choices=[], + ) + ) + elif default_grid is None: + raise ValueError( + f"model_config['continuous_states']['{name}'] is None, but no " + f"matching continuous_grid_functions['{name}'] was given. A `None` " + "grid means there is no default to fall back to -- it must be " + "paired with a continuous_grid_functions entry." + ) + else: + continuous_grid_functions_processed[name] = _make_constant_grid_func( + default_grid + ) + + return continuous_grid_functions_processed, list(continuous_grid_functions.keys()) + + +def _make_constant_grid_func(default_grid): + """Return a callable that ignores its input and always returns ``default_grid``. + + A factory, rather than defining the closure directly in the loop above, to avoid + Python's late-binding-closure-in-a-loop pitfall: a closure defined inside a loop + body captures the loop *variable*, not its value at that iteration, so every closure + created across iterations would end up returning whatever ``default_grid`` was left + holding after the *last* iteration. Binding it as this function's own parameter + forces a fresh value per call. + + """ + + def constant_grid_func(**kwargs): + return default_grid + + return constant_grid_func diff --git a/src/dcegm/pre_processing/model_functions/upper_evelope_wrapper.py b/src/dcegm/pre_processing/model_functions/upper_evelope_wrapper.py index 1833d1d2..f8e30418 100644 --- a/src/dcegm/pre_processing/model_functions/upper_evelope_wrapper.py +++ b/src/dcegm/pre_processing/model_functions/upper_evelope_wrapper.py @@ -4,6 +4,7 @@ def create_upper_envelope_function( model_config, + continuous_grid_functions, ): if len(model_config["choices"]) < 2: return no_upper_envelope_dummy_function @@ -63,6 +64,13 @@ def value_function( ) elif method == "druedahl_jorgensen": + # This state-choice's own assets_begin_of_period grid, evaluated on + # demand -- self-referential, this state-choice is solving/storing + # against its own common wealth grid (no parent/child ambiguity, unlike + # law_of_motion.py's grid selection for a transition *into* a state). + m_grid = continuous_grid_functions["assets_begin_of_period"]( + **state_choice_dict + ) return drued_jorg_jax( endog_grid=endog_grid, policy=policy, @@ -70,7 +78,7 @@ def value_function( expected_value_zero_savings=expected_value_zero_assets, value_function=value_function, value_function_kwargs=value_kwargs, - m_grid=model_config["continuous_states_info"]["assets_begin_of_period"], + m_grid=m_grid, ) else: diff --git a/src/dcegm/pre_processing/model_structure/continuous_state_grids.py b/src/dcegm/pre_processing/model_structure/continuous_state_grids.py new file mode 100644 index 00000000..216de161 --- /dev/null +++ b/src/dcegm/pre_processing/model_structure/continuous_state_grids.py @@ -0,0 +1,311 @@ +"""State-specific continuous state grids. + +Grids live on the *state-choice* space, not the bare state space: that's where +the solution itself lives (``value_solved``/``policy_solved``/``endog_grid_solved`` +are indexed by state-choice, see ``create_solution_container``), so a grid may +depend on the discrete state *and* on ``choice``. + +""" + +import numpy as np + +from dcegm.pre_processing.check_model_config import ( + resolve_tuning_params_from_assets_end_of_period, +) + + +def evaluate_state_specific_continuous_grids( + state_choice_space, + discrete_state_choice_names, + continuous_grid_functions, + state_specific_names, + continuous_states_info, +): + """Evaluate each user-supplied continuous grid function once per state-choice. + + Only evaluates ``state_specific_names`` -- the continuous states the user + actually made state-specific via the ``continuous_grid_functions`` argument to + ``create_model_dict`` (see ``process_continuous_grid_functions``); names left at + their default (constant) grid are skipped, since they are trivially consistent + across every state-choice and don't need checking. + + Also validates that every state-choice's grid has the expected length for that + name. For names with a declared default array in + ``model_config["continuous_states"]``, that's the declared length. For names + declared as ``None`` (no default array at all -- fully state-choice-specific), + there is nothing to read a length from yet, so it is *pinned* here instead: the + first state-choice's grid_func evaluation fixes the length, and every other + state-choice's grid is validated against that. Either way, storage containers + are rectangular (same number of grid points for every state-choice, only the + values differ), so a mismatched length would otherwise silently corrupt shapes + downstream instead of failing here, close to the user-supplied function that + caused it. + + This is a one-time NumPy-side pass over the full state-choice space, done once + during model-structure construction. The grids are used to build the + child-sharing consistency check (see + ``check_continuous_grid_consistency_across_shared_children`` below) and are not + stored in ``model_structure`` or threaded through solving. The pinned lengths + (for names declared as ``None``) *are* returned, for the caller to merge back + into ``continuous_states_info`` (e.g. ``n_continuous_state_combinations``, + ``n_total_wealth_grid``), since those couldn't be computed any earlier than + this, the first point a real state-choice exists to evaluate against. + + """ + if len(state_specific_names) == 0: + return {}, {} + + expected_lengths = _expected_grid_lengths(continuous_states_info) + resolved_lengths = {} + + grids_per_state_choice = {} + for name in state_specific_names: + grid_func = continuous_grid_functions[name] + pinning_size = name not in expected_lengths + + grid_rows = [] + for row in state_choice_space: + state_choice_dict = { + key: row[i] for i, key in enumerate(discrete_state_choice_names) + } + grid = np.asarray(grid_func(**state_choice_dict)) + + if pinning_size: + expected_lengths[name] = grid.shape[0] if grid.ndim == 1 else -1 + resolved_lengths[name] = expected_lengths[name] + pinning_size = False + + expected_length = expected_lengths[name] + if grid.ndim != 1 or grid.shape[0] != expected_length: + length_source = ( + f"the length of the first state-choice evaluated, since " + f"model_config['continuous_states']['{name}'] is None" + if name in resolved_lengths + else "the length declared in " + f"model_config['continuous_states']['{name}']" + ) + raise ValueError( + f"\n\nThe continuous grid function for '{name}' returned an " + f"array of shape {grid.shape} for the state-choice\n\n" + f"{state_choice_dict}\n\nbut every state-choice's grid for " + f"'{name}' must be a 1d array of length {expected_length} " + f"({length_source}). All state-choices must use grids of the " + "same size for a given continuous state; only the grid values " + "may vary." + ) + grid_rows.append(grid) + + grids_per_state_choice[name] = np.stack(grid_rows, axis=0) + + return grids_per_state_choice, resolved_lengths + + +def _expected_grid_lengths(continuous_states_info): + expected_lengths = { + name: len(grid) + for name, grid in continuous_states_info[ + "additional_continuous_state_grids" + ].items() + if grid is not None + } + if continuous_states_info["assets_grid_end_of_period"] is not None: + expected_lengths["assets_end_of_period"] = len( + continuous_states_info["assets_grid_end_of_period"] + ) + if continuous_states_info.get("assets_begin_of_period") is not None: + expected_lengths["assets_begin_of_period"] = len( + continuous_states_info["assets_begin_of_period"] + ) + return expected_lengths + + +def pin_state_specific_lengths_from_first_state_choice( + state_choice_space, + discrete_state_choice_names, + continuous_grid_functions, + state_specific_names, + continuous_states_info, +): + """Pin the lengths of grids declared as ``None`` from a single state-choice. + + Used on the load path (``load_model_dict``), where ``model_config`` is rebuilt + from the raw user config -- leaving those lengths unresolved again, see + ``check_model_config.py`` -- while the state-choice space comes back from the + pickle, so ``evaluate_state_specific_continuous_grids`` never runs. Every + state-choice's grid for a given name was already validated to have the same + length when the model was built and saved, so one state-choice suffices to + recover it here. + + """ + expected_lengths = _expected_grid_lengths(continuous_states_info) + deferred_names = [ + name for name in state_specific_names if name not in expected_lengths + ] + if len(deferred_names) == 0: + return {} + + first_state_choice = np.asarray(state_choice_space[0]) + state_choice_dict = { + key: first_state_choice[i] for i, key in enumerate(discrete_state_choice_names) + } + return { + name: np.asarray(continuous_grid_functions[name](**state_choice_dict)).shape[0] + for name in deferred_names + } + + +def merge_resolved_state_specific_lengths( + model_config, resolved_state_specific_lengths +): + """Merge grid lengths pinned by evaluation back into ``model_config``. + + Names declared as ``None`` in ``model_config["continuous_states"]`` have no default + array to compute ``n_continuous_state_combinations``/ ``n_total_wealth_grid`` from + ahead of time (see ``check_model_config.py``, where both are left ``None`` for this + reason) -- this fills them in now that the lengths have been pinned against a real + state-choice, the first point either becomes knowable. + + """ + if not resolved_state_specific_lengths: + return + + continuous_states_info = model_config["continuous_states_info"] + + if continuous_states_info["n_continuous_state_combinations"] is None: + lengths = [ + ( + resolved_state_specific_lengths[name] + if name in continuous_states_info["state_specific_size_names"] + else len( + continuous_states_info["additional_continuous_state_grids"][name] + ) + ) + for name in continuous_states_info["additional_continuous_state_names"] + ] + continuous_states_info["n_continuous_state_combinations"] = int( + np.prod(lengths) + ) + + if "assets_end_of_period" in resolved_state_specific_lengths: + continuous_states_info["n_assets_end_of_period"] = ( + resolved_state_specific_lengths["assets_end_of_period"] + ) + # The fues tuning params scale with this length, so they could not be + # resolved in check_model_config.py either. + resolve_tuning_params_from_assets_end_of_period( + tuning_params=model_config["upper_envelope"]["tuning_params"], + n_assets_end_of_period=resolved_state_specific_lengths[ + "assets_end_of_period" + ], + ) + + if model_config["n_total_wealth_grid"] is None: + if "assets_begin_of_period" in resolved_state_specific_lengths: + # Druedahl-Jorgensen stores on the begin-of-period assets grid, plus the + # expected value at zero assets. + model_config["n_total_wealth_grid"] = ( + resolved_state_specific_lengths["assets_begin_of_period"] + 1 + ) + elif model_config["upper_envelope"]["method"] == "fues": + model_config["n_total_wealth_grid"] = model_config["upper_envelope"][ + "tuning_params" + ]["n_total_wealth_grid"] + + +def check_continuous_grid_consistency_across_shared_children( + state_choice_space, + discrete_states_names, + map_state_choice_to_child_states, + grids_per_state_choice, +): + """Check that state-choices sharing a child state agree on their own grid. + + The batch-creation code (``pre_processing/batches/``) deduplicates children purely + by discrete-state index (``np.unique`` on ``map_state_choice_to_child_states``) and + computes the continuation-value interpolation once per unique child, reusing it for + every parent state-choice that maps to it. That shared computation is fed the + *parent's* own continuous grid (as "last period's grid", see ``law_of_motion.py``). + If two parent state-choices that share a child disagreed on their own grid, the + shared computation could only use one of them, silently corrupting the other's + continuation values. + + Grids live on the state-choice space (see + ``evaluate_state_specific_continuous_grids`` above), so each row's "own grid" is + available directly -- no collapsing to the parent's bare state needed, since + ``map_state_choice_to_child_states`` is already indexed by state-choice row (row + ``i`` of it is row ``i`` of ``state_choice_space``). + + We reuse the exact same grouping mechanism the batch code uses (``np.unique`` with + ``return_inverse``) so this check provably guards the invariant that optimization + depends on, rather than an approximation of it. + + """ + if len(grids_per_state_choice) == 0: + return + + invalid_idx = np.iinfo(map_state_choice_to_child_states.dtype).max + n_stochastic_states = map_state_choice_to_child_states.shape[1] + + for name, grid_of_state_choice in grids_per_state_choice.items(): + for k in range(n_stochastic_states): + child_idx_col = np.asarray(map_state_choice_to_child_states[:, k]) + valid_mask = child_idx_col != invalid_idx + + if not np.any(valid_mask): + continue + + row_idxs = np.flatnonzero(valid_mask) + child_idx_valid = child_idx_col[row_idxs] + own_grid_valid = grid_of_state_choice[row_idxs] + + _, first_occurrence, inverse = np.unique( + child_idx_valid, return_index=True, return_inverse=True + ) + representative_grid = own_grid_valid[first_occurrence][inverse] + + mismatch = np.any(own_grid_valid != representative_grid, axis=1) + if np.any(mismatch): + bad_local_idx = np.flatnonzero(mismatch)[0] + bad_row = row_idxs[bad_local_idx] + rep_row = row_idxs[first_occurrence[inverse[bad_local_idx]]] + child_state_idx = child_idx_valid[bad_local_idx] + + _raise_grid_mismatch_error( + name=name, + state_choice_space=state_choice_space, + discrete_states_names=discrete_states_names, + bad_row=bad_row, + rep_row=rep_row, + child_state_idx=child_state_idx, + ) + + +def _raise_grid_mismatch_error( + name, + state_choice_space, + discrete_states_names, + bad_row, + rep_row, + child_state_idx, +): + def _state_choice_dict(row_idx): + row = state_choice_space[row_idx] + example_dict = {key: row[i] for i, key in enumerate(discrete_states_names)} + example_dict["choice"] = row[-1] + return example_dict + + raise ValueError( + f"\n\nThe continuous state '{name}' is declared state-specific via " + "the continuous_grid_functions argument to create_model_dict, but the " + "following two state-choice combinations transition to the same child " + "state " + f"(child state index {child_state_idx}) while using different grids for " + f"'{name}':\n\n{_state_choice_dict(bad_row)}\n\nand\n\n" + f"{_state_choice_dict(rep_row)}\n\n" + "The grid for a continuous state may only depend on discrete-state " + "variables and/or the choice, in a way that agrees across every " + "state-choice combination transitioning to the same child -- e.g. " + "time-invariant states like sex/education, the period, or a choice that " + "passes through unchanged (or 1:1) from parent to child. Please revisit " + f"the grid function for '{name}' or the state space definition." + ) diff --git a/src/dcegm/pre_processing/model_structure/model_structure.py b/src/dcegm/pre_processing/model_structure/model_structure.py index eb24180b..c5552b87 100644 --- a/src/dcegm/pre_processing/model_structure/model_structure.py +++ b/src/dcegm/pre_processing/model_structure/model_structure.py @@ -26,36 +26,13 @@ def create_model_structure( - "transform_between_state_and_state_choice_vec" (callable) """ - # Create continuous state space - continuous_states_info = model_config["continuous_states_info"] - additional_continuous_state_grids = continuous_states_info.get( - "additional_continuous_state_grids", {} - ) - - if continuous_states_info["has_additional_continuous_state"]: - continuous_state_names = continuous_states_info[ - "additional_continuous_state_names" - ] - continuous_grids = [ - additional_continuous_state_grids[name] for name in continuous_state_names - ] - - continuous_state_mesh = jnp.meshgrid(*continuous_grids, indexing="ij") - continuous_state_space = { - name: grid.ravel() - for name, grid in zip(continuous_state_names, continuous_state_mesh) - } - else: - continuous_state_space = {"dummy_cont": jnp.zeros(1)} - - print("Starting state space creation") state_space_objects = create_state_space( model_config=model_config, sparsity_condition=model_funcs["sparsity_condition"], debugging=False, ) - print("State space created.\n") - print("Starting state-choice space creation and child state mapping.") + n_states = state_space_objects["state_space"].shape[0] + print(f" discrete state space: {n_states:,} states") state_choice_and_child_state_objects = ( create_state_choice_space_and_child_state_mapping( @@ -65,15 +42,23 @@ def create_model_structure( "next_period_deterministic_state" ], state_space_arrays=state_space_objects, + continuous_grid_functions=model_funcs["continuous_grid_functions"], + state_specific_continuous_grid_names=model_funcs[ + "state_specific_continuous_grid_names" + ], ) ) + n_state_choices = state_choice_and_child_state_objects["state_choice_space"].shape[ + 0 + ] + print(f" discrete state-choice space: {n_state_choices:,} state-choices") + state_space_objects.pop("map_state_to_index_with_proxy") - state_space_objects.pop("state_space_incl_proxies") + state_space_objects.pop("map_state_to_index_incl_proxies") model_structure = { **state_space_objects, **state_choice_and_child_state_objects, "choice_range": jnp.asarray(model_config["choices"]), - "continuous_state_space": continuous_state_space, } return jax.tree.map(create_array_with_smallest_int_dtype, model_structure) diff --git a/src/dcegm/pre_processing/model_structure/shared.py b/src/dcegm/pre_processing/model_structure/shared.py index 8fdb7d66..dc79b247 100644 --- a/src/dcegm/pre_processing/model_structure/shared.py +++ b/src/dcegm/pre_processing/model_structure/shared.py @@ -3,16 +3,13 @@ from dcegm.pre_processing.shared import get_smallest_int_type -def create_indexer_for_space(space, max_var_values=None): +def create_indexer_for_space(space, max_var_values): """Create indexer for space.""" # Indexer has always unsigned data type with integers starting at zero # Leave one additional value for the invalid number data_type = get_smallest_int_type(space.shape[0] + 1) max_value = np.iinfo(data_type).max - if max_var_values is None: - max_var_values = np.max(space, axis=0) - map_vars_to_index = np.full( max_var_values + 1, fill_value=max_value, dtype=data_type ) diff --git a/src/dcegm/pre_processing/model_structure/state_choice_space.py b/src/dcegm/pre_processing/model_structure/state_choice_space.py index 8116e0c5..88f09256 100644 --- a/src/dcegm/pre_processing/model_structure/state_choice_space.py +++ b/src/dcegm/pre_processing/model_structure/state_choice_space.py @@ -2,6 +2,11 @@ import numpy as np +from dcegm.pre_processing.model_structure.continuous_state_grids import ( + check_continuous_grid_consistency_across_shared_children, + evaluate_state_specific_continuous_grids, + merge_resolved_state_specific_lengths, +) from dcegm.pre_processing.model_structure.shared import create_indexer_for_space from dcegm.pre_processing.shared import get_smallest_int_type @@ -11,6 +16,8 @@ def create_state_choice_space_and_child_state_mapping( state_specific_choice_set, next_period_deterministic_state, state_space_arrays, + continuous_grid_functions, + state_specific_continuous_grid_names, ): """Create state choice space of all feasible state-choice combinations. @@ -66,6 +73,9 @@ def create_state_choice_space_and_child_state_mapping( stochastic_state_space = state_space_arrays["stochastic_state_space"] map_state_to_index_with_proxy = state_space_arrays["map_state_to_index_with_proxy"] map_state_to_index = state_space_arrays["map_state_to_index"] + map_state_to_index_incl_proxies = state_space_arrays[ + "map_state_to_index_incl_proxies" + ] state_space = state_space_arrays["state_space"] n_states, n_state_and_stochastic_variables = state_space.shape @@ -104,6 +114,16 @@ def create_state_choice_space_and_child_state_mapping( dtype=state_space_indexer_dtype, ) + # Parallel to the (proxied) child map above, but with each actual child's own + # distinct index (into state_space_incl_proxies). Used only for the + # grid-consistency check below, then discarded -- it is not returned or threaded + # into solving, so it adds no persistent memory. See map_state_to_index_incl_proxies. + map_state_choice_to_child_states_actual = np.full( + (n_states * n_choices, n_stochastic_states), + fill_value=np.iinfo(map_state_to_index_incl_proxies.dtype).max, + dtype=map_state_to_index_incl_proxies.dtype, + ) + stochastic_states_tuple = tuple( stochastic_state_space[:, i] for i in range(n_stochastic_vars) ) @@ -188,6 +208,9 @@ def create_state_choice_space_and_child_state_mapping( ) map_state_choice_to_child_states[idx, :] = child_idxs + map_state_choice_to_child_states_actual[idx, :] = ( + map_state_to_index_incl_proxies[states_next_tuple] + ) idx += 1 @@ -195,8 +218,13 @@ def create_state_choice_space_and_child_state_mapping( state_choice_space = state_choice_space_raw[:idx] map_state_choice_to_parent_state = map_state_choice_to_parent_state[:idx] map_state_choice_to_child_states = map_state_choice_to_child_states[:idx, :] + map_state_choice_to_child_states_actual = map_state_choice_to_child_states_actual[ + :idx, : + ] - map_state_choice_to_index, _ = create_indexer_for_space(state_choice_space) + map_state_choice_to_index, _ = create_indexer_for_space( + state_choice_space, max_var_values=np.max(state_choice_space, axis=0) + ) # Create indexer with proxy state_space_incl_proxies = state_space_arrays["state_space_incl_proxies"] @@ -223,6 +251,39 @@ def create_state_choice_space_and_child_state_mapping( for i, key in enumerate(discrete_states_names + ["choice"]) } + # Grids live on the state-choice space (that's where the solution itself + # lives, see continuous_state_grids.py), so this can only be evaluated once + # state_choice_space itself is built -- computed here rather than passed in + # from model_structure.py. + grids_per_state_choice, resolved_state_specific_lengths = ( + evaluate_state_specific_continuous_grids( + state_choice_space=state_choice_space, + discrete_state_choice_names=discrete_states_names + ["choice"], + continuous_grid_functions=continuous_grid_functions or {}, + state_specific_names=state_specific_continuous_grid_names or [], + continuous_states_info=model_config["continuous_states_info"], + ) + ) + merge_resolved_state_specific_lengths( + model_config=model_config, + resolved_state_specific_lengths=resolved_state_specific_lengths, + ) + + # Run the grid-consistency check against the *actual* child indices (each child + # its own, proxies not collapsed), not the proxied map used for solving. The + # proxy is a solution-reuse pointer: parents that transition to genuinely + # different children (e.g. death at different ages, all proxied to one + # last-period slot) must not be forced to share a grid. The proxied child's + # value is independent of the second continuous state (e.g. bequest depends on + # wealth, not experience), so reusing one representative parent's grid for it in + # the solve stays correct; see continuous_state_grids.py. + check_continuous_grid_consistency_across_shared_children( + state_choice_space=state_choice_space, + discrete_states_names=discrete_states_names, + map_state_choice_to_child_states=map_state_choice_to_child_states_actual, + grids_per_state_choice=grids_per_state_choice, + ) + test_child_state_mapping( model_config=model_config, state_choice_space=state_choice_space, @@ -238,6 +299,9 @@ def create_state_choice_space_and_child_state_mapping( "map_state_choice_to_index_with_proxy": map_state_choice_to_index_with_proxy, "map_state_choice_to_parent_state": map_state_choice_to_parent_state, "map_state_choice_to_child_states": map_state_choice_to_child_states, + "map_state_choice_to_child_states_actual": ( + map_state_choice_to_child_states_actual + ), } return dict_of_state_choice_space_objects diff --git a/src/dcegm/pre_processing/model_structure/state_space.py b/src/dcegm/pre_processing/model_structure/state_space.py index a0ff3f86..1e6c029b 100644 --- a/src/dcegm/pre_processing/model_structure/state_space.py +++ b/src/dcegm/pre_processing/model_structure/state_space.py @@ -13,7 +13,7 @@ from dcegm.pre_processing.shared import create_array_with_smallest_int_dtype -def create_state_space(model_config, sparsity_condition, debugging=False): +def create_state_space(model_config, sparsity_condition, debugging): """Create state space object and indexer. We need to add the convention for the state space objects. @@ -184,6 +184,16 @@ def create_state_space(model_config, sparsity_condition, debugging=False): else: map_state_to_index_with_proxy = map_state_to_index + # Indexer into ``state_space_incl_proxies`` (valid + proxied states), giving each + # actual child state a distinct index -- unlike ``map_state_to_index_with_proxy``, + # which collapses proxied states onto their proxy target. Used only for the + # grid-consistency check, so parents that transition to genuinely different child + # states (e.g. death at different ages, all proxied to one last-period state) are + # not forced to share a grid. The solve keeps using the proxy indexer. + map_state_to_index_incl_proxies, _ = create_indexer_for_space( + state_space_incl_proxies, max_var_values=max_values_unrestricted + ) + state_space_dict = { key: create_array_with_smallest_int_dtype(state_space[:, i]) for i, key in enumerate(discrete_states_names) @@ -199,6 +209,7 @@ def create_state_space(model_config, sparsity_condition, debugging=False): "state_space_dict": state_space_dict, "map_state_to_index": map_state_to_index, "map_state_to_index_with_proxy": map_state_to_index_with_proxy, + "map_state_to_index_incl_proxies": map_state_to_index_incl_proxies, "stochastic_state_space": stochastic_state_space, "stochastic_states_names": stochastic_state_names, "state_names_without_stochastic": state_names_without_stochastic, @@ -218,7 +229,9 @@ def create_state_space(model_config, sparsity_condition, debugging=False): array_of_states_proxied_to[:, i] for i in range(array_of_states_proxied_to.shape[1]) ) - full_indexer, _ = create_indexer_for_space(state_space_full) + full_indexer, _ = create_indexer_for_space( + state_space_full, max_var_values=max_values_unrestricted + ) idxs_proxied_to = full_indexer[tuple_of_states_proxied_from] debug_df["idxs_proxied_to"] = -9999 debug_df.loc[debug_df["is_proxied"], "idxs_proxied_to"] = idxs_proxied_to diff --git a/src/dcegm/pre_processing/model_structure/stochastic_states.py b/src/dcegm/pre_processing/model_structure/stochastic_states.py index bd631519..92b60b76 100644 --- a/src/dcegm/pre_processing/model_structure/stochastic_states.py +++ b/src/dcegm/pre_processing/model_structure/stochastic_states.py @@ -61,6 +61,7 @@ def process_stochastic_transitions( processed_exog_func = determine_function_arguments_and_partial_model_specs( func=func, model_specs=model_specs, + not_allowed_state_choices=[], ) func_list += [processed_exog_func] func_dict[name] = processed_exog_func @@ -119,7 +120,7 @@ def process_stochastic_model_specifications(model_config): def create_sparse_stochastic_trans_map( - model_structure, model_funcs, model_config_processed, from_saved=False + model_structure, model_funcs, model_config_processed, from_saved ): """Create sparse mapping from state-choice to stochastic states.""" state_choice_dict = model_structure["state_choice_space_dict"] @@ -222,23 +223,35 @@ def kronecker_with_index(idx): trans_vector = jnp.kron(trans_vector, func(idx)) return trans_vector - all_transitions = jax.vmap(kronecker_with_index)( - jnp.arange(n_state_choices), + # Only the NaN pattern is used below, so evaluate in blocks and keep booleans. + block_size = 10_000 + keep_mask = np.concatenate( + [ + ~np.isnan( + np.asarray( + jax.vmap(kronecker_with_index)( + jnp.arange(start, min(start + block_size, n_state_choices)) + ) + ) + ) + for start in range(0, n_state_choices, block_size) + ] ) - # Find non-NaN positions (states to keep) - keep_mask = ~np.isnan(all_transitions) - # Select non-NaN child states directly sparse_child_states_mapping = model_structure["map_state_choice_to_child_states"][ keep_mask ].reshape(keep_mask.shape[0], -1) + sparse_child_states_mapping_actual = model_structure[ + "map_state_choice_to_child_states_actual" + ][keep_mask].reshape(keep_mask.shape[0], -1) state_choice_dict_with_idx = { **state_choice_dict, "index": jnp.arange(n_state_choices), } return ( sparse_child_states_mapping, + sparse_child_states_mapping_actual, state_choice_dict_with_idx, compute_stochastic_transition_vec, trans_func_dict, diff --git a/src/dcegm/pre_processing/setup_model.py b/src/dcegm/pre_processing/setup_model.py index b8c54547..2aa7b944 100644 --- a/src/dcegm/pre_processing/setup_model.py +++ b/src/dcegm/pre_processing/setup_model.py @@ -4,13 +4,23 @@ import jax import jax.numpy as jnp -from dcegm.pre_processing.batches.batch_creation import create_batches_and_information +from dcegm.pre_processing.batches.batch_creation import ( + bundle_law_of_motion_arrays, + create_batches_and_information, +) from dcegm.pre_processing.check_model_config import check_model_config_and_process from dcegm.pre_processing.check_model_specs import extract_model_specs_info from dcegm.pre_processing.model_functions.process_model_functions import ( process_model_functions_and_extract_info, process_sparsity_condition, ) +from dcegm.pre_processing.model_functions.upper_evelope_wrapper import ( + create_upper_envelope_function, +) +from dcegm.pre_processing.model_structure.continuous_state_grids import ( + merge_resolved_state_specific_lengths, + pin_state_specific_lengths_from_first_state_choice, +) from dcegm.pre_processing.model_structure.model_structure import create_model_structure from dcegm.pre_processing.model_structure.state_space import create_state_space from dcegm.pre_processing.model_structure.stochastic_states import ( @@ -32,6 +42,7 @@ def create_model_dict( state_space_functions: Dict[str, Callable] = None, stochastic_states_transitions: Dict[str, Callable] = None, shock_functions: Dict[str, Callable] = None, + continuous_grid_functions: Dict[str, Callable] = None, debug_info: str = None, use_stochastic_sparsity=False, ): @@ -81,6 +92,7 @@ def create_model_dict( budget_constraint=budget_constraint, stochastic_states_transitions=stochastic_states_transitions, shock_functions=shock_functions, + continuous_grid_functions=continuous_grid_functions, ) specs_read_funcs, specs_params_info = extract_model_specs_info(model_specs) @@ -91,17 +103,29 @@ def create_model_dict( **specs_params_info, } + print("Setting up model:") model_structure = create_model_structure( model_config=model_config_processed, model_funcs=model_funcs, ) + # Only now are the continuous grids' sizes known: a grid declared as `None` is + # pinned by evaluating its grid function against a real state-choice, which + # create_model_structure above is the first thing to have. The upper envelope + # reads those sizes, so it is built here rather than alongside the other model + # functions. + model_funcs["compute_upper_envelope"] = create_upper_envelope_function( + model_config=model_config_processed, + continuous_grid_functions=model_funcs["continuous_grid_functions"], + ) + if use_stochastic_sparsity: n_stochastic_original = model_structure[ "map_state_choice_to_child_states" ].shape[1] ( model_structure["map_state_choice_to_child_states"], + model_structure["map_state_choice_to_child_states_actual"], model_structure["state_choice_space_dict"], model_funcs["compute_stochastic_transition_vec"], model_funcs["sparse_processed_stochastic_funcs"], @@ -113,7 +137,8 @@ def create_model_dict( ) n_sparse = model_structure["map_state_choice_to_child_states"].shape[1] print( - f"Stochastic transition mapping sparsified from {n_stochastic_original} to {n_sparse} " + f" stochastic transitions to children: {n_sparse:,} per state-choice " + f"(sparsified from {n_stochastic_original:,})" ) model_funcs["stochastic_state_mapping"] = create_stochastic_state_mapping( @@ -121,23 +146,30 @@ def create_model_dict( model_structure["stochastic_states_names"], ) - print("State, state-choice and child state mapping created.\n") - print("Start creating batches for the model.") - batch_info = create_batches_and_information( model_structure=model_structure, n_periods=model_config_processed["n_periods"], min_period_batch_segments=model_config_processed["min_period_batch_segments"], batch_mode=model_config_processed["batch_mode"], ) + # Fold each batch's per-branch law-of-motion index arrays into the single + # dict calc_law_of_motion reads, keeping only the branch this model takes. + # Model-static, so it is decided here at setup rather than per batch/period. + batch_info = bundle_law_of_motion_arrays( + batch_info, + transition_depends_on_choice=model_funcs["budget_depends_on_choice"], + state_space_dict=model_structure["state_space_dict"], + ) if not debug_info == "all": # Delete large arrays which is not needed. Not if all is requested # by the debug string. model_structure.pop("map_state_choice_to_child_states") model_structure.pop("map_state_choice_to_index") + model_structure.pop("map_state_choice_to_child_states_actual") + model_structure.pop("state_space_incl_proxies") batch_info = jax.tree.map(create_array_with_smallest_int_dtype, batch_info) - print("Model setup complete.\n") + print("Model setup complete.") return { "model_config": model_config_processed, "model_funcs": model_funcs, @@ -156,6 +188,7 @@ def create_model_dict_and_save( state_space_functions: Dict[str, Callable] = None, stochastic_states_transitions: Dict[str, Callable] = None, shock_functions: Dict[str, Callable] = None, + continuous_grid_functions: Dict[str, Callable] = None, path: str = "model.pkl", debug_info=None, use_stochastic_sparsity=False, @@ -176,6 +209,7 @@ def create_model_dict_and_save( state_space_functions=state_space_functions, stochastic_states_transitions=stochastic_states_transitions, shock_functions=shock_functions, + continuous_grid_functions=continuous_grid_functions, debug_info=debug_info, use_stochastic_sparsity=use_stochastic_sparsity, ) @@ -198,6 +232,7 @@ def load_model_dict( state_space_functions: Dict[str, Callable] = None, stochastic_states_transitions: Dict[str, Callable] = None, shock_functions: Dict[str, Callable] = None, + continuous_grid_functions: Dict[str, Callable] = None, path: str = "model.pkl", use_stochastic_sparsity=False, ): @@ -216,9 +251,38 @@ def load_model_dict( budget_constraint=budget_constraint, stochastic_states_transitions=stochastic_states_transitions, shock_functions=shock_functions, + continuous_grid_functions=continuous_grid_functions, ) ) + # model_config is rebuilt from the raw user config above, so the sizes of + # continuous states declared as None are unresolved again -- and the state-choice + # space that pins them comes from the pickle instead of being rebuilt, so nothing + # on this path would otherwise fill them in. + merge_resolved_state_specific_lengths( + model_config=model["model_config"], + resolved_state_specific_lengths=( + pin_state_specific_lengths_from_first_state_choice( + state_choice_space=model["model_structure"]["state_choice_space"], + discrete_state_choice_names=( + list(model["model_structure"]["discrete_states_names"]) + ["choice"] + ), + continuous_grid_functions=model["model_funcs"][ + "continuous_grid_functions" + ], + state_specific_names=model["model_funcs"][ + "state_specific_continuous_grid_names" + ], + continuous_states_info=model["model_config"]["continuous_states_info"], + ) + ), + ) + + model["model_funcs"]["compute_upper_envelope"] = create_upper_envelope_function( + model_config=model["model_config"], + continuous_grid_functions=model["model_funcs"]["continuous_grid_functions"], + ) + specs_read_funcs, specs_params_info = extract_model_specs_info(model_specs) model["model_funcs"]["read_funcs"] = specs_read_funcs diff --git a/src/dcegm/pre_processing/shared.py b/src/dcegm/pre_processing/shared.py index 9e8e4f4f..827c7b28 100644 --- a/src/dcegm/pre_processing/shared.py +++ b/src/dcegm/pre_processing/shared.py @@ -9,12 +9,9 @@ def determine_function_arguments_and_partial_model_specs( func, model_specs, - not_allowed_state_choices=None, + not_allowed_state_choices, ): signature = set(inspect.signature(func).parameters) - not_allowed_state_choices = ( - [] if not_allowed_state_choices is None else not_allowed_state_choices - ) partialed_func, signature = partial_model_specs_and_update_signature( func=func, diff --git a/src/dcegm/pre_processing/sol_container.py b/src/dcegm/pre_processing/sol_container.py index d8311024..0b8428e4 100644 --- a/src/dcegm/pre_processing/sol_container.py +++ b/src/dcegm/pre_processing/sol_container.py @@ -7,14 +7,15 @@ def create_solution_container( n_total_wealth_grid: int, n_state_choices: int, n_continuous_state_combinations: int, - store_endog_grid: bool = True, + store_endog_grid: bool, ): """Create solution containers for value, policy, and endog_grid. endog_grid is only allocated when store_endog_grid is True. When the Druedahl-Jorgensen upper envelope is used, the "endogenous" grid is actually the fixed exogenous grid, so storing it is skipped and callers read the exogenous - grid directly instead (see model_config["continuous_states_info"]["dj_wealth_grid"]). + grid directly instead: each reader recomputes the state-choice's own grid on + demand (see compute_own_dj_wealth_grid). """ value_solved = jnp.full( @@ -38,14 +39,3 @@ def create_solution_container( ) return value_solved, policy_solved, endog_grid_solved - - -def broadcast_dj_wealth_grid(continuous_states_info: Dict[str, Any], shape): - """Broadcast the fixed Druedahl-Jorgensen wealth grid to the given shape. - - Used in place of reading a stored endog_grid when - model_config["upper_envelope"]["skip_endog_grid_storage"] is True. Expects - continuous_states_info = model_config["continuous_states_info"]. - - """ - return jnp.broadcast_to(continuous_states_info["dj_wealth_grid"], shape) diff --git a/src/dcegm/simulation/random_keys.py b/src/dcegm/simulation/random_keys.py index a14e44ba..a5a8dd9b 100644 --- a/src/dcegm/simulation/random_keys.py +++ b/src/dcegm/simulation/random_keys.py @@ -3,7 +3,13 @@ def draw_random_keys_for_seed(n_agents, n_periods, taste_shock_scale_is_scalar, seed): - """Draw the random keys jax uses for a given seed.""" + """Draw the random keys jax uses for a given seed. + + Every key row belongs to the period that uses it: period ``t`` draws its own + income and taste shocks from row ``t``, and the stochastic-state key in row ``t`` + is the one it uses to transition into ``t+1``. + + """ # We start by determining the number of keys per period for the three stochastic components # of the model. We will draw all keys together and assign them then based on index. # First: The transition of the exogenous processes. This happens vectorized over all @@ -48,8 +54,10 @@ def draw_random_keys_for_seed(n_agents, n_periods, taste_shock_scale_is_scalar, "taste_shock_keys": sim_keys_draw[:-1, idx_3, :], } - # In the last period we only need the taste shock. + # The last period draws its own income and taste shocks like any other; it just + # never transitions onwards, so it needs no stochastic-state key. last_period_sim_keys = { + "income_shock_keys": sim_keys_draw[-1, idx_2, :], "taste_shock_keys": sim_keys_draw[-1, idx_3, :], } return sim_keys, last_period_sim_keys diff --git a/src/dcegm/simulation/sim_utils.py b/src/dcegm/simulation/sim_utils.py index e0d10d22..ccd5b39d 100644 --- a/src/dcegm/simulation/sim_utils.py +++ b/src/dcegm/simulation/sim_utils.py @@ -1,4 +1,5 @@ import jax +import jax.numpy as jnp import numpy as np import pandas as pd from jax import vmap @@ -16,11 +17,8 @@ def transition_to_next_period( choice, params, model_funcs_sim, - read_funcs, sim_keys, ): - n_agents = assets_end_of_period.shape[0] - stochastic_states_next_period = vmap( realize_stochastic_states, in_axes=(0, 0, 0, 0, None, None) )( @@ -57,20 +55,14 @@ def transition_to_next_period( discrete_states_beginning_of_period[key].dtype ) - income_shock_std = read_funcs["income_shock_std"](params) - income_shock_mean = read_funcs["income_shock_mean"](params) - - # Draw income shocks. - income_shocks_next_period = draw_normal_shocks( - key=sim_keys["income_shock_keys"], - num_agents=n_agents, - mean=income_shock_mean, - std=income_shock_std, - ) - - next_period_wealth = model_funcs_sim["compute_assets_begin_of_period"] - if continuous_state_beginning_of_period is not None: - + # The additional continuous states are settled here, at the end of the period: + # they never depend on the choice made in the period they belong to (see + # _check_choice_not_declared in process_model_functions.py), so the child's own + # state vector -- which already carries this period's choice as its + # lagged_choice -- is everything the law of motion needs. + if continuous_state_beginning_of_period is None: + continuous_state_next_period = None + else: continuous_state_next_period = calculate_second_continuous_state_for_all_agents( discrete_states_beginning_of_period=discrete_states_next_period, continuous_state_beginning_of_period=continuous_state_beginning_of_period, @@ -78,31 +70,104 @@ def transition_to_next_period( compute_continuous_state=model_funcs_sim["next_period_continuous_state"], ) - all_states_next_period = { - **discrete_states_next_period, - **continuous_state_next_period, - } - else: - all_states_next_period = discrete_states_next_period.copy() - continuous_state_next_period = None + # Neither next period's wealth nor its income shock is produced here. Both belong + # to the period that consumes them: each period draws its own shock and applies + # the budget equation at its own beginning (see draw_income_shocks_for_period and + # assets_begin_of_period_for_each_choice), because the budget equation may depend + # on that period's own choice -- which is not known yet at this point. + return discrete_states_next_period, continuous_state_next_period - assets_beginning_of_next_period, budget_aux = ( - calculate_assets_begin_of_period_for_all_agents( - states_beginning_of_period=all_states_next_period, - asset_grid_point_end_of_previous_period=assets_end_of_period, - income_shocks_of_period=income_shocks_next_period, + +def draw_income_shocks_for_period(n_agents, params, read_funcs, income_shock_keys): + """This period's own income shocks, drawn at its beginning. + + Drawn in the period that consumes them -- the one whose budget equation turns them + into beginning-of-period wealth -- and reported on that period's own row, so a + simulated wealth level can be reconciled against the budget equation from one row of + the output. + + """ + return draw_normal_shocks( + key=income_shock_keys, + num_agents=n_agents, + mean=read_funcs["income_shock_mean"](params), + std=read_funcs["income_shock_std"](params), + ) + + +def assets_begin_of_period_for_each_choice( + states_beginning_of_period, + continuous_state_beginning_of_period, + assets_end_of_previous_period, + income_shock, + choice_range, + params, + compute_assets_begin_of_period, + budget_depends_on_choice, +): + """Beginning-of-period wealth for every choice, before any choice is made. + + Applied at the *start* of a period, using the state that period was entered with + -- discrete states and additional continuous states alike, both already settled + at the end of last period -- plus last period's end-of-period assets and this + period's income shock. + + Returns arrays of shape ``(n_agents, n_choices)``. When the budget equation + declares ``choice`` the agent faces a genuinely different wealth per choice (e.g. + a choice-specific cost), and the discrete choice is then made by comparing each + choice's value *at its own wealth* -- which is exactly what the solved solution is + indexed by. When it does not, the single wealth is repeated across the choice axis + so everything downstream keeps one uniform shape, without paying for redundant + budget evaluations. + + """ + n_choices = len(choice_range) + states = states_beginning_of_period + if continuous_state_beginning_of_period is not None: + states = {**states, **continuous_state_beginning_of_period} + + def _for_choice(choice_value): + states_of_choice = { + **states, + "choice": jnp.full( + assets_end_of_previous_period.shape[0], + choice_value, + dtype=choice_range.dtype, + ), + } + return calculate_assets_begin_of_period_for_all_agents( + states_beginning_of_period=states_of_choice, + asset_grid_point_end_of_previous_period=assets_end_of_previous_period, + income_shocks_of_period=income_shock, params=params, - compute_assets_begin_of_period=next_period_wealth, + compute_assets_begin_of_period=compute_assets_begin_of_period, ) - ) - return ( - assets_beginning_of_next_period, - budget_aux, - discrete_states_next_period, - continuous_state_next_period, - income_shocks_next_period, - ) + if not budget_depends_on_choice: + assets, aux = calculate_assets_begin_of_period_for_all_agents( + states_beginning_of_period=states, + asset_grid_point_end_of_previous_period=assets_end_of_previous_period, + income_shocks_of_period=income_shock, + params=params, + compute_assets_begin_of_period=compute_assets_begin_of_period, + ) + assets = jnp.repeat(assets[:, None], n_choices, axis=1) + aux = { + key: jnp.repeat(val[:, None], n_choices, axis=1) for key, val in aux.items() + } + return assets, aux + + # Python loop over the (small, static) choice axis rather than vmap: choice is a + # concrete value here, so each iteration traces the user's budget equation with + # it already bound. + per_choice = [_for_choice(choice_value) for choice_value in choice_range] + assets = jnp.stack([a for a, _ in per_choice], axis=1) + aux_keys = per_choice[0][1].keys() + aux = { + key: jnp.stack([aux_j[key] for _, aux_j in per_choice], axis=1) + for key in aux_keys + } + return assets, aux def compute_final_utility_for_each_choice( @@ -118,7 +183,7 @@ def compute_final_utility_for_each_choice( return util -def draw_normal_shocks(key, num_agents, mean=0, std=1): +def draw_normal_shocks(key, num_agents, mean, std): return jax.random.normal(key=key, shape=(num_agents,)) * std + mean diff --git a/src/dcegm/simulation/simulate.py b/src/dcegm/simulation/simulate.py index e5bfefe2..3baa03df 100644 --- a/src/dcegm/simulation/simulate.py +++ b/src/dcegm/simulation/simulate.py @@ -13,13 +13,42 @@ ) from dcegm.simulation.random_keys import draw_random_keys_for_seed from dcegm.simulation.sim_utils import ( + assets_begin_of_period_for_each_choice, compute_final_utility_for_each_choice, + draw_income_shocks_for_period, transition_to_next_period, vectorized_utility, ) from dcegm.simulation.taste_shocks import draw_taste_shocks +def _check_initial_assets(states_initial): + """Reject the pre-``assets_end_of_previous_period`` initial-state contract. + + Passing ``assets_begin_of_period`` used to mean "period 0's wealth, take it as + given", which made period 0 the one period whose budget equation never ran. Now + every period builds its own wealth, so what the first period needs is what it + inherits, not what it ends up with -- and silently reading the old key would quietly + reinterpret its value. + + """ + if "assets_end_of_previous_period" in states_initial: + return + if "assets_begin_of_period" in states_initial: + raise ValueError( + "states_initial must contain 'assets_end_of_previous_period', not " + "'assets_begin_of_period'. Every period, the first included, computes " + "its beginning-of-period wealth from the budget equation, so supply " + "the assets agents carry into the first period; its wealth then " + "follows from the budget equation as in every other period." + ) + raise ValueError( + "states_initial must contain 'assets_end_of_previous_period': the assets " + "agents carry into the first period, from which the budget equation builds " + "that period's wealth." + ) + + def simulate_all_periods( states_initial, n_periods, @@ -31,7 +60,7 @@ def simulate_all_periods( model_structure, model_funcs, model_config, - alt_model_funcs_sim=None, + alt_model_funcs_sim, ): alt_model_funcs_sim = ( model_funcs if alt_model_funcs_sim is None else alt_model_funcs_sim @@ -45,8 +74,15 @@ def simulate_all_periods( for key, value in states_initial.items() if key in discrete_state_space } - states_initial_dtype["assets_begin_of_period"] = states_initial[ - "assets_begin_of_period" + # Every period, the first included, builds its own beginning-of-period wealth + # from the budget equation (see simulate_single_period), so the carry holds what + # that needs -- last period's end-of-period assets -- rather than a ready-made + # wealth level. That is why the user supplies assets_end_of_previous_period: + # agents enter period 0 out of a notional period -1 whose savings they carry, and + # period 0's wealth follows from the budget equation like any other period's. + _check_initial_assets(states_initial) + states_initial_dtype["assets_end_of_previous_period"] = states_initial[ + "assets_end_of_previous_period" ] if "dummy_stochastic" in model_structure["stochastic_states_names"]: @@ -64,6 +100,9 @@ def simulate_all_periods( if has_additional_continuous_state: for name in additional_continuous_state_names: + # Carried as this period's realised value: unlike wealth, it is settled + # at the end of the previous period, so the user's array is simply the + # first period's own state. states_initial_dtype[name] = states_initial[name] n_agents = len(states_initial["period"]) @@ -77,6 +116,7 @@ def simulate_all_periods( ], seed=seed, ) + read_funcs = model_funcs["read_funcs"] simulate_body = partial( simulate_single_period, @@ -87,7 +127,8 @@ def simulate_all_periods( model_structure_sol=model_structure, model_funcs_sim=alt_model_funcs_sim, compute_utility=model_funcs["compute_utility"], - read_funcs=model_funcs["read_funcs"], + read_funcs=read_funcs, + continuous_grid_functions=model_funcs["continuous_grid_functions"], model_config=model_config, ) @@ -110,6 +151,11 @@ def simulate_all_periods( compute_utility_final=model_funcs["compute_utility_final"], continuous_states_info=model_config["continuous_states_info"], model_structure_sol=model_structure, + compute_assets_begin_of_period=alt_model_funcs_sim[ + "compute_assets_begin_of_period" + ], + budget_depends_on_choice=alt_model_funcs_sim["budget_depends_on_choice"], + read_funcs=read_funcs, ) # Standard simulation output @@ -140,6 +186,7 @@ def simulate_single_period( model_funcs_sim, compute_utility, read_funcs, + continuous_grid_functions, model_config, ): @@ -151,25 +198,56 @@ def simulate_single_period( "additional_continuous_state_names" ] - if has_additional_continuous_state: - continuous_state_name = additional_continuous_state_names[0] + # The carry holds bookkeeping entries alongside the states, so always select + # the discrete states by name rather than taking the carry wholesale. + discrete_states_beginning_of_period = { + key: value + for key, value in states_beginning_of_period.items() + if key in model_structure_sol["discrete_states_names"] + } + choice_range = model_structure_sol["choice_range"] + n_agents = states_beginning_of_period["assets_end_of_previous_period"].shape[0] + # This period's own income shock, drawn here rather than inherited from last + # period's transition, so the shock, the wealth it builds and the row it is + # reported on all belong to the same period. + income_shock = draw_income_shocks_for_period( + n_agents=n_agents, + params=params, + read_funcs=read_funcs, + income_shock_keys=sim_keys["income_shock_keys"], + ) + + if has_additional_continuous_state: + # Already settled at the end of last period, unlike wealth: the continuous + # state may not depend on the choice about to be made (see + # _check_choice_not_declared in process_model_functions.py), so there is + # nothing left to resolve here. continuous_state_beginning_of_period = { name: states_beginning_of_period[name] for name in additional_continuous_state_names } - discrete_states_beginning_of_period = { - key: value - for key, value in states_beginning_of_period.items() - if key in model_structure_sol["discrete_states_names"] - } else: - discrete_states_beginning_of_period = states_beginning_of_period continuous_state_beginning_of_period = None - assets_begin_of_period = states_beginning_of_period["assets_begin_of_period"] - - choice_range = model_structure_sol["choice_range"] + # Law of motion at the *beginning* of the period, before any choice is made, + # and evaluated for every choice: a budget equation may declare "choice" and + # then hand the agent a different wealth per choice, which is exactly what the + # solved solution is indexed by. Shape (n_agents, n_choices). + assets_begin_of_period, budget_aux = assets_begin_of_period_for_each_choice( + states_beginning_of_period=discrete_states_beginning_of_period, + continuous_state_beginning_of_period=continuous_state_beginning_of_period, + assets_end_of_previous_period=states_beginning_of_period[ + "assets_end_of_previous_period" + ], + income_shock=income_shock, + choice_range=choice_range, + params=params, + compute_assets_begin_of_period=model_funcs_sim[ + "compute_assets_begin_of_period" + ], + budget_depends_on_choice=model_funcs_sim["budget_depends_on_choice"], + ) discount_factor = read_funcs["discount_factor"](params) # Interpolate policy and value function for all agents. @@ -187,23 +265,19 @@ def simulate_single_period( params=params, discrete_states_names=model_structure_sol["discrete_states_names"], compute_utility=compute_utility, - continuous_state_space=model_structure_sol["continuous_state_space"], - additional_continuous_state_grids=continuous_states_info[ - "additional_continuous_state_grids" - ], + continuous_grid_functions=continuous_grid_functions, upper_envelope_method=model_config["upper_envelope"]["method"], has_additional_continuous_state=has_additional_continuous_state, discount_factor=discount_factor, skip_endog_grid_storage=model_config["upper_envelope"][ "skip_endog_grid_storage" ], - dj_wealth_grid=continuous_states_info.get("dj_wealth_grid"), ) # Draw taste shocks and calculate final value. taste_shocks = draw_taste_shocks( params=params, - states_beginning_of_period=states_beginning_of_period, + states_beginning_of_period=discrete_states_beginning_of_period, n_choices=len(choice_range), taste_shock_function=model_funcs_sim["taste_shock_function"], taste_shock_keys=sim_keys["taste_shock_keys"], @@ -220,39 +294,52 @@ def simulate_single_period( )[:, 0] consumption = jnp.take_along_axis(policy, choice_index[:, None], axis=1)[:, 0] + # The wealth the agent actually faces is the one belonging to the choice made. + assets_realized = jnp.take_along_axis( + assets_begin_of_period, choice_index[:, None], axis=1 + )[:, 0] + budget_aux = { + key: jnp.take_along_axis(val, choice_index[:, None], axis=1)[:, 0] + for key, val in budget_aux.items() + } + if has_additional_continuous_state: + states_of_period = { + **discrete_states_beginning_of_period, + **continuous_state_beginning_of_period, + } + else: + states_of_period = discrete_states_beginning_of_period + utility_period = vmap(vectorized_utility, in_axes=(0, 0, 0, None, None))( consumption, - states_beginning_of_period, + states_of_period, choice, params, compute_utility, ) - savings_current_period = assets_begin_of_period - consumption - - ( - assets_beginning_of_next_period, - budget_aux, - discrete_states_next_period, - continuous_state_next_period, - income_shocks_next_period, - ) = transition_to_next_period( - discrete_states_beginning_of_period=discrete_states_beginning_of_period, - continuous_state_beginning_of_period=continuous_state_beginning_of_period, - assets_end_of_period=savings_current_period, - choice=choice, - params=params, - model_funcs_sim=model_funcs_sim, - read_funcs=read_funcs, - sim_keys=sim_keys, + savings_current_period = assets_realized - consumption + + discrete_states_next_period, continuous_state_next_period = ( + transition_to_next_period( + discrete_states_beginning_of_period=discrete_states_beginning_of_period, + continuous_state_beginning_of_period=continuous_state_beginning_of_period, + assets_end_of_period=savings_current_period, + choice=choice, + params=params, + model_funcs_sim=model_funcs_sim, + sim_keys=sim_keys, + ) ) states_next_period = discrete_states_next_period + # Carry what next period needs to run the budget equation at its beginning -- + # this period's end-of-period assets; it draws its own income shock -- alongside + # next period's already-settled continuous state. if has_additional_continuous_state: for name in additional_continuous_state_names: states_next_period[name] = continuous_state_next_period[name] - - states_next_period["assets_begin_of_period"] = assets_beginning_of_next_period + states_next_period["assets_end_of_previous_period"] = savings_current_period result = { "choice": choice, @@ -261,11 +348,13 @@ def simulate_single_period( "taste_shocks": taste_shocks, "value_max": value_max, "value_choice": values_across_choices, - "assets_begin_of_period": assets_begin_of_period, + "assets_begin_of_period": assets_realized, "savings": savings_current_period, - "income_shock": income_shocks_next_period, + "income_shock": income_shock, **budget_aux, - **states_beginning_of_period, + # Only the model's own states -- the carry additionally holds a bookkeeping + # entry (last period's end-of-period assets) that is not simulation output. + **states_of_period, } return states_next_period, result @@ -282,6 +371,9 @@ def simulate_final_period( compute_utility_final, continuous_states_info, model_structure_sol, + compute_assets_begin_of_period, + budget_depends_on_choice, + read_funcs, ): invalid_number = np.array( np.iinfo(map_state_choice_to_index.dtype).max, @@ -294,40 +386,68 @@ def simulate_final_period( for key, value in states_begin_of_final_period.items() if key in model_structure_sol["discrete_states_names"] } - assets_begin_of_final_period = states_begin_of_final_period[ - "assets_begin_of_period" - ] + # Drawn here, like every other period's (see simulate_single_period). + income_shock_final_period = draw_income_shocks_for_period( + n_agents=n_agents, + params=params, + read_funcs=read_funcs, + income_shock_keys=sim_keys["income_shock_keys"], + ) + # The final period applies the budget equation at its own beginning too, per + # choice, exactly like every other period (see simulate_single_period). Its + # continuous state arrives already settled from the previous period's + # transition. if continuous_states_info["has_additional_continuous_state"]: - additional_continuous_state_names = continuous_states_info[ - "additional_continuous_state_names" - ] - continuous_states_beginning_of_period = { + continuous_state_begin_of_final_period = { name: states_begin_of_final_period[name] - for name in additional_continuous_state_names - } - states_begin_of_final_period = { - **discrete_states_begin_last_period, - **continuous_states_beginning_of_period, + for name in continuous_states_info["additional_continuous_state_names"] } else: - states_begin_of_final_period = discrete_states_begin_last_period + continuous_state_begin_of_final_period = None + + assets_begin_of_final_period, budget_aux_final = ( + assets_begin_of_period_for_each_choice( + states_beginning_of_period=discrete_states_begin_last_period, + continuous_state_beginning_of_period=continuous_state_begin_of_final_period, + assets_end_of_previous_period=states_begin_of_final_period[ + "assets_end_of_previous_period" + ], + income_shock=income_shock_final_period, + choice_range=choice_range, + params=params, + compute_assets_begin_of_period=compute_assets_begin_of_period, + budget_depends_on_choice=budget_depends_on_choice, + ) + ) + + # Utility is evaluated per choice, at that choice's own wealth, so the state + # dict handed to it is repeated across the choice axis. Every state in it is + # choice-invariant -- only wealth differs by choice. + n_choices = len(choice_range) + states_begin_last_period = dict(discrete_states_begin_last_period) + if continuous_state_begin_of_final_period is not None: + states_begin_last_period.update(continuous_state_begin_of_final_period) + states_per_choice = { + key: jnp.repeat(val[:, None], n_choices, axis=1) + for key, val in states_begin_last_period.items() + } utilities_pre_taste_shock = vmap( vmap( compute_final_utility_for_each_choice, - in_axes=(None, 0, None, None, None), # choices + in_axes=(0, 0, 0, None, None), # choices ), in_axes=(0, None, 0, None, None), # agents )( - states_begin_of_final_period, + states_per_choice, choice_range, assets_begin_of_final_period, params, compute_utility_final, ) state_choice_indexes = get_state_choice_index_per_discrete_states( - states=states_begin_of_final_period, + states=discrete_states_begin_last_period, map_state_choice_to_index=map_state_choice_to_index, discrete_states_names=discrete_states_names, ) @@ -338,7 +458,7 @@ def simulate_final_period( # Draw taste shocks and calculate final value. taste_shocks = draw_taste_shocks( params=params, - states_beginning_of_period=states_begin_of_final_period, + states_beginning_of_period=discrete_states_begin_last_period, n_choices=len(choice_range), taste_shock_function=taste_shock_function, taste_shock_keys=sim_keys["taste_shock_keys"], @@ -355,17 +475,28 @@ def simulate_final_period( values_across_choices, choice_index[:, None], axis=1 )[:, 0] + # Everything is consumed in the final period, at the wealth belonging to the + # choice actually made. + assets_realized_final = jnp.take_along_axis( + assets_begin_of_final_period, choice_index[:, None], axis=1 + )[:, 0] + budget_aux_final = { + key: jnp.take_along_axis(val, choice_index[:, None], axis=1)[:, 0] + for key, val in budget_aux_final.items() + } + result = { "choice": choice, - "consumption": assets_begin_of_final_period, + "consumption": assets_realized_final, "utility": utility_period, "value_max": value_period, "value_choice": values_across_choices[np.newaxis], "taste_shocks": taste_shocks[np.newaxis, :, :], - "assets_begin_of_period": assets_begin_of_final_period, + "assets_begin_of_period": assets_realized_final, "savings": jnp.zeros_like(utility_period), - "income_shock": jnp.zeros(n_agents), - **states_begin_of_final_period, + "income_shock": income_shock_final_period, + **budget_aux_final, + **states_begin_last_period, } return result diff --git a/src/dcegm/solve_single_period.py b/src/dcegm/solve_single_period.py index e12a0d81..ac58d5fe 100644 --- a/src/dcegm/solve_single_period.py +++ b/src/dcegm/solve_single_period.py @@ -1,3 +1,5 @@ +from typing import Any, Callable, Dict, Optional, Tuple, Union + import jax.numpy as jnp from jax import vmap @@ -6,22 +8,102 @@ from dcegm.egm.solve_euler_equation import ( calculate_candidate_solutions_from_euler_equation, ) +from dcegm.law_of_motion import compute_own_continuous_grid_combos + +# The three solution containers threaded through the backward-induction scan as +# its carry: (value_solved, policy_solved, endog_grid_solved). The last is +# `None` when `skip_endog_grid_storage` is True. +SolutionCarry = Tuple[jnp.ndarray, jnp.ndarray, Optional[jnp.ndarray]] def solve_single_period( - carry, - xs, - params, - continuous_grids_info, - continuous_state_space, - income_shocks_scaled, - model_funcs, - income_shock_weights, - upper_envelope_method, - skip_endog_grid_storage, - debug_info, -): - """Solve a single period of the model using DCEGM.""" + carry: SolutionCarry, + xs: Tuple[Any, ...], + params: Dict[str, float], + continuous_grids_info: Dict[str, Any], + state_choice_space_dict: Dict[str, jnp.ndarray], + income_shocks_scaled: jnp.ndarray, + model_funcs: Dict[str, Any], + income_shock_weights: jnp.ndarray, + upper_envelope_method: str, + skip_endog_grid_storage: bool, + debug_info: Optional[Dict[str, bool]], +) -> Union[Tuple[SolutionCarry, Tuple[()]], Dict[str, jnp.ndarray]]: + """Solve one batch of state-choices -- the body of the backward induction scan. + + This is the ``f`` of the ``jax.lax.scan`` in ``backward_induction.py``. One + call solves every state-choice in one batch, reading the already-solved + children out of the carry and writing its own results back into it. Batches are + ordered so that a state-choice's children are always solved in an earlier + iteration (see ``pre_processing/batches/`` and the batching guide). + + The three EGM steps run in order: interpolate the children's continuation + values (``interpolate_value_and_marg_util``), aggregate them over choices and + income shocks, then invert the Euler equation and refine with the upper + envelope (both in ``solve_for_interpolated_values``). + + Args: + carry: The solution containers threaded through the scan, as + ``(value_solved, policy_solved, endog_grid_solved)``. Each is indexed + by *state-choice* and has shape ``(n_state_choices, + n_continuous_state_combinations, n_total_wealth_grid)``. They start out + filled by ``create_solution_container`` and are progressively filled + in, starting at the final period and working back to period 0. + ``endog_grid_solved`` is ``None`` when ``skip_endog_grid_storage`` is + True, because under Druedahl-Jorgensen every state-choice's + "endogenous" grid is by construction its own ``assets_begin_of_period`` + grid and is recomputed on demand (``compute_own_dj_wealth_grid``) + instead of stored. + xs: The per-batch slice produced by the scan. An 8-tuple; all index arrays + are into the global state-choice space unless noted: + + 0. ``state_choices_idxs`` -- the state-choices this batch solves, i.e. + where its results are written. + 1. ``child_state_choices_to_aggr_choice`` -- for each unique child + *state*, the positions of its choices in the deduplicated child + state-choice axis; used to aggregate over choices. Out-of-bounds + entries mark choices that state does not have. + 2. ``child_states_to_integrate_stochastic`` -- for each (row of this + batch, stochastic realisation), the position of the resulting child + *state*; used to integrate over the stochastic transition. + 3. ``child_state_choice_idxs_to_interp`` -- the deduplicated child + state-choices whose stored solution is read this iteration. + 4. ``child_state_idxs`` -- for each of those, the state it belongs to + (its own state, dropping the choice). + 5. ``state_choice_mat`` -- state-choice dict for this batch's own rows. + 6. ``state_choice_mat_child`` -- state-choice dict for the children in + (3). + 7. ``law_of_motion_arrays`` -- the child index/dedup arrays + ``calc_law_of_motion`` reads, carrying only the branch this model + takes. Assembled once at model setup and threaded through unchanged + (see ``bundle_law_of_motion_arrays`` in ``batch_creation.py``): a + ``rep_parent_state_choice_idx_per_child_state_choice`` when any + transition declares ``choice``, otherwise the per-unique-child-state + dedup arrays (``unique_child_states``, + ``rep_parent_state_choice_idx_per_child_state``, + ``state_row_for_state_choice``). The child state-choice dict itself + (6) is passed to ``calc_law_of_motion`` directly. + params: Model parameters. + continuous_grids_info: ``model_config["continuous_states_info"]``. + state_choice_space_dict: Full state-choice space; ``calc_law_of_motion`` + gathers the representative-parent index array in (7) out of this. + income_shocks_scaled: Quadrature points for the income shock, already + scaled by its mean and standard deviation. + model_funcs: Processed model functions. + income_shock_weights: Quadrature weights matching ``income_shocks_scaled``. + upper_envelope_method: ``"fues"`` or ``"druedahl_jorgensen"``. + skip_endog_grid_storage: Whether the endogenous grid is stored at all; see + ``carry`` above. + debug_info: ``None`` in the normal solve. When given, the function returns + a dict rather than a scan-shaped ``(carry, ())`` pair, optionally + including the pre-upper-envelope candidates. + + Returns: + ``(carry, ())`` with the updated solution containers -- the empty second + element because nothing is stacked per scan step. In debug mode a dict is + returned instead. + + """ value_solved, policy_solved, endog_grid_solved = carry ( @@ -32,8 +114,10 @@ def solve_single_period( child_state_idxs, state_choice_mat, state_choice_mat_child, + law_of_motion_arrays, ) = xs + value_child_state_choice = value_solved[child_state_choice_idxs_to_interp] policy_child_state_choice = policy_solved[child_state_choice_idxs_to_interp] endog_grid_child_state_choice = ( None @@ -44,30 +128,26 @@ def solve_single_period( # EGM step 1) value_interpolated, marginal_utility_interpolated = interpolate_value_and_marg_util( model_funcs=model_funcs, - state_choice_vec=state_choice_mat_child, + child_state_choices_with_proxy=state_choice_mat_child, continuous_grids_info=continuous_grids_info, income_shocks_scaled=income_shocks_scaled, endog_grid_child_state_choice=endog_grid_child_state_choice, - continuous_state_space=continuous_state_space, policy_child_state_choice=policy_child_state_choice, - value_child_state_choice=value_solved[child_state_choice_idxs_to_interp], + value_child_state_choice=value_child_state_choice, params=params, upper_envelope_method=upper_envelope_method, skip_endog_grid_storage=skip_endog_grid_storage, + law_of_motion_arrays=law_of_motion_arrays, + state_choice_space_dict=state_choice_space_dict, ) # Check if we have a scalar taste shock scale or state specific. Extract in each of the cases. - taste_shock_scale_is_scalar = model_funcs["taste_shock_function"][ - "taste_shock_scale_is_scalar" - ] + ts_function = model_funcs["taste_shock_function"] + taste_shock_scale_is_scalar = ts_function["taste_shock_scale_is_scalar"] if taste_shock_scale_is_scalar: - taste_shock_scale = model_funcs["taste_shock_function"][ - "read_out_taste_shock_scale" - ](params) + taste_shock_scale = ts_function["read_out_taste_shock_scale"](params) else: - taste_shock_scale_per_state_func = model_funcs["taste_shock_function"][ - "taste_shock_scale_per_state" - ] + taste_shock_scale_per_state_func = ts_function["taste_shock_scale_per_state"] taste_shock_scale = vmap(taste_shock_scale_per_state_func, in_axes=(0, None))( state_choice_mat_child, params ) @@ -83,7 +163,6 @@ def solve_single_period( taste_shock_scale_is_scalar=taste_shock_scale_is_scalar, income_shock_weights=income_shock_weights, continuous_grids_info=continuous_grids_info, - continuous_state_space=continuous_state_space, model_funcs=model_funcs, debug_info=debug_info, ) @@ -121,21 +200,64 @@ def solve_single_period( def solve_for_interpolated_values( - value_interpolated, - marginal_utility_interpolated, - state_choice_mat, - child_state_idxs, - states_to_choices_child_states, - params, - taste_shock_scale, - taste_shock_scale_is_scalar, - income_shock_weights, - continuous_grids_info, - continuous_state_space, - model_funcs, - debug_info, -): + value_interpolated: jnp.ndarray, + marginal_utility_interpolated: jnp.ndarray, + state_choice_mat: Dict[str, jnp.ndarray], + child_state_idxs: jnp.ndarray, + states_to_choices_child_states: jnp.ndarray, + params: Dict[str, float], + taste_shock_scale: Union[float, jnp.ndarray], + taste_shock_scale_is_scalar: bool, + income_shock_weights: jnp.ndarray, + continuous_grids_info: Dict[str, Any], + model_funcs: Dict[str, Any], + debug_info: Optional[Dict[str, bool]], +) -> Dict[str, jnp.ndarray]: + """EGM steps 2 and 3: aggregate continuation values, then invert the Euler eq. + + Split out from ``solve_single_period`` because the last two periods reach it by + a different route: ``final_periods.py`` computes the final period analytically + (consumption equals wealth) and then calls this directly for the second-to-last + period, rather than going through the scan. + + Takes the *interpolated* child continuation values -- one entry per child + state-choice, income shock and continuous-state combination -- and returns the + current period's solution: + + 1. ``aggregate_marg_utils_and_exp_values`` collapses the child-choice axis with + logit choice probabilities and the income-shock axis with quadrature + weights, giving one marginal utility and one expected value per child state. + 2. ``calculate_candidate_solutions_from_euler_equation`` gathers those to this + period's state-choices and inverts the Euler equation, producing candidate + (endogenous grid, policy, value) triples. + 3. ``run_upper_envelope`` discards the candidates that are not on the upper + envelope of the value correspondence -- the step that makes this DC-EGM + rather than plain EGM. + + Args: + value_interpolated: Child values, shape + ``(n_child_state_choices, n_continuous_combinations, n_wealth, + n_income_shocks)``. + marginal_utility_interpolated: Child marginal utilities, same shape. + state_choice_mat: State-choice dict for the rows being solved. + child_state_idxs: For each (row, stochastic realisation), the position of + the child state -- the stochastic integration map. + states_to_choices_child_states: For each child state, the positions of its + choices -- the choice aggregation map. + params: Model parameters. + taste_shock_scale: Scalar, or one value per state-choice. + taste_shock_scale_is_scalar: Which of the two it is. + income_shock_weights: Quadrature weights for the income shock. + continuous_grids_info: ``model_config["continuous_states_info"]``. + model_funcs: Processed model functions. + debug_info: When given with ``return_candidates``, the pre-upper-envelope + candidates are added to the returned dict. + Returns: + dict with ``"endog_grid"``, ``"policy"`` and ``"value"`` for the rows in + ``state_choice_mat``, plus the candidate arrays in debug mode. + + """ # EGM step 2) # Aggregate the marginal utilities and expected values over all child state-choice # combinations and income shock draws @@ -156,7 +278,6 @@ def solve_for_interpolated_values( expected_values, ) = calculate_candidate_solutions_from_euler_equation( continuous_grids_info=continuous_grids_info, - continuous_state_space=continuous_state_space, marg_util_next=marg_util, emax_next=emax, state_choice_mat=state_choice_mat, @@ -178,12 +299,13 @@ def solve_for_interpolated_values( policy_candidate=policy_candidate, value_candidate=value_candidate, expected_values=expected_values, - continuous_state_space=continuous_state_space, state_choice_mat=state_choice_mat, compute_utility=model_funcs["compute_utility"], params=params, discount_factor=discount_factor, compute_upper_envelope_for_state_choice=model_funcs["compute_upper_envelope"], + continuous_grid_functions=model_funcs["continuous_grid_functions"], + continuous_grids_info=continuous_grids_info, ) out_dict = { "endog_grid": endog_grid_state_choice, @@ -202,56 +324,121 @@ def solve_for_interpolated_values( def run_upper_envelope( - endog_grid_candidate, - policy_candidate, - value_candidate, - expected_values, - continuous_state_space, - state_choice_mat, - compute_utility, - params, - discount_factor, - compute_upper_envelope_for_state_choice, -): - """Run upper envelope to remove suboptimal candidates. + endog_grid_candidate: jnp.ndarray, + policy_candidate: jnp.ndarray, + value_candidate: jnp.ndarray, + expected_values: jnp.ndarray, + state_choice_mat: Dict[str, jnp.ndarray], + compute_utility: Callable, + params: Dict[str, float], + discount_factor: float, + compute_upper_envelope_for_state_choice: Callable, + continuous_grid_functions: Dict[str, Any], + continuous_grids_info: Dict[str, Any], +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """DC-EGM's refinement step: discard candidates not on the upper envelope. + + The last of the three EGM steps (see ``solve_for_interpolated_values``): + ``calculate_candidate_solutions_from_euler_equation`` (EGM step 3) can + produce a non-monotonic, self-intersecting candidate (endogenous grid, + policy, value) correspondence whenever the discrete choice set makes the + value function non-concave. This removes the candidates that are not on + the upper envelope of ``value_candidate`` -- the refinement that makes + this DC-EGM rather than plain EGM. + + Vectorized over all state-choice combinations. Builds each state-choice's own + continuous-state combo grid on demand, from its own identity (``state_choice_mat`` + is each row's own state-choice here, no representative-parent selection needed -- + same reasoning as ``solve_euler_equation.py``'s EGM step), instead of reusing one + grid shared across every state-choice -- a state-choice's own grid may differ once + continuous grids are state-choice-specific, and the upper-envelope refinement below + needs to agree with the EGM candidates it is refining. + + Args: + endog_grid_candidate: Candidate endogenous wealth grid from EGM step 3, + shape ``(n_state_choices, n_continuous_combinations, n_exog_savings)``. + policy_candidate: Candidate policy, same shape. + value_candidate: Candidate value, same shape. + expected_values: Expected value per state-choice and continuous + combination, used as the value at zero wealth; only the first + exogenous-savings entry is read (``[:, :, 0]``) since it is + constant along that axis. + state_choice_mat: State-choice dict for the rows being refined -- each + row's own identity, feeding both its own continuous grid and the + user's ``compute_utility``/``compute_upper_envelope_for_state_choice``. + compute_utility: User-supplied utility function. + params: Model parameters. + discount_factor: This period's discount factor. + compute_upper_envelope_for_state_choice: The model's compiled upper- + envelope routine (FUES or Druedahl-Jorgensen), applied per + continuous-state combination. + continuous_grid_functions: Functions returning each (possibly state- + choice-specific) continuous grid, used to build each row's own + combo grid on demand. + continuous_grids_info: ``model_config["continuous_states_info"]``. - Vectorized over all state-choice combinations. + Returns: + tuple ``(endog_grid, policy, value)``, each shape ``(n_state_choices, + n_continuous_combinations, n_total_wealth_grid)`` -- the refined + solution for every state-choice, on the fixed-width storage grid + (``model_config["upper_envelope"]["tuning_params"]["n_total_wealth_grid"]``) + the underlying FUES/Druedahl-Jorgensen routine pads or truncates to, + which need not equal ``n_exog_savings``. """ return vmap( - vmap( - compute_upper_envelope_for_state_choice, - in_axes=( - 0, - 0, - 0, - 0, - 0, - None, - None, - None, - None, - ), - ), - in_axes=( - 0, - 0, - 0, - 0, - None, - 0, - None, - None, - None, - ), + _run_upper_envelope_for_state_choice, + in_axes=(0, 0, 0, 0, 0, None, None, None, None, None, None, None), )( endog_grid_candidate, policy_candidate, value_candidate, expected_values[:, :, 0], - continuous_state_space, state_choice_mat, compute_utility, params, discount_factor, + compute_upper_envelope_for_state_choice, + continuous_grid_functions, + continuous_grids_info["additional_continuous_state_names"], + continuous_grids_info["has_additional_continuous_state"], + ) + + +def _run_upper_envelope_for_state_choice( + endog_grid_candidate_state_choice, + policy_candidate_state_choice, + value_candidate_state_choice, + expected_values_state_choice, + state_choice_vec, + compute_utility, + params, + discount_factor, + compute_upper_envelope_for_state_choice, + continuous_grid_functions, + additional_continuous_state_names, + has_additional_continuous_state, +): + if has_additional_continuous_state: + own_continuous_state_vec = compute_own_continuous_grid_combos( + state_choice_vec, + continuous_grid_functions, + additional_continuous_state_names, + ) + else: + own_continuous_state_vec = {"dummy_cont": jnp.zeros(1)} + + return vmap( + compute_upper_envelope_for_state_choice, + in_axes=(0, 0, 0, 0, 0, None, None, None, None), + )( + endog_grid_candidate_state_choice, + policy_candidate_state_choice, + value_candidate_state_choice, + expected_values_state_choice, + own_continuous_state_vec, + state_choice_vec, + compute_utility, + params, + discount_factor, ) diff --git a/src/dcegm/templates/fullmodel/model_funcs/budget_constraint.py b/src/dcegm/templates/fullmodel/model_funcs/budget_constraint.py index d6c46c45..0da87f87 100644 --- a/src/dcegm/templates/fullmodel/model_funcs/budget_constraint.py +++ b/src/dcegm/templates/fullmodel/model_funcs/budget_constraint.py @@ -6,7 +6,7 @@ def budget_constraint_exp( experience, already_retired, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, model_specs, ): @@ -24,7 +24,7 @@ def budget_constraint_exp( income_from_previous_period = _calc_stochastic_income( experience=exp_years, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, params=params, ) diff --git a/src/dcegm/templates/fullmodel/run_example.ipynb b/src/dcegm/templates/fullmodel/run_example.ipynb index ed9533c3..c6c98110 100644 --- a/src/dcegm/templates/fullmodel/run_example.ipynb +++ b/src/dcegm/templates/fullmodel/run_example.ipynb @@ -560,7 +560,7 @@ " \"job_offer\": jnp.ones(n_agents),\n", " \"survival\": jnp.ones(n_agents),\n", " \"experience\": jnp.ones(n_agents),\n", - " \"assets_begin_of_period\": jnp.ones(n_agents) * 10,\n", + " \"assets_end_of_previous_period\": jnp.ones(n_agents) * 10,\n", "}\n", "\n", "model_solved.simulate(states_initial=states_initial, seed=42)" diff --git a/src/dcegm/templates/fullmodel/run_example.py b/src/dcegm/templates/fullmodel/run_example.py index e63d4ab6..4d4182d7 100644 --- a/src/dcegm/templates/fullmodel/run_example.py +++ b/src/dcegm/templates/fullmodel/run_example.py @@ -87,7 +87,7 @@ "job_offer": jnp.ones(n_agents), "survival": jnp.ones(n_agents), "experience": jnp.ones(n_agents), - "assets_begin_of_period": jnp.ones(n_agents) * 10, + "assets_end_of_previous_period": jnp.ones(n_agents) * 10, } model_solved.simulate(states_initial=states_initial, seed=42) diff --git a/src/dcegm/templates/simplemodel/model_funcs/budget_constraint.py b/src/dcegm/templates/simplemodel/model_funcs/budget_constraint.py index 4d9cf49c..5637c7f0 100644 --- a/src/dcegm/templates/simplemodel/model_funcs/budget_constraint.py +++ b/src/dcegm/templates/simplemodel/model_funcs/budget_constraint.py @@ -8,7 +8,7 @@ def budget_constraint( period: int, lagged_choice: int, asset_end_of_previous_period: float, - income_shock_previous_period: float, + income_shock: float, model_specs: Dict[str, Any], params: Dict[str, float], ) -> float: @@ -22,10 +22,9 @@ def budget_constraint( denoting the current child state. savings_end_of_previous_period (float): One point on the exogenous savings grid carried over from the previous period. - income_shock_pervious_period (float): Stochastic shock on labor income; - may or may not be normally distributed. This float represents one - particular realization of the income_shock_draws carried over from the - previous period. + income_shock (float): Stochastic shock on labor income, realized in this + period; may or may not be normally distributed. This float represents + one particular realization of the income_shock_draws. params (dict): Dictionary containing model parameters. model_specs (dict): model_specs dictionary. @@ -37,7 +36,7 @@ def budget_constraint( income_from_previous_period = _calc_stochastic_income( period=period, lagged_choice=lagged_choice, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, min_age=model_specs["min_age"], constant=params["constant"], exp=params["exp"], @@ -83,10 +82,9 @@ def _calc_stochastic_income( Args: state (jnp.ndarray): 1d array of shape (n_state_variables,) denoting the current child state. - wage_shock (float): Stochastic shock on labor income; - may or may not be normally distributed. This float represents one - particular realization of the income_shock_draws carried over from - the previous period. + wage_shock (float): Stochastic shock on labor income, realized in this + period; may or may not be normally distributed. This float represents + one particular realization of the income_shock_draws. params (dict): Dictionary containing model parameters. Relevant here are the coefficients of the wage equation. model_specs (dict): model_specs dictionary. diff --git a/src/dcegm/templates/simplemodel/run_example.ipynb b/src/dcegm/templates/simplemodel/run_example.ipynb index 65452c58..cb262d8b 100644 --- a/src/dcegm/templates/simplemodel/run_example.ipynb +++ b/src/dcegm/templates/simplemodel/run_example.ipynb @@ -205,7 +205,7 @@ " \"period\": jnp.zeros(n_agents),\n", " \"lagged_choice\": jnp.zeros(n_agents), # all agents start as workers\n", " \"experience\": jnp.ones(n_agents),\n", - " \"assets_begin_of_period\": jnp.ones(n_agents) * 10,\n", + " \"assets_end_of_previous_period\": jnp.ones(n_agents) * 10,\n", "}\n", "\n", "sim_df = model_solved.simulate(states_initial=states_initial, seed=42)" diff --git a/src/dcegm/templates/simplemodel/run_example.py b/src/dcegm/templates/simplemodel/run_example.py index c8e77a92..8d4308ed 100644 --- a/src/dcegm/templates/simplemodel/run_example.py +++ b/src/dcegm/templates/simplemodel/run_example.py @@ -72,7 +72,7 @@ "period": jnp.zeros(n_agents), "lagged_choice": jnp.zeros(n_agents), # all agents start as workers "experience": jnp.ones(n_agents), - "assets_begin_of_period": jnp.ones(n_agents) * 10, + "assets_end_of_previous_period": jnp.ones(n_agents) * 10, } sim_df = model_solved.simulate(states_initial=states_initial, seed=42) diff --git a/src/dcegm/toy_models/cons_ret_model_dcegm_paper/budget_constraint.py b/src/dcegm/toy_models/cons_ret_model_dcegm_paper/budget_constraint.py index 9fc4f210..03e07eb4 100644 --- a/src/dcegm/toy_models/cons_ret_model_dcegm_paper/budget_constraint.py +++ b/src/dcegm/toy_models/cons_ret_model_dcegm_paper/budget_constraint.py @@ -8,7 +8,7 @@ def budget_constraint( period: int, lagged_choice: int, asset_end_of_previous_period: float, - income_shock_previous_period: float, + income_shock: float, model_specs: Dict[str, Any], params: Dict[str, float], ) -> float: @@ -22,10 +22,9 @@ def budget_constraint( denoting the current child state. asset_end_of_previous_period (float): One point on the exogenous savings grid carried over from the previous period. - income_shock_pervious_period (float): Stochastic shock on labor income; - may or may not be normally distributed. This float represents one - particular realization of the income_shock_draws carried over from the - previous period. + income_shock (float): Stochastic shock on labor income, realized in this + period; may or may not be normally distributed. This float represents + one particular realization of the income_shock_draws. params (dict): Dictionary containing model parameters. model_specs (dict): model_specs dictionary. @@ -37,7 +36,7 @@ def budget_constraint( income_from_previous_period = _calc_stochastic_income( period=period, lagged_choice=lagged_choice, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, min_age=model_specs["min_age"], constant=params["constant"], exp=params["exp"], @@ -83,10 +82,9 @@ def _calc_stochastic_income( Args: state (jnp.ndarray): 1d array of shape (n_state_variables,) denoting the current child state. - wage_shock (float): Stochastic shock on labor income; - may or may not be normally distributed. This float represents one - particular realization of the income_shock_draws carried over from - the previous period. + wage_shock (float): Stochastic shock on labor income, realized in this + period; may or may not be normally distributed. This float represents + one particular realization of the income_shock_draws. params (dict): Dictionary containing model parameters. Relevant here are the coefficients of the wage equation. model_specs (dict): model_specs dictionary. diff --git a/src/dcegm/toy_models/cons_ret_model_stochastic_ltc/budget_equation.py b/src/dcegm/toy_models/cons_ret_model_stochastic_ltc/budget_equation.py index acb56521..50dbbcf1 100644 --- a/src/dcegm/toy_models/cons_ret_model_stochastic_ltc/budget_equation.py +++ b/src/dcegm/toy_models/cons_ret_model_stochastic_ltc/budget_equation.py @@ -5,12 +5,12 @@ def budget_equation_with_ltc( ltc, lagged_choice, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, ): resource = ( (1 + params["interest_rate"]) * asset_end_of_previous_period - + (params["wage_avg"] + income_shock_previous_period) + + (params["wage_avg"] + income_shock) * (1 - lagged_choice) # if worked last period - ltc * params["ltc_cost"] ) diff --git a/src/dcegm/toy_models/cons_ret_model_with_cont_exp/budget_constraint.py b/src/dcegm/toy_models/cons_ret_model_with_cont_exp/budget_constraint.py index 729aa9df..fb23ea23 100644 --- a/src/dcegm/toy_models/cons_ret_model_with_cont_exp/budget_constraint.py +++ b/src/dcegm/toy_models/cons_ret_model_with_cont_exp/budget_constraint.py @@ -8,7 +8,7 @@ def budget_constraint_cont_exp( lagged_choice, experience, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, model_specs, ): @@ -19,6 +19,6 @@ def budget_constraint_cont_exp( lagged_choice=lagged_choice, experience=experience_years, asset_end_of_previous_period=asset_end_of_previous_period, - income_shock_previous_period=income_shock_previous_period, + income_shock=income_shock, params=params, ) diff --git a/src/dcegm/toy_models/cons_ret_model_with_exp/budget_constraint_with_exp.py b/src/dcegm/toy_models/cons_ret_model_with_exp/budget_constraint_with_exp.py index 0f415643..d509e70d 100644 --- a/src/dcegm/toy_models/cons_ret_model_with_exp/budget_constraint_with_exp.py +++ b/src/dcegm/toy_models/cons_ret_model_with_exp/budget_constraint_with_exp.py @@ -5,7 +5,7 @@ def budget_constraint_exp( lagged_choice, experience, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, ): @@ -13,7 +13,7 @@ def budget_constraint_exp( income_from_previous_period = _calc_stochastic_income( experience=experience, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, params=params, ) diff --git a/tests/resources/divorce_model/dcegm_functions.py b/tests/resources/divorce_model/dcegm_functions.py index 12518fd6..84365b67 100644 --- a/tests/resources/divorce_model/dcegm_functions.py +++ b/tests/resources/divorce_model/dcegm_functions.py @@ -61,7 +61,7 @@ def budget_constraint( lagged_choice, partner_state, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, model_specs, ): @@ -79,6 +79,39 @@ def budget_constraint( return jnp.maximum(wealth, params["consumption_floor"]) / multiplier +def budget_constraint_choice_dependent( + period, + lagged_choice, + choice, + partner_state, + asset_end_of_previous_period, + income_shock, + params, + model_specs, +): + """Same as ``budget_constraint``, plus a cost paid when *this* period's own choice + is to work. + + Declaring ``choice`` is what routes the solve down the per-state-choice law of + motion (``budget_depends_on_choice``); the subtracted cost is what makes that + routing observable in the solution -- mirrors ``reference.py``'s + ``resources_after_transition(..., choice_next=..., work_cost=...)``, which every + caller of ``budget_constraint`` here must be compared against instead of + ``resources_after_transition`` without those two arguments. + + """ + multiplier = jnp.where(partner_state == 1, 2.0, 1.0) + own_income = params["y_work"] * (lagged_choice == 0) + partner_income = params["y_partner"] * (partner_state == 1) + wealth = ( + multiplier * asset_end_of_previous_period * (1 + params["interest_rate"]) + + own_income + + partner_income + - params["work_cost"] * (choice == 0) + ) + return jnp.maximum(wealth, params["consumption_floor"]) / multiplier + + def feasible_choice_set(lagged_choice, model_specs): return np.arange(model_specs["n_choices"]) @@ -91,13 +124,33 @@ def partner_transition(partner_state, params): return jnp.array([1 - prob_married_next, prob_married_next]) -def build_and_solve(params, n_periods, a_grid): +def build_and_solve( + params, n_periods, a_grid, continuous_grid_functions=None, budget_fn=None +): + """Build and solve the divorce model. + + ``budget_fn`` defaults to the choice-independent ``budget_constraint``; pass + ``budget_constraint_choice_dependent`` (with ``params["work_cost"]`` set) to route + the solve down the per-state-choice law of motion instead -- see + ``reference.solve_reference``'s ``work_cost`` argument for the independent ground + truth this is checked against. + + """ + budget_fn = budget_constraint if budget_fn is None else budget_fn model_specs = {"n_periods": n_periods, "n_choices": 2} model_config = { "n_periods": n_periods, "choices": np.arange(2), "stochastic_states": {"partner_state": np.arange(2)}, - "continuous_states": {"assets_end_of_period": a_grid}, + # `None` once a grid function takes the name over; the pairing is strict in + # both directions (see process_continuous_grid_functions). + "continuous_states": { + "assets_end_of_period": ( + None + if "assets_end_of_period" in (continuous_grid_functions or {}) + else a_grid + ) + }, "n_quad_points": 5, } model = dcegm.setup_model( @@ -114,6 +167,7 @@ def build_and_solve(params, n_periods, a_grid): "utility": utility_final, "marginal_utility": marginal_utility_final, }, - budget_constraint=budget_constraint, + budget_constraint=budget_fn, + continuous_grid_functions=continuous_grid_functions, ) return model, model.solve(params) diff --git a/tests/resources/divorce_model/reference.py b/tests/resources/divorce_model/reference.py index 201c5f75..a63e55df 100644 --- a/tests/resources/divorce_model/reference.py +++ b/tests/resources/divorce_model/reference.py @@ -32,7 +32,15 @@ def partner_transition_np(partner_state, params): return np.array([1 - prob_married_next, prob_married_next]) -def resources_after_transition(a, partner_state_0, partner_state_1, own_income, params): +def resources_after_transition( + a, + partner_state_0, + partner_state_1, + own_income, + params, + choice_next=None, + work_cost=0.0, +): """Plain resources, with the individual's asset stock rescaled once for the period-to-period partner transition: halved on divorce (lose the ex-partner's share), doubled on marriage (new partner matches wealth), @@ -40,6 +48,13 @@ def resources_after_transition(a, partner_state_0, partner_state_1, own_income, partnered, not on income, not in utility. Partner income is added whenever partnered next period (partner_state_1), unscaled, exactly like own_income. + + ``choice_next``/``work_cost``: optional support for a choice-dependent + budget, mirroring dcegm's convention that a child's own choice may affect + its own beginning-of-period wealth (see + ``dcegm_functions.budget_constraint_choice_dependent``). Both default to a + no-op, so every existing caller is unaffected. + """ if partner_state_0 == 1 and partner_state_1 == 0: a = a / 2 # divorce @@ -48,6 +63,8 @@ def resources_after_transition(a, partner_state_0, partner_state_1, own_income, partner_income = params["y_partner"] if partner_state_1 == 1 else 0.0 wealth = a * (1 + params["interest_rate"]) + own_income + partner_income + if choice_next is not None: + wealth = wealth - work_cost * (choice_next == 0) return max(wealth, params["consumption_floor"]) @@ -149,13 +166,52 @@ def solve_reference_backward_induction(params, a_grid): return solved -def solve_reference(n_periods, params, a_grid): +def grid_for(a_grid, period, choice): + """Resolve the end-of-period assets grid for one (period, choice). + + `a_grid` is either a plain 1d array -- one grid shared by every + state-choice, the original behaviour -- or a callable + `(period, choice) -> 1d array`, which is what lets this reference mirror + dcegm's state-choice-specific + `continuous_grid_functions["assets_end_of_period"]`. + + `period` and `choice` are the only things such a grid may depend on in + this model, so they are the only things this takes. dcegm requires every + parent state-choice transitioning into a shared child to agree on its own + grid (`check_continuous_grid_consistency_across_shared_children`). + `partner_state` is stochastic, so both partner states parent the same + children; `lagged_choice` fails for the same reason (two parents with + different `lagged_choice` but the same `choice` share a child). `choice` + itself is safe because it passes 1:1 into the child's `lagged_choice`, + and `period` because every parent of a given child sits in the same one. + + """ + if callable(a_grid): + return np.asarray(a_grid(period=period, choice=choice)) + return np.asarray(a_grid) + + +def solve_reference(n_periods, params, a_grid, work_cost=0.0): """General n-period manual EGM backward induction, no upper envelope. Only the terminal period is analytic (consume everything). Every earlier - period is solved on `a_grid` and stored; the period before it reads - (interpolates) that stored grid for its own continuation value and - marginal utility, exactly as dcegm does internally. + period is solved on its own end-of-period assets grid (see `grid_for`) + and stored; the period before it reads (interpolates) that stored grid + for its own continuation value and marginal utility, exactly as dcegm + does internally. + + `a_grid` may be a single shared array or a per-state-choice callable -- + the latter is the reference counterpart of dcegm's + `continuous_grid_functions`. Because each (period, partner_state, choice) + stores its solution on whatever grid it was solved with, and the + continuation lookup interpolates that stored grid by wealth value, no + other part of this function needs to know which case it is in. + + `work_cost`: when nonzero, a child's own choice to work (`choice_next == + 0`) costs this much, deducted from its own beginning-of-period wealth -- + the reference counterpart of dcegm's + `dcegm_functions.budget_constraint_choice_dependent`. Defaults to 0.0 + (no-op), matching every existing caller's economics exactly. Returns solved[period][(partner_state, choice)] = dict with parallel arrays "endog_grid", "policy", "value", for period in @@ -233,31 +289,37 @@ def continuation_value_and_marg_util(period, wealth, partner_state, choice): transition_probs = partner_transition_np(partner_state_now, params) for choice_now in (0, 1): - endog_grid = np.empty_like(a_grid) - policy = np.empty_like(a_grid) - value = np.empty_like(a_grid) + own_a_grid = grid_for(a_grid, period, choice_now) + endog_grid = np.empty_like(own_a_grid) + policy = np.empty_like(own_a_grid) + value = np.empty_like(own_a_grid) # Income from choice_now (work vs not) arrives at the start # of period + 1, dcegm-style. income_next = params["y_work"] * (choice_now == 0) - for i, a_end in enumerate(a_grid): + for i, a_end in enumerate(own_a_grid): expected_marg_util = 0.0 expected_value = 0.0 for partner_state_next, prob_partner in enumerate(transition_probs): if prob_partner == 0.0: continue - wealth_next = resources_after_transition( - a_end, - partner_state_now, - partner_state_next, - income_next, - params, - ) choice_values = np.empty(2) choice_marg_utils = np.empty(2) for choice_next in (0, 1): + # A child's own choice may cost it its own + # beginning-of-period wealth (work_cost), so this + # is computed per choice_next, not once and shared. + wealth_next = resources_after_transition( + a_end, + partner_state_now, + partner_state_next, + income_next, + params, + choice_next=choice_next, + work_cost=work_cost, + ) v, mu = continuation_value_and_marg_util( period + 1, wealth_next, diff --git a/tests/sparse_death/budget.py b/tests/sparse_death/budget.py index d6c46c45..0da87f87 100644 --- a/tests/sparse_death/budget.py +++ b/tests/sparse_death/budget.py @@ -6,7 +6,7 @@ def budget_constraint_exp( experience, already_retired, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, model_specs, ): @@ -24,7 +24,7 @@ def budget_constraint_exp( income_from_previous_period = _calc_stochastic_income( experience=exp_years, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, params=params, ) diff --git a/tests/sparse_death/test_death_proxy_euler.py b/tests/sparse_death/test_death_proxy_euler.py new file mode 100644 index 00000000..ecc585ba --- /dev/null +++ b/tests/sparse_death/test_death_proxy_euler.py @@ -0,0 +1,253 @@ +"""Closed-form Euler-equation oracle for the death (value-reuse proxy) redirect. + +A three-period toy in which the agent dies with certainty after period 0. The period-0 +parent's only child is therefore a death state at period 1, which proxies across periods +to the single last-period (period 2) death slot -- the exact cross-period redirect the +decoupling has to get right. + +Death is consume-all, so the child's marginal value is analytic (``bequest'(wealth, +experience)``), which makes the parent's Euler equation closed form despite the proxy. +The bequest depends on experience, so the check is sensitive to the child's own (non- +proxy) continuation being read correctly. If the solve reproduces this Euler equation, +the proxy value redirect plus the non-proxy transition are correct. + +""" + +import jax.numpy as jnp +import numpy as np +from numpy.testing import assert_allclose +from scipy.special import roots_sh_legendre +from scipy.stats import norm + +import dcegm + +N_PERIODS = 3 +LAST = N_PERIODS - 1 + +_PARAMS = { + "discount_factor": 0.95, + "interest_rate": 0.05, + "income_shock_std": 0.1, + "income_shock_mean": 0.0, + "taste_shock_scale": 1.0, + "wage": 1.0, + "wage_growth": 0.5, + "consumption_floor": 1e-8, + "bequest_exp_slope": 0.8, + "delta": 0.0, +} +_MODEL_CONFIG = { + "n_periods": N_PERIODS, + "choices": np.arange(1, dtype=int), + "continuous_states": { + "assets_end_of_period": jnp.linspace(0, 20, 12, dtype=float), + "experience": jnp.linspace(0, 1, 3, dtype=float), + }, + "stochastic_states": {"survival": [0, 1]}, + "n_quad_points": 5, +} + + +# ===================================================================================== +# Model functions +# ===================================================================================== + + +def _utility(consumption, survival, experience, choice, params): + alive = jnp.log(consumption) - choice * params["delta"] + dead = (1 + params["bequest_exp_slope"] * experience) * jnp.log(consumption) + return jnp.where(survival == 0, dead, alive) + + +def _marginal_utility(consumption, survival, experience, choice, params): + dead = (1 + params["bequest_exp_slope"] * experience) / consumption + return jnp.where(survival == 0, dead, 1 / consumption) + + +def _inverse_marginal_utility(marginal_utility, params): + # Only alive (non-terminal) states invert the Euler equation; death is consume-all. + return 1 / marginal_utility + + +def _utility_final(wealth, survival, experience, params): + alive = jnp.log(wealth) + dead = (1 + params["bequest_exp_slope"] * experience) * jnp.log(wealth) + return jnp.where(survival == 0, dead, alive) + + +def _marginal_utility_final(wealth, survival, experience, params): + dead = (1 + params["bequest_exp_slope"] * experience) / wealth + return jnp.where(survival == 0, dead, 1 / wealth) + + +def _budget(period, asset_end_of_previous_period, income_shock, choice, params): + # Declares ``choice`` (numerically a no-op with a single choice) so the model + # takes the choice-dependent law-of-motion branch -- i.e. so the death child is + # read out of ``law_of_motion_arrays["child_state_choices"]`` rather than the + # coarser per-state dedup, exercising the branch this oracle targets. Income + # grows with ``period`` so the redirect is load-bearing: reading the death + # child at its real age (period 1) versus the proxy's period (the last period) + # would change the transition and break the Euler equation below. + wage = params["wage"] * (1 + params["wage_growth"] * period) + income = jnp.exp(income_shock) * wage + 0.0 * choice + wealth = (1 + params["interest_rate"]) * asset_end_of_previous_period + income + return jnp.maximum(wealth, params["consumption_floor"]) + + +def _next_period_experience(experience): + # Identity: keeps the child's experience on a grid node, so the death value is + # read at an exact node and the closed-form comparison is not clouded by + # experience-axis interpolation. + return experience + + +def _next_period_continuous_state(experience): + return {"experience": _next_period_experience(experience)} + + +def _next_period_deterministic_state(period, choice): + return {"period": period + 1, "lagged_choice": choice} + + +def _sparsity_condition(period, lagged_choice, survival): + # Living states are solved in place; every death state redirects to the single + # last-period death slot (the cross-period value-reuse proxy). + if survival == 0: + return {"period": LAST, "lagged_choice": lagged_choice, "survival": 0} + return True + + +def _state_specific_choice_set(period, lagged_choice): + return np.arange(1, dtype=int) + + +def _prob_survival(): + # Death next period with certainty, so the period-0 parent's only child is a + # death state (proxied), which makes its Euler equation closed form. + return jnp.array([1.0, 0.0]) + + +def _setup_and_solve(): + model = dcegm.setup_model( + model_specs={"n_periods": N_PERIODS}, + model_config=_MODEL_CONFIG, + utility_functions={ + "utility": _utility, + "marginal_utility": _marginal_utility, + "inverse_marginal_utility": _inverse_marginal_utility, + }, + utility_functions_final_period={ + "utility": _utility_final, + "marginal_utility": _marginal_utility_final, + }, + state_space_functions={ + "state_specific_choice_set": _state_specific_choice_set, + "next_period_continuous_state": _next_period_continuous_state, + "next_period_deterministic_state": _next_period_deterministic_state, + "sparsity_condition": _sparsity_condition, + }, + budget_constraint=_budget, + stochastic_states_transitions={"survival": _prob_survival}, + ) + return model, model.solve(params=_PARAMS) + + +# ===================================================================================== +# Closed-form Euler right-hand side +# ===================================================================================== + + +def _euler_rhs(savings, experience, params, draws, weights, child_period): + """Discounted expected marginal value of the (certain) death child. + + Death is consume-all, so the child consumes its whole beginning-of-period wealth and + its marginal value is ``bequest'(wealth, experience)``. Experience is unchanged into + the child (identity transition); the child's *own* age (``child_period``, not the + proxy's last-period identity) sets its wage -- the quantity that makes the redirect + load-bearing (see ``_budget``). + + """ + r = params["interest_rate"] + wage = params["wage"] * (1 + params["wage_growth"] * child_period) + marg_slope = 1 + params["bequest_exp_slope"] * experience + rhs = 0.0 + for draw, weight in zip(draws, weights): + income = np.exp(draw) * wage + wealth_child = max((1 + r) * savings + income, params["consumption_floor"]) + rhs += weight * marg_slope / wealth_child + return params["discount_factor"] * (1 + r) * rhs + + +def _assert_period0_euler_holds(model, solved): + """Check every period-0 parent's Euler equation against the closed-form RHS. + + Shared by the oracle test and the regression test below, which monkeypatches the + non-proxy child builder and expects this same check to fail. + + """ + quad_points, quad_weights = roots_sh_legendre(_MODEL_CONFIG["n_quad_points"]) + draws = ( + norm.ppf(quad_points) * _PARAMS["income_shock_std"] + + _PARAMS["income_shock_mean"] + ) + + scs = model.model_structure["state_choice_space"] + experience_grid = np.asarray(_MODEL_CONFIG["continuous_states"]["experience"]) + endog_grid = np.asarray(solved.endog_grid) + policy = np.asarray(solved.policy) + + period_0_state_choices = np.where(scs[:, 0] == 0)[0] + assert period_0_state_choices.size > 0 + + checked = 0 + for state_choice_idx in period_0_state_choices: + child_period = int(scs[state_choice_idx, 0]) + 1 # the death child's real age + for experience_idx, experience in enumerate(experience_grid): + for wealth_idx in range(endog_grid.shape[2]): + resources = endog_grid[state_choice_idx, experience_idx, wealth_idx] + consumption = policy[state_choice_idx, experience_idx, wealth_idx] + if not np.isfinite(resources) or resources <= 0: + continue + savings = resources - consumption + lhs = 1 / consumption # living parent's marginal utility + rhs = _euler_rhs( + savings, experience, _PARAMS, draws, quad_weights, child_period + ) + assert_allclose(lhs, rhs, atol=1e-6) + checked += 1 + assert checked > 0 + + +def test_death_proxy_satisfies_closed_form_euler(): + model, solved = _setup_and_solve() + _assert_period0_euler_holds(model, solved) + + +def test_death_proxy_regression_detects_broken_redirect(monkeypatch): + """The oracle above is only meaningful if it actually fails when the redirect is + broken. + + Monkeypatch the non-proxy child builder to fall back to the proxy identity (the + child the law of motion would see if the discrete-state decoupling in + ``single_segment.py`` were undone) and check the same Euler equation now fails, + since the death child's income would then be read at the proxy's (last) period + instead of its own. + + """ + import dcegm.pre_processing.batches.single_segment as single_segment + + def _proxy_passthrough( + state_space_dict, + unique_child_states, + state_row_for_state_choice, + proxy_child_state_choices, + ): + return proxy_child_state_choices + + monkeypatch.setattr( + single_segment, "build_no_proxy_child_state_choices", _proxy_passthrough + ) + + model, solved = _setup_and_solve() + with np.testing.assert_raises(AssertionError): + _assert_period0_euler_holds(model, solved) diff --git a/tests/sparse_death/test_sparse_death_period_grid.py b/tests/sparse_death/test_sparse_death_period_grid.py new file mode 100644 index 00000000..77df97bb --- /dev/null +++ b/tests/sparse_death/test_sparse_death_period_grid.py @@ -0,0 +1,202 @@ +"""A period-specific second-continuous grid builds and solves even though every death +state proxies to a single last-period death slot. + +The grid-consistency check groups by the actual child (each its own), so the proxy- +collapsed death children do not force a shared grid. The solve is decoupled from the +proxy too: the law of motion runs on the non-proxy child (its real age), and only the +value/policy lookup follows the proxy to its solved slot (see ``calc_law_of_motion`` / +``child_state_dedup.py``). + +The oracle at the bottom solves the model with death proxied and with death solved +explicitly per age and compares the living-state values: they agree even when the +bequest depends on experience, which validates the proxy value redirect. + +""" + +import jax.numpy as jnp +import numpy as np +import pytest + +import dcegm +from tests.sparse_death.budget import budget_constraint_exp +from tests.sparse_death.state_space import create_state_space_functions +from tests.sparse_death.stochastic_processes import job_offer, prob_survival +from tests.sparse_death.utility import ( + create_final_period_utility_function_dict, + create_utility_function_dict, +) + +N_PERIODS = 20 +N_EXP = 7 + +_MODEL_SPECS = { + "n_periods": N_PERIODS, + "n_choices": 3, + "min_ret_period": 5, + "max_ret_period": 10, + "fresh_bonus": 0.1, + "exp_scale": 20, +} +_PARAMS = { + "delta": 0.5, + "discount_factor": 0.95, + "taste_shock_scale": 1, + "income_shock_std": 1, + "income_shock_mean": 0.0, + "interest_rate": 0.05, + "constant": 1, + "exp": 0.1, + "exp_squared": -0.01, + "consumption_floor": 0.5, +} + + +def _base_config(): + return { + "min_period_batch_segments": [5, 12], + "n_periods": N_PERIODS, + "choices": np.arange(3, dtype=int), + "deterministic_states": {"already_retired": np.arange(2, dtype=int)}, + "continuous_states": { + "assets_end_of_period": jnp.arange(0, 100, 5, dtype=float), + "experience": jnp.linspace(0, 1, N_EXP, dtype=float), + }, + "stochastic_states": {"job_offer": [0, 1], "survival": [0, 1]}, + "n_quad_points": 5, + } + + +def _setup( + model_config, + continuous_grid_functions=None, + utility_functions=None, + utility_functions_final_period=None, + state_space_functions=None, +): + return dcegm.setup_model( + model_specs=_MODEL_SPECS, + model_config=model_config, + utility_functions=utility_functions or create_utility_function_dict(), + utility_functions_final_period=( + utility_functions_final_period + or create_final_period_utility_function_dict() + ), + state_space_functions=state_space_functions or create_state_space_functions(), + budget_constraint=budget_constraint_exp, + stochastic_states_transitions={ + "job_offer": job_offer, + "survival": prob_survival, + }, + continuous_grid_functions=continuous_grid_functions, + ) + + +def _period_specific_exp_grid(period): + # A genuinely period-dependent grid, spanning [0, 0.9] .. [0, 1.0] over the life + # cycle, so each age's death child keeps its own grid. + return jnp.linspace(0, 1, N_EXP) * (0.9 + 0.1 * period / (N_PERIODS - 1)) + + +def _constant_exp_grid(period): + # Ignores period: same grid as the declared-array baseline. + return jnp.linspace(0, 1, N_EXP) + + +def test_period_specific_grid_builds_and_solves(): + config = _base_config() + config["continuous_states"]["experience"] = None # supplied by the grid function + model = _setup(config, {"experience": _period_specific_exp_grid}) + solved = model.solve(params=_PARAMS) + assert np.isfinite( + np.asarray(solved.value)[~np.isnan(np.asarray(solved.value))] + ).all() + + +def test_constant_grid_function_matches_declared_array(): + # Sanity/no-regression: a period-*independent* grid supplied via + # continuous_grid_functions must reproduce the declared-array solve. + baseline = _setup(_base_config()).solve(params=_PARAMS) + + config = _base_config() + config["continuous_states"]["experience"] = None + via_func = _setup(config, {"experience": _constant_exp_grid}).solve(params=_PARAMS) + + for name in ("value", "policy", "endog_grid"): + a = np.asarray(getattr(baseline, name)) + b = np.asarray(getattr(via_func, name)) + m = np.isfinite(a) & np.isfinite(b) + # Relative: ``value`` differs up to ~1e7 across machines, where it hits ~1e22. + assert np.allclose(a[m], b[m], rtol=1e-10, atol=1e-10), name + + +def _bequest_utility_dict(bequest_slope): + """Utility whose death (survival == 0) branch is bequest = log(wealth) + + bequest_slope * experience -- a death value that depends on the second + continuous state. Marginal utility of consumption is unchanged (1 / c).""" + + def utility(consumption, survival, choice, experience, params): + alive = jnp.log(consumption) - choice * params["delta"] + dead = jnp.log(consumption) + bequest_slope * experience + return jnp.where(survival == 0, dead, alive) + + return { + "utility": utility, + "marginal_utility": lambda consumption, survival, choice, experience, params: 1 + / consumption, + "inverse_marginal_utility": lambda marginal_utility: 1 / marginal_utility, + } + + +def _bequest_final_utility_dict(bequest_slope): + return { + "utility": lambda wealth, experience, params: jnp.log(wealth) + + bequest_slope * experience, + "marginal_utility": lambda wealth, experience, params: 1 / wealth, + } + + +def _iter_batch_infos(batch_info): + for key, seg in batch_info.items(): + if not key.startswith("batches_info_segment_"): + continue + yield seg + if not seg.get("batches_cover_all", True): + yield seg["last_batch_info"] + + +def test_experience_dependent_bequest_solves_and_redirects_through_proxy(): + # A bequest that depends on experience solves with death proxied to one + # last-period slot, and the proxy structure holds: the value lookup collapses + # every death age onto that slot while the law-of-motion transition child keeps + # the child's real age. + config = _base_config() + config["continuous_states"]["experience"] = None + model = _setup( + config, + {"experience": _period_specific_exp_grid}, + utility_functions=_bequest_utility_dict(0.3), + utility_functions_final_period=_bequest_final_utility_dict(0.3), + ) + solved = model.solve(params=_PARAMS) + value = np.asarray(solved.value) + assert np.isfinite(value[~np.isnan(value)]).all() + + last_period = N_PERIODS - 1 + found_earlier_transition_age = False + for seg in _iter_batch_infos(model.batch_info): + proxy = seg["state_choices_childs"] + lma = seg["law_of_motion_arrays"] + dead = np.asarray(proxy["survival"]) == 0 + # The value/policy this row reads lives in the reduced space: every death + # child is redirected to a last-period death slot. + assert np.all(np.asarray(proxy["period"])[dead] == last_period) + # The transition, by contrast, runs on the child's real (non-proxy) age. + transition_period = np.take_along_axis( + np.asarray(lma["unique_child_states"]["period"]), + np.asarray(lma["state_row_for_state_choice"]), + axis=-1, + ) + found_earlier_transition_age |= bool( + np.any(transition_period[dead] < last_period) + ) + assert found_earlier_transition_age diff --git a/tests/test_asset_correction.py b/tests/test_asset_correction.py index 97b81dba..3fc5726d 100644 --- a/tests/test_asset_correction.py +++ b/tests/test_asset_correction.py @@ -46,7 +46,7 @@ def test_adjust_observed_assets_wealth_only(): period=period, lagged_choice=lagged_choice, asset_end_of_previous_period=asset_end, - income_shock_previous_period=0.0, + income_shock=0.0, model_specs=model_specs, params=params, ) @@ -97,7 +97,7 @@ def test_adjust_observed_assets_second_continuous_state(): lagged_choice=lagged_choice, experience=experience, asset_end_of_previous_period=asset_end, - income_shock_previous_period=0.0, + income_shock=0.0, params=params, model_specs=model_specs, ) diff --git a/tests/test_biased_sim.py b/tests/test_biased_sim.py index 5a87f925..34b7bc4f 100644 --- a/tests/test_biased_sim.py +++ b/tests/test_biased_sim.py @@ -114,7 +114,7 @@ def test_sim_and_sol_model(model_configs): "period": np.zeros(n_agents, dtype=int), "lagged_choice": np.zeros(n_agents, dtype=int), "married": initial_marriage_states, - "assets_begin_of_period": np.ones(n_agents, dtype=float) * 10, + "assets_end_of_previous_period": np.ones(n_agents, dtype=float) * 10, } df = model.solve_and_simulate( @@ -221,7 +221,12 @@ def test_sim_and_sol_model(model_configs): df_always_period_0_choice = df_always_period_0[ df_always_period_0["choice"] == choice ] - if df_always_period_0_choice.shape[0] > 0: + # A cell can hold a handful of agents (always-married *and* retiring in + # period 0 is rare), and a mean over those is noisier than the tolerance + # below by an order of magnitude. Assert only where the subsample is big + # enough to carry the comparison: measured gaps run ~0.05 from a few + # hundred agents upwards, against ~0.1-0.4 from single digits. + if df_always_period_0_choice.shape[0] >= 100: relevant_agents = df_always_period_0_choice.index.get_level_values( "agent" ) diff --git a/tests/test_budget_aux_with_second_continuous_state.py b/tests/test_budget_aux_with_second_continuous_state.py index b24b5e26..81fc86b6 100644 --- a/tests/test_budget_aux_with_second_continuous_state.py +++ b/tests/test_budget_aux_with_second_continuous_state.py @@ -38,7 +38,7 @@ def _wealth_beginning_of_period( asset_end_of_previous_period, lagged_choice, experience, - income_shock_previous_period, + income_shock, params, ): working = lagged_choice == 0 @@ -49,7 +49,7 @@ def _wealth_beginning_of_period( + params["exp"] * experience_years + params["exp_squared"] * experience_years**2 ) - income_from_previous_period = jnp.exp(labor_income + income_shock_previous_period) + income_from_previous_period = jnp.exp(labor_income + income_shock) wealth_beginning_of_period = ( income_from_previous_period * working @@ -63,7 +63,7 @@ def budget_constraint_without_aux( asset_end_of_previous_period, lagged_choice, experience, - income_shock_previous_period, + income_shock, params, ): return _wealth_beginning_of_period( @@ -71,7 +71,7 @@ def budget_constraint_without_aux( asset_end_of_previous_period, lagged_choice, experience, - income_shock_previous_period, + income_shock, params, ) @@ -81,7 +81,7 @@ def budget_constraint_with_aux( asset_end_of_previous_period, lagged_choice, experience, - income_shock_previous_period, + income_shock, params, ): wealth = _wealth_beginning_of_period( @@ -89,10 +89,10 @@ def budget_constraint_with_aux( asset_end_of_previous_period, lagged_choice, experience, - income_shock_previous_period, + income_shock, params, ) - aux_dict = {"income_shock": income_shock_previous_period} + aux_dict = {"income_shock": income_shock} return wealth, aux_dict diff --git a/tests/test_combined_solve.py b/tests/test_combined_solve.py new file mode 100644 index 00000000..cb34f3e8 --- /dev/null +++ b/tests/test_combined_solve.py @@ -0,0 +1,221 @@ +"""Tests for assembling a big model's solve from smaller sub-model solves. + +The model has an invariant "type" deterministic state (never changed by the +deterministic update, so it never crosses a transition). The big model spans both types; +the two sub-models each fix one type. Their state-choice spaces therefore partition the +big model's, and the assembled solve must reproduce the pooled solve. + +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from numpy.testing import assert_allclose + +import dcegm + +jax.config.update("jax_enable_x64", True) + + +# ==================================================================================== +# Model functions +# ==================================================================================== + + +def flow_util(consumption, choice, type, params): + rho = params["rho"] + # Work disutility differs by type, so the two types have different solutions. + disutility = params["delta"] * (choice == 0) * (1 + 0.5 * type) + return consumption ** (1 - rho) / (1 - rho) - disutility + + +def marginal_utility(consumption, params): + return consumption ** (-params["rho"]) + + +def inverse_marginal_utility(marginal_utility, params): + return marginal_utility ** (-1 / params["rho"]) + + +def final_period_utility(wealth, choice, params): + return wealth ** (1 - params["rho"]) / (1 - params["rho"]) + + +def marginal_final(wealth, choice, params): + return marginal_utility(wealth, params) + + +def state_specific_choice_set(period, lagged_choice, model_specs): + # Retirement (choice 1) is absorbing. + if lagged_choice == 1: + return np.array([1]) + return np.array([0, 1]) + + +def next_period_deterministic_state(period, choice, lagged_choice): + # "type" is deliberately not returned -> it is carried over unchanged (invariant). + return { + "period": period + 1, + "lagged_choice": choice, + } + + +def sparsity_condition(period, lagged_choice, type): + return True + + +def budget_constraint( + lagged_choice, + asset_end_of_previous_period, + income_shock, + params, +): + interest_factor = 1 + params["interest_rate"] + wage = params["wage_constant"] + income_shock + resource = interest_factor * asset_end_of_previous_period + wage * ( + lagged_choice == 0 + ) + return jnp.maximum(resource, 0.5) + + +utility_functions = { + "utility": flow_util, + "inverse_marginal_utility": inverse_marginal_utility, + "marginal_utility": marginal_utility, +} +utility_functions_final_period = { + "utility": final_period_utility, + "marginal_utility": marginal_final, +} +state_space_functions = { + "state_specific_choice_set": state_specific_choice_set, + "next_period_deterministic_state": next_period_deterministic_state, + "sparsity_condition": sparsity_condition, +} + +BUILD_KWARGS = dict( + model_specs={"n_choices": 2}, + utility_functions=utility_functions, + utility_functions_final_period=utility_functions_final_period, + state_space_functions=state_space_functions, + stochastic_states_transitions={}, + budget_constraint=budget_constraint, +) + + +def _model_config(type_grid): + return { + "n_periods": 5, + "choices": [0, 1], + "continuous_states": { + "assets_end_of_period": jnp.linspace(0, 50, 80), + }, + "deterministic_states": { + "type": np.asarray(type_grid, dtype=int), + }, + "n_quad_points": 5, + } + + +@pytest.fixture(scope="module") +def params(): + return { + "interest_rate": 0.03, + "wage_constant": 2.0, + "income_shock_std": 1.0, + "income_shock_mean": 0.0, + "taste_shock_scale": 1.0, + "discount_factor": 0.95, + "rho": 1.5, + "delta": 0.3, + } + + +@pytest.fixture(scope="module") +def big_config(): + return _model_config([0, 1]) + + +@pytest.fixture(scope="module") +def small_models(): + return [ + dcegm.setup_model(model_config=_model_config([0]), **BUILD_KWARGS), + dcegm.setup_model(model_config=_model_config([1]), **BUILD_KWARGS), + ] + + +def test_assembled_solve_matches_pooled(params, big_config, small_models): + big_model = dcegm.setup_model(model_config=big_config, **BUILD_KWARGS) + pooled = big_model.solve(params) + + solve_from_small = dcegm.get_solve_from_small_models( + small_models=small_models, + parallel=False, + model_config=big_config, + **BUILD_KWARGS, + ) + assembled = solve_from_small(params) + + assert_allclose(np.asarray(assembled.value), np.asarray(pooled.value)) + assert_allclose(np.asarray(assembled.policy), np.asarray(pooled.policy)) + assert_allclose(np.asarray(assembled.endog_grid), np.asarray(pooled.endog_grid)) + + +# The memory warning fires whenever there are fewer devices than sub-models, which is +# the normal single-device case here; what is under test is the assembled result. +@pytest.mark.filterwarnings("ignore:parallel=True dispatches") +def test_parallel_matches_pooled(params, big_config, small_models): + big_model = dcegm.setup_model(model_config=big_config, **BUILD_KWARGS) + pooled = big_model.solve(params) + + solve_from_small = dcegm.get_solve_from_small_models( + small_models=small_models, + parallel=True, + model_config=big_config, + **BUILD_KWARGS, + ) + assembled = solve_from_small(params) + + assert_allclose(np.asarray(assembled.value), np.asarray(pooled.value)) + assert_allclose(np.asarray(assembled.policy), np.asarray(pooled.policy)) + + +def test_overlapping_submodels_raise(big_config): + # Two identical full-span sub-models cover every row twice -> not a partition. + overlapping = [ + dcegm.setup_model(model_config=_model_config([0, 1]), **BUILD_KWARGS), + dcegm.setup_model(model_config=_model_config([0, 1]), **BUILD_KWARGS), + ] + with pytest.raises(ValueError, match="partition|sum to"): + dcegm.get_solve_from_small_models( + small_models=overlapping, + parallel=False, + model_config=big_config, + **BUILD_KWARGS, + ) + + +def test_incomplete_submodels_raise(big_config): + # Only one of the two types -> some big rows are uncovered. + incomplete = [dcegm.setup_model(model_config=_model_config([0]), **BUILD_KWARGS)] + with pytest.raises(ValueError, match="partition|sum to"): + dcegm.get_solve_from_small_models( + small_models=incomplete, + parallel=False, + model_config=big_config, + **BUILD_KWARGS, + ) + + +def test_mismatched_continuous_grid_raises(big_config, small_models): + # A big model with a different wealth grid is inconsistent with the sub-models. + bad_config = _model_config([0, 1]) + bad_config["continuous_states"]["assets_end_of_period"] = jnp.linspace(0, 50, 79) + with pytest.raises(ValueError, match="grid|n_total_wealth_grid"): + dcegm.get_solve_from_small_models( + small_models=small_models, + parallel=False, + model_config=bad_config, + **BUILD_KWARGS, + ) diff --git a/tests/test_divorce_toy_model.py b/tests/test_divorce_toy_model.py index a0b33367..9fd6f1a6 100644 --- a/tests/test_divorce_toy_model.py +++ b/tests/test_divorce_toy_model.py @@ -147,7 +147,7 @@ def wealth(asset, lagged_choice, partner_state): lagged_choice=lagged_choice, partner_state=partner_state, asset_end_of_previous_period=asset, - income_shock_previous_period=0.0, + income_shock=0.0, params=PARAMS, ) diff --git a/tests/test_initial_assets_contract.py b/tests/test_initial_assets_contract.py new file mode 100644 index 00000000..00b07369 --- /dev/null +++ b/tests/test_initial_assets_contract.py @@ -0,0 +1,174 @@ +"""The first period's wealth comes from the budget equation, like every other period. + +``states_initial`` carries ``assets_end_of_previous_period``: what agents bring into the +first period, not what they have to spend in it. These tests pin that down -- that +period 0's ``assets_begin_of_period`` is exactly what the model's own budget equation +produces from the supplied assets and period 0's own income shock, that the shock is the +one reported on period 0's own row, and that the old key is rejected rather than +silently reinterpreted. + +""" + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +import dcegm +import dcegm.toy_models as toy_models + +N_AGENTS = 50 +SEED = 4321 + + +@pytest.fixture(scope="module") +def solved_paper_model(): + model_funcs = toy_models.load_example_model_functions("dcegm_paper") + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config("dcegm_paper") + ) + model = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ) + return model, model.solve(params), params, model_specs, model_config + + +def _initial_states(assets_end_of_previous_period): + return { + "period": np.zeros(N_AGENTS, dtype=int), + "lagged_choice": np.zeros(N_AGENTS, dtype=int), + "assets_end_of_previous_period": assets_end_of_previous_period, + } + + +def test_first_period_wealth_is_the_budget_equation_applied_to_the_given_assets( + solved_paper_model, +): + _model, solved, params, model_specs, _model_config = solved_paper_model + assets_end_of_previous_period = np.linspace(1.0, 20.0, N_AGENTS) + + df = solved.simulate( + states_initial=_initial_states(assets_end_of_previous_period), seed=SEED + ) + + # Each period draws its own income shock and reports it on its own row, so + # period 0's wealth can be reconciled straight from period 0's output. + income_shock = df.xs(0, level="period")["income_shock"].to_numpy() + + budget_constraint = toy_models.load_example_model_functions("dcegm_paper")[ + "budget_constraint" + ] + expected = np.array( + [ + budget_constraint( + period=0, + lagged_choice=0, + asset_end_of_previous_period=assets_end_of_previous_period[i], + income_shock=income_shock[i], + model_specs=model_specs, + params=params, + ) + for i in range(N_AGENTS) + ] + ) + + simulated = df.xs(0, level="period")["assets_begin_of_period"].to_numpy() + assert_allclose(simulated, expected, rtol=1e-6) + + # Guards the test itself: the budget equation must actually move the number, or + # the assertion above would also pass if the supplied assets were used directly. + assert not np.allclose(expected, assets_end_of_previous_period) + + +def test_more_assets_carried_in_means_more_wealth_in_the_first_period( + solved_paper_model, +): + """The supplied assets are an input to period 0, not an unused placeholder.""" + _model, solved, _params, _model_specs, _model_config = solved_paper_model + + df_poor = solved.simulate( + states_initial=_initial_states(np.full(N_AGENTS, 5.0)), seed=SEED + ) + df_rich = solved.simulate( + states_initial=_initial_states(np.full(N_AGENTS, 25.0)), seed=SEED + ) + + assets_poor = df_poor.xs(0, level="period")["assets_begin_of_period"].to_numpy() + assets_rich = df_rich.xs(0, level="period")["assets_begin_of_period"].to_numpy() + assert np.all(assets_rich > assets_poor) + + +def test_old_initial_assets_key_is_rejected(solved_paper_model): + """``assets_begin_of_period`` meant something different, so it must not be read.""" + _model, solved, _params, _model_specs, _model_config = solved_paper_model + + states_initial = _initial_states(np.full(N_AGENTS, 10.0)) + states_initial["assets_begin_of_period"] = states_initial.pop( + "assets_end_of_previous_period" + ) + + with pytest.raises(ValueError, match="assets_end_of_previous_period"): + solved.simulate(states_initial=states_initial, seed=SEED) + + +def test_missing_initial_assets_key_is_rejected(solved_paper_model): + _model, solved, _params, _model_specs, _model_config = solved_paper_model + + states_initial = _initial_states(np.full(N_AGENTS, 10.0)) + del states_initial["assets_end_of_previous_period"] + + with pytest.raises(ValueError, match="assets_end_of_previous_period"): + solved.simulate(states_initial=states_initial, seed=SEED) + + +def test_each_period_reports_the_income_shock_that_built_its_own_wealth( + solved_paper_model, +): + """One row of the output is enough to reconcile that row's wealth. + + Each period draws its own income shock, feeds it to its own budget equation and + reports it on its own row. So for every period -- not just the first -- the budget + equation applied to the previous row's savings and this row's ``income_shock`` must + reproduce this row's ``assets_begin_of_period``. + + """ + _model, solved, params, model_specs, _model_config = solved_paper_model + assets_end_of_previous_period = np.linspace(1.0, 20.0, N_AGENTS) + + df = solved.simulate( + states_initial=_initial_states(assets_end_of_previous_period), seed=SEED + ) + budget_constraint = toy_models.load_example_model_functions("dcegm_paper")[ + "budget_constraint" + ] + + periods = sorted(df.index.get_level_values("period").unique()) + for period in periods: + this = df.xs(period, level="period") + carried_in = ( + assets_end_of_previous_period + if period == periods[0] + else df.xs(period - 1, level="period")["savings"].to_numpy() + ) + expected = np.array( + [ + budget_constraint( + period=period, + lagged_choice=int(this["lagged_choice"].to_numpy()[i]), + asset_end_of_previous_period=carried_in[i], + income_shock=this["income_shock"].to_numpy()[i], + model_specs=model_specs, + params=params, + ) + for i in range(N_AGENTS) + ] + ) + assert_allclose( + this["assets_begin_of_period"].to_numpy(), + expected, + rtol=1e-6, + err_msg=f"period {period}", + ) + + # The final period is included above only if it reports a real shock rather than + # a placeholder; guard that it does. + assert df.xs(periods[-1], level="period")["income_shock"].std() > 0 diff --git a/tests/test_interpnd_regular_own_grids.py b/tests/test_interpnd_regular_own_grids.py new file mode 100644 index 00000000..2a3cead3 --- /dev/null +++ b/tests/test_interpnd_regular_own_grids.py @@ -0,0 +1,254 @@ +"""Tests for the per-child-grid n-D regular interpolation (Druedahl-Jorgensen path). + +Companion to test_interpnd_regular.py, which validates the original shared-grid +contract (still used by direct callers -- see interpolation/interpnd_regular.py's +docstrings for why the two are kept separate). This file validates the "own grids" +variant used by the solve path once continuous grids are state-choice-specific: +each child is interpolated against its *own* grid, evaluated via a +continuous_grid_functions-style callable rather than read from a precomputed +per-child array, not one grid shared across all children. + +""" + +import jax.numpy as jnp +import numpy as np +from scipy.interpolate import RegularGridInterpolator + +from dcegm.interpolation.interpnd_regular import ( + interpnd_policy_and_value_for_child_states_on_own_regular_grids, +) + + +def _compute_utility(consumption, params, **kwargs): + return consumption ** params["u_scale"] + + +def _make_grid_func(grids_per_child): + # Fake continuous_grid_function: looks its child's own grid up by "child_idx" + # instead of computing it from economics -- lets these tests hand-pick + # per-child grids directly, same role a real grid_func plays in + # interpnd_policy_and_value_for_child_states_on_own_regular_grids's per-child + # vmap. + grids_per_child = jnp.asarray(grids_per_child) + + def grid_func(child_idx, **kwargs): + return grids_per_child[child_idx] + + return grid_func + + +def _scipy_expected_policy( + policy_grid_child_states, + exp_green_grids, + exp_red_grids, + wealth_grids, + continuous_state_child_states, + wealth_child_states, +): + n_child_state_choices, n_cont_combinations = continuous_state_child_states[ + "exp_green" + ].shape + expected = np.empty_like(wealth_child_states) + for i in range(n_child_state_choices): + policy_grid_nd = policy_grid_child_states[i].reshape( + exp_green_grids[i].size, exp_red_grids[i].size, wealth_grids[i].size + ) + interp = RegularGridInterpolator( + (exp_green_grids[i], exp_red_grids[i], wealth_grids[i]), + policy_grid_nd, + method="linear", + bounds_error=False, + fill_value=None, + ) + for j in range(n_cont_combinations): + for w in range(wealth_child_states.shape[2]): + for q in range(wealth_child_states.shape[3]): + point = np.array( + [ + continuous_state_child_states["exp_green"][i, j], + continuous_state_child_states["exp_red"][i, j], + wealth_child_states[i, j, w, q], + ] + ) + expected[i, j, w, q] = interp(point).item() + return expected + + +def test_interpnd_own_grids_matches_scipy_per_child_with_different_grids(): + # Two children with *different* grids for both continuous states and for + # wealth (child 1's grid is a scaled-up version of child 0's) -- this is + # exactly the scenario that would silently break if the interpolation axis + # were wrongly shared across children instead of read per child. + exp_green_grids = np.array( + [ + [0.0, 0.3, 0.8], + [0.0, 0.6, 1.6], + ] + ) + exp_red_grids = np.array( + [ + [0.0, 0.5], + [0.0, 1.0], + ] + ) + wealth_grids = np.array( + [ + [0.0, 2.0, 5.0, 9.0], + [0.0, 4.0, 10.0, 18.0], + ], + dtype=float, + ) + + n_child_state_choices = 2 + n_cont_combinations = exp_green_grids.shape[1] * exp_red_grids.shape[1] + n_wealth_eval = 4 + n_quad = 3 + + rng = np.random.default_rng(42) + + # Query points within each child's own grid range (including a bit outside for + # extrapolation coverage), NOT the other child's range. + continuous_state_child_states = { + "exp_green": np.stack( + [ + rng.uniform(-0.2, 1.0, size=n_cont_combinations), + rng.uniform(-0.2, 2.0, size=n_cont_combinations), + ] + ), + "exp_red": np.stack( + [ + rng.uniform(-0.2, 0.7, size=n_cont_combinations), + rng.uniform(-0.2, 1.4, size=n_cont_combinations), + ] + ), + } + wealth_child_states = np.stack( + [ + rng.uniform(-1.0, 10.0, size=(n_cont_combinations, n_wealth_eval, n_quad)), + rng.uniform(-2.0, 20.0, size=(n_cont_combinations, n_wealth_eval, n_quad)), + ] + ) + policy_grid_child_states = rng.normal( + size=(n_child_state_choices, n_cont_combinations, wealth_grids.shape[1]) + ) + # High value grid so policy interpolation dominates, no consume-all overwrite. + value_grid_child_states = np.full_like(policy_grid_child_states, 1e8) + + state_choice_child_states = { + "child_idx": jnp.arange(n_child_state_choices, dtype=jnp.int32), + "choice": jnp.zeros(n_child_state_choices, dtype=jnp.int32), + } + wealth_grid_func_ = _make_grid_func(wealth_grids) + + policy_out, _ = interpnd_policy_and_value_for_child_states_on_own_regular_grids( + continuous_grid_functions={ + "exp_green": _make_grid_func(exp_green_grids), + "exp_red": _make_grid_func(exp_red_grids), + }, + additional_continuous_state_names=["exp_green", "exp_red"], + wealth_grid_func=lambda state_dict: wealth_grid_func_(**state_dict), + policy_grid_child_states=jnp.asarray(policy_grid_child_states), + value_grid_child_states=jnp.asarray(value_grid_child_states), + continuous_state_child_states={ + k: jnp.asarray(v) for k, v in continuous_state_child_states.items() + }, + wealth_child_states=jnp.asarray(wealth_child_states), + state_choice_child_states=state_choice_child_states, + compute_utility=_compute_utility, + params={"u_scale": 2.0}, + discount_factor=0.95, + ) + + expected = _scipy_expected_policy( + policy_grid_child_states, + exp_green_grids, + exp_red_grids, + wealth_grids, + continuous_state_child_states, + wealth_child_states, + ) + np.testing.assert_allclose(np.asarray(policy_out), expected, rtol=1e-6, atol=1e-8) + + +def test_interpnd_own_grids_matches_shared_grid_path_when_grids_are_identical(): + # Regression check: when every child's own grid happens to be the same + # (today's global-grid behavior), the per-child-grid path must reproduce + # exactly what the original shared-grid function produces. + from dcegm.interpolation.interpnd_regular import ( + interpnd_policy_and_value_for_child_states_on_regular_grids, + ) + + exp_green_grid = np.array([0.0, 0.3, 0.8]) + exp_red_grid = np.array([0.0, 0.5]) + wealth_grid = np.array([0.0, 2.0, 5.0, 9.0], dtype=float) + + n_child_state_choices = 3 + n_cont_combinations = exp_green_grid.size * exp_red_grid.size + + rng = np.random.default_rng(7) + continuous_state_child_states = { + "exp_green": rng.uniform( + -0.2, 1.0, size=(n_child_state_choices, n_cont_combinations) + ), + "exp_red": rng.uniform( + -0.2, 0.7, size=(n_child_state_choices, n_cont_combinations) + ), + } + wealth_child_states = rng.uniform( + -1.0, 10.0, size=(n_child_state_choices, n_cont_combinations, 4, 3) + ) + policy_grid_child_states = rng.normal( + size=(n_child_state_choices, n_cont_combinations, wealth_grid.size) + ) + value_grid_child_states = rng.normal( + loc=-20.0, size=(n_child_state_choices, n_cont_combinations, wealth_grid.size) + ) + state_choice_child_states = { + "choice": jnp.zeros(n_child_state_choices, dtype=jnp.int32) + } + + shared_policy, shared_value = ( + interpnd_policy_and_value_for_child_states_on_regular_grids( + additional_continuous_state_grids={ + "exp_green": jnp.asarray(exp_green_grid), + "exp_red": jnp.asarray(exp_red_grid), + }, + wealth_grid=jnp.asarray(wealth_grid), + policy_grid_child_states=jnp.asarray(policy_grid_child_states), + value_grid_child_states=jnp.asarray(value_grid_child_states), + continuous_state_child_states={ + k: jnp.asarray(v) for k, v in continuous_state_child_states.items() + }, + wealth_child_states=jnp.asarray(wealth_child_states), + state_choice_child_states=state_choice_child_states, + compute_utility=_compute_utility, + params={"u_scale": 2.0}, + discount_factor=0.95, + ) + ) + + own_grids_policy, own_grids_value = ( + interpnd_policy_and_value_for_child_states_on_own_regular_grids( + continuous_grid_functions={ + "exp_green": lambda **kwargs: jnp.asarray(exp_green_grid), + "exp_red": lambda **kwargs: jnp.asarray(exp_red_grid), + }, + additional_continuous_state_names=["exp_green", "exp_red"], + wealth_grid_func=lambda state_dict: jnp.asarray(wealth_grid), + policy_grid_child_states=jnp.asarray(policy_grid_child_states), + value_grid_child_states=jnp.asarray(value_grid_child_states), + continuous_state_child_states={ + k: jnp.asarray(v) for k, v in continuous_state_child_states.items() + }, + wealth_child_states=jnp.asarray(wealth_child_states), + state_choice_child_states=state_choice_child_states, + compute_utility=_compute_utility, + params={"u_scale": 2.0}, + discount_factor=0.95, + ) + ) + + np.testing.assert_array_equal( + np.asarray(shared_policy), np.asarray(own_grids_policy) + ) + np.testing.assert_array_equal(np.asarray(shared_value), np.asarray(own_grids_value)) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 90ad52c6..3cc35d5a 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -157,6 +157,7 @@ def functional_form(x, y): compute_utility = determine_function_arguments_and_partial_model_specs( utility_crra, model_specs={}, + not_allowed_state_choices=[], ) test_cases[test_id]["test_points"] = test_points @@ -267,6 +268,7 @@ def test_interp2d_against_custom(test_cases, test_id): state_choice_vec={"choice": 0}, params=PARAMS, discount_factor=PARAMS["discount_factor"], + cont_state_name="continuous_state", ) ) diff --git a/tests/test_law_of_motion.py b/tests/test_law_of_motion.py index 5854387f..c58aae4f 100644 --- a/tests/test_law_of_motion.py +++ b/tests/test_law_of_motion.py @@ -31,7 +31,7 @@ def budget_constraint_based_on_experience( lagged_choice: int, continuous_state_beginning_of_period: float, asset_end_of_previous_period: float, - income_shock_previous_period: float, + income_shock: float, params: Dict[str, float], ) -> float: @@ -40,7 +40,7 @@ def budget_constraint_based_on_experience( wage = _calc_stochastic_income_for_experience( experience=experience_years, lagged_choice=lagged_choice, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, params=params, ) working_hours = _transform_lagged_choice_to_working_hours(lagged_choice) @@ -137,7 +137,7 @@ def test_get_beginning_of_period_wealth( wealth_beginning_of_period = budget_constraint( **child_state_dict, asset_end_of_previous_period=savings_grid[random_saving_scalar], - income_shock_previous_period=quad_points[random_shock_scalar], + income_shock=quad_points[random_shock_scalar], model_specs=model_specs, params=params, ) @@ -253,15 +253,18 @@ def _check_subset_matches_full(model, params): ) subset_result = calc_law_of_motion_for_state_choices( - state_choice_vec=state_choice_subset, - continuous_state_space=model_structure["continuous_state_space"], - assets_grid_end_of_period=continuous_states_info["assets_grid_end_of_period"], + child_state_choices=state_choice_subset, income_shocks_scaled=income_shocks_scaled, params=params, model_funcs=model_funcs, has_additional_continuous_states=continuous_states_info[ "has_additional_continuous_state" ], + additional_continuous_state_names=continuous_states_info[ + "additional_continuous_state_names" + ], + # No parent/child distinction being tested here. + representative_parent_state_choice_vec=state_choice_subset, ) expected_parent_states = map_state_choice_to_parent_state[test_idx] diff --git a/tests/test_load_save.py b/tests/test_load_save.py index bb9938d0..fc583199 100644 --- a/tests/test_load_save.py +++ b/tests/test_load_save.py @@ -114,7 +114,7 @@ def test_load_and_save_solution( states_initial = { "period": jnp.zeros(n_agents, dtype=int), "lagged_choice": jnp.zeros(n_agents, dtype=int), - "assets_begin_of_period": jnp.ones(n_agents, dtype=float) * 10, + "assets_end_of_previous_period": jnp.ones(n_agents, dtype=float) * 10, } seed = 132 @@ -137,3 +137,137 @@ def test_load_and_save_solution( import os os.remove("sol.pkl") + + +# ===================================================================================== +# Continuous grids declared as None in model_config +# +# Their size is only knowable once a real state-choice exists to evaluate the grid +# function against, so check_model_config.py leaves it unresolved. On the load path +# model_config is rebuilt from the raw user config while the state-choice space comes +# back from the pickle, so the size has to be pinned again there. +# ===================================================================================== + + +def _state_specific_experience_model(): + """``experience`` -- an additional continuous state -- supplied per state-choice.""" + model_funcs = toy_models.load_example_model_functions("with_cont_exp") + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config("with_cont_exp") + ) + model_config = dict(model_config) + model_config["continuous_states"] = dict(model_config["continuous_states"]) + grid = jnp.asarray(model_config["continuous_states"]["experience"]) + model_config["continuous_states"]["experience"] = None + + def experience_grid(period): + return grid + + return ( + model_funcs, + params, + model_specs, + model_config, + {"experience": experience_grid}, + ) + + +def _state_specific_wealth_grid_model(): + """``assets_begin_of_period`` -- the Druedahl-Jorgensen wealth grid.""" + model_funcs = toy_models.load_example_model_functions("with_exp") + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config("with_exp") + ) + model_config = dict(model_config) + model_config["continuous_states"] = dict(model_config["continuous_states"]) + model_config["continuous_states"]["assets_begin_of_period"] = None + model_config["upper_envelope"] = {"method": "druedahl_jorgensen"} + grid = jnp.linspace(0, 50, 50) + + def wealth_grid(period): + return grid + + return ( + model_funcs, + params, + model_specs, + model_config, + {"assets_begin_of_period": wealth_grid}, + ) + + +def _state_specific_assets_end_of_period_model(): + """``assets_end_of_period`` -- the one name whose deferral also moves tuning params. + + Its length sizes the fues tuning params, which the upper envelope closes over + before the load path pins anything, so this covers more than the other two. + + """ + model_funcs = toy_models.load_example_model_functions("dcegm_paper") + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config("dcegm_paper") + ) + model_config = dict(model_config) + model_config["continuous_states"] = dict(model_config["continuous_states"]) + grid = jnp.asarray(model_config["continuous_states"]["assets_end_of_period"]) + model_config["continuous_states"]["assets_end_of_period"] = None + + def assets_grid(period): + return grid + + return ( + model_funcs, + params, + model_specs, + model_config, + {"assets_end_of_period": assets_grid}, + ) + + +@pytest.mark.parametrize( + "model_loader", + [ + _state_specific_experience_model, + _state_specific_wealth_grid_model, + _state_specific_assets_end_of_period_model, + ], +) +def test_loaded_model_pins_deferred_continuous_grid_sizes(model_loader, tmp_path): + model_funcs, params, model_specs, model_config, continuous_grid_functions = ( + model_loader() + ) + path = str(tmp_path / "model.pkl") + + model_saved = dcegm.setup_model( + model_config=model_config, + model_specs=model_specs, + continuous_grid_functions=continuous_grid_functions, + model_save_path=path, + **model_funcs, + ) + model_loaded = dcegm.setup_model( + model_config=model_config, + model_specs=model_specs, + continuous_grid_functions=continuous_grid_functions, + model_load_path=path, + **model_funcs, + ) + + for key in ["n_total_wealth_grid"]: + assert model_loaded.model_config[key] == model_saved.model_config[key] + assert model_loaded.model_config[key] is not None + + combinations_key = "n_continuous_state_combinations" + assert ( + model_loaded.model_config["continuous_states_info"][combinations_key] + == model_saved.model_config["continuous_states_info"][combinations_key] + ) + assert ( + model_loaded.model_config["continuous_states_info"][combinations_key] + is not None + ) + + np.testing.assert_array_equal( + np.asarray(model_loaded.solve(params).value), + np.asarray(model_saved.solve(params).value), + ) diff --git a/tests/test_model_config.py b/tests/test_model_config.py index 5e2c2f15..934ac0d6 100644 --- a/tests/test_model_config.py +++ b/tests/test_model_config.py @@ -130,19 +130,45 @@ def test_skip_endog_grid_storage_false_for_fues(valid_model_config): assert options["upper_envelope"]["skip_endog_grid_storage"] is False -def test_dj_wealth_grid_and_skip_flag(valid_model_config): +def test_fues_rejects_assets_begin_of_period(valid_model_config): + # assets_begin_of_period is only meaningful for druedahl_jorgensen; declaring + # it together with fues (the default method here) used to be silently + # ignored rather than rejected. + valid_model_config["continuous_states"]["assets_begin_of_period"] = np.linspace( + 0, 10, 11 + ) + with pytest.raises(ValueError, match="only used by the 'druedahl_jorgensen'"): + check_model_config_and_process(valid_model_config) + + +def test_n_total_wealth_grid_and_skip_flag(valid_model_config): + # The Druedahl-Jorgensen wealth grid is a zero-wealth point plus the declared + # assets_begin_of_period, so its length is len + 1. No shared grid array is + # stored (each reader recomputes its own via compute_own_dj_wealth_grid); only + # the length, n_total_wealth_grid, is needed here -- to size the containers. valid_model_config["upper_envelope"] = {"method": "druedahl_jorgensen"} valid_model_config["continuous_states"]["assets_begin_of_period"] = np.linspace( 0, 10, 11 ) options = check_model_config_and_process(valid_model_config) - dj_wealth_grid = options["continuous_states_info"]["dj_wealth_grid"] - expected = np.concatenate( - ([0.0], valid_model_config["continuous_states"]["assets_begin_of_period"]) - ) - assert_array_equal(np.asarray(dj_wealth_grid), expected) - assert dj_wealth_grid.shape[0] == options["n_total_wealth_grid"] + assert options["n_total_wealth_grid"] == 11 + 1 + assert "dj_wealth_grid" not in options["continuous_states_info"] + assert options["upper_envelope"]["skip_endog_grid_storage"] is True + + +def test_none_assets_begin_of_period_defers_n_total_wealth_grid(valid_model_config): + # `None` means the grid is fully state-choice-specific -- no default array to + # read a length from, so n_total_wealth_grid is left unresolved (None) here, + # pinned later once a real state-choice can be evaluated against. + valid_model_config["upper_envelope"] = {"method": "druedahl_jorgensen"} + valid_model_config["continuous_states"]["assets_begin_of_period"] = None + options = check_model_config_and_process(valid_model_config) + + assert options["continuous_states_info"]["assets_begin_of_period"] is None + assert options["n_total_wealth_grid"] is None + # skip_endog_grid_storage itself only depends on method/n_choices, not on + # whether the grid is resolved yet, so it's still known immediately. assert options["upper_envelope"]["skip_endog_grid_storage"] is True diff --git a/tests/test_partial_and_interfaces.py b/tests/test_partial_and_interfaces.py index de6ec50e..a458b8e1 100644 --- a/tests/test_partial_and_interfaces.py +++ b/tests/test_partial_and_interfaces.py @@ -38,7 +38,7 @@ def _budget_constraint_two_cont( cont_a, cont_b, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, ): interest_factor = 1 + params["interest_rate"] @@ -46,7 +46,7 @@ def _budget_constraint_two_cont( params["wage_constant"] + params["wage_a"] * cont_a + params["wage_b"] * cont_b ) resource = interest_factor * asset_end_of_previous_period + ( - wage + income_shock_previous_period + wage + income_shock ) * (lagged_choice != 0) return jnp.maximum(resource, 0.5) diff --git a/tests/test_pre_processing.py b/tests/test_pre_processing.py index b8468cc5..b81485f4 100644 --- a/tests/test_pre_processing.py +++ b/tests/test_pre_processing.py @@ -78,6 +78,7 @@ def test_wrap_function(): util_processed = determine_function_arguments_and_partial_model_specs( utility_crra, options, + not_allowed_state_choices=[], ) calc_util = vmap(util_wrap, in_axes=(0, None, None))( @@ -212,6 +213,9 @@ def next_period_continuous_state( utility_functions=create_utility_function_dict(), utility_functions_final_period=create_final_period_utility_function_dict(), budget_constraint=budget_constraint, + stochastic_states_transitions=None, + shock_functions=None, + continuous_grid_functions=None, ) next_period_continuous_state = model_funcs["next_period_continuous_state"] diff --git a/tests/test_sim_varying_taste_shock_scale.py b/tests/test_sim_varying_taste_shock_scale.py index 838cd391..587dfb42 100644 --- a/tests/test_sim_varying_taste_shock_scale.py +++ b/tests/test_sim_varying_taste_shock_scale.py @@ -50,7 +50,7 @@ def model_setup(): "lagged_choice": np.zeros(n_agents), # all agents start as workers "married": np.zeros(n_agents), "ltc": np.zeros(n_agents), - "assets_begin_of_period": np.ones(n_agents) * 10, + "assets_end_of_previous_period": np.ones(n_agents) * 10, } df = model.get_solve_and_simulate_func(states_initial=initial_states, seed=seed)( diff --git a/tests/test_sim_with_aux.py b/tests/test_sim_with_aux.py index 3f642ea2..85913e3a 100644 --- a/tests/test_sim_with_aux.py +++ b/tests/test_sim_with_aux.py @@ -10,7 +10,7 @@ def budget_with_aux( period, lagged_choice, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, model_specs, params, ): @@ -18,7 +18,7 @@ def budget_with_aux( period, lagged_choice, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, model_specs, params, ) @@ -32,7 +32,7 @@ def budget_without_aux( period, lagged_choice, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, model_specs, params, ): @@ -40,7 +40,7 @@ def budget_without_aux( period, lagged_choice, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, model_specs, params, ) @@ -51,7 +51,7 @@ def budget_constraint_raw( period, lagged_choice, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, model_specs, params, ): @@ -59,7 +59,7 @@ def budget_constraint_raw( income_from_previous_period = _calc_stochastic_income( period=period, lagged_choice=lagged_choice, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, min_age=model_specs["min_age"], constant=params["constant"], exp=params["exp"], @@ -79,7 +79,7 @@ def budget_constraint_raw( return ( wealth_beginning_of_period, - income_shock_previous_period, + income_shock, income_from_previous_period, ) @@ -149,7 +149,7 @@ def test_sim_and_sol_model(model_config): states_initial = { "period": jnp.zeros(n_agents, dtype=int), "lagged_choice": jnp.zeros(n_agents, dtype=int), - "assets_begin_of_period": jnp.ones(n_agents, dtype=float) * 10, + "assets_end_of_previous_period": jnp.ones(n_agents, dtype=float) * 10, } n_periods = model_config["n_periods"] seed = 132 diff --git a/tests/test_simulate.py b/tests/test_simulate.py index 9b39e58e..a22fcad0 100644 --- a/tests/test_simulate.py +++ b/tests/test_simulate.py @@ -6,7 +6,7 @@ import jax.numpy as jnp import numpy as np import pytest -from numpy.testing import assert_array_almost_equal as aaae +from numpy.testing import assert_allclose import dcegm import dcegm.toy_models as toy_models @@ -72,7 +72,7 @@ def test_simulate(model_setup): "lagged_choice": np.zeros(n_agents), # all agents start as workers "married": np.zeros(n_agents), "ltc": np.zeros(n_agents), - "assets_begin_of_period": np.ones(n_agents) * 10, + "assets_end_of_previous_period": np.ones(n_agents) * 10, } df = model_setup["model_solved"].simulate( @@ -87,7 +87,11 @@ def test_simulate(model_setup): ) assert len(ids_violating_absorbing_retirement) == 0 - aaae(value_period_zero.mean(), expected.mean(), decimal=2) + # Monte-Carlo check of the Bellman identity: the realized continuation utility + # stands in for its expectation, so the two sides agree only up to sampling + # error. That error is ~0.03 across seeds at this many agents, hence the + # tolerance -- tightening it further needs more agents, not a smaller number. + assert_allclose(value_period_zero.mean(), expected.mean(), atol=0.1, rtol=0) def test_simulate_second_continuous_choice(model_setup): @@ -132,7 +136,7 @@ def test_simulate_second_continuous_choice(model_setup): "married": np.zeros(n_agents), "ltc": np.zeros(n_agents), "experience": np.ones(n_agents), - "assets_begin_of_period": np.ones(n_agents) * 10, + "assets_end_of_previous_period": np.ones(n_agents) * 10, } result = simulate_all_periods( @@ -146,6 +150,7 @@ def test_simulate_second_continuous_choice(model_setup): model_funcs=model_cont.model_funcs, model_config=model_cont.model_config, model_structure=model_cont.model_structure, + alt_model_funcs_sim=None, ) df = create_simulation_df(result) @@ -158,4 +163,8 @@ def test_simulate_second_continuous_choice(model_setup): ) assert len(ids_violating_absorbing_retirement) == 0 - aaae(value_period_zero.mean(), expected.mean(), decimal=2) + # Monte-Carlo check of the Bellman identity: the realized continuation utility + # stands in for its expectation, so the two sides agree only up to sampling + # error. That error is ~0.03 across seeds at this many agents, hence the + # tolerance -- tightening it further needs more agents, not a smaller number. + assert_allclose(value_period_zero.mean(), expected.mean(), atol=0.1, rtol=0) diff --git a/tests/test_simulate_continuous_state.py b/tests/test_simulate_continuous_state.py index 44fda7b6..09db79d7 100644 --- a/tests/test_simulate_continuous_state.py +++ b/tests/test_simulate_continuous_state.py @@ -87,7 +87,7 @@ def test_simulate_discrete_versus_continuous_experience(test_setup): "period": np.zeros(n_agents), "lagged_choice": np.zeros(n_agents), # all agents start as workers "experience": np.ones(n_agents), - "assets_begin_of_period": np.ones(n_agents) * 10, + "assets_end_of_previous_period": np.ones(n_agents) * 10, } df_disc = test_setup["model_solved_disc"].simulate( @@ -167,16 +167,25 @@ def test_simulation_passes_continuous_state_to_utility(model_solved_dj_exp_utili "period": np.zeros(n_agents, dtype=int), "lagged_choice": np.zeros(n_agents, dtype=int), "experience": experience_initial, - "assets_begin_of_period": np.array([10.0, 25.0, 10.0, 25.0]), + "assets_end_of_previous_period": np.array([10.0, 25.0, 10.0, 25.0]), } df = model_solved_dj_exp_utility.simulate(states_initial=states_initial, seed=111) df_first_period = df.xs(0, level="period") + # The first period's wealth comes out of its budget equation, so the interface is + # queried at the simulated states rather than at the assets carried in. + states_simulated = { + "period": np.zeros(n_agents, dtype=int), + "lagged_choice": np.zeros(n_agents, dtype=int), + "experience": df_first_period["experience"].to_numpy(), + "assets_begin_of_period": df_first_period["assets_begin_of_period"].to_numpy(), + } + for choice in range(N_DISCRETE_CHOICES): _policy, value = ( model_solved_dj_exp_utility.policy_and_value_for_states_and_choices( - states=states_initial, + states=states_simulated, choices=np.full(n_agents, choice), ) ) diff --git a/tests/test_state_space.py b/tests/test_state_space.py index b1c47aa6..523a7225 100644 --- a/tests/test_state_space.py +++ b/tests/test_state_space.py @@ -238,6 +238,7 @@ def test_state_space(): dict_of_state_space_objects = create_state_space( model_config=model_config_sparse, sparsity_condition=processed_sparsity_condition, + debugging=False, ) state_space = dict_of_state_space_objects["state_space"] diff --git a/tests/test_state_specific_grids_config.py b/tests/test_state_specific_grids_config.py new file mode 100644 index 00000000..7b70421f --- /dev/null +++ b/tests/test_state_specific_grids_config.py @@ -0,0 +1,776 @@ +"""Tests for state-specific continuous state grids: config plumbing and the child- +sharing consistency check. + +Grids live on the state-choice space, not the bare state space -- that's where the +solution itself lives (value_solved/policy_solved/endog_grid_solved are indexed by +state-choice). A grid may therefore depend on the discrete state *and* on "choice". + +""" + +import numpy as np +import pytest + +from dcegm.pre_processing.check_model_config import check_model_config_and_process +from dcegm.pre_processing.model_functions.process_model_functions import ( + process_continuous_grid_functions, + process_sparsity_condition, +) +from dcegm.pre_processing.model_structure.continuous_state_grids import ( + check_continuous_grid_consistency_across_shared_children, + evaluate_state_specific_continuous_grids, +) +from dcegm.pre_processing.model_structure.state_choice_space import ( + create_state_choice_space_and_child_state_mapping, +) +from dcegm.pre_processing.model_structure.state_space import create_state_space + + +def _base_model_config(): + return { + "n_periods": 3, + "n_quad_points": 1, + "choices": [0, 1], + "deterministic_states": {"group": [0, 1, 2]}, + "continuous_states": { + "assets_end_of_period": np.linspace(0, 1, 10), + "experience": np.linspace(0, 1, 4), + }, + } + + +# ===================================================================================== +# continuous_grid_functions plumbing +# +# continuous_grid_functions is a top-level argument to create_model_dict, following +# the same convention as shock_functions/stochastic_states_transitions -- functions +# are never nested inside model_config, which holds data/config only. +# ===================================================================================== + + +def test_continuous_grid_functions_absent_defaults_to_constant_grids(): + # Both "assets_end_of_period" and "experience" go through the constant-wrapper + # branch here (neither is state-specific), with two *different* default grids + # -- a regression test for a late-binding-closure-in-a-loop bug where every + # constant wrapper would have ended up returning the last-processed default + # grid instead of its own (see process_continuous_grid_functions). + processed = check_model_config_and_process(_base_model_config()) + continuous_grid_functions, state_specific_names = process_continuous_grid_functions( + continuous_grid_functions=None, model_config=processed, model_specs={} + ) + assert state_specific_names == [] + np.testing.assert_allclose( + continuous_grid_functions["experience"](period=0, lagged_choice=0, group=1), + np.linspace(0, 1, 4), + ) + np.testing.assert_allclose( + continuous_grid_functions["assets_end_of_period"]( + period=0, lagged_choice=0, group=1 + ), + np.linspace(0, 1, 10), + ) + + +def test_continuous_grid_functions_must_be_a_dict(): + processed = check_model_config_and_process(_base_model_config()) + with pytest.raises(ValueError, match="must be a dictionary"): + process_continuous_grid_functions( + continuous_grid_functions=[1, 2], model_config=processed, model_specs={} + ) + + +def test_continuous_grid_functions_rejects_unknown_state_name(): + processed = check_model_config_and_process(_base_model_config()) + with pytest.raises(ValueError, match="not a continuous state"): + process_continuous_grid_functions( + continuous_grid_functions={"not_a_state": lambda group: group}, + model_config=processed, + model_specs={}, + ) + + +def test_continuous_grid_functions_rejects_non_callable(): + processed = check_model_config_and_process(_base_model_config()) + with pytest.raises(ValueError, match="must be a callable"): + process_continuous_grid_functions( + continuous_grid_functions={"experience": np.zeros(4)}, + model_config=processed, + model_specs={}, + ) + + +def test_none_grid_without_matching_grid_function_raises(): + # A `None` grid means there is no default to fall back to -- it must be paired + # with a continuous_grid_functions entry, or there is nothing to build the grid + # from at all. + config = _base_model_config() + config["continuous_states"]["experience"] = None + processed = check_model_config_and_process(config) + with pytest.raises(ValueError, match="no matching continuous_grid_functions"): + process_continuous_grid_functions( + continuous_grid_functions=None, model_config=processed, model_specs={} + ) + + +def test_array_declared_grid_with_grid_function_raises(): + """The other direction of the None convention. + + A declared array is unused once a grid function takes over -- only its length ever + mattered, and that is now pinned by evaluating the grid function against a + representative state-choice (see continuous_state_grids.py). Leaving a real array + there is rejected rather than silently ignored, so the two directions together make + the convention unambiguous. + + """ + processed = check_model_config_and_process(_base_model_config()) + with pytest.raises(ValueError, match="is not None"): + process_continuous_grid_functions( + continuous_grid_functions={ + "experience": lambda group: np.linspace(0, 1, 4) + }, + model_config=processed, + model_specs={}, + ) + + +def test_none_grid_paired_with_grid_function_is_accepted(): + config = _base_model_config() + config["continuous_states"]["experience"] = None + processed = check_model_config_and_process(config) + + # No default array to read a length from -- pending until a representative + # state-choice can be evaluated against. + assert ( + processed["continuous_states_info"]["additional_continuous_state_grids"][ + "experience" + ] + is None + ) + assert processed["continuous_states_info"]["state_specific_size_names"] == [ + "experience" + ] + assert ( + processed["continuous_states_info"]["n_continuous_state_combinations"] is None + ) + + def grid_func(group): + return np.linspace(0, 1, 4) * (group + 1) + + continuous_grid_functions, state_specific_names = process_continuous_grid_functions( + continuous_grid_functions={"experience": grid_func}, + model_config=processed, + model_specs={}, + ) + assert state_specific_names == ["experience"] + np.testing.assert_allclose( + continuous_grid_functions["experience"]( + period=0, lagged_choice=0, group=1, choice=0 + ), + np.linspace(0, 1, 4) * 2, + ) + + +def _build_state_choice_space_pinning_grids(config, grid_functions): + """Process ``config`` and build the state-choice space, pinning deferred lengths. + + Returns the processed config, whose deferred sizes are filled in as a side effect of + building the state-choice space -- the only point at which a real state-choice + exists to evaluate a grid function against. + + """ + processed_config = check_model_config_and_process(config) + continuous_grid_functions, state_specific_names = process_continuous_grid_functions( + continuous_grid_functions=grid_functions, + model_config=processed_config, + model_specs={}, + ) + state_space_objects = create_state_space( + model_config=processed_config, + sparsity_condition=process_sparsity_condition( + state_space_functions={}, model_specs={} + ), + debugging=False, + ) + + def next_period_deterministic_state(**kwargs): + return { + "period": kwargs["period"] + 1, + "lagged_choice": kwargs["choice"], + "group": kwargs["group"], + } + + create_state_choice_space_and_child_state_mapping( + model_config=processed_config, + state_specific_choice_set=lambda **kwargs: np.array([0, 1]), + next_period_deterministic_state=next_period_deterministic_state, + state_space_arrays=state_space_objects, + continuous_grid_functions=continuous_grid_functions, + state_specific_continuous_grid_names=state_specific_names, + ) + return processed_config + + +def test_assets_end_of_period_follows_the_same_none_convention(): + """``assets_end_of_period`` is not exempt from the pairing rule. + + It used to be: ``check_model_config.py`` reads its length eagerly, for every + upper_envelope method, to size the fues tuning params, and that happens before + ``continuous_grid_functions`` is known. Those tuning params are now deferred + instead (see ``resolve_tuning_params_from_assets_end_of_period``), so the name + obeys the same convention as every other. + + """ + config = _base_model_config() + config["continuous_states"]["assets_end_of_period"] = None + processed = check_model_config_and_process(config) + + # Nothing to read a length from yet, so neither the grid nor anything derived + # from its length is resolved at this point. + assert processed["continuous_states_info"]["assets_grid_end_of_period"] is None + assert processed["continuous_states_info"]["n_assets_end_of_period"] is None + assert processed["n_total_wealth_grid"] is None + assert "n_total_wealth_grid" not in processed["upper_envelope"]["tuning_params"] + + with pytest.raises(ValueError, match="no matching continuous_grid_functions"): + process_continuous_grid_functions( + continuous_grid_functions=None, model_config=processed, model_specs={} + ) + + +def test_array_declared_assets_end_of_period_with_grid_function_raises(): + processed = check_model_config_and_process(_base_model_config()) + with pytest.raises(ValueError, match="is not None"): + process_continuous_grid_functions( + continuous_grid_functions={ + "assets_end_of_period": lambda group: np.linspace(0, 1, 10) + }, + model_config=processed, + model_specs={}, + ) + + +def test_state_specific_assets_end_of_period_resolves_fues_tuning_params(): + """The deferred fues sizes must come out exactly as the declared array's would. + + ``n_constrained_points_to_add`` and ``n_total_wealth_grid`` scale with the end-of- + period grid's length, so declaring the name ``None`` only works if they are filled + in once that length is pinned -- and the upper envelope closes over the tuning- + params dict before then, which is why they are mutated in place rather than + recomputed into a fresh dict. + + """ + declared = check_model_config_and_process(_base_model_config()) + + config = _base_model_config() + config["continuous_states"]["assets_end_of_period"] = None + resolved = _build_state_choice_space_pinning_grids( + config, + {"assets_end_of_period": lambda group: np.linspace(0, 1, 10)}, + ) + + assert resolved["continuous_states_info"]["n_assets_end_of_period"] == 10 + assert resolved["n_total_wealth_grid"] == declared["n_total_wealth_grid"] + for name in ("n_constrained_points_to_add", "n_total_wealth_grid"): + assert ( + resolved["upper_envelope"]["tuning_params"][name] + == declared["upper_envelope"]["tuning_params"][name] + ) + + +def test_none_grid_size_pinned_via_representative_state_choice(): + # No default array for "experience" at all -- n_continuous_state_combinations + # can't be known until create_state_choice_space_and_child_state_mapping + # evaluates the grid function against a real state-choice. + config = _base_model_config() + config["continuous_states"]["experience"] = None + processed_config = check_model_config_and_process(config) + assert ( + processed_config["continuous_states_info"]["n_continuous_state_combinations"] + is None + ) + + def grid_func(group): + return np.linspace(0, 1, 4) * (group + 1) + + continuous_grid_functions, state_specific_names = process_continuous_grid_functions( + continuous_grid_functions={"experience": grid_func}, + model_config=processed_config, + model_specs={}, + ) + + sparsity_condition = process_sparsity_condition( + state_space_functions={}, model_specs={} + ) + state_space_objects = create_state_space( + model_config=processed_config, + sparsity_condition=sparsity_condition, + debugging=False, + ) + + def next_period_deterministic_state(**kwargs): + return { + "period": kwargs["period"] + 1, + "lagged_choice": kwargs["choice"], + "group": kwargs["group"], + } + + def state_specific_choice_set(**kwargs): + return np.array([0, 1]) + + create_state_choice_space_and_child_state_mapping( + model_config=processed_config, + state_specific_choice_set=state_specific_choice_set, + next_period_deterministic_state=next_period_deterministic_state, + state_space_arrays=state_space_objects, + continuous_grid_functions=continuous_grid_functions, + state_specific_continuous_grid_names=state_specific_names, + ) + + # Pinned to 4 -- the length grid_func("experience") always returns, regardless + # of which state-choice it was evaluated against. + assert ( + processed_config["continuous_states_info"]["n_continuous_state_combinations"] + == 4 + ) + + +def test_none_assets_begin_of_period_size_pinned_via_representative_state_choice(): + # Same pinning mechanism, applied to assets_begin_of_period (Druedahl-Jorgensen's + # fixed common wealth grid) instead of an additional continuous state -- + # n_total_wealth_grid can't be known until a real state-choice is evaluated. + # _base_model_config() has two choices, so skip_endog_grid_storage is True and + # assets_begin_of_period state-specificity coexists fine with "experience" + # (an additional continuous state, left at its real default array here) -- see + # process_continuous_grid_functions and + # test_assets_begin_of_period_state_specific.py for the end-to-end version and + # the remaining (single-choice) restriction. + config = _base_model_config() + config["continuous_states"]["assets_begin_of_period"] = None + config["upper_envelope"] = {"method": "druedahl_jorgensen"} + processed_config = check_model_config_and_process(config) + assert processed_config["n_total_wealth_grid"] is None + + def assets_begin_of_period_grid_func(group): + return np.linspace(0, 10, 6) * (group + 1) + + continuous_grid_functions, state_specific_names = process_continuous_grid_functions( + continuous_grid_functions={ + "assets_begin_of_period": assets_begin_of_period_grid_func + }, + model_config=processed_config, + model_specs={}, + ) + + sparsity_condition = process_sparsity_condition( + state_space_functions={}, model_specs={} + ) + state_space_objects = create_state_space( + model_config=processed_config, + sparsity_condition=sparsity_condition, + debugging=False, + ) + + def next_period_deterministic_state(**kwargs): + return { + "period": kwargs["period"] + 1, + "lagged_choice": kwargs["choice"], + "group": kwargs["group"], + } + + def state_specific_choice_set(**kwargs): + return np.array([0, 1]) + + create_state_choice_space_and_child_state_mapping( + model_config=processed_config, + state_specific_choice_set=state_specific_choice_set, + next_period_deterministic_state=next_period_deterministic_state, + state_space_arrays=state_space_objects, + continuous_grid_functions=continuous_grid_functions, + state_specific_continuous_grid_names=state_specific_names, + ) + + # Pinned to len(grid) + 1 (expected value at zero wealth prepended), the same + # formula check_model_config.py uses for a directly-declared grid. No shared + # grid array is stored either way -- each reader recomputes its own via + # compute_own_dj_wealth_grid -- so only this length is resolved here. + assert processed_config["n_total_wealth_grid"] == 7 + + +def test_assets_begin_of_period_state_specific_with_single_choice_raises(): + # With a single discrete choice, skip_endog_grid_storage is False (see + # check_model_config.py): the Druedahl-Jorgensen upper envelope is skipped + # entirely, so the stored endogenous grid is not the fixed + # assets_begin_of_period grid there, and a state-specific + # assets_begin_of_period has nothing to plug into (unlike the >= 2-choices + # case exercised above and in test_assets_begin_of_period_state_specific.py). + config = _base_model_config() + config["choices"] = [0] + config["continuous_states"]["assets_begin_of_period"] = None + config["upper_envelope"] = {"method": "druedahl_jorgensen"} + processed_config = check_model_config_and_process(config) + assert not processed_config["upper_envelope"]["skip_endog_grid_storage"] + + def assets_begin_of_period_grid_func(group): + return np.linspace(0, 10, 6) * (group + 1) + + with pytest.raises(ValueError, match="skip_endog_grid_storage"): + process_continuous_grid_functions( + continuous_grid_functions={ + "assets_begin_of_period": assets_begin_of_period_grid_func + }, + model_config=processed_config, + model_specs={}, + ) + + +def test_process_continuous_grid_functions_wraps_state_specific_and_default_names(): + def grid_func(group): + return np.linspace(0, 1, 4) * (group + 1) + + config = _base_model_config() + config["continuous_states"]["experience"] = None + processed = check_model_config_and_process(config) + continuous_grid_functions, state_specific_names = process_continuous_grid_functions( + continuous_grid_functions={"experience": grid_func}, + model_config=processed, + model_specs={}, + ) + assert state_specific_names == ["experience"] + + # State-specific name: varies with the discrete state. + np.testing.assert_allclose( + continuous_grid_functions["experience"]( + period=0, lagged_choice=0, group=1, choice=0 + ), + np.linspace(0, 1, 4) * 2, + ) + np.testing.assert_allclose( + continuous_grid_functions["experience"]( + period=0, lagged_choice=0, group=2, choice=0 + ), + np.linspace(0, 1, 4) * 3, + ) + + # Default name: constant regardless of the discrete state passed in. + np.testing.assert_allclose( + continuous_grid_functions["assets_end_of_period"]( + period=0, lagged_choice=0, group=0, choice=0 + ), + np.linspace(0, 1, 10), + ) + np.testing.assert_allclose( + continuous_grid_functions["assets_end_of_period"]( + period=2, lagged_choice=1, group=2, choice=1 + ), + np.linspace(0, 1, 10), + ) + + +def test_process_continuous_grid_functions_allows_choice_dependence(): + # Grids live on the state-choice space, so "choice" is a legitimate parameter, + # unlike model_specs/state-only functions elsewhere in the codebase. + def grid_func(choice): + return np.linspace(0, 1, 4) * (choice + 1) + + config = _base_model_config() + config["continuous_states"]["experience"] = None + processed = check_model_config_and_process(config) + continuous_grid_functions, state_specific_names = process_continuous_grid_functions( + continuous_grid_functions={"experience": grid_func}, + model_config=processed, + model_specs={}, + ) + assert state_specific_names == ["experience"] + np.testing.assert_allclose( + continuous_grid_functions["experience"]( + period=0, lagged_choice=0, group=1, choice=1 + ), + np.linspace(0, 1, 4) * 2, + ) + + +def test_evaluate_state_specific_grids_rejects_wrong_length(): + state_choice_space = np.array([[0, 0, 0, 0], [0, 0, 1, 1]]) + discrete_state_choice_names = ["period", "lagged_choice", "group", "choice"] + continuous_states_info = { + "additional_continuous_state_grids": {"experience": np.linspace(0, 1, 4)}, + "assets_grid_end_of_period": np.linspace(0, 1, 10), + } + + def bad_grid_func(**kwargs): + return np.linspace(0, 1, 3) if kwargs["group"] == 0 else np.linspace(0, 1, 4) + + with pytest.raises(ValueError, match="must be a 1d array"): + evaluate_state_specific_continuous_grids( + state_choice_space=state_choice_space, + discrete_state_choice_names=discrete_state_choice_names, + continuous_grid_functions={"experience": bad_grid_func}, + state_specific_names=["experience"], + continuous_states_info=continuous_states_info, + ) + + +def test_evaluate_state_specific_grids_skips_names_without_user_callable(): + state_choice_space = np.array([[0, 0, 0, 0], [0, 0, 1, 1]]) + discrete_state_choice_names = ["period", "lagged_choice", "group", "choice"] + continuous_states_info = { + "additional_continuous_state_grids": {"experience": np.linspace(0, 1, 4)}, + "assets_grid_end_of_period": np.linspace(0, 1, 10), + } + + grids_per_state_choice, resolved_lengths = evaluate_state_specific_continuous_grids( + state_choice_space=state_choice_space, + discrete_state_choice_names=discrete_state_choice_names, + continuous_grid_functions={}, + state_specific_names=[], + continuous_states_info=continuous_states_info, + ) + assert grids_per_state_choice == {} + assert resolved_lengths == {} + + +# ===================================================================================== +# Consistency check across state-choices sharing a child state +# ===================================================================================== + + +def _build_state_choice_objects(grid_func_dict, group_resets_on_choice=False): + config = _base_model_config() + for name in grid_func_dict: + config["continuous_states"][name] = None + processed_config = check_model_config_and_process(config) + + continuous_grid_functions, state_specific_names = process_continuous_grid_functions( + continuous_grid_functions=grid_func_dict, + model_config=processed_config, + model_specs={}, + ) + + sparsity_condition = process_sparsity_condition( + state_space_functions={}, model_specs={} + ) + state_space_objects = create_state_space( + model_config=processed_config, + sparsity_condition=sparsity_condition, + debugging=False, + ) + + def next_period_deterministic_state(**kwargs): + # By default "group" is time invariant (passes through unchanged from + # parent to child), like sex/education would be. When + # group_resets_on_choice is True, it instead becomes a function of the + # choice -- i.e. it no longer passes through 1:1, so two different parents + # can share a child while disagreeing on their own "group". "choice" + # itself always passes through 1:1 into the child's own "lagged_choice" + # (enforced by check_endog_update_function in state_choice_space.py), so a + # grid depending only on "choice" can never violate the consistency check, + # regardless of what else is going on. + next_group = kwargs["choice"] % 3 if group_resets_on_choice else kwargs["group"] + return { + "period": kwargs["period"] + 1, + "lagged_choice": kwargs["choice"], + "group": next_group, + } + + def state_specific_choice_set(**kwargs): + return np.array([0, 1]) + + return ( + processed_config, + state_space_objects, + next_period_deterministic_state, + state_specific_choice_set, + continuous_grid_functions, + state_specific_names, + ) + + +def test_grid_depending_on_time_invariant_state_passes(): + def grid_func(group): + return np.linspace(0, 1, 4) * (group + 1) + + ( + processed_config, + state_space_objects, + next_period_deterministic_state, + state_specific_choice_set, + continuous_grid_functions, + state_specific_names, + ) = _build_state_choice_objects({"experience": grid_func}) + + # "group" is time invariant and shared by parent and child by construction, so + # any two parents sharing a child necessarily agree on it -- this must not raise. + create_state_choice_space_and_child_state_mapping( + model_config=processed_config, + state_specific_choice_set=state_specific_choice_set, + next_period_deterministic_state=next_period_deterministic_state, + state_space_arrays=state_space_objects, + continuous_grid_functions=continuous_grid_functions, + state_specific_continuous_grid_names=state_specific_names, + ) + + +# group_resets_on_choice pins the child's "group" to the choice made, so every state +# whose "group" disagrees with its "lagged_choice" is unreachable here by construction. +@pytest.mark.filterwarnings("ignore:(?s).*are not child states") +def test_grid_depending_on_choice_always_passes(): + # "choice" always passes through 1:1 into the child (it becomes the child's + # own "lagged_choice") -- so two parents sharing a child are guaranteed to + # have made the same choice, and a grid depending only on "choice" can never + # violate the consistency check. Combined with group_resets_on_choice=True, + # which *would* break a group-dependent grid (see the test below), to show + # this holds regardless of what else varies. + def grid_func(choice): + return np.linspace(0, 1, 4) * (choice + 1) + + ( + processed_config, + state_space_objects, + next_period_deterministic_state, + state_specific_choice_set, + continuous_grid_functions, + state_specific_names, + ) = _build_state_choice_objects( + {"experience": grid_func}, group_resets_on_choice=True + ) + + create_state_choice_space_and_child_state_mapping( + model_config=processed_config, + state_specific_choice_set=state_specific_choice_set, + next_period_deterministic_state=next_period_deterministic_state, + state_space_arrays=state_space_objects, + continuous_grid_functions=continuous_grid_functions, + state_specific_continuous_grid_names=state_specific_names, + ) + + +def test_grid_depending_on_lagged_choice_fails(): + def grid_func(lagged_choice): + return np.linspace(0, 1, 4) * (lagged_choice + 1) + + ( + processed_config, + state_space_objects, + next_period_deterministic_state, + state_specific_choice_set, + continuous_grid_functions, + state_specific_names, + ) = _build_state_choice_objects({"experience": grid_func}) + + # A parent with lagged_choice=0 and one with lagged_choice=1 both choosing + # choice=1 land on the same child (period+1, lagged_choice=1, same group), but + # disagree on their own grid (indexed by their own lagged_choice -- note this + # is the *parent's* pre-existing state, not "choice" itself, which always + # passes through consistently) -- must raise. + with pytest.raises(ValueError, match="different grids"): + create_state_choice_space_and_child_state_mapping( + model_config=processed_config, + state_specific_choice_set=state_specific_choice_set, + next_period_deterministic_state=next_period_deterministic_state, + state_space_arrays=state_space_objects, + continuous_grid_functions=continuous_grid_functions, + state_specific_continuous_grid_names=state_specific_names, + ) + + +def test_grid_on_state_that_does_not_pass_through_1to1_fails(): + def grid_func(group): + return np.linspace(0, 1, 4) * (group + 1) + + ( + processed_config, + state_space_objects, + next_period_deterministic_state, + state_specific_choice_set, + continuous_grid_functions, + state_specific_names, + ) = _build_state_choice_objects( + {"experience": grid_func}, group_resets_on_choice=True + ) + + # Now "group" is overwritten by (choice % 3) in the child, so two parents with + # different own "group" but the same choice share a child while disagreeing on + # their grid (indexed by their own "group") -- must raise. + with pytest.raises(ValueError, match="different grids"): + create_state_choice_space_and_child_state_mapping( + model_config=processed_config, + state_specific_choice_set=state_specific_choice_set, + next_period_deterministic_state=next_period_deterministic_state, + state_space_arrays=state_space_objects, + continuous_grid_functions=continuous_grid_functions, + state_specific_continuous_grid_names=state_specific_names, + ) + + +def test_no_state_specific_grids_is_a_no_op(): + ( + processed_config, + state_space_objects, + next_period_deterministic_state, + state_specific_choice_set, + _, + _, + ) = _build_state_choice_objects({}) + + create_state_choice_space_and_child_state_mapping( + model_config=processed_config, + state_specific_choice_set=state_specific_choice_set, + next_period_deterministic_state=next_period_deterministic_state, + state_space_arrays=state_space_objects, + continuous_grid_functions=None, + state_specific_continuous_grid_names=None, + ) + + +def test_check_function_directly_with_hand_built_mapping(): + # Two state-choice rows (0 and 1) share child state index 5, and disagree on + # their own grid for "experience" -- the check should name both rows. + state_choice_space = np.array( + [ + [0, 0, 0, 0], # period, lagged_choice, group, choice + [0, 1, 0, 0], + ] + ) + discrete_states_names = ["period", "lagged_choice", "group"] + map_state_choice_to_child_states = np.array([[5], [5]]) + grids_per_state_choice = { + "experience": np.array( + [ + [0.0, 1.0], # row 0's own grid + [0.0, 2.0], # row 1's own grid -- disagrees + ] + ) + } + + with pytest.raises(ValueError, match="child state index 5"): + check_continuous_grid_consistency_across_shared_children( + state_choice_space=state_choice_space, + discrete_states_names=discrete_states_names, + map_state_choice_to_child_states=map_state_choice_to_child_states, + grids_per_state_choice=grids_per_state_choice, + ) + + +def test_check_function_directly_passes_when_grids_agree(): + state_choice_space = np.array( + [ + [0, 0, 0, 0], + [0, 1, 0, 0], + ] + ) + discrete_states_names = ["period", "lagged_choice", "group"] + map_state_choice_to_child_states = np.array([[5], [5]]) + grids_per_state_choice = { + "experience": np.array( + [ + [0.0, 1.0], + [0.0, 1.0], + ] + ) + } + + check_continuous_grid_consistency_across_shared_children( + state_choice_space=state_choice_space, + discrete_states_names=discrete_states_names, + map_state_choice_to_child_states=map_state_choice_to_child_states, + grids_per_state_choice=grids_per_state_choice, + ) diff --git a/tests/test_state_specific_grids_end_to_end.py b/tests/test_state_specific_grids_end_to_end.py new file mode 100644 index 00000000..baf79889 --- /dev/null +++ b/tests/test_state_specific_grids_end_to_end.py @@ -0,0 +1,460 @@ +"""End-to-end tests for state-choice-specific continuous grids. + +All three grid names go through the same four checks, so they are written once and +parametrized over ``CASES`` rather than duplicated per name and per upper-envelope +method: + +``experience`` + An *additional continuous state*. Plays the law-of-motion-input role: a + representative parent's own grid feeds the child's transition. +``assets_end_of_period`` + The exogenous savings grid. Also a law-of-motion input, and additionally each + state-choice's own EGM grid. +``assets_begin_of_period`` + Druedahl-Jorgensen's fixed common wealth grid ("m_grid"). Not a law-of-motion + input at all -- it is the *output* storage grid, so it needs own-grid + (self-referential) threading rather than a representative parent. Only + reachable when ``skip_endog_grid_storage`` is True. + +The four checks: + +1. A grid function returning the model's *own default* must reproduce the plain + solve bit-for-bit -- the feature must be a no-op when it changes nothing. +2. A constant-but-*different* grid must reproduce a model declaring that same grid + directly. This is the strong one: it fails both if the wrong grid is used and if + the grid function is silently ignored. +3. Same, through ``simulate()`` -- a separate reader path from solve. +4. Same, through ``choice_values_for_states``/``choice_policies_for_states`` -- a + third reader path, distinct from both. + +Unit-level config validation lives in ``test_state_specific_grids_config.py``, the +law-of-motion mechanics in ``test_state_specific_grids_law_of_motion.py``, and +checks against independently computed solutions in +``test_state_specific_grids_reference.py``. + +""" + +import jax.numpy as jnp +import numpy as np +import pytest + +import dcegm +import dcegm.toy_models as toy_models +from dcegm.toy_models.cons_ret_model_dcegm_paper.utility_functions import ( + inverse_marginal_utility_crra, + marginal_utility_crra, +) + +N_AGENTS = 1_000 +N_QUERY_STATES = 20 + +# ===================================================================================== +# Model loaders +# ===================================================================================== + + +def _with_cont_exp_fues(): + """Continuous experience, FUES upper envelope (the toy model's default).""" + model_funcs = toy_models.load_example_model_functions("with_cont_exp") + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config("with_cont_exp") + ) + return model_funcs, params, model_specs, dict(model_config), "continuous" + + +def _with_cont_exp_dj(): + """Continuous experience, Druedahl-Jorgensen upper envelope.""" + model_funcs, params, model_specs, model_config, exp_kind = _with_cont_exp_fues() + model_config["continuous_states"] = dict(model_config["continuous_states"]) + model_config["continuous_states"]["assets_begin_of_period"] = jnp.linspace( + 0, 50, 50 + ) + model_config["upper_envelope"] = {"method": "druedahl_jorgensen"} + return model_funcs, params, model_specs, model_config, exp_kind + + +def _with_exp_dj(): + """Discrete experience -- no additional continuous state -- plus DJ. + + The only configuration in which a state-specific ``assets_begin_of_period`` is + supported alongside a single-dimensional wealth axis. + + """ + model_funcs = toy_models.load_example_model_functions("with_exp") + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config("with_exp") + ) + model_config = dict(model_config) + model_config["continuous_states"] = dict(model_config["continuous_states"]) + model_config["continuous_states"]["assets_begin_of_period"] = jnp.linspace( + 0, 50, 50 + ) + model_config["upper_envelope"] = {"method": "druedahl_jorgensen"} + return model_funcs, params, model_specs, model_config, "discrete" + + +# (grid name, loader). Every combination that is actually supported -- see the +# module docstring for why assets_begin_of_period has no FUES entry. +CASES = { + "experience__fues": ("experience", _with_cont_exp_fues), + "experience__dj": ("experience", _with_cont_exp_dj), + "assets_end_of_period__fues": ("assets_end_of_period", _with_cont_exp_fues), + "assets_end_of_period__dj": ("assets_end_of_period", _with_cont_exp_dj), + "assets_begin_of_period__dj": ("assets_begin_of_period", _with_exp_dj), + "assets_begin_of_period__dj_multidim": ( + "assets_begin_of_period", + _with_cont_exp_dj, + ), +} + + +# ===================================================================================== +# Shared helpers +# ===================================================================================== + + +def _declared_grid(model_config, name): + return jnp.asarray(model_config["continuous_states"][name]) + + +def _with_declared_grid(model_config, name, grid): + config = dict(model_config) + config["continuous_states"] = dict(model_config["continuous_states"]) + config["continuous_states"][name] = grid + return config + + +def _config_for_grid_function(model_config, name): + """Config to pair with a ``continuous_grid_functions`` entry for ``name``. + + A declared array is unused once a grid function takes over and must be set to + ``None`` to say so, for every name alike. See + ``process_continuous_grid_functions``. + + """ + return _with_declared_grid(model_config, name, None) + + +def _constant(grid): + def grid_func(period): + return grid + + return grid_func + + +def _initial_states(exp_kind, n): + experience = np.zeros(n, dtype=int) if exp_kind == "discrete" else np.ones(n) * 0.5 + return { + "period": np.zeros(n, dtype=int), + "lagged_choice": np.zeros(n, dtype=int), + "experience": experience, + "assets_end_of_previous_period": np.ones(n) * 10, + } + + +def _query_states(exp_kind, n): + experience = ( + np.arange(n, dtype=int) % 3 + if exp_kind == "discrete" + else np.linspace(0.0, 1.0, n) + ) + return { + "period": np.zeros(n, dtype=int), + "lagged_choice": np.zeros(n, dtype=int), + "experience": experience, + "assets_begin_of_period": np.ones(n) * 10, + } + + +_SOLVE_CACHE = {} + + +def _solve_pair(case, scale): + """Memoized ``_solve_pair_uncached``. + + Checks 2-4 all interrogate the *same* pair of solved models through three different + readers, so without this each ``(case, scale)`` would be solved once per check -- + four solves of every model instead of one. Solutions are read-only here, so sharing + them across tests is safe, and it is what keeps a six-case matrix affordable. + + """ + key = (case, scale) + if key not in _SOLVE_CACHE: + _SOLVE_CACHE[key] = _solve_pair_uncached(case, scale) + return _SOLVE_CACHE[key] + + +def _solve_pair_uncached(case, scale): + """Solve the same model twice: grid declared directly vs via a grid function. + + Returns ``(reference_solved, state_specific_solved, exp_kind)``. With ``scale == + 1.0`` the reference is just the unmodified model, which is check 1; with any other + scale it is a model declaring the scaled grid directly, which is checks 2-4. + + """ + name, loader = CASES[case] + model_funcs, params, model_specs, model_config, exp_kind = loader() + + target_grid = _declared_grid(model_config, name) * scale + + reference_config = ( + model_config + if scale == 1.0 + else _with_declared_grid(model_config, name, target_grid) + ) + reference_solved = dcegm.setup_model( + model_config=reference_config, model_specs=model_specs, **model_funcs + ).solve(params) + + state_specific_solved = dcegm.setup_model( + model_config=_config_for_grid_function(model_config, name), + model_specs=model_specs, + continuous_grid_functions={name: _constant(target_grid)}, + **model_funcs, + ).solve(params) + + return reference_solved, state_specific_solved, exp_kind + + +# ===================================================================================== +# The four checks, over every supported (grid name, upper envelope) combination +# ===================================================================================== + + +@pytest.mark.parametrize("case", list(CASES)) +def test_constant_grid_func_reproduces_default_solve_bit_for_bit(case): + """A grid function returning the default must change nothing at all.""" + reference_solved, state_specific_solved, _ = _solve_pair(case, scale=1.0) + + np.testing.assert_array_equal( + np.asarray(reference_solved.value), np.asarray(state_specific_solved.value) + ) + np.testing.assert_array_equal( + np.asarray(reference_solved.policy), np.asarray(state_specific_solved.policy) + ) + + +@pytest.mark.parametrize("case", list(CASES)) +def test_constant_but_different_grid_matches_direct_declaration(case): + """The strong check: a different grid, delivered two ways, must agree exactly.""" + reference_solved, state_specific_solved, _ = _solve_pair(case, scale=2.0) + + np.testing.assert_array_equal( + np.asarray(reference_solved.value), np.asarray(state_specific_solved.value) + ) + np.testing.assert_array_equal( + np.asarray(reference_solved.policy), np.asarray(state_specific_solved.policy) + ) + + +@pytest.mark.parametrize("case", list(CASES)) +def test_constant_but_different_grid_matches_direct_declaration_when_simulated(case): + """Same contract through ``simulate()``, a separate reader from solve. + + A bug confined to the simulation reader would leave the solve-level test above green + while silently corrupting simulated output. + + """ + reference_solved, state_specific_solved, exp_kind = _solve_pair(case, scale=2.0) + states_initial = _initial_states(exp_kind, N_AGENTS) + + df_reference = reference_solved.simulate(states_initial=states_initial, seed=111) + df_state_specific = state_specific_solved.simulate( + states_initial=states_initial, seed=111 + ) + + for column in df_reference.columns: + np.testing.assert_array_equal( + df_reference[column].to_numpy(), + df_state_specific[column].to_numpy(), + err_msg=f"column {column}", + ) + + +@pytest.mark.parametrize("case", list(CASES)) +def test_constant_but_different_grid_matches_direct_declaration_for_choice_queries( + case, +): + """Same contract through the ``choice_*_for_states`` readers, a third path.""" + reference_solved, state_specific_solved, exp_kind = _solve_pair(case, scale=2.0) + states = _query_states(exp_kind, N_QUERY_STATES) + + np.testing.assert_array_equal( + np.asarray(reference_solved.choice_values_for_states(states)), + np.asarray(state_specific_solved.choice_values_for_states(states)), + ) + np.testing.assert_array_equal( + np.asarray(reference_solved.choice_policies_for_states(states)), + np.asarray(state_specific_solved.choice_policies_for_states(states)), + ) + + +# ===================================================================================== +# Cases that do not fit the matrix +# ===================================================================================== + + +@pytest.mark.parametrize("case", ["experience__fues", "assets_end_of_period__fues"]) +def test_period_dependent_grid_solves_and_differs_from_default(case): + """A genuinely period-varying grid -- the parent/child distinction matters. + + Safe under ``check_continuous_grid_consistency_across_shared_children`` (period + increments deterministically for every parent of a given child), but the + parent's own grid differs from the child's. This is exactly what breaks if the + child's identity is wrongly used to select the grid. + + """ + name, loader = CASES[case] + model_funcs, params, model_specs, model_config, _ = loader() + default_grid = _declared_grid(model_config, name) + + def period_dependent_grid_func(period): + return default_grid * (1.0 + 0.1 * period) + + varying_value = np.asarray( + dcegm.setup_model( + model_config=_config_for_grid_function(model_config, name), + model_specs=model_specs, + continuous_grid_functions={name: period_dependent_grid_func}, + **model_funcs, + ) + .solve(params) + .value + ) + baseline_value = np.asarray( + dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ) + .solve(params) + .value + ) + + # dcegm pads variable-length endogenous grids with NaN to a common width; that + # padding pattern is a storage convention and must be unchanged, while the + # valid entries must be finite and actually different from the baseline. + baseline_nan = np.isnan(baseline_value) + varying_nan = np.isnan(varying_value) + np.testing.assert_array_equal(baseline_nan, varying_nan) + assert np.all(np.isfinite(varying_value[~varying_nan])) + assert not np.allclose(varying_value[~varying_nan], baseline_value[~baseline_nan]) + + +def test_state_specific_grid_by_group_matches_separate_per_group_models(): + """Ground truth: one model with a per-group grid == two single-group models. + + ``group`` enters neither utility, budget, nor law of motion, so the two groups' + economics are identical except for which grid their solution is stored and + interpolated on. A combined solve must therefore reproduce each separate solve + exactly -- an equivalence that holds only if the per-group grid is threaded + correctly all the way through. + + """ + model_funcs, params, model_specs, model_config, _ = _with_cont_exp_fues() + + default_grid = _declared_grid(model_config, "experience") + group1_grid = default_grid * 2.0 + + def grid_func(group): + return jnp.where(group == 0, default_grid, group1_grid) + + grouped_config = _with_declared_grid(model_config, "experience", None) + grouped_config["deterministic_states"] = {"group": [0, 1]} + grouped_solved = dcegm.setup_model( + model_config=grouped_config, + model_specs=model_specs, + continuous_grid_functions={"experience": grid_func}, + **model_funcs, + ).solve(params) + + solved_per_group = [ + dcegm.setup_model( + model_config=_with_declared_grid(model_config, "experience", grid), + model_specs=model_specs, + **model_funcs, + ).solve(params) + for grid in (default_grid, group1_grid) + ] + + base_states = _query_states("continuous", 15) + for group, solved_separately in enumerate(solved_per_group): + states = {**base_states, "group": np.full(15, group, dtype=int)} + np.testing.assert_array_equal( + np.asarray(grouped_solved.choice_values_for_states(states)), + np.asarray(solved_separately.choice_values_for_states(base_states)), + ) + np.testing.assert_array_equal( + np.asarray(grouped_solved.choice_policies_for_states(states)), + np.asarray(solved_separately.choice_policies_for_states(base_states)), + ) + + +def _utility_crra_with_experience(consumption, choice, experience, params): + """CRRA utility that depends on the continuous state *directly*. + + ``with_cont_exp``'s own utility takes only ``consumption`` and ``choice``, so + ``determine_function_arguments_and_partial_model_specs`` filters "experience" + out before it reaches the function -- meaning a wrong experience value would + have no effect on the output at all. This restores that dependence. + + """ + import jax + + utility_consumption = jax.lax.select( + jnp.allclose(params["rho"], 1), + jnp.log(consumption), + (consumption ** (1 - params["rho"]) - 1) / (1 - params["rho"]), + ) + return ( + utility_consumption + - (1 - choice) * params["delta"] + + params["exp"] * experience + ) + + +def test_upper_envelope_refinement_uses_each_state_choices_own_grid(): + """Regression test for the upper-envelope step reading the shared grid. + + ``run_upper_envelope`` refines the EGM candidates *after* they are generated, + and used to feed the model-wide default grid into the upper envelope's + ``value_function`` regardless of any per-state-choice override. + + Two things are needed to make that observable, which is why this cannot be + folded into the matrix above: a utility function that depends on "experience" + directly (see above), and the Druedahl-Jorgensen upper envelope, which + evaluates ``value_function`` unconditionally at every point. FUES only calls it + when a credit-constrained non-monotonicity is present, which this + parameterization does not trigger. + + """ + model_funcs, params, model_specs, model_config, _ = _with_cont_exp_dj() + model_funcs = dict(model_funcs) + model_funcs["utility_functions"] = { + "utility": _utility_crra_with_experience, + "marginal_utility": marginal_utility_crra, + "inverse_marginal_utility": inverse_marginal_utility_crra, + } + + scaled_grid = _declared_grid(model_config, "experience") * 2.0 + + reference_solved = dcegm.setup_model( + model_config=_with_declared_grid(model_config, "experience", scaled_grid), + model_specs=model_specs, + **model_funcs, + ).solve(params) + state_specific_solved = dcegm.setup_model( + model_config=_with_declared_grid(model_config, "experience", None), + model_specs=model_specs, + continuous_grid_functions={"experience": _constant(scaled_grid)}, + **model_funcs, + ).solve(params) + + np.testing.assert_array_equal( + np.asarray(reference_solved.value), np.asarray(state_specific_solved.value) + ) + np.testing.assert_array_equal( + np.asarray(reference_solved.policy), np.asarray(state_specific_solved.policy) + ) + np.testing.assert_array_equal( + np.asarray(reference_solved.endog_grid), + np.asarray(state_specific_solved.endog_grid), + ) diff --git a/tests/test_state_specific_grids_law_of_motion.py b/tests/test_state_specific_grids_law_of_motion.py new file mode 100644 index 00000000..45b09739 --- /dev/null +++ b/tests/test_state_specific_grids_law_of_motion.py @@ -0,0 +1,334 @@ +"""Law-of-motion tests for the grid/granularity behaviour added on this branch. + +Two closely related concerns live here, both unit-level tests of +``law_of_motion.py``: + +1. **Whose grid feeds the transition** -- the representative-parent selection. + A child's transition is evaluated over the *parent's* own continuous grid, + not the child's, which only matters once grids are state-choice-specific. + +2. **At what granularity it is evaluated** -- the state-level dedup. The + transition into a child does not depend on the child's own future choice, so + unless a user transition function declares ``choice``, + ``calc_law_of_motion`` evaluates once per unique child *state* and gathers the + result out instead of recomputing it per state-choice. That is purely a cost + optimization, so the central test is an equivalence check: forcing the + per-state-choice path by adding an otherwise-unused ``choice`` argument to the + budget equation must reproduce the deduplicated path bit-for-bit. + +End-to-end consequences of state-specific grids are in +``test_state_specific_grids_end_to_end.py``; the config layer is in +``test_state_specific_grids_config.py``. + +""" + +import jax.numpy as jnp +import numpy as np +from jax import vmap + +import dcegm +import dcegm.toy_models as toy_models +from dcegm.law_of_motion import ( + _continuous_state_next_period_for_one_state, + _get_continuous_state_next_period, + compute_own_continuous_grid_combos, +) +from dcegm.toy_models.cons_ret_model_with_cont_exp.budget_constraint import ( + budget_constraint_cont_exp, +) + +# ===================================================================================== +# Representative-parent grid selection +# ===================================================================================== + + +def test_compute_own_continuous_grid_combos_single_continuous_state(): + state_dict = {"group": jnp.array(1)} + + def grid_func(group): + return jnp.where(group == 0, jnp.array([0.0, 1.0]), jnp.array([10.0, 11.0])) + + combos = compute_own_continuous_grid_combos( + state_dict=state_dict, + continuous_grid_functions={"experience": grid_func}, + additional_continuous_state_names=["experience"], + ) + np.testing.assert_allclose(combos["experience"], jnp.array([10.0, 11.0])) + + +def test_compute_own_continuous_grid_combos_meshes_multiple_names(): + state_dict = {"group": jnp.array(0)} + grid_functions = { + "a": lambda group: jnp.array([0.0, 1.0]), + "b": lambda group: jnp.array([10.0, 20.0, 30.0]), + } + + combos = compute_own_continuous_grid_combos( + state_dict=state_dict, + continuous_grid_functions=grid_functions, + additional_continuous_state_names=["a", "b"], + ) + expected_a, expected_b = jnp.meshgrid( + jnp.array([0.0, 1.0]), jnp.array([10.0, 20.0, 30.0]), indexing="ij" + ) + np.testing.assert_allclose(combos["a"], expected_a.ravel()) + np.testing.assert_allclose(combos["b"], expected_b.ravel()) + + +def test_continuous_state_next_period_for_one_state_uses_each_states_own_grid(): + # Two states (group 0 and group 1), each with its own 2-point "experience" grid. + # Exercises the actual production function (vmapped exactly as + # _get_continuous_state_next_period does), not a standalone reference + # implementation. + state_space_dict = {"group": jnp.array([0, 1])} + + def grid_func(group): + return jnp.where(group == 0, jnp.array([0.0, 1.0]), jnp.array([10.0, 11.0])) + + def compute_continuous_state(group, experience, params): + return {"experience": experience + 100.0 * group} + + result = vmap( + _continuous_state_next_period_for_one_state, + in_axes=(0, 0, None, None, None, None), + )( + state_space_dict, + state_space_dict, # representative_parent_state_choice_vec == state_space_dict: no parent/child distinction being tested here, see the dedicated test for that below. + {"experience": grid_func}, + ["experience"], + {}, + compute_continuous_state, + ) + + # State 0 (group=0) transitions from its own grid [0, 1]. + np.testing.assert_allclose(result["experience"][0], jnp.array([0.0, 1.0])) + # State 1 (group=1) transitions from its own, different grid [10, 11] -- if the + # grid were wrongly shared/broadcast, this would come out as [10, 11] + 100, not + # distinguishable from state 0 having (wrongly) used the same grid. + np.testing.assert_allclose(result["experience"][1], jnp.array([110.0, 111.0])) + + +def test_continuous_state_next_period_for_one_state_constant_grid_matches_outer_product(): + # Regression check: when every state's grid_func returns the same grid + # (today's global-grid behavior), the per-state fused computation must + # reproduce exactly the outer-product result the old broadcast vmap produced. + state_space_dict = {"group": jnp.array([0, 1, 2])} + shared_grid = jnp.array([0.0, 1.0, 2.0, 3.0]) + + def grid_func(group): + return shared_grid + + def compute_continuous_state(group, experience, params): + return {"experience": experience + group} + + result = vmap( + _continuous_state_next_period_for_one_state, + in_axes=(0, 0, None, None, None, None), + )( + state_space_dict, + state_space_dict, + {"experience": grid_func}, + ["experience"], + {}, + compute_continuous_state, + ) + + expected = shared_grid[None, :] + jnp.array([0, 1, 2])[:, None] + np.testing.assert_allclose(result["experience"], expected) + + +def test_get_continuous_state_next_period_dummy_path_unaffected(): + # Models without an additional continuous state never touch + # continuous_grid_functions at all. + state_space_dict = {"group": jnp.array([0, 1])} + + result = _get_continuous_state_next_period( + has_additional_continuous_states=False, + child_state_choices=state_space_dict, + representative_last_period_parent_states=state_space_dict, + additional_continuous_state_names=[], + params={}, + model_funcs={}, + ) + assert result["dummy_cont"].shape == (2, 1) + + +def test_get_continuous_state_next_period_uses_representative_parent_not_state_space_dict(): + # state_space_dict is the CHILD's own identity (used for the transition function + # call itself); representative_parent_state_choice_vec is a representative PARENT's + # identity (used only to pick which grid to feed in). These must be allowed to + # differ -- this is the exact bug this test guards against: using the child's + # identity to select the grid instead of the parent's. + state_space_dict = {"group": jnp.array([9, 9])} # child's own group -- irrelevant + representative_parent_state_choice_vec = { + "group": jnp.array([0, 1]) + } # representative parent's group + + def grid_func(group): + return jnp.where(group == 0, jnp.array([0.0, 1.0]), jnp.array([10.0, 11.0])) + + def compute_continuous_state(group, experience, params): + return {"experience": experience} + + result = _get_continuous_state_next_period( + has_additional_continuous_states=True, + child_state_choices=state_space_dict, + representative_last_period_parent_states=representative_parent_state_choice_vec, + additional_continuous_state_names=["experience"], + params={}, + model_funcs={ + "continuous_grid_functions": {"experience": grid_func}, + "next_period_continuous_state": compute_continuous_state, + }, + ) + # If the child's identity (group=9 for both rows) had wrongly been used for grid + # selection, grid_func(9) would give [10, 11] for both rows. + np.testing.assert_allclose(result["experience"][0], jnp.array([0.0, 1.0])) + np.testing.assert_allclose(result["experience"][1], jnp.array([10.0, 11.0])) + + +# ===================================================================================== +# Evaluation granularity: state-level dedup vs per state-choice +# ===================================================================================== + + +def _load_with_cont_exp(): + model_funcs = toy_models.load_example_model_functions("with_cont_exp") + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config("with_cont_exp") + ) + return model_funcs, params, model_specs, model_config + + +def _budget_constraint_with_unused_choice( + period, + lagged_choice, + choice, + experience, + asset_end_of_previous_period, + income_shock, + params, + model_specs, +): + """Identical economics to budget_constraint_cont_exp, but declaring ``choice``. + + ``choice`` is deliberately unused in the body: it only changes which granularity the + law of motion is evaluated at (per state-choice rather than per unique child state), + which must not change the result. + + """ + return budget_constraint_cont_exp( + period=period, + lagged_choice=lagged_choice, + experience=experience, + asset_end_of_previous_period=asset_end_of_previous_period, + income_shock=income_shock, + params=params, + model_specs=model_specs, + ) + + +def test_default_toy_models_take_the_state_level_fast_path(): + # None of the shipped toy models' transition functions declare "choice", so + # they must all take the deduplicated path -- otherwise the rest of the test + # suite would never exercise it. + for name in ["with_cont_exp", "with_exp", "dcegm_paper"]: + model_funcs = toy_models.load_example_model_functions(name) + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config(name) + ) + model = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ) + assert not model.model_funcs["budget_depends_on_choice"], name + + +def test_choice_in_budget_signature_selects_the_state_choice_path(): + model_funcs, params, model_specs, model_config = _load_with_cont_exp() + model_funcs = dict(model_funcs) + model_funcs["budget_constraint"] = _budget_constraint_with_unused_choice + + model = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ) + assert model.model_funcs["budget_depends_on_choice"] + + +def test_state_level_and_state_choice_level_paths_agree_bit_for_bit(): + # The core equivalence check: a budget equation that declares (but ignores) + # "choice" forces the per-state-choice path, and must reproduce the + # deduplicated per-state path exactly. A dedup/gather bug -- e.g. a wrong + # state_row_for_state_choice mapping -- would show up here as mismatched + # value/policy, since it would feed each child the wrong state's transition. + model_funcs, params, model_specs, model_config = _load_with_cont_exp() + + fast_path_model = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ) + assert not fast_path_model.model_funcs["budget_depends_on_choice"] + fast_path_solved = fast_path_model.solve(params) + + slow_path_funcs = dict(model_funcs) + slow_path_funcs["budget_constraint"] = _budget_constraint_with_unused_choice + slow_path_model = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **slow_path_funcs + ) + assert slow_path_model.model_funcs["budget_depends_on_choice"] + slow_path_solved = slow_path_model.solve(params) + + np.testing.assert_array_almost_equal( + np.asarray(fast_path_solved.value), np.asarray(slow_path_solved.value) + ) + np.testing.assert_array_almost_equal( + np.asarray(fast_path_solved.policy), np.asarray(slow_path_solved.policy) + ) + np.testing.assert_array_almost_equal( + np.asarray(fast_path_solved.endog_grid), + np.asarray(slow_path_solved.endog_grid), + ) + + +def test_genuinely_choice_dependent_budget_differs_from_choice_free_one(): + # Sensitivity check for the test above: confirms a budget equation that + # actually *uses* choice produces a different solution, so the bit-for-bit + # agreement there reflects the two paths genuinely computing the same thing, + # not "choice" being unable to affect the budget at all. + model_funcs, params, model_specs, model_config = _load_with_cont_exp() + + def budget_constraint_using_choice( + period, + lagged_choice, + choice, + experience, + asset_end_of_previous_period, + income_shock, + params, + model_specs, + ): + base = budget_constraint_cont_exp( + period=period, + lagged_choice=lagged_choice, + experience=experience, + asset_end_of_previous_period=asset_end_of_previous_period, + income_shock=income_shock, + params=params, + model_specs=model_specs, + ) + return base + 0.5 * choice + + baseline_model = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ) + baseline_solved = baseline_model.solve(params) + + choice_dependent_funcs = dict(model_funcs) + choice_dependent_funcs["budget_constraint"] = budget_constraint_using_choice + choice_dependent_model = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **choice_dependent_funcs + ) + choice_dependent_solved = choice_dependent_model.solve(params) + + baseline_value = np.asarray(baseline_solved.value) + choice_dependent_value = np.asarray(choice_dependent_solved.value) + finite = np.isfinite(baseline_value) & np.isfinite(choice_dependent_value) + assert not np.allclose(baseline_value[finite], choice_dependent_value[finite]) diff --git a/tests/test_state_specific_grids_reference.py b/tests/test_state_specific_grids_reference.py new file mode 100644 index 00000000..e5b57f7c --- /dev/null +++ b/tests/test_state_specific_grids_reference.py @@ -0,0 +1,1068 @@ +"""Absolute checks: dcegm against independently implemented solutions. + +Every other test for the features on this branch is *relative* -- one dcegm solve +compared against another dcegm solve (two configurations that must agree, or two +evaluation granularities that must agree). Those catch regressions, but share a +blind spot: an error made identically on both sides is invisible to them. + +The tests here pin dcegm against solutions computed outside it, in two flavours: + +**Closed form** (first half). A two-period model where the Euler equation is exact: +period 1 is terminal, so consumption equals wealth and the marginal utility of +wealth is ``wealth ** -rho`` regardless of the period-1 choice. Period 0 then +satisfies, for each savings point ``a`` on that state-choice's *own* grid:: + + RHS = discount_factor * (1 + r) * E_shock[ sum_k P(k) * wealth_1(k) ** -rho ] + c_0 = RHS ** (-1 / rho) + wealth_0 = a + c_0 + +with ``P(k)`` the logit choice probabilities over period-1 choices (see +``aggregate_marg_utils_and_exp_values``). + +**Hand-solved reference** (second half). The divorce model's +``tests/resources/divorce_model/reference.py`` -- a plain-NumPy backward-induction +EGM sharing no code with the solver -- over four periods with stochastic partner +transitions. Its ``grid_for`` mirrors dcegm's ``continuous_grid_functions``. + +Two independent properties get checked throughout, and neither substitutes for the +other: + +1. **Which grid was used** -- ``endog_grid - policy`` recovers the end-of-period + assets grid a state-choice actually solved on. This is the load-bearing check + for grid threading; the value comparisons below provably cannot make it. +2. **Whether the economics are right** -- the Euler residual, or the policy/value + comparison against the reference. + +Not covered here: the Druedahl-Jorgensen ``assets_begin_of_period`` grid. There the +stored policy is *interpolated* onto the common wealth grid rather than being a raw +EGM solution, so it does not satisfy the Euler equation pointwise and has no closed +form of this kind. It is covered by the equivalence tests in +``test_state_specific_grids_end_to_end.py``. + +""" + +import jax.numpy as jnp +import numpy as np +import pytest +from numpy.testing import assert_allclose +from scipy.special import roots_sh_legendre +from scipy.stats import norm + +import dcegm +import tests.resources.divorce_model.dcegm_functions as dm +import tests.resources.divorce_model.reference as ref +from tests.test_divorce_toy_model import ( + A_GRID_MAX, + A_GRID_POINTS, + N_PERIODS, +) +from tests.test_divorce_toy_model import PARAMS as DIVORCE_PARAMS +from tests.test_divorce_toy_model import dcegm_raw_arrays + +# ===================================================================================== +# Part 1: closed-form Euler equation, two-period model +# ===================================================================================== + + +RHO = 1.5 +N_ASSET_POINTS = 12 +N_QUAD = 5 + +CLOSED_FORM_PARAMS = { + "discount_factor": 0.95, + "interest_rate": 0.04, + "rho": RHO, + "delta": 0.3, + "taste_shock_scale": 0.2, + "income_shock_std": 0.2, + "income_shock_mean": 0.0, +} + +MODEL_SPECS = {"n_choices": 2, "wage": 5.0, "pension": 2.0, "work_cost": 0.4} + +# Declared default grid, used by the solves that leave assets_end_of_period at its +# global default. A solve that supplies a continuous_grid_functions entry for the name +# must declare it as None instead (see process_continuous_grid_functions), which is +# what _assets_config below builds. +DEFAULT_ASSET_GRID = np.linspace(0.0, 20.0, N_ASSET_POINTS) + +# An additional continuous state, for test_state_specific_experience_grid_matches_closed_form +# below. Deliberately absent from utility/budget: experience is a pure pass-through +# combo axis here, so the closed-form Euler equation above stays exactly valid +# regardless of its value -- what this exercises is the *threading* of a second, +# independently state-specific grid, not a new closed-form derivation. +EXPERIENCE_DEFAULT_GRID = np.linspace(0.2, 0.8, 3) + + +# ===================================================================================== +# Model functions +# ===================================================================================== + + +def _utility(consumption, choice, params): + return consumption ** (1 - params["rho"]) / (1 - params["rho"]) - params[ + "delta" + ] * (choice == 0) + + +def _marginal_utility(consumption, params): + return consumption ** (-params["rho"]) + + +def _inverse_marginal_utility(marginal_utility, params): + return marginal_utility ** (-1 / params["rho"]) + + +def _utility_final(wealth, choice, params): + return wealth ** (1 - params["rho"]) / (1 - params["rho"]) - params["delta"] * ( + choice == 0 + ) + + +def _marginal_utility_final(wealth, params): + return wealth ** (-params["rho"]) + + +UTILITY_FUNCTIONS = { + "utility": _utility, + "marginal_utility": _marginal_utility, + "inverse_marginal_utility": _inverse_marginal_utility, +} +UTILITY_FUNCTIONS_FINAL = { + "utility": _utility_final, + "marginal_utility": _marginal_utility_final, +} + + +def _budget( + lagged_choice, + asset_end_of_previous_period, + income_shock, + params, + model_specs, +): + """Beginning-of-period wealth. + + Does not depend on the current choice. + + """ + income = jnp.where(lagged_choice == 0, model_specs["wage"], model_specs["pension"]) + return ( + (1 + params["interest_rate"]) * asset_end_of_previous_period + + income + + income_shock + ) + + +def _budget_choice_dependent( + lagged_choice, + choice, + asset_end_of_previous_period, + income_shock, + params, + model_specs, +): + """Same, plus a cost paid when this period's own choice is to work. + + Declaring ``choice`` is what routes the solve down the per-state-choice law of + motion (see _declares_choice), and the subtracted cost is what makes that routing + observable in the solution. + + """ + base = _budget( + lagged_choice=lagged_choice, + asset_end_of_previous_period=asset_end_of_previous_period, + income_shock=income_shock, + params=params, + model_specs=model_specs, + ) + return base - model_specs["work_cost"] * (choice == 0) + + +def _assets_grid_by_group(group): + """A genuinely different end-of-period assets grid per group.""" + return jnp.asarray(DEFAULT_ASSET_GRID) * (1.0 + 0.5 * group) + + +def _experience_grid_by_group(group): + """A genuinely different experience grid per group, independent of the assets + one.""" + return jnp.asarray(EXPERIENCE_DEFAULT_GRID) * (1.0 + 0.3 * group) + + +def _next_experience(period, lagged_choice, experience, model_specs): + """Deterministic, choice-free experience transition -- not part of the closed + form.""" + return {"experience": experience + 0.1 * (lagged_choice == 0)} + + +# ===================================================================================== +# Closed form +# ===================================================================================== + + +def _quadrature(): + quad_points, quad_weights = roots_sh_legendre(N_QUAD) + quad_draws = ( + norm.ppf(quad_points) * CLOSED_FORM_PARAMS["income_shock_std"] + + CLOSED_FORM_PARAMS["income_shock_mean"] + ) + return np.asarray(quad_draws), np.asarray(quad_weights) + + +def _closed_form_consumption(savings, period_zero_choice, choices, budget_fn): + """Period-0 consumption implied by the Euler equation, for one savings point. + + ``period_zero_choice`` becomes the period-1 state's ``lagged_choice``. The inner + loop is over period-1 choices, which matter only when the budget itself depends on + the current choice; otherwise every ``k`` gives the same wealth and the logit + weights sum back to one. + + """ + quad_draws, quad_weights = _quadrature() + + wealth_next = np.empty((len(choices), len(quad_draws))) + for k_idx, choice_next in enumerate(choices): + wealth_next[k_idx] = np.asarray( + budget_fn( + lagged_choice=period_zero_choice, + choice=choice_next, + asset_end_of_previous_period=savings, + income_shock=quad_draws, + params=CLOSED_FORM_PARAMS, + model_specs=MODEL_SPECS, + ) + ) + + values_next = np.empty_like(wealth_next) + for k_idx, choice_next in enumerate(choices): + values_next[k_idx] = np.asarray( + _utility_final( + wealth=wealth_next[k_idx], choice=choice_next, params=CLOSED_FORM_PARAMS + ) + ) + marg_util_next = wealth_next ** (-RHO) + + # Logit choice probabilities over period-1 choices, rescaled by the max the + # same way calculate_choice_probs_and_unsqueezed_logsum does. + scale = CLOSED_FORM_PARAMS["taste_shock_scale"] + rescaled = np.exp((values_next - values_next.max(axis=0, keepdims=True)) / scale) + choice_probs = rescaled / rescaled.sum(axis=0, keepdims=True) + + marg_util_aggregated = (choice_probs * marg_util_next).sum(axis=0) + marg_util_integrated = marg_util_aggregated @ quad_weights + + rhs = ( + marg_util_integrated + * (1 + CLOSED_FORM_PARAMS["interest_rate"]) + * CLOSED_FORM_PARAMS["discount_factor"] + ) + return rhs ** (-1 / RHO) + + +def _budget_ignoring_choice(**kwargs): + kwargs.pop("choice", None) + return _budget(**kwargs) + + +# ===================================================================================== +# Model setup +# ===================================================================================== + + +def _assets_config(continuous_grid_functions): + """``None`` once a grid function takes the name over, the default array otherwise. + + The pairing is strict in both directions, so this has to follow whatever + ``continuous_grid_functions`` the caller passes. + + """ + grid_funcs = continuous_grid_functions or {} + return None if "assets_end_of_period" in grid_funcs else DEFAULT_ASSET_GRID + + +def _solve(choices, budget_fn, continuous_grid_functions=None): + model_config = { + "n_periods": 2, + "choices": choices, + "deterministic_states": {"group": [0, 1]}, + "continuous_states": { + "assets_end_of_period": _assets_config(continuous_grid_functions) + }, + "n_quad_points": N_QUAD, + } + model = dcegm.setup_model( + model_config=model_config, + model_specs=MODEL_SPECS, + utility_functions=UTILITY_FUNCTIONS, + utility_functions_final_period=UTILITY_FUNCTIONS_FINAL, + budget_constraint=budget_fn, + continuous_grid_functions=continuous_grid_functions, + ) + return model, model.solve(CLOSED_FORM_PARAMS) + + +def _solve_with_experience(choices, continuous_grid_functions): + """Same closed-form model as ``_solve``, plus an additional continuous state. + + Separate from ``_solve`` (rather than adding an optional ``state_space_functions`` + there) so the five existing tests using ``_solve`` keep the simplest possible model + and cannot be affected by this addition. + + ``experience`` is declared as ``None`` -- required whenever a + ``continuous_grid_functions`` entry takes over a name (see + ``process_continuous_grid_functions``); its length is then pinned from the first + state-choice's own grid evaluation. + + """ + model_config = { + "n_periods": 2, + "choices": choices, + "deterministic_states": {"group": [0, 1]}, + "continuous_states": { + "assets_end_of_period": _assets_config(continuous_grid_functions), + "experience": None, + }, + "n_quad_points": N_QUAD, + } + model = dcegm.setup_model( + model_config=model_config, + model_specs=MODEL_SPECS, + utility_functions=UTILITY_FUNCTIONS, + utility_functions_final_period=UTILITY_FUNCTIONS_FINAL, + budget_constraint=_budget, + state_space_functions={"next_period_continuous_state": _next_experience}, + continuous_grid_functions=continuous_grid_functions, + ) + return model, model.solve(CLOSED_FORM_PARAMS) + + +def _period_zero_state_choices(model): + state_choice_space = model.model_structure["state_choice_space"] + names = model.model_structure["discrete_states_names"] + ["choice"] + idxs = np.where(state_choice_space[:, 0] == 0)[0] + rows = [ + {name: int(state_choice_space[i, col]) for col, name in enumerate(names)} + for i in idxs + ] + return idxs, rows + + +# ===================================================================================== +# Tests +# ===================================================================================== + + +def test_state_specific_assets_grid_matches_closed_form(): + """The headline case: a per-group end-of-period assets grid. + + Checks both that each state-choice solved on *its own* grid, and that the resulting + consumption satisfies the closed-form Euler equation on that grid. + + """ + choices = [0, 1] + model, solved = _solve( + choices=choices, + budget_fn=_budget, + continuous_grid_functions={"assets_end_of_period": _assets_grid_by_group}, + ) + endog_grid = np.asarray(solved.endog_grid) + policy = np.asarray(solved.policy) + + idxs, rows = _period_zero_state_choices(model) + n_checked = 0 + + for state_choice_idx, row in zip(idxs, rows): + own_grid = np.asarray(_assets_grid_by_group(group=row["group"])) + + # Column 0 is the zero-wealth point prepended by the solver; the raw EGM + # points follow. + grid_row = endog_grid[state_choice_idx, 0, 1 : len(own_grid) + 1] + policy_row = policy[state_choice_idx, 0, 1 : len(own_grid) + 1] + + # 1) endog_grid = savings + consumption, so this recovers the savings grid + # the state-choice actually solved on. A wrong grid shows up here even + # though the Euler residual below would not notice. + implied_savings = grid_row - policy_row + assert_allclose(implied_savings, own_grid, atol=1e-6) + + # 2) The economics, against the independently computed Euler equation. + for point, (savings, consumption) in enumerate(zip(own_grid, policy_row)): + if savings <= 1e-10: + continue # zero-savings point: corner, not an interior Euler solution + expected = _closed_form_consumption( + savings=savings, + period_zero_choice=row["choice"], + choices=choices, + budget_fn=_budget_ignoring_choice, + ) + assert_allclose( + consumption, expected, rtol=1e-6, err_msg=f"{row} @ {point}" + ) + n_checked += 1 + + assert n_checked > 0 + + +def test_default_grid_matches_closed_form_without_grid_functions(): + """Sensitivity anchor: the same closed form must pin the *unmodified* solver. + + If this failed, a failure of the test above could not be attributed to the state- + specific grid machinery -- it would just mean the closed form is wrong. + + """ + choices = [0, 1] + model, solved = _solve(choices=choices, budget_fn=_budget) + endog_grid = np.asarray(solved.endog_grid) + policy = np.asarray(solved.policy) + + idxs, rows = _period_zero_state_choices(model) + + for state_choice_idx, row in zip(idxs, rows): + grid_row = endog_grid[state_choice_idx, 0, 1 : N_ASSET_POINTS + 1] + policy_row = policy[state_choice_idx, 0, 1 : N_ASSET_POINTS + 1] + assert_allclose(grid_row - policy_row, DEFAULT_ASSET_GRID, atol=1e-6) + + for savings, consumption in zip(DEFAULT_ASSET_GRID, policy_row): + if savings <= 1e-10: + continue + expected = _closed_form_consumption( + savings=savings, + period_zero_choice=row["choice"], + choices=choices, + budget_fn=_budget_ignoring_choice, + ) + assert_allclose(consumption, expected, rtol=1e-6) + + +def test_choice_dependent_budget_matches_closed_form(): + """A budget equation that reads ``choice`` must be solved with it. + + The closed form here aggregates over period-1 choices with logit weights, because + the budget makes next period's wealth differ by that choice. Solving while ignoring + ``choice`` would produce a different consumption policy -- which is exactly what the + companion test below confirms. + + """ + choices = [0, 1] + model, solved = _solve(choices=choices, budget_fn=_budget_choice_dependent) + assert model.model_funcs["budget_depends_on_choice"] + + endog_grid = np.asarray(solved.endog_grid) + policy = np.asarray(solved.policy) + idxs, rows = _period_zero_state_choices(model) + + n_checked = 0 + for state_choice_idx, row in zip(idxs, rows): + grid_row = endog_grid[state_choice_idx, 0, 1 : N_ASSET_POINTS + 1] + policy_row = policy[state_choice_idx, 0, 1 : N_ASSET_POINTS + 1] + assert_allclose(grid_row - policy_row, DEFAULT_ASSET_GRID, atol=1e-6) + + for savings, consumption in zip(DEFAULT_ASSET_GRID, policy_row): + if savings <= 1e-10: + continue + expected = _closed_form_consumption( + savings=savings, + period_zero_choice=row["choice"], + choices=choices, + budget_fn=_budget_choice_dependent, + ) + assert_allclose(consumption, expected, rtol=1e-6) + n_checked += 1 + + assert n_checked > 0 + + +def test_choice_dependent_closed_form_differs_from_choice_free_one(): + """Sensitivity anchor for the test above. + + Confirms the choice-dependent closed form is genuinely a different number, so that + test would fail if the solver silently dropped ``choice`` from the budget equation. + + """ + savings = 5.0 + with_choice = _closed_form_consumption( + savings=savings, + period_zero_choice=0, + choices=[0, 1], + budget_fn=_budget_choice_dependent, + ) + without_choice = _closed_form_consumption( + savings=savings, + period_zero_choice=0, + choices=[0, 1], + budget_fn=_budget_ignoring_choice, + ) + assert not np.isclose(with_choice, without_choice) + + +def test_state_specific_grid_and_choice_dependent_budget_together(): + """Both new features at once, still against the closed form. + + This combination routes through the per-state-choice law of motion *and* the per- + state-choice grid selection simultaneously, which no other test covers. + + """ + choices = [0, 1] + model, solved = _solve( + choices=choices, + budget_fn=_budget_choice_dependent, + continuous_grid_functions={"assets_end_of_period": _assets_grid_by_group}, + ) + endog_grid = np.asarray(solved.endog_grid) + policy = np.asarray(solved.policy) + idxs, rows = _period_zero_state_choices(model) + + n_checked = 0 + for state_choice_idx, row in zip(idxs, rows): + own_grid = np.asarray(_assets_grid_by_group(group=row["group"])) + grid_row = endog_grid[state_choice_idx, 0, 1 : len(own_grid) + 1] + policy_row = policy[state_choice_idx, 0, 1 : len(own_grid) + 1] + assert_allclose(grid_row - policy_row, own_grid, atol=1e-6) + + for savings, consumption in zip(own_grid, policy_row): + if savings <= 1e-10: + continue + expected = _closed_form_consumption( + savings=savings, + period_zero_choice=row["choice"], + choices=choices, + budget_fn=_budget_choice_dependent, + ) + assert_allclose(consumption, expected, rtol=1e-6) + n_checked += 1 + + assert n_checked > 0 + + +def test_state_specific_experience_grid_matches_closed_form(): + """A second, independently state-specific grid (an additional continuous state). + + ``experience`` is a pure pass-through here (it enters neither utility nor the + budget), so the closed form from Part 1 is completely unaffected by its value -- + what this actually exercises is that two state-specific grids threaded through the + same law-of-motion call (see ``calc_law_of_motion``) don't corrupt each other's + combo indexing. Checked at every one of ``experience``'s own combo points, not just + the first, since a combo-axis indexing bug would typically only show up away from + index 0. + + """ + choices = [0, 1] + model, solved = _solve_with_experience( + choices=choices, + continuous_grid_functions={ + "assets_end_of_period": _assets_grid_by_group, + "experience": _experience_grid_by_group, + }, + ) + endog_grid = np.asarray(solved.endog_grid) + policy = np.asarray(solved.policy) + + idxs, rows = _period_zero_state_choices(model) + n_checked = 0 + + for state_choice_idx, row in zip(idxs, rows): + own_asset_grid = np.asarray(_assets_grid_by_group(group=row["group"])) + own_experience_grid = np.asarray(_experience_grid_by_group(group=row["group"])) + + for combo_idx in range(len(own_experience_grid)): + grid_row = endog_grid[ + state_choice_idx, combo_idx, 1 : len(own_asset_grid) + 1 + ] + policy_row = policy[ + state_choice_idx, combo_idx, 1 : len(own_asset_grid) + 1 + ] + + # 1) Still solving on its own assets grid, for every experience combo -- + # not just leaking the right values at combo 0. + implied_savings = grid_row - policy_row + assert_allclose(implied_savings, own_asset_grid, atol=1e-6) + + # 2) The economics are unaffected by experience: same closed form as + # Part 1, evaluated at every combo point. + for savings, consumption in zip(own_asset_grid, policy_row): + if savings <= 1e-10: + continue + expected = _closed_form_consumption( + savings=savings, + period_zero_choice=row["choice"], + choices=choices, + budget_fn=_budget_ignoring_choice, + ) + assert_allclose( + consumption, + expected, + rtol=1e-6, + err_msg=f"{row} @ combo {combo_idx}, savings {savings}", + ) + n_checked += 1 + + assert n_checked > 0 + + +def test_experience_grid_varying_by_group_is_actually_state_specific(): + """Sensitivity anchor: confirm the experience grids really do differ by group. + + Without this, ``test_state_specific_experience_grid_matches_closed_form`` would + still pass if ``_experience_grid_by_group`` happened to collapse to one shared + grid -- it would just be re-testing the default (non-state-specific) path. + + """ + grid_group_0 = np.asarray(_experience_grid_by_group(group=0)) + grid_group_1 = np.asarray(_experience_grid_by_group(group=1)) + assert not np.allclose(grid_group_0, grid_group_1) + assert not np.allclose(grid_group_1, EXPERIENCE_DEFAULT_GRID) + + +# ===================================================================================== +# Part 2: hand-solved n-period reference (divorce model) +# ===================================================================================== + + +BASE_GRID = np.linspace(0.0, 2 * A_GRID_MAX, 2 * A_GRID_POINTS) + + +def _grid_scale(period, choice): + """Scale factor making each (period, choice) own a genuinely different grid.""" + return 1.0 + 0.20 * period + 0.35 * choice + + +def assets_grid(period, choice): + """The state-choice-specific grid, used by *both* sides. + + dcegm calls this with keyword arguments filtered to its signature; the reference + calls it through ``ref.grid_for``, also by keyword. Sharing one definition removes + any chance of the two sides silently drifting apart -- which would show up as a + comparison failure that looks like a solver bug. + + ``period`` and ``choice`` are the only admissible dependencies here; see + ``ref.grid_for`` for why ``partner_state`` and ``lagged_choice`` are not. + + """ + return BASE_GRID * _grid_scale(period, choice) + + +def _solve_both(): + model, solved = dm.build_and_solve( + DIVORCE_PARAMS, + n_periods=N_PERIODS, + a_grid=BASE_GRID, + continuous_grid_functions={"assets_end_of_period": assets_grid}, + ) + ref_solved = ref.solve_reference(N_PERIODS, DIVORCE_PARAMS, assets_grid) + return model, solved, ref_solved + + +@pytest.mark.parametrize("work0", [0, 1]) +@pytest.mark.parametrize("partner_state", [0, 1]) +def test_state_specific_grid_matches_hand_solved_reference(partner_state, work0): + """Dcegm with a per-(period, choice) assets grid must reproduce the reference. + + What this establishes: with state-specific grids active, dcegm still solves + the *right model* -- checked against an implementation sharing no code with + the solver, over 4 periods with stochastic partner transitions. + + What it deliberately does not establish: *which* grid was used. Measured + directly (see the note below), a reference solved on the wrong -- shared -- + grid lands within 2.3e-4 to 4.5e-3 relative of dcegm, while the correct-grid + reference lands within 1.2e-4 to 4.5e-4; the ranges overlap, so no tolerance + separates them. That is inherent rather than a weakness of the setup: the + policy is the same underlying function of wealth either way, and sampling it + on a different grid then interpolating recovers nearly the same values. + ``test_dcegm_solves_on_its_own_declared_grid`` below is what pins the grid. + + """ + model, solved, ref_solved = _solve_both() + + endog_dcegm, policy_dcegm, value_dcegm = dcegm_raw_arrays( + model, solved, period=0, work0=work0, partner_state_0=partner_state + ) + # The divorce model stores individual wealth; a married agent's reference + # solve is in household units, hence the factor of two (same convention as + # test_dcegm_policy_matches_hand_solved_reference_n_periods). + scale = 2.0 if partner_state == 1 else 1.0 + + # Restrict to end-of-period assets whose *continuation* lookup next period is + # genuinely interpolated on both sides. A married agent is queried at + # 2*a_end against period 1's own grid, which tops out at + # BASE_GRID.max() * _grid_scale(1, choice); past that both dcegm and the + # reference extrapolate, by different rules, and comparing there would be + # measuring the extrapolation rules rather than the solve. + interior_max = BASE_GRID.max() * _grid_scale(1, 0) / scale + a0_end_dcegm = endog_dcegm - policy_dcegm + keep = (a0_end_dcegm > 1.0) & (a0_end_dcegm < 0.8 * interior_max) + assert keep.sum() > 50 + + ref_period0 = ref_solved[0][(partner_state, work0)] + ref_own_grid = ref.grid_for(assets_grid, 0, work0) + + ref_policy = ( + np.interp(scale * a0_end_dcegm[keep], ref_own_grid, ref_period0["policy"]) + / scale + ) + ref_value = np.interp( + scale * a0_end_dcegm[keep], ref_own_grid, ref_period0["value"] + ) + + np.testing.assert_allclose(policy_dcegm[keep], ref_policy, rtol=1e-3) + np.testing.assert_allclose(value_dcegm[keep], ref_value, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("work0", [0, 1]) +@pytest.mark.parametrize("partner_state", [0, 1]) +def test_dcegm_solves_on_its_own_declared_grid(partner_state, work0): + """The savings points dcegm actually used are that state-choice's own grid. + + ``endog_grid = savings + consumption``, so ``endog_grid - policy`` recovers + the end-of-period assets grid the state-choice was solved on. This is the + load-bearing test for grid threading, and unlike the policy comparison above + it discriminates sharply: solving the same model *without* the grid functions + leaves 385 of this state-choice's 400 own grid points unaccounted for. + + """ + model, solved, _ = _solve_both() + endog_dcegm, policy_dcegm, _ = dcegm_raw_arrays( + model, solved, period=0, work0=work0, partner_state_0=partner_state + ) + own_grid = np.asarray(assets_grid(period=0, choice=work0)) + + implied_savings = endog_dcegm - policy_dcegm + # dcegm may prepend a natural-borrowing-constraint point below the exogenous + # grid, so match the declared grid as a subset rather than elementwise. + for a_end in own_grid: + assert np.isclose( + implied_savings, a_end, atol=1e-6 + ).any(), f"savings point {a_end} of this state-choice's own grid was not solved" + + +def test_grid_varying_only_by_period_and_choice_is_actually_state_specific(): + """Sensitivity anchor: confirm the grids really do differ across state-choices. + + Without this, the two tests above would still pass if ``assets_grid`` happened to + collapse to one shared grid -- they would just be re-testing the default path. + + """ + grid_p0_work = np.asarray(assets_grid(period=0, choice=0)) + grid_p0_retire = np.asarray(assets_grid(period=0, choice=1)) + grid_p1_work = np.asarray(assets_grid(period=1, choice=0)) + + # Varies with choice, and with period. + assert not np.allclose(grid_p0_work, grid_p0_retire) + assert not np.allclose(grid_p0_work, grid_p1_work) + # At least one state-choice differs from the declared default, so the solve + # cannot be silently falling back to it. + assert not np.allclose(grid_p0_retire, BASE_GRID) + + +def test_grid_depending_on_stochastic_partner_state_is_rejected(): + """A grid varying with the stochastic partner_state must not be accepted. + + Both partner states are parents of the same children here, so they would disagree on + the grid feeding a shared continuation value. dcegm has to reject this at build time + rather than silently pick one -- documenting the boundary of what the reference + above is allowed to mirror. + + """ + + def bad_grid(period, partner_state): + return BASE_GRID * (1.0 + 0.5 * partner_state) + + with pytest.raises(ValueError, match="different grids"): + dm.build_and_solve( + DIVORCE_PARAMS, + n_periods=N_PERIODS, + a_grid=BASE_GRID, + continuous_grid_functions={"assets_end_of_period": bad_grid}, + ) + + +# ===================================================================================== +# Part 2b: choice-dependent budget, hand-solved n-period reference +# +# The closed form (Part 1) covers this only for a 2-period model; every multi-period +# check of a choice-dependent budget elsewhere (test_choice_dependent_budget_... +# _simulates_... below) is dcegm-vs-dcegm, not against ground truth. This closes that +# gap the same way Part 2 does for state-specific grids: an independent hand-rolled +# EGM, here with reference.solve_reference's work_cost argument standing in for +# dcegm_functions.budget_constraint_choice_dependent. +# ===================================================================================== + + +WORK_COST = 5.0 +CHOICE_DEPENDENT_PARAMS = {**DIVORCE_PARAMS, "work_cost": WORK_COST} + + +def _solve_both_choice_dependent(): + model, solved = dm.build_and_solve( + CHOICE_DEPENDENT_PARAMS, + n_periods=N_PERIODS, + a_grid=BASE_GRID, + budget_fn=dm.budget_constraint_choice_dependent, + ) + ref_solved = ref.solve_reference( + N_PERIODS, CHOICE_DEPENDENT_PARAMS, BASE_GRID, work_cost=WORK_COST + ) + return model, solved, ref_solved + + +@pytest.mark.parametrize("work0", [0, 1]) +@pytest.mark.parametrize("partner_state", [0, 1]) +def test_choice_dependent_budget_matches_hand_solved_reference(partner_state, work0): + """A choice-dependent budget, over the full n-period model, against ground truth. + + Mirrors ``test_state_specific_grid_matches_hand_solved_reference``, but for the + *other* feature that routes through the per-state-choice law of motion + (``budget_depends_on_choice``) instead of a state-specific + grid. A shared (non-state-specific) grid is used deliberately, to isolate this + feature rather than re-testing the grid-threading combination Part 2 already + covers. + + """ + model, solved, ref_solved = _solve_both_choice_dependent() + assert model.model_funcs["budget_depends_on_choice"] + + endog_dcegm, policy_dcegm, value_dcegm = dcegm_raw_arrays( + model, solved, period=0, work0=work0, partner_state_0=partner_state + ) + # Household- vs individual-units convention, same as + # test_state_specific_grid_matches_hand_solved_reference. + scale = 2.0 if partner_state == 1 else 1.0 + + # Interior points only -- see test_state_specific_grid_matches_hand_solved_ + # reference for why extrapolation-edge points are excluded. + a0_end_dcegm = endog_dcegm - policy_dcegm + interior_max = BASE_GRID.max() / scale + keep = (a0_end_dcegm > 1.0) & (a0_end_dcegm < 0.8 * interior_max) + assert keep.sum() > 50 + + ref_period0 = ref_solved[0][(partner_state, work0)] + + ref_policy = ( + np.interp(scale * a0_end_dcegm[keep], BASE_GRID, ref_period0["policy"]) / scale + ) + ref_value = np.interp(scale * a0_end_dcegm[keep], BASE_GRID, ref_period0["value"]) + + np.testing.assert_allclose(policy_dcegm[keep], ref_policy, rtol=1e-3) + np.testing.assert_allclose(value_dcegm[keep], ref_value, rtol=1e-3, atol=1e-3) + + +def test_choice_dependent_budget_hand_solved_reference_differs_from_choice_free_one(): + """Sensitivity anchor: confirm ``work_cost`` actually changes the reference. + + Without this, the test above would still pass if ``work_cost`` were silently + ignored by ``resources_after_transition`` -- it would just be re-testing + ``test_state_specific_grid_matches_hand_solved_reference`` with extra, unused + plumbing. + + """ + with_cost = ref.solve_reference( + N_PERIODS, CHOICE_DEPENDENT_PARAMS, BASE_GRID, work_cost=WORK_COST + ) + without_cost = ref.solve_reference( + N_PERIODS, CHOICE_DEPENDENT_PARAMS, BASE_GRID, work_cost=0.0 + ) + + policy_with = with_cost[0][(0, 0)]["policy"] + policy_without = without_cost[0][(0, 0)]["policy"] + assert not np.allclose(policy_with, policy_without) + + +def _load_with_cont_exp(): + import dcegm.toy_models as toy_models + + model_funcs = toy_models.load_example_model_functions("with_cont_exp") + params, model_specs, model_config = ( + toy_models.load_example_params_model_specs_and_config("with_cont_exp") + ) + return model_funcs, params, model_specs, model_config + + +def _budget_with_unused_choice( + period, + lagged_choice, + choice, + experience, + asset_end_of_previous_period, + income_shock, + params, + model_specs, +): + """Identical economics to the toy model's budget, but declaring ``choice``. + + ``choice`` is deliberately unused: declaring it is what routes the model down the + per-choice law of motion, which must not change the answer. + + """ + from dcegm.toy_models.cons_ret_model_with_cont_exp.budget_constraint import ( + budget_constraint_cont_exp, + ) + + return budget_constraint_cont_exp( + period=period, + lagged_choice=lagged_choice, + experience=experience, + asset_end_of_previous_period=asset_end_of_previous_period, + income_shock=income_shock, + params=params, + model_specs=model_specs, + ) + + +# ===================================================================================== +# Part 3: choice-dependent budget through simulate() +# ===================================================================================== + + +def test_choice_dependent_budget_applies_in_the_first_period_too(): + """A budget declaring ``choice`` must work in simulate(), not just solve(). + + The budget equation is applied at the *start* of each period, for every choice, + before the choice is drawn -- so the agent can face a genuinely different wealth + per choice and pick by comparing each choice's value at its own wealth. The first + period is no exception: the user supplies the assets carried into it, and its + wealth is built from those by the same budget equation as every other period's. + + Before this, such a model solved fine and then crashed in simulate() with + ``KeyError: 'choice'``, because simulation computed a single wealth from the + next period's state -- which carries ``lagged_choice`` but not ``choice``. + + """ + model_funcs, params, model_specs, model_config = _load_with_cont_exp() + model_funcs = dict(model_funcs) + model_funcs["budget_constraint"] = _budget_with_unused_choice + + solved = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ).solve(params) + + n_agents = 200 + given_assets = 10.0 + states_initial = { + "period": np.zeros(n_agents, dtype=int), + "lagged_choice": np.zeros(n_agents, dtype=int), + "experience": np.ones(n_agents) * 0.5, + "assets_end_of_previous_period": np.ones(n_agents) * given_assets, + } + df = solved.simulate(states_initial=states_initial, seed=1) + + assert np.all(np.isfinite(df["assets_begin_of_period"])) + # The supplied number is what agents carry in, not what they have to spend: the + # budget equation turns it into period 0's wealth, as in every other period. + first_period_assets = df.xs(0, level=0)["assets_begin_of_period"].to_numpy() + assert not np.allclose(first_period_assets, given_assets) + + +def test_choice_dependent_budget_matches_choice_free_one_when_the_cost_is_zero(): + """Sensitivity anchor for the simulation path. + + A budget that declares ``choice`` but does not use it routes through the per-choice + machinery (``budget_depends_on_choice`` is True), while an otherwise-identical + budget that omits ``choice`` takes the broadcast path. Both describe the same + economics, so simulated output must agree exactly -- which would fail if the per- + choice path mis-selected the realized column. + + """ + model_funcs, params, model_specs, model_config = _load_with_cont_exp() + + plain = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ).solve(params) + + with_choice_funcs = dict(model_funcs) + with_choice_funcs["budget_constraint"] = _budget_with_unused_choice + with_choice = dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **with_choice_funcs + ).solve(params) + + n_agents = 200 + states_initial = { + "period": np.zeros(n_agents, dtype=int), + "lagged_choice": np.zeros(n_agents, dtype=int), + "experience": np.ones(n_agents) * 0.5, + "assets_end_of_previous_period": np.ones(n_agents) * 10, + } + df_plain = plain.simulate(states_initial=states_initial, seed=7) + df_with_choice = with_choice.simulate(states_initial=states_initial, seed=7) + + for column in df_plain.columns: + np.testing.assert_allclose( + df_plain[column].to_numpy(), + df_with_choice[column].to_numpy(), + err_msg=f"column {column}", + ) + + +def _next_experience_declaring_choice( + period, lagged_choice, choice, experience, model_specs +): + """Identical to the toy model's transition, but declaring ``choice``.""" + max_experience_period = period + model_specs["max_init_experience"] + return { + "experience": (1 / max_experience_period) + * ((max_experience_period - 1) * experience + (lagged_choice == 0)) + } + + +def _with_continuous_state_transition(model_funcs, transition): + model_funcs = dict(model_funcs) + state_space_functions = dict(model_funcs["state_space_functions"]) + state_space_functions["next_period_continuous_state"] = transition + model_funcs["state_space_functions"] = state_space_functions + return model_funcs + + +def test_continuous_state_transition_declaring_choice_is_rejected(): + """Only the budget equation may depend on the choice made in its own period. + + Every other state is fixed before the choice, so ``next_period_continuous_state`` is + settled at the end of the previous period and can never see the choice about to be + made. Declaring it is a modelling error, not a supported path, so it has to fail + loudly at model build rather than be silently ignored. + + """ + model_funcs, _params, model_specs, model_config = _load_with_cont_exp() + + with pytest.raises(ValueError, match="declares 'choice'"): + dcegm.setup_model( + model_config=model_config, + model_specs=model_specs, + **_with_continuous_state_transition( + model_funcs, _next_experience_declaring_choice + ), + ) + + +def test_budget_declaring_the_old_income_shock_name_is_rejected(): + """``income_shock_previous_period`` was renamed to ``income_shock``. + + The shock is realized in the period whose beginning-of-period wealth the budget + equation computes. The old name generalised a property of the shipped example + models -- income earned against ``lagged_choice`` -- into the framework, where it + does not hold. Left unchecked the rename would surface as a bare ``KeyError`` out + of the signature filter, so it has to fail with something that says what to do. + + """ + model_funcs, _params, model_specs, model_config = _load_with_cont_exp() + + def budget_with_old_name( + lagged_choice, + asset_end_of_previous_period, + income_shock_previous_period, + params, + model_specs, + ): + return _budget( + lagged_choice=lagged_choice, + asset_end_of_previous_period=asset_end_of_previous_period, + income_shock=income_shock_previous_period, + params=params, + model_specs=model_specs, + ) + + model_funcs = dict(model_funcs) + model_funcs["budget_constraint"] = budget_with_old_name + + with pytest.raises(ValueError, match="renamed to 'income_shock'"): + dcegm.setup_model( + model_config=model_config, model_specs=model_specs, **model_funcs + ) diff --git a/tests/test_two_occupation_model.py b/tests/test_two_occupation_model.py index d07db308..c68bb21d 100644 --- a/tests/test_two_occupation_model.py +++ b/tests/test_two_occupation_model.py @@ -87,7 +87,7 @@ def budget_constraint_cont_exp( exp_green, exp_red, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, ): interest_factor = 1 + params["interest_rate"] @@ -98,8 +98,8 @@ def budget_constraint_cont_exp( ) resource = ( interest_factor * asset_end_of_previous_period - + (wage + income_shock_previous_period) * (lagged_choice != 2) - + (wage + income_shock_previous_period) * 0.5 * (lagged_choice == 2) + + (wage + income_shock) * (lagged_choice != 2) + + (wage + income_shock) * 0.5 * (lagged_choice == 2) ) return jnp.maximum(resource, 0.5) @@ -153,7 +153,7 @@ def budget_constraint_discrete_exp( exp_green, exp_red, asset_end_of_previous_period, - income_shock_previous_period, + income_shock, params, ): interest_factor = 1 + params["interest_rate"] @@ -164,8 +164,8 @@ def budget_constraint_discrete_exp( ) resource = ( interest_factor * asset_end_of_previous_period - + (wage + income_shock_previous_period) * (lagged_choice != 2) - + (wage + income_shock_previous_period) * 0.5 * (lagged_choice == 2) + + (wage + income_shock) * (lagged_choice != 2) + + (wage + income_shock) * 0.5 * (lagged_choice == 2) ) return jnp.maximum(resource, 0.5) @@ -505,7 +505,7 @@ def initial_states(): n_agents = 100 return { "n_agents": n_agents, - "assets_begin_of_period": jnp.ones(n_agents), + "assets_end_of_previous_period": jnp.ones(n_agents), "exp_green": jnp.zeros(n_agents), "exp_red": jnp.zeros(n_agents), "lagged_choice": jnp.zeros(n_agents), @@ -851,13 +851,16 @@ def test_exact_vs_offgrid_simulation(df_cont_exp, df_offgrid): atol_mean=1e-2, atol_max=1.01e-2, ) + # Bounds on interpolation error between an exact experience grid and one with + # step 1.8, over the region the agents actually visit (wealth ~1.3-8.8 here). + # Measured gaps are ~0.021 mean / ~0.040 max. assert_sim_means_close( df_cont_exp, df_offgrid, "consumption", group_by="period", - atol_mean=2e-2, - atol_max=3e-2, + atol_mean=3e-2, + atol_max=5e-2, ) diff --git a/tests/test_two_period_continuous_experience.py b/tests/test_two_period_continuous_experience.py index 0a2389d3..19709091 100644 --- a/tests/test_two_period_continuous_experience.py +++ b/tests/test_two_period_continuous_experience.py @@ -70,7 +70,7 @@ def marginal_utility_weighted( lagged_consumption=lagged_consumption, lagged_choice=lagged_choice, experience=exp_new, - income_shock_previous_period=income_shock, + income_shock=income_shock, params=params, ) model_functions = toy_models.load_example_model_functions("dcegm_paper") @@ -115,7 +115,7 @@ def budget_constraint_continuous( lagged_consumption: float, lagged_choice: int, experience: float, - income_shock_previous_period: float, + income_shock: float, params: Dict[str, float], ) -> float: @@ -125,7 +125,7 @@ def budget_constraint_continuous( income_from_previous_period = calc_stochastic_income( experience=experience_years, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, params=params, ) @@ -143,7 +143,7 @@ def budget_constraint_continuous_dcegm( asset_end_of_previous_period: float, lagged_choice: int, experience: float, - income_shock_previous_period: float, + income_shock: float, params: Dict[str, float], ) -> float: @@ -153,7 +153,7 @@ def budget_constraint_continuous_dcegm( income_from_previous_period = calc_stochastic_income( experience=experience_years, - wage_shock=income_shock_previous_period, + wage_shock=income_shock, params=params, ) @@ -281,15 +281,10 @@ def create_test_inputs(): value_interp_final_period, marginal_utility_final_last_period, ) = solve_final_period( - idx_state_choices_final_period=last_two_period_batch_info_cont[ - "idx_state_choices_final_period" - ], - state_choice_mat_final_period=last_two_period_batch_info_cont[ - "state_choice_mat_final_period" - ], + batch_info=last_two_period_batch_info_cont, + model_structure=model.model_structure, income_shocks_scaled=income_shocks_scaled, continuous_states_info=model_config["continuous_states_info"], - model_structure=model.model_structure, params=params, upper_envelope_method=model_config["upper_envelope"]["method"], skip_endog_grid_storage=model_config["upper_envelope"][ @@ -318,7 +313,6 @@ def create_test_inputs(): taste_shock_scale_is_scalar=True, income_shock_weights=income_shock_weights, continuous_grids_info=model_config["continuous_states_info"], - continuous_state_space=model.model_structure["continuous_state_space"], model_funcs=model_funcs_cont, debug_info=None, ) @@ -457,9 +451,9 @@ def _get_solve_last_two_periods_args(model, params, has_second_continuous_state) income_shock_draws_unscaled * income_shock_std + income_shock_mean ) - n_continuous_state_combinations = model_structure["continuous_state_space"][ - next(iter(model_structure["continuous_state_space"])) - ].shape[0] + n_continuous_state_combinations = model_config["continuous_states_info"][ + "n_continuous_state_combinations" + ] ( value_solved, policy_solved, @@ -469,6 +463,7 @@ def _get_solve_last_two_periods_args(model, params, has_second_continuous_state) # Read out grid size n_total_wealth_grid=model_config["n_total_wealth_grid"], n_state_choices=model_structure["state_choice_space"].shape[0], + store_endog_grid=True, ) return (