econ-ark · sbenthall · Nov 6, 2020 · Nov 6, 2020 · Nov 7, 2020 · Nov 7, 2020
diff --git a/HARK/ConsumptionSaving/ConsPortfolioFrameModel.py b/HARK/ConsumptionSaving/ConsPortfolioFrameModel.py
@@ -0,0 +1,147 @@
+"""
+This file contains classes and functions for representing,
+solving, and simulating agents who must allocate their resources
+among consumption, saving in a risk-free asset (with a low return),
+and saving in a risky asset (with higher average return).
+
+This file also demonstrates a "frame" model architecture.
+"""
+import numpy as np
+from scipy.optimize import minimize_scalar
+from copy import deepcopy
+from HARK import HARKobject, NullFunc, Frame, FrameAgentType  # Basic HARK features
+from HARK.ConsumptionSaving.ConsIndShockModel import (
+    IndShockConsumerType,  # PortfolioConsumerType inherits from it
+    ValueFunc,  # For representing 1D value function
+    MargValueFunc,  # For representing 1D marginal value function
+    utility,  # CRRA utility function
+    utility_inv,  # Inverse CRRA utility function
+    utilityP,  # CRRA marginal utility function
+    utility_invP,  # Derivative of inverse CRRA utility function
+    utilityP_inv,  # Inverse CRRA marginal utility function
+    init_idiosyncratic_shocks,  # Baseline dictionary to build on
+)
+from HARK.ConsumptionSaving.ConsGenIncProcessModel import (
+    ValueFunc2D,  # For representing 2D value function
+    MargValueFunc2D,  # For representing 2D marginal value function
+)
+
+from HARK.ConsumptionSaving.ConsPortfolioModel import (
+    init_portfolio,
+    solveConsPortfolio,
+    PortfolioConsumerType,
+    PortfolioSolution
+)
+
+from HARK.distribution import combineIndepDstns
+from HARK.distribution import Lognormal, MeanOneLogNormal, Bernoulli  # Random draws for simulating agents
+from HARK.interpolation import (
+    LinearInterp,  # Piecewise linear interpolation
+    CubicInterp,  # Piecewise cubic interpolation
+    LinearInterpOnInterp1D,  # Interpolator over 1D interpolations
+    BilinearInterp,  # 2D interpolator
+    ConstantFunction,  # Interpolator-like class that returns constant value
+    IdentityFunction,  # Interpolator-like class that returns one of its arguments
+)
+
+class PortfolioConsumerFrameType(FrameAgentType, PortfolioConsumerType):
+    """
+    A consumer type with a portfolio choice, using Frame architecture.
+
+    A subclass of PortfolioConsumerType for now.
+    This is mainly to keep the _solver_ logic intact.
+    """
+
+    # values for aggregate variables
+    # to be set when simulation initializes.
+    # currently not doing anything because still using old
+    # initializeSim()
+    aggregate_init_values = {
+        'PlvlAggNow' : 1.0
+    }
+
+    def birth_aNrmNow(self, N):
+        return Lognormal(
+            mu=self.aNrmInitMean,
+            sigma=self.aNrmInitStd,
+            seed=self.RNG.randint(0, 2 ** 31 - 1),
+        ).draw(N)
+
+    def birth_pLvlNow(self, N):
+        pLvlInitMeanNow = self.pLvlInitMean + np.log(
+            self.state_now["PlvlAggNow"]
+        )  # Account for newer cohorts having higher permanent income
+
+        return Lognormal(
+            pLvlInitMeanNow,
+            self.pLvlInitStd,
+            seed=self.RNG.randint(0, 2 ** 31 - 1)
+        ).draw(N)
+
+
+    # values to assign to agents at birth
+    birth_values = {
+        'ShareNow' : 0,
+        'AdjustNow' : False,
+        'aNrmNow' : birth_aNrmNow,
+        'pLvlNow' : birth_pLvlNow
+    }
+
+    def transition_ShareNow(self, **context):
+        ShareNow = np.zeros(self.AgentCount) + np.nan
+
+        # Loop over each period of the cycle, getting controls separately depending on "age"
+        for t in range(self.T_cycle):
+            these = t == self.t_cycle
+
+            # Get controls for agents who *can* adjust their portfolio share
+            those = np.logical_and(these, self.shocks['AdjustNow'])
+
+            ShareNow[those] = self.solution[t].ShareFuncAdj(self.state_now['mNrmNow'][those])
+
+            # Get Controls for agents who *can't* adjust their portfolio share
+            those = np.logical_and(
+                these,
+                np.logical_not(self.shocks['AdjustNow']))
+            ShareNow[those] = self.solution[t].ShareFuncFxd(
+                context['mNrmNow'][those], ShareNow[those]
+            )
+
+        # redundant for now
+        self.controls["ShareNow"] = ShareNow
+
+        return ShareNow
+
+    def transition_cNrmNow(self, **context):
+        cNrmNow = np.zeros(self.AgentCount) + np.nan
+        ShareNow = self.controls["ShareNow"]
+
+        # Loop over each period of the cycle, getting controls separately depending on "age"
+        for t in range(self.T_cycle):
+            these = t == self.t_cycle
+
+            # Get controls for agents who *can* adjust their portfolio share
+            those = np.logical_and(these, context['AdjustNow'])
+            cNrmNow[those] = self.solution[t].cFuncAdj(context['mNrmNow'][those])
+
+            # Get Controls for agents who *can't* adjust their portfolio share
+            those = np.logical_and(
+                these,
+                np.logical_not(context['AdjustNow']))
+            cNrmNow[those] = self.solution[t].cFuncFxd(
+                context['mNrmNow'][those], ShareNow[those]
+            )
+
+        # Store controls as attributes of self
+        # redundant for now
+        self.controls["cNrmNow"] = cNrmNow
+
+        return cNrmNow
+
+    frames = {
+        ('RiskyNow','AdjustNow') : PortfolioConsumerType.getShocks,
+        ('pLvlNow', 'PlvlAggNow', 'bNrmNow', 'mNrmNow') : PortfolioConsumerType.getStates,
+        ('ShareNow') : transition_ShareNow,
+        ('cNrmNow') : transition_cNrmNow,
+        ('aNrmNow', 'aNrmNow') : PortfolioConsumerType.getPostStates
+    }
diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py
@@ -444,10 +444,7 @@ def simBirth(self, which_agents):
         None
         """
         IndShockConsumerType.simBirth(self, which_agents)
-        # Checking for control variable attribute here
-        # because we have not namespaced controls yet
-        if hasattr(self, 'ShareNow'):
-            self.ShareNow[which_agents] = 0
+        self.controls['ShareNow'][which_agents] = 0
         # here a shock is being used as a 'post state'
         self.shocks['AdjustNow'][which_agents] = False
 

diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioFrameModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioFrameModel.py
@@ -0,0 +1,120 @@
+import HARK.ConsumptionSaving.ConsPortfolioFrameModel as cpfm
+import numpy as np
+import unittest
+
+
+class PortfolioConsumerTypeTestCase(unittest.TestCase):
+    def setUp(self):
+        # Create portfolio choice consumer type
+        self.pcct = cpfm.PortfolioConsumerFrameType()
+        self.pcct.cycles = 0
+
+        # Solve the model under the given parameters
+
+        self.pcct.solve()
+
+class UnitsPortfolioConsumerTypeTestCase(PortfolioConsumerTypeTestCase):
+    def test_simOnePeriod(self):
+
+        self.pcct.T_sim = 30
+        self.pcct.AgentCount = 10
+        self.pcct.track_vars += ['aNrmNow']
+        self.pcct.initializeSim()
+
+        self.assertFalse(
+            np.any(self.pcct.shocks['AdjustNow'])
+        )
+
+        self.pcct.simOnePeriod()
+
+        self.assertAlmostEqual(
+            self.pcct.controls["ShareNow"][0],
+            0.8627164488246847
+        )
+        self.assertAlmostEqual(
+            self.pcct.controls["cNrmNow"][0],
+            1.67874799
+        )
+
+class SimulatePortfolioConsumerTypeTestCase(PortfolioConsumerTypeTestCase):
+
+    def test_simulation(self):
+
+        self.pcct.T_sim = 30
+        self.pcct.AgentCount = 10
+        self.pcct.track_vars += [
+            'mNrmNow',
+            'cNrmNow',
+            'ShareNow',
+            'aNrmNow',
+            'RiskyNow',
+            'RportNow',
+            'AdjustNow',
+            'PermShkNow',
+            'bNrmNow'
+        ]
+        self.pcct.initializeSim()
+
+        self.pcct.simulate()
+
+        self.assertAlmostEqual(
+            self.pcct.history['mNrmNow'][0][0], 9.70233892
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['cNrmNow'][0][0], 1.6787479894848298
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['ShareNow'][0][0], 0.8627164488246847
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['aNrmNow'][0][0], 8.023590930905383
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['AdjustNow'][0][0], 1.0
+        )
+
+
+        # the next period
+        self.assertAlmostEqual(
+            self.pcct.history['RiskyNow'][1][0], 0.8950304697526602
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['RportNow'][1][0], 0.9135595661654792
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['AdjustNow'][1][0], 1.0
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['PermShkNow'][1][0], 1.0050166461586711
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['bNrmNow'][1][0], 7.293439643953855
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['mNrmNow'][1][0], 8.287859049575047
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['cNrmNow'][1][0], 1.5773607434989751
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['ShareNow'][1][0], 0.9337608822146805
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['aNrmNow'][1][0], 6.710498306076072
+        )
+
+        self.assertAlmostEqual(
+            self.pcct.history['aNrmNow'][15][0], 5.304746367434934
+        )