Fixed trace colors

Author: lucasimi

app/streamlit_app.py

@@ -697,14 +697,20 @@ def mapper_figure_section(df_X, df_y, mapper_plot):
         scaleratio=1,
     )
 
+    _enhance_fig(mapper_plot, mapper_fig, df_X, df_y, agg)
+
+    return mapper_fig
+
+
+def _enhance_fig(mapper_plot, mapper_fig, df_X, df_y, agg):
     menu_cmap = {}
-    if dim == 2:
+    if mapper_plot.dim == 2:
         menu_cmap = _set_cmap_buttons_2d(mapper_fig)
         menu_feature = _set_feature_select_button_2d(
             mapper_plot, mapper_fig, df_X, df_y, agg
         )
 
-    elif dim == 3:
+    elif mapper_plot.dim == 3:
         menu_cmap = _set_cmap_buttons_3d(mapper_fig)
         menu_feature = _set_feature_select_button_3d(
             mapper_plot, mapper_fig, df_X, df_y, agg
@@ -718,8 +724,6 @@ def mapper_figure_section(df_X, df_y, mapper_plot):
         uirevision="constant",
     )
 
-    return mapper_fig
-
 
 def _set_cmap_buttons_2d(fig):
     return dict(
@@ -900,7 +904,7 @@ def _set_feature_select_button_3d(mapper_plot, mapper_fig, df_X, df_y, agg):
                             )
                         ],
                     },
-                    [1],  # Trace indices
+                    [0, 1],
                 ],
             )
             for col_feat in col_feats

tests/test_unit_cover.py

@@ -3,11 +3,7 @@
 import numpy as np
 
 from tdamapper.core import TrivialCover
-from tdamapper.cover import (
-    BallCover,
-    CubicalCover,
-    KNNCover,
-)
+from tdamapper.cover import BallCover, CubicalCover, KNNCover
 
 
 def dataset(dim=1, num=10000):