diff --git a/mapclassify/tests/test_rgba.py b/mapclassify/tests/test_rgba.py index 02f4bee..c2e6f13 100644 --- a/mapclassify/tests/test_rgba.py +++ b/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'])) diff --git a/mapclassify/util.py b/mapclassify/util.py index 21879ae..8b29d6f 100644 --- a/mapclassify/util.py +++ b/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)