Add 'test lite' mode to tests

.github/workflows/ci_cd.yaml

@@ -42,6 +42,8 @@ jobs:
           POLICYENGINE_US_DATA_GITHUB_TOKEN: ${{ secrets.POLICYENGINE_US_DATA_GITHUB_TOKEN }}
       - name: Build datasets
         run: make data
+        env:
+          TEST_LITE: true
       - name: Run tests
         run: pytest
       - name: Test documentation builds

policyengine_us_data/datasets/cps/enhanced_cps.py

@@ -15,6 +15,7 @@
     CPS_2024,
 )
 import torch
+import os
 
 
 def reweight(
@@ -51,7 +52,7 @@ def loss(weights):
     optimizer = torch.optim.Adam([weights], lr=1e-2)
     from tqdm import trange
 
-    iterator = trange(10_000)
+    iterator = trange(10_000) if not os.environ.get("TEST_LITE") else trange(100)
     for i in iterator:
         optimizer.zero_grad()
         l, worst_name, worst_val = loss(torch.exp(weights))

policyengine_us_data/datasets/cps/extended_cps.py

@@ -4,6 +4,7 @@
 from .cps import *
 from ..puf import *
 import pandas as pd
+import os
 
 # These are sorted by magnitude.
 # First 15 contain 90%.
@@ -110,6 +111,9 @@ def generate(self):
         data = cps_sim.dataset.load_dataset()
         new_data = {}
 
+        if os.environ.get("TEST_LITE"):
+            IMPUTED_VARIABLES = IMPUTED_VARIABLES[:1]
+
         for variable in list(data) + IMPUTED_VARIABLES:
             variable_metadata = cps_sim.tax_benefit_system.variables.get(
                 variable