Merge pull request #1363 from alanlujan91/rewards

Update Cobb-Douglas rewards

Documentation/CHANGELOG.md

@@ -33,6 +33,7 @@ Release Date: TBD
 - Suppress warning from calc_stable_points when it would be raised by inapplicable AgentType subclasses. [1493](https://github.com/econ-ark/HARK/pull/1493)
 - Fixes notation errors in IndShockConsumerType.make_euler_error_func from prior changes. [1495](https://github.com/econ-ark/HARK/pull/1495)
 - Fixes typos in IdentityFunction interpolator class. [1492](https://github.com/econ-ark/HARK/pull/1492)
+- Expands functionality of Cobb-Douglas aggregator for CRRA utility. [1363](https://github.com/econ-ark/HARK/pull/1363)
 
 ### 0.15.1
 

HARK/rewards.py

@@ -723,6 +723,40 @@ def inv(self, *args, **kwargs):
         return self.inverse(*args, **kwargs)
 
 
+def CDutility(c, d, c_share, d_bar):
+    return c**c_share * (d + d_bar) ** (1 - c_share)
+
+
+def CDutilityPc(c, d, c_share, d_bar):
+    return c_share * ((d + d_bar) / c) ** (1 - c_share)
+
+
+def CDutilityPd(c, d, c_share, d_bar):
+    return (1 - c_share) * (c / (d + d_bar)) ** c_share
+
+
+def CDutilityPc_inv(uc, d, c_share, d_bar):
+    return (d + d_bar) * (uc / c_share) ** (1 / (1 - c_share))
+
+
+def CRRACDutility(c, d, c_share, d_bar, CRRA):
+    return CDutility(c, d, c_share, d_bar) ** (1 - CRRA) / (1 - CRRA)
+
+
+def CRRACDutilityPc(c, d, c_share, d_bar, CRRA):
+    return c_share / c * CDutility(c, d, c_share, d_bar) ** (1 - CRRA)
+
+
+def CRRACDutilityPd(c, d, c_share, d_bar, CRRA):
+    return (1 - c_share) / (d + d_bar) * CDutility(c, d, c_share, d_bar) ** (1 - CRRA)
+
+
+def CRRACDutilityPc_inv(uc, d, c_share, d_bar, CRRA):
+    return (c_share / uc * (d + d_bar) ** (c_share * CRRA - c_share - CRRA + 1)) ** (
+        1 / (c_share * CRRA - c_share + 1)
+    )
+
+
 class UtilityFuncCRRA(UtilityFunction):
     """
     A class for representing a CRRA utility function.
@@ -877,37 +911,36 @@ class UtilityFuncCobbDouglas(UtilityFunction):
 
     Parameters
     ----------
-    EOS : float
-        The coefficient for elasticity of substitution.
-    factor : float
-        Factor productivity parameter. (e.g. TFP in production function)
+    c_share : float
+        Share parameter for consumption. Must be between 0 and 1.
+    d_bar : float
+        Intercept parameter for durable consumption. Must be non-negative.
     """
 
-    distance_criteria = ["EOS", "factor"]
-
-    def __init__(self, EOS, factor=1.0):
-        self.EOS = np.asarray(EOS)
-        self.factor = factor
-
-        assert np.isclose(
-            np.sum(self.EOS), 1.0
-        ), """The sum of the elasticity of substitution
-        parameters must be less than or equal to 1."""
+    distance_criteria = ["c_share", "d_bar"]
 
-        assert factor > 0, "Factor must be positive."
+    def __init__(self, c_share, d_bar=0):
+        self.c_share = c_share
+        self.d_bar = d_bar
 
-        self.dim = len(self.EOS)  # number of goods
+        assert 0 <= c_share <= 1, "Share parameter must be between 0 and 1."
 
-    def __call__(self, x):
+    def __call__(self, c, d):
         """
         Evaluate the utility function at a given level of consumption c.
         """
-        assert self.EOS.size == x.shape[-1], "x must be compatible with EOS"
-        cobb_douglas(x, self.EOS, self.factor)
+        return CDutility(c, d, self.c_share, self.d_bar)
 
-    def derivative(self, x, args=()):
-        assert self.EOS.size == x.shape[-1], "x must be compatible with EOS"
-        return cobb_douglas_pn(x, self.EOS, self.factor, args)
+    def derivative(self, c, d, axis=0):
+        if axis == 0:
+            return CDutilityPc(c, d, self.c_share, self.d_bar)
+        elif axis == 1:
+            return CDutilityPd(c, d, self.c_share, self.d_bar)
+        else:
+            raise ValueError(f"Axis must be 0 or 1, not {axis}")
+
+    def inverse(self, uc, d):
+        return CDutilityPc_inv(uc, d, self.c_share, self.d_bar)
 
 
 class UtilityFuncCobbDouglasCRRA(UtilityFuncCobbDouglas):
@@ -918,22 +951,33 @@ class UtilityFuncCobbDouglasCRRA(UtilityFuncCobbDouglas):
 
     Parameters
     ----------
-    EOS : float
-        The coefficient for elasticity of substitution.
-    factor : float
-        Factor productivity parameter. (e.g. TFP in production function)
+    c_share : float
+        Share parameter for consumption. Must be between 0 and 1.
+    d_bar : float
+        Intercept parameter for durable consumption. Must be non-negative.
     CRRA: float
         Coefficient of relative risk aversion.
     """
 
-    distance_criteria = ["EOS", "factor", "CRRA"]
+    distance_criteria = ["c_share", "d_bar", "CRRA"]
 
-    def __init__(self, EOS, factor, CRRA):
-        super().__init__(EOS, factor)
+    def __init__(self, c_share, CRRA, d_bar=0):
+        super().__init__(c_share, d_bar)
         self.CRRA = CRRA
 
-    def __call__(self, x):
-        return CRRAutility(cobb_douglas(x, self.EOS, self.factor), self.CRRA)
+    def __call__(self, c, d):
+        return CRRACDutility(c, d, self.c_share, self.d_bar, self.CRRA)
+
+    def derivative(self, c, d, axis=0):
+        if axis == 0:
+            return CRRACDutilityPc(c, d, self.c_share, self.d_bar, self.CRRA)
+        elif axis == 1:
+            return CRRACDutilityPd(c, d, self.c_share, self.d_bar, self.CRRA)
+        else:
+            raise ValueError(f"Axis must be 0 or 1, not {axis}")
+
+    def inverse(self, uc, d):
+        return CRRACDutilityPc_inv(uc, d, self.c_share, self.d_bar, self.CRRA)
 
 
 class UtilityFuncConstElastSubs(UtilityFunction):