Properly filter UK subnational simulation results

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    changed:
+    - Filter nationwide subnational results when user is running sim at subnational level

policyengine/outputs/macro/comparison/calculate_economy_comparison.py

@@ -3,6 +3,10 @@
 from microdf import MicroSeries
 import numpy as np
 from policyengine.utils.data_download import download
+from policyengine.utils.uk_geography import (
+    should_zero_constituency,
+    should_zero_local_authority,
+)
 import pandas as pd
 import h5py
 from pydantic import BaseModel
@@ -692,7 +696,10 @@ class UKConstituencyBreakdownWithValues(BaseModel):
 
 
 def uk_constituency_breakdown(
-    baseline: SingleEconomy, reform: SingleEconomy, country_id: str
+    baseline: SingleEconomy,
+    reform: SingleEconomy,
+    country_id: str,
+    region: str | None = None,
 ) -> UKConstituencyBreakdown:
     if country_id != "uk":
         return None
@@ -701,8 +708,8 @@ def uk_constituency_breakdown(
         "by_constituency": {},
         "outcomes_by_region": {},
     }
-    for region in ["uk", "england", "scotland", "wales", "northern_ireland"]:
-        output["outcomes_by_region"][region] = {
+    for region_ in ["uk", "england", "scotland", "wales", "northern_ireland"]:
+        output["outcomes_by_region"][region_] = {
             "Gain more than 5%": 0,
             "Gain less than 5%": 0,
             "No change": 0,
@@ -732,45 +739,53 @@ def uk_constituency_breakdown(
     for i in range(len(constituency_names)):
         name: str = constituency_names.iloc[i]["name"]
         code: str = constituency_names.iloc[i]["code"]
-        weight: np.ndarray = weights[i]
-        baseline_income = MicroSeries(baseline_hnet, weights=weight)
-        reform_income = MicroSeries(reform_hnet, weights=weight)
-        average_household_income_change: float = (
-            reform_income.sum() - baseline_income.sum()
-        ) / baseline_income.count()
-        percent_household_income_change: float = (
-            reform_income.sum() / baseline_income.sum() - 1
-        )
+
+        if should_zero_constituency(region, code, name):
+            average_household_income_change = 0.0
+            percent_household_income_change = 0.0
+        else:
+            weight: np.ndarray = weights[i]
+            baseline_income = MicroSeries(baseline_hnet, weights=weight)
+            reform_income = MicroSeries(reform_hnet, weights=weight)
+            average_household_income_change = (
+                reform_income.sum() - baseline_income.sum()
+            ) / baseline_income.count()
+            percent_household_income_change = (
+                reform_income.sum() / baseline_income.sum() - 1
+            )
+
         output["by_constituency"][name] = {
             "average_household_income_change": average_household_income_change,
             "relative_household_income_change": percent_household_income_change,
             "x": int(constituency_names.iloc[i]["x"]),  # Geographic positions
             "y": int(constituency_names.iloc[i]["y"]),
         }
 
-        regions = ["uk"]
-        if "E" in code:
-            regions.append("england")
-        elif "S" in code:
-            regions.append("scotland")
-        elif "W" in code:
-            regions.append("wales")
-        elif "N" in code:
-            regions.append("northern_ireland")
-
-        if percent_household_income_change > 0.05:
-            bucket = "Gain more than 5%"
-        elif percent_household_income_change > 1e-3:
-            bucket = "Gain less than 5%"
-        elif percent_household_income_change > -1e-3:
-            bucket = "No change"
-        elif percent_household_income_change > -0.05:
-            bucket = "Lose less than 5%"
-        else:
-            bucket = "Lose more than 5%"
+        # Only count non-zeroed constituencies in outcomes_by_region
+        if not should_zero_constituency(region, code, name):
+            regions = ["uk"]
+            if "E" in code:
+                regions.append("england")
+            elif "S" in code:
+                regions.append("scotland")
+            elif "W" in code:
+                regions.append("wales")
+            elif "N" in code:
+                regions.append("northern_ireland")
+
+            if percent_household_income_change > 0.05:
+                bucket = "Gain more than 5%"
+            elif percent_household_income_change > 1e-3:
+                bucket = "Gain less than 5%"
+            elif percent_household_income_change > -1e-3:
+                bucket = "No change"
+            elif percent_household_income_change > -0.05:
+                bucket = "Lose less than 5%"
+            else:
+                bucket = "Lose more than 5%"
 
-        for region_ in regions:
-            output["outcomes_by_region"][region_][bucket] += 1
+            for region_ in regions:
+                output["outcomes_by_region"][region_][bucket] += 1
 
     return UKConstituencyBreakdownWithValues(**output)
 
@@ -790,7 +805,10 @@ class UKLocalAuthorityBreakdownWithValues(BaseModel):
 
 
 def uk_local_authority_breakdown(
-    baseline: SingleEconomy, reform: SingleEconomy, country_id: str
+    baseline: SingleEconomy,
+    reform: SingleEconomy,
+    country_id: str,
+    region: str | None = None,
 ) -> UKLocalAuthorityBreakdown:
     if country_id != "uk":
         return None
@@ -822,15 +840,21 @@ def uk_local_authority_breakdown(
     for i in range(len(local_authority_names)):
         name: str = local_authority_names.iloc[i]["name"]
         code: str = local_authority_names.iloc[i]["code"]
-        weight: np.ndarray = weights[i]
-        baseline_income = MicroSeries(baseline_hnet, weights=weight)
-        reform_income = MicroSeries(reform_hnet, weights=weight)
-        average_household_income_change: float = (
-            reform_income.sum() - baseline_income.sum()
-        ) / baseline_income.count()
-        percent_household_income_change: float = (
-            reform_income.sum() / baseline_income.sum() - 1
-        )
+
+        if should_zero_local_authority(region, code, name):
+            average_household_income_change = 0.0
+            percent_household_income_change = 0.0
+        else:
+            weight: np.ndarray = weights[i]
+            baseline_income = MicroSeries(baseline_hnet, weights=weight)
+            reform_income = MicroSeries(reform_hnet, weights=weight)
+            average_household_income_change = (
+                reform_income.sum() - baseline_income.sum()
+            ) / baseline_income.count()
+            percent_household_income_change = (
+                reform_income.sum() / baseline_income.sum() - 1
+            )
+
         output["by_local_authority"][name] = {
             "average_household_income_change": average_household_income_change,
             "relative_household_income_change": percent_household_income_change,
@@ -840,8 +864,6 @@ def uk_local_authority_breakdown(
             "y": int(local_authority_names.iloc[i]["y"]),
         }
 
-        # Note: Country-level aggregation and bucketing logic removed for local authorities
-
     return UKLocalAuthorityBreakdownWithValues(**output)
 
 
@@ -901,10 +923,12 @@ def calculate_economy_comparison(
     intra_decile_impact_data = intra_decile_impact(baseline, reform)
     labor_supply_response_data = labor_supply_response(baseline, reform)
     constituency_impact_data: UKConstituencyBreakdown = (
-        uk_constituency_breakdown(baseline, reform, country_id)
+        uk_constituency_breakdown(baseline, reform, country_id, options.region)
     )
     local_authority_impact_data: UKLocalAuthorityBreakdown = (
-        uk_local_authority_breakdown(baseline, reform, country_id)
+        uk_local_authority_breakdown(
+            baseline, reform, country_id, options.region
+        )
     )
     wealth_decile_impact_data = wealth_decile_impact(
         baseline, reform, country_id

policyengine/utils/uk_geography.py

@@ -0,0 +1,90 @@
+"""Utilities for UK geographic region filtering."""
+
+from typing import Literal
+
+UKRegionType = Literal["uk", "country", "constituency", "local_authority"]
+
+UK_REGION_TYPES: tuple[UKRegionType, ...] = (
+    "uk",
+    "country",
+    "constituency",
+    "local_authority",
+)
+
+
+def determine_uk_region_type(region: str | None) -> UKRegionType:
+    """
+    Determine the type of UK region from a region string.
+
+    Args:
+        region: A region string (e.g., "country/scotland", "constituency/Aberdeen North",
+                "local_authority/leicester") or None.
+
+    Returns:
+        One of "uk", "country", "constituency", or "local_authority".
+
+    Raises:
+        ValueError: If the region prefix is not a valid UK region type.
+    """
+    if region is None:
+        return "uk"
+
+    prefix = region.split("/")[0]
+    if prefix not in UK_REGION_TYPES:
+        raise ValueError(
+            f"Invalid UK region type: '{prefix}'. "
+            f"Expected one of: {list(UK_REGION_TYPES)}"
+        )
+
+    return prefix
+
+
+def get_country_from_code(code: str) -> str | None:
+    """Get country name from geographic code prefix (E, S, W, N)."""
+    prefix_map = {
+        "E": "england",
+        "S": "scotland",
+        "W": "wales",
+        "N": "northern_ireland",
+    }
+    return prefix_map.get(code[0])
+
+
+def should_zero_constituency(region: str | None, code: str, name: str) -> bool:
+    """Return True if this constituency's impacts should be zeroed out."""
+    region_type = determine_uk_region_type(region)
+
+    if region_type == "uk":
+        return False
+    # region is guaranteed to be non-None for non-uk region types
+    assert region is not None
+    if region_type == "country":
+        target = region.split("/")[1]
+        return get_country_from_code(code) != target
+    if region_type == "constituency":
+        target = region.split("/")[1]
+        return code != target and name != target
+    if region_type == "local_authority":
+        return True
+    return False
+
+
+def should_zero_local_authority(
+    region: str | None, code: str, name: str
+) -> bool:
+    """Return True if this local authority's impacts should be zeroed out."""
+    region_type = determine_uk_region_type(region)
+
+    if region_type == "uk":
+        return False
+    # region is guaranteed to be non-None for non-uk region types
+    assert region is not None
+    if region_type == "country":
+        target = region.split("/")[1]
+        return get_country_from_code(code) != target
+    if region_type == "local_authority":
+        target = region.split("/")[1]
+        return code != target and name != target
+    if region_type == "constituency":
+        return True
+    return False