Run lite versions of imputations and calibrations in tests (#251)

* Run lite versions of imputations and calibrations in tests
Fixes #250

* Downsample SCF data too

* Add more logging

* Don't tune hyperparameters on PR tests

* Use correct env var

* Increase epochs

* Impute all extended CPS variables

Author: nikhilwoodruff

.github/workflows/pr_code_changes.yaml

@@ -42,7 +42,7 @@ jobs:
           - name: Build datasets
             run: make data
             env:
-              LITE_MODE: true
+              TEST_LITE: true
           - name: Run tests
             run: pytest
           - name: Test documentation builds

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    fixed:
+    - Runtime for tests reduced.

policyengine_us_data/datasets/cps/cps.py

@@ -14,6 +14,8 @@
 from policyengine_us_data.utils import QRF
 import logging
 
+test_lite = os.environ.get("TEST_LITE")
+
 
 class CPS(Dataset):
     name = "cps"
@@ -49,21 +51,33 @@ def generate(self):
             raw_data[entity] for entity in ENTITIES
         ]
 
+        logging.info("Adding ID variables")
         add_id_variables(cps, person, tax_unit, family, spm_unit, household)
+        logging.info("Adding personal variables")
         add_personal_variables(cps, person)
+        logging.info("Adding personal income variables")
         add_personal_income_variables(cps, person, self.raw_cps.time_period)
+        logging.info("Adding previous year income variables")
         add_previous_year_income(self, cps)
+        logging.info("Adding SSN card type")
         add_ssn_card_type(cps, person)
+        logging.info("Adding family variables")
         add_spm_variables(cps, spm_unit)
+        logging.info("Adding household variables")
         add_household_variables(cps, household)
+        logging.info("Adding rent")
         add_rent(self, cps, person, household)
+        logging.info("Adding auto loan balance")
         add_auto_loan_balance(self, cps)
+        logging.info("Adding tips")
         add_tips(self, cps)
+        logging.info("Added all variables")
 
         raw_data.close()
         self.save_dataset(cps)
-
+        logging.info("Adding takeup")
         add_takeup(self)
+        logging.info("Downsampling")
 
         # Downsample
         if self.frac is not None and self.frac < 1.0:
@@ -146,7 +160,9 @@ def add_rent(self, cps: h5py.File, person: DataFrame, household: DataFrame):
         },
         na_action="ignore",
     ).fillna(train_df.tenure_type)
-    train_df = train_df[train_df.is_household_head].sample(100_000)
+    train_df = train_df[train_df.is_household_head].sample(
+        100_000 if not test_lite else 1_000
+    )
     inference_df = cps_sim.calculate_dataframe(PREDICTORS)
     mask = inference_df.is_household_head.values
     inference_df = inference_df[mask]
@@ -290,7 +306,7 @@ def add_auto_loan_balance(self, cps: h5py.File) -> None:
     donor_data = donor_data.loc[
         np.random.choice(
             donor_data.index,
-            size=100_000,
+            size=100_000 if not test_lite else 1_000,
             replace=True,
             p=donor_data.household_weight / donor_data.household_weight.sum(),
         )
@@ -303,7 +319,7 @@ def add_auto_loan_balance(self, cps: h5py.File) -> None:
         X_train=donor_data,
         predictors=PREDICTORS,
         imputed_variables=IMPUTED_VARIABLES,
-        tune_hyperparameters=True,
+        tune_hyperparameters=not test_lite,
     )
 
     imputations = fitted_model.predict(X_test=receiver_data)

policyengine_us_data/datasets/cps/enhanced_cps.py

@@ -14,6 +14,7 @@
     CPS_2019,
     CPS_2024,
 )
+import os
 
 try:
     import torch
@@ -66,7 +67,7 @@ def dropout_weights(weights, p):
 
     start_loss = None
 
-    iterator = trange(5_000)
+    iterator = trange(5_000 if not os.environ.get("TEST_LITE") else 1_000)
     for i in iterator:
         optimizer.zero_grad()
         weights_ = dropout_weights(weights, dropout_rate)
@@ -88,6 +89,9 @@ def train_previous_year_income_model():
 
     sim = Microsimulation(dataset=CPS_2019)
 
+    if os.environ.get("TEST_LITE"):
+        sim.subsample(1_000)
+
     VARIABLES = [
         "previous_year_income_available",
         "employment_income",

policyengine_us_data/datasets/cps/extended_cps.py

@@ -71,9 +71,6 @@
     "deductible_mortgage_interest",
 ]
 
-if os.environ.get("TEST_LITE"):
-    IMPUTED_VARIABLES = IMPUTED_VARIABLES[:7]
-
 
 class ExtendedCPS(Dataset):
     cps: Type[CPS]
@@ -86,8 +83,8 @@ def generate(self):
         cps_sim = Microsimulation(dataset=self.cps)
         puf_sim = Microsimulation(dataset=self.puf)
 
-        if os.environ.get("LITE_MODE"):
-            puf_sim.subsample(10_000)
+        if os.environ.get("TEST_LITE"):
+            puf_sim.subsample(1_000)
 
         INPUTS = [
             "age",

policyengine_us_data/datasets/puf/puf.py

@@ -9,6 +9,7 @@
 from policyengine_us_data.utils.uprating import (
     create_policyengine_uprating_factors_table,
 )
+import os
 
 rng = np.random.default_rng(seed=64)
 
@@ -18,6 +19,8 @@ def impute_pension_contributions_to_puf(puf_df):
     from policyengine_us_data.datasets.cps import CPS_2021
 
     cps = Microsimulation(dataset=CPS_2021)
+    if os.environ.get("TEST_LITE"):
+        cps.subsample(1_000)
     cps_df = cps.calculate_dataframe(
         ["employment_income", "household_weight", "pre_tax_contributions"]
     )
@@ -46,6 +49,11 @@ def impute_missing_demographics(
         .fillna(0)
     )
 
+    if os.environ.get("TEST_LITE"):
+        puf_with_demographics = puf_with_demographics.sample(
+            n=1_000, random_state=0
+        )
+
     DEMOGRAPHIC_VARIABLES = [
         "AGEDP1",
         "AGEDP2",

policyengine_us_data/datasets/sipp/sipp.py

@@ -6,6 +6,9 @@
 from policyengine_us_data.storage import STORAGE_FOLDER
 import pickle
 from huggingface_hub import hf_hub_download
+import os
+
+test_lite = os.environ.get("TEST_LITE")
 
 
 def train_tip_model():
@@ -100,7 +103,7 @@ def train_tip_model():
     sipp = sipp.loc[
         np.random.choice(
             sipp.index,
-            size=100_000,
+            size=100_000 if not test_lite else 1_000,
             replace=True,
             p=sipp.household_weight / sipp.household_weight.sum(),
         )