From df4a911041015617c0a1a096246d8cff752e6b13 Mon Sep 17 00:00:00 2001 From: Maximilian Blesch Date: Sun, 18 Feb 2024 10:20:34 -0500 Subject: [PATCH] Fix bug in value aggregation (#95) --- src/dcegm/solve.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/dcegm/solve.py b/src/dcegm/solve.py index e0c14159..d79412fa 100644 --- a/src/dcegm/solve.py +++ b/src/dcegm/solve.py @@ -260,6 +260,9 @@ def backward_induction( resources_period = resources_beginning_of_period[ state_objects_period["idx_parent_states"] ] + reshape_state_choice_vec_to_mat_prev_period = period_specific_state_objects[ + period + 1 + ]["reshape_state_choice_vec_to_mat"] ( endog_grid_period, policy_left_period, @@ -272,6 +275,7 @@ def backward_induction( marg_util_interpolated_previous_period=marg_util_interpolated_next_period, params=params, state_objects=state_objects_period, + reshape_state_choice_vec_to_mat_prev_period=reshape_state_choice_vec_to_mat_prev_period, exog_savings_grid=exog_savings_grid, resources_period=resources_period, income_shock_weights=income_shock_weights, @@ -292,6 +296,7 @@ def solve_single_period( marg_util_interpolated_previous_period: jnp.ndarray, params: Dict[str, float], state_objects: Dict[str, np.ndarray], + reshape_state_choice_vec_to_mat_prev_period: np.ndarray, exog_savings_grid: np.ndarray, resources_period: jnp.ndarray, income_shock_weights: jnp.ndarray, @@ -305,9 +310,7 @@ def solve_single_period( marg_util, emax = aggregate_marg_utils_and_exp_values( value_state_choice_specific=value_interpolated_previous_period, marg_util_state_choice_specific=marg_util_interpolated_previous_period, - reshape_state_choice_vec_to_mat=state_objects[ - "reshape_state_choice_vec_to_mat" - ], + reshape_state_choice_vec_to_mat=reshape_state_choice_vec_to_mat_prev_period, taste_shock_scale=taste_shock_scale, income_shock_weights=income_shock_weights, )