Impute pregnancy

Fixes #162

Author: PavelMakarchuk

policyengine_us_data/datasets/cps/cps.py

@@ -229,8 +229,23 @@ def add_personal_variables(cps: h5py.File, person: DataFrame) -> None:
         person.A_AGE == 80,
         # NB: randint is inclusive of first argument, exclusive of second.
         np.random.randint(80, 85, len(person)),
-        person.A_AGE,
+        np.where(
+            person.A_AGE == 0,
+            np.where(
+                np.random.randint(0, 2, len(person)), # Random number of 0 or 1
+                # If 1 is flipped, select a random number between -0.75 and 0
+                # This will represent the pregnany month 
+                # At -0.75 the pregnancy month is 0 and at -0.0001 the pregnancy month is 9
+                np.random.uniform(-0.75, 0, len(person)), 
+                # If 0 is flipped, the child is a newborn at the age of 0 to 1 
+                np.random.uniform(0, 1, len(person)),
+            ),
+            person.A_AGE,
+        )
     )
+    cps["is_pregnant"] = (cps["age"] >= -0.75) & (cps["age"] < 0)
+    cps["is_newborn"] = (cps["age"] >= 0) & (cps["age"] < 1)
+
     # A_SEX is 1 -> male, 2 -> female.
     cps["is_female"] = person.A_SEX == 2
     # "Is...blind or does...have serious difficulty seeing even when Wearing

policyengine_us_data/utils/loss.py

@@ -254,6 +254,8 @@ def build_loss_matrix(dataset: type, time_period):
         # Rough estimate, not CPS derived
         "real_estate_taxes": 400e9,  # Rough estimate between 350bn and 600bn total property tax collections
         "rent": 735e9,  # ACS total uprated by CPI
+        "is_newborn": 3_491_679, # ACS total of people aged 0
+        "is_pregnant": 2_618_759, # 75% of the ACS total of people aged 0
     }
 
     for variable_name, target in HARD_CODED_TOTALS.items():
@@ -340,6 +342,19 @@ def build_loss_matrix(dataset: type, time_period):
         )
         targets_array.append(row["population_under_5"])
 
+    # Population by number of newborns and pregancies
+
+    age = sim.calculate("age").values
+    newborns = (age >= 0) & (age < 1)
+    label = "census/newborns"
+    loss_matrix[label] = sim.map_result(newborns, "person", "household")
+    targets_array.append(HARD_CODED_TOTALS["is_newborn"])
+
+    pregnancies = (age >= -0.75) & (age < 0)
+    label = "census/pregnancies"
+    loss_matrix[label] = sim.map_result(pregnancies, "person", "household")
+    targets_array.append(HARD_CODED_TOTALS["is_pregnant"])
+
     if any(loss_matrix.isna().sum() > 0):
         raise ValueError("Some targets are missing from the loss matrix")