Repository: MaxGhenis/synthimpute

Issue:
Add observation weights to rf_impute
`sample_weight` in `fit` should do it: https://scikit-learn.org/stable/modules/generated/sklearn.ensemble.RandomForestRegressor.html#sklearn.ensemble.RandomForestRegressor.fit

Discussion:

Patch:
diff --git a/synthimpute/rf_impute.py b/synthimpute/rf_impute.py
index bdec8fc..06b22a4 100644
--- a/synthimpute/rf_impute.py
+++ b/synthimpute/rf_impute.py
@@ -45,7 +45,7 @@ def rf_quantile(m, X, q):
 
 
 def rf_impute(x_train, y_train, x_new, x_cols=None, random_state=None,
-              **kwargs):
+              sample_weight_train=None, **kwargs):
     """Impute labels from a training set to a new data set using 
        random forests quantile regression.
        
@@ -57,6 +57,8 @@ def rf_impute(x_train, y_train, x_new, x_cols=None, random_state=None,
             x_train (these must also be in x_new).
         random_state: Optional random seed passed to RandomForestRegressor and
             for uniform distribution of quantiles.
+        sample_weight_train: Vector indicating the weights associated with each
+            row of x_train/y_train. Defaults to None.
         **kwargs: Other args passed to RandomForestRegressor, e.g. 
             `n_estimators=50`.  rf_impute uses all RandomForestRegressor
             defaults unless otherwise specified.
@@ -65,7 +67,10 @@ def rf_impute(x_train, y_train, x_new, x_cols=None, random_state=None,
         Imputed labels for new_x.
     """
     rf = ensemble.RandomForestRegressor(random_state=random_state, **kwargs)
-    rf.fit(x_train, y_train)
+    if sample_weight_train is None:
+        rf.fit(x_train, y_train)
+    else:
+        rf.fit(x_train, y_train, sample_weight=sample_weight_train)
     if random_state is not None:
         np.random.seed(random_state)
     quantiles = np.random.rand(y_train.size)  # Uniform distribution.


Test patch:
diff --git a/synthimpute/tests/test_rf_impute.py b/synthimpute/tests/test_rf_impute.py
index c02b6f0..ce99542 100644
--- a/synthimpute/tests/test_rf_impute.py
+++ b/synthimpute/tests/test_rf_impute.py
@@ -12,7 +12,14 @@ def test_rf_impute():
                       'x2': np.random.randn(n)})
     # Construct example relationship.
     y = x.x1 + np.power(x.x2, 3) + np.random.randn(n)
-    si.rf_impute(x.iloc[:N_TRAIN], y.iloc[:N_TRAIN], x.iloc[N_TRAIN:])
+    # Split into test and train.
+    x_train = x.iloc[:N_TRAIN]
+    y_train = y.iloc[:N_TRAIN]
+    x_test = x.iloc[N_TRAIN:]
+    si.rf_impute(x_train, y_train, x_test)
     # Try with some args.
-    si.rf_impute(x.iloc[:N_TRAIN], y.iloc[:N_TRAIN], x.iloc[N_TRAIN:],
+    si.rf_impute(x_train, y_train, x_test,
                  random_state=10, n_estimators=200)
+    # Try with sample_weight_train.
+    si.rf_impute(x_train, y_train, x_test,
+                 sample_weight_train=np.random.randn(N_TRAIN))