Child states including trans probs calc.

Author: MaxBlesch

src/dcegm/interfaces/model_class.py

@@ -322,6 +322,17 @@ 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."""
+        child_states_df = self.get_child_states(state, choice)
+
+        trans_probs = self.model_funcs["compute_stochastic_transition_vec"](
+            params=params, choice=choice, **state
+        )
+        child_states_df["trans_probs"] = trans_probs
+        return child_states_df
+
     def compute_law_of_motions(self, params):
         return calc_cont_grids_next_period(
             params=params,

src/dcegm/solve_single_period.py

@@ -122,7 +122,7 @@ def solve_for_interpolated_values(
     debug_info,
 ):
     # EGM step 2)
-    # Aggregate the marginal utilities and expected values over all state-choice
+    # Aggregate the marginal utilities and expected values over all child state-choice
     # combinations and income shock draws
     marg_util, emax = aggregate_marg_utils_and_exp_values(
         value_state_choice_specific=value_interpolated,