Skip to content

Instantly share code, notes, and snippets.

@MaxGhenis

MaxGhenis/axiom_co_snap_ecps_parity.ipynb

Created May 1, 2026 16:50
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save MaxGhenis/785d167274ca772122a3dc48c4a48c6e to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save MaxGhenis/785d167274ca772122a3dc48c4a48c6e to your computer and use it in GitHub Desktop.
Download ZIP
Axiom Colorado SNAP eCPS parity demo
Display the source blob
Display the rendered blob
Raw
axiom_co_snap_ecps_parity.ipynb
{
"cells": [
{
"cell_type": "markdown",
"id": "d42ab816",
"metadata": {},
"source": [
"# Axiom Colorado SNAP eCPS parity demo\n",
"\n",
"This executed notebook shows Axiom RuleSpec computing Colorado SNAP benefits, then compares the Colorado SNAP composition against PolicyEngine enhanced CPS records.\n",
"\n",
"The parity check compares Axiom `snap_regular_month_allotment` to PolicyEngine `snap_normal_allotment`. That avoids two unrelated effects: eCPS does not include application-date facts for initial-month proration, and PolicyEngine top-level `snap` includes microsimulation take-up adjustments."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "c156fd9b",
"metadata": {
"execution": {
"iopub.execute_input": "2026-05-01T16:49:50.130416Z",
"iopub.status.busy": "2026-05-01T16:49:50.130351Z",
"iopub.status.idle": "2026-05-01T16:49:50.272642Z",
"shell.execute_reply": "2026-05-01T16:49:50.272168Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Repository commits:\n",
" axiom-rules: 1442b35\n",
" rules-us: 79be93e\n",
" rules-us-co: 1694a8a\n",
"\n",
"PolicyEngine version used by parity command:\n",
" policyengine-us: 1.680.0\n"
]
}
],
"source": [
"from pathlib import Path\n",
"import json\n",
"import os\n",
"import subprocess\n",
"import tempfile\n",
"\n",
"import yaml\n",
"from IPython.display import Markdown, display\n",
"\n",
"REPO = Path(\"/Users/maxghenis/TheAxiomFoundation/rules-us-co\")\n",
"AXIOM_RULES = REPO.parent / \"axiom-rules\"\n",
"RULES_US = REPO.parent / \"rules-us\"\n",
"PROGRAM = REPO / \"policies/cdhs/snap/fy-2026-benefit-calculation.yaml\"\n",
"TEST_TEMPLATE = PROGRAM.with_name(f\"{PROGRAM.stem}.test.yaml\")\n",
"AXIOM_BINARY = AXIOM_RULES / \"target/debug/axiom-rules\"\n",
"\n",
"def run(args, **kwargs):\n",
" return subprocess.run(args, text=True, capture_output=True, check=True, **kwargs).stdout.strip()\n",
"\n",
"def git(repo, *args):\n",
" return run([\"git\", \"-C\", str(repo), *args])\n",
"\n",
"commits = {\n",
" \"axiom-rules\": git(AXIOM_RULES, \"rev-parse\", \"--short\", \"HEAD\"),\n",
" \"rules-us\": git(RULES_US, \"rev-parse\", \"--short\", \"HEAD\"),\n",
" \"rules-us-co\": git(REPO, \"rev-parse\", \"--short\", \"HEAD\"),\n",
"}\n",
"pe_version = run([\n",
" \"uv\", \"run\", \"--with\", \"policyengine-us\", \"python\", \"-c\",\n",
" \"import importlib.metadata; print(importlib.metadata.version('policyengine-us'))\",\n",
"])\n",
"\n",
"print(\"Repository commits:\")\n",
"for name, sha in commits.items():\n",
" print(f\" {name}: {sha}\")\n",
"print(\"\\nPolicyEngine version used by parity command:\")\n",
"print(f\" policyengine-us: {pe_version}\")"
]
},
{
"cell_type": "markdown",
"id": "68f228db",
"metadata": {},
"source": [
"## Single household calculation using the Rust engine\n",
"\n",
"Compile the Colorado SNAP FY 2026 composition and run one household through the Axiom Rust engine. This uses the same facts as the composition test: one adult, $1,000/month earned income, $500/month shelter cost, and Colorado HCUA eligibility."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "b18bd1d1",
"metadata": {
"execution": {
"iopub.execute_input": "2026-05-01T16:49:50.274179Z",
"iopub.status.busy": "2026-05-01T16:49:50.274052Z",
"iopub.status.idle": "2026-05-01T16:49:50.335016Z",
"shell.execute_reply": "2026-05-01T16:49:50.334600Z"
}
},
"outputs": [
{
"data": {
"text/markdown": [
"| Output | Value |\n",
"|---|---:|\n",
"| `gross_income` | `1000.0` |\n",
"| `standard_deduction` | `209.0` |\n",
"| `earned_income_deduction` | `200.0` |\n",
"| `snap_standard_utility_allowance` | `594.0` |\n",
"| `shelter_costs` | `1094.0` |\n",
"| `excess_shelter_deduction` | `744.0` |\n",
"| `snap_net_income` | `0.0` |\n",
"| `snap_maximum_allotment` | `298.0` |\n",
"| `snap_eligible` | `holds` |\n",
"| `snap_regular_month_allotment` | `298.0` |"
],
"text/plain": [
"<IPython.core.display.Markdown object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"def scalar_value(value):\n",
" if isinstance(value, bool):\n",
" return {\"kind\": \"bool\", \"value\": value}\n",
" if isinstance(value, int) and not isinstance(value, bool):\n",
" return {\"kind\": \"integer\", \"value\": value}\n",
" if isinstance(value, float):\n",
" return {\"kind\": \"decimal\", \"value\": f\"{value:.6f}\".rstrip(\"0\").rstrip(\".\") or \"0\"}\n",
" if isinstance(value, str) and len(value) == 10 and value[4] == \"-\" and value[7] == \"-\":\n",
" return {\"kind\": \"date\", \"value\": value}\n",
" if isinstance(value, str):\n",
" return {\"kind\": \"text\", \"value\": value}\n",
" raise TypeError(value)\n",
"\n",
"def output_to_python(output):\n",
" if output[\"kind\"] == \"judgment\":\n",
" return output[\"outcome\"]\n",
" value = output[\"value\"]\n",
" if value[\"kind\"] == \"decimal\":\n",
" return float(value[\"value\"])\n",
" return value[\"value\"]\n",
"\n",
"base_inputs = yaml.safe_load(TEST_TEMPLATE.read_text())[0][\"input\"]\n",
"case_inputs = dict(base_inputs)\n",
"case_inputs.update(\n",
" {\n",
" \"household_size\": 1,\n",
" \"employee_wages_received\": 1000,\n",
" \"other_gain_or_benefit_payments\": 0,\n",
" \"elderly_or_disabled_household\": False,\n",
" \"household_shelter_costs_incurred\": 500,\n",
" \"household_incurred_or_anticipated_heating_or_cooling_costs_separate_from_rent_or_mortgage\": True,\n",
" \"initial_application_month\": False,\n",
" \"application_date\": \"2026-01-01\",\n",
" }\n",
")\n",
"\n",
"outputs_to_query = [\n",
" \"gross_income\",\n",
" \"standard_deduction\",\n",
" \"earned_income_deduction\",\n",
" \"snap_standard_utility_allowance\",\n",
" \"shelter_costs\",\n",
" \"excess_shelter_deduction\",\n",
" \"snap_net_income\",\n",
" \"snap_maximum_allotment\",\n",
" \"snap_eligible\",\n",
" \"snap_regular_month_allotment\",\n",
"]\n",
"\n",
"period = {\"period_kind\": \"month\", \"start\": \"2026-01-01\", \"end\": \"2026-01-31\", \"name\": \"2026-01\"}\n",
"interval = {\"start\": \"2026-01-01\", \"end\": \"2026-01-31\"}\n",
"request = {\n",
" \"mode\": \"fast\",\n",
" \"dataset\": {\n",
" \"inputs\": [\n",
" {\n",
" \"name\": name,\n",
" \"entity\": \"Household\",\n",
" \"entity_id\": \"demo-household\",\n",
" \"interval\": interval,\n",
" \"value\": scalar_value(value),\n",
" }\n",
" for name, value in case_inputs.items()\n",
" ],\n",
" \"relations\": [],\n",
" },\n",
" \"queries\": [{\"entity_id\": \"demo-household\", \"period\": period, \"outputs\": outputs_to_query}],\n",
"}\n",
"\n",
"DEMO_DIR = Path(tempfile.mkdtemp(prefix=\"axiom-snap-demo-\"))\n",
"DEMO_ARTIFACT = DEMO_DIR / \"program.compiled.json\"\n",
"DEMO_REQUEST = DEMO_DIR / \"request.json\"\n",
"DEMO_REQUEST.write_text(json.dumps(request))\n",
"\n",
"subprocess.run(\n",
" [str(AXIOM_BINARY), \"compile\", \"--program\", str(PROGRAM), \"--output\", str(DEMO_ARTIFACT)],\n",
" check=True,\n",
" text=True,\n",
" capture_output=True,\n",
")\n",
"result = subprocess.run(\n",
" [str(AXIOM_BINARY), \"run-compiled\", \"--artifact\", str(DEMO_ARTIFACT)],\n",
" input=json.dumps(request),\n",
" check=True,\n",
" text=True,\n",
" capture_output=True,\n",
")\n",
"\n",
"payload = json.loads(result.stdout)\n",
"calculated = {\n",
" name: output_to_python(payload[\"results\"][0][\"outputs\"][name])\n",
" for name in outputs_to_query\n",
"}\n",
"\n",
"rows = \"\\n\".join(f\"| `{name}` | `{value}` |\" for name, value in calculated.items())\n",
"display(Markdown(\"| Output | Value |\\n|---|---:|\\n\" + rows))"
]
},
{
"cell_type": "markdown",
"id": "c70533f8",
"metadata": {},
"source": [
"## JavaScript caller\n",
"\n",
"There is not yet a polished native npm package, but JavaScript can use the Rust engine today through the JSON CLI interface. This Node snippet reuses the compiled artifact and request from the previous cell."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "4ad68f1c",
"metadata": {
"execution": {
"iopub.execute_input": "2026-05-01T16:49:50.336102Z",
"iopub.status.busy": "2026-05-01T16:49:50.336020Z",
"iopub.status.idle": "2026-05-01T16:49:50.406612Z",
"shell.execute_reply": "2026-05-01T16:49:50.406181Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{\n",
" \"gross_income\": \"1000\",\n",
" \"snap_net_income\": \"0\",\n",
" \"snap_regular_month_allotment\": \"298\"\n",
"}\n",
"\n"
]
}
],
"source": [
"node_script = \"\"\"\n",
"const fs = require(\"node:fs\");\n",
"const { spawnSync } = require(\"node:child_process\");\n",
"\n",
"const binary = process.env.AXIOM_BINARY;\n",
"const artifact = process.env.DEMO_ARTIFACT;\n",
"const request = fs.readFileSync(process.env.DEMO_REQUEST, \"utf8\");\n",
"const result = spawnSync(binary, [\"run-compiled\", \"--artifact\", artifact], {\n",
" input: request,\n",
" encoding: \"utf8\",\n",
"});\n",
"\n",
"if (result.status !== 0) {\n",
" console.error(result.stderr || result.stdout);\n",
" process.exit(result.status || 1);\n",
"}\n",
"\n",
"const payload = JSON.parse(result.stdout);\n",
"const outputs = payload.results[0].outputs;\n",
"function value(output) {\n",
" if (output.kind === \"judgment\") return output.outcome;\n",
" return output.value.value;\n",
"}\n",
"console.log(JSON.stringify({\n",
" gross_income: value(outputs.gross_income),\n",
" snap_net_income: value(outputs.snap_net_income),\n",
" snap_regular_month_allotment: value(outputs.snap_regular_month_allotment),\n",
"}, null, 2));\n",
"\"\"\"\n",
"\n",
"js_result = subprocess.run(\n",
" [\"node\", \"-e\", node_script],\n",
" env={\n",
" **os.environ,\n",
" \"AXIOM_BINARY\": str(AXIOM_BINARY),\n",
" \"DEMO_ARTIFACT\": str(DEMO_ARTIFACT),\n",
" \"DEMO_REQUEST\": str(DEMO_REQUEST),\n",
" },\n",
" text=True,\n",
" capture_output=True,\n",
" check=True,\n",
")\n",
"print(js_result.stdout)"
]
},
{
"cell_type": "markdown",
"id": "9a86df17",
"metadata": {},
"source": [
"## Full Colorado eCPS parity run\n",
"\n",
"Now run the full Colorado enhanced CPS comparison. The `--project-policyengine-utility-allowance` option maps PolicyEngine's eCPS utility allowance type and aggregate deduction values into the closest Colorado input facts, because eCPS does not expose all raw utility and deduction facts. The live Axiom computation still derives eligibility, net income, utility allowance, shelter deduction, and allotment from encoded RuleSpec rules."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "6a4ccf86",
"metadata": {
"execution": {
"iopub.execute_input": "2026-05-01T16:49:50.407805Z",
"iopub.status.busy": "2026-05-01T16:49:50.407715Z",
"iopub.status.idle": "2026-05-01T16:50:42.138387Z",
"shell.execute_reply": "2026-05-01T16:50:42.137843Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loading PolicyEngine enhanced CPS...\n",
"Projecting 762 CO eCPS SPM units...\n",
"Skipped 64 CO eCPS SPM units with SNAP unit size < 1.\n",
"Compiling Colorado SNAP RuleSpec composition...\n",
"Running Axiom Rules over projected eCPS records...\n",
"\n",
"Compared 762 PolicyEngine eCPS SPM units\n",
"Tolerance: $1.50\n",
"Matches: 762/762 (100.0%)\n",
"Mean absolute difference: $0.06\n",
"Max absolute difference: $1.20\n",
"\n",
"Top 8 differences:\n",
" spm=25363001 PE=$68.20 Axiom=$67.00 diff=$-1.20 eligible PE=True Axiom=holds gross PE=$1748.00 Axiom=$1748.00 net PE=$767.00 Axiom=$767.00 utility PE=$594.00 Axiom=$594.00\n",
" spm=104993002 PE=$143.20 Axiom=$142.00 diff=$-1.20 eligible PE=True Axiom=holds gross PE=$1773.72 Axiom=$1773.72 net PE=$516.70 Axiom=$516.70 utility PE=$594.00 Axiom=$594.00\n",
" spm=40225001 PE=$28.20 Axiom=$27.00 diff=$-1.20 eligible PE=True Axiom=holds gross PE=$2572.25 Axiom=$2572.25 net PE=$1726.86 Axiom=$1726.86 utility PE=$594.00 Axiom=$594.00\n",
" spm=70323001 PE=$77.20 Axiom=$76.00 diff=$-1.20 eligible PE=True Axiom=holds gross PE=$4099.91 Axiom=$4099.91 net PE=$3056.93 Axiom=$3056.93 utility PE=$594.00 Axiom=$594.00\n",
" spm=54787001 PE=$51.10 Axiom=$50.00 diff=$-1.10 eligible PE=True Axiom=holds gross PE=$1329.36 Axiom=$1329.36 net PE=$823.82 Axiom=$823.82 utility PE=$594.00 Axiom=$594.00\n",
" spm=31770001 PE=$173.10 Axiom=$172.00 diff=$-1.10 eligible PE=True Axiom=holds gross PE=$2735.77 Axiom=$2735.77 net PE=$1243.79 Axiom=$1243.79 utility PE=$594.00 Axiom=$594.00\n",
" spm=35617001 PE=$89.10 Axiom=$88.00 diff=$-1.10 eligible PE=True Axiom=holds gross PE=$2944.32 Axiom=$2944.32 net PE=$1523.47 Axiom=$1523.47 utility PE=$594.00 Axiom=$594.00\n",
" spm=44167001 PE=$195.10 Axiom=$194.00 diff=$-1.10 eligible PE=True Axiom=holds gross PE=$3826.78 Axiom=$3826.78 net PE=$2663.69 Axiom=$2663.69 utility PE=$594.00 Axiom=$594.00\n",
"\n",
"STDERR:\n",
"Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n",
"\n"
]
}
],
"source": [
"cmd = [\n",
" \"uv\",\n",
" \"run\",\n",
" \"--with\",\n",
" \"policyengine-us\",\n",
" \"--with\",\n",
" \"pyyaml\",\n",
" \"python\",\n",
" str(REPO / \"scripts/compare_snap_policyengine_ecps.py\"),\n",
" \"--project-policyengine-utility-allowance\",\n",
" \"--fail-on-mismatch\",\n",
" \"--max-differences\",\n",
" \"8\",\n",
"]\n",
"completed = subprocess.run(cmd, cwd=REPO, text=True, capture_output=True)\n",
"print(completed.stdout)\n",
"if completed.stderr:\n",
" print(\"STDERR:\")\n",
" print(completed.stderr)\n",
"completed.check_returncode()"
]
},
{
"cell_type": "markdown",
"id": "8a2ae29f",
"metadata": {},
"source": [
"## Takeaway\n",
"\n",
"The parity run is a direct check that Axiom can calculate Colorado SNAP from encoded federal and Colorado rules at PolicyEngine-level fidelity over real eCPS records. Remaining differences are within the notebook's rounding tolerance because Axiom floors final allotments to whole dollars while PolicyEngine's normal allotment can retain fractional dollars."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.14.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment