Equations in Models

Documentation/CHANGELOG.md

@@ -14,6 +14,7 @@ Release Date: TBD
 
 ### Major Changes
 
+- Moves `HARK.core.Model` to new `HARK.model` submodule, and allow `Model` object to be configured with equations, which are dictionaries of callables that take model variables as arguments. [#1292](https://github.com/econ-ark/HARK/pull/1292)
 - Adds `HARK.core.AgentPopulation` class to represent a population of agents with ex-ante heterogeneous parametrizations as distributions. [#1237](https://github.com/econ-ark/HARK/pull/1237)
 
 ### Minor Changes

Documentation/reference/index.rst

@@ -13,6 +13,7 @@ API Reference
    tools/frame
    tools/helpers
    tools/interpolation
+   tools/model
    tools/numba_tools
    tools/parallel
    tools/rewards

Documentation/reference/tools/model.rst

@@ -0,0 +1,7 @@
+Model
+----------
+
+.. automodule:: HARK.model
+   :members:
+   :undoc-members:
+   :show-inheritance:

HARK/ConsumptionSaving/ConsAggShockModel.py

@@ -140,12 +140,7 @@ def __init__(self, **kwds):
         params = init_agg_shocks.copy()
         params.update(kwds)
 
-        AgentType.__init__(
-            self,
-            solution_terminal=deepcopy(IndShockConsumerType.solution_terminal_),
-            pseudo_terminal=False,
-            **params
-        )
+        super().__init__( **params)
 
         # Add consumer-type specific objects, copying to create independent versions
         self.time_vary = deepcopy(IndShockConsumerType.time_vary_)

HARK/ConsumptionSaving/ConsIndShockModel.py

@@ -52,6 +52,7 @@
     MargValueFuncCRRA,
     ValueFuncCRRA,
 )
+from HARK.model import Control
 from HARK.metric import MetricObject
 from HARK.rewards import (
     CRRAutility,
@@ -1563,6 +1564,14 @@ def prepare_to_calc_EndOfPrdvP(self):
     # Do Perfect Foresight MIT Shock: Forces Newborns to follow solution path of the agent he/she replaced when True
 }
 
+PerfForesightConsumerType_dynamics = {
+    # need dynamic equation for Rnrm here. It's going to get overwritten in downstream models
+    'bNrm' : lambda Rnrm, aNrm : Rnrm * aNrm,
+    'mNrm' : lambda bNrm, theta : bNrm + theta,
+    'cNrm' : Control(['mNrm']),
+    'aNrm' : lambda mNrm, cNrm : mNrm - cNrm
+}
+
 
 class PerfForesightConsumerType(AgentType):
     """
@@ -1616,6 +1625,7 @@ def __init__(self, verbose=1, quiet=False, **kwds):
         set_verbosity_level((4 - verbose) * 10)
 
         self.update_Rfree()  # update interest rate if time varying
+        self.equations.update(PerfForesightConsumerType_dynamics)
 
     def pre_solve(self):
         self.update_solution_terminal()  # Solve the terminal period problem
@@ -1855,9 +1865,9 @@ def transition(self):
         PlvlAggNow = self.state_prev["PlvlAgg"] * self.PermShkAggNow
         # "Effective" interest factor on normalized assets
         ReffNow = RfreeNow / self.shocks["PermShk"]
-        bNrmNow = ReffNow * aNrmPrev  # Bank balances before labor income
+        bNrmNow = self.equations['bNrm'](ReffNow, aNrmPrev)  # Bank balances before labor income
         # Market resources after income
-        mNrmNow = bNrmNow + self.shocks["TranShk"]
+        mNrmNow = self.equations['mNrm'](bNrmNow, self.shocks["TranShk"])
 
         return pLvlNow, PlvlAggNow, bNrmNow, mNrmNow, None
 
@@ -1899,7 +1909,7 @@ def get_poststates(self):
         None
         """
         # should this be "Now", or "Prev"?!?
-        self.state_now["aNrm"] = self.state_now["mNrm"] - self.controls["cNrm"]
+        self.state_now["aNrm"] = self.equations['aNrm'](self.state_now['mNrm'], self.controls['cNrm'])
         # Useful in some cases to precalculate asset level
         self.state_now["aLvl"] = self.state_now["aNrm"] * self.state_now["pLvl"]
 
@@ -2108,6 +2118,15 @@ def check_conditions(self, verbose=None):
     }
 )
 
+IndShockConsumerType_dynamics = {
+    **PerfForesightConsumerType_dynamics,
+    'G' : lambda gamma, psi : gamma * psi,
+    'Rnrm' : lambda R, G : R / G,
+    'bNrm' : lambda Rnrm, aNrm : Rnrm * aNrm,
+    'mNrm' : lambda bNrm, theta : bNrm + theta,
+    'cNrm' : Control(['mNrm']),
+    'aNrm' : lambda mNrm, cNrm : mNrm - cNrm
+}
 
 class IndShockConsumerType(PerfForesightConsumerType):
     """
@@ -2149,6 +2168,8 @@ def __init__(self, verbose=1, quiet=False, **kwds):
 
         self.update()  # Make assets grid, income process, terminal solution
 
+        self.equations.update(IndShockConsumerType_dynamics)
+
     def update_income_process(self):
         """
         Updates this agent's income process based on his own attributes.
@@ -3230,7 +3251,7 @@ def check_conditions(self, verbose=None):
             self.PermGroFac[0] * self.InvEx_PermShkInv
         )  # [url]/#PGroAdj
 
-        self.thorn = (self.Rfree * self.DiscFac) ** (1 / self.CRRA)
+        self.thorn =  (self.parameters['Rfree'] * self.parameters['DiscFac']) ** (1 / self.parameters['CRRA'])
 
         # self.Ex_RNrm           = self.Rfree*Ex_PermShkInv/(self.PermGroFac[0]*self.LivPrb[0])
         self.GPFRaw = self.thorn / (self.PermGroFac[0])  # [url]/#GPF
@@ -3621,7 +3642,7 @@ def __init__(self, **kwds):
         params.update(kwds)
 
         # Initialize a basic AgentType
-        PerfForesightConsumerType.__init__(self, **params)
+        super().__init__(**params)
 
         # Add consumer-type specific objects, copying to create independent versions
         self.solve_one_period = make_one_period_oo_solver(ConsKinkedRsolver)

HARK/core.py

@@ -23,77 +23,10 @@
     TimeVaryingDiscreteDistribution,
     combine_indep_dstns,
 )
+from HARK.model import Model
 from HARK.parallel import multi_thread_commands, multi_thread_commands_fake
 from HARK.utilities import NullFunc, get_arg_names
 
-
-class Model:
-    """
-    A class with special handling of parameters assignment.
-    """
-
-    def assign_parameters(self, **kwds):
-        """
-        Assign an arbitrary number of attributes to this agent.
-
-        Parameters
-        ----------
-        **kwds : keyword arguments
-            Any number of keyword arguments of the form key=value.  Each value
-            will be assigned to the attribute named in self.
-
-        Returns
-        -------
-        none
-        """
-        self.parameters.update(kwds)
-        for key in kwds:
-            setattr(self, key, kwds[key])
-
-    def get_parameter(self, name):
-        """
-        Returns a parameter of this model
-
-        Parameters
-        ----------
-        name : string
-            The name of the parameter to get
-
-        Returns
-        -------
-        value :
-            The value of the parameter
-        """
-        return self.parameters[name]
-
-    def __eq__(self, other):
-        if isinstance(other, type(self)):
-            return self.parameters == other.parameters
-
-        return NotImplemented
-
-    def __init__(self):
-        if not hasattr(self, "parameters"):
-            self.parameters = {}
-
-    def __str__(self):
-        type_ = type(self)
-        module = type_.__module__
-        qualname = type_.__qualname__
-
-        s = f"<{module}.{qualname} object at {hex(id(self))}.\n"
-        s += "Parameters:"
-
-        for p in self.parameters:
-            s += f"\n{p}: {self.parameters[p]}"
-
-        s += ">"
-        return s
-
-    def __repr__(self):
-        return self.__str__()
-
-
 class AgentType(Model):
     """
     A superclass for economic agents in the HARK framework. Each model should

HARK/model.py

@@ -0,0 +1,110 @@
+"""
+Models in the abstract.
+"""
+from typing import Any, Callable, Mapping, Optional, Sequence, Union
+
+class Control():
+    """
+    A class used to indicate that a variable is a control variable.
+
+    Parameters
+    -----------
+    policy_args : Sequence[str]
+        A sequence of variable names, which refer to the values that are the arguments
+        to decision rules for this control variable.
+
+    """
+
+    def __init__(
+            self,
+            policy_args : Sequence[str]
+            ):
+        self.policy_args = policy_args
+
+class Model:
+    """
+    An economic model.
+    This object should contain the information about an environment's dynamics.
+
+    Parameters
+    ----------
+    equations : Mapping(str, Union[Callable, Control])
+        A mapping from model variable names (as strings) to transition equations governing
+        these variables.
+    parameters : Optional[Mapping(str, Any)]
+        A mapping from parameters names (strings) to parameter values.
+    options : Mapping(str, Any)
+        A mapping from options (str) to option values.
+    """
+
+    def assign_parameters(self, **kwds):
+        """
+        Assign an arbitrary number of attributes to this agent.
+
+        Parameters
+        ----------
+        **kwds : keyword arguments
+            Any number of keyword arguments of the form key=value.  Each value
+            will be assigned to the attribute named in self.
+
+        Returns
+        -------
+        none
+        """
+        self.parameters.update(kwds)
+        for key in kwds:
+            setattr(self, key, kwds[key])
+
+    def get_parameter(self, name: str):
+        """
+        Returns a parameter of this model
+
+        Parameters
+        ----------
+        name : string
+            The name of the parameter to get
+
+        Returns
+        -------
+        value :
+            The value of the parameter
+        """
+        return self.parameters[name]
+
+    def __eq__(self, other):
+        if isinstance(other, type(self)):
+            return (self.parameters == other.parameters) and (self.equations == other.equations)
+
+        return NotImplemented
+
+    def __init__(
+            self,
+            equations : Mapping[str, Union[Callable, Control]] = {},
+            parameters : Optional[Mapping[str, Any]] = None,
+            options : Mapping[str, Any] = {}
+            ):
+
+        self.equations = equations
+        self.options = options
+        if not hasattr(self, "parameters"):
+            self.parameters = {}
+            if parameters is not None:
+                self.assign_parameters(**parameters)
+
+
+    def __str__(self):
+        type_ = type(self)
+        module = type_.__module__
+        qualname = type_.__qualname__
+
+        s = f"<{module}.{qualname} object at {hex(id(self))}.\n"
+        s += "Parameters:"
+
+        for p in self.parameters:
+            s += f"\n{p}: {self.parameters[p]}"
+
+        s += ">"
+        return s
+
+    def __repr__(self):
+        return self.__str__()

HARK/tests/test_model.py

@@ -0,0 +1,47 @@
+from HARK.model import Control, Model
+import unittest
+
+
+class test_Model(unittest.TestCase):
+    def setUp(self):
+
+        equations_a = {
+            'm' : lambda a, r : a * r,
+            'c' : Control(['m']),
+            'a' : lambda m, c : m - c
+        }
+
+        parameters_a = {
+            'r' : 1.02
+        }
+
+        parameters_b = {
+            'r' : 1.03
+        }
+
+        equations_c = {
+            'm' : lambda a, r : a * r,
+            'c' : Control(['m']),
+            'a' : lambda m, c : m - 2 * c
+        }
+
+        # similar test to distance_metric
+        self.model_a = Model(
+            equations = equations_a,
+            parameters = parameters_a
+        )
+
+        self.model_b = Model(
+            equations = equations_a,
+            parameters = parameters_b
+        )
+
+        self.model_c = Model(
+            equations = equations_c,
+            parameters = parameters_a
+        )
+
+    def test_eq(self):
+        self.assertEqual(self.model_a, self.model_a)
+        self.assertNotEqual(self.model_a, self.model_b)
+        self.assertNotEqual(self.model_a, self.model_c)