refactor ConsRepAgentModel in line with new DiscreteDistribution class
@@ -10,7 +10,7 @@ from builtins import str
from builtins import range
import numpy as np
from HARK.interpolation import LinearInterp
-from HARK.simulation import drawUniform, drawDiscrete
+from HARK.simulation import drawUniform
from HARK.ConsumptionSaving.ConsIndShockModel import IndShockConsumerType,\
ConsumerSolution,MargValueFunc, init_idiosyncratic_shocks
@@ -51,9 +51,9 @@ def solveConsRepAgent(solution_next,DiscFac,CRRA,IncomeDstn,CapShare,DeprFac,Per
'''
# Unpack next period's solution and the income distribution
vPfuncNext = solution_next.vPfunc
- ShkPrbsNext = IncomeDstn[0]
- PermShkValsNext = IncomeDstn[1]
- TranShkValsNext = IncomeDstn[2]
+ ShkPrbsNext = IncomeDstn.pmf
+ PermShkValsNext = IncomeDstn.X[0]
+ TranShkValsNext = IncomeDstn.X[1]
# Make tiled versions of end-of-period assets, shocks, and probabilities
aNrmNow = aXtraGrid
@@ -141,9 +141,9 @@ def solveConsRepAgentMarkov(solution_next,MrkvArray,DiscFac,CRRA,IncomeDstn,CapS
for j in range(StateCount):
# Define next-period-state conditional objects
vPfuncNext = solution_next.vPfunc[j]
- ShkPrbsNext = IncomeDstn[j][0]
- PermShkValsNext = IncomeDstn[j][1]
- TranShkValsNext = IncomeDstn[j][2]
+ ShkPrbsNext = IncomeDstn[j].pmf
+ PermShkValsNext = IncomeDstn[j].X[0]
+ TranShkValsNext = IncomeDstn[j].X[1]
# Make tiled versions of end-of-period assets, shocks, and probabilities
ShkCount = ShkPrbsNext.size
@@ -315,14 +315,11 @@ class RepAgentMarkovConsumerType(RepAgentConsumerType):
i = self.MrkvNow[0]
IncomeDstnNow = self.IncomeDstn[t-1][i] # set current income distribution
PermGroFacNow = self.PermGroFac[t-1][i] # and permanent growth factor
- Indices = np.arange(IncomeDstnNow[0].size) # just a list of integers
# Get random draws of income shocks from the discrete distribution
- EventDraw = drawDiscrete(N=1,
- X=Indices,
- P=IncomeDstnNow[0],
+ EventDraw = IncomeDstnNow.draw_events(1,
seed=self.RNG.randint(0,2**31-1))
- PermShkNow = IncomeDstnNow[1][EventDraw]*PermGroFacNow # permanent "shock" includes expected growth
- TranShkNow = IncomeDstnNow[2][EventDraw]
+ PermShkNow = IncomeDstnNow.X[0][EventDraw]*PermGroFacNow # permanent "shock" includes expected growth
+ TranShkNow = IncomeDstnNow.X[1][EventDraw]
self.PermShkNow = np.array(PermShkNow)
self.TranShkNow = np.array(TranShkNow)
