Consumption Saving Refactor - move examples to notebooks - tests boilerplate

HARK/ConsumptionSaving/ConsAggShockModel.py

@@ -19,9 +19,12 @@
 from HARK.ConsumptionSaving.ConsIndShockModel import ConsumerSolution, IndShockConsumerType
 from HARK import HARKobject, Market, AgentType
 from copy import deepcopy
-import matplotlib.pyplot as plt
+# import matplotlib.pyplot as plt
 import HARK.ConsumptionSaving.ConsumerParameters as Params
 
+__all__ = ['MargValueFunc2D', 'AggShockConsumerType', 'AggShockMarkovConsumerType',
+'CobbDouglasEconomy', 'SmallOpenEconomy', 'CobbDouglasMarkovEconomy',
+'SmallOpenMarkovEconomy', 'CobbDouglasAggVars', 'AggregateSavingRule', 'AggShocksDynamicRule']
 
 utility = CRRAutility
 utilityP = CRRAutilityP
@@ -1790,198 +1793,3 @@ def __init__(self, AFunc):
         '''
         self.AFunc = AFunc
         self.distance_criteria = ['AFunc']
-
-
-###############################################################################
-
-def main():
-    from time import clock
-    from HARK.utilities import plotFuncs
-
-    def mystr(number): return "{:.4f}".format(number)
-
-    solve_agg_shocks_micro = False  # Solve an AggShockConsumerType's microeconomic problem
-    solve_agg_shocks_market = True  # Solve for the equilibrium aggregate saving rule in a CobbDouglasEconomy
-
-    solve_markov_micro = False  # Solve an AggShockMarkovConsumerType's microeconomic problem
-    solve_markov_market = True  # Solve for the equilibrium aggregate saving rule in a CobbDouglasMarkovEconomy
-    solve_krusell_smith = True  # Solve a simple Krusell-Smith-style two state, two shock model
-    solve_poly_state = False    # Solve a CobbDouglasEconomy with many states, potentially utilizing the "state jumper"
-
-    # EXAMPLE IMPLEMENTATIONS OF AggShockConsumerType ###
-
-    if solve_agg_shocks_micro or solve_agg_shocks_market:
-        # Make an aggregate shocks consumer type
-        AggShockExample = AggShockConsumerType()
-        AggShockExample.cycles = 0
-
-        # Make a Cobb-Douglas economy for the agents
-        EconomyExample = CobbDouglasEconomy(agents=[AggShockExample])
-        EconomyExample.makeAggShkHist()  # Simulate a history of aggregate shocks
-
-        # Have the consumers inherit relevant objects from the economy
-        AggShockExample.getEconomyData(EconomyExample)
-
-    if solve_agg_shocks_micro:
-        # Solve the microeconomic model for the aggregate shocks example type (and display results)
-        t_start = clock()
-        AggShockExample.solve()
-        t_end = clock()
-        print('Solving an aggregate shocks consumer took ' + mystr(t_end-t_start) + ' seconds.')
-        print('Consumption function at each aggregate market resources-to-labor ratio gridpoint:')
-        m_grid = np.linspace(0, 10, 200)
-        AggShockExample.unpackcFunc()
-        for M in AggShockExample.Mgrid.tolist():
-            mMin = AggShockExample.solution[0].mNrmMin(M)
-            c_at_this_M = AggShockExample.cFunc[0](m_grid+mMin, M*np.ones_like(m_grid))
-            plt.plot(m_grid+mMin, c_at_this_M)
-        plt.ylim(0., None)
-        plt.show()
-
-    if solve_agg_shocks_market:
-        # Solve the "macroeconomic" model by searching for a "fixed point dynamic rule"
-        t_start = clock()
-        print('Now solving for the equilibrium of a Cobb-Douglas economy.  This might take a few minutes...')
-        EconomyExample.solve()
-        t_end = clock()
-        print('Solving the "macroeconomic" aggregate shocks model took ' + str(t_end - t_start) + ' seconds.')
-
-        print('Aggregate savings as a function of aggregate market resources:')
-        plotFuncs(EconomyExample.AFunc, 0, 2*EconomyExample.kSS)
-        print('Consumption function at each aggregate market resources gridpoint (in general equilibrium):')
-        AggShockExample.unpackcFunc()
-        m_grid = np.linspace(0, 10, 200)
-        AggShockExample.unpackcFunc()
-        for M in AggShockExample.Mgrid.tolist():
-            mMin = AggShockExample.solution[0].mNrmMin(M)
-            c_at_this_M = AggShockExample.cFunc[0](m_grid+mMin, M*np.ones_like(m_grid))
-            plt.plot(m_grid+mMin, c_at_this_M)
-        plt.ylim(0., None)
-        plt.show()
-
-    # EXAMPLE IMPLEMENTATIONS OF AggShockMarkovConsumerType #
-
-    if solve_markov_micro or solve_markov_market or solve_krusell_smith:
-        # Make a Markov aggregate shocks consumer type
-        AggShockMrkvExample = AggShockMarkovConsumerType()
-        AggShockMrkvExample.IncomeDstn[0] = 2*[AggShockMrkvExample.IncomeDstn[0]]
-        AggShockMrkvExample.cycles = 0
-
-        # Make a Cobb-Douglas economy for the agents
-        MrkvEconomyExample = CobbDouglasMarkovEconomy(agents=[AggShockMrkvExample])
-        MrkvEconomyExample.DampingFac = 0.2  # Turn down damping
-        MrkvEconomyExample.makeAggShkHist()  # Simulate a history of aggregate shocks
-        AggShockMrkvExample.getEconomyData(
-            MrkvEconomyExample)  # Have the consumers inherit relevant objects from the economy
-
-    if solve_markov_micro:
-        # Solve the microeconomic model for the Markov aggregate shocks example type (and display results)
-        t_start = clock()
-        AggShockMrkvExample.solve()
-        t_end = clock()
-        print('Solving an aggregate shocks Markov consumer took ' + mystr(t_end-t_start) + ' seconds.')
-
-        print('Consumption function at each aggregate market \
-               resources-to-labor ratio gridpoint (for each macro state):')
-        m_grid = np.linspace(0, 10, 200)
-        AggShockMrkvExample.unpackcFunc()
-        for i in range(2):
-            for M in AggShockMrkvExample.Mgrid.tolist():
-                mMin = AggShockMrkvExample.solution[0].mNrmMin[i](M)
-                c_at_this_M = AggShockMrkvExample.cFunc[0][i](m_grid+mMin, M*np.ones_like(m_grid))
-                plt.plot(m_grid+mMin, c_at_this_M)
-            plt.ylim(0., None)
-            plt.show()
-
-    if solve_markov_market:
-        # Solve the "macroeconomic" model by searching for a "fixed point dynamic rule"
-        t_start = clock()
-        print('Now solving a two-state Markov economy.  This should take a few minutes...')
-        MrkvEconomyExample.solve()
-        t_end = clock()
-        print('Solving the "macroeconomic" aggregate shocks model took ' + str(t_end - t_start) + ' seconds.')
-
-        print('Consumption function at each aggregate market \
-               resources-to-labor ratio gridpoint (for each macro state):')
-        m_grid = np.linspace(0, 10, 200)
-        AggShockMrkvExample.unpackcFunc()
-        for i in range(2):
-            for M in AggShockMrkvExample.Mgrid.tolist():
-                mMin = AggShockMrkvExample.solution[0].mNrmMin[i](M)
-                c_at_this_M = AggShockMrkvExample.cFunc[0][i](m_grid+mMin, M*np.ones_like(m_grid))
-                plt.plot(m_grid+mMin, c_at_this_M)
-            plt.ylim(0., None)
-            plt.show()
-
-    if solve_krusell_smith:
-        # Make a Krusell-Smith agent type
-        # NOTE: These agents aren't exactly like KS, as they don't have serially correlated unemployment
-        KSexampleType = deepcopy(AggShockMrkvExample)
-        KSexampleType.IncomeDstn[0] = [[np.array([0.96, 0.04]), np.array([1.0, 1.0]), np.array([1.0/0.96, 0.0])],
-                                       [np.array([0.90, 0.10]), np.array([1.0, 1.0]), np.array([1.0/0.90, 0.0])]]
-
-        # Make a KS economy
-        KSeconomy = deepcopy(MrkvEconomyExample)
-        KSeconomy.agents = [KSexampleType]
-        KSeconomy.AggShkDstn = [[np.array([1.0]), np.array([1.0]), np.array([1.05])],
-                                [np.array([1.0]), np.array([1.0]), np.array([0.95])]]
-        KSeconomy.PermGroFacAgg = [1.0, 1.0]
-        KSexampleType.getEconomyData(KSeconomy)
-        KSeconomy.makeAggShkHist()
-
-        # Solve the K-S model
-        t_start = clock()
-        print('Now solving a Krusell-Smith-style economy.  This should take about a minute...')
-        KSeconomy.solve()
-        t_end = clock()
-        print('Solving the Krusell-Smith model took ' + str(t_end - t_start) + ' seconds.')
-
-    if solve_poly_state:
-        StateCount = 15     # Number of Markov states
-        GrowthAvg = 1.01    # Average permanent income growth factor
-        GrowthWidth = 0.02  # PermGroFacAgg deviates from PermGroFacAgg in this range
-        Persistence = 0.90  # Probability of staying in the same Markov state
-        PermGroFacAgg = np.linspace(GrowthAvg-GrowthWidth, GrowthAvg+GrowthWidth, num=StateCount)
-
-        # Make the Markov array with chosen states and persistence
-        PolyMrkvArray = np.zeros((StateCount, StateCount))
-        for i in range(StateCount):
-            for j in range(StateCount):
-                if i == j:
-                    PolyMrkvArray[i, j] = Persistence
-                elif (i == (j-1)) or (i == (j+1)):
-                    PolyMrkvArray[i, j] = 0.5*(1.0 - Persistence)
-        PolyMrkvArray[0, 0] += 0.5*(1.0 - Persistence)
-        PolyMrkvArray[StateCount-1, StateCount-1] += 0.5*(1.0 - Persistence)
-
-        # Make a consumer type to inhabit the economy
-        PolyStateExample = AggShockMarkovConsumerType()
-        PolyStateExample.MrkvArray = PolyMrkvArray
-        PolyStateExample.PermGroFacAgg = PermGroFacAgg
-        PolyStateExample.IncomeDstn[0] = StateCount*[PolyStateExample.IncomeDstn[0]]
-        PolyStateExample.cycles = 0
-
-        # Make a Cobb-Douglas economy for the agents
-        PolyStateEconomy = CobbDouglasMarkovEconomy(agents=[PolyStateExample])
-        PolyStateEconomy.MrkvArray = PolyMrkvArray
-        PolyStateEconomy.PermGroFacAgg = PermGroFacAgg
-        PolyStateEconomy.PermShkAggStd = StateCount*[0.006]
-        PolyStateEconomy.TranShkAggStd = StateCount*[0.003]
-        PolyStateEconomy.slope_prev = StateCount*[1.0]
-        PolyStateEconomy.intercept_prev = StateCount*[0.0]
-        PolyStateEconomy.update()
-        PolyStateEconomy.makeAggShkDstn()
-        PolyStateEconomy.makeAggShkHist()  # Simulate a history of aggregate shocks
-        PolyStateExample.getEconomyData(
-            PolyStateEconomy)  # Have the consumers inherit relevant objects from the economy
-
-        # Solve the many state model
-        t_start = clock()
-        print('Now solving an economy with ' + str(StateCount) + ' Markov states.  This might take a while...')
-        PolyStateEconomy.solve()
-        t_end = clock()
-        print('Solving a model with ' + str(StateCount) + ' states took ' + str(t_end - t_start) + ' seconds.')
-
-
-if __name__ == '__main__':
-    main()

HARK/ConsumptionSaving/ConsGenIncProcessModel.py

@@ -20,6 +20,10 @@
 from HARK.ConsumptionSaving.ConsIndShockModel import ConsIndShockSetup, ConsumerSolution, IndShockConsumerType
 import HARK.ConsumptionSaving.ConsumerParameters as Params
 
+__all__ = ['ValueFunc2D', 'MargValueFunc2D', 'MargMargValueFunc2D', 'pLvlFuncAR1',
+'ConsGenIncProcessSolver', 'GenIncProcessConsumerType',
+'IndShockExplicitPermIncConsumerType', 'PersistentShockConsumerType']
+
 utility = CRRAutility
 utilityP = CRRAutilityP
 utilityPP = CRRAutilityPP
@@ -1310,150 +1314,3 @@ def updatepLvlNextFunc(self):
         self.addToTimeVary('pLvlNextFunc')
         if not orig_time:
             self.timeRev()
-
-
-###############################################################################
-
-def main():
-    from HARK.utilities import plotFuncs
-    from time import clock
-    import matplotlib.pyplot as plt
-    def mystr(number): return "{:.4f}".format(number)
-
-    do_simulation = False
-
-    # Display information about the pLvlGrid used in these examples
-    print('The infinite horizon examples presented here use a grid of persistent income levels (pLvlGrid)')
-    print('based on percentiles of the long run distribution of pLvl for the given parameters. These percentiles')
-    print('are specified in the attribute pLvlPctiles. Here, the lowest percentile is ' +
-          str(Params.init_explicit_perm_inc['pLvlPctiles'][0]*100) + ' and the highest')
-    print('percentile is ' + str(Params.init_explicit_perm_inc['pLvlPctiles'][-1]*100) + '.\n')
-
-    # Make and solve an example "explicit permanent income" consumer with idiosyncratic shocks
-    ExplicitExample = IndShockExplicitPermIncConsumerType()
-    t_start = clock()
-    ExplicitExample.solve()
-    t_end = clock()
-    print('Solving an explicit permanent income consumer took ' + mystr(t_end-t_start) + ' seconds.')
-
-    # Plot the consumption function at various permanent income levels
-    print('Consumption function by pLvl for explicit permanent income consumer:')
-    pLvlGrid = ExplicitExample.pLvlGrid[0]
-    mLvlGrid = np.linspace(0, 20, 300)
-    for p in pLvlGrid:
-        M_temp = mLvlGrid + ExplicitExample.solution[0].mLvlMin(p)
-        C = ExplicitExample.solution[0].cFunc(M_temp, p*np.ones_like(M_temp))
-        plt.plot(M_temp, C)
-    plt.xlim(0., 20.)
-    plt.ylim(0., None)
-    plt.xlabel('Market resource level mLvl')
-    plt.ylabel('Consumption level cLvl')
-    plt.show()
-
-    # Now solve the *exact same* problem, but with the permanent income normalization
-    NormalizedExample = IndShockConsumerType(**Params.init_explicit_perm_inc)
-    t_start = clock()
-    NormalizedExample.solve()
-    t_end = clock()
-    print('Solving the equivalent problem with permanent income normalized out took ' +
-          mystr(t_end-t_start) + ' seconds.')
-
-    # Show that the normalized consumption function for the "explicit permanent income" consumer
-    # is almost identical for every permanent income level (and the same as the normalized problem's
-    # cFunc), but is less accurate due to extrapolation outside the bounds of pLvlGrid.
-    print('Normalized consumption function by pLvl for explicit permanent income consumer:')
-    pLvlGrid = ExplicitExample.pLvlGrid[0]
-    mNrmGrid = np.linspace(0, 20, 300)
-    for p in pLvlGrid:
-        M_temp = mNrmGrid*p + ExplicitExample.solution[0].mLvlMin(p)
-        C = ExplicitExample.solution[0].cFunc(M_temp, p*np.ones_like(M_temp))
-        plt.plot(M_temp/p, C/p)
-    plt.xlim(0., 20.)
-    plt.ylim(0., None)
-    plt.xlabel('Normalized market resources mNrm')
-    plt.ylabel('Normalized consumption cNrm')
-    plt.show()
-    print('Consumption function for normalized problem (without explicit permanent income):')
-    mNrmMin = NormalizedExample.solution[0].mNrmMin
-    plotFuncs(NormalizedExample.solution[0].cFunc, mNrmMin, mNrmMin+20)
-
-    print('The "explicit permanent income" solution deviates from the solution to the normalized problem because')
-    print('of errors from extrapolating beyond the bounds of the pLvlGrid. The error is largest for pLvl values')
-    print('near the upper and lower bounds, and propagates toward the center of the distribution.\n')
-
-    # Plot the value function at various permanent income levels
-    if ExplicitExample.vFuncBool:
-        pGrid = np.linspace(0.1, 3.0, 24)
-        M = np.linspace(0.001, 5, 300)
-        for p in pGrid:
-            M_temp = M+ExplicitExample.solution[0].mLvlMin(p)
-            C = ExplicitExample.solution[0].vFunc(M_temp, p*np.ones_like(M_temp))
-            plt.plot(M_temp, C)
-        plt.ylim([-200, 0])
-        plt.xlabel('Market resource level mLvl')
-        plt.ylabel('Value v')
-        plt.show()
-
-    # Simulate some data
-    if do_simulation:
-        ExplicitExample.T_sim = 500
-        ExplicitExample.track_vars = ['mLvlNow', 'cLvlNow', 'pLvlNow']
-        ExplicitExample.makeShockHistory()  # This is optional
-        ExplicitExample.initializeSim()
-        ExplicitExample.simulate()
-        plt.plot(np.mean(ExplicitExample.mLvlNow_hist, axis=1))
-        plt.xlabel('Simulated time period')
-        plt.ylabel('Average market resources mLvl')
-        plt.show()
-
-###############################################################################
-
-    # Make and solve an example "persistent idisyncratic shocks" consumer
-    PersistentExample = PersistentShockConsumerType(**Params.init_persistent_shocks)
-    t_start = clock()
-    PersistentExample.solve()
-    t_end = clock()
-    print('Solving a persistent income shocks consumer took ' + mystr(t_end-t_start) + ' seconds.')
-
-    # Plot the consumption function at various levels of persistent income pLvl
-    print('Consumption function by persistent income level pLvl for a consumer with AR1 coefficient of '
-          + str(PersistentExample.PrstIncCorr) + ':')
-    pLvlGrid = PersistentExample.pLvlGrid[0]
-    mLvlGrid = np.linspace(0, 20, 300)
-    for p in pLvlGrid:
-        M_temp = mLvlGrid + PersistentExample.solution[0].mLvlMin(p)
-        C = PersistentExample.solution[0].cFunc(M_temp, p*np.ones_like(M_temp))
-        plt.plot(M_temp, C)
-    plt.xlim(0., 20.)
-    plt.ylim(0., None)
-    plt.xlabel('Market resource level mLvl')
-    plt.ylabel('Consumption level cLvl')
-    plt.show()
-
-    # Plot the value function at various persistent income levels
-    if PersistentExample.vFuncBool:
-        pGrid = PersistentExample.pLvlGrid[0]
-        M = np.linspace(0.001, 5, 300)
-        for p in pGrid:
-            M_temp = M+PersistentExample.solution[0].mLvlMin(p)
-            C = PersistentExample.solution[0].vFunc(M_temp, p*np.ones_like(M_temp))
-            plt.plot(M_temp, C)
-        plt.ylim([-200, 0])
-        plt.xlabel('Market resource level mLvl')
-        plt.ylabel('Value v')
-        plt.show()
-
-    # Simulate some data
-    if do_simulation:
-        PersistentExample.T_sim = 500
-        PersistentExample.track_vars = ['mLvlNow', 'cLvlNow', 'pLvlNow']
-        PersistentExample.initializeSim()
-        PersistentExample.simulate()
-        plt.plot(np.mean(PersistentExample.mLvlNow_hist, axis=1))
-        plt.xlabel('Simulated time period')
-        plt.ylabel('Average market resources mLvl')
-        plt.show()
-
-
-if __name__ == '__main__':
-    main()