Merge pull request #44 from PolicyEngine/api-parity

Add documentation and chart functions

policyengine/constants.py

@@ -0,0 +1,20 @@
+"""Mainly simulation options and parameters."""
+
+# Datasets
+
+ENHANCED_FRS = "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5"
+FRS = "hf://policyengine/policyengine-uk-data/frs_2022_23.h5"
+
+ENHANCED_CPS = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"
+CPS = "hf://policyengine/policyengine-us-data/cps_2023.h5"
+POOLED_CPS = "hf://policyengine/policyengine-us-data/pooled_3_year_cps_2023.h5"
+
+DATASETS = {
+    "uk": {"frs": FRS, "enhanced_frs": ENHANCED_FRS},
+    "us": {"cps": CPS, "enhanced_cps": ENHANCED_CPS, "pooled_cps": POOLED_CPS},
+}
+
+DEFAULT_DATASETS = {
+    "uk": ENHANCED_FRS,
+    "us": CPS,
+}

policyengine/outputs/household/single/net_income.py

@@ -1,2 +1,2 @@
 def net_income(simulation):
-    return simulation.baseline.calculate("household_net_income").sum()
+    return simulation.selected.calculate("household_net_income").sum()

policyengine/outputs/macro/comparison/budget.py

@@ -8,7 +8,13 @@ def budget(simulation: Simulation):
         simulation (Simulation): The simulation for which the revenue impact is to be calculated.
 
     Returns:
-        float: The revenue impact of the simulation.
+        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")

policyengine/outputs/macro/comparison/decile/income.py

@@ -3,6 +3,16 @@
 
 
 def income(simulation: Simulation):
+    """Calculate the impact of the reform on income deciles.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the impact details with the following keys:
+            - relative (dict): A dictionary with keys representing deciles and values as relative income changes.
+            - average (dict): A dictionary with keys representing deciles and values as average income changes.
+    """
     baseline = simulation.calculate("macro/baseline")
     reform = simulation.calculate("macro/reform")
 

policyengine/outputs/macro/comparison/decile/wealth.py

@@ -3,6 +3,16 @@
 
 
 def wealth(simulation: Simulation):
+    """Calculate the impact of the reform on wealth deciles.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the impact details with the following keys:
+            - relative (dict): A dictionary with keys representing deciles and values as relative income changes.
+            - average (dict): A dictionary with keys representing deciles and values as average income changes.
+    """
     if simulation.country != "uk":
         return {}
 

policyengine/outputs/macro/comparison/detailed_budget.py

@@ -2,6 +2,17 @@
 
 
 def detailed_budget(simulation: Simulation):
+    """Calculate the detailed budgetary impact of the given simulation.
+
+    Args:
+        simulation (Simulation): The simulation for which the budgetary impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the detailed budgetary impact for each program with the following keys:
+            - baseline (float): The baseline budgetary impact of the program.
+            - reform (float): The reform budgetary impact of the program.
+            - difference (float): The difference between the reform and baseline budgetary impacts.
+    """
     baseline = simulation.calculate("macro/baseline")
     reform = simulation.calculate("macro/reform")
     result = {}

policyengine/outputs/macro/comparison/inequality.py

@@ -2,6 +2,17 @@
 
 
 def inequality(simulation: Simulation):
+    """Calculate the impact of the reform on inequality.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the inequality impact details with the following keys:
+            - gini (dict): A dictionary with baseline and reform Gini coefficients.
+            - top_10_pct_share (dict): A dictionary with baseline and reform top 10% income share.
+            - top_1_pct_share (dict): A dictionary with baseline and reform top 1% income share.
+    """
     baseline = simulation.calculate("macro/baseline")["household"][
         "inequality"
     ]

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

@@ -0,0 +1,27 @@
+from policyengine import Simulation
+
+
+def data(simulation: Simulation) -> dict:
+    if not simulation.options.get("include_constituencies"):
+        return {}
+
+    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

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

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

policyengine/outputs/macro/comparison/poverty/age.py

@@ -3,6 +3,16 @@
 
 
 def age(simulation: Simulation):
+    """Calculate the impact of the reform on poverty by age.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the poverty and deep poverty impact details with the following keys:
+            - poverty (dict): A dictionary with keys representing age groups and values as dictionaries with baseline and reform poverty rates.
+            - deep_poverty (dict): A dictionary with keys representing age groups and values as dictionaries with baseline and reform deep poverty rates.
+    """
     baseline = simulation.calculate("macro/baseline")["household"]["finance"]
     reform = simulation.calculate("macro/reform")["household"]["finance"]
     baseline_demographics = simulation.calculate("macro/baseline")[

policyengine/outputs/macro/comparison/poverty/gender.py

@@ -3,6 +3,16 @@
 
 
 def gender(simulation: Simulation):
+    """Calculate the impact of the reform on poverty by gender.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the poverty and deep poverty impact details with the following keys:
+            - poverty (dict): A dictionary with keys representing genders and values as dictionaries with baseline and reform poverty rates.
+            - deep_poverty (dict): A dictionary with keys representing genders and values as dictionaries with baseline and reform deep poverty rates.
+    """
     baseline = simulation.calculate("macro/baseline")["household"]["finance"]
     reform = simulation.calculate("macro/reform")["household"]["finance"]
     baseline_demographics = simulation.calculate("macro/baseline")[

policyengine/outputs/macro/comparison/poverty/race.py

@@ -3,6 +3,15 @@
 
 
 def race(simulation: Simulation):
+    """Calculate the impact of the reform on poverty by race.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the poverty impact details with the following keys:
+            - poverty (dict): A dictionary with keys representing races and values as dictionaries with baseline and reform poverty rates.
+    """
     baseline = simulation.calculate("macro/baseline")["household"]["finance"]
     reform = simulation.calculate("macro/reform")["household"]["finance"]
     baseline_demographics = simulation.calculate("macro/baseline")[

policyengine/outputs/macro/comparison/winners/income_decile.py

@@ -4,6 +4,16 @@
 
 
 def income_decile(simulation: Simulation):
+    """Calculate the impact of the reform on income deciles.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the impact details with the following keys:
+            - deciles (dict): A dictionary with keys representing outcome labels and values as lists of percentages for each decile.
+            - all (dict): A dictionary with keys representing outcome labels and values as overall percentages.
+    """
     baseline = simulation.calculate("macro/baseline")
     reform = simulation.calculate("macro/reform")
 

policyengine/outputs/macro/comparison/winners/wealth_decile.py

@@ -4,6 +4,16 @@
 
 
 def wealth_decile(simulation: Simulation):
+    """Calculate the impact of the reform on wealth deciles.
+
+    Args:
+        simulation (Simulation): The simulation for which the impact is to be calculated.
+
+    Returns:
+        dict: A dictionary containing the impact details with the following keys:
+            - deciles (dict): A dictionary with keys representing outcome labels and values as lists of percentages for each decile.
+            - all (dict): A dictionary with keys representing outcome labels and values as overall percentages.
+    """
     baseline = simulation.calculate("macro/baseline")
     reform = simulation.calculate("macro/reform")
 

policyengine/outputs/macro/single/gov/balance.py

@@ -2,17 +2,15 @@
 
 
 def balance(simulation: Simulation) -> dict:
-    sim = simulation.baseline
+    sim = simulation.selected
     if simulation.country == "uk":
         total_tax = sim.calculate("gov_tax").sum()
         total_spending = sim.calculate("gov_spending").sum()
         total_state_tax = 0
     elif simulation.country == "us":
         total_tax = sim.calculate("household_tax").sum()
         total_spending = sim.calculate("household_benefits").sum()
-        total_state_tax = simulation.calculate(
-            "household_state_income_tax"
-        ).sum()
+        total_state_tax = sim.calculate("household_state_income_tax").sum()
     return {
         "total_tax": total_tax,
         "total_spending": total_spending,