From c2d8a7a35c4bdf4a04830ea6f7ea4ab1464b57d4 Mon Sep 17 00:00:00 2001
From: Nikhil Woodruff <35577657+nikhilwoodruff@users.noreply.github.com>
Date: Thu, 24 Oct 2024 12:03:57 +0100
Subject: [PATCH] Fix subsampling bug (#297)

* Fix Subsample doesn't work with behavioural response simulations #296

* Versioning

* Add safety check
---
 changelog_entry.yaml                        | 4 ++++
 policyengine_core/simulations/simulation.py | 8 ++++++++
 2 files changed, 12 insertions(+)

diff --git a/changelog_entry.yaml b/changelog_entry.yaml
index e69de29b..1fa7f744 100644
--- a/changelog_entry.yaml
+++ b/changelog_entry.yaml
@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    fixed:
+    - Bug causing subsample to not work in some situations.
diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py
index 482c90e4..a18312c2 100644
--- a/policyengine_core/simulations/simulation.py
+++ b/policyengine_core/simulations/simulation.py
@@ -1549,6 +1549,14 @@ def subsample(
         # Update the dataset and rebuild the simulation
         self.dataset = Dataset.from_dataframe(df, self.dataset.time_period)
         self.build_from_dataset()
+
+        # Ensure the baseline branch has the new data.
+        if "baseline" in self.branches:
+            baseline_tax_benefit_system = self.branches[
+                "baseline"
+            ].tax_benefit_system
+            self.branches["baseline"] = self.clone()
+            self.branches["tax_benefit_system"] = baseline_tax_benefit_system
         return self