Shift to branch of US

PolicyEngine · Sep 22, 2024 · 5a3f94d · 5a3f94d
1 parent c35a21c
commit 5a3f94d
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 13 deletions.
diff --git a/policyengine_us_data/datasets/cps/enhanced_cps.py b/policyengine_us_data/datasets/cps/enhanced_cps.py
@@ -99,13 +99,13 @@ def train_previous_year_income_model():
     df = sim.calculate_dataframe(VARIABLES + OUTPUTS, 2019, map_to="person")
     df_train = df[df.previous_year_income_available]
 
-    from survey_enhance import Imputation
+    from policyengine_us_data.utils import QRF
 
-    income_last_year = Imputation()
+    income_last_year = QRF()
     X = df_train[VARIABLES[1:]]
     y = df_train[OUTPUTS]
 
-    income_last_year.train(X, y)
+    income_last_year.fit(X, y)
 
     return income_last_year
 

diff --git a/policyengine_us_data/datasets/cps/extended_cps.py b/policyengine_us_data/datasets/cps/extended_cps.py
@@ -101,7 +101,7 @@ def generate(self):
         y = pd.DataFrame(columns=IMPUTED_VARIABLES, index=X.index)
 
         model = QRF()
-        model.train(
+        model.fit(
             X_train,
             y_train,
         )

diff --git a/policyengine_us_data/datasets/puf/puf.py b/policyengine_us_data/datasets/puf/puf.py
@@ -22,14 +22,13 @@ def impute_pension_contributions_to_puf(puf_df):
         ["employment_income", "household_weight", "pre_tax_contributions"]
     )
 
-    from survey_enhance import Imputation
+    from policyengine_us_data.utils import QRF
 
-    pension_contributions = Imputation()
+    pension_contributions = QRF()
 
-    pension_contributions.train(
+    pension_contributions.fit(
         X=cps_df[["employment_income"]],
         Y=cps_df[["pre_tax_contributions"]],
-        sample_weight=cps_df["household_weight"],
     )
     return pension_contributions.predict(
         X=puf_df[["employment_income"]],
@@ -39,7 +38,7 @@ def impute_pension_contributions_to_puf(puf_df):
 def impute_missing_demographics(
     puf: pd.DataFrame, demographics: pd.DataFrame
 ) -> pd.DataFrame:
-    from survey_enhance import Imputation
+    from policyengine_us_data.utils import QRF
 
     puf_with_demographics = (
         puf[puf.RECID.isin(demographics.RECID)]
@@ -63,9 +62,9 @@ def impute_missing_demographics(
         "XTOT",
     ]
 
-    demographics_from_puf = Imputation()
+    demographics_from_puf = QRF()
 
-    demographics_from_puf.train(
+    demographics_from_puf.fit(
         X=puf_with_demographics[NON_DEMOGRAPHIC_VARIABLES],
         Y=puf_with_demographics[DEMOGRAPHIC_VARIABLES],
     )

diff --git a/pyproject.toml b/pyproject.toml
@@ -22,8 +22,8 @@ dependencies = [
 dev = [
     "black",
     "pytest",
-    "policyengine_us==1.71.1",
-    "survey_enhance",
+    "policyengine_us @ git+https://github.com/policyengine/policyengine-us@us-rent-data",
+    "quantile-forest",
     "torch",
     "tables",
     "tabulate",