Baseline simulation created before simulation modifier is applied

Author: janton10

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    fixed:
+      - Baseline simulation created before simulation modifier is applied.

policyengine_uk/simulation.py

@@ -80,9 +80,7 @@ def __init__(
         self.invalidated_caches = set()
         self.debug: bool = False
         self.trace: bool = trace
-        self.tracer: SimpleTracer = (
-            SimpleTracer() if not trace else FullTracer()
-        )
+        self.tracer: SimpleTracer = SimpleTracer() if not trace else FullTracer()
         self.opt_out_cache: bool = False
         self.max_spiral_loops: int = 10
         self.memory_config = None
@@ -125,10 +123,6 @@ def __init__(
         self.move_values("capital_gains", "capital_gains_before_response")
         self.move_values("employment_income", "employment_income_before_lsr")
 
-        if scenario is not None:
-            if scenario.simulation_modifier is not None:
-                scenario.simulation_modifier(self)
-
         if scenario is not None:
             self.baseline = Simulation(
                 scenario=None,
@@ -139,6 +133,10 @@ def __init__(
         else:
             self.baseline = self.clone()
 
+        if scenario is not None:
+            if scenario.simulation_modifier is not None:
+                scenario.simulation_modifier(self)
+
         self.calculated_periods = []
 
     def reset_calculations(self):
@@ -157,9 +155,7 @@ def apply_parameter_changes(self, changes: dict):
         self.tax_benefit_system.reset_parameters()
 
         for parameter in changes:
-            p: Parameter = self.tax_benefit_system.parameters.get_child(
-                parameter
-            )
+            p: Parameter = self.tax_benefit_system.parameters.get_child(parameter)
             if isinstance(changes[parameter], dict):
                 # Time-period specific changes
                 for time_period in changes[parameter]:
@@ -178,9 +174,7 @@ def build_from_situation(self, situation: Dict) -> None:
         Args:
             situation: Dictionary describing household composition and characteristics
         """
-        self.build_from_populations(
-            self.tax_benefit_system.instantiate_entities()
-        )
+        self.build_from_populations(self.tax_benefit_system.instantiate_entities())
         from policyengine_core.simulations.simulation_builder import (
             SimulationBuilder,
         )  # Import here to avoid circular dependency
@@ -204,9 +198,7 @@ def build_from_url(self, url: str) -> None:
             ValueError: If URL is not a HuggingFace URL
         """
         if "hf://" not in url:
-            raise ValueError(
-                f"Non-HuggingFace URLs are currently not supported."
-            )
+            raise ValueError(f"Non-HuggingFace URLs are currently not supported.")
 
         # Parse HuggingFace URL components
         owner, repo, filename = url.split("/")[-3:]
@@ -284,9 +276,7 @@ def build_from_dataset(self, dataset: Dataset) -> None:
         Args:
             dataset: PolicyEngine Dataset object containing simulation data
         """
-        data: Dict[str, Dict[str, Union[float, int, str]]] = (
-            dataset.load_dataset()
-        )
+        data: Dict[str, Dict[str, Union[float, int, str]]] = dataset.load_dataset()
 
         first_variable = data[list(data.keys())[0]]
         first_time_period = list(first_variable.keys())[0]
@@ -315,9 +305,7 @@ def get_first_array(variable_name: str) -> np.ndarray:
             for time_period in data[variable]:
                 if variable not in self.tax_benefit_system.variables:
                     continue
-                self.set_input(
-                    variable, time_period, data[variable][time_period]
-                )
+                self.set_input(variable, time_period, data[variable][time_period])
 
         # Now convert to the new UKSingleYearDataset
         self.input_variables = self.get_known_variables()
@@ -330,9 +318,7 @@ def get_first_array(variable_name: str) -> np.ndarray:
         self.build_from_multi_year_dataset(multi_year_dataset)
         self.dataset = multi_year_dataset
 
-    def build_from_single_year_dataset(
-        self, dataset: UKSingleYearDataset
-    ) -> None:
+    def build_from_single_year_dataset(self, dataset: UKSingleYearDataset) -> None:
         """Build simulation from a single-year UK dataset.
 
         Args:
@@ -343,9 +329,7 @@ def build_from_single_year_dataset(
         self.build_from_multi_year_dataset(dataset)
         self.dataset = dataset
 
-    def build_from_multi_year_dataset(
-        self, dataset: UKMultiYearDataset
-    ) -> None:
+    def build_from_multi_year_dataset(self, dataset: UKMultiYearDataset) -> None:
         """Build simulation from a multi-year UK dataset.
 
         Args:
@@ -506,9 +490,7 @@ def get_variable_dependencies(self, variable_name: str, depth: int = 1):
                 trace=True,
             )
             self._example_simulation.calculate("hbai_household_net_income")
-            for (
-                variable
-            ) in self._example_simulation.tax_benefit_system.variables:
+            for variable in self._example_simulation.tax_benefit_system.variables:
                 try:
                     dependencies = get_variable_dependencies(
                         variable, self._example_simulation
@@ -529,9 +511,7 @@ def get_variable_dependencies(self, variable_name: str, depth: int = 1):
         all_dependencies = set(dependencies)
         for dep in dependencies:
             try:
-                sub_dependencies = self.get_variable_dependencies(
-                    dep, depth - 1
-                )
+                sub_dependencies = self.get_variable_dependencies(dep, depth - 1)
                 all_dependencies.update(sub_dependencies)
             except ValueError:
                 continue
@@ -562,13 +542,9 @@ def compare(
         df_self = self.calculate_dataframe(variables, period=period)
         df_other = other.calculate_dataframe(variables, period=period)
 
-        df_combined = pd.concat(
-            [df_self, df_other], axis=1, keys=["self", "other"]
-        )
+        df_combined = pd.concat([df_self, df_other], axis=1, keys=["self", "other"])
         # Reset columns, then order by variable name first then self/other
-        df_combined.columns = [
-            f"{var}_{entity}" for entity, var in df_combined.columns
-        ]
+        df_combined.columns = [f"{var}_{entity}" for entity, var in df_combined.columns]
         # Add change where applicable (numeric only)
         for variable in variables:
             self_col = f"{variable}_self"