Add shock discretization functionality

[mj023]

This PR is duplicate of #170. Merging the previous PR with the main branch became unfeasible because there were too many merge conflicts.

This PR will add a way to specify multiple common shock types directly through the interface of LCM.

To accomplish this, a new ShockGrid grid class will be added. The parameters for the shocks are passed via a new fixed_parameters attribute during the model creation. During the solution the shocks will be treated like normal DiscreteGrids, but during the simulation, the next values will be drawn from the shocks undiscretized probability distribution. Here the shocks transition probabilities will be linearly interpolated, as the discretized transition probabilities are only calculated at the gridpoints, but subjects can land between gridpoints.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[timmens]

Very nice! I have a few general comments. Excited to discuss this soon! 🙂

Bugs

1. ShockGrid not exported - Listed in all but not imported in src/lcm/init.py

API Design

2. Inconsistent pass pattern - Built-in shock transitions use def next_x(): pass while all other transitions return values. Consider inferring transitions automatically from ShockGrid - if a state has a ShockGrid, generate its transition implicitly (no next_ function needed). This also removes the redundant distribution_type in both grid and decorator. Only custom shocks would need a next function then. Counter argument would be that this breaks the "contract" that each state variable (which a shock is) has a next function.
3. ShockGrid.get_coordinate breaks interface - Takes extra params argument unlike other ContinuousGrid subclasses (hence the # ty: ignore). Could consider a different abstraction?

Implementation

4. ShockGrid.to_jax() returns placeholders - Returns [0, 1, 2, ...] until update_sas_with_shocks fills in real values. This implicit contract could cause subtle bugs.
5. InternalRegime now mutable - Changed from frozen=True to frozen=False. Could use dataclasses.replace() instead to preserve immutability?
6. Python loops in discretization - _fill_tauchen and rouwenhorst use Python loops. Could vectorize with jax.vmap?

Minor

7. Outdated docstring - ShockGrid docstring mentions nonexistent start/stop attributes

[hmgaudecker]

Plan:

• Add an exogenous argument to Model.__init__(), which takes {"next_fn" = {p0: a, p1: b}} as an argument when next_fn is decorated with @lcm.mark.stochastic. This means that we can construct the grid upfront and have just one state of InternalRegime
• Don't think much about changing this ex post yet (something like self.exogenous | params, which we eventually want)
• For MY example: Set sig_x = 1 and multiply with a param element inside the function using it.
• Finetuning will be done in a PR solving ENH: Improved flow -> solve(Model, params), simulate(Model, params, initial_states, solution) #219

[timmens]

This is starting to look really really good!

[mj023]

I think I was able to make most of the requested changes. There is a new fixed_param attribute through which the parameters for the shocks are passed and there is a Shock Class which can be used as an easy way to get grids and transition probs if they are somehow needed by the user, for example to draw initial states etc.

I also tested the solution with the MY Model and there are only extremely small differences during the solution, probably caused by the use of map_coordinates instead of the old integer indexes. Solution and simulation are also still running in the same amount of time as before.

When creating the tests for the shocks, I found a bug with the RNG-Keys used during the simulation. For some reason, the names of the next_state-functions sometimes get swapped, which will also lead to some keys being swapped around. I put a band-aid solution in place (sorting names), but I will open an Issue so this can be fixed in a different PR. I also had to skip the tests on GPU, because using FP32 gives different RNG-Keys than using FP64. Our current tests for stochastic transitions just seem to be very resistant to changes in the keys.

[hmgaudecker]

This is coming along very nicely, thanks a lot!!! I only looked at the interface so far. Looks neat, though a bit more streamlining could be helpful (maybe documentation is enough).

I also like the fixed_params attribute as a start, we'll flesh it out elsewhere then.

Thanks for tracking down those bugs, they sound like they would have become a nightmare!

[hmgaudecker]

I am a little lost with the naming (productivity_shock and prod_shock) and structure here. If I understand correctly,

• productivity_shock sets the grid and via next_productivity_shock its transition
• prod_shock is only for setting the initial distribution?

However, why do we need to repeat the distribution_type, n_points, and shock_params arguments? Wouldn't an interface Shock(shock_grid, shock_transition) be doable? Thinking DRY here.

[mj023]

This is mostly caused by the new fixed_param. The ShockGrid can only be fully initialized when the Model gets created, because it's missing the params. The only other way to do this would be to repeat the params, eg when creating the ShockGrid or the Shock and then in the fixed params. Right now the Shock Class is therefore only a streamlined way to access the internal shock discretization functions. Therefore the repetition can at least be avoided if one does not need the Shock outside the model.

The next_state function would also optimally be eliminated for Shock states, but this currently clashes with the way we process transitions. This should be part of a different PR that eliminates the duplication of all next_state functions for all regimes.