masster 0.5.9__py3-none-any.whl → 0.5.10__py3-none-any.whl

masster/sample/plot.py

@@ -43,6 +43,7 @@ See Also:
 """
 
 import os
+import warnings
 
 import datashader as ds
 import holoviews as hv
@@ -55,18 +56,11 @@ import polars as pl
 from bokeh.models import HoverTool
 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
+from cmap import Colormap
 
 # Parameters removed - using hardcoded defaults
-
-
-hv.extension("bokeh")
+# hv.extension("bokeh")
 
 
 def _process_cmap(cmap, fallback="viridis", logger=None):
@@ -85,8 +79,8 @@ def _process_cmap(cmap, fallback="viridis", logger=None):
     if cmap is None:
         cmap = "viridis"
     elif cmap == "grey":
-        cmap = "Greys256"
-    
+        cmap = "greys"
+        
     # If cmap package is not available, fall back to process_cmap
     if Colormap is None:
         if logger:
@@ -205,6 +199,108 @@ def _display_plot(plot_object, layout=None):
         return None
 
 
+def _export_with_webdriver_manager(plot_obj, filename, format_type, logger=None):
+    """
+    Export plot to PNG or SVG using webdriver-manager for automatic driver management.
+    
+    Parameters:
+        plot_obj: Bokeh plot object or holoviews object to export
+        filename: Output filename 
+        format_type: Either "png" or "svg"
+        logger: Logger for error reporting (optional)
+        
+    Returns:
+        bool: True if export successful, False otherwise
+    """
+    try:
+        # Convert holoviews to bokeh if needed
+        if hasattr(plot_obj, 'opts'):  # Likely a holoviews object
+            import holoviews as hv
+            bokeh_plot = hv.render(plot_obj)
+        else:
+            bokeh_plot = plot_obj
+            
+        # Try webdriver-manager export first
+        try:
+            from webdriver_manager.chrome import ChromeDriverManager
+            from selenium import webdriver
+            from selenium.webdriver.chrome.service import Service
+            from selenium.webdriver.chrome.options import Options
+            
+            # Set up Chrome options for headless operation
+            chrome_options = Options()
+            chrome_options.add_argument("--headless")
+            chrome_options.add_argument("--no-sandbox")
+            chrome_options.add_argument("--disable-dev-shm-usage")
+            chrome_options.add_argument("--disable-gpu")
+            
+            # Use webdriver-manager to automatically get the correct ChromeDriver
+            service = Service(ChromeDriverManager().install())
+            driver = webdriver.Chrome(service=service, options=chrome_options)
+            
+            # Export with managed webdriver
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", category=UserWarning)
+                # Filter out bokeh.io.export warnings specifically
+                warnings.filterwarnings("ignore", module="bokeh.io.export")
+                
+                if format_type == "png":
+                    from bokeh.io import export_png
+                    export_png(bokeh_plot, filename=filename, webdriver=driver)
+                elif format_type == "svg":
+                    from bokeh.io import export_svg
+                    export_svg(bokeh_plot, filename=filename, webdriver=driver)
+                else:
+                    raise ValueError(f"Unsupported format: {format_type}")
+                
+            driver.quit()
+            return True
+            
+        except ImportError:
+            if logger:
+                logger.debug(f"webdriver-manager not available, using default {format_type.upper()} export")
+            # Fall back to default export
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", category=UserWarning)
+                # Filter out bokeh.io.export warnings specifically
+                warnings.filterwarnings("ignore", module="bokeh.io.export")
+                
+                if format_type == "png":
+                    from bokeh.io import export_png
+                    export_png(bokeh_plot, filename=filename)
+                elif format_type == "svg":
+                    from bokeh.io import export_svg
+                    export_svg(bokeh_plot, filename=filename)
+            return True
+            
+        except Exception as e:
+            if logger:
+                logger.debug(f"{format_type.upper()} export with webdriver-manager failed: {e}, using default {format_type.upper()} export")
+            try:
+                # Final fallback to default export
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", category=UserWarning)
+                    # Filter out bokeh.io.export warnings specifically
+                    warnings.filterwarnings("ignore", module="bokeh.io.export")
+                    
+                    if format_type == "png":
+                        from bokeh.io import export_png
+                        export_png(bokeh_plot, filename=filename)
+                    elif format_type == "svg":
+                        from bokeh.io import export_svg
+                        export_svg(bokeh_plot, filename=filename)
+                return True
+            except Exception as e2:
+                if logger:
+                    logger.error(f"{format_type.upper()} export failed: {e2}")
+                return False
+                
+    except Exception as e:
+        if logger:
+            logger.error(f"Export preparation failed: {e}")
+        return False
+
+
 def _handle_sample_plot_output(self, plot_obj, filename=None, plot_type="bokeh"):
     """
     Helper function to handle consistent save/display behavior for sample plots.
@@ -236,16 +332,11 @@ def _handle_sample_plot_output(self, plot_obj, filename=None, plot_type="bokeh")
                 save(plot_obj)
             self.logger.success(f"Plot saved to: {abs_filename}")
         elif filename.endswith(".png"):
-            try:
-                if plot_type == "bokeh":
-                    from bokeh.io.export import export_png
-                    export_png(plot_obj, filename=filename)
-                elif plot_type in ["panel", "holoviews"]:
-                    import holoviews as hv
-                    hv.save(plot_obj, filename, fmt="png")
+            success = _export_with_webdriver_manager(plot_obj, filename, "png", self.logger)
+            if success:
                 self.logger.success(f"Plot saved to: {abs_filename}")
-            except Exception:
-                # Fall back to HTML if PNG export not available
+            else:
+                # Fall back to HTML if PNG export fails completely
                 html_filename = filename.replace('.png', '.html')
                 abs_html_filename = os.path.abspath(html_filename)
                 if plot_type == "panel":
@@ -259,6 +350,25 @@ def _handle_sample_plot_output(self, plot_obj, filename=None, plot_type="bokeh")
                     output_file(html_filename)
                     save(plot_obj)
                 self.logger.warning(f"PNG export not available, saved as HTML instead: {abs_html_filename}")
+        elif filename.endswith(".svg"):
+            success = _export_with_webdriver_manager(plot_obj, filename, "svg", self.logger)
+            if success:
+                self.logger.success(f"Plot saved to: {abs_filename}")
+            else:
+                # Fall back to HTML if SVG export fails completely
+                html_filename = filename.replace('.svg', '.html')
+                abs_html_filename = os.path.abspath(html_filename)
+                if plot_type == "panel":
+                    plot_obj.save(html_filename, embed=True)  # type: ignore[attr-defined]
+                elif plot_type == "holoviews":
+                    import panel
+                    panel.panel(plot_obj).save(html_filename, embed=True)  # type: ignore[attr-defined]
+                elif plot_type == "bokeh":
+                    from bokeh.plotting import output_file
+                    from bokeh.io import save
+                    output_file(html_filename)
+                    save(plot_obj)
+                self.logger.warning(f"SVG export not available, saved as HTML instead: {abs_html_filename}")
         elif filename.endswith(".pdf"):
             # Try to save as PDF, fall back to HTML if not available
             try:
@@ -444,6 +554,472 @@ def plot_chrom(
     self._handle_sample_plot_output(layout, filename, "panel")
 
 
+def _create_raster_plot(sample, mz_range=None, rt_range=None, raster_cmap='greys', 
+                       raster_log=True, raster_min=1, raster_dynamic=True, raster_threshold=0.8, raster_max_px=8,
+                       width=750, height=600, filename=None):
+    """Create the raster plot layer from MS1 data."""
+    # Process colormap using the cmap package with proper error handling
+    raster_cmap_processed = _process_cmap(raster_cmap if raster_cmap is not None else 'greys', fallback="greys", logger=sample.logger)
+
+    # get columns rt, mz, inty from sample.ms1_df, It's polars DataFrame
+    spectradf = sample.ms1_df.to_pandas()
+
+    # remove any inty<raster_min
+    spectradf = spectradf[spectradf["inty"] >= raster_min]
+    # keep only rt, mz, and inty
+    spectradf = spectradf[["rt", "mz", "inty"]]
+    if mz_range is not None:
+        spectradf = spectradf[(spectradf["mz"] >= mz_range[0]) & (spectradf["mz"] <= mz_range[1])]
+    if rt_range is not None:
+        spectradf = spectradf[(spectradf["rt"] >= rt_range[0]) & (spectradf["rt"] <= rt_range[1])]
+
+    maxrt = spectradf["rt"].max()
+    minrt = spectradf["rt"].min()
+    maxmz = spectradf["mz"].max()
+    minmz = spectradf["mz"].min()
+
+    def new_bounds_hook(plot, elem):
+        x_range = plot.state.x_range
+        y_range = plot.state.y_range
+        x_range.bounds = minrt, maxrt
+        y_range.bounds = minmz, maxmz
+
+    points = hv.Points(
+        spectradf,
+        kdims=["rt", "mz"],
+        vdims=["inty"],
+        label="MS1 survey scans",
+    ).opts(
+        fontsize={"title": 16, "labels": 14, "xticks": 6, "yticks": 12},
+        color=np.log(dim("inty")),
+        colorbar=True,
+        cmap="Magma",
+        tools=["hover"],
+    )
+
+    if filename is not None:
+        dyn = False
+        if not filename.endswith(".html"):
+            raster_dynamic = False
+
+    dyn = raster_dynamic
+    raster = hd.rasterize(
+        points,
+        aggregator=ds.max("inty"),
+        interpolation="bilinear",
+        dynamic=dyn,
+    ).opts(
+        active_tools=["box_zoom"],
+        cmap=raster_cmap_processed,
+        tools=["hover"],
+        hooks=[new_bounds_hook],
+        width=width,
+        height=height,
+        cnorm="log" if raster_log else "linear",
+        xlabel="Retention time (s)",
+        ylabel="m/z",
+        colorbar=True,
+        colorbar_position="right",
+        axiswise=True
+    )
+    raster = hd.dynspread(
+        raster,
+        threshold=raster_threshold,
+        how="add",
+        shape="square",
+        max_px=raster_max_px,
+    )
+    
+    return raster
+
+
+def _load_and_merge_oracle_data(sample, oracle_folder, link_by_feature_uid, 
+                               min_id_level, max_id_level, min_ms_level):
+    """Load oracle data and merge with features."""
+    if sample.features_df is None:
+        sample.logger.error("Cannot plot 2D oracle: features_df is not available")
+        return None
+        
+    feats = sample.features_df.clone()
+    sample.logger.debug(f"Features data shape: {len(feats)} rows")
+
+    # Convert to pandas for oracle operations that require pandas functionality
+    if hasattr(feats, "to_pandas"):
+        feats = feats.to_pandas()
+
+    # check if annotationfile is not None
+    if oracle_folder is None:
+        sample.logger.info("No oracle folder provided, plotting features only")
+        return None
+    
+    # try to read the annotationfile as a csv file and add it to feats
+    oracle_file_path = os.path.join(oracle_folder, "diag", "summary_by_feature.csv")
+    sample.logger.debug(f"Loading oracle data from: {oracle_file_path}")
+    try:
+        oracle_data = pd.read_csv(oracle_file_path)
+        sample.logger.info(f"Oracle data loaded successfully with {len(oracle_data)} rows")
+    except Exception as e:
+        sample.logger.error(f"Could not read {oracle_file_path}: {e}")
+        return None
+
+    if link_by_feature_uid:
+        cols_to_keep = [
+            "title", "scan_idx", "mslevel", "hits", "id_level", "id_label",
+            "id_ion", "id_class", "id_evidence", "score", "score2",
+        ]
+        oracle_data = oracle_data[cols_to_keep]
+        
+        # extract feature_uid from title. It begins with "uid:XYZ,"
+        sample.logger.debug("Extracting feature UIDs from oracle titles using pattern 'uid:(\\d+)'")
+        oracle_data["feature_uid"] = oracle_data["title"].str.extract(r"uid:(\d+)")
+        oracle_data["feature_uid"] = oracle_data["feature_uid"].astype(int)
+        
+        # sort by id_level, remove duplicate feature_uid, keep the first one
+        sample.logger.debug("Sorting by ID level and removing duplicates")
+        oracle_data = oracle_data.sort_values(by=["id_level"], ascending=False)
+        oracle_data = oracle_data.drop_duplicates(subset=["feature_uid"], keep="first")
+        sample.logger.debug(f"After deduplication: {len(oracle_data)} unique oracle annotations")
+    else:
+        cols_to_keep = [
+            "precursor", "rt", "title", "scan_idx", "mslevel", "hits", "id_level",
+            "id_label", "id_ion", "id_class", "id_evidence", "score", "score2",
+        ]
+        oracle_data = oracle_data[cols_to_keep]
+        oracle_data["feature_uid"] = None
+        
+        # iterate over the rows and find the feature_uid in feats by looking at the closest rt and mz
+        for i, row in oracle_data.iterrows():
+            candidates = feats[
+                (abs(feats["rt"] - row["rt"]) < 1) & (abs(feats["mz"] - row["precursor"]) < 0.005)
+            ].copy()
+            if len(candidates) > 0:
+                # sort by delta rt
+                candidates["delta_rt"] = abs(candidates["rt"] - row["rt"])
+                candidates = candidates.sort_values(by=["delta_rt"])
+                oracle_data.at[i, "feature_uid"] = candidates["feature_uid"].values[0]
+        # remove precursor and rt columns
+        oracle_data = oracle_data.drop(columns=["precursor", "rt"])
+
+    # Merge features with oracle data
+    sample.logger.debug(f"Merging {len(feats)} features with oracle data")
+    feats = feats.merge(oracle_data, how="left", on="feature_uid")
+    sample.logger.debug(f"After merge: {len(feats)} total features")
+
+    # filter feats by id_level
+    initial_count = len(feats)
+    if min_id_level is not None:
+        feats = feats[(feats["id_level"] >= min_id_level)]
+        sample.logger.debug(f"After min_id_level filter ({min_id_level}): {len(feats)} features")
+    if max_id_level is not None:
+        feats = feats[(feats["id_level"] <= max_id_level)]
+        sample.logger.debug(f"After max_id_level filter ({max_id_level}): {len(feats)} features")
+    if min_ms_level is not None:
+        feats = feats[(feats["mslevel"] >= min_ms_level)]
+        sample.logger.debug(f"After min_ms_level filter ({min_ms_level}): {len(feats)} features")
+    
+    sample.logger.info(f"Feature filtering complete: {initial_count} → {len(feats)} features remaining")
+    return feats
+
+
+def _setup_color_mapping(sample, feats, colorby, cmap, legend_groups=None):
+    """Set up categorical color mapping for features."""
+    import matplotlib.colors as mcolors
+    
+    feats["color"] = "black"  # Default fallback color
+    cvalues = None
+    color_column = "color"  # Default to fixed color
+    colors = []
+    
+    # Determine which column to use for categorical coloring
+    if colorby in ["class", "hg", "id_class", "id_hg"]:
+        categorical_column = "id_class"
+        # replace nans with 'mix'
+        feats[categorical_column] = feats[categorical_column].fillna("mix")
+    elif colorby in ["ion", "id_ion"]:
+        categorical_column = "id_ion"
+        feats[categorical_column] = feats[categorical_column].fillna("mix")
+    elif colorby in ["evidence", "id_evidence"]:
+        categorical_column = "id_evidence" 
+        feats[categorical_column] = feats[categorical_column].fillna("mix")
+    elif colorby in ["level", "id_level"]:
+        categorical_column = "id_level"
+        feats[categorical_column] = feats[categorical_column].fillna("mix")
+    else:
+        categorical_column = None
+    
+    if categorical_column is not None:
+        # Use provided legend_groups or derive from data
+        if legend_groups is not None:
+            # Use all specified groups to ensure consistent legend/coloring
+            cvalues = legend_groups[:]  # Copy the list
+            # Ensure 'mix' is always present as the last group if not already included
+            if 'mix' not in cvalues:
+                cvalues.append('mix')
+            sample.logger.info(f"Using provided legend_groups for legend: {cvalues}")
+            
+            # Check which provided groups actually have data
+            present_groups = feats[categorical_column].unique()
+            missing_groups = [grp for grp in cvalues if grp not in present_groups]
+            if missing_groups:
+                sample.logger.warning(f"Provided legend_groups not found in data: {missing_groups}")
+            sample.logger.info(f"Groups present in data: {sorted(present_groups)}")
+            
+            # Assign any points not in legend_groups to 'mix'
+            feats.loc[~feats[categorical_column].isin(cvalues[:-1]), categorical_column] = 'mix'
+        else:
+            # Original behavior: use only groups present in data
+            cvalues = feats[categorical_column].unique()
+            # sort alphabetically
+            cvalues = sorted(cvalues)
+            # flip the strings left to right
+            fcvalues = [cvalues[i][::-1] for i in range(len(cvalues))]
+            # sort in alphabetical order the flipped strings and return the index
+            idx = np.argsort(fcvalues)
+            # apply to cvalues
+            cvalues = [cvalues[i] for i in idx]
+            sample.logger.info(f"Using groups derived from data: {cvalues}")
+
+        color_column = categorical_column  # Use categorical coloring
+
+    # Process colormap for categorical data
+    if cvalues is not None:
+        num_colors = len(cvalues)
+        
+        # Use colormap for categorical data - use _process_cmap for proper handling
+        try:
+            colormap = Colormap(cmap)
+            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 - handle different color formats
+                if hasattr(color, '__len__') and len(color) >= 3:
+                    # It's an array-like color (RGB or RGBA)
+                    colors.append(mcolors.to_hex(color[:3]))
+                else:
+                    # It's a single value, convert to RGB
+                    colors.append(mcolors.to_hex([color, color, color]))
+        except (AttributeError, ValueError, TypeError):
+            # Fallback to using _process_cmap if direct Colormap fails
+            cmap_palette = _process_cmap(cmap, fallback="viridis", logger=sample.logger)
+            # Sample colors from the palette
+            colors = []
+            for i in range(num_colors):
+                idx = int(i * (len(cmap_palette) - 1) / (num_colors - 1)) if num_colors > 1 else len(cmap_palette) // 2
+                colors.append(cmap_palette[idx])
+
+        # Create a mapping from class name to color to ensure consistent color assignment
+        # Each class gets the same color based on its position in the cvalues list
+        class_to_color = {class_name: colors[i] for i, class_name in enumerate(cvalues)}
+
+        # assign color to each row based on colorby category
+        feats["color"] = "black"
+        for class_name, color in class_to_color.items():
+            if colorby in ["class", "hg", "id_class", "id_hg"]:
+                feats.loc[feats["id_class"] == class_name, "color"] = color
+            elif colorby in ["ion", "id_ion"]:
+                feats.loc[feats["id_ion"] == class_name, "color"] = color
+            elif colorby in ["id_evidence", "ms2_evidence"]:
+                feats.loc[feats["id_evidence"] == class_name, "color"] = color
+
+    return cvalues, color_column, colors
+
+
+def _create_feature_overlay(sample, raster, feats, cvalues, color_column, colors,
+                           markersize, title, legend):
+    """Create feature overlay with identified and unidentified features."""
+    # replace NaN with 0 in id_level
+    feats["id_level"] = feats["id_level"].fillna(0)
+    
+    # Create unified visualization with all features in single layer
+    # This avoids the multiple layer legend conflicts that cause dark colors and shared toggling
+    sample.logger.debug("Creating unified feature visualization with categorical coloring")
+    
+    # Prepare categorical coloring for identified features only (id_level >= 1)
+    identified_feats = feats[feats["id_level"] >= 1].copy() if len(feats[feats["id_level"] >= 1]) > 0 else pd.DataFrame()
+    unidentified_feats = feats[feats["id_level"] < 1].copy() if len(feats[feats["id_level"] < 1]) > 0 else pd.DataFrame()
+    
+    overlay = raster
+    
+    # Single layer for identified features with categorical coloring
+    if len(identified_feats) > 0 and cvalues is not None:
+        # Create proper confidence-based marker styling
+        identified_feats["marker_style"] = identified_feats["id_level"].apply(
+            lambda x: "circle" if x >= 2 else "circle_cross"
+        )
+        identified_feats["fill_alpha"] = identified_feats["id_level"].apply(
+            lambda x: 1.0 if x >= 2 else 0.3  # Full opacity for high conf, transparent for medium
+        )
+        
+        oracle_hover_identified = HoverTool(
+            tooltips=[
+                ("rt", "@rt"),
+                ("m/z", "@mz{0.0000}"),
+                ("feature_uid", "@feature_uid"),
+                ("id_level", "@id_level"),
+                ("id_class", "@id_class"),
+                ("id_label", "@id_label"),
+                ("id_ion", "@id_ion"),
+                ("id_evidence", "@id_evidence"),
+                ("score", "@score"),
+                ("score2", "@score2"),
+            ],
+        )
+        
+        # Create completely separate overlay elements for each category
+        overlays_to_combine = [raster]  # Start with raster base
+        
+        for i, category in enumerate(cvalues):
+            category_data = identified_feats[identified_feats[color_column] == category].copy()
+            if len(category_data) > 0:
+                # Create a completely separate Points element for this category
+                category_points = hv.Points(
+                    category_data,
+                    kdims=["rt", "mz"],
+                    vdims=[
+                        "inty", "feature_uid", "id_level", "id_class", "id_label",
+                        "id_ion", "id_evidence", "score", "score2", "fill_alpha"
+                    ],
+                    label=str(category)  # This becomes the legend label
+                ).options(
+                    color=colors[i],  # Use pre-computed hex color for this category
+                    marker="circle",
+                    size=markersize,
+                    alpha="fill_alpha",
+                    tools=[oracle_hover_identified],
+                    show_legend=True,
+                )
+                overlays_to_combine.append(category_points)
+            else:
+                # Create empty Points element for categories with no data to ensure they appear in legend
+                empty_data = pd.DataFrame(columns=['rt', 'mz', 'inty', 'feature_uid', 'id_level', 
+                                                  'id_class', 'id_label', 'id_ion', 'id_evidence', 
+                                                  'score', 'score2', 'fill_alpha'])
+                category_points = hv.Points(
+                    empty_data,
+                    kdims=["rt", "mz"],
+                    vdims=[
+                        "inty", "feature_uid", "id_level", "id_class", "id_label",
+                        "id_ion", "id_evidence", "score", "score2", "fill_alpha"
+                    ],
+                    label=str(category)  # This becomes the legend label
+                ).options(
+                    color=colors[i],  # Use pre-computed hex color for this category
+                    marker="circle",
+                    size=markersize,
+                    alpha=1.0,
+                    tools=[oracle_hover_identified],
+                    show_legend=True,
+                )
+                overlays_to_combine.append(category_points)
+        
+        # Combine all overlays
+        overlay = overlays_to_combine[0]  # Start with raster
+        for layer in overlays_to_combine[1:]:
+            overlay = overlay * layer
+            
+    else:
+        # No categorical data - just set overlay to raster
+        overlay = raster
+        
+    # Separate layer for unidentified features (always black crosses)
+    if len(unidentified_feats) > 0:
+        oracle_hover_no_id = HoverTool(
+            tooltips=[
+                ("rt", "@rt"),
+                ("m/z", "@mz{0.0000}"),
+                ("feature_uid", "@feature_uid"),
+                ("id_level", "@id_level"),
+            ],
+        )
+        
+        feature_points_no_id = hv.Points(
+            unidentified_feats,
+            kdims=["rt", "mz"],
+            vdims=["inty", "feature_uid", "id_level"],
+        ).options(
+            color="black",
+            marker="x", 
+            size=markersize,
+            alpha=1.0,
+            tools=[oracle_hover_no_id],
+            show_legend=False,
+        )
+        
+        overlay = overlay * feature_points_no_id
+
+    if title is not None:
+        sample.logger.debug(f"Setting plot title: {title}")
+        overlay = overlay.opts(title=title)
+
+    # Configure legend if requested and categorical coloring is available
+    if legend is not None and cvalues is not None and len(cvalues) > 1:
+        sample.logger.debug(f"Configuring integrated legend at '{legend}' position with {len(cvalues)} categories: {cvalues}")
+        
+        # Map legend position parameter to HoloViews legend position
+        legend_position_map = {
+            "top_right": "top_right",
+            "top_left": "top_left", 
+            "bottom_right": "bottom_right",
+            "bottom_left": "bottom_left",
+            "right": "right",
+            "left": "left",
+            "top": "top",
+            "bottom": "bottom"
+        }
+        
+        hv_legend_pos = legend_position_map.get(legend, "bottom_right")
+        
+        # Apply legend configuration to the overlay
+        overlay = overlay.opts(
+            legend_position=hv_legend_pos,
+            legend_opts={'title': '', 'padding': 2, 'spacing': 2}
+        )
+        
+        sample.logger.debug(f"Applied integrated legend at position '{hv_legend_pos}'")
+    elif legend is None:
+        # Explicitly hide legend when legend=None
+        overlay = overlay.opts(show_legend=False)
+        sample.logger.debug("Legend hidden (legend=None)")
+        
+    return overlay
+
+
+def _handle_output(sample, overlay, filename):
+    """Handle plot export or display."""
+    if filename is not None:
+        # if filename includes .html, save the layout to an HTML file
+        if filename.endswith(".html"):
+            # For HoloViews overlay, we need to convert to Panel for saving
+            panel.Column(overlay).save(filename, embed=True)
+        elif filename.endswith(".svg"):
+            success = _export_with_webdriver_manager(overlay, filename, "svg", sample.logger)
+            if success:
+                sample.logger.success(f"SVG exported: {os.path.abspath(filename)}")
+            else:
+                sample.logger.warning(f"SVG export failed: {os.path.abspath(filename)}")
+        elif filename.endswith(".png"):
+            success = _export_with_webdriver_manager(overlay, filename, "png", sample.logger)
+            if success:
+                sample.logger.success(f"PNG exported: {os.path.abspath(filename)}")
+            else:
+                sample.logger.warning(f"PNG export failed: {os.path.abspath(filename)}")
+        else:
+            # Default to PNG for any other format
+            png_filename = filename + ".png" if not filename.endswith(('.png', '.svg', '.html')) else filename
+            success = _export_with_webdriver_manager(overlay, png_filename, "png", sample.logger)
+            if success:
+                sample.logger.success(f"PNG exported: {os.path.abspath(png_filename)}")
+            else:
+                sample.logger.warning(f"PNG export failed: {os.path.abspath(png_filename)}")
+    else:
+        # Create a Panel layout for consistent alignment with plot_2d()
+        layout = panel.Column(overlay)
+        # Return the Panel layout (consistent with plot_2d behavior)
+        return layout
+
+
 def plot_2d(
     self,
     filename=None,
@@ -457,6 +1033,8 @@ def plot_2d(
     marker="circle",
     markersize=5,
     size="static",
+    raster_log=True,
+    raster_min=1,
     raster_dynamic=True,
     raster_max_px=8,
     raster_threshold=0.8,
@@ -499,6 +1077,10 @@ def plot_2d(
             - 'dynamic': Uses coordinate-based sizing that scales with zoom level (markers get larger when zooming in)
             - 'static': Uses screen-based sizing that remains constant regardless of zoom level
             - 'slider': Provides an interactive slider to dynamically adjust marker size
+        raster_log (bool, default True):
+            Use logarithmic scaling for raster intensity (True) or linear scaling (False).
+        raster_min (float, default 1):
+            Minimum intensity threshold for raster data filtering.
         raster_dynamic (bool, default True):
             Whether to use dynamic rasterization for the background point cloud.
         raster_max_px (int, default 8):
@@ -531,8 +1113,8 @@ def plot_2d(
 
     # get columns rt, mz, inty from self.ms1_df, It's polars DataFrame
     spectradf = self.ms1_df.select(["rt", "mz", "inty"])
-    # remove any inty<1
-    spectradf = spectradf.filter(pl.col("inty") >= 1)
+    # remove any inty<raster_min
+    spectradf = spectradf.filter(pl.col("inty") >= raster_min)
     # keep only rt, mz, and inty
     spectradf = spectradf.select(["rt", "mz", "inty"])
     if mz_range is not None:
@@ -652,7 +1234,7 @@ def plot_2d(
         hooks=[new_bounds_hook],
         width=width,
         height=height,
-        cnorm="log",
+        cnorm="log" if raster_log else "linear",
         xlabel="Retention time (s)",
         ylabel="m/z",
         colorbar=True,
@@ -1008,492 +1590,128 @@ def plot_2d(
 def plot_2d_oracle(
     self,
     oracle_folder=None,
-    link_by_feature_uid=None,
+    link_by_feature_uid=True,
+    min_id_level=1,
+    max_id_level=4,
+    min_ms_level=2,
     colorby="hg",
-    filename=None,
-    min_id_level=None,
-    max_id_level=None,
-    min_ms_level=None,
-    title=None,
-    cmap=None,
-    markersize=10,
+    legend_groups=None,
+    markersize=5,
+    cmap='Turbo',
+    raster_cmap='grey',
+    raster_log=True,
+    raster_min=1,
     raster_dynamic=True,
     raster_max_px=8,
     raster_threshold=0.8,
     mz_range=None,
     rt_range=None,
+    width=750,
+    height=600,
+    filename=None,
+    title=None,
+    legend="bottom_right",
 ):
     """
-    Plot a 2D overlay visualization of MS1 survey scans and feature annotations, including oracle annotation data if provided.
+    Plot a 2D visualization combining MS1 raster data and oracle-annotated features.
 
-    This function reads the primary mass spectrometry data, applies filtering, processes oracle annotation data (if provided),
-    and produces an interactive plot combining various data layers. The visualization includes rasterized MS1 data and feature
-    points colored by annotation.
+    Creates an interactive plot overlaying MS1 survey scan data with feature annotations 
+    from oracle files. Features are colored categorically based on identification class,
+    ion type, or evidence level.
 
     Parameters:
-        self: The object instance containing MS1 and feature data.
-        oracle_folder (str, optional): Path to the oracle folder containing the annotation file
-            (expected at "<oracle_folder>/diag/summary_by_feature.csv"). If None, oracle data is not used.
-        link_by_feature_uid (bool, optional): Whether to link features by their IDs in the overlay.
-        colorby (str, optional): Parameter that determines the color assignment for annotated features.
-            Expected values include 'hg', 'class', 'id_class', or 'id_hg'. Default is 'hg'.
-        filename (str, optional): Name of the file where the plot should be saved. If provided and ends with
-            ".html", the panel layout is saved as an interactive HTML file; otherwise, the output is saved as a PNG.
-        min_id_level (int, optional): Minimum identification level for oracle annotations to include.
-        max_id_level (int, optional): Maximum identification level for oracle annotations to include.
-        min_ms_level (int, optional): Minimum MS level for features to include.
-        title (str, optional): Title to be displayed on the resulting plot. Default is None.
-        cmap (str, optional): Colormap to be used for the rasterized plot. Acceptable values include None, "grey",
-            "iridescent", or other valid colormap names. Default is None. When None, 'Greys256' is used.
-        markersize (int, optional): Marker size for feature points in the overlay. Default is 10.
-        raster_dynamic (bool, optional): If True, enables dynamic rasterization of the overlay. If filename is provided
-            and does not end with ".html", raster_dynamic is set to False. Default is True.
-        raster_max_px (int, optional): Maximum pixel size for dynamic rasterization. Default is 8.
-        raster_threshold (float, optional): Threshold for dynamic raster spread. Default is 0.8.
-        mz_range (tuple, optional): m/z range for filtering MS1 data.
-        rt_range (tuple, optional): Retention time range for filtering MS1 data.
+        oracle_folder (str, optional): Path to oracle folder containing 
+            "diag/summary_by_feature.csv". Required for oracle annotations.
+        link_by_feature_uid (bool): Whether to link features by UID (True) or by m/z/RT proximity.
+        min_id_level (int): Minimum identification confidence level to include.
+        max_id_level (int): Maximum identification confidence level to include.
+        min_ms_level (int): Minimum MS level for features to include.
+        colorby (str): Feature coloring scheme - "id_class", "id_ion", "id_evidence", etc.
+        legend_groups (list, optional): List of groups to include in legend and coloring scheme.
+            If provided, legend will show exactly these groups. 'mix' is automatically added
+            as the last group to contain points not matching other groups. Works for all
+            categorical coloring types (id_class, id_ion, id_evidence, etc.).
+            If None (default), all groups present in the data will be shown without filtering.
+            All specified classes will appear in the legend even if no features are present.
+        markersize (int): Size of feature markers.
+        cmap (str): Colormap name for categorical coloring.
+        raster_cmap (str): Colormap for MS1 raster background.
+        raster_log (bool): Use logarithmic scaling for raster intensity (True) or linear scaling (False).
+        raster_min (float): Minimum intensity threshold for raster data filtering.
+        raster_dynamic (bool): Enable dynamic rasterization.
+        raster_threshold (float): Dynamic raster spread threshold.
+        raster_max_px (int): Maximum pixel size for rasterization.
+        mz_range (tuple, optional): m/z range filter (min, max).
+        rt_range (tuple, optional): Retention time range filter (min, max).
+        width/height (int): Plot dimensions in pixels.
+        filename (str, optional): Export filename (.html/.svg/.png). If None, displays inline.
+        title (str, optional): Plot title.
+        legend (str, optional): Legend position ("top_right", "bottom_left", etc.) or None.
 
     Returns:
-        None
-
-    The function either displays the interactive panel layout or saves the visualization to a file based on
-    the provided filename. If the primary file object or feature data is missing, the function prints an
-    informative message and returns without plotting.
+        HoloViews layout for display (if filename is None), otherwise None.
     """
 
-    if self.file_obj is None:
-        print("Please load a file first.")
-        return
-
-    # 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()
-
-    # remove any inty<1
-    spectradf = spectradf[spectradf["inty"] >= 1]
-    # keep only rt, mz, and inty
-    spectradf = spectradf[["rt", "mz", "inty"]]
-    if mz_range is not None:
-        spectradf = spectradf[(spectradf["mz"] >= mz_range[0]) & (spectradf["mz"] <= mz_range[1])]
-    if rt_range is not None:
-        spectradf = spectradf[(spectradf["rt"] >= rt_range[0]) & (spectradf["rt"] <= rt_range[1])]
-
-    maxrt = spectradf["rt"].max()
-    minrt = spectradf["rt"].min()
-    maxmz = spectradf["mz"].max()
-    minmz = spectradf["mz"].min()
-
-    def new_bounds_hook(plot, elem):
-        x_range = plot.state.x_range
-        y_range = plot.state.y_range
-        x_range.bounds = minrt, maxrt
-        y_range.bounds = minmz, maxmz
-
-    points = hv.Points(
-        spectradf,
-        kdims=["rt", "mz"],
-        vdims=["inty"],
-        label="MS1 survey scans",
-    ).opts(
-        fontsize={"title": 16, "labels": 14, "xticks": 6, "yticks": 12},
-        color=np.log(dim("inty")),
-        colorbar=True,
-        cmap="Magma",
-        tools=["hover"],
-    )
-
-    if filename is not None:
-        dyn = False
-        if not filename.endswith(".html"):
-            raster_dynamic = False
-
-    dyn = raster_dynamic
-    raster = hd.rasterize(
-        points,
-        aggregator=ds.max("inty"),
-        interpolation="bilinear",
-        dynamic=dyn,  # alpha=10,                min_alpha=0,
-    ).opts(
-        active_tools=["box_zoom"],
-        cmap=cmap_palette,
-        tools=["hover"],
-        hooks=[new_bounds_hook],
-        width=1000,
-        height=1000,
-        cnorm="log",
-        xlabel="Retention time (s)",
-        ylabel="m/z",
-        colorbar=True,
-        colorbar_position="right",
-        axiswise=True,
-    )
-    raster = hd.dynspread(
-        raster,
-        threshold=raster_threshold,
-        how="add",
-        shape="square",
-        max_px=raster_max_px,
-    )
+    self.logger.info(f"Starting plot_2d_oracle with oracle_folder: {oracle_folder}")
+    self.logger.debug(f"Parameters - link_by_feature_uid: {link_by_feature_uid}, min_id_level: {min_id_level}, max_id_level: {max_id_level}")
+    self.logger.debug(f"Plot parameters - colorby: {colorby}, markersize: {markersize}, filename: {filename}")
 
+    # Early validation
     if self.features_df is None:
+        self.logger.error("Cannot plot 2D oracle: features_df is not available")
         return
-    feats = self.features_df.clone()
-
-    # Convert to pandas for oracle operations that require pandas functionality
-    if hasattr(feats, "to_pandas"):
-        feats = feats.to_pandas()
-
-    # check if annotationfile is not None
+        
     if oracle_folder is None:
-        return
-    # try to read the annotationfile as a csv file and add it to feats
-    try:
-        oracle_data = pd.read_csv(
-            os.path.join(oracle_folder, "diag", "summary_by_feature.csv"),
-        )
-    except Exception:
-        print(f"Could not read {oracle_folder}/diag/summary_by_feature.csv")
+        self.logger.info("No oracle folder provided, plotting features only")
         return
 
-    if link_by_feature_uid:
-        # scan_idx	slaw_id	slaw_ms2_id	mz	rt	level	formula	ion	species	name	rarity	lib_id	hg	mod	lib	score	score2	score_db	score_db_data	ms2_tic	ms2_evidence	ms2_matched_n	ms2_missed_n	ms2_matched	ms2_missed	ms2_top1
-        cols_to_keep = [
-            "title",
-            "scan_idx",
-            "mslevel",
-            "hits",
-            "id_level",
-            "id_label",
-            "id_ion",
-            "id_class",
-            "id_evidence",
-            "score",
-            "score2",
-        ]
-        oracle_data = oracle_data[cols_to_keep]
-        # extract feature_uid from title. It begins with "fid:XYZ;"
-        oracle_data["feature_uid"] = oracle_data["title"].str.extract(r"fid:(\d+)")
-        oracle_data["feature_uid"] = oracle_data["feature_uid"].astype(int)
-        # sort by id_level, remove duplicate feature_uid, keep the first one
-        oracle_data = oracle_data.sort_values(by=["id_level"], ascending=False)
-        oracle_data = oracle_data.drop_duplicates(subset=["feature_uid"], keep="first")
-    else:
-        cols_to_keep = [
-            "precursor",
-            "rt",
-            "title",
-            "scan_idx",
-            "mslevel",
-            "hits",
-            "id_level",
-            "id_label",
-            "id_ion",
-            "id_class",
-            "id_evidence",
-            "score",
-            "score2",
-        ]
-        # link
-        oracle_data = oracle_data[cols_to_keep]
-        oracle_data["feature_uid"] = None
-        # iterate over the rows and find the feature_uid in feats by looking at the closest rt and mz
-        for i, row in oracle_data.iterrows():
-            candidates = feats[
-                (abs(feats["rt"] - row["rt"]) < 1) & (abs(feats["mz"] - row["precursor"]) < 0.005)
-            ].copy()
-            if len(candidates) > 0:
-                # sort by delta rt
-                candidates["delta_rt"] = abs(candidates["rt"] - row["rt"])
-                candidates = candidates.sort_values(by=["delta_rt"])
-                oracle_data.at[i, "feature_uid"] = candidates["feature_uid"].values[0]
-        # remove precursor and rt columns
-        oracle_data = oracle_data.drop(columns=["precursor", "rt"])
-
-    feats = feats.merge(oracle_data, how="left", on="feature_uid")
-
-    # filter feats by id_level
-    if min_id_level is not None:
-        feats = feats[(feats["id_level"] >= min_id_level)]
-    if max_id_level is not None:
-        feats = feats[(feats["id_level"] <= max_id_level)]
-    if min_ms_level is not None:
-        feats = feats[(feats["mslevel"] >= min_ms_level)]
-
-    feats["color"] = "black"
-
-    cvalues = None
-    if colorby in ["class", "hg", "id_class", "id_hg"]:
-        # replace nans in feats['id_class'] with 'mix'
-        feats["id_class"] = feats["id_class"].fillna("mix")
-        cvalues = feats["id_class"].unique()
-        # sort alphabetically
-        cvalues = sorted(cvalues)
-        # flip the strings left to right
-        fcvalues = [cvalues[i][::-1] for i in range(len(cvalues))]
-        # sort in alphabetical order the flipped strings and return the index
-        idx = np.argsort(fcvalues)
-        # apply to cvalues
-        cvalues = [cvalues[i] for i in idx]
-    elif colorby in ["ion", "id_ion"]:
-        cvalues = feats["id_ion"].unique()
-    elif colorby in ["id_evidence", "ms2_evidence"]:
-        cvalues = feats["id_evidence"].unique()
-
-    if cvalues is not None:
-        num_colors = len(cvalues)
-        
-        # 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"
-
-        for i, c in enumerate(cvalues):
-            if colorby in ["class", "hg", "id_class", "id_hg"]:
-                feats.loc[feats["id_class"] == c, "color"] = colors[i]
-            elif colorby in ["ion", "id_ion"]:
-                feats.loc[feats["id_ion"] == c, "color"] = colors[i]
-            elif colorby in ["id_evidence", "ms2_evidence"]:
-                feats.loc[feats["id_evidence"] == c, "color"] = colors[i]
-
-    # replace NaN with 0 in id_level
-    feats["id_level"] = feats["id_level"].fillna(0)
-    # feature_points_1 are all features with column ms2_scans not null
-    feature_points_1 = None
-    feat_df = feats.copy()
-    feat_df = feat_df[feat_df["id_level"] == 2]
-
-    oracle_hover_1 = HoverTool(
-        tooltips=[
-            ("rt", "@rt"),
-            ("m/z", "@mz{0.0000}"),
-            ("feature_uid", "@feature_uid"),
-            ("id_level", "@id_level"),
-            ("id_class", "@id_class"),
-            ("id_label", "@id_label"),
-            ("id_ion", "@id_ion"),
-            ("id_evidence", "@id_evidence"),
-            ("score", "@score"),
-            ("score2", "@score2"),
-        ],
-    )
-    feature_points_1 = hv.Points(
-        feat_df,
-        kdims=["rt", "mz"],
-        vdims=[
-            "inty",
-            "feature_uid",
-            "id_level",
-            "id_class",
-            "id_label",
-            "id_ion",
-            "id_evidence",
-            "score",
-            "score2",
-            "color",
-        ],
-        label="ID by MS2",
-    ).options(
-        color="color",
-        marker="circle",
-        size=markersize,
-        fill_alpha=1.0,
-        tools=[oracle_hover_1],
+    # Create raster plot layer
+    raster = _create_raster_plot(
+        self,
+        mz_range=mz_range,
+        rt_range=rt_range,
+        raster_cmap=raster_cmap,
+        raster_log=raster_log,
+        raster_min=raster_min,
+        raster_dynamic=raster_dynamic,
+        raster_threshold=raster_threshold,
+        raster_max_px=raster_max_px,
+        width=width,
+        height=height,
+        filename=filename
     )
 
-    # feature_points_2 are all features that have ms2_scans not null and id_level ==1
-    feature_points_2 = None
-    feat_df = feats.copy()
-    feat_df = feat_df[(feat_df["ms2_scans"].notnull()) & (feat_df["id_level"] == 1)]
-    if len(feat_df) > 0:
-        oracle_hover_2 = HoverTool(
-            tooltips=[
-                ("rt", "@rt"),
-                ("m/z", "@mz{0.0000}"),
-                ("feature_uid", "@feature_uid"),
-                ("id_level", "@id_level"),
-                ("id_label", "@id_label"),
-                ("id_ion", "@id_ion"),
-                ("id_class", "@id_class"),
-            ],
-        )
-        feature_points_2 = hv.Points(
-            feat_df,
-            kdims=["rt", "mz"],
-            vdims=[
-                "inty",
-                "feature_uid",
-                "id_level",
-                "id_label",
-                "id_ion",
-                "id_class",
-                "color",
-            ],
-            label="ID by MS1, with MS2",
-        ).options(
-            color="color",
-            marker="circle",
-            size=markersize,
-            fill_alpha=0.0,
-            tools=[oracle_hover_2],
-        )
-
-    # feature_points_3 are all features that have ms2_scans null and id_level ==1
-    feature_points_3 = None
-    feat_df = feats.copy()
-    feat_df = feat_df[(feat_df["ms2_scans"].isnull()) & (feat_df["id_level"] == 1)]
-    if len(feat_df) > 0:
-        oracle_hover_3 = HoverTool(
-            tooltips=[
-                ("rt", "@rt"),
-                ("m/z", "@mz{0.0000}"),
-                ("feature_uid", "@feature_uid"),
-                ("id_level", "@id_level"),
-                ("id_label", "@id_label"),
-                ("id_ion", "@id_ion"),
-                ("id_class", "@id_class"),
-            ],
-        )
-        feature_points_3 = hv.Points(
-            feat_df,
-            kdims=["rt", "mz"],
-            vdims=[
-                "inty",
-                "feature_uid",
-                "id_level",
-                "id_label",
-                "id_ion",
-                "id_class",
-                "color",
-            ],
-            label="ID by MS1, no MS2",
-        ).options(
-            color="color",
-            marker="diamond",
-            size=markersize,
-            fill_alpha=0.0,
-            tools=[oracle_hover_3],
-        )
-
-    # feature_points_4 are all features that have ms2_scans null and id_level ==0
-    feature_points_4 = None
-    feat_df = feats.copy()
-    feat_df = feat_df[(feat_df["ms2_scans"].notnull()) & (feat_df["id_level"] < 1)]
-    if len(feat_df) > 0:
-        oracle_hover_4 = HoverTool(
-            tooltips=[
-                ("rt", "@rt"),
-                ("m/z", "@mz{0.0000}"),
-                ("feature_uid", "@feature_uid"),
-                ("inty", "@inty"),
-            ],
-        )
-        feature_points_4 = hv.Points(
-            feat_df,
-            kdims=["rt", "mz"],
-            vdims=["inty", "feature_uid"],
-            label="No ID, with MS2",
-        ).options(
-            color="gray",
-            marker="circle",
-            size=markersize,
-            fill_alpha=0.0,
-            tools=[oracle_hover_4],
-        )
-
-    # feature_points_5 are all features that have ms2_scans null and id_level ==0
-    feature_points_5 = None
-    feat_df = feats.copy()
-    feat_df = feat_df[(feat_df["ms2_scans"].isnull()) & (feat_df["id_level"] < 1)]
-    if len(feat_df) > 0:
-        oracle_hover_5 = HoverTool(
-            tooltips=[
-                ("rt", "@rt"),
-                ("m/z", "@mz{0.0000}"),
-                ("feature_uid", "@feature_uid"),
-                ("inty", "@inty"),
-            ],
-        )
-        feature_points_5 = hv.Points(
-            feat_df,
-            kdims=["rt", "mz"],
-            vdims=["inty", "feature_uid"],
-            label="No ID, no MS2",
-        ).options(
-            color="gray",
-            marker="diamond",
-            fill_alpha=0.0,
-            size=markersize,
-            tools=[oracle_hover_5],
-        )
-
-    overlay = raster
-
-    if feature_points_1 is not None:
-        overlay = overlay * feature_points_1
-    if feature_points_2 is not None:
-        overlay = overlay * feature_points_2
-    if feature_points_3 is not None:
-        overlay = overlay * feature_points_3
-    if feature_points_4 is not None:
-        overlay = overlay * feature_points_4
-    # if not show_only_features_with_ms2:
-    if feature_points_5 is not None:
-        overlay = overlay * feature_points_5
-
-    if title is not None:
-        overlay = overlay.opts(title=title)
+    # Load and process oracle data
+    feats = _load_and_merge_oracle_data(
+        self,
+        oracle_folder=oracle_folder,
+        link_by_feature_uid=link_by_feature_uid,
+        min_id_level=min_id_level,
+        max_id_level=max_id_level,
+        min_ms_level=min_ms_level
+    )
+    
+    if feats is None:
+        return
 
-    # Create a panel layout
-    layout = panel.Column(overlay)
+    # Set up color scheme and categorical mapping
+    cvalues, color_column, colors = _setup_color_mapping(self, feats, colorby, cmap, legend_groups)
+
+    # Create feature overlay with all visualization elements
+    overlay = _create_feature_overlay(
+        self,
+        raster=raster,
+        feats=feats,
+        cvalues=cvalues,
+        color_column=color_column,
+        colors=colors,
+        markersize=markersize,
+        title=title,
+        legend=legend
+    )
 
-    if filename is not None:
-        # if filename includes .html, save the panel layout to an HTML file
-        if filename.endswith(".html"):
-            layout.save(filename, embed=True)
-        else:
-            # save the panel layout as a png
-            hv.save(overlay, filename, fmt="png")
-    else:
-        # Check if we're in a notebook environment and display appropriately
-        return _display_plot(overlay, layout)
+    # Handle output: export or display
+    return _handle_output(self, overlay, filename)
 
 
 def plot_ms2_eic(
@@ -1756,96 +1974,6 @@ def plot_ms2_cycle(
         max_px=raster_max_px,
     )
 
-    """
-    feature_points_1 = None
-    feature_points_2 = None
-    feature_points_3 = None
-    feature_points_4 = None
-    feature_points_iso = None
-    # Plot features as red dots if features is True
-    if self.features_df is not None and show_features:
-        feats = self.features_df.clone()
-        # Convert to pandas for operations that require pandas functionality
-        if hasattr(feats, 'to_pandas'):
-            feats = feats.to_pandas()
-        # if ms2_scans is not null, keep only the first element of the list
-        feats['ms2_scans'] = feats['ms2_scans'].apply(lambda x: x[0] if type(x) == list else x)
-        # keep only iso==0, i.e. the main
-        feats = feats[feats['iso']==0]
-        # find features with ms2_scans not None  and iso==0
-        features_df = feats[feats['ms2_scans'].notnull()]
-        feature_points_1 = hv.Points(
-        features_df, kdims=["rt", "mz"], vdims=["feature_uid", "inty", "quality", "rt_delta", "ms2_scans"], label="Features with MS2 data"
-        ).options(
-        color=color_1,
-        marker=marker,
-        size=size_1,
-        tools=["hover"],
-        )
-        # find features without MS2 data
-        features_df = feats[feats['ms2_scans'].isnull()]
-        feature_points_2 = hv.Points(
-        features_df, kdims=["rt", "mz"], vdims=["feature_uid", "inty", "quality", "rt_delta"], label="Features without MS2 data"
-        ).options(
-        color='red',
-        size=size_2,
-        marker=marker,
-        tools=["hover"],
-        )
-
-        if show_isotopes:
-            feats = self.features_df
-            features_df = feats[feats['iso']>0]
-            feature_points_iso = hv.Points(
-            features_df, kdims=["rt", "mz"], vdims=["feature_uid", "inty", "quality", "rt_delta", "iso", "iso_of"], label="Isotopes"
-            ).options(
-            color='violet',
-            marker=marker,
-            size=size_1,
-            tools=["hover"],
-            )
-    if show_ms2:
-        # find all self.scans_df with mslevel 2 that are not linked to a feature
-        ms2_orphan = self.scans_df.filter(pl.col('ms_level')==2).filter(pl.col("feature_uid")<0)
-
-        if len(ms2_orphan) > 0:
-            # pandalize
-            ms2 = ms2_orphan.to_pandas()
-            feature_points_3 = hv.Points(
-            ms2, kdims=["rt", "prec_mz"], vdims=["index", "inty_tot", "bl"], label="Orphan MS2 scans"
-            ).options(
-            color=color_2,
-            marker='x',
-            size=size_2,
-            tools=["hover"],
-            )
-
-        ms2_linked = self.scans_df.filter(pl.col('ms_level')==2).filter(pl.col("feature_uid")>=0)
-        if len(ms2_linked) > 0:
-            # pandalize
-            ms2 = ms2_linked.to_pandas()
-            feature_points_4 = hv.Points(
-            ms2, kdims=["rt", "prec_mz"], vdims=["index", "inty_tot", "bl"], label="Linked MS2 scans"
-            ).options(
-            color=color_1,
-            marker='x',
-            size=size_2,
-            tools=["hover"],
-            )
-
-
-    if feature_points_4 is not None:
-        overlay = overlay * feature_points_4
-    if feature_points_3 is not None:
-        overlay = overlay * feature_points_3
-    if feature_points_1 is not None:
-        overlay = overlay * feature_points_1
-    if not show_only_features_with_ms2:
-        if feature_points_2 is not None:
-            overlay = overlay * feature_points_2
-    if feature_points_iso is not None:
-        overlay = overlay * feature_points_iso
-    """
     if title is not None:
         overlay = overlay.opts(title=title)