Merge pull request #47 from PolicyEngine/dev

Add budget charts

CHANGELOG.md

@@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.1.0] - 2024-12-02 12:12:59
+
+### Added
+
+- Budget chart.
+- Budget window chart and data.
+
 ## [2.0.0] - 2024-12-01 12:38:17
 
 ### Added
@@ -25,5 +32,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 
 
+[2.1.0]: https://github.com/PolicyEngine/policyengine.py/compare/2.0.0...2.1.0
 [2.0.0]: https://github.com/PolicyEngine/policyengine.py/compare/1.0.1...2.0.0
 [1.0.1]: https://github.com/PolicyEngine/policyengine.py/compare/1.0.0...1.0.1

changelog.yaml

@@ -14,3 +14,9 @@
     - Major version bump to succeed old versions of the package under a different
       repo.
   date: 2024-12-01 12:38:17
+- bump: minor
+  changes:
+    added:
+    - Budget chart.
+    - Budget window chart and data.
+  date: 2024-12-02 12:12:59

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

@@ -0,0 +1,109 @@
+from policyengine import Simulation
+import plotly.graph_objects as go
+import plotly.express as px
+import pandas as pd
+from policyengine.utils.charts import *
+
+
+def general(simulation: Simulation, chart: bool = False):
+    """Calculate the budgetary impact of the given simulation.
+
+    Args:
+        simulation (Simulation): The simulation for which the revenue impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the budgetary impact details with the following keys:
+            - budgetary_impact (float): The overall budgetary impact.
+            - tax_revenue_impact (float): The impact on tax revenue.
+            - state_tax_revenue_impact (float): The impact on state tax revenue.
+            - benefit_spending_impact (float): The impact on benefit spending.
+            - households (int): The number of households.
+            - baseline_net_income (float): The total net income in the baseline scenario.
+    """
+    baseline = simulation.calculate("macro/baseline")
+    reform = simulation.calculate("macro/reform")
+
+    tax_revenue_impact = (
+        reform["gov"]["balance"]["total_tax"]
+        - baseline["gov"]["balance"]["total_tax"]
+    )
+    state_tax_revenue_impact = (
+        reform["gov"]["balance"]["total_state_tax"]
+        - baseline["gov"]["balance"]["total_state_tax"]
+    )
+    benefit_spending_impact = (
+        reform["gov"]["balance"]["total_spending"]
+        - baseline["gov"]["balance"]["total_spending"]
+    )
+    budgetary_impact = tax_revenue_impact - benefit_spending_impact
+    households = sum(baseline["household"]["demographics"]["household_weight"])
+    baseline_net_income = baseline["household"]["finance"]["total_net_income"]
+    result = dict(
+        budgetary_impact=budgetary_impact,
+        tax_revenue_impact=tax_revenue_impact,
+        state_tax_revenue_impact=state_tax_revenue_impact,
+        benefit_spending_impact=benefit_spending_impact,
+        households=households,
+        baseline_net_income=baseline_net_income,
+    )
+    if chart:
+        return budget_chart(simulation, result)
+    else:
+        return result
+
+
+def budget_chart(simulation: Simulation, data: dict) -> go.Figure:
+    if simulation.country == "uk":
+        x = ["Tax revenues", "Benefit spending", "Budgetary impact"]
+        y = [
+            data["tax_revenue_impact"],
+            -data["benefit_spending_impact"],
+            data["budgetary_impact"],
+        ]
+    else:
+        x = [
+            "Federal tax revenues",
+            "State tax revenues",
+            "Benefit spending",
+            "Budgetary impact",
+        ]
+        y = [
+            data["tax_revenue_impact"] - data["state_tax_revenue_impact"],
+            data["state_tax_revenue_impact"],
+            -data["benefit_spending_impact"],
+            data["budgetary_impact"],
+        ]
+    fig = go.Figure(
+        data=[
+            go.Waterfall(
+                x=x,
+                y=[i / 1e9 for i in y],
+                orientation="v",
+                measure=["relative"] * (len(x) - 1) + ["total"],
+                text=[
+                    (
+                        "+" + str(round(val / 1e9, 1))
+                        if val > 0
+                        else str(round(val / 1e9, 1))
+                    )
+                    for val in y
+                ],
+                textposition="inside",
+                increasing={"marker": {"color": BLUE}},
+                decreasing={"marker": {"color": DARK_GRAY}},
+                totals={"marker": {"color": BLUE if y[-1] > 0 else DARK_GRAY}},
+                connector={
+                    "line": {"color": DARK_GRAY, "width": 2, "dash": "dot"}
+                },
+            )
+        ]
+    )
+
+    fig.update_layout(
+        title="Budgetary impact by government revenue and spending",
+        xaxis_title="",
+        yaxis_title="Budgetary impact (£ billions)",
+        yaxis_tickformat=",.0f",
+    )
+
+    return format_fig(fig)

policyengine/outputs/macro/comparison/detailed_budget.py → policyengine/outputs/macro/comparison/budget/programs.py

@@ -1,7 +1,7 @@
 from policyengine import Simulation
 
 
-def detailed_budget(simulation: Simulation):
+def programs(simulation: Simulation):
     """Calculate the detailed budgetary impact of the given simulation.
 
     Args:

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

@@ -0,0 +1,68 @@
+from policyengine import Simulation
+import plotly.graph_objects as go
+import plotly.express as px
+import pandas as pd
+from policyengine.utils.charts import *
+
+
+def window(
+    simulation: Simulation,
+    chart: bool = False,
+    federal_only: bool = False,
+    count_years: int = 1,
+):
+    if count_years == 1:
+        kwargs = {}
+    else:
+        kwargs = {"count_years": count_years}
+    baseline = simulation.calculate(
+        "macro/baseline/gov/budget_window", **kwargs
+    )
+    reform = simulation.calculate("macro/reform/gov/budget_window", **kwargs)
+    total_budget_effect = [
+        (y - x) / 1e9
+        for x, y in zip(baseline["total_budget"], reform["total_budget"])
+    ]
+    federal_budget_effect = [
+        (y - x) / 1e9
+        for x, y in zip(
+            baseline["total_federal_budget"], reform["total_federal_budget"]
+        )
+    ]
+
+    result = dict(
+        total_budget=total_budget_effect,
+        federal_budget=federal_budget_effect,
+    )
+
+    if chart:
+        return budget_window_chart(
+            result["federal_budget"]
+            if federal_only
+            else result["total_budget"]
+        )
+    else:
+        return result
+
+
+def budget_window_chart(budget_effect) -> go.Figure:
+    fig = go.Figure(
+        data=[
+            go.Bar(
+                y=budget_effect,
+                x=list(map(str, range(2025, 2025 + 10))),
+                marker=dict(
+                    color=[BLUE if y > 0 else DARK_GRAY for y in budget_effect]
+                ),
+                text=[
+                    f"{'+' if y >= 0 else ''}{y:.1f}" for y in budget_effect
+                ],
+            )
+        ]
+    ).update_layout(
+        title="Budgetary impact by year",
+        xaxis_title="Year",
+        yaxis_title="Budgetary impact (£ billions)",
+    )
+
+    return format_fig(fig)

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

@@ -5,6 +5,46 @@
 from policyengine.utils.charts import *
 
 
+def parliamentary_constituencies(
+    simulation: Simulation,
+    chart: bool = False,
+    variable: str = None,
+    aggregator: str = None,
+    relative: 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,
@@ -38,9 +78,6 @@ def heatmap(
         version=None,
     )
     constituency_names = pd.read_csv(constituency_names_file_path)
-    hex_map_locations = pd.read_csv(
-        "/Users/nikhilwoodruff/uk-local-area-calibration/policyengine_uk_local_areas/hex_map/hex_map_2024.csv"
-    ).set_index("code")
 
     if variable is None:
         variable = "household_net_income"
@@ -60,12 +97,6 @@ def heatmap(
                 - constituency_baseline[constituency][variable]
             )
 
-    constituency_names["x"] = hex_map_locations.loc[
-        constituency_names["code"]
-    ]["x"].values
-    constituency_names["y"] = hex_map_locations.loc[
-        constituency_names["code"]
-    ]["y"].values
     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