Skip to content

Commit

Permalink
rowidcf correction (#46)
Browse files Browse the repository at this point in the history
* rowidcf bug correction

* test number of unique rowidcf
  • Loading branch information
LamAdr authored Nov 23, 2023
1 parent ce6c16b commit ad449a6
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 1 deletion.
3 changes: 2 additions & 1 deletion marginaleffects/datagrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,8 @@ def datagridcf(model=None, newdata=None, **kwargs):
result = newdata.join(df_cross, how = "cross")

# Create rowid and rowidcf
result = result.with_columns(pl.Series(range(result.shape[0])).alias("rowidcf"))
rowidcf = [i for i in range(newdata.shape[0]) for _ in range(result.select(pl.count()).item() // newdata.select(pl.count()).item())]
result = result.with_columns(pl.Series(rowidcf).alias("rowidcf"))

result.datagrid_explicit = list(kwargs.keys())

Expand Down
1 change: 1 addition & 0 deletions tests/test_datagrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,4 @@ def test_simple_grid():
def test_cf():
assert datagrid(newdata = mtcars, mpg = 32).shape[0] == 1
assert datagridcf(newdata = mtcars, mpg = [30, 32]).shape[0] == 64
assert datagridcf(newdata = mtcars, mpg = [30, 32]).unique("rowidcf").shape[0] == 32

0 comments on commit ad449a6

Please sign in to comment.