[BUG] zero_inflated_regressor.py

koaning / scikit-lego

Extra blocks for scikit-learn pipelines.

https://koaning.github.io/scikit-lego/

MIT License

1.28k stars 117 forks source link

Closed jcoding2022 closed 6 months ago

jcoding2022 commented 6 months ago

The ZeroInflatedRegressor class consists of a classifier and a regressor and it says in the documentation that "The regressor is only trained on examples where the target is non-zero, which makes it easier for it to focus." In the code, however, the regressor is actually trained on examples where the classifier-predicted target is non-zero.

Just replace line 111 with the following code:
```
non_zero_indices = np.where(y != 0)[0]
```

FBruzzesi commented 6 months ago

Oh gosh! Thanks for reporting this! I don't think it was intended, but will double check and get back to you or fix it in the next release 👌