put loading initial values to newborns with read_shocks into sim_birth

econ-ark · Oct 17, 2023 · 799de9a · 799de9a
1 parent 0c52963
commit 799de9a
Show file tree

Hide file tree

Showing 2 changed files with 256 additions and 158 deletions.
diff --git a/HARK/simulation/monte_carlo.py b/HARK/simulation/monte_carlo.py
@@ -384,21 +384,7 @@ def get_mortality(self):
         """
         who_dies = self.vars_now['live'] <= 0
 
-        if self.read_shocks:
-            # Instead of simulating births, assign the saved newborn initial conditions
-            if np.sum(who_dies) > 0:
-                for var_name in self.initial:
-                    self.vars_now[var_name][
-                        who_dies
-                    ] = self.newborn_init_history[var_name][
-                        self.t_sim, who_dies
-                    ]
-
-                # Reset ages of newborns
-                self.t_age[who_dies] = 0
-                self.t_cycle[who_dies] = 0
-        else:
-            self.sim_birth(who_dies)
+        self.sim_birth(who_dies)
 
         self.who_dies = who_dies
         return None
@@ -418,20 +404,26 @@ def sim_birth(self, which_agents):
         None
         """
         if self.read_shocks:
+            t = self.t_sim - 1 if self.t_sim > 0 else 0
             initial_vals = {
-                init_var: self.newborn_init_history[init_var][self.t_sim, :]
+                init_var: self.newborn_init_history[init_var][t, which_agents]
                 for init_var
                 in self.initial
-                }
+            }
+
         else:
             initial_vals = draw_shocks(
                 self.initial,
                 np.zeros(which_agents.sum())
             )
 
-        for varn in initial_vals:
-            self.vars_now[varn][which_agents] = initial_vals[varn]
-            self.newborn_init_history[varn][self.t_sim, which_agents] = initial_vals[varn]
+        if np.sum(which_agents) > 0:
+            for varn in initial_vals:
+                self.vars_now[varn][which_agents] = initial_vals[varn]
+                self.newborn_init_history[varn][self.t_sim, which_agents] = initial_vals[varn]
+
+        self.t_age[which_agents] = 0
+        self.t_cycle[which_agents] = 0
 
     def simulate(self, sim_periods=None):
         """