|
6 | 6 | import polars as pl |
7 | 7 | from polars.series.series import ArrayLike |
8 | 8 |
|
9 | | -from .bin import doane, freedman_diaconis, rice, scott, sqrt, sturges |
| 9 | +from .bin import auto, doane, freedman_diaconis, rice, scott, sqrt, sturges |
10 | 10 |
|
11 | | -BinMethod = Literal["doane", "fd", "rice", "sturges", "scott", "sqrt"] |
| 11 | +BinMethod = Literal["auto", "doane", "fd", "rice", "sturges", "scott", "sqrt"] |
12 | 12 |
|
13 | 13 |
|
14 | 14 | def _bin_count(x: pl.Series, bin_count: int | BinMethod) -> int: |
15 | 15 | if isinstance(bin_count, int): |
16 | 16 | return bin_count |
| 17 | + elif bin_count == "auto": |
| 18 | + return auto(x) |
17 | 19 | elif bin_count == "doane": |
18 | 20 | return doane(x) |
19 | 21 | elif bin_count == "fd": |
@@ -178,7 +180,7 @@ def psi( |
178 | 180 | current: ArrayLike, |
179 | 181 | *, |
180 | 182 | bins: list[float] | None = None, |
181 | | - bin_count: int | BinMethod = "fd", |
| 183 | + bin_count: int | BinMethod = "auto", |
182 | 184 | include_nulls: bool = True, |
183 | 185 | epsilon: float | None = 1e-4, |
184 | 186 | ) -> float: |
@@ -208,6 +210,7 @@ def psi( |
208 | 210 | If an integer, the number of bins. It can also be a string corresponding to an |
209 | 211 | auto-binning method, by default "fd". The possible methods are |
210 | 212 |
|
| 213 | + - "auto", see [rapidstats.bin.auto][] |
211 | 214 | - "doane", see [rapidstats.bin.doane][] |
212 | 215 | - "fd", see [rapidstats.bin.freedman_diaconis][] |
213 | 216 | - "rice", see [rapidstats.bin.rice][] |
|
0 commit comments