Fix tax function indexing

.github/workflows/deploy_docs.yml

@@ -16,9 +16,10 @@ jobs:
       - name: Setup Miniconda
         uses: conda-incubator/setup-miniconda@v2
         with:
-          activate-environment: ogcore-dev
+          miniforge-variant: Mambaforge
+          activate-environment: ogusa-dev
           environment-file: environment.yml
-          python-version: 3.7
+          python-version: "3.10"
           auto-activate-base: false
 
       - name: Build # Build Jupyter Book

.github/workflows/docs_check.yml

@@ -13,9 +13,10 @@ jobs:
       - name: Setup Miniconda
         uses: conda-incubator/setup-miniconda@v2
         with:
-          activate-environment: ogcore-dev
+          miniforge-variant: Mambaforge
+          activate-environment: ogusa-dev
           environment-file: environment.yml
-          python-version: 3.7
+          python-version: "3.10"
           auto-activate-base: false
 
       - name: Build # Build Jupyter Book

environment.yml

@@ -2,7 +2,7 @@ name: ogcore-dev
 channels:
 - conda-forge
 dependencies:
-- python>=3.7.7, <3.11  # This restriction can be removed as soon as we test Python 3.11
+- python>=3.7.7, <=3.11  # This restriction can be removed as soon as we test Python 3.11
 - ipython
 - setuptools
 - scipy>=1.7.1

ogcore/firm.py

@@ -1,4 +1,5 @@
 import numpy as np
+from ogcore import aggregates as aggr
 
 """
 ------------------------------------------------------------------------
@@ -679,3 +680,30 @@ def solve_L(Y, K, K_g, p, method, m=-1):
         ) ** (epsilon / (epsilon - 1))
 
     return L
+
+
+def adj_cost(K, Kp1, p, method):
+    r"""
+    Firm capital adjstment costs
+
+    ..math::
+        \Psi(K_{t}, K_{t+1}) = \frac{\psi}{2}\biggr(\frac{\biggr(\frac{I_{t}}{K_{t}}-\mu\biggl)^{2}}{\frac{I_{t}}{K_{t}}}\biggl)
+
+    Args:
+        K (array-like): Current period capital stock
+        Kp1 (array-like): One-period ahead capital stock
+        p (OG-USA Parameters class object): Model parameters
+        method (str): 'SS' or 'TPI'
+
+    Returns
+        Psi (array-like): Capital adjustment costs per unit of investment
+    """
+    if method == "SS":
+        ac_method = "total_ss"
+    else:
+        ac_method = "total_tpi"
+    Inv = aggr.get_I(None, Kp1, K, p, ac_method)
+
+    Psi = ((p.psi / 2) * (Inv / K - p.mu) ** 2) / (Inv / K)
+
+    return Psi

ogcore/parameter_plots.py

@@ -1118,7 +1118,7 @@ def plot_2D_taxfunc(
     # get tax rates for each point in the income support and plot
     fig, ax = plt.subplots()
     for i, tax_params in enumerate(tax_param_list):
-        tax_params = tax_params[rate_key][s][t]
+        tax_params = tax_params[rate_key][t][s]
         rates = txfunc.get_tax_rates(
             tax_params,
             X,

ogcore/parameters.py

@@ -387,6 +387,11 @@ def compute_default_params(self):
                 if len(tax_to_set[0]) > self.S:
                     for t, v in enumerate(tax_to_set):
                         tax_to_set[t] = tax_to_set[t][: self.S]
+                if len(tax_to_set[0]) < self.S:
+                    tax_params_to_add = [tax_to_set[:][-1]] * (
+                        self.S - len(tax_to_set[0])
+                    )
+                    tax_to_set[0].extend(tax_params_to_add)
                 setattr(self, item, tax_to_set)
             else:
                 print(

setup.py

@@ -23,7 +23,7 @@
         ]
     },
     include_packages=True,
-    python_requires=">=3.7.7, <3.11",
+    python_requires=">=3.7.7, <=3.11",
     install_requires=[
         "scipy>=1.7.1",
         "pandas>=1.2.5",

tests/test_firm.py

@@ -938,3 +938,60 @@ def test_solve_L(Y, K, Kg, p, method, expected):
     """
     L = firm.solve_L(Y, K, Kg, p, method)
     assert np.allclose(L, expected, atol=1e-6)
+
+
+p1 = Specifications()
+p1.psi = 4.0
+p1.g_n_ss = 0.01
+p1.g_y = 0.03
+p1.delta = 0.05
+p1.mu = 0.090759079
+K_1 = 5
+Kp1_1 = 5
+expected_Psi_1 = 0.0
+expected_dPsidK_1 = 0.0
+expected_dPsidKp1_1 = 0.0
+
+p2 = Specifications()
+p2.psi = 2.0
+p2.g_n_ss = 0.0
+p2.g_y = 0.03
+p2.delta = 0.05
+p2.mu = 0.05
+K_2 = 6
+Kp1_2 = 6
+expected_Psi_2 = 0.011527985
+expected_dPsidK_2 = -0.122196836
+expected_dPsidKp1_2 = 0.102296044
+
+
+p3 = Specifications()
+p3.psi = 4.0
+p3.g_n_ss = 0.0
+p3.g_n = np.array([-0.01, 0.02, 0.03, 0.0])
+p3.T = 3
+p3.g_y = 0.04
+p3.delta = 0.05
+p3.mu = 0.05
+K_3 = np.array([4, 4.5, 5.5])
+Kp1_3 = np.array([4.5, 5.5, 5])
+expected_Psi_3 = np.array([0.309124823, 0.534408906, -1.520508524])
+expected_dPsidK_3 = np.array([-0.805820108, -0.846107505, 2.657143029])
+expected_dPsidKp1_3 = np.array([0.479061039, 0.43588367, -62.31580895])
+
+
+@pytest.mark.parametrize(
+    "K,Kp1,p,method,expected",
+    [
+        (K_1, Kp1_1, p1, "SS", expected_Psi_1),
+        (K_2, Kp1_2, p2, "SS", expected_Psi_2),
+        (K_3, Kp1_3, p3, "TPI", expected_Psi_3),
+    ],
+    ids=["Zero cost", "Non-zero cost", "TPI"],
+)
+def test_adj_cost(K, Kp1, p, method, expected):
+    """
+    Test of the firm capital adjustment cost function.
+    """
+    test_val = firm.adj_cost(K, Kp1, p, method)
+    assert np.allclose(test_val, expected)

tests/test_parameters.py

@@ -186,3 +186,12 @@ def test_conditional_validator():
     new_specs = {"budget_balance": True, "baseline_spending": True}
     specs.update_specifications(new_specs, raise_errors=False)
     assert len(specs.errors) > 0
+
+
+def test_expand_taxfunc_params():
+    specs = Specifications()
+    new_specs = {"etr_params": [[[0.35]]]}
+    specs.update_specifications(new_specs)
+    assert len(specs.etr_params) == specs.T + specs.S
+    assert len(specs.etr_params[0]) == specs.S
+    assert specs.etr_params[0][0][0] == 0.35