Repository: nubank/fklearn (language: python)

Issue:
Double/Debiased Resusing Models when Curried
### Code sample
A common pattern in fklearn is to partially define a training function to create a `fit_fn` that will later be reused. For example, let's say we want to do feature selection. We can create a `fit_fn` passing all the parameters to a learner but the `feature_columns`

```python
fit_fn = non_parametric_double_ml_learner(
    df=train_set_eng.head(10000),
    treatment_column = treatment,
    outcome_column = outcome,
    debias_model = LGBMRegressor(),
    denoise_model = LGBMRegressor(),
    final_model = LGBMRegressor(),
    prediction_column= "dml_elast_prediction",
)
```
Then, we can try estimating the above model with different feature sets. For example, we can fit a model with only 3  (`m1`) features and another one with 6 features (`m2`). 

```python
f1 = features[:3]
f2 = features[:6]

m1 = fit_fn(feature_columns=f1)[0]
m2 = fit_fn(feature_columns=f2)[0]
````

However, once we try to make predictions with the first model, we get the following error

```python
m1(train_set_eng[f1].head())["dml_elast_prediction"]
```
```
ValueError: Number of features of the model must match the input. Model n_features_ is 6 and input n_features is 3 
```

So, `m1` is requiring 6 features? But we only trained it with 3? What is happening?

### Problem description
The problem is that fklearn is reusing the same instance of the model parameters in the learner function

```python
debias_model = LGBMRegressor(),
denoise_model = LGBMRegressor(),
final_model = LGBMRegressor(),
```

That way, since we train `m2` after `m1` and we reuse instance across both models, `m1` ends up having the model instance from `m2`, which requires 6 features. 

### Possible solutions
One easy solution is to use `from sklearn.base import clone` in the learner [here](https://github.com/nubank/fklearn/blob/ca1d8bd5b0667a5b9076fe02cc840f31e52b4ae1/src/fklearn/causal/cate_learning/double_machine_learning.py#L117)

```python
debias_model = GradientBoostingRegressor() if not debias_model else clone(debias_model)
denoise_model = GradientBoostingRegressor() if not denoise_model else clone(denoise_model)
final_model = GradientBoostingRegressor() if not final_model else clone(final_model)
```

That ensures we are always creating a new instance when calling this learner function. 

This is probably happening here too 
https://github.com/nubank/fklearn/blob/ca1d8bd5b0667a5b9076fe02cc840f31e52b4ae1/src/fklearn/causal/cate_learning/double_machine_learning.py#L29


Discussion / hints:

Fix patch:
diff --git a/src/fklearn/causal/cate_learning/double_machine_learning.py b/src/fklearn/causal/cate_learning/double_machine_learning.py
index d6aa0d54..fe6a9799 100644
--- a/src/fklearn/causal/cate_learning/double_machine_learning.py
+++ b/src/fklearn/causal/cate_learning/double_machine_learning.py
@@ -6,6 +6,7 @@
 from sklearn.base import RegressorMixin
 from sklearn.model_selection import KFold
 from sklearn.ensemble import GradientBoostingRegressor
+from sklearn.base import clone
 from toolz import curry
 from typing import Union
 
@@ -25,7 +26,7 @@ def _cv_estimate(model: RegressorMixin,
     cv_pred = pd.Series(np.nan, index=train_data.index)
 
     for train, test in cv.split(train_data):
-        m = model.fit(train_data[features].iloc[train], train_data[y].iloc[train])
+        m = clone(model, safe=False).fit(train_data[features].iloc[train], train_data[y].iloc[train])
         cv_pred.iloc[test] = m.predict(train_data[features].iloc[test])
         models += [m]
 
@@ -114,9 +115,9 @@ def non_parametric_double_ml_learner(df: pd.DataFrame,
 
     features = feature_columns if not encode_extra_cols else expand_features_encoded(df, feature_columns)
 
-    debias_model = GradientBoostingRegressor() if not debias_model else debias_model
-    denoise_model = GradientBoostingRegressor() if not denoise_model else denoise_model
-    final_model = GradientBoostingRegressor() if not final_model else final_model
+    debias_model = GradientBoostingRegressor() if debias_model is None else clone(debias_model, safe=False)
+    denoise_model = GradientBoostingRegressor() if denoise_model is None else clone(denoise_model, safe=False)
+    final_model = GradientBoostingRegressor() if final_model is None else clone(final_model, safe=False)
 
     t_hat, mts = _cv_estimate(debias_model, df,
                               features if debias_feature_columns is None else debias_feature_columns,
@@ -151,7 +152,7 @@ def p(new_df: pd.DataFrame) -> pd.DataFrame:
         'prediction_column': prediction_column,
         'package': "sklearn",
         'package_version': sk_version,
-        'feature_importance': dict(zip(features, model_final_fitted.feature_importances_)),
+        'feature_importance': None,
         'training_samples': len(df)},
         'debias_models': mts,
         'denoise_models': mys,


Test patch:
diff --git a/tests/causal/cate_learning/test_double_machine_learning.py b/tests/causal/cate_learning/test_double_machine_learning.py
index 3a226da6..1d3b2c9b 100644
--- a/tests/causal/cate_learning/test_double_machine_learning.py
+++ b/tests/causal/cate_learning/test_double_machine_learning.py
@@ -2,6 +2,8 @@
 
 import pandas as pd
 import numpy as np
+from sklearn.ensemble import RandomForestRegressor
+
 from fklearn.causal.cate_learning.double_machine_learning import non_parametric_double_ml_learner, _cv_estimate
 from sklearn.linear_model import LinearRegression
 
@@ -56,3 +58,30 @@ def test_non_parametric_double_ml_learner():
 
     assert Counter(expected_col_train) == Counter(pred_train.columns.tolist())
     assert Counter(expected_col_test) == Counter(pred_test.columns.tolist())
+
+
+def test_non_parametric_double_ml_learner_curry():
+
+    df_train = pd.DataFrame(dict(
+        # TE = 1       TE = 2      TE = 3
+        x1=[1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],
+        x2=[1, 2, 3, 1, 3, 1, 4, 1, 4, 1, 4, 1],
+        t=[1 + 1, 1 + 1, 2 + 1, 3 + 1, 1 + 2, 1 + 2, 2 + 2, 3 + 2, 1 + 3, 1 + 3, 2 + 3, 3 + 3],
+        y=[1 - 1, 1 - 1, 2 - 1, 3 - 1, 1 - 2, 1 - 2, 3 - 2, 5 - 2, 1 - 3, 1 - 3, 4 - 3, 7 - 3],
+    ))
+
+    np.random.seed(123)
+
+    fit_fn = non_parametric_double_ml_learner(df=df_train,
+                                              treatment_column="t",
+                                              outcome_column="y",
+                                              debias_model=RandomForestRegressor(),
+                                              denoise_model=RandomForestRegressor(),
+                                              final_model=RandomForestRegressor(),
+                                              prediction_column="test_prediction")
+
+    m1, m1_df, _ = fit_fn(feature_columns=["x1"])
+    m2, m2_df, _ = fit_fn(feature_columns=["x1", "x2"])
+
+    pd.testing.assert_frame_equal(m1_df, m1(df_train))
+    pd.testing.assert_frame_equal(m2_df, m2(df_train))