Adding debug code.

Author: MaxBlesch

src/dcegm/backward_induction.py

@@ -96,7 +96,9 @@ def backward_induction(
         policy_solved,
         endog_grid_solved,
     ) = create_solution_container(
-        model_config=model_config,
+        continuous_states_info=model_config["continuous_states_info"],
+        # Read out grid size
+        n_total_wealth_grid=model_config["tuning_params"]["n_total_wealth_grid"],
         n_state_choices=model_structure["state_choice_space"].shape[0],
     )
 

src/dcegm/interfaces/inspect_solution.py

@@ -39,7 +39,7 @@ def partially_solve(
         model_funcs=model_funcs,
     )
 
-    last_relevant_period = model_config["n_periods"] - n_periods - 1
+    last_relevant_period = model_config["n_periods"] - n_periods
 
     relevant_state_choices_mask = (
         model_structure["state_choice_space"][:, 0] >= last_relevant_period
@@ -53,14 +53,61 @@ def partially_solve(
         policy_solved,
         endog_grid_solved,
     ) = create_solution_container(
-        model_config=model_config,
+        continuous_states_info=model_config["continuous_states_info"],
+        # Read out grid size
+        n_total_wealth_grid=model_config["tuning_params"]["n_total_wealth_grid"],
         n_state_choices=relevant_state_choice_space.shape[0],
     )
 
     if return_candidates:
+        n_assets_end_of_period = model_config["continuous_states_info"][
+            "assets_grid_end_of_period"
+        ].shape[0]
         (value_candidates, policy_candidates, endog_grid_candidates) = (
             create_solution_container(
-                model_config=model_config,
+                continuous_states_info=model_config["continuous_states_info"],
+                n_total_wealth_grid=n_assets_end_of_period,
                 n_state_choices=relevant_state_choice_space.shape[0],
             )
         )
+
+    # Create debug information
+    debug_info = {
+        "return_candidates": return_candidates,
+        "rescale_idx": np.where(relevant_state_choices_mask)[0].min(),
+    }
+    (
+        value_solved,
+        policy_solved,
+        endog_grid_solved,
+        value_candidates_second_last,
+        policy_candidates_second_last,
+        endog_grid_candidates_second_last,
+    ) = solve_last_two_periods(
+        params=params,
+        continuous_states_info=continuous_states_info,
+        cont_grids_next_period=cont_grids_next_period,
+        income_shock_weights=income_shock_weights,
+        model_funcs=model_funcs,
+        last_two_period_batch_info=batch_info["last_two_period_info"],
+        value_solved=value_solved,
+        policy_solved=policy_solved,
+        endog_grid_solved=endog_grid_solved,
+        debug_info=debug_info,
+    )
+    if return_candidates:
+        idx_second_last = batch_info["last_two_period_info"][
+            "idx_state_choices_second_last_period"
+        ]
+        idx_second_last_rescaled = idx_second_last - debug_info["rescale_idx"]
+        value_candidates = value_candidates.at[idx_second_last_rescaled, ...].set(
+            value_candidates_second_last
+        )
+        policy_candidates = policy_candidates.at[idx_second_last_rescaled, ...].set(
+            policy_candidates_second_last,
+        )
+        endog_grid_candidates = endog_grid_candidates.at[
+            idx_second_last_rescaled, ...
+        ].set(endog_grid_candidates_second_last)
+
+    return value_solved, policy_solved, endog_grid_solved

src/dcegm/pre_processing/sol_container.py

@@ -4,16 +4,11 @@
 
 
 def create_solution_container(
-    model_config: Dict[str, Any],
+    continuous_states_info: Dict[str, Any],
+    n_total_wealth_grid: int,
     n_state_choices: int,
 ):
     """Create solution containers for value, policy, and endog_grid."""
-
-    # Read out grid size
-    n_total_wealth_grid = model_config["tuning_params"]["n_total_wealth_grid"]
-
-    # Check if second continuous state exists and read out array size
-    continuous_states_info = model_config["continuous_states_info"]
     if continuous_states_info["second_continuous_exists"]:
         n_second_continuous_grid = continuous_states_info["n_second_continuous_grid"]
 

tests/test_two_period_continuous_experience.py

@@ -449,7 +449,8 @@ def _get_solve_last_two_periods_args(model, params, has_second_continuous_state)
 
     # Create solution containers for value, policy, and endogenous grids
     value_solved, policy_solved, endog_grid_solved = create_solution_container(
-        model_config=model_config,
+        continuous_states_info=model_config["continuous_states_info"],
+        n_total_wealth_grid=model_config["tuning_params"]["n_total_wealth_grid"],
         n_state_choices=model_structure["state_choice_space"].shape[0],
     )