Reorganize code blocks for readability
@@ -1082,26 +1082,8 @@ def solve_RiskyContrib_Cns(
# STEP ONE
# Find end-of-period (continuation) value function and its derivatives.
- # The "inversion" machinery can deal with assets of 0 even if there is a
- # natural borrowing constraint, so include zeros.
- nNrmGrid = np.concatenate([np.array([0.0]), nNrmGrid])
- aNrmGrid = np.concatenate([np.array([0.0]), aXtraGrid])
-
- # Create tiled arrays with conforming dimensions. These are used
- # to compute expectations at every grid combinations
- # Convention will be (a,n,s)
- aNrm_tiled, nNrm_tiled, Share_tiled = np.meshgrid(
- aNrmGrid, nNrmGrid, ShareGrid, indexing="ij"
- )
-
- # Evaluate realizations of the derivatives and levels of next period's
- # value function
-
- # The agent who can adjust starts at the "contrib" stage, the one who can't
- # starts at the Fxd stage.
-
- # We are interested in marginal values before the realization of the
- # adjustment random variable. Compute those objects
+ # Start by constructing functions for next-period's pre-adjustment-shock
+ # expected value functions
if AdjustPrb < 1.0:
dvdm_next = lambda m, n, s: AdjustPrb * dvdmFunc_Reb_Adj_next(m, n) + (
@@ -1128,7 +1110,8 @@ def solve_RiskyContrib_Cns(
if vFuncBool:
v_next = lambda m, n, s: vFunc_Reb_Adj_next(m, n)
-
+ # Now construct a function that evaluates and discounts them given a
+ # vector of return and income shocks and an end-of-period state
def end_of_period_derivs(shocks,a,nTil,s):
"""
Computes the end-of-period derivatives (and optionally the value) of the
@@ -1182,21 +1165,22 @@ def solve_RiskyContrib_Cns(
else:
return np.stack([end_of_prd_dvda, end_of_prd_dvdn, end_of_prd_dvds])
- # Find end of period derivatives and value as discounted expectations of
- # next period's derivatives and value.
- # Create a function to recover the derivatives (and possible value) of the
- # end of period value function conditional on states and shocks.
- end_of_period_ds_func = lambda shocks, a, n, s: end_of_period_derivs(
- shocks,
- a,
- n,
- s
+ # Now find the expected values on a (a, nTil, s) grid
+
+ # The "inversion" machinery can deal with assets of 0 even if there is a
+ # natural borrowing constraint, so include zeros.
+ nNrmGrid = np.concatenate([np.array([0.0]), nNrmGrid])
+ aNrmGrid = np.concatenate([np.array([0.0]), aXtraGrid])
+
+ # Create tiled arrays with conforming dimensions.
+ aNrm_tiled, nNrm_tiled, Share_tiled = np.meshgrid(
+ aNrmGrid, nNrmGrid, ShareGrid, indexing="ij"
)
- # Then integrate over shocks
- # (eop = end of period)
+ # Find end of period derivatives and value as expectations of (discounted)
+ # next period's derivatives and value.
eop_derivs = calc_expectation(
- ShockDstn, end_of_period_ds_func, aNrm_tiled, nNrm_tiled, Share_tiled
+ ShockDstn, end_of_period_derivs, aNrm_tiled, nNrm_tiled, Share_tiled
)[:, :, :, :, 0]
# Unpack results
