masster 0.5.4__py3-none-any.whl → 0.5.6__py3-none-any.whl

masster/sample/adducts.py

@@ -473,7 +473,7 @@ def find_adducts(self, **kwargs):
     self.logger.debug(f"Min probability threshold: {min_probability}")
 
     # Generate comprehensive adduct specifications using the Sample method
-    adducts_df = _get_adducts(self, 
+    adducts_df = self._get_adducts( 
         adducts_list=adducts_list,
         charge_min=charge_min,
         charge_max=charge_max,

masster/sample/helpers.py

@@ -358,7 +358,7 @@ def get_eic(self, mz, mz_tol=None):
         return None
 
 
-def select(
+def features_select(
     self,
     mz=None,
     rt=None,
@@ -372,6 +372,7 @@ def select(
     height_scaled=None,
     prominence=None,
     height=None,
+    uids=None,
 ):
     """
     Select features based on specified criteria and return the filtered DataFrame.
@@ -389,6 +390,7 @@ def select(
         height_scaled: scaled height filter (tuple for range, single value for minimum)
         prominence: prominence filter (tuple for range, single value for minimum)
         height: height filter (tuple for range, single value for minimum)
+        uids: feature UID filter (list of feature UIDs, polars/pandas DataFrame with feature_uid/feature_id column, or None for all features)
 
     Returns:
         polars.DataFrame: Filtered features DataFrame
@@ -398,6 +400,20 @@ def select(
         # self.logger.info("No features found. R")
         return
     feats = self.features_df.clone()
+    
+    # Filter by feature UIDs if provided
+    if uids is not None:
+        feature_uids_to_keep = self._get_feature_uids(features=uids, verbose=True)
+        if not feature_uids_to_keep:
+            self.logger.warning("No valid feature UIDs provided.")
+            return feats.limit(0)  # Return empty DataFrame with same structure
+        
+        feats_len_before_filter = len(feats)
+        feats = feats.filter(pl.col("feature_uid").is_in(feature_uids_to_keep))
+        self.logger.debug(
+            f"Selected features by UIDs. Features removed: {feats_len_before_filter - len(feats)}",
+        )
+    
     if coherence is not None:
         has_coherence = "chrom_coherence" in self.features_df.columns
         if not has_coherence:

masster/sample/plot.py

@@ -57,12 +57,78 @@ from holoviews import dim
 from holoviews.plotting.util import process_cmap
 from matplotlib.colors import rgb2hex
 
+# Import cmap for colormap handling
+try:
+    from cmap import Colormap
+except ImportError:
+    Colormap = None
+
 # Parameters removed - using hardcoded defaults
 
 
 hv.extension("bokeh")
 
 
+def _process_cmap(cmap, fallback="viridis", logger=None):
+    """
+    Process colormap using the cmap package, similar to study's implementation.
+    
+    Parameters:
+        cmap: Colormap specification (string name, cmap.Colormap object, or None)
+        fallback: Fallback colormap name if cmap processing fails
+        logger: Logger for warnings (optional)
+        
+    Returns:
+        list: List of hex color strings for the colormap
+    """
+    # Handle None case
+    if cmap is None:
+        cmap = "viridis"
+    elif cmap == "grey":
+        cmap = "Greys256"
+    elif cmap == "iridescent":
+        cmap = "iridescent_r"
+    
+    # If cmap package is not available, fall back to process_cmap
+    if Colormap is None:
+        if logger:
+            logger.warning("cmap package not available, using holoviews process_cmap")
+        return process_cmap(cmap, provider="bokeh")
+    
+    try:
+        # Handle colormap using cmap.Colormap
+        if isinstance(cmap, str):
+            colormap = Colormap(cmap)
+            # Generate 256 colors and convert to hex
+            import matplotlib.colors as mcolors
+            colors = colormap(np.linspace(0, 1, 256))
+            palette = [mcolors.rgb2hex(color) for color in colors]
+        else:
+            colormap = cmap
+            # Try to use to_bokeh() method first
+            try:
+                palette = colormap.to_bokeh()
+                # Ensure we got a color palette, not another mapper
+                if not isinstance(palette, (list, tuple)):
+                    # Fall back to generating colors manually
+                    import matplotlib.colors as mcolors
+                    colors = colormap(np.linspace(0, 1, 256))
+                    palette = [mcolors.rgb2hex(color) for color in colors]
+            except AttributeError:
+                # Fall back to generating colors manually
+                import matplotlib.colors as mcolors
+                colors = colormap(np.linspace(0, 1, 256))
+                palette = [mcolors.rgb2hex(color) for color in colors]
+        
+        return palette
+        
+    except (AttributeError, ValueError, TypeError) as e:
+        # Fallback to process_cmap if cmap interpretation fails
+        if logger:
+            logger.warning(f"Could not interpret colormap '{cmap}': {e}, falling back to {fallback}")
+        return process_cmap(fallback, provider="bokeh")
+
+
 def _is_notebook_environment():
     """
     Detect if code is running in a notebook environment (Jupyter, JupyterLab, or Marimo).
@@ -462,10 +528,8 @@ def plot_2d(
         self.logger.error("No MS1 data available.")
         return
 
-    if cmap is None:
-        cmap = "iridescent_r"
-    elif cmap == "grey":
-        cmap = "Greys256"
+    # Process colormap using the cmap package
+    cmap_palette = _process_cmap(cmap, fallback="iridescent_r", logger=self.logger)
 
     # get columns rt, mz, inty from self.ms1_df, It's polars DataFrame
     spectradf = self.ms1_df.select(["rt", "mz", "inty"])
@@ -585,7 +649,7 @@ def plot_2d(
         dynamic=dyn,  # alpha=10,                min_alpha=0,
     ).opts(
         active_tools=["box_zoom"],
-        cmap=process_cmap(cmap, provider="bokeh"),  # blues
+        cmap=cmap_palette,
         tools=["hover"],
         hooks=[new_bounds_hook],
         width=width,
@@ -1003,10 +1067,8 @@ def plot_2d_oracle(
         print("Please load a file first.")
         return
 
-    if cmap is None or cmap == "grey":
-        cmap = "Greys256"
-    elif cmap == "iridescent":
-        cmap = "iridescent_r"
+    # Process colormap using the cmap package
+    cmap_palette = _process_cmap(cmap, fallback="Greys256", logger=self.logger)
 
     # get columns rt, mz, inty from self.ms1_df, It's polars DataFrame
     spectradf = self.ms1_df.to_pandas()
@@ -1057,7 +1119,7 @@ def plot_2d_oracle(
         dynamic=dyn,  # alpha=10,                min_alpha=0,
     ).opts(
         active_tools=["box_zoom"],
-        cmap=process_cmap(cmap, provider="bokeh"),  # blues
+        cmap=cmap_palette,
         tools=["hover"],
         hooks=[new_bounds_hook],
         width=1000,
@@ -1183,13 +1245,45 @@ def plot_2d_oracle(
 
     if cvalues is not None:
         num_colors = len(cvalues)
-        cmap = "rainbow"
-        cmap_provider = "colorcet"
-        cm = process_cmap(cmap, ncolors=num_colors, provider=cmap_provider)
-        colors = [
-            rgb2hex(cm[int(i * (len(cm) - 1) / (num_colors - 1))]) if num_colors > 1 else rgb2hex(cm[0])
-            for i in range(num_colors)
-        ]
+        
+        # Use cmap package for categorical colormap
+        try:
+            if Colormap is not None:
+                # Use rainbow colormap for categorical data
+                colormap = Colormap("rainbow")
+                colors = []
+                for i in range(num_colors):
+                    # Generate evenly spaced colors across the colormap
+                    t = i / (num_colors - 1) if num_colors > 1 else 0.5
+                    color = colormap(t)
+                    # Convert to hex
+                    import matplotlib.colors as mcolors
+                    # Convert color to hex - handle different color formats
+                    if hasattr(color, '__len__') and len(color) >= 3:
+                        # It's an array-like color (RGB or RGBA)
+                        colors.append(mcolors.rgb2hex(color[:3]))
+                    else:
+                        # It's a single value, convert to RGB
+                        colors.append(mcolors.rgb2hex([color, color, color]))
+            else:
+                # Fallback to original method
+                cmap = "rainbow"
+                cmap_provider = "colorcet"
+                cm = process_cmap(cmap, ncolors=num_colors, provider=cmap_provider)
+                colors = [
+                    rgb2hex(cm[int(i * (len(cm) - 1) / (num_colors - 1))]) if num_colors > 1 else rgb2hex(cm[0])
+                    for i in range(num_colors)
+                ]
+        except Exception:
+            # Final fallback to original method
+            cmap = "rainbow"
+            cmap_provider = "colorcet" 
+            cm = process_cmap(cmap, ncolors=num_colors, provider=cmap_provider)
+            colors = [
+                rgb2hex(cm[int(i * (len(cm) - 1) / (num_colors - 1))]) if num_colors > 1 else rgb2hex(cm[0])
+                for i in range(num_colors)
+            ]
+        
         # assign color to each row based on id_class. If id_class is null, assign 'black'
         feats["color"] = "black"
 
@@ -1577,10 +1671,8 @@ def plot_ms2_cycle(
         print("Cycle number not found in scans_df.")
         return
 
-    if cmap is None:
-        cmap = "iridescent_r"
-    elif cmap == "grey":
-        cmap = "Greys256"
+    # Process colormap using the cmap package
+    cmap_palette = _process_cmap(cmap, fallback="iridescent_r", logger=self.logger)
 
     # find all scans in cycle
     scans = self.scans_df.filter(pl.col("cycle") == cycle)
@@ -1645,7 +1737,7 @@ def plot_ms2_cycle(
         dynamic=raster_dynamic,  # alpha=10,                min_alpha=0,
     ).opts(
         active_tools=["box_zoom"],
-        cmap=process_cmap(cmap, provider="bokeh"),  # blues
+        cmap=cmap_palette,
         tools=["hover"],
         hooks=[new_bounds_hook],
         width=1000,

masster/sample/sample.py

@@ -57,10 +57,11 @@ from masster.sample.helpers import _get_scan_uids
 from masster.sample.helpers import _get_feature_uids
 # from masster.sample.helpers import _features_sync - made internal only
 from masster.sample.adducts import find_adducts
+from masster.sample.adducts import _get_adducts
 # Removed _get_adducts - only used in study modules
 from masster.sample.helpers import features_delete
 from masster.sample.helpers import features_filter
-from masster.sample.helpers import select
+from masster.sample.helpers import features_select
 from masster.sample.helpers import select_closest_scan
 from masster.sample.helpers import get_dda_stats
 from masster.sample.helpers import get_feature
@@ -110,6 +111,7 @@ from masster.sample.save import export_chrom
 from masster.sample.save import export_dda_stats
 from masster.sample.save import export_features
 from masster.sample.save import export_mgf
+from masster.sample.save import export_xlsx
 from masster.sample.save import save
 # Removed internal-only import: _save_featureXML
 
@@ -139,7 +141,6 @@ class Sample:
 
                      Core initialization parameters:
                      - file (str, optional): The file path or file object to be loaded
-                     - ondisk (bool): Whether to keep data on disk or load into memory. Default is False
                      - label (str, optional): An optional label to identify the file or dataset
                      - log_level (str): The logging level to be set for the logger. Defaults to 'INFO'
                      - log_label (str, optional): Optional label for the logger
@@ -221,11 +222,12 @@ class Sample:
     save = save
     find_features = find_features
     find_adducts = find_adducts
+    _get_adducts= _get_adducts
     find_iso = find_iso
     find_ms2 = find_ms2
     get_spectrum = get_spectrum
     filter = features_filter
-    select = select
+    select = features_select
     features_filter = filter  # New function that keeps only specified features
     filter_features = filter
     features_select = select
@@ -238,6 +240,7 @@ class Sample:
     get_parameters_property = get_parameters_property
     set_parameters_property = set_parameters_property
     export_features = export_features
+    export_xlsx = export_xlsx
     export_mgf = export_mgf
     export_chrom = export_chrom
     export_dda_stats = export_dda_stats

masster/sample/save.py

@@ -105,7 +105,7 @@ def save(self, filename=None):
     self._save_sample5(filename=filename)
     self.file_path = filename
 
-
+'''
 def _save_featureXML(self, filename="features.featureXML"):
     if self._oms_features_map is None:
         self.logger.warning("No features found.")
@@ -114,7 +114,7 @@ def _save_featureXML(self, filename="features.featureXML"):
     fh.store(filename, self._oms_features_map)
     self.logger.debug(f"Features Map saved to {filename}")
 
-
+'''
 def export_features(self, filename="features.csv"):
     """
     Export the features DataFrame to a CSV or Excel file.
@@ -411,6 +411,11 @@ def export_mgf(
             rt_str = f"{rt:.2f}"
             mz_str = f"{mz:.4f}"
 
+            # Initialize charge for this feature
+            charge = preferred_charge
+            if row["charge"] is not None and row["charge"] != 0:
+                charge = row["charge"]
+
             # Skip features without MS2 data (unless include_all_ms1 is True, but we already handled MS1 above)
             if row["ms2_scans"] is None:
                 skip = skip + 1
@@ -822,6 +827,60 @@ def export_dda_stats(self, filename="stats.csv"):
     self.logger.info(f"DDA statistics exported to {filename}")
 
 
+def export_xlsx(self, filename="features.xlsx"):
+    """
+    Export the features DataFrame to an Excel file.
+
+    This method exports the features DataFrame (features_df) to an Excel (.xlsx) file.
+    Columns with data types 'List' or 'Object' are excluded from the export to ensure
+    compatibility with Excel format. A boolean column 'has_ms2' is added to indicate
+    whether MS2 data is available for each feature.
+
+    Parameters:
+        filename (str): The path to the output Excel file. Must end with '.xlsx' or '.xls'.
+                        Defaults to 'features.xlsx'.
+
+    Raises:
+        ValueError: If filename doesn't end with '.xlsx' or '.xls'
+
+    Side Effects:
+        Writes the exported data to the specified Excel file and logs the export operation.
+    """
+    if self.features_df is None:
+        self.logger.warning("No features found. Cannot export to Excel.")
+        return
+
+    # Validate filename extension
+    if not filename.lower().endswith(('.xlsx', '.xls')):
+        raise ValueError("Filename must end with '.xlsx' or '.xls' for Excel export")
+    
+    filename = os.path.abspath(filename)
+    
+    # Clone the DataFrame to avoid modifying the original
+    clean_df = self.features_df.clone()
+    
+    # Add a column has_ms2=True if column ms2_scans is not None
+    if "ms2_scans" in clean_df.columns:
+        clean_df = clean_df.with_columns(
+            (pl.col("ms2_scans").is_not_null()).alias("has_ms2")
+        )
+    
+    # Filter out columns with List or Object data types that can't be exported to Excel
+    exportable_columns = [
+        col for col in clean_df.columns
+        if clean_df[col].dtype not in (pl.List, pl.Object)
+    ]
+    
+    clean_df = clean_df.select(exportable_columns)
+    
+    # Convert to pandas and export to Excel
+    pandas_df = clean_df.to_pandas()
+    pandas_df.to_excel(filename, index=False)
+    
+    self.logger.info(f"Features exported to {filename} (Excel format)")
+    self.logger.debug(f"Exported {len(clean_df)} features with {len(exportable_columns)} columns")
+
+
 def export_chrom(self, filename="chrom.csv"):
     # saves self.chrom_df to a csv file. Remove the scan_uid and chrom columns if the file already exists
     if self.chrom_df is None:

masster/spectrum.py

@@ -827,8 +827,11 @@ class Spectrum:
                 inty = inty[idx]
             p.line(mz, inty, line_color="black", legend_label=label)
         else:
-            data = self.to_dict()
-            data = {key: val for key, val in data.items() if isinstance(val, np.ndarray) and val.size == mz.size}
+            # Build data dictionary from spectrum attributes (numpy arrays)
+            data = {}
+            for key, val in self.__dict__.items():
+                if isinstance(val, np.ndarray) and val.size == mz.size:
+                    data[key] = val
             if ylog:
                 data["zeros"] = np.ones_like(mz)
             else:

masster/study/helpers.py

@@ -500,7 +500,7 @@ def align_reset(self):
 # TODO I don't get this param
 def get_consensus(self, quant="chrom_area"):
     if self.consensus_df is None:
-        self.logger.error("No consensus map found.")
+        self.logger.error("No consensus found.")
         return None
 
     # Convert Polars DataFrame to pandas for this operation since the result is used for export
@@ -613,7 +613,7 @@ def get_gaps_matrix(self, uids=None, samples=None):
     import polars as pl
     
     if self.consensus_df is None or self.consensus_df.is_empty():
-        self.logger.error("No consensus map found.")
+        self.logger.error("No consensus found.")
         return None
     
     if self.consensus_mapping_df is None or self.consensus_mapping_df.is_empty():

masster/study/plot.py

@@ -564,6 +564,10 @@ def plot_consensus_2d(
     Parameters:
         filename (str, optional): Path to save the plot
         colorby (str): Column name to use for color mapping (default: "number_samples")
+                      Automatically detects if column contains categorical (string) or 
+                      numeric data and applies appropriate color mapping:
+                      - Categorical: Uses factor_cmap with distinct colors and legend
+                      - Numeric: Uses LinearColorMapper with continuous colorbar
         sizeby (str): Column name to use for size mapping (default: "inty_mean")
         markersize (int): Base marker size (default: 6)
         scaling (str): Controls whether points scale with zoom. Options:
@@ -645,12 +649,13 @@ def plot_consensus_2d(
     from bokeh.models import HoverTool
     from bokeh.models import LinearColorMapper
     from bokeh.io.export import export_png
+    from bokeh.transform import factor_cmap
 
     try:
         from bokeh.models import ColorBar  # type: ignore[attr-defined]
     except ImportError:
         from bokeh.models.annotations import ColorBar
-    from bokeh.palettes import viridis
+    from bokeh.palettes import viridis, Category20
 
     # Import cmap for colormap handling
     from cmap import Colormap
@@ -695,61 +700,144 @@ def plot_consensus_2d(
         self.logger.warning(f"Could not interpret colormap '{cmap}': {e}, falling back to viridis")
         palette = viridis(256)
 
-    color_mapper = LinearColorMapper(
-        palette=palette,
-        low=data[colorby].min(),
-        high=data[colorby].max(),
+    # Check if colorby column contains categorical data (string/object)
+    colorby_values = data[colorby].to_list()
+    is_categorical = (
+        data_pd[colorby].dtype in ["object", "string", "category"] or
+        isinstance(colorby_values[0], str) if colorby_values else False
     )
+    
+    if is_categorical:
+        # Handle categorical coloring
+        # Use natural order of unique values - don't sort to preserve correct legend mapping
+        # Sorting would break the correspondence between legend labels and point colors
+        unique_values = [v for v in data_pd[colorby].unique() if v is not None]
+        
+        if len(unique_values) <= 20:
+            palette = Category20[min(20, max(3, len(unique_values)))]
+        else:
+            # For many categories, use a subset of the viridis palette
+            palette = viridis(min(256, len(unique_values)))
+        
+        color_mapper = factor_cmap(colorby, palette, unique_values)
+    else:
+        # Handle numeric coloring with LinearColorMapper
+        color_mapper = LinearColorMapper(
+            palette=palette,
+            low=data[colorby].min(),
+            high=data[colorby].max(),
+        )
     # scatter plot rt vs mz
     p = bp.figure(
         width=width,
         height=height,
-        title="Consensus map",
+        title=f"Consensus features, colored by {colorby}",
     )
-    p.xaxis.axis_label = "Retention Time (min)"
-    p.yaxis.axis_label = "m/z"
+    p.xaxis.axis_label = "RT [s]"
+    p.yaxis.axis_label = "m/z [Th]"
     scatter_renderer: Any = None
-    if scaling.lower() in ["dyn", "dynamic"]:
-        # Calculate appropriate radius for dynamic scaling based on data range
-        rt_range = data["rt"].max() - data["rt"].min()
-        mz_range = data["mz"].max() - data["mz"].min()
-        # Use a fraction of the smaller dimension for radius, similar to sample plotting
-        dynamic_radius = min(rt_range, mz_range) * 0.0005 * markersize
+    if is_categorical:
+        # For categorical data, create separate renderers for each category
+        # This enables proper legend interactivity where each category can be toggled independently
+        unique_values = [v for v in data_pd[colorby].unique() if v is not None]
+        
+        if len(unique_values) <= 20:
+            palette = Category20[min(20, max(3, len(unique_values)))]
+        else:
+            palette = viridis(min(256, len(unique_values)))
+        
+        # Create a separate renderer for each category
+        for i, category in enumerate(unique_values):
+            # Filter data for this category
+            category_data = data.filter(pl.col(colorby) == category)
+            category_data_pd = category_data.to_pandas()
+            category_source = bp.ColumnDataSource(category_data_pd)
+            
+            color = palette[i % len(palette)]
+            
+            if scaling.lower() in ["dyn", "dynamic"]:
+                # Calculate appropriate radius for dynamic scaling
+                rt_range = data["rt"].max() - data["rt"].min()
+                mz_range = data["mz"].max() - data["mz"].min()
+                dynamic_radius = min(rt_range, mz_range) * 0.0005 * markersize
+                
+                renderer = p.circle(
+                    x="rt",
+                    y="mz", 
+                    radius=dynamic_radius,
+                    fill_color=color,
+                    line_color=None,
+                    alpha=alpha,
+                    source=category_source,
+                    legend_label=str(category),
+                )
+            else:
+                renderer = p.scatter(
+                    x="rt",
+                    y="mz",
+                    size="markersize", 
+                    fill_color=color,
+                    line_color=None,
+                    alpha=alpha,
+                    source=category_source,
+                    legend_label=str(category),
+                )
+        
+        # No single scatter_renderer for categorical data
+        scatter_renderer = None
         
-        scatter_renderer = p.circle(
-            x="rt",
-            y="mz",
-            radius=dynamic_radius,
-            fill_color={"field": colorby, "transform": color_mapper},
-            line_color=None,
-            alpha=alpha,
-            source=source,
-        )
     else:
-        scatter_renderer = p.scatter(
-            x="rt",
-            y="mz",
-            size="markersize",
-            fill_color={"field": colorby, "transform": color_mapper},
-            line_color=None,
-            alpha=alpha,
-            source=source,
-        )
+        # Handle numeric coloring - single renderer with color mapping
+        if scaling.lower() in ["dyn", "dynamic"]:
+            # Calculate appropriate radius for dynamic scaling
+            rt_range = data["rt"].max() - data["rt"].min()
+            mz_range = data["mz"].max() - data["mz"].min()
+            dynamic_radius = min(rt_range, mz_range) * 0.0005 * markersize
+            
+            scatter_renderer = p.circle(
+                x="rt",
+                y="mz",
+                radius=dynamic_radius,
+                fill_color={"field": colorby, "transform": color_mapper},
+                line_color=None,
+                alpha=alpha,
+                source=source,
+            )
+        else:
+            scatter_renderer = p.scatter(
+                x="rt",
+                y="mz",
+                size="markersize",
+                fill_color={"field": colorby, "transform": color_mapper},
+                line_color=None,
+                alpha=alpha,
+                source=source,
+            )
     # add hover tool
-    # Start with base tooltips
+    # Start with base tooltips - rt and mz moved to top, removed consensus_id and iso_mean
     tooltips = [
+        ("rt", "@rt"),
+        ("mz", "@mz"),
         ("consensus_uid", "@consensus_uid"),
-        ("consensus_id", "@consensus_id"),
         ("number_samples", "@number_samples"),
         ("number_ms2", "@number_ms2"),
-        ("rt", "@rt"),
-        ("mz", "@mz"),
         ("inty_mean", "@inty_mean"),
-        ("iso_mean", "@iso_mean"),
         ("coherence_mean", "@chrom_coherence_mean"),
         ("prominence_scaled_mean", "@chrom_prominence_scaled_mean"),
     ]
     
+    # Add adduct_top if it exists in data
+    if "adduct_top" in data.columns:
+        tooltips.append(("adduct_top", "@adduct_top"))
+    
+    # Add id_top_name if it exists in data
+    if "id_top_name" in data.columns:
+        tooltips.append(("id_top_name", "@id_top_name"))
+    
+    # Add id_top_adduct if it exists in data
+    if "id_top_adduct" in data.columns:
+        tooltips.append(("id_top_adduct", "@id_top_adduct"))
+    
     # Add id_top_* columns if they exist and have non-null values
     id_top_columns = ["id_top_name", "id_top_class", "id_top_adduct", "id_top_score"]
     for col in id_top_columns:
@@ -764,19 +852,28 @@ def plot_consensus_2d(
     
     hover = HoverTool(
         tooltips=tooltips,
-        renderers=[scatter_renderer],
     )
+    # For categorical data, hover will work on all renderers automatically
+    # For numeric data, specify the single renderer
+    if not is_categorical and scatter_renderer:
+        hover.renderers = [scatter_renderer]
+    
     p.add_tools(hover)
 
-    # add colorbar
-    color_bar = ColorBar(
-        color_mapper=color_mapper,
-        label_standoff=12,
-        location=(0, 0),
-        title=colorby,
-        ticker=BasicTicker(desired_num_ticks=8),
-    )
-    p.add_layout(color_bar, "right")
+    # add colorbar only for numeric data (LinearColorMapper)
+    if not is_categorical:
+        color_bar = ColorBar(
+            color_mapper=color_mapper,
+            label_standoff=12,
+            location=(0, 0),
+            title=colorby,
+            ticker=BasicTicker(desired_num_ticks=8),
+        )
+        p.add_layout(color_bar, "right")
+    else:
+        # For categorical data, configure the legend that was automatically created
+        p.legend.location = "top_right"
+        p.legend.click_policy = "hide"
 
     if filename is not None:
         # Convert relative paths to absolute paths using study folder as base