Extend sparsity to exogenous processes

docs/source/Tutorials/two_period_model_tutorial.ipynb

@@ -296,7 +296,7 @@
    "source": [
     "#### State space functions\n",
     "\n",
-    "Next we define state space functions ```create_state_space``` and ```get_state_specific_choice_set```. They can be directly imported the dcegm package, but we display them here, in order to explain how they work."
+    "Next we define state space functions ```create_state_space``` and ```state_specific_choice_set```. They can be directly imported the dcegm package, but we display them here, in order to explain how they work."
    ]
   },
   {
@@ -402,7 +402,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def get_state_specific_choice_set(\n",
+    "def state_specific_choice_set(\n",
     "    state: np.ndarray,\n",
     "    state_space: np.ndarray,  # noqa: U100\n",
     "    indexer: np.ndarray,\n",
@@ -460,7 +460,7 @@
     }
    ],
    "source": [
-    "choice_set_5 = get_state_specific_choice_set(_state_space[4], _state_space, _indexer)\n",
+    "choice_set_5 = state_specific_choice_set(_state_space[4], _state_space, _indexer)\n",
     "choice_set_5"
    ]
   },
@@ -479,7 +479,7 @@
    "source": [
     "state_space_functions = {\n",
     "    \"create_state_space\": create_state_space,\n",
-    "    \"get_state_specific_choice_set\": get_state_specific_choice_set,\n",
+    "    \"state_specific_choice_set\": state_specific_choice_set,\n",
     "}"
    ]
   },

src/dcegm/final_periods.py

@@ -18,7 +18,7 @@ def solve_last_two_periods(
     income_shock_weights: jnp.ndarray,
     exog_grids: Dict[str, jnp.ndarray],
     model_funcs: Dict[str, Callable],
-    batch_info,
+    last_two_period_batch_info,
     value_solved,
     policy_solved,
     endog_grid_solved,
@@ -39,7 +39,7 @@ def solve_last_two_periods(
         compute_utility (callable): User supplied utility function.
         compute_marginal_utility (callable): User supplied marginal utility
             function.
-        batch_info (dict): Dictionary containing information about the batch
+        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
@@ -56,9 +56,15 @@ 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"],
-        idx_parent_states_final_period=batch_info["idxs_parent_states_final_period"],
-        state_choice_mat_final_period=batch_info["state_choice_mat_final_period"],
+        idx_state_choices_final_period=last_two_period_batch_info[
+            "idx_state_choices_final_period"
+        ],
+        idx_parent_states_final_period=last_two_period_batch_info[
+            "idxs_parent_states_final_period"
+        ],
+        state_choice_mat_final_period=last_two_period_batch_info[
+            "state_choice_mat_final_period"
+        ],
         cont_grids_next_period=cont_grids_next_period,
         exog_grids=exog_grids,
         params=params,
@@ -72,9 +78,13 @@ def solve_last_two_periods(
     endog_grid, policy, value = solve_for_interpolated_values(
         value_interpolated=value_interp_final_period,
         marginal_utility_interpolated=marginal_utility_final_last_period,
-        state_choice_mat=batch_info["state_choice_mat_second_last_period"],
-        child_state_idxs=batch_info["child_states_second_last_period"],
-        states_to_choices_child_states=batch_info["state_to_choices_final_period"],
+        state_choice_mat=last_two_period_batch_info[
+            "state_choice_mat_second_last_period"
+        ],
+        child_state_idxs=last_two_period_batch_info["child_states_second_last_period"],
+        states_to_choices_child_states=last_two_period_batch_info[
+            "state_to_choices_final_period"
+        ],
         params=params,
         taste_shock_scale=taste_shock_scale,
         income_shock_weights=income_shock_weights,
@@ -83,7 +93,7 @@ def solve_last_two_periods(
         has_second_continuous_state=has_second_continuous_state,
     )
 
-    idx_second_last = batch_info["idx_state_choices_second_last_period"]
+    idx_second_last = last_two_period_batch_info["idx_state_choices_second_last_period"]
 
     value_solved = value_solved.at[idx_second_last, ...].set(value)
     policy_solved = policy_solved.at[idx_second_last, ...].set(policy)

src/dcegm/interface.py

@@ -5,6 +5,11 @@
     interp_policy_on_wealth,
     interp_value_on_wealth,
 )
+from dcegm.interpolation.interp2d import (
+    interp2d_policy_and_value_on_wealth_and_regular_grid,
+    interp2d_policy_on_wealth_and_regular_grid,
+    interp2d_value_on_wealth_and_regular_grid,
+)
 
 
 def policy_and_value_for_state_choice_vec(
@@ -47,14 +52,13 @@ def policy_and_value_for_state_choice_vec(
 
 
 def value_for_state_choice_vec(
-    state_choice_vec,
-    wealth,
-    map_state_choice_to_index,
-    discrete_states_names,
     endog_grid_solved,
     value_solved,
-    compute_utility,
     params,
+    model,
+    state_choice_vec,
+    wealth,
+    second_continous=None,
 ):
     """Get policy and value for a given state and choice vector.
 
@@ -67,20 +71,36 @@ def value_for_state_choice_vec(
         Tuple[float, float]: Policy and value for the given state and choice vector.
 
     """
+    map_state_choice_to_index = model["model_structure"]["map_state_choice_to_index"]
+    discrete_states_names = model["model_structure"]["discrete_states_names"]
+    compute_utility = model["model_funcs"]["compute_utility"]
+
     state_choice_tuple = tuple(
         state_choice_vec[st] for st in discrete_states_names + ["choice"]
     )
 
     state_choice_index = map_state_choice_to_index[state_choice_tuple]
 
-    value = interp_value_on_wealth(
-        wealth=wealth,
-        endog_grid=jnp.take(endog_grid_solved, state_choice_index, axis=0),
-        value=jnp.take(value_solved, state_choice_index, axis=0),
-        compute_utility=compute_utility,
-        state_choice_vec=state_choice_vec,
-        params=params,
-    )
+    if second_continous is None:
+        value = interp_value_on_wealth(
+            wealth=wealth,
+            endog_grid=jnp.take(endog_grid_solved, state_choice_index, axis=0),
+            value=jnp.take(value_solved, state_choice_index, axis=0),
+            compute_utility=compute_utility,
+            state_choice_vec=state_choice_vec,
+            params=params,
+        )
+    else:
+        value = interp2d_value_on_wealth_and_regular_grid(
+            regular_grid=model["options"]["exog_grids"]["second_continuous"],
+            wealth_grid=jnp.take(endog_grid_solved, state_choice_index, axis=0),
+            value_grid=jnp.take(value_solved, state_choice_index, axis=0),
+            regular_point_to_interp=second_continous,
+            wealth_point_to_interp=wealth,
+            compute_utility=compute_utility,
+            state_choice_vec=state_choice_vec,
+            params=params,
+        )
     return value
 
 

src/dcegm/law_of_motion.py

@@ -14,7 +14,7 @@ def calc_cont_grids_next_period(
             discrete_states_beginning_of_period=state_space_dict,
             continuous_grid=exog_grids["second_continuous"],
             params=params,
-            compute_continuous_state=model_funcs["update_continuous_state"],
+            compute_continuous_state=model_funcs["next_period_continuous_state"],
         )
 
         # Extra dimension for continuous state