Black formatted aggregates.py and household.py and updated parameters.py

rickecon · Sep 6, 2024 · cd728f5 · cd728f5
1 parent 881ad75
commit cd728f5
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 19 deletions.
diff --git a/ogcore/aggregates.py b/ogcore/aggregates.py
@@ -261,19 +261,16 @@ def get_RM(Y, p, method):
         RM = np.zeros_like(Y)
         RM[0] = p.alpha_RM_1 * Y[0]
         for t in range(1, p.tG1):
-            RM[t] = (
-                (1 + p.g_RM[t]) / (np.exp(p.g_y) * (1 + p.g_n[t]))
-            ) * RM[t - 1]
+            RM[t] = ((1 + p.g_RM[t]) / (np.exp(p.g_y) * (1 + p.g_n[t]))) * RM[
+                t - 1
+            ]
         for t in range(p.tG1, p.tG2 - 1):
-            RM[t] = (
-                p.rho_RM * p.alpha_RM_T * Y[t]
-                + (
-                    (1 - p.rho_RM)
-                    * ((1 + p.g_RM[t]) / (np.exp(p.g_y) * (1 + p.g_n[t])))
-                    * RM[t - 1]
-                )
+            RM[t] = p.rho_RM * p.alpha_RM_T * Y[t] + (
+                (1 - p.rho_RM)
+                * ((1 + p.g_RM[t]) / (np.exp(p.g_y) * (1 + p.g_n[t])))
+                * RM[t - 1]
             )
-        RM[p.tG2 - 1:] = p.alpha_RM_T * Y[p.tG2 - 1:]
+        RM[p.tG2 - 1 :] = p.alpha_RM_T * Y[p.tG2 - 1 :]
 
     return RM
 

diff --git a/ogcore/household.py b/ogcore/household.py
@@ -314,7 +314,11 @@ def get_cons(r_p, w, p_tilde, b, b_splus1, n, bq, rm, net_tax, e, p):
 
     """
     cons = (
-        (1 + r_p) * b + w * e * n + bq + rm - net_tax
+        (1 + r_p) * b
+        + w * e * n
+        + bq
+        + rm
+        - net_tax
         - b_splus1 * np.exp(p.g_y)
     ) / p_tilde
 

diff --git a/ogcore/parameters.py b/ogcore/parameters.py
@@ -244,14 +244,20 @@ def compute_default_params(self):
             )
             setattr(self, item, tax_to_set_out)
 
-        # Try to deal with size of eta.  It may vary by S, J, T, but
-        # want to allow user to enter one that varies by only S, S and J,
+        # Try to deal with size of eta and eta_RM. They may vary by S, J, T,
+        # but want to allow user to enter one that varies by only S, S and J,
         # S and T, or T and S and J.
-        param_in = getattr(self, "eta")
-        param_out = extrapolate_array(
-            param_in, dims=(self.T + self.S, self.S, self.J), item="eta"
-        )
-        setattr(self, "eta", param_out)
+        eta_params_to_TP = [
+            "eta",
+            "eta_RM",
+        ]
+        for item in eta_params_to_TP:
+            param_in = getattr(self, item)
+            param_out = extrapolate_array(
+                param_in, dims=(self.T + self.S, self.S, self.J), item=item
+            )
+            setattr(self, item, param_out)
+        # extrapolate lifetime ability e matrix over time dimension
         param_in = getattr(self, "e")
         param_out = extrapolate_array(
             param_in, dims=(self.T, self.S, self.J), item="e"