Add budget charts and data

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.py

@@ -6,12 +6,12 @@
 
 
 def parliamentary_constituencies(
-        simulation: Simulation, 
-        chart: bool = False, 
-        variable: str = None,
-        aggregator: str = None,
-        relative: bool = None,
-    ) -> dict:
+    simulation: Simulation,
+    chart: bool = False,
+    variable: str = None,
+    aggregator: str = None,
+    relative: bool = None,
+) -> dict:
     if not simulation.options.get("include_constituencies"):
         return {}
 
@@ -45,9 +45,11 @@ def parliamentary_constituencies(
     return result
 
 
-
 def heatmap(
     simulation: Simulation,
+    variable: str = None,
+    aggregator: str = None,
+    relative: bool = None,
 ) -> dict:
     if not simulation.options.get("include_constituencies"):
         return {}