Skip to content

Commit

Permalink
Fix: Make PolyFit stat robust to missing data (#3010)
Browse files Browse the repository at this point in the history
* Fix: Make PolyFit stat robust to missing data

Fixes #2992

* Update release notes
  • Loading branch information
mwaskom committed Sep 12, 2022
1 parent 15713b3 commit c5c5f4a
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 1 deletion.
5 changes: 5 additions & 0 deletions doc/whatsnew/v0.12.1.rst
@@ -0,0 +1,5 @@

v0.12.1 (Unreleased)
--------------------

- |Fix| Make :class:`objects.PolyFit` robust to missing data (:pr:`3010`).
5 changes: 4 additions & 1 deletion seaborn/_stats/regression.py
Expand Up @@ -38,7 +38,10 @@ def _fit_predict(self, data):

def __call__(self, data, groupby, orient, scales):

return groupby.apply(data, self._fit_predict)
return (
groupby
.apply(data.dropna(subset=["x", "y"]), self._fit_predict)
)


@dataclass
Expand Down
9 changes: 9 additions & 0 deletions tests/_stats/test_regression.py
Expand Up @@ -4,6 +4,7 @@

import pytest
from numpy.testing import assert_array_equal, assert_array_almost_equal
from pandas.testing import assert_frame_equal

from seaborn._core.groupby import GroupBy
from seaborn._stats.regression import PolyFit
Expand Down Expand Up @@ -50,3 +51,11 @@ def test_one_grouper(self, df):
grid = np.linspace(part["x"].min(), part["x"].max(), gridsize)
assert_array_equal(part["x"], grid)
assert part["y"].diff().diff().dropna().abs().gt(0).all()

def test_missing_data(self, df):

groupby = GroupBy(["group"])
df.iloc[5:10] = np.nan
res1 = PolyFit()(df[["x", "y"]], groupby, "x", {})
res2 = PolyFit()(df[["x", "y"]].dropna(), groupby, "x", {})
assert_frame_equal(res1, res2)

0 comments on commit c5c5f4a

Please sign in to comment.