Add basic HH functions

ogcore/household.py

@@ -14,6 +14,23 @@
 ------------------------------------------------------------------------
 """
 
+def mu_c(c, sigma):
+    r"""
+    Compute the marginal utility of consumption.
+
+    .. math::
+        MU_{c} = c^{-\sigma}
+
+    Args:
+        c (array_like): household consumption
+        sigma (scalar): coefficient of relative risk aversion
+
+    Returns:
+        output (array_like): marginal utility of consumption
+
+    """
+
+    return c**(-sigma)
 
 def marg_ut_cons(c, sigma):
     r"""
@@ -45,6 +62,25 @@ def marg_ut_cons(c, sigma):
     return output
 
 
+def mu_n(n, chi_n, p):
+    r"""
+    Compute the marginal disutility of labor.
+
+    .. math::
+        MDU_{n} = \chi^n_{s} n^{\frac{1}{\theta}}
+
+    Args:
+        n (array_like): household labor supply
+        chi_n (array_like): utility weights on disutility of labor
+        p (OG-Core Specifications object): model parameters
+
+    Returns:
+        output (array_like): marginal disutility of labor supply
+
+    """
+    return chi_n * n ** (1 / p.theta)
+
+
 def marg_ut_labor(n, chi_n, p):
     r"""
     Compute the marginal disutility of labor.
@@ -135,6 +171,77 @@ def marg_ut_labor(n, chi_n, p):
     output = np.squeeze(output)
     return output
 
+def mu_b(b, chi_b, p):
+    r"""
+    Compute the marginal utility of savings.
+
+    .. math::
+        MU_{b} = \chi^b_{j}b_{j,s,t}^{-\sigma}
+
+    Args:
+        b (array_like): household savings
+        chi_b (array_like): utility weights on savings
+        p (OG-Core Specifications object): model parameters
+
+    Returns:
+        output (array_like): marginal utility of savings
+
+    """
+    return chi_b * b**(-p.sigma)
+
+def inv_mu_c(value, sigma):
+    r"""
+    Compute the inverse of the marginal utility of consumption.
+
+    .. math::
+        c = \left(\frac{1}{val}\right)^{-1/\sigma}
+
+    Args:
+        value (array_like): marginal utility of consumption
+        sigma (scalar): coefficient of relative risk aversion
+
+    Returns:
+        output (array_like): household consumption
+
+    """
+    return value ** (-1 / sigma)
+
+def inv_mu_n(value, chi_n, p):
+    r"""
+    Compute the inverse of the marginal disutility of labor.
+
+    .. math::
+        n = \left(\frac{val}{\chi^n_{s}}\right)^{\theta}
+
+    Args:
+        value (array_like): marginal disutility of labor
+        chi_n (array_like): utility weights on disutility of labor
+        p (OG-Core Specifications object): model parameters
+
+    Returns:
+        output (array_like): household labor supply
+
+    """
+    return (value / chi_n) ** p.theta
+
+def inv_mu_b(value, chi_b, p):
+    r"""
+    Compute the inverse of the marginal utility of savings.
+
+    .. math::
+        b = \left(\frac{MU_b}{\chi^b_{j}}\right)^{-1/\sigma}
+
+    Args:
+        value (array_like): marginal utility of savings
+        chi_b (array_like): utility weights on savings
+        p (OG-Core Specifications object): model parameters
+
+    Returns:
+        output (array_like): household savings
+
+    """
+    return (value / chi_b) ** (-1 / p.sigma)
+
 
 def get_bq(BQ, j, p, method):
     r"""
@@ -238,6 +345,35 @@ def get_tr(TR, j, p, method):
     return tr
 
 
+def get_b_splus1(b, r, w, p_tilde, c, n, bq, net_tax, e, p):
+    r"""
+    Calculate next period household savings from the budget constraint.
+
+    .. math::
+        b_{j,s+1,t+1} = \frac{(1 + r_{t+1})b_{j,s,t} + w_{t+1}e_{j,s}n_{j,s,t}
+        + bq_{j,s,t} + tr_{j,s,t} - T_{j,s,t} - c_{j,s,t}}{p_{t+1}}
+        
+
+    Args:
+        b (Numpy array): household savings one period ahead
+        r (array_like): the real interest rate
+        w (array_like): the real wage rate
+        p_tilde (array_like): the ratio of real GDP to nominal GDP
+        n (Numpy array): household labor supply
+        bq (Numpy array): household bequests received
+        net_tax (Numpy array): household net taxes paid
+        e (Numpy array): effective labor units
+        p (OG-Core Specifications object): model parameters
+
+    Returns:
+        b_splus1 (Numpy array): household savings
+
+    """
+    b_splus1 = ((1+r) * b + w * e * n + bq - net_tax - c * p_tilde) * np.exp(-p.g_y)
+
+    return b_splus1
+
+
 def get_cons(r, w, p_tilde, b, b_splus1, n, bq, net_tax, e, p):
     r"""
     Calculate household consumption.
@@ -312,6 +448,44 @@ def get_ci(c_s, p_i, p_tilde, tau_c, alpha_c, method="SS"):
     return c_si
 
 
+def savings_Euler_inverse(
+        r, w, p_tilde, b_splus1, c_splus1, bq, factor, tr, ubi, theta, rho, etr_params, mtry_params, t, j, p, method
+):
+    r"""
+    Compute current consumption from next period's consumption and savings using the savings Euler equation.
+
+    .. math::
+        c_{j,s,t} = \Bigl( \tilde{p}_{t} e^{-\sigma g_y}
+        \biggl[\chi^b_j\rho_s(b_{j,s+1,t+1})^{-\sigma} +
+        \beta_j\bigl(1 - \rho_s\bigr)\Bigl(\frac{1 + r_{t+1}
+        \bigl[1 - \tau^{mtry}_{s+1,t+1}\bigr]}{\tilde{p}_{t+1}}\Bigr)
+        (c_{j,s+1,t+1})^{-\sigma}\biggr] \Bigr)^{-1/\sigma}
+
+    Args:
+        r (array_like): the real interest rate
+        w (array_like): the real wage rate
+        p_tilde (array_like): composite good price
+        b_splus1 (Numpy array): household savings one period ahead
+        c_splus1 (Numpy array): household consumption one period ahead
+        bq (Numpy array): household bequests received
+        factor (scalar): scaling factor converting model units to dollars
+        tr (Numpy array): government transfers to household
+        ubi (Numpy array): universal basic income payment
+        theta (Numpy array): social security replacement rate for each
+            lifetime income group
+        rho (Numpy array): mortality rates
+        etr_params (list): parameters of the effective tax rate
+            functions
+        mtry_params (list): parameters of the marginal tax rate
+            on capital income functions
+        t (int): model period
+        j (int): index of ability type
+        p (OG-Core Specifications object): model parameters
+        method (str): adjusts calculation dimensions based on 'SS' or
+            'TPI'
+    """
+
+
 def FOC_savings(
     r,
     w,

tests/test_household.py

@@ -18,67 +18,107 @@
 @pytest.mark.parametrize(
     "c,sigma,expected", test_data, ids=["Scalar 0", "Scalar 1", "Vector"]
 )
-def test_marg_ut_cons(c, sigma, expected):
+def test_mu_c(c, sigma, expected):
     # Test marginal utility of consumption calculation
-    test_value = household.marg_ut_cons(c, sigma)
+    test_value = household.mu_c(c, sigma)
 
     assert np.allclose(test_value, expected)
 
 
+
+test_data = [
+    (10, 1, .1),
+    (55.90169944, 2.5, 0.2),
+    (
+        np.array([9.18958684, 0.002913041, 0.273217159]),
+        3.2,
+        np.array([0.5, 6.2, 1.5]),
+    ),
+]
+
+@pytest.mark.parametrize(
+    "value,sigma,expected", test_data, ids=["Scalar 0", "Scalar 1", "Vector"]
+)
+def test_inv_mu_c(value, sigma, expected):
+    # Test inverse marginal utility of consumption calculation
+    test_value = household.inv_mu_c(value, sigma)
+
+    assert np.allclose(test_value, expected)
+
+
+
 # Tuples in order: n, p, expected result
 p1 = Specifications()
-p1.b_ellipse = 0.527
-p1.upsilon = 1.497
-p1.ltilde = 1.0
+p1.theta = 0.4
 p1.chi_n = 3.3
 
 p2 = Specifications()
-p2.b_ellipse = 0.527
-p2.upsilon = 0.9
-p2.ltilde = 1.0
-p2.chi_n = 3.3
+p2.theta = 0.4
+p2.chi_n = 2.0
 
 p3 = Specifications()
-p3.b_ellipse = 0.527
-p3.upsilon = 0.9
-p3.ltilde = 2.3
+p3.theta = 0.5
 p3.chi_n = 3.3
 
 p4 = Specifications()
-p4.b_ellipse = 2.6
-p4.upsilon = 1.497
-p4.ltilde = 1.0
+p4.theta = 1.5
 p4.chi_n = 3.3
 
 test_data = [
-    (0.87, p1, 2.825570309),
-    (0.0, p1, 0.0009117852028298067),
-    (0.99999, p1, 69.52423604),
-    (0.00001, p1, 0.005692782),
-    (0.8, p2, 1.471592068),
-    (0.8, p3, 0.795937549),
-    (0.8, p4, 11.66354267),
+    (0.87, p1, 2.329764758),
+    (0.0, p1, 0.0),
+    (0.99999, p1, 3.299917501),
+    (0.001, p1, 0.0000001043551628),
+    (0.8, p2, 1.144866804),
+    (0.8, p3, 2.112),
+    (0.8, p4, 2.843853791),
     (
-        np.array([[0.8, 0.9, 0.3], [0.5, 0.2, 0.99]]),
+        np.array([[0.87, 0.0], [0.99999, 0.001]]),
         p1,
         np.array(
             [
-                [2.364110379, 3.126796062, 1.014935377],
-                [1.4248841, 0.806333875, 6.987729463],
+                [2.329764758, 0.0],
+                [3.299917501, 0.0000001043551628],
             ]
         ),
     ),
 ]
 
-
 @pytest.mark.parametrize(
     "n,params,expected",
     test_data,
     ids=["1", "2", "3", "4", "5", "6", "7", "8"],
 )
-def test_marg_ut_labor(n, params, expected):
+def test_mu_n(n, params, expected):
     # Test marginal utility of labor calculation
-    test_value = household.marg_ut_labor(n, params.chi_n, params)
+    test_value = household.mu_n(n, params.chi_n, params)
+
+    assert np.allclose(test_value, expected)
+
+
+test_data = [
+    (2.329764758, p1, 0.87),
+    (0.0, p1, 0.0),
+    (3.299917501, p1, 0.99999),
+    (0.0000001043551628, p1, 0.001),
+    (1.144866804, p2, 0.8),
+    (2.112, p3, 0.8),
+    (2.843853791, p4, 0.8),
+    (
+        np.array([[2.329764758, 0.0], [3.299917501, 0.0000001043551628]]),
+        p1,
+        np.array([[0.87, 0.0], [0.99999, 0.001]]),
+    ),
+]
+
+@pytest.mark.parametrize(
+    "value,params,expected",
+    test_data,
+    ids=["1", "2", "3", "4", "5", "6", "7", "8"],
+)
+def test_inv_mu_n(value, params, expected):
+    # Test inverse marginal utility of labor calculation
+    test_value = household.inv_mu_n(value, params.chi_n, params)
 
     assert np.allclose(test_value, expected)