ploomber / sklearn-evaluation

Machine learning model evaluation made easy: plots, tables, HTML reports, experiment tracking and Jupyter notebook analysis.
https://sklearn-evaluation.ploomber.io
Apache License 2.0
455 stars 54 forks source link

[ENH:] Cumulative Gain Chart: Allow single line version #345

Open emrynHofmannElephant opened 5 months ago

emrynHofmannElephant commented 5 months ago

Just wondering if there could be a small change to the function here: https://github.com/ploomber/sklearn-evaluation/blob/96456eed27a77abb69cead1f4d9c483d45fdaf61/src/sklearn_evaluation/plot/cumulative_gain_lift_curve.py#L99

So that it allows for one set of predicted probabilities (instead of requiring two). In particular this would require a change from y_score to y_pred, and likely either the removal of the second call to _cumulative_gain_curve() on line 164 & 171.

As this function is limited to the binary case anyway, the second line is redundant, as it is just the inverse probability of the other. Therefore, in case both lines are needed, a simple input check (y_score.shape[1]) & calculation could generate the second set instead?