Reduce code duplication

Author: nikhilwoodruff

policyengine/simulation.py

@@ -289,6 +289,8 @@ def _initialise_simulations(self):
             situation=self.data if not macro else None,
             reform=self._parsed_baseline,
         )
+        if self.time_period is not None:
+            self.baseline_sim.default_calculation_period = self.time_period
 
         if "region" in self.options and isinstance(
             self.baseline_sim, CountryMicrosimulation
@@ -299,8 +301,6 @@ def _initialise_simulations(self):
                 self.options["region"],
                 reform=self.baseline_sim.reform,
             )
-        if self.time_period is not None:
-            self.baseline_sim.default_calculation_period = self.time_period
 
         if "subsample" in self.options:
             self.baseline_sim = self.baseline_sim.subsample(
@@ -315,6 +315,9 @@ def _initialise_simulations(self):
                 situation=self.data if not macro else None,
                 reform=self._parsed_reform,
             )
+
+            if self.time_period is not None:
+                self.reformed_sim.default_calculation_period = self.time_period
             if "region" in self.options and isinstance(
                 self.reformed_sim, CountryMicrosimulation
             ):
@@ -330,9 +333,6 @@ def _initialise_simulations(self):
                     self.options["subsample"]
                 )
 
-            if self.time_period is not None:
-                self.reformed_sim.default_calculation_period = self.time_period
-
             # Set the 'baseline tax-benefit system' to be the actual baseline. For example, when working out an individual's
             # baseline MTR, it should use the actual policy baseline, not always current law.
 
@@ -404,6 +404,11 @@ def _apply_region_to_simulation(
                 with h5py.File(weights_file_path, "r") as f:
                     weights = f[str(self.time_period)][...]
 
+                print(
+                    weights[constituency_id],
+                    simulation.default_calculation_period,
+                )
+
                 simulation.calculate("household_net_income")
 
                 simulation.set_input(
@@ -419,7 +424,7 @@ def _apply_region_to_simulation(
                     local_folder=None,
                     version=None,
                 )
-                la_names_file_path = Path(constituency_names_file_path)
+                la_names_file_path = Path(la_names_file_path)
                 la_names = pd.read_csv(la_names_file_path)
                 if la in la_names.code.values:
                     la_id = la_names[la_names.code == la].index[0]

policyengine/utils/maps.py

@@ -7,28 +7,36 @@
 from policyengine.utils.charts import *
 
 
-def plot_hex_map(value_by_area_name: dict, location_type: str) -> dict:
+def get_location_options_table(location_type: str) -> pd.DataFrame:
     if location_type == "parliamentary_constituencies":
         area_names_file_path = download(
             repo="policyengine/policyengine-uk-data",
             repo_filename="constituencies_2024.csv",
             local_folder=None,
             version=None,
         )
-        x_bounds = [30, 85]
-        y_bounds = [-50, 2]
     elif location_type == "local_authorities":
         area_names_file_path = download(
             repo="policyengine/policyengine-uk-data",
             repo_filename="local_authorities_2021.csv",
             local_folder=None,
             version=None,
         )
+    df = pd.read_csv(area_names_file_path)
+    return df
+
+
+def plot_hex_map(value_by_area_name: dict, location_type: str) -> dict:
+    if location_type == "parliamentary_constituencies":
+        x_bounds = [30, 85]
+        y_bounds = [-50, 2]
+    elif location_type == "local_authorities":
         x_bounds = [-20, 25]
         y_bounds = [-10, 35]
     else:
         raise ValueError("Invalid location_type: " + location_type)
-    area_names = pd.read_csv(area_names_file_path)
+
+    area_names = get_location_options_table(location_type)
 
     x_range = area_names["x"].max() - area_names["x"].min()
     y_range = area_names["y"].max() - area_names["y"].min()