Merge pull request #51 from PolicyEngine/dev-improv

Improve constituency mapping tools

Author: nikhilwoodruff

docs/_config.yml

@@ -1,7 +1,7 @@
 # _config.yml
-title: PolicyEngine
 author: PolicyEngine
-copyright: "2022"
+title: PolicyEngine
+copyright: "2025"
 logo: logo.png
 
 execute:
@@ -19,11 +19,4 @@ sphinx:
     html_theme: furo
     pygments_style: default
     html_css_files:
-      - style.css
-  # extra_extensions:
-  #   - sphinx.ext.autodoc
-  #   - sphinxarg.ext
-  #   - sphinx.ext.viewcode
-  #   - sphinx.ext.napoleon
-  #   - sphinx_math_dollar
-  #   - sphinx.ext.mathjax
+      - style.css

policyengine/outputs/macro/comparison/budget/general.py

@@ -104,6 +104,8 @@ def budget_chart(simulation: Simulation, data: dict) -> go.Figure:
         xaxis_title="",
         yaxis_title="Budgetary impact (£ billions)",
         yaxis_tickformat=",.0f",
+        uniformtext_minsize=12,
+        uniformtext_mode="hide",
     )
 
-    return format_fig(fig)
+    return format_fig(fig, simulation.country)

policyengine/outputs/macro/comparison/local_areas/parliamentary_constituencies.py

@@ -3,165 +3,44 @@
 from policyengine.utils.huggingface import download
 import plotly.express as px
 from policyengine.utils.charts import *
+from policyengine.utils.constituency_maps import plot_hex_map
+from typing import Callable
+from policyengine_core import Microsimulation
+from microdf import MicroSeries
 
 
 def parliamentary_constituencies(
     simulation: Simulation,
     chart: bool = False,
-    variable: str = None,
-    aggregator: str = None,
-    relative: bool = None,
+    metric: Callable[[Microsimulation], MicroSeries] = None,
+    comparator: bool = None,
 ) -> dict:
     if not simulation.options.get("include_constituencies"):
         return {}
 
-    if chart:
-        return heatmap(
-            simulation=simulation,
-            variable=variable,
-            aggregator=aggregator,
-            relative=relative,
-        )
-
-    constituency_baseline = simulation.calculate(
-        "macro/baseline/gov/local_areas/parliamentary_constituencies"
-    )
-    constituency_reform = simulation.calculate(
-        "macro/reform/gov/local_areas/parliamentary_constituencies"
-    )
-
-    result = {}
-
-    for constituency in constituency_baseline:
-        result[constituency] = {}
-        for key in constituency_baseline[constituency]:
-            result[constituency][key] = {
-                "change": constituency_reform[constituency][key]
-                - constituency_baseline[constituency][key],
-                "baseline": constituency_baseline[constituency][key],
-                "reform": constituency_reform[constituency][key],
-            }
-
-    return result
-
-
-def heatmap(
-    simulation: Simulation,
-    variable: str = None,
-    aggregator: str = None,
-    relative: bool = None,
-) -> dict:
-    if not simulation.options.get("include_constituencies"):
-        return {}
-
-    options = {}
+    kwargs = {}
+    if metric is not None:
+        kwargs["metric"] = metric
 
-    if variable is not None:
-        options["variables"] = [variable]
-    if aggregator is not None:
-        options["aggregator"] = aggregator
+    if comparator is None:
+        comparator = lambda x, y: (y / x) - 1
 
     constituency_baseline = simulation.calculate(
-        "macro/baseline/gov/local_areas/parliamentary_constituencies",
-        **options,
+        "macro/baseline/gov/local_areas/parliamentary_constituencies", **kwargs
     )
     constituency_reform = simulation.calculate(
-        "macro/reform/gov/local_areas/parliamentary_constituencies", **options
+        "macro/reform/gov/local_areas/parliamentary_constituencies", **kwargs
     )
 
     result = {}
 
-    constituency_names_file_path = download(
-        repo="policyengine/policyengine-uk-data",
-        repo_filename="constituencies_2024.csv",
-        local_folder=None,
-        version=None,
-    )
-    constituency_names = pd.read_csv(constituency_names_file_path)
-
-    if variable is None:
-        variable = "household_net_income"
-    if relative is None:
-        relative = True
-
     for constituency in constituency_baseline:
-        if relative:
-            result[constituency] = (
-                constituency_reform[constituency][variable]
-                / constituency_baseline[constituency][variable]
-                - 1
-            )
-        else:
-            result[constituency] = (
-                constituency_reform[constituency][variable]
-                - constituency_baseline[constituency][variable]
-            )
-
-    x_range = constituency_names["x"].max() - constituency_names["x"].min()
-    y_range = constituency_names["y"].max() - constituency_names["y"].min()
-    # Expand x range to preserve aspect ratio
-    expanded_lower_x_range = -(y_range - x_range) / 2
-    expanded_upper_x_range = x_range - expanded_lower_x_range
-    constituency_names.x = (
-        constituency_names.x - (constituency_names.y % 2 == 0) * 0.5
-    )
-    constituency_names["Relative change"] = (
-        pd.Series(list(result.values()), index=list(result.keys()))
-        .loc[constituency_names["name"]]
-        .values
-    )
-
-    label = simulation.baseline.tax_benefit_system.variables[variable].label
-
-    fig = px.scatter(
-        constituency_names,
-        x="x",
-        y="y",
-        color="Relative change",
-        hover_name="name",
-        title=f"{'Relative change' if relative else 'Change'} in {label} by parliamentary constituency",
-    )
-
-    format_fig(fig)
-
-    # Show hexagons on scatter points
-
-    fig.update_traces(
-        marker=dict(
-            symbol="hexagon", line=dict(width=0, color="lightgray"), size=15
+        result[constituency] = comparator(
+            constituency_baseline[constituency],
+            constituency_reform[constituency],
         )
-    )
-    fig.update_layout(
-        xaxis_tickvals=[],
-        xaxis_title="",
-        yaxis_tickvals=[],
-        yaxis_title="",
-        xaxis_range=[expanded_lower_x_range, expanded_upper_x_range],
-        yaxis_range=[
-            constituency_names["y"].min(),
-            constituency_names["y"].max(),
-        ],
-    ).update_traces(marker_size=10).update_layout(
-        xaxis_range=[30, 85], yaxis_range=[-50, 2]
-    )
-
-    x_min = fig.data[0]["marker"]["color"].min()
-    x_max = fig.data[0]["marker"]["color"].max()
-    max_abs = max(abs(x_min), abs(x_max))
 
-    fig.update_layout(
-        coloraxis=dict(
-            cmin=-max_abs,
-            cmax=max_abs,
-            colorscale=[
-                [0, DARK_GRAY],
-                [0.5, "lightgray"],
-                [1, BLUE],
-            ],
-            colorbar=dict(
-                tickformat=".0%" if relative else ",.0f",
-            ),
-        )
-    )
+    if chart:
+        return plot_hex_map(result)
 
-    return fig
+    return result

policyengine/outputs/macro/single/gov/local_areas/parliamentary_constituencies.py

@@ -1,8 +1,11 @@
 from policyengine import Simulation
 from policyengine.utils.huggingface import download
 import h5py
-from microdf import MicroDataFrame
+from microdf import MicroSeries
 import pandas as pd
+from typing import Callable
+from policyengine_core import Microsimulation
+from policyengine.utils.constituency_maps import plot_hex_map
 
 DEFAULT_VARIABLES = [
     "household_net_income",
@@ -11,11 +14,21 @@
 
 def parliamentary_constituencies(
     simulation: Simulation,
-    variables: list = DEFAULT_VARIABLES,
-    aggregator: str = "sum",
+    metric: Callable[[Microsimulation], MicroSeries] = None,
+    chart: bool = False,
 ) -> dict:
+    """Calculate the impact of the reform on parliamentary constituencies.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+        custom_function (Callable[[Microsimulation], [float]]): A custom function to calculate the impact. This must be called on a Microsimulation and return a float (we will call it for each constituency weight set).
+
+    """
     if not simulation.options.get("include_constituencies"):
         return {}
+
+    if metric is None:
+        metric = lambda sim: sim.calculate("household_net_income").median()
     weights_file_path = download(
         repo="policyengine/policyengine-uk-data",
         repo_filename="parliamentary_constituency_weights.h5",
@@ -33,15 +46,40 @@ def parliamentary_constituencies(
     with h5py.File(weights_file_path, "r") as f:
         weights = f[str(simulation.time_period)][...]
 
-    sim_df = simulation.selected.calculate_dataframe(variables)
-
     result = {}
 
+    sim = simulation.selected
+    original_hh_weight = sim.calculate("household_weight").values
+
     for constituency_id in range(weights.shape[0]):
-        weighted_df = MicroDataFrame(sim_df, weights=weights[constituency_id])
+        sim.set_input(
+            "household_weight",
+            sim.default_calculation_period,
+            weights[constituency_id],
+        )
+        sim.get_holder("person_weight").delete_arrays(
+            sim.default_calculation_period
+        )
+        sim.get_holder("benunit_weight").delete_arrays(
+            sim.default_calculation_period
+        )
+        calculation_result = metric(simulation.selected)
         code = constituency_names.code.iloc[constituency_id]
         result[constituency_names.set_index("code").loc[code]["name"]] = (
-            getattr(weighted_df, aggregator)().to_dict()
+            calculation_result
         )
 
+    sim.get_holder("person_weight").delete_arrays(
+        sim.default_calculation_period
+    )
+    sim.get_holder("benunit_weight").delete_arrays(
+        sim.default_calculation_period
+    )
+    sim.set_input(
+        "household_weight", sim.default_calculation_period, original_hh_weight
+    )
+
+    if chart:
+        return plot_hex_map(result)
+
     return result

policyengine/outputs/macro/single/household/income_distribution.py

@@ -29,7 +29,7 @@ def income_distribution_chart(simulation: Simulation) -> go.Figure:
         color_discrete_sequence=[BLUE],
     )
     fig.update_layout(
-        title="Number of household by net income band",
+        title="Number of households by net income band",
         xaxis_title="Household net income",
         yaxis_title="Number of households",
     )

policyengine/simulation.py

@@ -39,6 +39,8 @@ class Simulation:
     """The tax-benefit simulation for the baseline scenario."""
     reformed: CountrySimulation = None
     """The tax-benefit simulation for the reformed scenario."""
+    selected: CountryMicrosimulation = None
+    """The selected simulation for the current calculation."""
     verbose: bool = False
     """Whether to print out progress messages."""
 
@@ -74,6 +76,11 @@ def __init__(
         elif isinstance(reform, int):
             reform = Reform.from_api(reform, country_id=country)
 
+        if isinstance(baseline, dict):
+            baseline = Reform.from_dict(baseline, country_id=country)
+        elif isinstance(baseline, int):
+            baseline = Reform.from_api(baseline, country_id=country)
+
         self.baseline = baseline
         self.reform = reform