Repository: PolicyEngine/policyengine-us-data
Pull request #208: Don't downsample Pooled CPS
Author: login: anth-volk
Created: 2025-04-11T22:10:41Z | Merged: 2025-04-22T18:24:26Z

Fixes #207.

This PR introduces a _very_ naïve means of preventing downsampling when constructing the Pooled 3-Year CPS dataset, which we use by default for states due to states' limited number of observations in any given single year CPS set. This PR creates three new datasets - essentially the old un-downsampled 2021, '22, and '23 CPS - then uses those to construct the Pooled 3-Year while maintaining the default downsampled versions of these, as well. 

This code should be improved separately, and if this code is still here, unchanged, on July 1, 2025, someone should open a ticket.

Linked issue #207: Stop downsampling pooled 3-year CPS
Linked issue comments: user: login: anth-volk
created_at: 2025-04-11T20:40:04Z
body: To confirm: looks like we are subsampling the pooled 3-year CPS. The mechanism works like this:
* In the base `CPS` class, we downsample [here](https://github.com/PolicyEngine/policyengine-us-data/blob/fe4acf9a1cc75881eff29aa80f58d79b5121e7cf/policyengine_us_data/datasets/cps/cps.py#L61) when generating any dataset that inherits from this class
* The `CPS_2021`, `CPS_2022`, and `CPS_2023` datasets, which inherit from `CPS`, merely specify some generation inputs, meaning that they are subsampled when generated [here](https://github.com/PolicyEngine/policyengine-us-data/blob/fe4acf9a1cc75881eff29aa80f58d79b5121e7cf/policyengine_us_data/datasets/cps/cps.py#L641)
* The `PooledCPS` class merely merges together datasets, and thus the `Pooled_3_Year_CPS_2023`, which inherits from it, just merges together the subsampled 2021, 2022, and 2023 datasets
author_association: COLLABORATOR

Diff:
diff --git a/changelog_entry.yaml b/changelog_entry.yaml
index e69de29bb..6ff4a1976 100644
--- a/changelog_entry.yaml
+++ b/changelog_entry.yaml
@@ -0,0 +1,8 @@
+- bump: patch
+  changes:
+    added:
+    - Non-downsampled versions of the 2021, 2022, and 2023 CPS datasets
+    changed:
+    - Modified downsampling method within CPS base dataset class
+    - Pooled 3-Year CPS generation uses the non-downsampled versions of the 2021, 2022, and 2023 CPS datasets
+    - Downsampling method attempts to preserve original dtype values
\ No newline at end of file
diff --git a/policyengine_us_data/datasets/cps/cps.py b/policyengine_us_data/datasets/cps/cps.py
index 02e04e418..26579b826 100644
--- a/policyengine_us_data/datasets/cps/cps.py
+++ b/policyengine_us_data/datasets/cps/cps.py
@@ -20,10 +20,15 @@ class CPS(Dataset):
     raw_cps: Type[CensusCPS] = None
     previous_year_raw_cps: Type[CensusCPS] = None
     data_format = Dataset.ARRAYS
+    frac: float | None = 1
 
     def generate(self):
         """Generates the Current Population Survey dataset for PolicyEngine US microsimulations.
         Technical documentation and codebook here: https://www2.census.gov/programs-surveys/cps/techdocs/cpsmar21.pdf
+
+        Args:
+            frac (float, optional): Fraction of the dataset to keep. Defaults to 1. Example: To downsample to 25% of dataset,
+                set frac=0.25.
         """
 
         if self.raw_cps is None:
@@ -57,19 +62,42 @@ def generate(self):
         add_takeup(self)
 
         # Downsample
+        if self.frac is not None and self.frac < 1.0:
+            self.downsample(frac=self.frac)
 
-        self.downsample(fraction=0.5)
-
-    def downsample(self, fraction: float = 0.5):
+    def downsample(self, frac: float):
         from policyengine_us import Microsimulation
 
-        sim = Microsimulation(dataset=self)
-        sim.subsample(frac=fraction)
+        # Store original dtypes before modifying
         original_data: dict = self.load_dataset()
+        original_dtypes = {
+            key: original_data[key].dtype for key in original_data
+        }
+
+        sim = Microsimulation(dataset=self)
+        sim.subsample(frac=frac)
+
         for key in original_data:
             if key not in sim.tax_benefit_system.variables:
                 continue
-            original_data[key] = sim.calculate(key).values
+            values = sim.calculate(key).values
+
+            # Preserve the original dtype if possible
+            if (
+                key in original_dtypes
+                and hasattr(values, "dtype")
+                and values.dtype != original_dtypes[key]
+            ):
+                try:
+                    original_data[key] = values.astype(original_dtypes[key])
+                except:
+                    # If conversion fails, log it but continue
+                    print(
+                        f"Warning: Could not convert {key} back to {original_dtypes[key]}"
+                    )
+                    original_data[key] = values
+            else:
+                original_data[key] = values
 
         self.save_dataset(original_data)
 
@@ -627,6 +655,7 @@ class CPS_2019(CPS):
     previous_year_raw_cps = CensusCPS_2018
     file_path = STORAGE_FOLDER / "cps_2019.h5"
     time_period = 2019
+    frac = 0.5
 
 
 class CPS_2020(CPS):
@@ -636,6 +665,7 @@ class CPS_2020(CPS):
     previous_year_raw_cps = CensusCPS_2019
     file_path = STORAGE_FOLDER / "cps_2020.h5"
     time_period = 2020
+    frac = 0.5
 
 
 class CPS_2021(CPS):
@@ -645,6 +675,7 @@ class CPS_2021(CPS):
     previous_year_raw_cps = CensusCPS_2020
     file_path = STORAGE_FOLDER / "cps_2021_v1_6_1.h5"
     time_period = 2021
+    frac = 0.5
 
 
 class CPS_2022(CPS):
@@ -654,6 +685,7 @@ class CPS_2022(CPS):
     previous_year_raw_cps = CensusCPS_2021
     file_path = STORAGE_FOLDER / "cps_2022_v1_6_1.h5"
     time_period = 2022
+    frac = 0.5
 
 
 class CPS_2023(CPS):
@@ -663,6 +695,7 @@ class CPS_2023(CPS):
     previous_year_raw_cps = CensusCPS_2022
     file_path = STORAGE_FOLDER / "cps_2023.h5"
     time_period = 2023
+    frac = 0.5
 
 
 class CPS_2024(CPS):
@@ -671,6 +704,37 @@ class CPS_2024(CPS):
     file_path = STORAGE_FOLDER / "cps_2024.h5"
     time_period = 2024
     url = "release://policyengine/policyengine-us-data/1.13.0/cps_2024.h5"
+    frac = 0.5
+
+
+# The below datasets are a very naïve way of preventing downsampling in the
+# Pooled 3-Year CPS. They should be replaced by a more sustainable approach.
+# If these are still here on July 1, 2025, please open an issue and raise at standup.
+class CPS_2021_Full(CPS):
+    name = "cps_2021_full"
+    label = "CPS 2021 (full)"
+    raw_cps = CensusCPS_2021
+    previous_year_raw_cps = CensusCPS_2020
+    file_path = STORAGE_FOLDER / "cps_2021_full.h5"
+    time_period = 2021
+
+
+class CPS_2022_Full(CPS):
+    name = "cps_2022_full"
+    label = "CPS 2022 (full)"
+    raw_cps = CensusCPS_2022
+    previous_year_raw_cps = CensusCPS_2021
+    file_path = STORAGE_FOLDER / "cps_2022_full.h5"
+    time_period = 2022
+
+
+class CPS_2023_Full(CPS):
+    name = "cps_2023_full"
+    label = "CPS 2023 (full)"
+    raw_cps = CensusCPS_2023
+    previous_year_raw_cps = CensusCPS_2022
+    file_path = STORAGE_FOLDER / "cps_2023_full.h5"
+    time_period = 2023
 
 
 class PooledCPS(Dataset):
@@ -724,9 +788,9 @@ class Pooled_3_Year_CPS_2023(PooledCPS):
     name = "pooled_3_year_cps_2023"
     file_path = STORAGE_FOLDER / "pooled_3_year_cps_2023.h5"
     input_datasets = [
-        CPS_2021,
-        CPS_2022,
-        CPS_2023,
+        CPS_2021_Full,
+        CPS_2022_Full,
+        CPS_2023_Full,
     ]
     time_period = 2023
     url = "hf://policyengine/policyengine-us-data/pooled_3_year_cps_2023.h5"
@@ -737,4 +801,7 @@ class Pooled_3_Year_CPS_2023(PooledCPS):
     CPS_2022().generate()
     CPS_2023().generate()
     CPS_2024().generate()
+    CPS_2021_Full().generate()
+    CPS_2022_Full().generate()
+    CPS_2023_Full().generate()
     Pooled_3_Year_CPS_2023().generate()


Reviews:
user: login: MaxGhenis
body: sure why not
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T23:59:11Z
user: login: nikhilwoodruff
body: Just a few minor fixes, then LGTM
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T14:52:35Z
user: login: anth-volk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T20:58:39Z
user: login: MaxGhenis
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T21:06:59Z
user: login: anth-volk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T19:16:46Z
user: login: nikhilwoodruff
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T15:26:30Z

Review comments:
diff_hunk: @@ -673,6 +711,39 @@ class CPS_2024(CPS):
     url = "release://policyengine/policyengine-us-data/1.13.0/cps_2024.h5"
 
 
+# The below datasets are a very naïve way of preventing downsampling in the
+# Pooled 3-Year CPS. They should be replaced by a more sustainable approach.
+# If these are still here on July 1, 2025, please open an issue and raise at standup.
+class CPS_2021_Not_Downsampled(CPS):
path: policyengine_us_data/datasets/cps/cps.py
user: login: nikhilwoodruff
body: Make this `CPS_2021_Full`
created_at: 2025-04-14T14:49:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -58,18 +59,55 @@ def generate(self):
 
         # Downsample
 
-        self.downsample(fraction=0.5)
+        if self.downsample_by_half:
+            self.downsample(fraction=0.5)
+
+    # def downsample(self, fraction: float = 0.5):
+    #     from policyengine_us import Microsimulation
+
+    #     sim = Microsimulation(dataset=self)
+    #     sim.subsample(frac=fraction)
+    #     original_data: dict = self.load_dataset()
+    #     for key in original_data:
+    #         if key not in sim.tax_benefit_system.variables:
+    #             continue
+    #         original_data[key] = sim.calculate(key).values
+
+    #     self.save_dataset(original_data)
path: policyengine_us_data/datasets/cps/cps.py
user: login: nikhilwoodruff
body: Remove
created_at: 2025-04-14T14:50:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,6 +20,7 @@ class CPS(Dataset):
     raw_cps: Type[CensusCPS] = None
     previous_year_raw_cps: Type[CensusCPS] = None
     data_format = Dataset.ARRAYS
+    downsample_by_half: bool = True
path: policyengine_us_data/datasets/cps/cps.py
user: login: nikhilwoodruff
body: Make this a percentage
created_at: 2025-04-14T14:52:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,6 +20,7 @@ class CPS(Dataset):
     raw_cps: Type[CensusCPS] = None
     previous_year_raw_cps: Type[CensusCPS] = None
     data_format = Dataset.ARRAYS
+    downsample_by_half: bool = True
path: policyengine_us_data/datasets/cps/cps.py
user: login: anth-volk
body: Readjusted how this is done. I'd love you to confirm that if the downsampling fraction is `1.0`, that indicates "no downsampling," and that a theoretical fraction of `0.0` would mean "return an empty dataset".
created_at: 2025-04-14T20:58:39Z
author_association: COLLABORATOR
diff_hunk: @@ -20,6 +20,7 @@ class CPS(Dataset):
     raw_cps: Type[CensusCPS] = None
     previous_year_raw_cps: Type[CensusCPS] = None
     data_format = Dataset.ARRAYS
+    downsample_by_half: bool = True
path: policyengine_us_data/datasets/cps/cps.py
user: login: MaxGhenis
body: Let's call it `sample_frac` (or even just `frac`) to match [pandas](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.sample.html). Defaults to 1. If None, equals 1 for full sample.
created_at: 2025-04-14T21:06:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,6 +20,7 @@ class CPS(Dataset):
     raw_cps: Type[CensusCPS] = None
     previous_year_raw_cps: Type[CensusCPS] = None
     data_format = Dataset.ARRAYS
+    downsample_by_half: bool = True
path: policyengine_us_data/datasets/cps/cps.py
user: login: anth-volk
body: I've done the following:
* Renamed to `frac` (in line with `-core`)
* Renamed the arg within `self.downsample()` to also be `frac`
* Applied a `frac` of 0.5 to all standard CPS datasets
* Removed a `frac` for the `CPS_20XX_Full` instances
created_at: 2025-04-16T19:16:46Z
author_association: COLLABORATOR

PR comments:
user: login: anth-volk
created_at: 2025-04-11T22:19:01Z
body: Tests not running due to expiry of GitHub PAT
author_association: COLLABORATOR
user: login: anth-volk
created_at: 2025-04-12T00:39:06Z
body: Tests now failing due to some sort of county FIPS error; unclear what's going on
author_association: COLLABORATOR
user: login: anth-volk
created_at: 2025-04-12T01:04:07Z
body: Attempting to fix in most recent push
author_association: COLLABORATOR