more tests

Author: jdebacker

ogcore/pensions.py

@@ -197,14 +197,14 @@ def DB_amount(w, e, n, j, p):
         w_S = np.append((p.w_preTP * np.ones(p.S))[:(-per_rmn)], w)
         n_S = np.append(p.n_preTP[:(-per_rmn), j], n)
 
-        DB_s = np.zeros(p.retirement_age)
+        DB_s = np.zeros(p.retire)
         DB = np.zeros(p.S)
         # TODO: we set a rep_rate_py in params, but not rep_rate.  What is it???
         DB = DB_1dim_loop(
             w_S,
             p.e[:, j],
             n_S,
-            p.retirement_age,
+            p.retire,
             p.S,
             p.g_y,
             L_inc_avg_s,
@@ -219,7 +219,7 @@ def DB_amount(w, e, n, j, p):
 
     else:
         if np.ndim(n) == 1:
-            DB_s = np.zeros(p.retirement_age)
+            DB_s = np.zeros(p.retire)
             DB = np.zeros(p.S)
             DB = DB_1dim_loop(
                 w,
@@ -238,14 +238,14 @@ def DB_amount(w, e, n, j, p):
             )
 
         elif np.ndim(n) == 2:
-            DB_sj = np.zeros((p.retirement_age, p.J))
+            DB_sj = np.zeros((p.retire, p.J))
             DB = np.zeros((p.S, p.J))
             L_inc_avg_sj = np.zeros((p.last_career_yrs, p.J))
             DB = DB_2dim_loop(
                 w,
                 e,
                 n,
-                p.retirement_age,
+                p.retire,
                 p.S,
                 p.g_y,
                 L_inc_avg_sj,
@@ -553,7 +553,7 @@ def delta_ret(self, r, Y, p):
 
 @numba.jit
 def deriv_DB_loop(
-    w, e, S, S_ret, per_rmn, d_theta, last_career_yrs, rep_rate_py, yr_contr
+    w, e, S, S_ret, per_rmn, last_career_yrs, rep_rate_py, yr_contr
 ):
     d_theta = np.zeros(per_rmn)
     num_per_retire = S - S_ret
@@ -626,9 +626,9 @@ def PS_2dim_loop(w, e, n, S_ret, S, J, g_y, vpoint, factor, L_inc_avg_sj, PS):
             L_inc_avg_sj[s, :] = (
                 w[s] / np.exp(g_y * (u - s)) * e[s, :] * n[s, :]
             )
-        PS[u, :] = (
-            MONTHS_IN_A_YEAR * vpoint * L_inc_avg_sj.sum(axis=0)
-        ) / (factor * THOUSAND)
+        PS[u, :] = (MONTHS_IN_A_YEAR * vpoint * L_inc_avg_sj.sum(axis=0)) / (
+            factor * THOUSAND
+        )
 
     return PS
 

tests/test_pensions.py

@@ -79,7 +79,7 @@ def test_replacement_rate_vals(n, w, factor, j, p_in, expected):
 # })
 p.S = 7
 p.rep_rate_py = 0.2
-p.retirement_age = 4
+p.retire = 4
 p.last_career_yrs = 3
 p.yr_contr = 4
 p.g_y = 0.03
@@ -89,24 +89,136 @@ def test_replacement_rate_vals(n, w, factor, j, p_in, expected):
 n = np.array([0.4, 0.45, 0.4, 0.42, 0.3, 0.2, 0.2])
 L_inc_avg = np.zeros(0)
 L_inc_avg_s = np.zeros(p.last_career_yrs)
-DB_s = np.zeros(p.retirement_age)
 DB = np.zeros(p.S)
-DB_loop_expected1 = np.array([0, 0, 0, 0, 0.337864778, 0.327879365, 0.318189065])
-args1 = w, e, n, p.retirement_age, p.S, p.g_y, L_inc_avg_s, L_inc_avg, DB_s, DB
+DB_loop_expected1 = np.array(
+    [0, 0, 0, 0, 0.337864778, 0.327879365, 0.318189065]
+)
+args1 = (
+    w,
+    e,
+    n,
+    p.retire,
+    p.S,
+    p.g_y,
+    L_inc_avg_s,
+    L_inc_avg,
+    DB,
+    p.last_career_yrs,
+    p.rep_rate_py,
+    p.yr_contr,
+)
 
 test_data = [(args1, DB_loop_expected1)]
-#             (classes2, args2, NDC_expected2)]
 
-@pytest.mark.parametrize('args,DB_loop_expected', test_data,
-                         ids=['SS/Complete'])
+
+@pytest.mark.parametrize(
+    "args,DB_loop_expected", test_data, ids=["SS/Complete"]
+)
 def test_DB_1dim_loop(args, DB_loop_expected):
     """
     Test of the pensions.DB_1dim_loop() function.
     """
 
-    w, e, n, S_ret, S, g_y, L_inc_avg_s, L_inc_avg, DB_s, DB = args
+    (
+        w,
+        e,
+        n,
+        S_ret,
+        S,
+        g_y,
+        L_inc_avg_s,
+        L_inc_avg,
+        DB,
+        last_career_yrs,
+        rep_rate_py,
+        yr_contr,
+    ) = args
     DB_loop = pensions.DB_1dim_loop(
-        w, e, n, S_ret, S, g_y, L_inc_avg_s, L_inc_avg, DB,
-        p.last_career_yrs,
-        p.rep_rate_py, p.yr_contr)
-    assert (np.allclose(DB_loop, DB_loop_expected))
+        w,
+        e,
+        n,
+        S_ret,
+        S,
+        g_y,
+        L_inc_avg_s,
+        L_inc_avg,
+        DB,
+        last_career_yrs,
+        rep_rate_py,
+        yr_contr,
+    )
+    assert np.allclose(DB_loop, DB_loop_expected)
+
+
+p = Specifications()
+p.S = 7
+p.retire = 4
+per_rmn = p.S
+p.last_career_yrs = 3
+p.yr_contr = p.retire
+p.rep_rate_py = 0.2
+p.g_y = 0.03
+w = np.array([1.2, 1.1, 1.21, 1, 1.01, 0.99, 0.8])
+e = np.array([1.1, 1.11, 0.9, 0.87, 0.87, 0.7, 0.6])
+deriv_DB_loop_expected = np.array(
+    [0.352, 0.3256, 0.2904, 0.232, 0.0, 0.0, 0.0]
+)
+d_theta_empty = np.zeros_like(w)
+args3 = (
+    w,
+    e,
+    p.S,
+    p.retire,
+    per_rmn,
+    p.last_career_yrs,
+    p.rep_rate_py,
+    p.yr_contr,
+)
+
+test_data = [(args3, deriv_DB_loop_expected)]
+
+
+@pytest.mark.parametrize("args,deriv_DB_loop_expected", test_data)
+def test_deriv_DB_loop(args, deriv_DB_loop_expected):
+    """
+    Test of the pensions.deriv_DB_loop() function.
+    """
+    (w, e, S, retire, per_rmn, last_career_yrs, rep_rate_py, yr_contr) = args
+    deriv_DB_loop = pensions.deriv_DB_loop(
+        w, e, S, retire, per_rmn, last_career_yrs, rep_rate_py, yr_contr
+    )
+
+    assert np.allclose(deriv_DB_loop, deriv_DB_loop_expected)
+
+
+p = Specifications()
+p.S = 7
+p.retire = 4
+p.vpoint = 0.4
+w = np.array([1.2, 1.1, 1.21, 1, 1.01, 0.99, 0.8])
+e = np.array([1.1, 1.11, 0.9, 0.87, 0.87, 0.7, 0.6])
+p.g_y = 0.03
+factor = 2
+d_theta_empty = np.zeros_like(w)
+deriv_PS_loop_expected1 = np.array(
+    [0.003168, 0.0029304, 0.0026136, 0.002088, 0, 0, 0]
+)
+args3 = (w, e, p.S, p.retire, per_rmn, d_theta_empty, p.vpoint, factor)
+
+test_data = [(args3, deriv_PS_loop_expected1)]
+
+
+@pytest.mark.parametrize(
+    "args,deriv_PS_loop_expected", test_data, ids=["SS/Complete"]
+)
+def test_deriv_PS_loop(args, deriv_PS_loop_expected):
+    """
+    Test of the pensions.deriv_PS_loop() function.
+    """
+    (w, e, S, retire, per_rmn, d_theta_empty, vpoint, factor) = args
+
+    deriv_PS_loop = pensions.deriv_PS_loop(
+        w, e, S, retire, per_rmn, d_theta_empty, vpoint, factor
+    )
+
+    assert np.allclose(deriv_PS_loop, deriv_PS_loop_expected)