Labeled income process

Documentation/CHANGELOG.md

@@ -34,6 +34,7 @@ Release Date: TBD
 * Adds Methods to calculate Heterogenous Agent Jacobian matrices. [#1185](https://github.com/econ-ark/HARK/pull/1185)
 * Enhances `combine_indep_dstns` to work with labeled distributions (`DiscreteDistributionLabeled`). [#1191](htttps://github.com/econ-ark/HARK/pull/1191)
 * Updates the `numpy` random generator from `RandomState` to `Generator`. [#1193](https://github.com/econ-ark/HARK/pull/1193)
+* Turns the income and income+return distributions into `DiscreteDistributionLabeled` objects. [#1189](https://github.com/econ-ark/HARK/pull/1189)
 
 ### Minor Changes
 

HARK/ConsumptionSaving/ConsIndShockModel.py

@@ -35,6 +35,7 @@
 from HARK.datasets.SCF.WealthIncomeDist.SCFDistTools import income_wealth_dists_from_scf
 from HARK.distribution import (
     DiscreteDistribution,
+    DiscreteDistributionLabeled,
     IndexDistribution,
     Lognormal,
     MeanOneLogNormal,
@@ -867,7 +868,7 @@ def m_nrm_next(self, shocks, a_nrm, Rfree):
         float
            normalized market resources in the next period
         """
-        return Rfree / (self.PermGroFac * shocks[0]) * a_nrm + shocks[1]
+        return Rfree / (self.PermGroFac * shocks['PermShk']) * a_nrm + shocks['TranShk']
 
     def calc_EndOfPrdvP(self):
         """
@@ -886,7 +887,7 @@ def calc_EndOfPrdvP(self):
         """
 
         def vp_next(shocks, a_nrm, Rfree):
-            return shocks[0] ** (-self.CRRA) * self.vPfuncNext(
+            return shocks['PermShk'] ** (-self.CRRA) * self.vPfuncNext(
                 self.m_nrm_next(shocks, a_nrm, Rfree)
             )
 
@@ -1132,7 +1133,7 @@ def make_cubic_cFunc(self, mNrm, cNrm):
         """
 
         def vpp_next(shocks, a_nrm, Rfree):
-            return shocks[0] ** (-self.CRRA - 1.0) * self.vPPfuncNext(
+            return shocks['PermShk'] ** (-self.CRRA - 1.0) * self.vPPfuncNext(
                 self.m_nrm_next(shocks, a_nrm, Rfree)
             )
 
@@ -1170,7 +1171,7 @@ def make_EndOfPrdvFunc(self, EndOfPrdvP):
 
         def v_lvl_next(shocks, a_nrm, Rfree):
             return (
-                shocks[0] ** (1.0 - self.CRRA) * self.PermGroFac ** (1.0 - self.CRRA)
+                shocks['PermShk'] ** (1.0 - self.CRRA) * self.PermGroFac ** (1.0 - self.CRRA)
             ) * self.vFuncNext(self.m_nrm_next(shocks, a_nrm, Rfree))
 
         EndOfPrdv = self.DiscFacEff * expected(
@@ -3569,7 +3570,7 @@ def __init__(self, sigma, UnempPrb, IncUnemp, n_approx, seed=0):
         super().__init__(pmv=dstn_approx.pmv, atoms=dstn_approx.atoms, seed=seed)
 
 
-class BufferStockIncShkDstn(DiscreteDistribution):
+class BufferStockIncShkDstn(DiscreteDistributionLabeled):
     """
     A one-period distribution object for the joint distribution of income
     shocks (permanent and transitory), as modeled in the Buffer Stock Theory
@@ -3628,7 +3629,13 @@ def __init__(
 
         joint_dstn = combine_indep_dstns(perm_dstn, tran_dstn)
 
-        super().__init__(pmv=joint_dstn.pmv, atoms=joint_dstn.atoms, seed=seed)
+        super().__init__(
+            name='Joint distribution of permanent and transitory shocks to income',
+            var_names=['PermShk','TranShk'],
+            pmv=joint_dstn.pmv,
+            data=joint_dstn.atoms,
+            seed=seed
+        )
 
 
 # Make a dictionary to specify a "kinked R" idiosyncratic shock consumer

HARK/ConsumptionSaving/ConsMarkovModel.py

@@ -323,7 +323,7 @@ def calc_EndOfPrdvPP(self):
         """
 
         def vpp_next(shocks, a_nrm, Rfree):
-            return shocks[0] ** (-self.CRRA - 1.0) * self.vPPfuncNext(
+            return shocks['PermShk'] ** (-self.CRRA - 1.0) * self.vPPfuncNext(
                 self.m_nrm_next(shocks, a_nrm, Rfree)
             )
 
@@ -332,7 +332,7 @@ def vpp_next(shocks, a_nrm, Rfree):
             * self.Rfree
             * self.Rfree
             * self.PermGroFac ** (-self.CRRA - 1.0)
-            * calc_expectation(self.IncShkDstn, vpp_next, self.aNrmNow, self.Rfree)
+            * self.IncShkDstn.expected(vpp_next, self.aNrmNow, self.Rfree)
         )
         return EndOfPrdvPP
 

HARK/ConsumptionSaving/ConsPortfolioModel.py

@@ -530,7 +530,7 @@ def m_nrm_next(self, shocks, b_nrm_next):
         Calculate future realizations of market resources
         """
 
-        return b_nrm_next / (shocks[0] * self.PermGroFac) + shocks[1]
+        return b_nrm_next / (shocks["PermShk"] * self.PermGroFac) + shocks["TranShk"]
 
     def calc_EndOfPrdvP(self):
         """
@@ -548,7 +548,9 @@ def dvdb_dist(shocks, b_nrm, Share_next):
 
             dvdmAdj_next = self.vPfuncAdj_next(mNrm_next)
             if self.AdjustPrb < 1.0:
-                dvdmFxd_next = self.dvdmFuncFxd_next(mNrm_next, Share_next)
+                # Expand to the same dimensions as mNrm
+                Share_next_expanded = Share_next + np.zeros_like(mNrm_next)
+                dvdmFxd_next = self.dvdmFuncFxd_next(mNrm_next, Share_next_expanded)
                 # Combine by adjustment probability
                 dvdm_next = (
                     self.AdjustPrb * dvdmAdj_next
@@ -557,7 +559,7 @@ def dvdb_dist(shocks, b_nrm, Share_next):
             else:  # Don't bother evaluating if there's no chance that portfolio share is fixed
                 dvdm_next = dvdmAdj_next
 
-            return (shocks[0] * self.PermGroFac) ** (-self.CRRA) * dvdm_next
+            return (shocks["PermShk"] * self.PermGroFac) ** (-self.CRRA) * dvdm_next
 
         def dvds_dist(shocks, b_nrm, Share_next):
             """
@@ -568,7 +570,9 @@ def dvds_dist(shocks, b_nrm, Share_next):
             # No marginal value of Share if it's a free choice!
             dvdsAdj_next = np.zeros_like(mNrm_next)
             if self.AdjustPrb < 1.0:
-                dvdsFxd_next = self.dvdsFuncFxd_next(mNrm_next, Share_next)
+                # Expand to the same dimensions as mNrm
+                Share_next_expanded = Share_next + np.zeros_like(mNrm_next)
+                dvdsFxd_next = self.dvdsFuncFxd_next(mNrm_next, Share_next_expanded)
                 # Combine by adjustment probability
                 dvds_next = (
                     self.AdjustPrb * dvdsAdj_next
@@ -577,11 +581,13 @@ def dvds_dist(shocks, b_nrm, Share_next):
             else:  # Don't bother evaluating if there's no chance that portfolio share is fixed
                 dvds_next = dvdsAdj_next
 
-            return (shocks[0] * self.PermGroFac) ** (1.0 - self.CRRA) * dvds_next
+            return (shocks["PermShk"] * self.PermGroFac) ** (
+                1.0 - self.CRRA
+            ) * dvds_next
 
         # Calculate intermediate marginal value of bank balances by taking expectations over income shocks
-        dvdb_intermed = calc_expectation(
-            self.IncShkDstn, dvdb_dist, self.bNrmNext, self.ShareNext
+        dvdb_intermed = self.IncShkDstn.expected(
+            dvdb_dist, self.bNrmNext, self.ShareNext
         )
 
         dvdbNvrs_intermed = self.uPinv(dvdb_intermed)
@@ -591,8 +597,8 @@ def dvds_dist(shocks, b_nrm, Share_next):
         dvdbFunc_intermed = MargValueFuncCRRA(dvdbNvrsFunc_intermed, self.CRRA)
 
         # Calculate intermediate marginal value of risky portfolio share by taking expectations
-        dvds_intermed = calc_expectation(
-            self.IncShkDstn, dvds_dist, self.bNrmNext, self.ShareNext
+        dvds_intermed = self.IncShkDstn.expected(
+            dvds_dist, self.bNrmNext, self.ShareNext
         )
 
         dvdsFunc_intermed = BilinearInterp(dvds_intermed, self.bNrmGrid, self.ShareGrid)
@@ -611,7 +617,10 @@ def EndOfPrddvda_dist(shock, a_nrm, Share_next):
             Rport = self.Rfree + Share_next * Rxs
             b_nrm_next = Rport * a_nrm
 
-            return Rport * dvdbFunc_intermed(b_nrm_next, Share_next)
+            # Ensure shape concordance
+            Share_next_rep = Share_next + np.zeros_like(b_nrm_next)
+
+            return Rport * dvdbFunc_intermed(b_nrm_next, Share_next_rep)
 
         def EndOfPrddvds_dist(shock, a_nrm, Share_next):
 
@@ -620,16 +629,19 @@ def EndOfPrddvds_dist(shock, a_nrm, Share_next):
             Rport = self.Rfree + Share_next * Rxs
             b_nrm_next = Rport * a_nrm
 
+            # Make the shares match the dimension of b, so that it can be vectorized
+            Share_next_expand = Share_next + np.zeros_like(b_nrm_next)
+
             return Rxs * a_nrm * dvdbFunc_intermed(
-                b_nrm_next, Share_next
-            ) + dvdsFunc_intermed(b_nrm_next, Share_next)
+                b_nrm_next, Share_next_expand
+            ) + dvdsFunc_intermed(b_nrm_next, Share_next_expand)
 
         # Calculate end-of-period marginal value of assets by taking expectations
         self.EndOfPrddvda = (
             self.DiscFac
             * self.LivPrb
-            * calc_expectation(
-                self.RiskyDstn, EndOfPrddvda_dist, self.aNrm_tiled, self.ShareNext
+            * self.RiskyDstn.expected(
+                EndOfPrddvda_dist, self.aNrm_tiled, self.ShareNext
             )
         )
 
@@ -639,8 +651,8 @@ def EndOfPrddvds_dist(shock, a_nrm, Share_next):
         self.EndOfPrddvds = (
             self.DiscFac
             * self.LivPrb
-            * calc_expectation(
-                self.RiskyDstn, EndOfPrddvds_dist, self.aNrm_tiled, self.ShareNext
+            * self.RiskyDstn.expected(
+                EndOfPrddvds_dist, self.aNrm_tiled, self.ShareNext
             )
         )
 
@@ -787,11 +799,11 @@ def v_intermed_dist(shocks, b_nrm, Share_next):
             else:  # Don't bother evaluating if there's no chance that portfolio share is fixed
                 v_next = vAdj_next
 
-            return (shocks[0] * self.PermGroFac) ** (1.0 - self.CRRA) * v_next
+            return (shocks["PermShk"] * self.PermGroFac) ** (1.0 - self.CRRA) * v_next
 
         # Calculate intermediate value by taking expectations over income shocks
-        v_intermed = calc_expectation(
-            self.IncShkDstn, v_intermed_dist, self.bNrmNext, self.ShareNext
+        v_intermed = self.IncShkDstn.expected(
+            v_intermed_dist, self.bNrmNext, self.ShareNext
         )
 
         vNvrs_intermed = self.uinv(v_intermed)
@@ -807,15 +819,17 @@ def EndOfPrdv_dist(shock, a_nrm, Share_next):
             Rport = self.Rfree + Share_next * Rxs
             b_nrm_next = Rport * a_nrm
 
-            return vFunc_intermed(b_nrm_next, Share_next)
+            # Make an extended share_next of the same dimension as b_nrm so
+            # that the function can be vectorized
+            Share_next_extended = Share_next + np.zeros_like(b_nrm_next)
+
+            return vFunc_intermed(b_nrm_next, Share_next_extended)
 
         # Calculate end-of-period value by taking expectations
         self.EndOfPrdv = (
             self.DiscFac
             * self.LivPrb
-            * calc_expectation(
-                self.RiskyDstn, EndOfPrdv_dist, self.aNrm_tiled, self.ShareNext
-            )
+            * self.RiskyDstn.expected(EndOfPrdv_dist, self.aNrm_tiled, self.ShareNext)
         )
 
         self.EndOfPrdvNvrs = self.uinv(self.EndOfPrdv)
@@ -1050,14 +1064,16 @@ def r_port(self, shocks, share):
         Calculate future realizations of market resources
         """
 
-        return (1.0 - share) * self.Rfree + share * shocks[2]
+        return (1.0 - share) * self.Rfree + share * shocks["Risky"]
 
     def m_nrm_next(self, shocks, a_nrm, r_port):
         """
         Calculate future realizations of market resources
         """
 
-        return r_port * a_nrm / (shocks[0] * self.PermGroFac) + shocks[1]
+        return (
+            r_port * a_nrm / (shocks["PermShk"] * self.PermGroFac) + shocks["TranShk"]
+        )
 
     def calc_EndOfPrdvP(self):
         """
@@ -1107,27 +1123,32 @@ def EndOfPrddvda_dists(shocks, a_nrm, shares):
             r_port = self.r_port(shocks, shares)
             m_nrm_next = self.m_nrm_next(shocks, a_nrm, r_port)
 
+            # Expand shares to the shape of m so that operations can be vectorized
+            shares_expanded = shares + np.zeros_like(m_nrm_next)
+
             return (
-                r_port * self.uP(shocks[0] * self.PermGroFac) * dvdm(m_nrm_next, shares)
+                r_port
+                * self.uP(shocks["PermShk"] * self.PermGroFac)
+                * dvdm(m_nrm_next, shares_expanded)
             )
 
         def EndOfPrddvds_dist(shocks, a_nrm, shares):
-            Rxs = shocks[2] - self.Rfree
+            Rxs = shocks["Risky"] - self.Rfree
             r_port = self.r_port(shocks, shares)
             m_nrm_next = self.m_nrm_next(shocks, a_nrm, r_port)
 
-            return Rxs * a_nrm * self.uP(shocks[0] * self.PermGroFac) * dvdm(
+            return Rxs * a_nrm * self.uP(shocks["PermShk"] * self.PermGroFac) * dvdm(
                 m_nrm_next, shares
-            ) + (shocks[0] * self.PermGroFac) ** (1.0 - self.CRRA) * dvds(
+            ) + (shocks["PermShk"] * self.PermGroFac) ** (1.0 - self.CRRA) * dvds(
                 m_nrm_next, shares
             )
 
         # Calculate end-of-period marginal value of assets by taking expectations
         self.EndOfPrddvda = (
             self.DiscFac
             * self.LivPrb
-            * calc_expectation(
-                self.ShockDstn, EndOfPrddvda_dists, self.aNrm_tiled, self.Share_tiled
+            * self.ShockDstn.expected(
+                EndOfPrddvda_dists, self.aNrm_tiled, self.Share_tiled
             )
         )
 
@@ -1137,8 +1158,8 @@ def EndOfPrddvds_dist(shocks, a_nrm, shares):
         self.EndOfPrddvds = (
             self.DiscFac
             * self.LivPrb
-            * calc_expectation(
-                self.ShockDstn, EndOfPrddvds_dist, self.aNrm_tiled, self.Share_tiled
+            * self.ShockDstn.expected(
+                EndOfPrddvds_dist, self.aNrm_tiled, self.Share_tiled
             )
         )
 
@@ -1159,14 +1180,12 @@ def v_dist(shocks, a_nrm, shares):
             else:  # Don't bother evaluating if there's no chance that portfolio share is fixed
                 v_next = vAdj_next
 
-            return (shocks[0] * self.PermGroFac) ** (1.0 - self.CRRA) * v_next
+            return (shocks["PermShk"] * self.PermGroFac) ** (1.0 - self.CRRA) * v_next
 
         self.EndOfPrdv = (
             self.DiscFac
             * self.LivPrb
-            * calc_expectation(
-                self.ShockDstn, v_dist, self.aNrm_tiled, self.Share_tiled
-            )
+            * self.ShockDstn.expected(v_dist, self.aNrm_tiled, self.Share_tiled)
         )
 
         self.EndOfPrdvNvrs = self.uinv(self.EndOfPrdv)

HARK/ConsumptionSaving/ConsRiskyAssetModel.py

@@ -19,6 +19,7 @@
 )
 from HARK.distribution import (
     DiscreteDistribution,
+    DiscreteDistributionLabeled,
     IndexDistribution,
     Lognormal,
     Bernoulli,
@@ -157,16 +158,38 @@ def update_ShockDstn(self):
         -------
         None
         """
+
+        # Create placeholder distributions
         if "RiskyDstn" in self.time_vary:
-            self.ShockDstn = [
+            dstn_list = [
                 combine_indep_dstns(self.IncShkDstn[t], self.RiskyDstn[t])
                 for t in range(self.T_cycle)
             ]
         else:
-            self.ShockDstn = [
+            dstn_list = [
                 combine_indep_dstns(self.IncShkDstn[t], self.RiskyDstn)
                 for t in range(self.T_cycle)
             ]
+
+        # Names of the variables (hedging for the unlikely case that in
+        # some index of IncShkDstn variables are in a switched order)
+        names_list = [
+            list(self.IncShkDstn[t].dataset.data_vars.keys()) + ["Risky"]
+            for t in range(self.T_cycle)
+        ]
+
+        conditional = {
+            "pmv": [x.pmv for x in dstn_list],
+            "data": [x.atoms for x in dstn_list],
+            "var_names": names_list,
+        }
+
+        # Now create the actual distribution using the index and labeled class
+        self.ShockDstn = IndexDistribution(
+            engine=DiscreteDistributionLabeled,
+            conditional=conditional,
+        )
+
         self.add_to_time_vary("ShockDstn")
 
         # Mark whether the risky returns and income shocks are independent (they are)