From 3cfd20a47ca5aca4b05122ffc286d284411f3191 Mon Sep 17 00:00:00 2001
From: jdebacker <jason.debacker@gmail.com>
Date: Thu, 18 Apr 2024 15:29:54 -0400
Subject: [PATCH] update to account for arrays

---
 ccc/calcfunctions.py | 23 ++++++++++++++++++++---
 ccc/calculator.py    |  2 +-
 ccc/constants.py     |  8 ++++----
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/ccc/calcfunctions.py b/ccc/calcfunctions.py
index b8363dd7..61485462 100644
--- a/ccc/calcfunctions.py
+++ b/ccc/calcfunctions.py
@@ -222,7 +222,20 @@ def npv_tax_depr(df, r, pi, land_expensing):
     return z
 
 
-def eq_coc(delta, z, w, u, u_d, inv_tax_credit, psi, nu, pi, r, re_credit=None, asset_type="None"):
+def eq_coc(
+    delta,
+    z,
+    w,
+    u,
+    u_d,
+    inv_tax_credit,
+    psi,
+    nu,
+    pi,
+    r,
+    re_credit=None,
+    asset_code="None",
+):
     r"""
     Compute the cost of capital
 
@@ -250,8 +263,12 @@ def eq_coc(delta, z, w, u, u_d, inv_tax_credit, psi, nu, pi, r, re_credit=None,
 
     """
     # case for assets eligible for R&E credit
-    if (asset_type in RE_ASSETS) and (re_credit is not None):
-        inv_tax_credit += re_credit
+    if (asset_code is not None) and (re_credit is not None) & isinstance(
+        delta, np.ndarray
+    ):
+        idx = [element in RE_ASSETS for element in asset_code]
+        inv_tax_credit = np.ones_like(delta)
+        inv_tax_credit[idx] += re_credit
     rho = (
         ((r - pi + delta) / (1 - u))
         * (1 - inv_tax_credit * nu - u_d * z * (1 - psi * inv_tax_credit))
diff --git a/ccc/calculator.py b/ccc/calculator.py
index 269ad95d..e20cb349 100644
--- a/ccc/calculator.py
+++ b/ccc/calculator.py
@@ -203,7 +203,7 @@ def calc_base(self):
                     self.__p.inflation_rate,
                     self.__p.r[t][f],
                     self.__p.re_credit,
-                    dfs[t]["asset_type"] #TODO: check that this is consistent with list if constants
+                    dfs[t]["bea_asset_code"],
                 )
                 if not self.__p.inventory_expensing:
                     idx = dfs[t]["asset_name"] == "Inventories"
diff --git a/ccc/constants.py b/ccc/constants.py
index 1e5c274e..03275974 100644
--- a/ccc/constants.py
+++ b/ccc/constants.py
@@ -23,13 +23,13 @@
 # TODO: perhaps make as a dict so that can vary across years?
 # And if policy variant, maybe move to default params?
 RE_ASSETS = [
-    "Credit-eligible own account software",
-    "Credit-eligible research and development",
-    "Wind and solar power structures",
+    "ENS3",  # "Credit-eligible own account software",
+    "RD70",  # "Credit-eligible research and development",
+    "SU60",  # "Wind and solar power structures",
 ]
 
 RE_INDUSTRIES = [
-    "Computer_and_Electronic_Product_Manufacturing",
+    "3340",  # "Computer_and_Electronic_Product_Manufacturing",
 ]
 
 MAJOR_IND_ORDERED = [