Fix clean method in case var is nan.

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).
 

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)

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