From 909bd161e2d370806fc80aa3617ddbf7e6252977 Mon Sep 17 00:00:00 2001
From: Mateo VG <mateovegi@gmail.com>
Date: Tue, 4 Apr 2023 09:58:53 -0400
Subject: [PATCH] Implement feature and test

---
 HARK/ConsumptionSaving/ConsRiskyAssetModel.py | 19 ++++++++---
 .../tests/test_ConsPortfolioModel.py          | 33 +++++++++++++++++++
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/HARK/ConsumptionSaving/ConsRiskyAssetModel.py b/HARK/ConsumptionSaving/ConsRiskyAssetModel.py
index cf42faf12..6bcf1f0d2 100644
--- a/HARK/ConsumptionSaving/ConsRiskyAssetModel.py
+++ b/HARK/ConsumptionSaving/ConsRiskyAssetModel.py
@@ -61,6 +61,11 @@ def __init__(self, verbose=False, quiet=False, **kwds):
         if not hasattr(self, "PortfolioBisect"):
             self.PortfolioBisect = False
 
+        # Boolean determines whether, when simulating a given time period,
+        # all agents will draw the same risky return factor (true by default)
+        if not hasattr(self, "sim_common_Rriksy"):
+            self.sim_common_Rriksy = True
+
         # Initialize a basic consumer type
         IndShockConsumerType.__init__(self, verbose=verbose, quiet=quiet, **kwds)
 
@@ -319,9 +324,15 @@ def get_Risky(self):
             RiskyAvg = self.RiskyAvg
             RiskyStd = self.RiskyStd
 
-        self.shocks["Risky"] = Lognormal.from_mean_std(
-            RiskyAvg, RiskyStd, seed=self.RNG.integers(0, 2**31 - 1)
-        ).draw(1)
+        # Draw either a common economy-wide return, or one for each agent
+        if self.sim_common_Rriksy:
+            self.shocks["Risky"] = Lognormal.from_mean_std(
+                RiskyAvg, RiskyStd, seed=self.RNG.integers(0, 2**31 - 1)
+            ).draw(1)
+        else:
+            self.shocks["Risky"] = Lognormal.from_mean_std(
+                RiskyAvg, RiskyStd, seed=self.RNG.integers(0, 2**31 - 1)
+            ).draw(self.AgentCount)
 
     def get_Adjust(self):
         """
@@ -1307,7 +1318,7 @@ def v_next(shocks, a_nrm):
     "AdjustPrb": 1.0,
     # When simulating the model, should all agents get the same risky return in
     # a given period?
-    "sim_common_Rriksy": True, 
+    "sim_common_Rriksy": True,
 }
 
 # Make a dictionary to specify a risky asset consumer type
diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py
index c66fd87d5..331d09bc9 100644
--- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py
+++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py
@@ -222,3 +222,36 @@ def test_discrete_and_joint(self):
 
         # Solve model under given parameters
         self.discrete_and_joint.solve()
+
+
+class testRiskyReturnDim(PortfolioConsumerTypeTestCase):
+    def test_simulation(self):
+        # Setup
+        self.pcct.T_sim = 30
+        self.pcct.AgentCount = 10
+        self.pcct.track_vars += [
+            "mNrm",
+            "cNrm",
+            "Risky",
+        ]
+        # Common (default) simulation
+        self.pcct.initialize_sim()
+        self.pcct.simulate()
+        # Assety that all columns of Risky are the same
+        self.assertTrue(
+            np.all(
+                self.pcct.history["Risky"]
+                == self.pcct.history["Risky"][:, 0][:, np.newaxis]
+            )
+        )
+        # Agent specific simulation
+        self.pcct.sim_common_Rriksy = False
+        self.pcct.initialize_sim()
+        self.pcct.simulate()
+        # Assety that all columns of Risky are not the same
+        self.assertFalse(
+            np.all(
+                self.pcct.history["Risky"]
+                == self.pcct.history["Risky"][:, 0][:, np.newaxis]
+            )
+        )