Merge pull request #41 from PolicyEngine/api-parity

Add some API impacts

.github/workflows/pr.yml

@@ -21,7 +21,7 @@ jobs:
           pip install furo
       
       - name: Generate documentation
-        run: make docs
+        run: make documentation
 
       - name: Check documentation build
         run: |

.github/workflows/push.yml

@@ -21,8 +21,6 @@ jobs:
           pip install furo
           pip install sphinx-argparse
 
-        
-      
       - name: Generate documentation
         run: make documentation
 

docs/_toc.yml

@@ -3,4 +3,7 @@ root: index
 parts:
   - caption: Maintenance
     chapters:
-    - file: maintaining
+    - file: maintaining
+  - caption: Schema
+    chapters:
+    - file: schema

docs/schema.ipynb

@@ -0,0 +1,140 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Schema\n",
+    "\n",
+    "This page contains examples of the schema of the returned simulation outputs for given types of simulation. Each subsection specifies the country, scope and whether a reform has been passed. The schema is given in YAML format.\n",
+    "\n",
+    "## UK, macro, reform-comparison"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "```yaml\n",
+       "macro:\n",
+       "  baseline:\n",
+       "    gov:\n",
+       "      balance:\n",
+       "        total_gov_spending: <class 'numpy.float64'>\n",
+       "        total_state_tax: <class 'int'>\n",
+       "        total_tax_revenue: <class 'numpy.float64'>\n",
+       "      programs:\n",
+       "        child_benefit: <class 'numpy.float64'>\n",
+       "        council_tax: <class 'numpy.float64'>\n",
+       "        fuel_duty: <class 'numpy.float64'>\n",
+       "        income_tax: <class 'numpy.float64'>\n",
+       "        national_insurance: <class 'numpy.float64'>\n",
+       "        ni_employer: <class 'numpy.float64'>\n",
+       "        pension_credit: <class 'numpy.float64'>\n",
+       "        state_pension: <class 'numpy.float64'>\n",
+       "        tax_credits: <class 'numpy.float64'>\n",
+       "        universal_credit: <class 'numpy.float64'>\n",
+       "        vat: <class 'numpy.float64'>\n",
+       "    household:\n",
+       "      demographic_values:\n",
+       "        household_count_people: <class 'list'>\n",
+       "        household_weight: <class 'list'>\n",
+       "        person_weight: <class 'list'>\n",
+       "  comparison:\n",
+       "    revenue_impact:\n",
+       "      tax_revenues: <class 'numpy.float64'>\n",
+       "  reform:\n",
+       "    gov:\n",
+       "      balance:\n",
+       "        total_gov_spending: <class 'numpy.float64'>\n",
+       "        total_state_tax: <class 'int'>\n",
+       "        total_tax_revenue: <class 'numpy.float64'>\n",
+       "      programs:\n",
+       "        child_benefit: <class 'numpy.float64'>\n",
+       "        council_tax: <class 'numpy.float64'>\n",
+       "        fuel_duty: <class 'numpy.float64'>\n",
+       "        income_tax: <class 'numpy.float64'>\n",
+       "        national_insurance: <class 'numpy.float64'>\n",
+       "        ni_employer: <class 'numpy.float64'>\n",
+       "        pension_credit: <class 'numpy.float64'>\n",
+       "        state_pension: <class 'numpy.float64'>\n",
+       "        tax_credits: <class 'numpy.float64'>\n",
+       "        universal_credit: <class 'numpy.float64'>\n",
+       "        vat: <class 'numpy.float64'>\n",
+       "    household:\n",
+       "      demographic_values:\n",
+       "        household_count_people: <class 'list'>\n",
+       "        household_weight: <class 'list'>\n",
+       "        person_weight: <class 'list'>\n",
+       "\n",
+       "```"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from policyengine import Simulation\n",
+    "from IPython.display import Markdown\n",
+    "import yaml\n",
+    "\n",
+    "def replace_value_with_dtype(tree):\n",
+    "    for key in tree:\n",
+    "        if isinstance(tree[key], dict):\n",
+    "            replace_value_with_dtype(tree[key])\n",
+    "        else:\n",
+    "            tree[key] = str(type(tree[key]))\n",
+    "    return tree\n",
+    "\n",
+    "sim = Simulation(\n",
+    "    country=\"uk\",\n",
+    "    scope=\"macro\",\n",
+    "    time_period=2025,\n",
+    "    reform={\n",
+    "        \"gov.hmrc.income_tax.allowances.personal_allowance.amount\": {\n",
+    "            \"2025\": 0,\n",
+    "        }\n",
+    "    }\n",
+    ")\n",
+    "sim.calculate(\"macro\")\n",
+    "tree = replace_value_with_dtype(sim.outputs)\n",
+    "\n",
+    "Markdown('```yaml\\n' + yaml.dump(tree, indent=2) + '\\n```')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

policyengine/outputs/macro/comparison/revenue_impact.py

@@ -10,6 +10,13 @@ def revenue_impact(simulation: Simulation):
     Returns:
         float: The revenue impact of the simulation.
     """
-    tax_revenue_baseline = simulation.calculate("macro/baseline/tax_revenue")
-    tax_revenue_reform = simulation.calculate("macro/reform/tax_revenue")
-    return tax_revenue_reform - tax_revenue_baseline
+    tax_revenue_baseline = simulation.calculate(
+        "macro/baseline/gov/balance"
+    )["total_tax_revenue"]
+    tax_revenue_reform = simulation.calculate(
+        "macro/reform/gov/balance"
+    )["total_tax_revenue"]
+    tax_revenue_impact = tax_revenue_reform - tax_revenue_baseline
+    return {
+        "tax_revenues": tax_revenue_impact,
+    }

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

@@ -0,0 +1,20 @@
+from policyengine import Simulation
+
+
+def balance(simulation: Simulation) -> dict:
+    sim = simulation.baseline
+    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()
+    return {
+        "total_tax_revenue": total_tax,
+        "total_gov_spending": total_spending,
+        "total_state_tax": total_state_tax,
+    }

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

@@ -0,0 +1,35 @@
+from policyengine import Simulation
+from dataclasses import dataclass
+
+
+@dataclass
+class UKProgram:
+    name: str
+    is_positive: bool
+
+
+class UKPrograms:
+    PROGRAMS = [
+        UKProgram("income_tax", True),
+        UKProgram("national_insurance", True),
+        UKProgram("vat", True),
+        UKProgram("council_tax", True),
+        UKProgram("fuel_duty", True),
+        UKProgram("tax_credits", False),
+        UKProgram("universal_credit", False),
+        UKProgram("child_benefit", False),
+        UKProgram("state_pension", False),
+        UKProgram("pension_credit", False),
+        UKProgram("ni_employer", True),
+    ]
+
+
+def programs(simulation: Simulation) -> dict:
+    if simulation.country == "uk":
+        return {
+            program.name: simulation.baseline.calculate(
+                program.name, map_to="household"
+            ).sum()
+            * (1 if program.is_positive else -1)
+            for program in UKPrograms.PROGRAMS
+        }

policyengine/outputs/macro/single/household/demographic_values.py

@@ -0,0 +1,15 @@
+from policyengine import Simulation
+
+
+def demographic_values(simulation: Simulation) -> dict:
+    sim = simulation.baseline
+    household_count_people = (
+        sim.calculate("household_count_people").astype(int).tolist()
+    )
+    person_weight = sim.calculate("person_weight").astype(float).tolist()
+    household_weight = sim.calculate("household_weight").astype(float).tolist()
+    return {
+        "household_count_people": household_count_people,
+        "person_weight": person_weight,
+        "household_weight": household_weight,
+    }

policyengine/simulation.py

@@ -73,9 +73,15 @@ def calculate(self, output: str):
         """
         if output.endswith("/"):
             output = output[:-1]
+
+        if output == "":
+            output = list(self.outputs.keys())[0]
 
         node = self.outputs
 
+        for child_key in output.split("/")[:-1]:
+            node = node[child_key]
+
         parent = node
         child_key = output.split("/")[-1]
         node = parent[child_key]