adding distributeParams

HARK/ConsumptionSaving/tests/test_ConsAggShockModel.py

@@ -1,24 +1,37 @@
+from HARK import distributeParams
 from HARK.ConsumptionSaving.ConsAggShockModel import AggShockConsumerType, CobbDouglasEconomy, AggShockMarkovConsumerType, CobbDouglasMarkovEconomy
+from HARK.distribution import Uniform
 import numpy as np
 import unittest
 
 class testAggShockConsumerType(unittest.TestCase):
 
     def setUp(self):
-        self.agent = AggShockConsumerType()
-        self.agent.cycles = 0
+        agent = AggShockConsumerType()
+        agent.AgentCount = 60
+        agent.cycles = 0
+
+        self.agents = distributeParams(agent,
+                                       'DiscFac',
+                                       3,
+                                       Uniform(bot=.96,
+                                              top=.98)
+                                       )
 
         self.economy = EconomyExample = CobbDouglasEconomy(
-            agents=[self.agent])
+            agents=self.agents)
+
+    def test_distributeParams(self):
+        self.assertEqual(self.agents[1].AgentCount, 20)
 
     def test_economy(self):
         # Make a Cobb-Douglas economy for the agents
         self.economy.makeAggShkHist()  # Simulate a history of aggregate shocks
 
         # Have the consumers inherit relevant objects from the economy
-        self.agent.getEconomyData(self.economy)
+        self.agents[0].getEconomyData(self.economy)
 
-        self.agent.solve()
+        self.agents[0].solve()
 
 class testAggShockMarkovConsumerType(unittest.TestCase):
 

HARK/core.py

@@ -1175,3 +1175,35 @@ def updateDynamics(self):
             for this_type in self.agents:
                 setattr(this_type, var_name, this_obj)
         return dynamics
+
+def distributeParams(agent, param_name,param_count,distribution):
+    '''
+    Distributes heterogeneous values of one parameter to the AgentTypes in self.agents.
+    Parameters
+    ----------
+    agent: AgentType
+        An agent to clone.
+    param_name : string
+        Name of the parameter to be assigned.
+    param_count : int
+        Number of different values the parameter will take on.
+    distribution : Distribution
+        A distribution.
+
+    Returns
+    -------
+    agent_set : [AgentType]
+        A list of param_count agents, ex ante heterogeneous with 
+        respect to param_name. The AgentCount of the original
+        will be split between the agents of the returned
+        list in proportion to the given distribution.
+    '''
+    param_dist = distribution.approx(N=param_count)
+
+    agent_set = [deepcopy(agent) for i in range(param_count)]
+
+    for j in range(param_count):
+        agent_set[j].AgentCount = int(agent.AgentCount*param_dist.pmf[j])
+        agent_set[j].__dict__[param_name] = param_dist.X[j]
+
+    return agent_set