use lists in dy rev table

PSLmodels · Oct 15, 2023 · 805aa62 · 805aa62
1 parent 03e0421
commit 805aa62
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 12 deletions.
diff --git a/ogcore/output_tables.py b/ogcore/output_tables.py
@@ -605,18 +605,27 @@ def dynamic_revenue_decomposition(
         year_list.append("SS")
     table_dict = {"Year": year_list}
     T, S, J = base_params.T, base_params.S, base_params.J
-    base_etr_params_4D = np.tile(
-        base_params.etr_params[:T, :, :].reshape(
-            T, S, 1, base_params.etr_params.shape[2]
-        ),
-        (1, 1, J, 1),
-    )
-    reform_etr_params_4D = np.tile(
-        reform_params.etr_params[:T, :, :].reshape(
-            T, S, 1, reform_params.etr_params.shape[2]
-        ),
-        (1, 1, J, 1),
-    )
+    num_params = len(base_params.etr_params[0][0])
+    base_etr_params_4D = [
+        [
+            [
+                [base_params.etr_params[t][s][i] for i in range(num_params)]
+                for j in range(J)
+            ]
+            for s in range(S)
+        ]
+        for t in range(T)
+    ]
+    reform_etr_params_4D = [
+        [
+            [
+                [reform_params.etr_params[t][s][i] for i in range(num_params)]
+                for j in range(J)
+            ]
+            for s in range(S)
+        ]
+        for t in range(T)
+    ]
     tax_rev_dict = {"indiv": {}, "biz": {}, "total": {}}
     indiv_liab = {}
     # Baseline IIT + payroll tax liability

diff --git a/tests/test_output_tables.py b/tests/test_output_tables.py
@@ -149,6 +149,13 @@ def test_dynamic_revenue_decomposition(include_business_tax, full_break_out):
     reform_params.capital_income_tax_noncompliance_rate = np.zeros(
         (reform_params.T, reform_params.J)
     )
+    # check if tax parameters are a numpy array
+    # this is relevant for cached parameter arrays saved before
+    # tax params were put in lists
+    if isinstance(base_params.etr_params, np.ndarray):
+        base_params.etr_params = base_params.etr_params.tolist()
+    if isinstance(reform_params.etr_params, np.ndarray):
+        reform_params.etr_params = reform_params.etr_params.tolist()
     df = output_tables.dynamic_revenue_decomposition(
         base_params,
         base_tpi,