Make non-OO consumption-saving solvers

HARK/ConsumptionSaving/ConsPortfolioModel.py

@@ -29,6 +29,8 @@
     ValueFuncCRRA,
 )
 from HARK.metric import MetricObject
+from HARK.rewards import UtilityFuncCRRA
+from HARK.distribution import expected
 
 
 # Define a class to represent the single period solution of the portfolio choice problem
@@ -180,6 +182,7 @@
         else:
             solver = ConsPortfolioJointDistSolver
         self.solve_one_period = make_one_period_oo_solver(solver)
+        #self.solve_one_period = solve_one_period_ConsPortfolio
 
         self.update()
 
@@ -328,6 +331,382 @@
 
         # Set the solver for the portfolio model, and update various constructed attributes
         self.solve_one_period = make_one_period_oo_solver(ConsSequentialPortfolioSolver)
+
+
+def solve_one_period_ConsPortfolio( solution_next,
+                                    ShockDstn,
+                                    IncShkDstn,
+                                    RiskyDstn,
+                                    LivPrb,
+                                    DiscFac,
+                                    CRRA,
+                                    Rfree,
+                                    PermGroFac,
+                                    BoroCnstArt,
+                                    aXtraGrid,
+                                    ShareGrid,
+                                    AdjustPrb,
+                                    ShareLimit,
+                                    vFuncBool,
+                                    DiscreteShareBool,
+                                    IndepDstnBool,):
+    """
+    Solve one period of a consumption-saving problem with portfolio allocation
+    between a riskless and risky asset. This function handles various sub-cases
+    or variations on the problem, including the possibility that the agent does
+    not necessarily get to update their portfolio share in every period, or that
+    they must choose a discrete rather than continuous risky share.
+
+    Parameters
+    ----------
+    solution_next : PortfolioSolution
+        Solution to next period's problem.
+    ShockDstn : Distribution
+        Joint distribution of permanent income shocks, transitory income shocks,
+        and risky returns.  This is only used if the input IndepDstnBool is False,
+        indicating that income and return distributions can't be assumed to be
+        independent.
+    IncShkDstn : Distribution
+        Discrete distribution of permanent income shocks and transitory income
+        shocks. This is only used if the input IndepDstnBool is True, indicating
+        that income and return distributions are independent.
+    RiskyDstn : Distribution
+       Distribution of risky asset returns. This is only used if the input
+       IndepDstnBool is True, indicating that income and return distributions
+       are independent.
+    LivPrb : float
+        Survival probability; likelihood of being alive at the beginning of
+        the succeeding period.
+    DiscFac : float
+        Intertemporal discount factor for future utility.
+    CRRA : float
+        Coefficient of relative risk aversion.
+    Rfree : float
+        Risk free interest factor on end-of-period assets.
+    PermGroFac : float
+        Expected permanent income growth factor at the end of this period.
+    BoroCnstArt: float or None
+        Borrowing constraint for the minimum allowable assets to end the
+        period with.  In this model, it is *required* to be zero.
+    aXtraGrid: np.array
+        Array of "extra" end-of-period asset values-- assets above the
+        absolute minimum acceptable level.
+    ShareGrid : np.array
+        Array of risky portfolio shares on which to define the interpolation
+        of the consumption function when Share is fixed. Also used when the
+        risky share choice is specified as discrete rather than continuous.
+    AdjustPrb : float
+        Probability that the agent will be able to update his portfolio share.
+    ShareLimit : float
+        Limiting lower bound of risky portfolio share as mNrm approaches infinity.
+    vFuncBool: boolean
+        An indicator for whether the value function should be computed and
+        included in the reported solution.
+    DiscreteShareBool : bool
+        Indicator for whether risky portfolio share should be optimized on the
+        continuous [0,1] interval using the FOC (False), or instead only selected
+        from the discrete set of values in ShareGrid (True).  If True, then
+        vFuncBool must also be True.
+    IndepDstnBool : bool
+        Indicator for whether the income and risky return distributions are in-
+        dependent of each other, which can speed up the expectations step.
+
+    Returns
+    -------
+    solution_now : PortfolioSolution
+        Solution to this period's problem.
+    """
+    # Make sure the individual is liquidity constrained.  Allowing a consumer to
+    # borrow *and* invest in an asset with unbounded (negative) returns is a bad mix.
+    if BoroCnstArt != 0.0:
+        raise ValueError("PortfolioConsumerType must have BoroCnstArt=0.0!")
+
+    # Make sure that if risky portfolio share is optimized only discretely, then
+    # the value function is also constructed (else this task would be impossible).
+    if DiscreteShareBool and (not vFuncBool):
+        raise ValueError(
+            "PortfolioConsumerType requires vFuncBool to be True when DiscreteShareBool is True!"
+        )
+
+    # Define the current period utility function and effective discount factor
+    uFunc = UtilityFuncCRRA(CRRA)
+    DiscFacEff = DiscFac * LivPrb  # "effective" discount factor
+
+    # Unpack next period's solution for easier access
+    vPfuncAdj_next = solution_next.vPfuncAdj
+    dvdmFuncFxd_next = solution_next.dvdmFuncFxd
+    dvdsFuncFxd_next = solution_next.dvdsFuncFxd
+    vFuncAdj_next = solution_next.vFuncAdj
+    vFuncFxd_next = solution_next.vFuncFxd
+
+    # Set a flag for whether the natural borrowing constraint is zero, which
+    # depends on whether the smallest transitory income shock is zero
+    BoroCnstNat_iszero = np.min(IncShkDstn.atoms[1]) == 0.0   
+
+    # Prepare to calculate end-of-period marginal values by creating an array
+    # of market resources that the agent could have next period, considering
+    # the grid of end-of-period assets and the distribution of shocks he might
+    # experience next period.
+
+    # Unpack the risky return shock distribution
+    Risky_next = RiskyDstn.atoms
+    RiskyMax = np.max(Risky_next)
+    RiskyMin = np.min(Risky_next)
+
+    # bNrm represents R*a, balances after asset return shocks but before income.
+    # This just uses the highest risky return as a rough shifter for the aXtraGrid.
+    if BoroCnstNat_iszero:
+        aNrmGrid = aXtraGrid
+        bNrmGrid = np.insert(RiskyMax * aXtraGrid, 0, RiskyMin * aXtraGrid[0])
+    else:
+        # Add an asset point at exactly zero
+        aNrmGrid = np.insert(aXtraGrid, 0, 0.0)
+        bNrmGrid = RiskyMax * np.insert(aXtraGrid, 0, 0.0)
+
+    # Get grid and shock sizes, for easier indexing
+    aNrmCount = aNrmGrid.size
+    ShareCount = ShareGrid.size
+
+    # Make tiled arrays to calculate future realizations of mNrm and Share when integrating over IncShkDstn
+    bNrmNext, ShareNext = np.meshgrid(bNrmGrid, ShareGrid, indexing="ij")
+
+    # Define functions that are used internally to evaluate future realizations
+    def calc_mNrm_next(S, b):
+        """
+        Calculate future realizations of market resources mNrm from the income
+        shock distribution S and normalized bank balances b.
+        """
+        return b / (S["PermShk"] * PermGroFac) + S["TranShk"]
+
+    def calc_dvdm_next(S, b, z):
+        """
+        Evaluate realizations of marginal value of market resources next period,
+        based on the income distribution S, values of bank balances bNrm, and
+        values of the risky share z.
+        """
+        mNrm_next = calc_mNrm_next(S, b)
+        dvdmAdj_next = vPfuncAdj_next(mNrm_next)
+
+        if AdjustPrb < 1.0:
+            # Expand to the same dimensions as mNrm
+            Share_next_expanded = z + np.zeros_like(mNrm_next)
+            dvdmFxd_next = dvdmFuncFxd_next(mNrm_next, Share_next_expanded)
+            # Combine by adjustment probability
+            dvdm_next = AdjustPrb * dvdmAdj_next + (1.0 - AdjustPrb) * dvdmFxd_next
+        else:  # Don't bother evaluating if there's no chance that portfolio share is fixed
+            dvdm_next = dvdmAdj_next
+
+        dvdm_next = (S["PermShk"] * PermGroFac) ** (-CRRA) * dvdm_next
+        return dvdm_next
+
+    def calc_dvds_next(S, b, z):
+        """
+        Evaluate realizations of marginal value of risky share next period, based
+        on the income distribution S, values of bank balances bNrm, and values of
+        the risky share z.
+        """
+        mNrm_next = calc_mNrm_next(S, b)
+
+        # No marginal value of Share if it's a free choice!
+        dvdsAdj_next = np.zeros_like(mNrm_next)
+
+        if AdjustPrb < 1.0:
+            # Expand to the same dimensions as mNrm
+            Share_next_expanded = z + np.zeros_like(mNrm_next)
+            dvdsFxd_next = dvdsFuncFxd_next(mNrm_next, Share_next_expanded)
+            # Combine by adjustment probability
+            dvds_next = AdjustPrb * dvdsAdj_next + (1.0 - AdjustPrb) * dvdsFxd_next
+        else:  # Don't bother evaluating if there's no chance that portfolio share is fixed
+            dvds_next = dvdsAdj_next
+
+        dvds_next = (S["PermShk"] * PermGroFac) ** (1.0 - CRRA) * dvds_next
+        return dvds_next
+
+    # Calculate end-of-period marginal value of assets and shares at each point
+    # in aNrm and ShareGrid. Does so by taking expectation of next period marginal
+    # values across income and risky return shocks.
+
+    # Calculate intermediate marginal value of bank balances by taking expectations over income shocks
+    dvdb_intermed = expected(calc_dvdm_next, IncShkDstn, args=(bNrmNext, ShareNext))
+    dvdbNvrs_intermed = uFunc.derinv(dvdb_intermed, order=(1, 0))
+    dvdbNvrsFunc_intermed = BilinearInterp(dvdbNvrs_intermed, bNrmGrid, ShareGrid)
+    dvdbFunc_intermed = MargValueFuncCRRA(dvdbNvrsFunc_intermed, CRRA)
+
+    # Calculate intermediate marginal value of risky portfolio share by taking expectations over income shocks
+    dvds_intermed = expected(calc_dvds_next, IncShkDstn, args=(bNrmNext, ShareNext))
+    dvdsFunc_intermed = BilinearInterp(dvds_intermed, bNrmGrid, ShareGrid)
+
+    # Make tiled arrays to calculate future realizations of bNrm and Share when integrating over RiskyDstn
+    aNrmNow, ShareNext = np.meshgrid(aNrmGrid, ShareGrid, indexing="ij")
+
+    # Define functions for calculating end-of-period marginal value
+    def calc_EndOfPrd_dvda(S, a, z):
+        '''
+        Compute end-of-period marginal value of assets at values a, conditional
+        on risky asset return S and risky share z.
+        '''
+        # Calculate future realizations of bank balances bNrm
+        Rxs = S - Rfree  # Excess returns
+        Rport = Rfree + z * Rxs  # Portfolio return
+        bNrm_next = Rport * a
+
+        # Ensure shape concordance
+        z_rep = z + np.zeros_like(bNrm_next)
+
+        # Calculate and return dvda
+        EndOfPrd_dvda = Rport * dvdbFunc_intermed(bNrm_next, z_rep)
+        return EndOfPrd_dvda
+
+    def EndOfPrddvds_dist(S, a, z):
+        '''
+        Compute end-of-period marginal value of risky share at values a, conditional
+        on risky asset return S and risky share z.
+        '''
+        # Calculate future realizations of bank balances bNrm
+        Rxs = S - Rfree  # Excess returns
+        Rport = Rfree + z * Rxs  # Portfolio return
+        bNrm_next = Rport * a
+
+        # Make the shares match the dimension of b, so that it can be vectorized
+        z_rep = z + np.zeros_like(bNrm_next)
+
+        # Calculate and return dvds
+        EndOfPrd_dvds = Rxs * a * dvdbFunc_intermed(bNrm_next, z_rep) + dvdsFunc_intermed(bNrm_next, z_rep)
+        return EndOfPrd_dvds
+
+    # Evaluate realizations of value and marginal value after asset returns are realized
+
+    # Calculate end-of-period marginal value of assets by taking expectations
+    EndOfPrd_dvda = DiscFacEff * expected(calc_EndOfPrd_dvda, RiskyDstn, args=(aNrmNow, ShareNext))
+    EndOfPrd_dvdaNvrs = uFunc.derinv(EndOfPrd_dvda)
+
+    # Calculate end-of-period marginal value of risky portfolio share by taking expectations
+    EndOfPrd_dvds = DiscFacEff * expected(EndOfPrddvds_dist, RiskyDstn, args=(aNrmNow, ShareNext))
+
+    # Now find the optimal (continuous) risky share on [0,1] by solving the first
+    # order condition EndOfPrd_dvds == 0.
+    FOC_s = EndOfPrd_dvds  # Relabel for convenient typing
+
+    # For each value of aNrm, find the value of Share such that FOC_s == 0
+    crossing = np.logical_and(FOC_s[:, 1:] <= 0.0, FOC_s[:, :-1] >= 0.0)
+    share_idx = np.argmax(crossing, axis=1)
+    # This represents the index of the segment of the share grid where dvds flips
+    # from positive to negative, indicating that there's a zero *on* the segment
+
+    # Calculate the fractional distance between those share gridpoints where the
+    # zero should be found, assuming a linear function; call it alpha
+    a_idx = np.arange(aNrmCount)
+    bot_s = ShareGrid[share_idx]
+    top_s = ShareGrid[share_idx + 1]
+    bot_f = FOC_s[a_idx, share_idx]
+    top_f = FOC_s[a_idx, share_idx + 1]
+    bot_c = EndOfPrd_dvdaNvrs[a_idx, share_idx]
+    top_c = EndOfPrd_dvdaNvrs[a_idx, share_idx + 1]
+    alpha = 1.0 - top_f / (top_f - bot_f)
+
+    # Calculate the continuous optimal risky share and optimal consumption
+    ShareAdj_now = (1.0 - alpha) * bot_s + alpha * top_s
+    cNrmAdj_now = (1.0 - alpha) * bot_c + alpha * top_c
+
+    # If agent wants to put more than 100% into risky asset, he is constrained.
+    # Likewise if he wants to put less than 0% into risky asset, he is constrained.
+    constrained_top = FOC_s[:, -1] > 0.0
+    constrained_bot = FOC_s[:, 0] < 0.0
+
+    # Apply those constraints to both risky share and consumption (but lower
+    # constraint should never be relevant)
+    ShareAdj_now[constrained_top] = 1.0
+    ShareAdj_now[constrained_bot] = 0.0
+    cNrmAdj_now[constrained_top] = EndOfPrd_dvdaNvrs[constrained_top, -1]
+    cNrmAdj_now[constrained_bot] = EndOfPrd_dvdaNvrs[constrained_bot, 0]
+
+    # When the natural borrowing constraint is *not* zero, then aNrm=0 is in the
+    # grid, but there's no way to "optimize" the portfolio if a=0, and consumption
+    # can't depend on the risky share if it doesn't meaningfully exist. Apply
+    # a small fix to the bottom gridpoint (aNrm=0) when this happens.
+    if (not BoroCnstNat_iszero):
+        ShareAdj_now[0] = 1.0
+        cNrmAdj_now[0] = EndOfPrd_dvdaNvrs[0, -1]
+
+    # Construct functions characterizing the solution for this period
+
+    # Calculate the endogenous mNrm gridpoints when the agent adjusts his portfolio,
+    # then construct the consumption function when the agent can adjust his share
+    mNrmAdj_now = np.insert(aNrmGrid + cNrmAdj_now, 0, 0.0)
+    cNrmAdj_now = np.insert(cNrmAdj_now, 0, 0.0)
+    cFuncAdj_now = LinearInterp(mNrmAdj_now, cNrmAdj_now)
+
+    # Construct the marginal value (of mNrm) function when the agent can adjust
+    vPfuncAdj_now = MargValueFuncCRRA(cFuncAdj_now, CRRA)
+
+    # Construct the consumption function when the agent *can't* adjust the risky
+    # share, as well as the marginal value of Share function
+    cFuncFxd_by_Share = []
+    dvdsFuncFxd_by_Share = []
+    for j in range(ShareCount):
+        cNrmFxd_temp = np.insert(EndOfPrd_dvdaNvrs[:, j], 0, 0.0)
+        mNrmFxd_temp = np.insert(aNrmGrid + cNrmFxd_temp[1:], 0, 0.0)
+        dvdsFxd_temp = np.insert(EndOfPrd_dvds[:, j], 0, EndOfPrd_dvds[0, j])
+        cFuncFxd_by_Share.append(LinearInterp(mNrmFxd_temp, cNrmFxd_temp))
+        dvdsFuncFxd_by_Share.append(LinearInterp(mNrmFxd_temp, dvdsFxd_temp))
+    cFuncFxd_now = LinearInterpOnInterp1D(cFuncFxd_by_Share, ShareGrid)
+    dvdsFuncFxd_now = LinearInterpOnInterp1D(dvdsFuncFxd_by_Share, ShareGrid)
+
+    # The share function when the agent can't adjust his portfolio is trivial
+    ShareFuncFxd_now = IdentityFunction(i_dim=1, n_dims=2)
+
+    # Construct the marginal value of mNrm function when the agent can't adjust his share
+    dvdmFuncFxd_now = MargValueFuncCRRA(cFuncFxd_now, CRRA)
+
+    # Construct the optimal risky share function when adjusting is possible
+    if BoroCnstNat_iszero:
+        Share_lower_bound = ShareLimit
+    else:
+        Share_lower_bound = 1.0
+    ShareAdj_now = np.insert(ShareAdj_now, 0, Share_lower_bound)
+    ShareFuncAdj_now = LinearInterp(mNrmAdj_now, ShareAdj_now, ShareLimit, 0.0)
+
+    # This is a point at which (a,c,share) have consistent length. Take the
+    # snapshot for storing the grid and values in the solution.
+    save_points = {
+        "a": deepcopy(aNrmGrid),
+        "eop_dvda_adj": uFunc.der(cNrmAdj_now),
+        "share_adj": deepcopy(ShareAdj_now),
+        "share_grid": deepcopy(ShareGrid),
+        "eop_dvda_fxd": uFunc.der(EndOfPrd_dvda),
+        "eop_dvds_fxd": EndOfPrd_dvds,
+    }
+
+    # Add the value function if requested TODO
+    if vFuncBool:
+        vFuncAdj_now = NullFunc()
+        vFuncFxd_now = NullFunc()
+    else:  # If vFuncBool is False, fill in dummy values
+        vFuncAdj_now = NullFunc()
+        vFuncFxd_now = NullFunc()
+
+    # Package and return the solution
+    solution_now = PortfolioSolution(
+                        cFuncAdj=cFuncAdj_now,
+                        ShareFuncAdj=ShareFuncAdj_now,
+                        vPfuncAdj=vPfuncAdj_now,
+                        vFuncAdj=vFuncAdj_now,
+                        cFuncFxd=cFuncFxd_now,
+                        ShareFuncFxd=ShareFuncFxd_now,
+                        dvdmFuncFxd=dvdmFuncFxd_now,
+                        dvdsFuncFxd=dvdsFuncFxd_now,
+                        vFuncFxd=vFuncFxd_now,
+                        AdjPrb=AdjustPrb,
+                        # WHAT IS THIS STUFF FOR??
+                        aGrid=save_points["a"],
+                        Share_adj=save_points["share_adj"],
+                        EndOfPrddvda_adj=save_points["eop_dvda_adj"],
+                        ShareGrid=save_points["share_grid"],
+                        EndOfPrddvda_fxd=save_points["eop_dvda_fxd"],
+                        EndOfPrddvds_fxd=save_points["eop_dvds_fxd"],
+
+                        )
+    return solution_now
 
 
 class ConsPortfolioSolver(MetricObject):