Use generic 'units' keys for easier data stacking

- Changed keys from 'total_ca_households'/'total_ca_tax_units' to 'total_ca_units'
- Changed 'households_with_rebate'/'tax_units_with_rebate' to 'units_with_rebate'
- Added 'unit_type' field to distinguish between household and tax_unit
- Results can now be easily stacked into a single DataFrame
- Demonstrated with combined results DataFrame example

Author: MaxGhenis

us/states/ca/aei_rebate/aei_rebate_analysis.ipynb

@@ -184,7 +184,7 @@
     }
    },
    "outputs": [],
-   "source": "def calculate_rebate_base_statistics(sim, unit_type=\"household\", year=2026):\n    \"\"\"\n    Calculate AEI rebate base program statistics for California households or tax units.\n    \n    Args:\n        sim: Microsimulation object with reform applied\n        unit_type: Either \"household\" or \"tax_unit\"\n        year: Year to calculate for\n    \n    Returns:\n        Dictionary with rebate base statistics\n    \"\"\"\n    print(f\"Calculating {unit_type} statistics for {year}...\")\n    \n    if unit_type == \"household\":\n        # Calculate rebate base for all households\n        rebate_base = sim.calculate(\"ca_aei_rebate_base\", year)\n        \n        # Filter for California households only\n        household_state = sim.calculate(\"state_code\", year, map_to=\"household\")\n        ca_mask = household_state == \"CA\"\n        \n        # Apply CA filter\n        ca_rebate_base = rebate_base[ca_mask]\n        total_ca_units = ca_mask.sum()\n        \n    else:  # tax_unit\n        # Calculate rebate base for all tax units (defined_for gives 0 for non-CA)\n        rebate_base = sim.calculate(\"ca_aei_rebate_base_tax_unit\", year)\n        \n        # Use calculate_dataframe to get household-level data\n        household_df = sim.calculate_dataframe(\n            [\"household_id\", \"state_code\"],\n            year,\n            map_to=\"household\"\n        )\n        \n        # Get tax unit data\n        tax_unit_df = sim.calculate_dataframe(\n            [\"tax_unit_id\", \"tax_unit_household_id\"],\n            year\n        )\n        \n        # Merge to get state for each tax unit\n        tax_unit_with_state = tax_unit_df.merge(\n            household_df[[\"household_id\", \"state_code\"]],\n            left_on=\"tax_unit_household_id\",\n            right_on=\"household_id\",\n            how=\"left\"\n        )\n        \n        # Create a boolean MicroSeries for CA tax units\n        ca_tax_unit_mask = tax_unit_with_state[\"state_code\"] == \"CA\"\n        total_ca_units = ca_tax_unit_mask.sum()\n        \n        # For tax units, we use all rebates (defined_for already filters to CA)\n        ca_rebate_base = rebate_base\n    \n    # Calculate statistics (MicroSeries already contain weights)\n    units_with_rebate = (ca_rebate_base > 0).sum()\n    total_rebate_base = ca_rebate_base.sum()\n    average_rebate_base = ca_rebate_base[ca_rebate_base > 0].mean() if units_with_rebate > 0 else 0\n    \n    return {\n        f\"total_ca_{unit_type}s\": total_ca_units,\n        f\"{unit_type}s_with_rebate\": units_with_rebate,\n        \"rebate_percentage\": units_with_rebate / total_ca_units,\n        \"average_rebate_base\": average_rebate_base,\n        \"total_rebate_base\": total_rebate_base,\n    }\n\n# Create simulation once\nprint(\"Loading data and creating simulation...\")\nreform = create_aei_reform()\nsim = Microsimulation(\n    dataset=\"hf://policyengine/policyengine-us-data/pooled_3_year_cps_2023.h5\",\n    reform=reform\n)\n\n# Calculate both household and tax unit results using the same simulation\nhousehold_results = calculate_rebate_base_statistics(sim, \"household\", SIMULATION_YEAR)\ntax_unit_results = calculate_rebate_base_statistics(sim, \"tax_unit\", SIMULATION_YEAR)"
+   "source": "def calculate_rebate_base_statistics(sim, unit_type=\"household\", year=2026):\n    \"\"\"\n    Calculate AEI rebate base program statistics for California households or tax units.\n    \n    Args:\n        sim: Microsimulation object with reform applied\n        unit_type: Either \"household\" or \"tax_unit\"\n        year: Year to calculate for\n    \n    Returns:\n        Dictionary with rebate base statistics\n    \"\"\"\n    print(f\"Calculating {unit_type} statistics for {year}...\")\n    \n    if unit_type == \"household\":\n        # Calculate rebate base for all households\n        rebate_base = sim.calculate(\"ca_aei_rebate_base\", year)\n        \n        # Filter for California households only\n        household_state = sim.calculate(\"state_code\", year, map_to=\"household\")\n        ca_mask = household_state == \"CA\"\n        \n        # Apply CA filter\n        ca_rebate_base = rebate_base[ca_mask]\n        total_ca_units = ca_mask.sum()\n        \n    else:  # tax_unit\n        # Calculate rebate base for all tax units (defined_for gives 0 for non-CA)\n        rebate_base = sim.calculate(\"ca_aei_rebate_base_tax_unit\", year)\n        \n        # Use calculate_dataframe to get household-level data\n        household_df = sim.calculate_dataframe(\n            [\"household_id\", \"state_code\"],\n            year,\n            map_to=\"household\"\n        )\n        \n        # Get tax unit data\n        tax_unit_df = sim.calculate_dataframe(\n            [\"tax_unit_id\", \"tax_unit_household_id\"],\n            year\n        )\n        \n        # Merge to get state for each tax unit\n        tax_unit_with_state = tax_unit_df.merge(\n            household_df[[\"household_id\", \"state_code\"]],\n            left_on=\"tax_unit_household_id\",\n            right_on=\"household_id\",\n            how=\"left\"\n        )\n        \n        # Create a boolean MicroSeries for CA tax units\n        ca_tax_unit_mask = tax_unit_with_state[\"state_code\"] == \"CA\"\n        total_ca_units = ca_tax_unit_mask.sum()\n        \n        # For tax units, we use all rebates (defined_for already filters to CA)\n        ca_rebate_base = rebate_base\n    \n    # Calculate statistics (MicroSeries already contain weights)\n    units_with_rebate = (ca_rebate_base > 0).sum()\n    total_rebate_base = ca_rebate_base.sum()\n    average_rebate_base = ca_rebate_base[ca_rebate_base > 0].mean() if units_with_rebate > 0 else 0\n    \n    return {\n        \"unit_type\": unit_type,\n        \"total_ca_units\": total_ca_units,\n        \"units_with_rebate\": units_with_rebate,\n        \"rebate_percentage\": units_with_rebate / total_ca_units,\n        \"average_rebate_base\": average_rebate_base,\n        \"total_rebate_base\": total_rebate_base,\n    }\n\n# Create simulation once\nprint(\"Loading data and creating simulation...\")\nreform = create_aei_reform()\nsim = Microsimulation(\n    dataset=\"hf://policyengine/policyengine-us-data/pooled_3_year_cps_2023.h5\",\n    reform=reform\n)\n\n# Calculate both household and tax unit results using the same simulation\nhousehold_results = calculate_rebate_base_statistics(sim, \"household\", SIMULATION_YEAR)\ntax_unit_results = calculate_rebate_base_statistics(sim, \"tax_unit\", SIMULATION_YEAR)"
   },
   {
    "cell_type": "markdown",
@@ -203,11 +203,18 @@
     }
    },
    "outputs": [],
-   "source": "# Display household results\nhousehold_table = pd.DataFrame([\n    {'Metric': 'Total CA Households', 'Value': f\"{household_results['total_ca_households']/1e6:.1f}M\"},\n    {'Metric': 'Households with Rebate', 'Value': f\"{household_results['households_with_rebate']/1e6:.1f}M ({household_results['rebate_percentage']:.0%})\"},\n    {'Metric': 'Average Rebate Base', 'Value': f\"${household_results['average_rebate_base']:,.0f}\"},\n    {'Metric': 'Total Rebate Base', 'Value': f\"${household_results['total_rebate_base']/1e9:.1f}B\"},\n])\n\nprint(f\"\\n=== HOUSEHOLD-LEVEL RESULTS ({SIMULATION_YEAR}) ===\")\nhousehold_table"
+   "source": "# Display household results\nhousehold_table = pd.DataFrame([\n    {'Metric': 'Total CA Households', 'Value': f\"{household_results['total_ca_units']/1e6:.1f}M\"},\n    {'Metric': 'Households with Rebate', 'Value': f\"{household_results['units_with_rebate']/1e6:.1f}M ({household_results['rebate_percentage']:.0%})\"},\n    {'Metric': 'Average Rebate Base', 'Value': f\"${household_results['average_rebate_base']:,.0f}\"},\n    {'Metric': 'Total Rebate Base', 'Value': f\"${household_results['total_rebate_base']/1e9:.1f}B\"},\n])\n\nprint(f\"\\n=== HOUSEHOLD-LEVEL RESULTS ({SIMULATION_YEAR}) ===\")\nhousehold_table"
   },
   {
    "cell_type": "code",
-   "source": "# Display tax unit results  \ntax_unit_table = pd.DataFrame([\n    {'Metric': 'Total CA Tax Units', 'Value': f\"{tax_unit_results['total_ca_tax_units']/1e6:.1f}M\"},\n    {'Metric': 'Tax Units with Rebate', 'Value': f\"{tax_unit_results['tax_units_with_rebate']/1e6:.1f}M ({tax_unit_results['rebate_percentage']:.0%})\"},\n    {'Metric': 'Average Rebate Base', 'Value': f\"${tax_unit_results['average_rebate_base']:,.0f}\"},\n    {'Metric': 'Total Rebate Base', 'Value': f\"${tax_unit_results['total_rebate_base']/1e9:.1f}B\"},\n])\n\nprint(f\"\\n=== TAX UNIT-LEVEL RESULTS ({SIMULATION_YEAR}) ===\")\ntax_unit_table",
+   "source": "# Display tax unit results  \ntax_unit_table = pd.DataFrame([\n    {'Metric': 'Total CA Tax Units', 'Value': f\"{tax_unit_results['total_ca_units']/1e6:.1f}M\"},\n    {'Metric': 'Tax Units with Rebate', 'Value': f\"{tax_unit_results['units_with_rebate']/1e6:.1f}M ({tax_unit_results['rebate_percentage']:.0%})\"},\n    {'Metric': 'Average Rebate Base', 'Value': f\"${tax_unit_results['average_rebate_base']:,.0f}\"},\n    {'Metric': 'Total Rebate Base', 'Value': f\"${tax_unit_results['total_rebate_base']/1e9:.1f}B\"},\n])\n\nprint(f\"\\n=== TAX UNIT-LEVEL RESULTS ({SIMULATION_YEAR}) ===\")\ntax_unit_table",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": "# Example: Stack results into a single DataFrame\nresults_df = pd.DataFrame([household_results, tax_unit_results])\nprint(\"\\n=== COMBINED RESULTS ===\")\nresults_df",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -236,20 +243,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": [
-    "#### VAT Formula\n",
-    "The rebate base (X) is used in the VAT rate calculation:\n",
-    "```\n",
-    "t = Rs/(Cp - X - T + Ro)\n",
-    "```\n",
-    "Where:\n",
-    "- **t** = VAT rate to be determined\n",
-    "- **Rs** = Revenue target\n",
-    "- **Cp** = Private consumption\n",
-    "- **X** = Total rebate base (calculated above)\n",
-    "- **T** = Existing taxes\n",
-    "- **Ro** = Other revenue"
-   ]
+   "source": "# Create summary table\nsummary_data = {\n    'Approach': ['Household-Level', 'Tax Unit-Level'],\n    'Total Rebate Base (X)': [\n        f\"${household_results['total_rebate_base']/1e9:.1f}B\",\n        f\"${tax_unit_results['total_rebate_base']/1e9:.1f}B\"\n    ],\n    'Coverage': [\n        f\"{household_results['units_with_rebate']/1e6:.1f}M ({household_results['rebate_percentage']:.0%})\",\n        f\"{tax_unit_results['units_with_rebate']/1e6:.1f}M ({tax_unit_results['rebate_percentage']:.0%})\"\n    ],\n    'VAT Formula Value': [\n        f\"${household_results['total_rebate_base']:,.0f}\",\n        f\"${tax_unit_results['total_rebate_base']:,.0f}\"\n    ]\n}\n\nsummary_df = pd.DataFrame(summary_data)\nprint(f\"=== SUMMARY OF REBATE BASE CALCULATIONS ({SIMULATION_YEAR}) ===\")\nsummary_df"
   },
   {
    "cell_type": "markdown",
@@ -408,7 +402,7 @@
     }
    },
    "outputs": [],
-   "source": "# Create comparison table\ncomparison_data = {\n    'Metric': [\n        'Total Count',\n        'Eligible Units (with rebate base)',\n        'Total Rebate Base',\n        'Average Rebate Base (eligible only)'\n    ],\n    'Households': [\n        f\"{household_results['total_ca_households']/1e6:.1f}M\",\n        f\"{household_results['households_with_rebate']/1e6:.1f}M ({household_results['rebate_percentage']:.0%})\",\n        f\"${household_results['total_rebate_base']/1e9:.1f}B\",\n        f\"${household_results['average_rebate_base']:,.0f}\"\n    ],\n    'Tax Units': [\n        f\"{tax_unit_results['total_ca_tax_units']/1e6:.1f}M\",\n        f\"{tax_unit_results['tax_units_with_rebate']/1e6:.1f}M ({tax_unit_results['rebate_percentage']:.0%})\",\n        f\"${tax_unit_results['total_rebate_base']/1e9:.1f}B\",\n        f\"${tax_unit_results['average_rebate_base']:,.0f}\"\n    ]\n}\n\ncomparison_df = pd.DataFrame(comparison_data)\n\n# Calculate ratios for all metrics\nratios = [\n    f\"{tax_unit_results['total_ca_tax_units']/household_results['total_ca_households']:.2f}x\",\n    f\"{tax_unit_results['tax_units_with_rebate']/household_results['households_with_rebate']:.2f}x\",\n    f\"{tax_unit_results['total_rebate_base']/household_results['total_rebate_base']:.2f}x\",\n    f\"{tax_unit_results['average_rebate_base']/household_results['average_rebate_base']:.2f}x\"\n]\n\ncomparison_df['Ratio (TU/HH)'] = ratios\n\nprint(f\"=== HOUSEHOLD VS TAX UNIT ANALYSIS ({SIMULATION_YEAR}) ===\")\ncomparison_df"
+   "source": "# Create comparison table\ncomparison_data = {\n    'Metric': [\n        'Total Count',\n        'Eligible Units (with rebate base)',\n        'Total Rebate Base',\n        'Average Rebate Base (eligible only)'\n    ],\n    'Households': [\n        f\"{household_results['total_ca_units']/1e6:.1f}M\",\n        f\"{household_results['units_with_rebate']/1e6:.1f}M ({household_results['rebate_percentage']:.0%})\",\n        f\"${household_results['total_rebate_base']/1e9:.1f}B\",\n        f\"${household_results['average_rebate_base']:,.0f}\"\n    ],\n    'Tax Units': [\n        f\"{tax_unit_results['total_ca_units']/1e6:.1f}M\",\n        f\"{tax_unit_results['units_with_rebate']/1e6:.1f}M ({tax_unit_results['rebate_percentage']:.0%})\",\n        f\"${tax_unit_results['total_rebate_base']/1e9:.1f}B\",\n        f\"${tax_unit_results['average_rebate_base']:,.0f}\"\n    ]\n}\n\ncomparison_df = pd.DataFrame(comparison_data)\n\n# Calculate ratios for all metrics\nratios = [\n    f\"{tax_unit_results['total_ca_units']/household_results['total_ca_units']:.2f}x\",\n    f\"{tax_unit_results['units_with_rebate']/household_results['units_with_rebate']:.2f}x\",\n    f\"{tax_unit_results['total_rebate_base']/household_results['total_rebate_base']:.2f}x\",\n    f\"{tax_unit_results['average_rebate_base']/household_results['average_rebate_base']:.2f}x\"\n]\n\ncomparison_df['Ratio (TU/HH)'] = ratios\n\nprint(f\"=== HOUSEHOLD VS TAX UNIT ANALYSIS ({SIMULATION_YEAR}) ===\")\ncomparison_df"
   },
   {
    "cell_type": "markdown",