policyengine-us 1.360.0__py3-none-any.whl → 1.437.0__py3-none-any.whl

policyengine_us/tests/test_batched.py

@@ -0,0 +1,369 @@
+#!/usr/bin/env python
+"""
+Generic test runner that can split tests into batches for memory management.
+Can be used for any test folder with configurable batch count.
+"""
+
+import subprocess
+import sys
+import os
+import gc
+import time
+import argparse
+from pathlib import Path
+from typing import List, Dict
+
+
+def count_yaml_files(directory: Path) -> int:
+    """Count YAML files in a directory recursively."""
+    if not directory.exists():
+        return 0
+    return len(list(directory.rglob("*.yaml")))
+
+
+def get_test_directories(base_path: Path) -> Dict[str, int]:
+    """Get all subdirectories and their test counts."""
+    dir_counts = {}
+
+    # Check for yaml files directly in base directory
+    root_count = len(list(base_path.glob("*.yaml")))
+    if root_count > 0:
+        dir_counts["."] = root_count
+
+    # Get all subdirectories with their test counts
+    for item in base_path.iterdir():
+        if item.is_dir():
+            yaml_count = count_yaml_files(item)
+            if yaml_count > 0:
+                dir_counts[item.name] = yaml_count
+
+    return dir_counts
+
+
+def split_into_batches(base_path: Path, num_batches: int) -> List[List[str]]:
+    """
+    Split test directories into specified number of batches.
+    Special handling for baseline tests to separate states.
+    Special handling for contrib tests to divide by folder count.
+    """
+    # Special handling for contrib tests - each folder is its own batch
+    # Only apply to policy/contrib (structural tests), not baseline/contrib
+    if str(base_path).endswith("policy/contrib"):
+        # Get all subdirectories and sort them alphabetically
+        subdirs = sorted(
+            [item for item in base_path.iterdir() if item.is_dir()]
+        )
+
+        # Get root level YAML files and sort them
+        root_files = sorted(list(base_path.glob("*.yaml")))
+
+        # Create one batch per subdirectory
+        batches = []
+        for subdir in subdirs:
+            batches.append([str(subdir)])
+
+        # If there are root files, group them together in their own batch
+        if root_files:
+            root_batch = [str(file) for file in root_files]
+            batches.append(root_batch)
+
+        return batches
+
+    # Special handling for reform tests - run all together in one batch
+    if "reform" in str(base_path):
+        return [[str(base_path)]]
+
+    # Special handling for baseline tests with 2 batches
+    if "baseline" in str(base_path) and num_batches == 2:
+        states_path = base_path / "gov" / "states"
+        if states_path.exists() and count_yaml_files(states_path) > 0:
+            # Batch 1: Only states
+            batch1 = [str(states_path)]
+
+            # Batch 2: Everything else (excluding states, household, and contrib)
+            batch2 = []
+
+            # Add root level files if any
+            for yaml_file in base_path.glob("*.yaml"):
+                batch2.append(str(yaml_file))
+
+            # Add all directories except gov/states, household, and contrib
+            for item in base_path.iterdir():
+                if item.is_dir():
+                    # Skip household and contrib directories (they'll be run separately)
+                    if item.name in ["household", "contrib"]:
+                        continue
+                    elif item.name == "gov":
+                        # Add gov subdirectories except states
+                        for gov_item in item.iterdir():
+                            if gov_item.is_dir() and gov_item.name != "states":
+                                batch2.append(str(gov_item))
+                            elif gov_item.suffix == ".yaml":
+                                batch2.append(str(gov_item))
+                    else:
+                        # Other non-gov directories
+                        batch2.append(str(item))
+
+            return [batch1, batch2] if batch2 else [batch1]
+
+    # Default behavior for non-baseline or different batch counts
+    dir_counts = get_test_directories(base_path)
+
+    if num_batches <= 0:
+        num_batches = 1
+
+    # If only 1 batch, return everything
+    if num_batches == 1:
+        return [[str(base_path)]]
+
+    # Sort directories by test count (largest first)
+    sorted_dirs = sorted(dir_counts.items(), key=lambda x: x[1], reverse=True)
+
+    # Initialize batches
+    batches = [[] for _ in range(num_batches)]
+    batch_counts = [0] * num_batches
+
+    # Distribute directories to batches (greedy algorithm - add to smallest batch)
+    for dir_name, count in sorted_dirs:
+        # Find batch with fewest tests
+        min_batch_idx = batch_counts.index(min(batch_counts))
+
+        # Add directory to that batch
+        if dir_name == ".":
+            # Root level files - add individually
+            for yaml_file in base_path.glob("*.yaml"):
+                batches[min_batch_idx].append(str(yaml_file))
+        else:
+            batches[min_batch_idx].append(str(base_path / dir_name))
+
+        batch_counts[min_batch_idx] += count
+
+    # Filter out empty batches
+    return [batch for batch in batches if batch]
+
+
+def run_batch(test_paths: List[str], batch_name: str) -> Dict:
+    """Run a batch of tests in an isolated subprocess."""
+
+    python_exe = sys.executable
+
+    start_time = time.time()
+
+    # Build command - direct policyengine-core with timeout protection
+    cmd = (
+        [
+            python_exe,
+            "-m",
+            "policyengine_core.scripts.policyengine_command",
+            "test",
+        ]
+        + test_paths
+        + ["-c", "policyengine_us"]
+    )
+
+    print(f"    Running {batch_name}...")
+    print(f"    Paths: {len(test_paths)} items")
+    print()
+
+    # Use Popen for more control over process lifecycle
+    process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,  # Line buffered
+    )
+
+    try:
+        test_completed = False
+        test_passed = False
+        output_lines = []
+
+        # Monitor output line by line
+        while True:
+            line = process.stdout.readline()
+            if not line:
+                # No more output, check if process is done
+                poll_result = process.poll()
+                if poll_result is not None:
+                    # Process terminated
+                    break
+                # Process still running but no output
+                time.sleep(0.1)
+                continue
+
+            # Print line in real-time
+            print(line, end="")
+            output_lines.append(line)
+
+            # Detect pytest completion
+            # Look for patterns like "====== 5638 passed in 491.24s ======"
+            # or "====== 2 failed, 5636 passed in 500s ======"
+            import re
+
+            if re.search(r"=+.*\d+\s+(passed|failed).*in\s+[\d.]+s.*=+", line):
+                test_completed = True
+                # Check if tests passed (no failures mentioned or 0 failed)
+                if "failed" not in line or "0 failed" in line:
+                    test_passed = True
+                else:
+                    test_passed = False
+
+                print(f"\n    Tests completed, terminating process...")
+
+                # Give 5 seconds grace period for cleanup
+                time.sleep(5)
+
+                # Terminate the process
+                process.terminate()
+                try:
+                    process.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    # Force kill if it won't terminate
+                    print(f"    Force killing process...")
+                    process.kill()
+                    process.wait()
+                break
+
+        # If we didn't detect completion, wait for process with timeout
+        if not test_completed:
+            try:
+                # Wait up to 30 minutes total
+                elapsed = time.time() - start_time
+                remaining_timeout = max(1800 - elapsed, 1)
+                process.wait(timeout=remaining_timeout)
+            except subprocess.TimeoutExpired:
+                print(f"\n    ⏱️ Timeout - terminating process...")
+                process.terminate()
+                try:
+                    process.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    process.kill()
+                    process.wait()
+
+                elapsed = time.time() - start_time
+                return {
+                    "elapsed": elapsed,
+                    "status": "timeout",
+                }
+
+        elapsed = time.time() - start_time
+
+        if test_completed:
+            print(f"\n    Batch completed in {elapsed:.1f}s")
+            return {
+                "elapsed": elapsed,
+                "status": "passed" if test_passed else "failed",
+            }
+        else:
+            # Process ended without detecting test completion
+            returncode = process.poll()
+            print(f"\n    Batch completed in {elapsed:.1f}s")
+            return {
+                "elapsed": elapsed,
+                "status": "passed" if returncode == 0 else "failed",
+            }
+
+    except Exception as e:
+        elapsed = time.time() - start_time
+        print(f"\n    ❌ Error: {str(e)[:100]}")
+
+        # Clean up process if still running
+        try:
+            process.terminate()
+            process.wait(timeout=5)
+        except:
+            try:
+                process.kill()
+                process.wait()
+            except:
+                pass
+
+        return {
+            "elapsed": elapsed,
+            "status": "error",
+        }
+
+
+def main():
+    """Main entry point for the generic test runner."""
+
+    parser = argparse.ArgumentParser(
+        description="Run tests in batches with memory cleanup between batches"
+    )
+    parser.add_argument("test_path", help="Path to the test directory")
+    parser.add_argument(
+        "--batches",
+        type=int,
+        default=2,
+        help="Number of batches to split tests into (default: 2)",
+    )
+
+    args = parser.parse_args()
+
+    if not os.path.exists("policyengine_us"):
+        print("Error: Must run from PolicyEngine US root directory")
+        sys.exit(1)
+
+    test_path = Path(args.test_path)
+    if not test_path.exists():
+        print(f"Error: Test path does not exist: {args.test_path}")
+        sys.exit(1)
+
+    print(f"PolicyEngine Test Runner")
+    print("=" * 60)
+    print(f"Test path: {test_path}")
+    print(f"Requested batches: {args.batches}")
+
+    # Count total tests
+    total_tests = count_yaml_files(test_path)
+    print(f"Total test files: {total_tests}")
+
+    # Split into batches
+    batches = split_into_batches(test_path, args.batches)
+    if len(batches) != args.batches:
+        print(
+            f"Actual batches: {len(batches)} (optimized for {total_tests} files)"
+        )
+    else:
+        print(f"Actual batches: {len(batches)}")
+    print("=" * 60)
+
+    # Run batches
+    all_failed = False
+    total_elapsed = 0
+
+    for i, batch_paths in enumerate(batches, 1):
+        print(f"\n[Batch {i}/{len(batches)}]")
+
+        # Show what's in this batch
+        batch_test_count = sum(
+            count_yaml_files(Path(p)) if Path(p).is_dir() else 1
+            for p in batch_paths
+        )
+        print(f"  Test files: ~{batch_test_count}")
+        print("-" * 60)
+
+        result = run_batch(batch_paths, f"Batch {i}")
+        total_elapsed += result["elapsed"]
+
+        if result["status"] != "passed":
+            all_failed = True
+
+        # Memory cleanup after each batch
+        print("  Cleaning up memory...")
+        gc.collect()
+
+    # Final summary
+    print("\n" + "=" * 60)
+    print("TEST SUMMARY")
+    print("=" * 60)
+    print(f"Total test files: {total_tests}")
+    print(f"Total time: {total_elapsed:.1f}s")
+
+    # Exit code
+    sys.exit(1 if all_failed else 0)
+
+
+if __name__ == "__main__":
+    main()

policyengine_us/tools/taxsim/generate_taxsim_tests.py

@@ -5,7 +5,6 @@ import pandas as pd
 import subprocess
 import stat
 from io import StringIO
-from policyengine_us_data import CPS_2024
 from policyengine_core.data import Dataset
 from policyengine_core.taxbenefitsystems import TaxBenefitSystem
 from policyengine_us import Microsimulation
@@ -284,6 +283,9 @@ class TaxSim35:
         return test_str
 
 
+DEFAULT_DATASET = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"
+
+
 if __name__ == "__main__":
     parser = ArgumentParser(description="Generate tests from TAXSIM35")
     parser.add_argument(
@@ -299,6 +301,6 @@ if __name__ == "__main__":
     taxsim = TaxSim35()
     taxsim.OUTPUT_VARIABLES = args.outputs
     result = taxsim.generate_from_microsimulation(
-        CPS_2024, args.year, number=args.num
+        DEFAULT_DATASET, args.year, number=args.num
     )
     print(result)

policyengine_us/variables/gov/aca/ptc/aca_ptc.py

@@ -13,7 +13,9 @@ class aca_ptc(Variable):
     def formula(tax_unit, period, parameters):
         plan_cost = tax_unit("slcsp", period)
         income = tax_unit("aca_magi", period)
-        applicable_figure = tax_unit("aca_ptc_phase_out_rate", period)
+        applicable_figure = tax_unit(
+            "aca_required_contribution_percentage", period
+        )
         takes_up_aca_if_eligible = tax_unit("takes_up_aca_if_eligible", period)
         return (
             max_(0, plan_cost - income * applicable_figure)

policyengine_us/variables/gov/aca/ptc/{aca_ptc_phase_out_rate.py → aca_required_contribution_percentage.py}

@@ -1,7 +1,7 @@
 from policyengine_us.model_api import *
 
 
-class aca_ptc_phase_out_rate(Variable):
+class aca_required_contribution_percentage(Variable):
     value_type = float
     entity = TaxUnit
     label = "ACA PTC phase-out rate (i.e., IRS Form 8962 'applicable figure')"
@@ -11,5 +11,5 @@ class aca_ptc_phase_out_rate(Variable):
 
     def formula(tax_unit, period, parameters):
         magi_frac = tax_unit("aca_magi_fraction", period)
-        p = parameters(period).gov.aca.ptc_phase_out_rate
+        p = parameters(period).gov.aca.required_contribution_percentage
         return np.interp(magi_frac, p.thresholds, p.amounts)

policyengine_us/variables/gov/ed/pell_grant/efc/pell_grant_simplified_formula_applies.py

@@ -12,6 +12,5 @@ class pell_grant_simplified_formula_applies(Variable):
         head_income = tax_unit("pell_grant_primary_income", period)
         p = parameters(period).gov.ed.pell_grant.efc.simplified
         income_eligible = head_income < p.income_limit
-        total_benefits = add(tax_unit, period, p.benefits)
-        has_benefits = np.any(total_benefits > 0)
+        has_benefits = add(tax_unit, period, p.benefits) > 0
         return income_eligible & has_benefits & p.applies

policyengine_us/variables/gov/fcc/lifeline/is_lifeline_eligible.py

@@ -21,10 +21,9 @@ class is_lifeline_eligible(Variable):
         )
         categorically_eligible = np.where(
             is_on_tribal_land,
-            np.any(tribal_lifeline_programs),
-            np.any(non_tribal_lifeline_programs),
+            tribal_lifeline_programs > 0,
+            non_tribal_lifeline_programs > 0,
         )
-        fpg_ratio = spm_unit("fcc_fpg_ratio", period)
-        fpg_limit = p.fpg_limit
-        fpg_eligible = fpg_ratio <= fpg_limit
-        return categorically_eligible | fpg_eligible
+        # Use the new unified income eligibility variable
+        income_eligible = spm_unit("is_lifeline_income_eligible", period)
+        return categorically_eligible | income_eligible

policyengine_us/variables/gov/fcc/lifeline/is_lifeline_income_eligible.py

@@ -0,0 +1,35 @@
+from policyengine_us.model_api import *
+
+
+class is_lifeline_income_eligible(Variable):
+    value_type = bool
+    entity = SPMUnit
+    label = "Income-eligible for Lifeline"
+    documentation = (
+        "Meets income requirements for Lifeline (federal or state-expanded)"
+    )
+    definition_period = YEAR
+    reference = (
+        "https://www.law.cornell.edu/cfr/text/47/54.409",
+        "https://statutes.capitol.texas.gov/Docs/UT/htm/UT.55.htm#55.015",
+    )
+
+    def formula(spm_unit, period, parameters):
+        p = parameters(period).gov.fcc.lifeline
+        fpg_ratio = spm_unit("fcc_fpg_ratio", period)
+
+        # Some states expand eligibility beyond federal limit
+        state = spm_unit.household("state_code_str", period)
+        is_tx = state == "TX"
+
+        # Texas expands to 150% FPG per Texas Utilities Code § 55.015(d-1)
+        tx_p = parameters(period).gov.states.tx.uct.lifeline
+
+        # Use state-specific limit where applicable, otherwise federal (135%)
+        # Convert to float32 to match fpg_ratio dtype for comparison
+        # Without this, float32(1.35) > float64(1.35) due to precision differences
+        fpg_limit = where(is_tx, tx_p.fpg_limit, p.fpg_limit).astype(
+            np.float32
+        )
+
+        return fpg_ratio <= fpg_limit

policyengine_us/variables/gov/fcc/lifeline/lifeline.py

@@ -33,11 +33,14 @@ class lifeline(Variable):
                     base_amount,
                 )
 
-        max_monthly_amount = (
+        federal_monthly_amount = (
             base_amount
             + is_rural_tribal * p.fcc.lifeline.amount.rural_tribal_supplement
         )
-        max_amount = max_monthly_amount * MONTHS_IN_YEAR
+        # Get TX supplement for the specific month (now properly defined as MONTH period)
+        tx_supplement = spm_unit("tx_lifeline_supplement", period.first_month)
+        total_monthly_amount = federal_monthly_amount + tx_supplement
+        max_amount = total_monthly_amount * MONTHS_IN_YEAR
         phone_broadband_cost = add(
             spm_unit, period, ["phone_cost", "broadband_cost"]
         )

policyengine_us/variables/gov/hhs/head_start/early_head_start.py

@@ -4,6 +4,18 @@ from policyengine_us.model_api import *
 class early_head_start(Variable):
     value_type = float
     entity = Person
-    label = "Amount of Early Head Start benefit"
+    label = "Early Head Start"
+    unit = USD
     definition_period = YEAR
     defined_for = "is_early_head_start_eligible"
+    reference = "https://headstart.gov/program-data/article/head-start-program-facts-fiscal-year-2022"
+
+    def formula(person, period, parameters):
+        p = parameters(period).gov.hhs.head_start.early_head_start
+        state = person.household("state_code_str", period)
+        spending = p.spending[state]
+        enrollment = p.enrollment[state]
+        mask = enrollment > 0
+        result = np.zeros_like(p.spending[state])
+        result[mask] = spending[mask] / enrollment[mask]
+        return result

policyengine_us/variables/gov/hhs/head_start/head_start.py

@@ -4,6 +4,18 @@ from policyengine_us.model_api import *
 class head_start(Variable):
     value_type = float
     entity = Person
-    label = "Amount of Head Start benefit"
+    label = "Head Start"
+    unit = USD
     definition_period = YEAR
     defined_for = "is_head_start_eligible"
+    reference = "https://headstart.gov/program-data/article/head-start-program-facts-fiscal-year-2022"
+
+    def formula(person, period, parameters):
+        p = parameters(period).gov.hhs.head_start
+        state = person.household("state_code_str", period)
+        spending = p.spending[state]
+        enrollment = p.enrollment[state]
+        mask = enrollment > 0
+        result = np.zeros_like(p.spending[state])
+        result[mask] = spending[mask] / enrollment[mask]
+        return result

policyengine_us/variables/gov/hhs/head_start/is_head_start_categorically_eligible.py

@@ -12,7 +12,5 @@ class is_head_start_categorically_eligible(Variable):
     )
 
     def formula(person, period, parameters):
-        tax_unit = person.tax_unit
         p = parameters(period).gov.hhs.head_start
-        programs = add(tax_unit, period, p.categorical_eligibility)
-        return np.any(programs)
+        return add(person.tax_unit, period, p.categorical_eligibility) > 0

policyengine_us/variables/gov/hhs/medicaid/eligibility/is_medicaid_eligible.py

@@ -14,17 +14,11 @@ class is_medicaid_eligible(Variable):
     def formula(person, period, parameters):
         category = person("medicaid_category", period)
         categorically_eligible = category != category.possible_values.NONE
-        istatus = person("immigration_status", period)
-        undocumented = istatus == istatus.possible_values.UNDOCUMENTED
-        state = person.household("state_code_str", period)
-        p = parameters(period).gov.hhs.medicaid.eligibility
-        state_covers_undocumented = p.undocumented_immigrant[state].astype(
-            bool
-        )
-        immigration_status_eligible = (
-            ~undocumented | undocumented & state_covers_undocumented
+        immigration_status_eligible = person(
+            "is_medicaid_immigration_status_eligible", period
         )
         ca_ffyp_eligible = person("ca_ffyp_eligible", period)
+        p = parameters(period).gov.hhs.medicaid.eligibility
         if p.work_requirements.applies:
             work_requirement_eligible = person(
                 "medicaid_work_requirement_eligible", period

policyengine_us/variables/gov/hhs/medicaid/eligibility/is_medicaid_immigration_status_eligible.py

@@ -0,0 +1,44 @@
+from policyengine_us.model_api import *
+
+
+class is_medicaid_immigration_status_eligible(Variable):
+    value_type = bool
+    entity = Person
+    label = "Person is eligible for Medicaid due to immigration status"
+    definition_period = YEAR
+    reference = [
+        "https://www.law.cornell.edu/uscode/text/42/1396b#v",
+        "https://www.law.cornell.edu/uscode/text/8/1641",
+        "https://www.kff.org/racial-equity-and-health-policy/fact-sheet/key-facts-on-health-coverage-of-immigrants/",
+    ]
+
+    def formula(person, period, parameters):
+        p = parameters(period).gov.hhs.medicaid.eligibility
+        immigration_status = person("immigration_status", period)
+        immigration_status_str = immigration_status.decode_to_str()
+
+        # Check if immigration status is in the eligible list
+        eligible_immigration_status = np.isin(
+            immigration_status_str, p.eligible_immigration_statuses
+        )
+
+        # Special handling for undocumented immigrants in states that cover them
+        undocumented = (
+            immigration_status
+            == immigration_status.possible_values.UNDOCUMENTED
+        )
+        state = person.household("state_code_str", period)
+        state_covers_undocumented = p.undocumented_immigrant[state].astype(
+            bool
+        )
+        undocumented_eligible = undocumented & state_covers_undocumented
+
+        ca_eligible_regardless_of_immigration_status = person(
+            "is_ca_medicaid_immigration_status_eligible", period
+        )
+
+        return (
+            eligible_immigration_status
+            | undocumented_eligible
+            | ca_eligible_regardless_of_immigration_status
+        )

policyengine_us/variables/gov/hhs/tanf/cash/income/tanf_gross_earned_income.py

@@ -0,0 +1,13 @@
+from policyengine_us.model_api import *
+
+
+class tanf_gross_earned_income(Variable):
+    value_type = float
+    entity = Person
+    label = (
+        "Temporary Assistance for Needy Families (TANF) gross earned income"
+    )
+    unit = USD
+    definition_period = MONTH
+
+    adds = "gov.hhs.tanf.cash.income.sources.earned"

policyengine_us/variables/gov/hhs/tanf/cash/income/tanf_gross_unearned_income.py

@@ -0,0 +1,13 @@
+from policyengine_us.model_api import *
+
+
+class tanf_gross_unearned_income(Variable):
+    value_type = float
+    entity = Person
+    label = (
+        "Temporary Assistance for Needy Families (TANF) gross unearned income"
+    )
+    unit = USD
+    definition_period = MONTH
+
+    adds = "gov.hhs.tanf.cash.income.sources.unearned"

policyengine_us/variables/gov/hhs/tanf/tanf_fpg.py

@@ -0,0 +1,28 @@
+from policyengine_us.model_api import *
+
+
+class tanf_fpg(Variable):
+    value_type = float
+    entity = SPMUnit
+    label = "TANF federal poverty guideline"
+    unit = USD
+    documentation = (
+        "The federal poverty guideline used to determine TANF eligibility."
+    )
+    definition_period = MONTH
+    reference = "https://www.law.cornell.edu/uscode/text/42/9902"
+
+    def formula(spm_unit, period, parameters):
+        # We can use tanf_unit_size for more accurate model in the future
+        n = spm_unit("spm_unit_size", period.this_year)
+        state_group = spm_unit.household("state_group_str", period.this_year)
+        year = period.start.year
+        month = period.start.month
+        if month >= 10:
+            instant_str = f"{year}-10-01"
+        else:
+            instant_str = f"{year - 1}-10-01"
+        p_fpg = parameters(instant_str).gov.hhs.fpg
+        p1 = p_fpg.first_person[state_group] / MONTHS_IN_YEAR
+        pn = p_fpg.additional_person[state_group] / MONTHS_IN_YEAR
+        return p1 + pn * (n - 1)