Merge pull request #201 from alanlujan91/i195

[WIP] create explode_agents method

sharkfin/population.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from dataclasses import dataclass
 from functools import partial
 from pprint import pprint
@@ -180,7 +181,9 @@ def agent_data(self):
         data_calls = {
             "aLvl": lambda a: a.state_now["aLvl"],
             "mNrm": lambda a: a.state_now["mNrm"],
-            "cNrm": lambda a: a.controls["cNrm"] if "cNrm" in a.controls else np.full(a.AgentCount, np.nan),
+            "cNrm": lambda a: a.controls["cNrm"]
+            if "cNrm" in a.controls
+            else np.full(a.AgentCount, np.nan),
             "mNrm_ratio_StE": lambda a: a.state_now["mNrm"] / a.mNrmStE,
         }
 
@@ -205,18 +208,13 @@ def class_stats(self, store=False):
         if self.ex_ante_hetero_params is None or len(self.ex_ante_hetero_params) == 0:
             cs = agent_data.aggregate(["mean", "std"])
 
-            mean_data = cs.loc['mean'].to_dict()
-            std_data = cs.loc['std'].to_dict()
+            mean_data = cs.loc["mean"].to_dict()
+            std_data = cs.loc["std"].to_dict()
 
             # this collapse the data into one row with appropriate column names
-            all_data = {k + '_mean' : [mean_data[k]] for k in mean_data}
-            all_data.update({
-                k + '_std' : [std_data[k]]
-                for k
-                in std_data
-            })
-            all_data['label'] = ["all"]
-
+            all_data = {k + "_mean": [mean_data[k]] for k in mean_data}
+            all_data.update({k + "_std": [std_data[k]] for k in std_data})
+            all_data["label"] = ["all"]
 
             cs = pd.DataFrame.from_dict(all_data)
 
@@ -261,6 +259,21 @@ def solve_distributed_agents(self):
         for agent in self.agents:
             agent.solve()
 
+    def explode_agents(self, num):
+
+        exploded_agents = []
+        exploded_dicts = []
+
+        for i, agent in enumerate(self.agents):
+            for j in range(num):
+                exploded_agents.append(deepcopy(agent))
+                exploded_dicts.append(deepcopy(self.agent_dicts[i]))
+
+        self.agents = exploded_agents
+        self.agent_dicts = exploded_dicts
+
+        self.create_database()
+
     def unpack_solutions(self):
         self.solution = [agent.solution for agent in self.agents]
 

sharkfin/tests/test_simulation.py

@@ -1,18 +1,16 @@
+import numpy as np
 from HARK.ConsumptionSaving.ConsPortfolioModel import SequentialPortfolioConsumerType
+from pytest import approx
+
 from sharkfin.broker import *
 from sharkfin.expectations import *
 from sharkfin.population import *
 from sharkfin.simulation import *
-from simulate.parameters import (
-    WHITESHARK,
-    build_population
-)
-import numpy as np
-
-from pytest import approx
+from simulate.parameters import LUCAS0, WHITESHARK, build_population
 
 ## MARKET SIMULATIONS
 
+
 def test_market_simulation():
     """
     Sets up and runs a MarketSimulation with no population.
@@ -51,6 +49,7 @@ def test_calibration_simulation():
 
     assert len(data["prices"]) == 2
 
+
 def test_series_simulation():
     """
     Sets up and runs an agent population simulation
@@ -63,7 +62,19 @@ def test_series_simulation():
     market = None
 
     sim = SeriesSimulation(q=q, r=r, market=market)
-    sim.simulate(burn_in=2, series=[(10000, 0), (10000, 0), (10000, 0), (10000, 0), (0,10000), (0, 10000), (0, 10000), (0, 10000)])
+    sim.simulate(
+        burn_in=2,
+        series=[
+            (10000, 0),
+            (10000, 0),
+            (10000, 0),
+            (10000, 0),
+            (0, 10000),
+            (0, 10000),
+            (0, 10000),
+            (0, 10000),
+        ],
+    )
 
     assert sim.broker.buy_sell_history[2] == (10000, 0)
     # assert(len(sim.history['buy_sell']) == 3) # need the padded day
@@ -74,6 +85,7 @@ def test_series_simulation():
 
 ## MACRO SIMULATIONS
 
+
 def test_macro_simulation():
     """
     Sets up and runs an simulation with an agent population.
@@ -82,12 +94,12 @@ def test_macro_simulation():
     pop = build_population(
         SequentialPortfolioConsumerType,
         WHITESHARK,
-        rng = np.random.default_rng(1)
-        )
+        rng=np.random.default_rng(1),
+    )
 
     # arguments to attention simulation
 
-    #a = 0.2
+    # a = 0.2
     q = 1
     r = 30
     market = None
@@ -97,7 +109,7 @@ def test_macro_simulation():
     attsim = MacroSimulation(
         pop,
         FinanceModel,
-        #a=a,
+        # a=a,
         q=q,
         r=r,
         market=market,
@@ -119,6 +131,7 @@ def test_macro_simulation():
 
     assert len(data["prices"]) == 30
 
+
 def test_attention_simulation():
     """
     Sets up and runs an agent population simulation
@@ -128,9 +141,8 @@ def test_attention_simulation():
     pop = build_population(
         SequentialPortfolioConsumerType,
         WHITESHARK,
-        rng = np.random.default_rng(1)
-        )
-
+        rng=np.random.default_rng(1),
+    )
 
     # arguments to attention simulation
 
@@ -149,9 +161,7 @@ def test_attention_simulation():
         r=r,
         market=market,
         days_per_quarter=days_per_quarter,
-        fm_args = {
-            'p1' : 0.5
-        }
+        fm_args={"p1": 0.5},
     )
     attsim.simulate(burn_in=20)
 
@@ -167,7 +177,7 @@ def test_attention_simulation():
     assert attsim.days_per_quarter == days_per_quarter
     assert attsim.fm.days_per_quarter == days_per_quarter
 
-    assert sim_stats['end_day'] == 30
+    assert sim_stats["end_day"] == 30
 
     data = attsim.daily_data()
 
@@ -184,3 +194,49 @@ def test_attention_simulation():
     assert ror_mean_1 == approx(ror_mean_2)
 
 
+def test_lucas0_simulation():
+    """
+    Sets up and runs an simulation with an agent population.
+    """
+    # initialize population
+    pop = build_population(
+        SequentialPortfolioConsumerType,
+        LUCAS0,
+        rng=np.random.default_rng(1),
+    )
+
+    assert len(pop.agent_database.index) == LUCAS0["num_per_type"]
+
+    # arguments to attention simulation
+
+    # a = 0.2
+    q = 1
+    r = 30
+    market = None
+
+    days_per_quarter = 30
+
+    attsim = MacroSimulation(
+        pop,
+        FinanceModel,
+        # a=a,
+        q=q,
+        r=r,
+        market=market,
+        days_per_quarter=days_per_quarter,
+    )
+    attsim.simulate(burn_in=20)
+
+    ## testing for existence of this class stat
+    attsim.pop.class_stats()["mNrm_ratio_StE_mean"]
+
+    attsim.daily_data()["sell_macro"]
+
+    attsim.sim_stats()
+
+    assert attsim.days_per_quarter == days_per_quarter
+    assert attsim.fm.days_per_quarter == days_per_quarter
+
+    data = attsim.daily_data()
+
+    assert len(data["prices"]) == 30

simulate/parameters.py

@@ -8,6 +8,9 @@
 
 
 def build_population(agent_type, parameters, rng=None, dphm=1500):
+
+    num_per_type = parameters.get("num_per_type", 1)
+
     pop = AgentPopulation(
         agent_type(), parameters, rng=rng, dollars_per_hark_money_unit=dphm
     )
@@ -21,6 +24,8 @@ def build_population(agent_type, parameters, rng=None, dphm=1500):
 
     pop.solve(merge_by=parameters["ex_post"])
 
+    pop.explode_agents(num_per_type)
+
     # initialize population model
     pop.init_simulation()
 
@@ -99,6 +104,6 @@ def build_population(agent_type, parameters, rng=None, dphm=1500):
 
 lucas0_parameter_dict = lucas0_agent_population_params
 lucas0_parameter_dict["AgentCount"] = 10  # TODO: What should this be?
-
+lucas0_parameter_dict["num_per_type"] = 10
 
 LUCAS0 = lucas0_parameter_dict