handle nans in colormap (#273)

ecoscope/analysis/classifier.py

@@ -105,11 +105,22 @@ def apply_color_map(dataframe, input_column_name, cmap, output_column_name=None)
         cmap = mpl.colormaps[cmap]
         cmap = cmap.resampled(dataframe[input_column_name].nunique())
 
-        cmap_colors = cmap(range(dataframe[input_column_name].nunique()))
+        if pd.api.types.is_numeric_dtype(dataframe[input_column_name].dtype):
+            cmap_colors = []
+            val_min = dataframe[input_column_name].min()
+            val_max = dataframe[input_column_name].max()
+            for val in dataframe[input_column_name].unique():
+                cmap_colors.append(cmap((val - val_min) / (val_max - val_min)))
+        else:
+            cmap_colors = cmap(range(len(dataframe[input_column_name].unique())))
+
+        color_list = []
+        for color in cmap_colors:
+            color_list.append(tuple([round(val * 255) for val in color]))
 
         # convert to hex first to put values in range(0,255), then to an RGBA tuple
         cmap = pd.Series(
-            [hex_to_rgba(mpl.colors.to_hex(color)) for color in cmap_colors],
+            color_list,
             index=dataframe[input_column_name].unique(),
         )
 

ecoscope/analysis/feature_density.py

@@ -1,3 +1,6 @@
+import numpy as np
+
+
 def calculate_feature_density(selection, grid, geometry_type="point"):
     def clip_density(cell):
         if geometry_type == "point":
@@ -12,5 +15,5 @@ def clip_density(cell):
             raise ValueError("Unsupported geometry type")
 
     grid["density"] = grid.geometry.apply(clip_density)
-    # grid["density"] = grid["density"].replace(0, np.nan)  # Set 0's to nan so they don't draw on map
+    grid["density"] = grid["density"].replace(0, np.nan)  # Set 0's to nan so they don't draw on map
     return grid

tests/test_classifier.py

@@ -1,5 +1,6 @@
 import pytest
 import pandas as pd
+import numpy as np
 from ecoscope.base import Trajectory
 from ecoscope.analysis.classifier import apply_classification, apply_color_map
 
@@ -52,15 +53,26 @@ def test_apply_colormap(sample_df, cmap):
     apply_classification(sample_df, input_column_name="value", scheme="equal_interval")
     apply_color_map(sample_df, "value_classified", cmap, output_column_name="colormap")
 
-    assert len(set(sample_df["colormap"].unique())) == len(sample_df["value_classified"].unique())
+    assert len(sample_df["colormap"].unique()) == len(sample_df["value_classified"].unique())
+
+
+def test_apply_colormap_with_nan():
+    df = pd.DataFrame(
+        data={"value": [1, 2, 3, 4, np.nan]},
+        index=["A", "B", "C", "D", "E"],
+    )
+    apply_color_map(df, "value", "viridis", output_column_name="colormap")
+
+    assert len(df["colormap"].unique()) == len(df["value"].unique())
+    assert df.loc["E"]["colormap"] == (0, 0, 0, 0)
 
 
 def test_apply_colormap_k2(sample_df):
     apply_classification(sample_df, input_column_name="value", scheme="equal_interval", k=2)
     cmap = "viridis"
     apply_color_map(sample_df, "value_classified", cmap, output_column_name="colormap")
 
-    assert len(set(sample_df["colormap"].unique())) == len(sample_df["value_classified"].unique())
+    assert len(sample_df["colormap"].unique()) == len(sample_df["value_classified"].unique())
 
 
 def test_apply_colormap_user_defined(movebank_relocations):
@@ -79,7 +91,7 @@ def test_apply_colormap_user_defined(movebank_relocations):
     ]
 
     apply_color_map(trajectory, "speed_bins", cmap)
-    assert len(set(trajectory["speed_bins_colormap"].unique())) == len(trajectory["speed_bins"].unique())
+    assert len(trajectory["speed_bins_colormap"].unique()) == len(trajectory["speed_bins"].unique())
 
 
 def test_apply_colormap_cmap_user_defined_bad(movebank_relocations):