Merge pull request #229 from knaaptime/colorarray

fix nan handling in color array

mapclassify/tests/test_rgba.py

@@ -6,14 +6,28 @@
 world = geopandas.read_file(
     "https://naciscdn.org/naturalearth/110m/cultural/ne_110m_admin_0_countries.zip"
 )
+world['area'] = world.area
 
+# columns are equivalent except for nan in the first position
+world['nanarea'] = world.area
+world.loc[0, 'nanarea'] = np.nan
 
 def test_rgba():
-    colors = get_color_array(world.area, cmap="viridis")[0]
-    assert_array_equal(colors, np.array([68, 1, 84, 255])) 
-
+    colors = get_color_array(world['area'], cmap="viridis")[-1]
+    assert_array_equal(colors, np.array([94, 201,  97, 255])) 
 
-def test_rgba():
-    colors = get_color_array(world.area, cmap="viridis", as_hex=True)[0]
-    assert_array_equal(colors,'#440154') 
+def test_rgba_hex():
+    colors = get_color_array(world['area'], cmap="viridis", as_hex=True)[-1]
+    assert_array_equal(colors,'#5ec961') 
 
+def test_rgba_nan():
+    colors = get_color_array(world['nanarea'], cmap="viridis", nan_color=[0,0,0,0])
+    # should be nan_color
+    assert_array_equal(colors[0], np.array([0, 0, 0, 0])) 
+    # still a cmap color
+    assert_array_equal(colors[-1], np.array([94, 201,  97, 255])) 
+
+def test_rgba_nan_hex():
+    colors = get_color_array(world['nanarea'], cmap="viridis",nan_color=[0,0,0,0], as_hex=True)
+    assert_array_equal(colors[0], np.array(['#000000'])) 
+    assert_array_equal(colors[-1], np.array(['#5ec961'])) 

mapclassify/util.py

@@ -59,6 +59,7 @@ def get_color_array(
     v = pd.Series(values, dtype=object)
     legit_indices = v[~v.isna()].index.values
     legit_vals = v.dropna().values
+    bogus_indices = v[v.isna()].index.values  # stash these for use later
     # transform (non-NaN) values into class bins
     bins = _classify(legit_vals, scheme=scheme, **kwargs).yb
 
@@ -69,13 +70,16 @@ def get_color_array(
     # generate RBGA array and convert to series
     rgbas = colormaps[cmap](normalized_vals, bytes=True, alpha=alpha)
     colors = pd.Series(list(rgbas), index=legit_indices).apply(np.array)
+    nan_colors = pd.Series(
+        [nan_color for i in range(len(bogus_indices))], index=bogus_indices
+    ).apply(lambda x: np.array(x).astype(np.uint8))
 
-    # put colors in their correct places and fill empty with designated color
+    # put colors in their correct places and fill empty with specified color
     v.update(colors)
-    v = v.fillna(f"{nan_color}").apply(np.array)
+    v.update(nan_colors)
 
     # convert to hexcolors if preferred
     if as_hex:
-        colors = colors.apply(lambda x: to_hex(x / 255.0))
+        colors = v.apply(lambda x: to_hex(x / 255.0))
         return colors.values
     return np.stack(v.values)