update tests, double agents

sbenthall · Mar 17, 2023 · 7e18a5d · 7e18a5d
1 parent 4ba8316
commit 7e18a5d
Showing 1 changed file with 28 additions and 18 deletions.
diff --git a/sharkfin/tests/test_simulation.py b/sharkfin/tests/test_simulation.py
@@ -3,16 +3,14 @@
 from sharkfin.expectations import *
 from sharkfin.population import *
 from sharkfin.simulation import *
-from simulate.parameters import (
-    WHITESHARK,
-    build_population
-)
+from simulate.parameters import WHITESHARK, build_population
 import numpy as np
 
 from pytest import approx
 
 ## 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,13 @@ def test_macro_simulation():
     pop = build_population(
         SequentialPortfolioConsumerType,
         WHITESHARK,
-        rng = np.random.default_rng(1)
-        )
+        rng=np.random.default_rng(1),
+        num_per_type=2,
+    )
 
     # arguments to attention simulation
 
-    #a = 0.2
+    # a = 0.2
     q = 1
     r = 30
     market = None
@@ -97,7 +110,7 @@ def test_macro_simulation():
     attsim = MacroSimulation(
         pop,
         FinanceModel,
-        #a=a,
+        # a=a,
         q=q,
         r=r,
         market=market,
@@ -119,6 +132,7 @@ def test_macro_simulation():
 
     assert len(data["prices"]) == 30
 
+
 def test_attention_simulation():
     """
     Sets up and runs an agent population simulation
@@ -128,9 +142,9 @@ def test_attention_simulation():
     pop = build_population(
         SequentialPortfolioConsumerType,
         WHITESHARK,
-        rng = np.random.default_rng(1)
-        )
-
+        rng=np.random.default_rng(1),
+        num_per_type=2,
+    )
 
     # arguments to attention simulation
 
@@ -149,9 +163,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 +179,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()
 
@@ -182,5 +194,3 @@ def test_attention_simulation():
     ror_mean_2 = attsim.ror_mean()
 
     assert ror_mean_1 == approx(ror_mean_2)
-
-