From 8b9f9e35bc421fe67e59d8e428dd54a6a4f21164 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 21:10:41 +0100 Subject: [PATCH 1/5] feat: Add congressional district impact to economy analysis Create congressional_district_impacts table and wire per-district income change computation into US economy comparison (local and Modal). Add congressional_district_impact field to EconomicImpactResponse. Co-Authored-By: Claude Opus 4.6 --- ...dd_congressional_district_impacts_table.py | 48 ++++++++++++++ src/policyengine_api/api/analysis.py | 63 +++++++++++++++++++ src/policyengine_api/modal_app.py | 35 +++++++++++ src/policyengine_api/models/__init__.py | 8 +++ .../models/congressional_district_impact.py | 42 +++++++++++++ 5 files changed, 196 insertions(+) create mode 100644 alembic/versions/20260220_a5e4144467e5_add_congressional_district_impacts_table.py create mode 100644 src/policyengine_api/models/congressional_district_impact.py diff --git a/alembic/versions/20260220_a5e4144467e5_add_congressional_district_impacts_table.py b/alembic/versions/20260220_a5e4144467e5_add_congressional_district_impacts_table.py new file mode 100644 index 0000000..5e04628 --- /dev/null +++ b/alembic/versions/20260220_a5e4144467e5_add_congressional_district_impacts_table.py @@ -0,0 +1,48 @@ +"""add_congressional_district_impacts_table + +Revision ID: a5e4144467e5 +Revises: 75e5dca14603 +Create Date: 2026-02-20 20:52:53.243197 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'a5e4144467e5' +down_revision: Union[str, Sequence[str], None] = '75e5dca14603' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('congressional_district_impacts', + sa.Column('baseline_simulation_id', sa.Uuid(), nullable=False), + sa.Column('reform_simulation_id', sa.Uuid(), nullable=False), + sa.Column('report_id', sa.Uuid(), nullable=True), + sa.Column('district_geoid', sa.Integer(), nullable=False), + sa.Column('state_fips', sa.Integer(), nullable=False), + sa.Column('district_number', sa.Integer(), nullable=False), + sa.Column('average_household_income_change', sa.Float(), nullable=False), + sa.Column('relative_household_income_change', sa.Float(), nullable=False), + sa.Column('population', sa.Float(), nullable=False), + sa.Column('id', sa.Uuid(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(['baseline_simulation_id'], ['simulations.id'], ), + sa.ForeignKeyConstraint(['reform_simulation_id'], ['simulations.id'], ), + sa.ForeignKeyConstraint(['report_id'], ['reports.id'], ), + sa.PrimaryKeyConstraint('id') + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table('congressional_district_impacts') + # ### end Alembic commands ### diff --git a/src/policyengine_api/api/analysis.py b/src/policyengine_api/api/analysis.py index 33fb27a..cfb1766 100644 --- a/src/policyengine_api/api/analysis.py +++ b/src/policyengine_api/api/analysis.py @@ -28,6 +28,8 @@ from policyengine_api.models import ( BudgetSummary, BudgetSummaryRead, + CongressionalDistrictImpact, + CongressionalDistrictImpactRead, Dataset, DecileImpact, DecileImpactRead, @@ -147,6 +149,7 @@ class EconomicImpactResponse(BaseModel): budget_summary: list[BudgetSummaryRead] | None = None intra_decile: list[IntraDecileImpactRead] | None = None detailed_budget: dict[str, dict[str, float | None]] | None = None + congressional_district_impact: list[CongressionalDistrictImpactRead] | None = None def _get_model_version( @@ -291,6 +294,7 @@ def _build_response( budget_summary_records = None intra_decile_records = None detailed_budget = None + district_impact_records = None if report.status == ReportStatus.COMPLETED: # Fetch decile impacts for this report @@ -438,6 +442,34 @@ def _build_response( for r in intra_rows ] + # Fetch congressional district impact records for this report + district_rows = session.exec( + select(CongressionalDistrictImpact).where( + CongressionalDistrictImpact.report_id == report.id + ) + ).all() + if district_rows: + district_impact_records = [ + CongressionalDistrictImpactRead( + id=d.id, + created_at=d.created_at, + baseline_simulation_id=d.baseline_simulation_id, + reform_simulation_id=d.reform_simulation_id, + report_id=d.report_id, + district_geoid=d.district_geoid, + state_fips=d.state_fips, + district_number=d.district_number, + average_household_income_change=_safe_float( + d.average_household_income_change + ), + relative_household_income_change=_safe_float( + d.relative_household_income_change + ), + population=_safe_float(d.population), + ) + for d in district_rows + ] + region_info = None if region: region_info = RegionInfo( @@ -471,6 +503,7 @@ def _build_response( budget_summary=budget_summary_records, intra_decile=intra_decile_records, detailed_budget=detailed_budget, + congressional_district_impact=district_impact_records, ) @@ -1289,6 +1322,36 @@ def build_dynamic(dynamic_id): ) session.add(record) + # Calculate congressional district impact + from policyengine.outputs.congressional_district_impact import ( + compute_us_congressional_district_impacts, + ) + + try: + district_impact = compute_us_congressional_district_impacts( + pe_baseline_sim, pe_reform_sim + ) + if district_impact.district_results: + for dr in district_impact.district_results: + record = CongressionalDistrictImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + district_geoid=dr["district_geoid"], + state_fips=dr["state_fips"], + district_number=dr["district_number"], + average_household_income_change=dr[ + "average_household_income_change" + ], + relative_household_income_change=dr[ + "relative_household_income_change" + ], + population=dr["population"], + ) + session.add(record) + except KeyError: + pass # congressional_district_geoid not in dataset + # Mark completed baseline_sim.status = SimulationStatus.COMPLETED baseline_sim.completed_at = datetime.now(timezone.utc) diff --git a/src/policyengine_api/modal_app.py b/src/policyengine_api/modal_app.py index 3699010..063114c 100644 --- a/src/policyengine_api/modal_app.py +++ b/src/policyengine_api/modal_app.py @@ -1672,6 +1672,7 @@ def economy_comparison_us(job_id: str, traceparent: str | None = None) -> None: # Import models inline from policyengine_api.models import ( BudgetSummary, + CongressionalDistrictImpact, Dataset, DecileImpact, Inequality, @@ -2135,6 +2136,40 @@ def economy_comparison_us(job_id: str, traceparent: str | None = None) -> None: ) session.add(record) + # Calculate congressional district impact + from policyengine.outputs.congressional_district_impact import ( + compute_us_congressional_district_impacts, + ) + + try: + district_impact = ( + compute_us_congressional_district_impacts( + pe_baseline_sim, pe_reform_sim + ) + ) + if district_impact.district_results: + for dr in district_impact.district_results: + record = CongressionalDistrictImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + district_geoid=dr["district_geoid"], + state_fips=dr["state_fips"], + district_number=dr[ + "district_number" + ], + average_household_income_change=dr[ + "average_household_income_change" + ], + relative_household_income_change=dr[ + "relative_household_income_change" + ], + population=dr["population"], + ) + session.add(record) + except KeyError: + pass # congressional_district_geoid not in dataset + # Mark simulations and report as completed baseline_sim.status = SimulationStatus.COMPLETED baseline_sim.completed_at = datetime.now(timezone.utc) diff --git a/src/policyengine_api/models/__init__.py b/src/policyengine_api/models/__init__.py index 257d03a..a342c53 100644 --- a/src/policyengine_api/models/__init__.py +++ b/src/policyengine_api/models/__init__.py @@ -5,6 +5,11 @@ BudgetSummaryCreate, BudgetSummaryRead, ) +from .congressional_district_impact import ( + CongressionalDistrictImpact, + CongressionalDistrictImpactCreate, + CongressionalDistrictImpactRead, +) from .change_aggregate import ( ChangeAggregate, ChangeAggregateCreate, @@ -100,6 +105,9 @@ "AggregateOutputRead", "AggregateStatus", "AggregateType", + "CongressionalDistrictImpact", + "CongressionalDistrictImpactCreate", + "CongressionalDistrictImpactRead", "ChangeAggregate", "ChangeAggregateCreate", "ChangeAggregateRead", diff --git a/src/policyengine_api/models/congressional_district_impact.py b/src/policyengine_api/models/congressional_district_impact.py new file mode 100644 index 0000000..f5bbc24 --- /dev/null +++ b/src/policyengine_api/models/congressional_district_impact.py @@ -0,0 +1,42 @@ +"""Congressional district impact output model.""" + +from datetime import datetime, timezone +from uuid import UUID, uuid4 + +from sqlmodel import Field, SQLModel + + +class CongressionalDistrictImpactBase(SQLModel): + """Base congressional district impact fields.""" + + baseline_simulation_id: UUID = Field(foreign_key="simulations.id") + reform_simulation_id: UUID = Field(foreign_key="simulations.id") + report_id: UUID | None = Field(default=None, foreign_key="reports.id") + district_geoid: int + state_fips: int + district_number: int + average_household_income_change: float + relative_household_income_change: float + population: float + + +class CongressionalDistrictImpact(CongressionalDistrictImpactBase, table=True): + """Congressional district impact database model.""" + + __tablename__ = "congressional_district_impacts" + + id: UUID = Field(default_factory=uuid4, primary_key=True) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + + +class CongressionalDistrictImpactCreate(CongressionalDistrictImpactBase): + """Schema for creating congressional district impacts.""" + + pass + + +class CongressionalDistrictImpactRead(CongressionalDistrictImpactBase): + """Schema for reading congressional district impacts.""" + + id: UUID + created_at: datetime From 87c6d945c5ce442996a62f7642a83df43cf6b404 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 22:08:59 +0100 Subject: [PATCH 2/5] feat: Add UK constituency impact to economy comparison Adds ConstituencyImpact DB model, Alembic migration, and wires constituency computation into both local and Modal UK economy comparison paths. Uses GCS-hosted weight matrix and constituency CSV. Co-Authored-By: Claude Opus 4.6 --- ...eb90bd74_add_constituency_impacts_table.py | 50 ++++++++++++ src/policyengine_api/api/analysis.py | 78 +++++++++++++++++++ src/policyengine_api/modal_app.py | 55 +++++++++++++ src/policyengine_api/models/__init__.py | 8 ++ .../models/constituency_impact.py | 43 ++++++++++ 5 files changed, 234 insertions(+) create mode 100644 alembic/versions/20260220_83ceeb90bd74_add_constituency_impacts_table.py create mode 100644 src/policyengine_api/models/constituency_impact.py diff --git a/alembic/versions/20260220_83ceeb90bd74_add_constituency_impacts_table.py b/alembic/versions/20260220_83ceeb90bd74_add_constituency_impacts_table.py new file mode 100644 index 0000000..c69410c --- /dev/null +++ b/alembic/versions/20260220_83ceeb90bd74_add_constituency_impacts_table.py @@ -0,0 +1,50 @@ +"""add_constituency_impacts_table + +Revision ID: 83ceeb90bd74 +Revises: a5e4144467e5 +Create Date: 2026-02-20 21:44:14.780195 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +# revision identifiers, used by Alembic. +revision: str = '83ceeb90bd74' +down_revision: Union[str, Sequence[str], None] = 'a5e4144467e5' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('constituency_impacts', + sa.Column('baseline_simulation_id', sa.Uuid(), nullable=False), + sa.Column('reform_simulation_id', sa.Uuid(), nullable=False), + sa.Column('report_id', sa.Uuid(), nullable=True), + sa.Column('constituency_code', sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column('constituency_name', sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column('x', sa.Integer(), nullable=False), + sa.Column('y', sa.Integer(), nullable=False), + sa.Column('average_household_income_change', sa.Float(), nullable=False), + sa.Column('relative_household_income_change', sa.Float(), nullable=False), + sa.Column('population', sa.Float(), nullable=False), + sa.Column('id', sa.Uuid(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(['baseline_simulation_id'], ['simulations.id'], ), + sa.ForeignKeyConstraint(['reform_simulation_id'], ['simulations.id'], ), + sa.ForeignKeyConstraint(['report_id'], ['reports.id'], ), + sa.PrimaryKeyConstraint('id') + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table('constituency_impacts') + # ### end Alembic commands ### diff --git a/src/policyengine_api/api/analysis.py b/src/policyengine_api/api/analysis.py index cfb1766..b45a111 100644 --- a/src/policyengine_api/api/analysis.py +++ b/src/policyengine_api/api/analysis.py @@ -30,6 +30,8 @@ BudgetSummaryRead, CongressionalDistrictImpact, CongressionalDistrictImpactRead, + ConstituencyImpact, + ConstituencyImpactRead, Dataset, DecileImpact, DecileImpactRead, @@ -150,6 +152,7 @@ class EconomicImpactResponse(BaseModel): intra_decile: list[IntraDecileImpactRead] | None = None detailed_budget: dict[str, dict[str, float | None]] | None = None congressional_district_impact: list[CongressionalDistrictImpactRead] | None = None + constituency_impact: list[ConstituencyImpactRead] | None = None def _get_model_version( @@ -295,6 +298,7 @@ def _build_response( intra_decile_records = None detailed_budget = None district_impact_records = None + constituency_impact_records = None if report.status == ReportStatus.COMPLETED: # Fetch decile impacts for this report @@ -470,6 +474,35 @@ def _build_response( for d in district_rows ] + # Fetch constituency impact records for this report + constituency_rows = session.exec( + select(ConstituencyImpact).where( + ConstituencyImpact.report_id == report.id + ) + ).all() + if constituency_rows: + constituency_impact_records = [ + ConstituencyImpactRead( + id=c.id, + created_at=c.created_at, + baseline_simulation_id=c.baseline_simulation_id, + reform_simulation_id=c.reform_simulation_id, + report_id=c.report_id, + constituency_code=c.constituency_code, + constituency_name=c.constituency_name, + x=c.x, + y=c.y, + average_household_income_change=_safe_float( + c.average_household_income_change + ), + relative_household_income_change=_safe_float( + c.relative_household_income_change + ), + population=_safe_float(c.population), + ) + for c in constituency_rows + ] + region_info = None if region: region_info = RegionInfo( @@ -504,6 +537,7 @@ def _build_response( intra_decile=intra_decile_records, detailed_budget=detailed_budget, congressional_district_impact=district_impact_records, + constituency_impact=constituency_impact_records, ) @@ -912,6 +946,50 @@ def build_dynamic(dynamic_id): ) session.add(record) + # Calculate constituency impact (UK only, requires weight matrix) + from policyengine.outputs.constituency_impact import ( + compute_uk_constituency_impacts, + ) + + try: + from policyengine_core.tools.google_cloud import download as gcs_download + + weight_matrix_path = gcs_download( + gcs_bucket="policyengine-uk-data-private", + gcs_key="parliamentary_constituency_weights.h5", + ) + constituency_csv_path = gcs_download( + gcs_bucket="policyengine-uk-data-private", + gcs_key="constituencies_2024.csv", + ) + constituency_impact = compute_uk_constituency_impacts( + pe_baseline_sim, + pe_reform_sim, + weight_matrix_path=weight_matrix_path, + constituency_csv_path=constituency_csv_path, + ) + if constituency_impact.constituency_results: + for cr in constituency_impact.constituency_results: + record = ConstituencyImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + constituency_code=cr["constituency_code"], + constituency_name=cr["constituency_name"], + x=cr["x"], + y=cr["y"], + average_household_income_change=cr[ + "average_household_income_change" + ], + relative_household_income_change=cr[ + "relative_household_income_change" + ], + population=cr["population"], + ) + session.add(record) + except Exception: + pass # Weight matrix not available, skip constituency impact + # Mark completed baseline_sim.status = SimulationStatus.COMPLETED baseline_sim.completed_at = datetime.now(timezone.utc) diff --git a/src/policyengine_api/modal_app.py b/src/policyengine_api/modal_app.py index 063114c..365eeee 100644 --- a/src/policyengine_api/modal_app.py +++ b/src/policyengine_api/modal_app.py @@ -1131,6 +1131,7 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None: # Import models inline from policyengine_api.models import ( BudgetSummary, + ConstituencyImpact, Dataset, DecileImpact, Inequality, @@ -1605,6 +1606,60 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None: ) session.add(record) + # Calculate constituency impact + from policyengine.outputs.constituency_impact import ( + compute_uk_constituency_impacts, + ) + + try: + from policyengine_core.tools.google_cloud import ( + download as gcs_download, + ) + + weight_matrix_path = gcs_download( + gcs_bucket="policyengine-uk-data-private", + gcs_key="parliamentary_constituency_weights.h5", + ) + constituency_csv_path = gcs_download( + gcs_bucket="policyengine-uk-data-private", + gcs_key="constituencies_2024.csv", + ) + constituency_impact = ( + compute_uk_constituency_impacts( + pe_baseline_sim, + pe_reform_sim, + weight_matrix_path=weight_matrix_path, + constituency_csv_path=constituency_csv_path, + ) + ) + if constituency_impact.constituency_results: + for cr in ( + constituency_impact.constituency_results + ): + record = ConstituencyImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + constituency_code=cr[ + "constituency_code" + ], + constituency_name=cr[ + "constituency_name" + ], + x=cr["x"], + y=cr["y"], + average_household_income_change=cr[ + "average_household_income_change" + ], + relative_household_income_change=cr[ + "relative_household_income_change" + ], + population=cr["population"], + ) + session.add(record) + except Exception: + pass # Weight matrix not available, skip + # Mark simulations and report as completed baseline_sim.status = SimulationStatus.COMPLETED baseline_sim.completed_at = datetime.now(timezone.utc) diff --git a/src/policyengine_api/models/__init__.py b/src/policyengine_api/models/__init__.py index a342c53..9f4896f 100644 --- a/src/policyengine_api/models/__init__.py +++ b/src/policyengine_api/models/__init__.py @@ -10,6 +10,11 @@ CongressionalDistrictImpactCreate, CongressionalDistrictImpactRead, ) +from .constituency_impact import ( + ConstituencyImpact, + ConstituencyImpactCreate, + ConstituencyImpactRead, +) from .change_aggregate import ( ChangeAggregate, ChangeAggregateCreate, @@ -108,6 +113,9 @@ "CongressionalDistrictImpact", "CongressionalDistrictImpactCreate", "CongressionalDistrictImpactRead", + "ConstituencyImpact", + "ConstituencyImpactCreate", + "ConstituencyImpactRead", "ChangeAggregate", "ChangeAggregateCreate", "ChangeAggregateRead", diff --git a/src/policyengine_api/models/constituency_impact.py b/src/policyengine_api/models/constituency_impact.py new file mode 100644 index 0000000..b44d1ad --- /dev/null +++ b/src/policyengine_api/models/constituency_impact.py @@ -0,0 +1,43 @@ +"""UK parliamentary constituency impact output model.""" + +from datetime import datetime, timezone +from uuid import UUID, uuid4 + +from sqlmodel import Field, SQLModel + + +class ConstituencyImpactBase(SQLModel): + """Base constituency impact fields.""" + + baseline_simulation_id: UUID = Field(foreign_key="simulations.id") + reform_simulation_id: UUID = Field(foreign_key="simulations.id") + report_id: UUID | None = Field(default=None, foreign_key="reports.id") + constituency_code: str + constituency_name: str + x: int + y: int + average_household_income_change: float + relative_household_income_change: float + population: float + + +class ConstituencyImpact(ConstituencyImpactBase, table=True): + """Constituency impact database model.""" + + __tablename__ = "constituency_impacts" + + id: UUID = Field(default_factory=uuid4, primary_key=True) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + + +class ConstituencyImpactCreate(ConstituencyImpactBase): + """Schema for creating constituency impacts.""" + + pass + + +class ConstituencyImpactRead(ConstituencyImpactBase): + """Schema for reading constituency impacts.""" + + id: UUID + created_at: datetime From bb9458a807778299bb21576c5efb4279a02d06e5 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 22:23:34 +0100 Subject: [PATCH 3/5] feat: Add UK local authority impact to economy comparison Adds LocalAuthorityImpact DB model, Alembic migration, and wires computation into both local and Modal UK economy comparison paths. Uses GCS-hosted weight matrix and local authority CSV. Co-Authored-By: Claude Opus 4.6 --- ...8d272_add_local_authority_impacts_table.py | 50 ++++++++++++ src/policyengine_api/api/analysis.py | 78 +++++++++++++++++++ src/policyengine_api/modal_app.py | 53 +++++++++++++ src/policyengine_api/models/__init__.py | 8 ++ .../models/local_authority_impact.py | 43 ++++++++++ 5 files changed, 232 insertions(+) create mode 100644 alembic/versions/20260220_a4ee5758d272_add_local_authority_impacts_table.py create mode 100644 src/policyengine_api/models/local_authority_impact.py diff --git a/alembic/versions/20260220_a4ee5758d272_add_local_authority_impacts_table.py b/alembic/versions/20260220_a4ee5758d272_add_local_authority_impacts_table.py new file mode 100644 index 0000000..01e8d64 --- /dev/null +++ b/alembic/versions/20260220_a4ee5758d272_add_local_authority_impacts_table.py @@ -0,0 +1,50 @@ +"""add_local_authority_impacts_table + +Revision ID: a4ee5758d272 +Revises: 83ceeb90bd74 +Create Date: 2026-02-20 22:20:00.037965 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +# revision identifiers, used by Alembic. +revision: str = 'a4ee5758d272' +down_revision: Union[str, Sequence[str], None] = '83ceeb90bd74' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('local_authority_impacts', + sa.Column('baseline_simulation_id', sa.Uuid(), nullable=False), + sa.Column('reform_simulation_id', sa.Uuid(), nullable=False), + sa.Column('report_id', sa.Uuid(), nullable=True), + sa.Column('local_authority_code', sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column('local_authority_name', sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column('x', sa.Integer(), nullable=False), + sa.Column('y', sa.Integer(), nullable=False), + sa.Column('average_household_income_change', sa.Float(), nullable=False), + sa.Column('relative_household_income_change', sa.Float(), nullable=False), + sa.Column('population', sa.Float(), nullable=False), + sa.Column('id', sa.Uuid(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(['baseline_simulation_id'], ['simulations.id'], ), + sa.ForeignKeyConstraint(['reform_simulation_id'], ['simulations.id'], ), + sa.ForeignKeyConstraint(['report_id'], ['reports.id'], ), + sa.PrimaryKeyConstraint('id') + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table('local_authority_impacts') + # ### end Alembic commands ### diff --git a/src/policyengine_api/api/analysis.py b/src/policyengine_api/api/analysis.py index b45a111..3df7fc5 100644 --- a/src/policyengine_api/api/analysis.py +++ b/src/policyengine_api/api/analysis.py @@ -32,6 +32,8 @@ CongressionalDistrictImpactRead, ConstituencyImpact, ConstituencyImpactRead, + LocalAuthorityImpact, + LocalAuthorityImpactRead, Dataset, DecileImpact, DecileImpactRead, @@ -153,6 +155,7 @@ class EconomicImpactResponse(BaseModel): detailed_budget: dict[str, dict[str, float | None]] | None = None congressional_district_impact: list[CongressionalDistrictImpactRead] | None = None constituency_impact: list[ConstituencyImpactRead] | None = None + local_authority_impact: list[LocalAuthorityImpactRead] | None = None def _get_model_version( @@ -299,6 +302,7 @@ def _build_response( detailed_budget = None district_impact_records = None constituency_impact_records = None + local_authority_impact_records = None if report.status == ReportStatus.COMPLETED: # Fetch decile impacts for this report @@ -503,6 +507,35 @@ def _build_response( for c in constituency_rows ] + # Fetch local authority impact records for this report + la_rows = session.exec( + select(LocalAuthorityImpact).where( + LocalAuthorityImpact.report_id == report.id + ) + ).all() + if la_rows: + local_authority_impact_records = [ + LocalAuthorityImpactRead( + id=la.id, + created_at=la.created_at, + baseline_simulation_id=la.baseline_simulation_id, + reform_simulation_id=la.reform_simulation_id, + report_id=la.report_id, + local_authority_code=la.local_authority_code, + local_authority_name=la.local_authority_name, + x=la.x, + y=la.y, + average_household_income_change=_safe_float( + la.average_household_income_change + ), + relative_household_income_change=_safe_float( + la.relative_household_income_change + ), + population=_safe_float(la.population), + ) + for la in la_rows + ] + region_info = None if region: region_info = RegionInfo( @@ -538,6 +571,7 @@ def _build_response( detailed_budget=detailed_budget, congressional_district_impact=district_impact_records, constituency_impact=constituency_impact_records, + local_authority_impact=local_authority_impact_records, ) @@ -990,6 +1024,50 @@ def build_dynamic(dynamic_id): except Exception: pass # Weight matrix not available, skip constituency impact + # Calculate local authority impact (UK only, requires weight matrix) + from policyengine.outputs.local_authority_impact import ( + compute_uk_local_authority_impacts, + ) + + try: + from policyengine_core.tools.google_cloud import download as gcs_download + + la_weight_matrix_path = gcs_download( + gcs_bucket="policyengine-uk-data-private", + gcs_key="local_authority_weights.h5", + ) + la_csv_path = gcs_download( + gcs_bucket="policyengine-uk-data-private", + gcs_key="local_authorities_2021.csv", + ) + la_impact = compute_uk_local_authority_impacts( + pe_baseline_sim, + pe_reform_sim, + weight_matrix_path=la_weight_matrix_path, + local_authority_csv_path=la_csv_path, + ) + if la_impact.local_authority_results: + for lr in la_impact.local_authority_results: + record = LocalAuthorityImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + local_authority_code=lr["local_authority_code"], + local_authority_name=lr["local_authority_name"], + x=lr["x"], + y=lr["y"], + average_household_income_change=lr[ + "average_household_income_change" + ], + relative_household_income_change=lr[ + "relative_household_income_change" + ], + population=lr["population"], + ) + session.add(record) + except Exception: + pass # Weight matrix not available, skip local authority impact + # Mark completed baseline_sim.status = SimulationStatus.COMPLETED baseline_sim.completed_at = datetime.now(timezone.utc) diff --git a/src/policyengine_api/modal_app.py b/src/policyengine_api/modal_app.py index 365eeee..7202293 100644 --- a/src/policyengine_api/modal_app.py +++ b/src/policyengine_api/modal_app.py @@ -1132,6 +1132,7 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None: from policyengine_api.models import ( BudgetSummary, ConstituencyImpact, + LocalAuthorityImpact, Dataset, DecileImpact, Inequality, @@ -1660,6 +1661,58 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None: except Exception: pass # Weight matrix not available, skip + # Calculate local authority impact + from policyengine.outputs.local_authority_impact import ( + compute_uk_local_authority_impacts, + ) + + try: + from policyengine_core.tools.google_cloud import ( + download as gcs_download_la, + ) + + la_weight_matrix_path = gcs_download_la( + gcs_bucket="policyengine-uk-data-private", + gcs_key="local_authority_weights.h5", + ) + la_csv_path = gcs_download_la( + gcs_bucket="policyengine-uk-data-private", + gcs_key="local_authorities_2021.csv", + ) + la_impact = compute_uk_local_authority_impacts( + pe_baseline_sim, + pe_reform_sim, + weight_matrix_path=la_weight_matrix_path, + local_authority_csv_path=la_csv_path, + ) + if la_impact.local_authority_results: + for lr in ( + la_impact.local_authority_results + ): + record = LocalAuthorityImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + local_authority_code=lr[ + "local_authority_code" + ], + local_authority_name=lr[ + "local_authority_name" + ], + x=lr["x"], + y=lr["y"], + average_household_income_change=lr[ + "average_household_income_change" + ], + relative_household_income_change=lr[ + "relative_household_income_change" + ], + population=lr["population"], + ) + session.add(record) + except Exception: + pass # Weight matrix not available, skip + # Mark simulations and report as completed baseline_sim.status = SimulationStatus.COMPLETED baseline_sim.completed_at = datetime.now(timezone.utc) diff --git a/src/policyengine_api/models/__init__.py b/src/policyengine_api/models/__init__.py index 9f4896f..c4fba8e 100644 --- a/src/policyengine_api/models/__init__.py +++ b/src/policyengine_api/models/__init__.py @@ -15,6 +15,11 @@ ConstituencyImpactCreate, ConstituencyImpactRead, ) +from .local_authority_impact import ( + LocalAuthorityImpact, + LocalAuthorityImpactCreate, + LocalAuthorityImpactRead, +) from .change_aggregate import ( ChangeAggregate, ChangeAggregateCreate, @@ -116,6 +121,9 @@ "ConstituencyImpact", "ConstituencyImpactCreate", "ConstituencyImpactRead", + "LocalAuthorityImpact", + "LocalAuthorityImpactCreate", + "LocalAuthorityImpactRead", "ChangeAggregate", "ChangeAggregateCreate", "ChangeAggregateRead", diff --git a/src/policyengine_api/models/local_authority_impact.py b/src/policyengine_api/models/local_authority_impact.py new file mode 100644 index 0000000..7ba5812 --- /dev/null +++ b/src/policyengine_api/models/local_authority_impact.py @@ -0,0 +1,43 @@ +"""UK local authority impact output model.""" + +from datetime import datetime, timezone +from uuid import UUID, uuid4 + +from sqlmodel import Field, SQLModel + + +class LocalAuthorityImpactBase(SQLModel): + """Base local authority impact fields.""" + + baseline_simulation_id: UUID = Field(foreign_key="simulations.id") + reform_simulation_id: UUID = Field(foreign_key="simulations.id") + report_id: UUID | None = Field(default=None, foreign_key="reports.id") + local_authority_code: str + local_authority_name: str + x: int + y: int + average_household_income_change: float + relative_household_income_change: float + population: float + + +class LocalAuthorityImpact(LocalAuthorityImpactBase, table=True): + """Local authority impact database model.""" + + __tablename__ = "local_authority_impacts" + + id: UUID = Field(default_factory=uuid4, primary_key=True) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + + +class LocalAuthorityImpactCreate(LocalAuthorityImpactBase): + """Schema for creating local authority impacts.""" + + pass + + +class LocalAuthorityImpactRead(LocalAuthorityImpactBase): + """Schema for reading local authority impacts.""" + + id: UUID + created_at: datetime From d29b76b35e46b67d70e2b8ffbfaa3f8650aa8d1c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 23:31:03 +0100 Subject: [PATCH 4/5] feat: Add wealth decile impact and migrate intra-decile to policyengine.py Migrate all intra-decile computation (UK+US, local+Modal) from the API's inline intra_decile.py helper to policyengine.py's new IntraDecileImpact output class. Add wealth decile impact and intra-wealth-decile impact for UK economy comparisons, using DecileImpact with decile_variable= "household_wealth_decile". Add decile_type column to intra_decile_impacts table to distinguish income vs wealth records. Co-Authored-By: Claude Opus 4.6 --- ...add_decile_type_to_intra_decile_impacts.py | 33 +++ src/policyengine_api/api/analysis.py | 199 ++++++++++++++---- src/policyengine_api/modal_app.py | 151 +++++++------ .../models/intra_decile_impact.py | 1 + 4 files changed, 278 insertions(+), 106 deletions(-) create mode 100644 alembic/versions/20260220_8d54837f0fcd_add_decile_type_to_intra_decile_impacts.py diff --git a/alembic/versions/20260220_8d54837f0fcd_add_decile_type_to_intra_decile_impacts.py b/alembic/versions/20260220_8d54837f0fcd_add_decile_type_to_intra_decile_impacts.py new file mode 100644 index 0000000..7b18b78 --- /dev/null +++ b/alembic/versions/20260220_8d54837f0fcd_add_decile_type_to_intra_decile_impacts.py @@ -0,0 +1,33 @@ +"""add_decile_type_to_intra_decile_impacts + +Revision ID: 8d54837f0fcd +Revises: a4ee5758d272 +Create Date: 2026-02-20 22:50:00.125733 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +# revision identifiers, used by Alembic. +revision: str = '8d54837f0fcd' +down_revision: Union[str, Sequence[str], None] = 'a4ee5758d272' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('intra_decile_impacts', sa.Column('decile_type', sqlmodel.sql.sqltypes.AutoString(), server_default='income', nullable=False)) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column('intra_decile_impacts', 'decile_type') + # ### end Alembic commands ### diff --git a/src/policyengine_api/api/analysis.py b/src/policyengine_api/api/analysis.py index 3df7fc5..76fb30f 100644 --- a/src/policyengine_api/api/analysis.py +++ b/src/policyengine_api/api/analysis.py @@ -156,6 +156,8 @@ class EconomicImpactResponse(BaseModel): congressional_district_impact: list[CongressionalDistrictImpactRead] | None = None constituency_impact: list[ConstituencyImpactRead] | None = None local_authority_impact: list[LocalAuthorityImpactRead] | None = None + wealth_decile: list[DecileImpactRead] | None = None + intra_wealth_decile: list[IntraDecileImpactRead] | None = None def _get_model_version( @@ -303,6 +305,8 @@ def _build_response( district_impact_records = None constituency_impact_records = None local_authority_impact_records = None + wealth_decile_records = None + intra_wealth_decile_records = None if report.status == ReportStatus.COMPLETED: # Fetch decile impacts for this report @@ -536,6 +540,62 @@ def _build_response( for la in la_rows ] + # Fetch wealth decile impact records (UK only) + wealth_decile_rows = session.exec( + select(DecileImpact).where( + DecileImpact.report_id == report.id, + DecileImpact.income_variable == "household_wealth_decile", + ) + ).all() + if wealth_decile_rows: + wealth_decile_records = [ + DecileImpactRead( + id=d.id, + created_at=d.created_at, + baseline_simulation_id=d.baseline_simulation_id, + reform_simulation_id=d.reform_simulation_id, + report_id=d.report_id, + income_variable=d.income_variable, + entity=d.entity, + decile=d.decile, + quantiles=d.quantiles, + baseline_mean=_safe_float(d.baseline_mean), + reform_mean=_safe_float(d.reform_mean), + absolute_change=_safe_float(d.absolute_change), + relative_change=_safe_float(d.relative_change), + count_better_off=_safe_float(d.count_better_off), + count_worse_off=_safe_float(d.count_worse_off), + count_no_change=_safe_float(d.count_no_change), + ) + for d in wealth_decile_rows + ] + + # Fetch intra-wealth-decile records (UK only) + intra_wealth_rows = session.exec( + select(IntraDecileImpact).where( + IntraDecileImpact.report_id == report.id, + IntraDecileImpact.decile_type == "wealth", + ) + ).all() + if intra_wealth_rows: + intra_wealth_decile_records = [ + IntraDecileImpactRead( + id=r.id, + created_at=r.created_at, + baseline_simulation_id=r.baseline_simulation_id, + reform_simulation_id=r.reform_simulation_id, + report_id=r.report_id, + decile_type=r.decile_type, + decile=r.decile, + lose_more_than_5pct=_safe_float(r.lose_more_than_5pct), + lose_less_than_5pct=_safe_float(r.lose_less_than_5pct), + no_change=_safe_float(r.no_change), + gain_less_than_5pct=_safe_float(r.gain_less_than_5pct), + gain_more_than_5pct=_safe_float(r.gain_more_than_5pct), + ) + for r in intra_wealth_rows + ] + region_info = None if region: region_info = RegionInfo( @@ -572,6 +632,8 @@ def _build_response( congressional_district_impact=district_impact_records, constituency_impact=constituency_impact_records, local_authority_impact=local_authority_impact_records, + wealth_decile=wealth_decile_records, + intra_wealth_decile=intra_wealth_decile_records, ) @@ -950,33 +1012,27 @@ def build_dynamic(dynamic_id): session.add(budget_record) # Calculate intra-decile impact (5-category income change distribution) - from policyengine_api.api.intra_decile import compute_intra_decile - - baseline_hh_data = { - k: pe_baseline_sim.output_dataset.data.household[k].values - for k in [ - "household_net_income", - "household_weight", - "household_count_people", - "household_income_decile", - ] - } - reform_hh_data = { - k: pe_reform_sim.output_dataset.data.household[k].values - for k in [ - "household_net_income", - "household_weight", - "household_count_people", - "household_income_decile", - ] - } - intra_decile_rows = compute_intra_decile(baseline_hh_data, reform_hh_data) - for row in intra_decile_rows: + from policyengine.outputs.intra_decile_impact import ( + compute_intra_decile_impacts as pe_compute_intra_decile, + ) + + intra_decile_results = pe_compute_intra_decile( + baseline_simulation=pe_baseline_sim, + reform_simulation=pe_reform_sim, + income_variable="household_net_income", + entity="household", + ) + for r in intra_decile_results.outputs: record = IntraDecileImpact( baseline_simulation_id=baseline_sim.id, reform_simulation_id=reform_sim.id, report_id=report.id, - **row, + decile=r.decile, + lose_more_than_5pct=r.lose_more_than_5pct, + lose_less_than_5pct=r.lose_less_than_5pct, + no_change=r.no_change, + gain_less_than_5pct=r.gain_less_than_5pct, + gain_more_than_5pct=r.gain_more_than_5pct, ) session.add(record) @@ -1068,6 +1124,63 @@ def build_dynamic(dynamic_id): except Exception: pass # Weight matrix not available, skip local authority impact + # Calculate wealth decile impact (UK only) + try: + from policyengine.outputs.decile_impact import ( + DecileImpact as PEDecileImpact, + ) + + PEDecileImpact.model_rebuild(_types_namespace={"Simulation": PESimulation}) + for decile_num in range(1, 11): + wealth_di = PEDecileImpact( + baseline_simulation=pe_baseline_sim, + reform_simulation=pe_reform_sim, + income_variable="household_net_income", + decile_variable="household_wealth_decile", + entity="household", + decile=decile_num, + ) + wealth_di.run() + record = DecileImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + income_variable="household_wealth_decile", + entity="household", + decile=decile_num, + quantiles=10, + baseline_mean=wealth_di.baseline_mean, + reform_mean=wealth_di.reform_mean, + absolute_change=wealth_di.absolute_change, + relative_change=wealth_di.relative_change, + ) + session.add(record) + + # Calculate intra-wealth-decile impact + intra_wealth_results = pe_compute_intra_decile( + baseline_simulation=pe_baseline_sim, + reform_simulation=pe_reform_sim, + income_variable="household_net_income", + decile_variable="household_wealth_decile", + entity="household", + ) + for r in intra_wealth_results.outputs: + record = IntraDecileImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + decile_type="wealth", + decile=r.decile, + lose_more_than_5pct=r.lose_more_than_5pct, + lose_less_than_5pct=r.lose_less_than_5pct, + no_change=r.no_change, + gain_less_than_5pct=r.gain_less_than_5pct, + gain_more_than_5pct=r.gain_more_than_5pct, + ) + session.add(record) + except (KeyError, Exception): + pass # household_wealth_decile not available (US), skip + # Mark completed baseline_sim.status = SimulationStatus.COMPLETED baseline_sim.completed_at = datetime.now(timezone.utc) @@ -1448,33 +1561,27 @@ def build_dynamic(dynamic_id): session.add(budget_record) # Calculate intra-decile impact (5-category income change distribution) - from policyengine_api.api.intra_decile import compute_intra_decile - - baseline_hh_data = { - k: pe_baseline_sim.output_dataset.data.household[k].values - for k in [ - "household_net_income", - "household_weight", - "household_count_people", - "household_income_decile", - ] - } - reform_hh_data = { - k: pe_reform_sim.output_dataset.data.household[k].values - for k in [ - "household_net_income", - "household_weight", - "household_count_people", - "household_income_decile", - ] - } - intra_decile_rows = compute_intra_decile(baseline_hh_data, reform_hh_data) - for row in intra_decile_rows: + from policyengine.outputs.intra_decile_impact import ( + compute_intra_decile_impacts as pe_compute_intra_decile_us, + ) + + intra_decile_results_us = pe_compute_intra_decile_us( + baseline_simulation=pe_baseline_sim, + reform_simulation=pe_reform_sim, + income_variable="household_net_income", + entity="household", + ) + for r in intra_decile_results_us.outputs: record = IntraDecileImpact( baseline_simulation_id=baseline_sim.id, reform_simulation_id=reform_sim.id, report_id=report.id, - **row, + decile=r.decile, + lose_more_than_5pct=r.lose_more_than_5pct, + lose_less_than_5pct=r.lose_less_than_5pct, + no_change=r.no_change, + gain_less_than_5pct=r.gain_less_than_5pct, + gain_more_than_5pct=r.gain_more_than_5pct, ) session.add(record) diff --git a/src/policyengine_api/modal_app.py b/src/policyengine_api/modal_app.py index 7202293..16130ac 100644 --- a/src/policyengine_api/modal_app.py +++ b/src/policyengine_api/modal_app.py @@ -1569,41 +1569,27 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None: session.add(budget_record) # Calculate intra-decile impact - from policyengine_api.api.intra_decile import ( - compute_intra_decile, - ) - - baseline_hh_data = { - k: pe_baseline_sim.output_dataset.data.household[ - k - ].values - for k in [ - "household_net_income", - "household_weight", - "household_count_people", - "household_income_decile", - ] - } - reform_hh_data = { - k: pe_reform_sim.output_dataset.data.household[ - k - ].values - for k in [ - "household_net_income", - "household_weight", - "household_count_people", - "household_income_decile", - ] - } - intra_decile_rows = compute_intra_decile( - baseline_hh_data, reform_hh_data + from policyengine.outputs.intra_decile_impact import ( + compute_intra_decile_impacts as pe_compute_intra_decile, + ) + + intra_decile_results = pe_compute_intra_decile( + baseline_simulation=pe_baseline_sim, + reform_simulation=pe_reform_sim, + income_variable="household_net_income", + entity="household", ) - for row in intra_decile_rows: + for r in intra_decile_results.outputs: record = IntraDecileImpact( baseline_simulation_id=baseline_sim.id, reform_simulation_id=reform_sim.id, report_id=report.id, - **row, + decile=r.decile, + lose_more_than_5pct=r.lose_more_than_5pct, + lose_less_than_5pct=r.lose_less_than_5pct, + no_change=r.no_change, + gain_less_than_5pct=r.gain_less_than_5pct, + gain_more_than_5pct=r.gain_more_than_5pct, ) session.add(record) @@ -1713,6 +1699,65 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None: except Exception: pass # Weight matrix not available, skip + # Calculate wealth decile impact (UK only) + try: + from policyengine.outputs.decile_impact import ( + DecileImpact as PEDecileImpact, + ) + + PEDecileImpact.model_rebuild( + _types_namespace={"Simulation": PESimulation} + ) + for decile_num in range(1, 11): + wealth_di = PEDecileImpact( + baseline_simulation=pe_baseline_sim, + reform_simulation=pe_reform_sim, + income_variable="household_net_income", + decile_variable="household_wealth_decile", + entity="household", + decile=decile_num, + ) + wealth_di.run() + record = DecileImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + income_variable="household_wealth_decile", + entity="household", + decile=decile_num, + quantiles=10, + baseline_mean=wealth_di.baseline_mean, + reform_mean=wealth_di.reform_mean, + absolute_change=wealth_di.absolute_change, + relative_change=wealth_di.relative_change, + ) + session.add(record) + + # Calculate intra-wealth-decile impact + intra_wealth_results = pe_compute_intra_decile( + baseline_simulation=pe_baseline_sim, + reform_simulation=pe_reform_sim, + income_variable="household_net_income", + decile_variable="household_wealth_decile", + entity="household", + ) + for r in intra_wealth_results.outputs: + record = IntraDecileImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + decile_type="wealth", + decile=r.decile, + lose_more_than_5pct=r.lose_more_than_5pct, + lose_less_than_5pct=r.lose_less_than_5pct, + no_change=r.no_change, + gain_less_than_5pct=r.gain_less_than_5pct, + gain_more_than_5pct=r.gain_more_than_5pct, + ) + session.add(record) + except (KeyError, Exception): + pass # household_wealth_decile not available, skip + # Mark simulations and report as completed baseline_sim.status = SimulationStatus.COMPLETED baseline_sim.completed_at = datetime.now(timezone.utc) @@ -2206,41 +2251,27 @@ def economy_comparison_us(job_id: str, traceparent: str | None = None) -> None: session.add(budget_record) # Calculate intra-decile impact - from policyengine_api.api.intra_decile import ( - compute_intra_decile, - ) - - baseline_hh_data = { - k: pe_baseline_sim.output_dataset.data.household[ - k - ].values - for k in [ - "household_net_income", - "household_weight", - "household_count_people", - "household_income_decile", - ] - } - reform_hh_data = { - k: pe_reform_sim.output_dataset.data.household[ - k - ].values - for k in [ - "household_net_income", - "household_weight", - "household_count_people", - "household_income_decile", - ] - } - intra_decile_rows = compute_intra_decile( - baseline_hh_data, reform_hh_data + from policyengine.outputs.intra_decile_impact import ( + compute_intra_decile_impacts as pe_compute_intra_decile_us, + ) + + intra_decile_results_us = pe_compute_intra_decile_us( + baseline_simulation=pe_baseline_sim, + reform_simulation=pe_reform_sim, + income_variable="household_net_income", + entity="household", ) - for row in intra_decile_rows: + for r in intra_decile_results_us.outputs: record = IntraDecileImpact( baseline_simulation_id=baseline_sim.id, reform_simulation_id=reform_sim.id, report_id=report.id, - **row, + decile=r.decile, + lose_more_than_5pct=r.lose_more_than_5pct, + lose_less_than_5pct=r.lose_less_than_5pct, + no_change=r.no_change, + gain_less_than_5pct=r.gain_less_than_5pct, + gain_more_than_5pct=r.gain_more_than_5pct, ) session.add(record) diff --git a/src/policyengine_api/models/intra_decile_impact.py b/src/policyengine_api/models/intra_decile_impact.py index 42c8173..f3cfbde 100644 --- a/src/policyengine_api/models/intra_decile_impact.py +++ b/src/policyengine_api/models/intra_decile_impact.py @@ -27,6 +27,7 @@ class IntraDecileImpactBase(SQLModel): baseline_simulation_id: UUID = Field(foreign_key="simulations.id") reform_simulation_id: UUID = Field(foreign_key="simulations.id") report_id: UUID | None = Field(default=None, foreign_key="reports.id") + decile_type: str = Field(default="income") # "income" or "wealth" decile: int # 1-10 for individual deciles, 0 for overall average lose_more_than_5pct: float | None = None lose_less_than_5pct: float | None = None From 3427073dc6758484556180b82988481d46d1cbc5 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 21 Feb 2026 00:13:45 +0100 Subject: [PATCH 5/5] test: Add tests for Phase 3 economic impact response fields Extend fixtures with factory functions for congressional district, constituency, local authority, wealth decile, and intra-wealth-decile records. Add test classes verifying _build_response() populates all new fields correctly. Co-Authored-By: Claude Opus 4.6 --- .../fixtures_economic_impact_response.py | 170 +++++++++++++++ tests/test_economic_impact_response.py | 200 ++++++++++++++++++ 2 files changed, 370 insertions(+) diff --git a/test_fixtures/fixtures_economic_impact_response.py b/test_fixtures/fixtures_economic_impact_response.py index 4d78371..66f3ce6 100644 --- a/test_fixtures/fixtures_economic_impact_response.py +++ b/test_fixtures/fixtures_economic_impact_response.py @@ -11,10 +11,13 @@ from policyengine_api.models import ( BudgetSummary, + CongressionalDistrictImpact, + ConstituencyImpact, Dataset, DecileImpact, Inequality, IntraDecileImpact, + LocalAuthorityImpact, Poverty, ProgramStatistics, Report, @@ -297,6 +300,168 @@ def add_program_statistics_records( return records +def add_congressional_district_records( + session: Session, + report: Report, + baseline_sim: Simulation, + reform_sim: Simulation, +) -> list[CongressionalDistrictImpact]: + """Add congressional district impact records.""" + records = [] + districts = [ + {"district_geoid": 101, "state_fips": 1, "district_number": 1}, + {"district_geoid": 602, "state_fips": 6, "district_number": 2}, + ] + for d in districts: + rec = CongressionalDistrictImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + district_geoid=d["district_geoid"], + state_fips=d["state_fips"], + district_number=d["district_number"], + average_household_income_change=500.0, + relative_household_income_change=0.01, + population=100000.0, + ) + session.add(rec) + records.append(rec) + session.commit() + return records + + +SAMPLE_DISTRICT_COUNT = 2 + + +def add_constituency_records( + session: Session, + report: Report, + baseline_sim: Simulation, + reform_sim: Simulation, +) -> list[ConstituencyImpact]: + """Add UK constituency impact records.""" + records = [] + constituencies = [ + {"code": "E14000530", "name": "Birmingham, Ladywood", "x": 410, "y": 290}, + {"code": "E14000639", "name": "Cities of London and Westminster", "x": 530, "y": 180}, + ] + for c in constituencies: + rec = ConstituencyImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + constituency_code=c["code"], + constituency_name=c["name"], + x=c["x"], + y=c["y"], + average_household_income_change=300.0, + relative_household_income_change=0.008, + population=80000.0, + ) + session.add(rec) + records.append(rec) + session.commit() + return records + + +SAMPLE_CONSTITUENCY_COUNT = 2 + + +def add_local_authority_records( + session: Session, + report: Report, + baseline_sim: Simulation, + reform_sim: Simulation, +) -> list[LocalAuthorityImpact]: + """Add UK local authority impact records.""" + records = [] + las = [ + {"code": "E09000001", "name": "City of London", "x": 532, "y": 181}, + {"code": "E09000002", "name": "Barking and Dagenham", "x": 549, "y": 186}, + ] + for la in las: + rec = LocalAuthorityImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + local_authority_code=la["code"], + local_authority_name=la["name"], + x=la["x"], + y=la["y"], + average_household_income_change=400.0, + relative_household_income_change=0.012, + population=50000.0, + ) + session.add(rec) + records.append(rec) + session.commit() + return records + + +SAMPLE_LA_COUNT = 2 + + +def add_wealth_decile_records( + session: Session, + report: Report, + baseline_sim: Simulation, + reform_sim: Simulation, +) -> list[DecileImpact]: + """Add 10 wealth decile impact records (income_variable=household_wealth_decile).""" + records = [] + for decile_num in range(1, 11): + rec = DecileImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + income_variable="household_wealth_decile", + entity="household", + decile=decile_num, + quantiles=10, + baseline_mean=float(10000 * decile_num), + reform_mean=float(10000 * decile_num + 500), + absolute_change=500.0, + relative_change=500.0 / (10000 * decile_num), + ) + session.add(rec) + records.append(rec) + session.commit() + return records + + +SAMPLE_WEALTH_DECILE_COUNT = 10 + + +def add_intra_wealth_decile_records( + session: Session, + report: Report, + baseline_sim: Simulation, + reform_sim: Simulation, +) -> list[IntraDecileImpact]: + """Add 11 intra-wealth-decile records (decile_type='wealth').""" + records = [] + for decile_num in list(range(1, 11)) + [0]: + rec = IntraDecileImpact( + baseline_simulation_id=baseline_sim.id, + reform_simulation_id=reform_sim.id, + report_id=report.id, + decile_type="wealth", + decile=decile_num, + lose_more_than_5pct=0.0, + lose_less_than_5pct=0.1, + no_change=0.5, + gain_less_than_5pct=0.3, + gain_more_than_5pct=0.1, + ) + session.add(rec) + records.append(rec) + session.commit() + return records + + +SAMPLE_INTRA_WEALTH_DECILE_COUNT = 11 + + # --------------------------------------------------------------------------- # Composite: fully populated report # --------------------------------------------------------------------------- @@ -313,4 +478,9 @@ def create_fully_populated_report( add_budget_summary_records(session, report, baseline_sim, reform_sim) add_intra_decile_records(session, report, baseline_sim, reform_sim) add_program_statistics_records(session, report, baseline_sim, reform_sim) + add_congressional_district_records(session, report, baseline_sim, reform_sim) + add_constituency_records(session, report, baseline_sim, reform_sim) + add_local_authority_records(session, report, baseline_sim, reform_sim) + add_wealth_decile_records(session, report, baseline_sim, reform_sim) + add_intra_wealth_decile_records(session, report, baseline_sim, reform_sim) return report, baseline_sim, reform_sim diff --git a/tests/test_economic_impact_response.py b/tests/test_economic_impact_response.py index 9401224..2c24c16 100644 --- a/tests/test_economic_impact_response.py +++ b/tests/test_economic_impact_response.py @@ -12,19 +12,29 @@ BUDGET_VARIABLES_UK, INTRA_DECILE_DECILE_COUNT, SAMPLE_BOTTOM_50_SHARE, + SAMPLE_CONSTITUENCY_COUNT, + SAMPLE_DISTRICT_COUNT, SAMPLE_GINI, SAMPLE_INEQUALITY_INCOME_VAR, + SAMPLE_INTRA_WEALTH_DECILE_COUNT, + SAMPLE_LA_COUNT, SAMPLE_POVERTY_TYPES, SAMPLE_TOP_1_SHARE, SAMPLE_TOP_10_SHARE, + SAMPLE_WEALTH_DECILE_COUNT, UK_PROGRAM_COUNT, UK_PROGRAMS, add_budget_summary_records, + add_congressional_district_records, + add_constituency_records, add_inequality_records, add_intra_decile_records, + add_intra_wealth_decile_records, + add_local_authority_records, add_poverty_by_age_records, add_poverty_records, add_program_statistics_records, + add_wealth_decile_records, create_fully_populated_report, create_report_with_simulations, ) @@ -462,6 +472,11 @@ def test__given_fully_populated_report__then_all_fields_present(self, session): assert response.intra_decile is not None assert response.program_statistics is not None assert response.detailed_budget is not None + assert response.congressional_district_impact is not None + assert response.constituency_impact is not None + assert response.local_authority_impact is not None + assert response.wealth_decile is not None + assert response.intra_wealth_decile is not None def test__given_fully_populated_report__then_report_id_matches(self, session): # Given @@ -483,3 +498,188 @@ def test__given_fully_populated_report__then_simulation_ids_match(self, session) # Then assert response.baseline_simulation.id == baseline_sim.id assert response.reform_simulation.id == reform_sim.id + + +# --------------------------------------------------------------------------- +# _build_response — congressional_district_impact +# --------------------------------------------------------------------------- + + +class TestBuildResponseCongressionalDistrict: + """Tests for congressional_district_impact in _build_response output.""" + + def test__given_district_records__then_correct_count(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_congressional_district_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.congressional_district_impact is not None + assert len(response.congressional_district_impact) == SAMPLE_DISTRICT_COUNT + + def test__given_no_district_records__then_field_is_none(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.congressional_district_impact is None + + def test__given_district_records__then_geoid_fields_populated(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_congressional_district_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + for d in response.congressional_district_impact: + assert d.district_geoid > 0 + assert d.state_fips >= 0 + assert d.population > 0 + + +# --------------------------------------------------------------------------- +# _build_response — constituency_impact +# --------------------------------------------------------------------------- + + +class TestBuildResponseConstituency: + """Tests for constituency_impact in _build_response output.""" + + def test__given_constituency_records__then_correct_count(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_constituency_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.constituency_impact is not None + assert len(response.constituency_impact) == SAMPLE_CONSTITUENCY_COUNT + + def test__given_no_constituency_records__then_field_is_none(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.constituency_impact is None + + def test__given_constituency_records__then_code_and_name_populated(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_constituency_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + for c in response.constituency_impact: + assert c.constituency_code is not None + assert c.constituency_name is not None + assert c.population > 0 + + +# --------------------------------------------------------------------------- +# _build_response — local_authority_impact +# --------------------------------------------------------------------------- + + +class TestBuildResponseLocalAuthority: + """Tests for local_authority_impact in _build_response output.""" + + def test__given_la_records__then_correct_count(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_local_authority_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.local_authority_impact is not None + assert len(response.local_authority_impact) == SAMPLE_LA_COUNT + + def test__given_no_la_records__then_field_is_none(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.local_authority_impact is None + + def test__given_la_records__then_code_and_name_populated(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_local_authority_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + for la in response.local_authority_impact: + assert la.local_authority_code is not None + assert la.local_authority_name is not None + assert la.population > 0 + + +# --------------------------------------------------------------------------- +# _build_response — wealth_decile +# --------------------------------------------------------------------------- + + +class TestBuildResponseWealthDecile: + """Tests for wealth_decile in _build_response output.""" + + def test__given_wealth_decile_records__then_correct_count(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_wealth_decile_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.wealth_decile is not None + assert len(response.wealth_decile) == SAMPLE_WEALTH_DECILE_COUNT + + def test__given_no_wealth_decile_records__then_field_is_none(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.wealth_decile is None + + def test__given_wealth_decile_records__then_income_variable_is_wealth(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_wealth_decile_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + for d in response.wealth_decile: + assert d.income_variable == "household_wealth_decile" + + +# --------------------------------------------------------------------------- +# _build_response — intra_wealth_decile +# --------------------------------------------------------------------------- + + +class TestBuildResponseIntraWealthDecile: + """Tests for intra_wealth_decile in _build_response output.""" + + def test__given_intra_wealth_records__then_correct_count(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_intra_wealth_decile_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.intra_wealth_decile is not None + assert len(response.intra_wealth_decile) == SAMPLE_INTRA_WEALTH_DECILE_COUNT + + def test__given_no_intra_wealth_records__then_field_is_none(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + + response = _build_response(report, baseline_sim, reform_sim, session) + + assert response.intra_wealth_decile is None + + def test__given_intra_wealth_records__then_decile_type_is_wealth(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_intra_wealth_decile_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + for r in response.intra_wealth_decile: + assert r.decile_type == "wealth" + + def test__given_intra_wealth_records__then_overall_row_present(self, session): + report, baseline_sim, reform_sim = create_report_with_simulations(session) + add_intra_wealth_decile_records(session, report, baseline_sim, reform_sim) + + response = _build_response(report, baseline_sim, reform_sim, session) + + decile_numbers = {r.decile for r in response.intra_wealth_decile} + assert 0 in decile_numbers