test updates to get_I_g

jdebacker · Sep 12, 2024 · 45476d6 · 45476d6
1 parent be7d5ae
commit 45476d6
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 7 deletions.
diff --git a/ogcore/fiscal.py b/ogcore/fiscal.py
@@ -375,7 +375,7 @@ def get_r_gov(r, p, method):
     return r_gov
 
 
-def get_I_g(Y, p, method="SS"):
+def get_I_g(Y, Ig_baseline, p, method="SS"):
     r"""
     Find investment in public capital
 
@@ -384,6 +384,8 @@ def get_I_g(Y, p, method="SS"):
 
     Args:
         Y (array_like): aggregate output
+        Ig_baseline (array_like): public infrastructure investment in
+            the baseliine simulation
         p (OG-Core Specifications object): model parameters
         method (str): either 'SS' for steady-state or 'TPI' for transition path
 
@@ -392,9 +394,9 @@ def get_I_g(Y, p, method="SS"):
     """
     if p.baseline_spending:
         if method == "SS":
-            I_g = p.alpha_bs_I * Y
+            I_g = p.alpha_bs_I[-1] * Ig_baseline
         else:
-            I_g = p.alpha_bs_I[: p.T] * Y
+            I_g = p.alpha_bs_I[: p.T] * Ig_baseline
     else:
         if method == "SS":
             I_g = p.alpha_I[-1] * Y

diff --git a/tests/test_fiscal.py b/tests/test_fiscal.py
@@ -276,14 +276,37 @@ def test_get_r_gov(r, p, method, r_gov_expected):
     assert np.allclose(r_gov, r_gov_expected)
 
 
-def test_get_I_g():
+@pytest.mark.parametrize(
+    "baseline_spending,Ig_baseline,method",
+    [
+        (True, 3.5, "SS"),
+        (False, [3.5, 4.2], "TPI"),
+        (False, None, "SS"),
+        (False, None, "TPI"),
+    ],
+    ids=["baseline spend, SS", "baseline spend, TPI", "SS", "TPI"],
+)
+def test_get_I_g(baseline_spending, Ig_baseline, method):
     """
     Test function to determine investment in public capital
     """
+    p = Specifications()
+    p.baseline_spending = baseline_spending
     Y = np.array([0.2, 4.0])
-    alpha_I = np.array([0.1, 0.5])
-    expected = np.array([0.02, 2.0])
-    test_val = fiscal.get_I_g(Y, alpha_I)
+    p.alpha_I = np.array([0.1, 0.5])
+    p.alpha_bs_I = np.array([0.1, 1.0])
+    if method == "SS":
+        test_val = fiscal.get_I_g(Y[-1], Ig_baseline, p, method)
+        if baseline_spending:
+            expected = np.array([3.5])
+        else:
+            expected = np.array([2.0])
+    else:
+        test_val = fiscal.get_I_g(Y, Ig_baseline, p, method)
+        if baseline_spending:
+            expected = np.array([0.35, 4.2])
+        else:
+            expected = np.array([0.02, 2.0])
 
     assert np.allclose(test_val, expected)