Make year a parameter in calculate_rebate_base_statistics()

MaxGhenis · MaxGhenis · commit 9430df82fcf9 · 2025-09-15T14:40:06.000-04:00
- Add year parameter to function (default 2026)
- Pass SIMULATION_YEAR constant when calling function
- Makes function more flexible for different year calculations
- Improves reusability and testing capabilities
diff --git a/us/states/ca/aei_rebate/aei_rebate_analysis.ipynb b/us/states/ca/aei_rebate/aei_rebate_analysis.ipynb
@@ -184,7 +184,7 @@
     }
    },
    "outputs": [],
-   "source": "def calculate_rebate_base_statistics(sim, unit_type=\"household\"):\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    \n    Returns:\n        Dictionary with rebate base statistics\n    \"\"\"\n    print(f\"Calculating {unit_type} statistics...\")\n    \n    if unit_type == \"household\":\n        # Calculate rebate base for all households\n        rebate_base = sim.calculate(\"ca_aei_rebate_base\", SIMULATION_YEAR)\n        \n        # Filter for California households only\n        household_state = sim.calculate(\"state_code\", SIMULATION_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\", SIMULATION_YEAR)\n        \n        # Use calculate_dataframe to get household-level data\n        household_df = sim.calculate_dataframe(\n            [\"household_id\", \"state_code\"],\n            SIMULATION_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            SIMULATION_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\")\ntax_unit_results = calculate_rebate_base_statistics(sim, \"tax_unit\")"
+   "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)"
   },
   {
    "cell_type": "markdown",

Original file line number	Diff line number	Diff line change
`@@ -184,7 +184,7 @@`
`184`	`184`	`}`
`185`	`185`	`},`
`186`	`186`	`"outputs": [],`
`187`		- "source": "def calculate_rebate_base_statistics(sim, unit_type=\"household\"):\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 \n Returns:\n Dictionary with rebate base statistics\n \"\"\"\n print(f\"Calculating {unit_type} statistics...\")\n \n if unit_type == \"household\":\n # Calculate rebate base for all households\n rebate_base = sim.calculate(\"ca_aei_rebate_base\", SIMULATION_YEAR)\n \n # Filter for California households only\n household_state = sim.calculate(\"state_code\", SIMULATION_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\", SIMULATION_YEAR)\n \n # Use calculate_dataframe to get household-level data\n household_df = sim.calculate_dataframe(\n [\"household_id\", \"state_code\"],\n SIMULATION_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 SIMULATION_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\")\ntax_unit_results = calculate_rebate_base_statistics(sim, \"tax_unit\")"
	`187`	+ "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)"
`188`	`188`	`},`
`189`	`189`	`{`
`190`	`190`	`"cell_type": "markdown",`