Simplify rebate calculation and rename for consistency

- Remove redundant where statement (phaseout_percentage already handles cutoff)
- Use lowercase for phaseout_width (not a constant)
- Rename ca_aei_rebate_base to ca_aei_rebate_base_household for consistency

Author: MaxGhenis

us/states/ca/aei_rebate/aei_rebate_analysis.ipynb

@@ -86,7 +86,7 @@
     }
    },
    "outputs": [],
-   "source": "def create_aei_reform():\n    \"\"\"\n    Creates a PolicyEngine reform that adds the California AEI rebate base variables.\n    \n    Returns:\n        Reform: PolicyEngine reform with ca_aei_rebate_base variables\n    \"\"\"\n    \n    class household_fpg(Variable):\n        value_type = float\n        entity = Household\n        label = \"Household's federal poverty guideline\"\n        definition_period = YEAR\n        unit = USD\n\n        def formula(household, period, parameters):\n            n = household(\"household_size\", period)\n            state_group = household(\"state_group_str\", period)\n            p_fpg = parameters(period).gov.hhs.fpg\n            p1 = p_fpg.first_person[state_group]\n            pn = p_fpg.additional_person[state_group]\n            return p1 + pn * (n - 1)\n    \n    class ca_aei_rebate_base_tax_unit(Variable):\n        value_type = float\n        entity = TaxUnit\n        label = \"California AEI rebate base (tax unit version)\"\n        unit = USD\n        definition_period = YEAR\n        defined_for = StateCode.CA\n\n        def formula(tax_unit, period, parameters):\n            # Use tax unit's own AGI\n            income = tax_unit(\"adjusted_gross_income\", period)\n            fpg = tax_unit(\"tax_unit_fpg\", period)\n            income_to_fpg_ratio = where(fpg > 0, income / fpg, np.inf)\n\n            # Phase-out parameters\n            PHASEOUT_START = 1.5   # 150% FPG\n            PHASEOUT_END = 1.75    # 175% FPG\n            phaseout_width = PHASEOUT_END - PHASEOUT_START\n\n            # Phase-out calculation\n            excess = max_(income_to_fpg_ratio - PHASEOUT_START, 0)\n            phaseout_percentage = min_(1, excess / phaseout_width)\n\n            return where(\n                income_to_fpg_ratio <= PHASEOUT_END,\n                fpg * (1 - phaseout_percentage),\n                0\n            )\n\n    class ca_aei_rebate_base(Variable):\n        value_type = float\n        entity = Household\n        label = \"California AEI rebate base\"\n        unit = USD\n        definition_period = YEAR\n\n        def formula(household, period, parameters):\n            # Sum AGI from all tax units in the household\n            income = household.sum(household.members.tax_unit(\"adjusted_gross_income\", period))\n            fpg = household(\"household_fpg\", period)\n            income_to_fpg_ratio = where(fpg > 0, income / fpg, np.inf)\n\n            # Phase-out parameters\n            PHASEOUT_START = 1.5   # 150% FPG\n            PHASEOUT_END = 1.75    # 175% FPG\n            phaseout_width = PHASEOUT_END - PHASEOUT_START\n\n            # Phase-out calculation\n            excess = max_(income_to_fpg_ratio - PHASEOUT_START, 0)\n            phaseout_percentage = min_(1, excess / phaseout_width)\n\n            return where(\n                income_to_fpg_ratio <= PHASEOUT_END,\n                fpg * (1 - phaseout_percentage),\n                0\n            )\n\n    class AEIReform(Reform):\n        def apply(self):\n            self.update_variable(household_fpg)\n            self.update_variable(ca_aei_rebate_base_tax_unit)\n            self.update_variable(ca_aei_rebate_base)\n    \n    return AEIReform\n\nprint(\"Reform defined successfully\")"
+   "source": "def create_aei_reform():\n    \"\"\"\n    Creates a PolicyEngine reform that adds the California AEI rebate base variables.\n    \n    Returns:\n        Reform: PolicyEngine reform with ca_aei_rebate_base variables\n    \"\"\"\n    \n    class household_fpg(Variable):\n        value_type = float\n        entity = Household\n        label = \"Household's federal poverty guideline\"\n        definition_period = YEAR\n        unit = USD\n\n        def formula(household, period, parameters):\n            n = household(\"household_size\", period)\n            state_group = household(\"state_group_str\", period)\n            p_fpg = parameters(period).gov.hhs.fpg\n            p1 = p_fpg.first_person[state_group]\n            pn = p_fpg.additional_person[state_group]\n            return p1 + pn * (n - 1)\n    \n    class ca_aei_rebate_base_tax_unit(Variable):\n        value_type = float\n        entity = TaxUnit\n        label = \"California AEI rebate base (tax unit version)\"\n        unit = USD\n        definition_period = YEAR\n        defined_for = StateCode.CA\n\n        def formula(tax_unit, period, parameters):\n            # Use tax unit's own AGI\n            income = tax_unit(\"adjusted_gross_income\", period)\n            fpg = tax_unit(\"tax_unit_fpg\", period)\n            income_to_fpg_ratio = where(fpg > 0, income / fpg, np.inf)\n\n            # Phase-out parameters\n            PHASEOUT_START = 1.5   # 150% FPG\n            PHASEOUT_END = 1.75    # 175% FPG\n            phaseout_width = PHASEOUT_END - PHASEOUT_START\n\n            # Phase-out calculation\n            excess = max_(income_to_fpg_ratio - PHASEOUT_START, 0)\n            phaseout_percentage = min_(1, excess / phaseout_width)\n\n            return fpg * (1 - phaseout_percentage)\n\n    class ca_aei_rebate_base_household(Variable):\n        value_type = float\n        entity = Household\n        label = \"California AEI rebate base (household version)\"\n        unit = USD\n        definition_period = YEAR\n\n        def formula(household, period, parameters):\n            # Sum AGI from all tax units in the household\n            income = household.sum(household.members.tax_unit(\"adjusted_gross_income\", period))\n            fpg = household(\"household_fpg\", period)\n            income_to_fpg_ratio = where(fpg > 0, income / fpg, np.inf)\n\n            # Phase-out parameters\n            PHASEOUT_START = 1.5   # 150% FPG\n            PHASEOUT_END = 1.75    # 175% FPG\n            phaseout_width = PHASEOUT_END - PHASEOUT_START\n\n            # Phase-out calculation\n            excess = max_(income_to_fpg_ratio - PHASEOUT_START, 0)\n            phaseout_percentage = min_(1, excess / phaseout_width)\n\n            return fpg * (1 - phaseout_percentage)\n\n    class AEIReform(Reform):\n        def apply(self):\n            self.update_variable(household_fpg)\n            self.update_variable(ca_aei_rebate_base_tax_unit)\n            self.update_variable(ca_aei_rebate_base_household)\n    \n    return AEIReform\n\nprint(\"Reform defined successfully\")"
   },
   {
    "cell_type": "markdown",
@@ -97,7 +97,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2025-09-15T18:56:46.866675Z",
@@ -106,114 +106,8 @@
      "shell.execute_reply": "2025-09-15T18:56:58.429120Z"
     }
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Loading data and creating simulation...\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Calculating household statistics for 2026...\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Calculating tax_unit statistics for 2026...\n"
-     ]
-    }
-   ],
-   "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\n",
-    "print(\"Loading data and creating simulation...\")\n",
-    "reform = create_aei_reform()\n",
-    "sim = 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\n",
-    "household_results = calculate_rebate_base_statistics(sim, \"household\", SIMULATION_YEAR)\n",
-    "tax_unit_results = calculate_rebate_base_statistics(sim, \"tax_unit\", SIMULATION_YEAR)"
-   ]
+   "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_household\", 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",