From feb80e9502125cd962362c90b8c013b090922bb5 Mon Sep 17 00:00:00 2001
From: 15b3 <30496251+15b3@users.noreply.github.com>
Date: Wed, 20 Oct 2021 22:33:54 +0900
Subject: [PATCH] Fix clean method in case var is nan.

---
 CONTRIBUTORS.md        |  1 +
 tests/test_windrose.py | 34 ++++++++++++++++++++++++++++++++++
 windrose/windrose.py   |  6 +++---
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
index 8e4eeb5..73ad0f2 100644
--- a/CONTRIBUTORS.md
+++ b/CONTRIBUTORS.md
@@ -14,6 +14,7 @@
 - Lionel Roubeyrie - LIMAIR - https://github.com/LionelR - original author
 - Bruno Ruas De Pinho - https://github.com/brunorpinho
 - Samuël Weber - Institut des Géosciences de l'Environnement, Grenoble, France, https://github.com/weber-s
+- 15b3 - https://github.com/15b3
 
 [Full Github contributors list](https://github.com/python-windrose/windrose/graphs/contributors).
 
diff --git a/tests/test_windrose.py b/tests/test_windrose.py
index 54a783b..3647918 100644
--- a/tests/test_windrose.py
+++ b/tests/test_windrose.py
@@ -22,9 +22,12 @@
 # from windrose import FIGSIZE_DEFAULT, DPI_DEFAULT
 from windrose import wrbar, wrbox, wrcontour, wrcontourf, wrpdf, wrscatter
 from windrose import plot_windrose
+from windrose import clean, clean_df
 
 import numpy as np
 import pandas as pd
+from numpy.testing import assert_allclose
+from pandas.testing import assert_frame_equal
 
 
 # Create wind speed and direction variables
@@ -208,3 +211,34 @@ def test_windrose_pd_not_default_names():
 #     #df = df.reset_index()
 #     #df = df.set_index([by, 'Timestamp'])
 #     #print(df)
+
+
+def test_windrose_clean():
+    direction = np.array(
+        [1.0, 1.0, 1.0, np.nan, np.nan, np.nan]
+    )
+    var = np.array(
+        [2.0, 0.0, np.nan, 2.0, 0.0, np.nan]
+    )
+    actual_direction, actual_var = clean(direction, var)
+    expected_direction = np.array([1.0])
+    expected_var = np.array([2.0])
+    assert_allclose(actual_direction, expected_direction)
+    assert_allclose(actual_var, expected_var)
+
+
+def test_windrose_clean_df():
+    df = pd.DataFrame(
+        {
+            "direction": [1.0, 1.0, 1.0, np.nan, np.nan, np.nan],
+            "speed": [2.0, 0.0, np.nan, 2.0, 0.0, np.nan],
+        }
+    )
+    actual_df = clean_df(df)
+    expected_df = pd.DataFrame(
+        {
+            "direction": [1.0],
+            "speed": [2.0],
+        }
+    )
+    assert_frame_equal(actual_df, expected_df)
diff --git a/windrose/windrose.py b/windrose/windrose.py
index 0fb4292..cf3354a 100644
--- a/windrose/windrose.py
+++ b/windrose/windrose.py
@@ -842,7 +842,7 @@ def clean_df(df, var=VAR_DEFAULT, direction=DIR_DEFAULT):
     removed from DataFrame
     if a direction is nan, this row is also removed from DataFrame
     """
-    return df[df[var].notnull() & df[var] != 0 & df[direction].notnull()]
+    return df[df[var].notnull() & (df[var] != 0) & df[direction].notnull()]
 
 
 def clean(direction, var, index=False):
@@ -853,8 +853,8 @@ def clean(direction, var, index=False):
     if a direction is nan, data is also removed from both array
     """
     dirmask = np.isfinite(direction)
-    varmask = var != 0 & np.isfinite(var)
-    mask = dirmask * varmask
+    varmask = (var != 0) & np.isfinite(var)
+    mask = dirmask & varmask
     if index is None:
         index = np.arange(mask.sum())
         return direction[mask], var[mask], index