openms-insight 0.1.0__py3-none-any.whl

openms_insight/components/lineplot.py

@@ -0,0 +1,492 @@
+"""Line plot component using Plotly.js."""
+
+from typing import Any, Dict, Optional
+
+import polars as pl
+
+from ..core.base import BaseComponent
+from ..core.registry import register_component
+from ..preprocessing.filtering import filter_and_collect_cached
+
+
+@register_component("lineplot")
+class LinePlot(BaseComponent):
+    """
+    Interactive stick plot component using Plotly.js.
+
+    Features:
+    - Stick-style peak visualization (vertical lines from baseline)
+    - Highlighting of selected data points
+    - Annotations with labels (mass labels, charge states, etc.)
+    - Zoom controls with auto-fit to highlighted data
+    - SVG export
+    - Click-to-select peaks with gold highlighting
+    - Cross-component linking via filters and interactivity
+
+    LinePlots can have separate `filters` and `interactivity` mappings:
+    - `filters`: Which selections filter this plot's data
+    - `interactivity`: What selection is set when a peak is clicked
+
+    Example:
+        # Plot filters by spectrum, clicking selects a peak
+        plot = LinePlot(
+            cache_id="peaks_plot",
+            data=peaks_df,
+            filters={'spectrum': 'scan_id'},
+            interactivity={'my_selection': 'mass'},
+            x_column='mass',
+            y_column='intensity',
+        )
+    """
+
+    _component_type: str = "lineplot"
+
+    def __init__(
+        self,
+        cache_id: str,
+        data: Optional[pl.LazyFrame] = None,
+        filters: Optional[Dict[str, str]] = None,
+        filter_defaults: Optional[Dict[str, Any]] = None,
+        interactivity: Optional[Dict[str, str]] = None,
+        cache_path: str = ".",
+        regenerate_cache: bool = False,
+        x_column: str = 'x',
+        y_column: str = 'y',
+        title: Optional[str] = None,
+        x_label: Optional[str] = None,
+        y_label: Optional[str] = None,
+        highlight_column: Optional[str] = None,
+        annotation_column: Optional[str] = None,
+        styling: Optional[Dict[str, Any]] = None,
+        config: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ):
+        """
+        Initialize the LinePlot component.
+
+        Args:
+            cache_id: Unique identifier for this component's cache (MANDATORY).
+                Creates a folder {cache_path}/{cache_id}/ for cached data.
+            data: Polars LazyFrame with plot data. Optional if cache exists.
+            filters: Mapping of identifier names to column names for filtering.
+                Example: {'spectrum': 'scan_id'}
+                When 'spectrum' selection exists, plot shows only data where
+                scan_id equals the selected value.
+            filter_defaults: Default values for filters when state is None.
+                Example: {'identification': -1}
+                When 'identification' selection is None, filter uses -1 instead.
+                This enables showing unannotated data when no identification selected.
+            interactivity: Mapping of identifier names to column names for clicks.
+                Example: {'my_selection': 'mass'}
+                When a peak is clicked, sets 'my_selection' to that peak's mass.
+                The selected peak is highlighted in gold (selectedColor).
+            cache_path: Base path for cache storage. Default "." (current dir).
+            regenerate_cache: If True, regenerate cache even if valid cache exists.
+            x_column: Column name for x-axis values
+            y_column: Column name for y-axis values
+            title: Plot title
+            x_label: X-axis label (defaults to x_column)
+            y_label: Y-axis label (defaults to y_column)
+            highlight_column: Optional column name containing boolean/int
+                              indicating which points to highlight
+            annotation_column: Optional column name containing text annotations
+                               to display on highlighted points
+            styling: Style configuration dict with keys:
+                - highlightColor: Color for highlighted points (default: '#E4572E')
+                - selectedColor: Color for clicked/selected peak (default: '#F3A712')
+                - unhighlightedColor: Color for normal points (default: 'lightblue')
+                - annotationBackground: Background color for annotations
+            config: Additional Plotly config options
+            **kwargs: Additional configuration options
+        """
+        self._x_column = x_column
+        self._y_column = y_column
+        self._title = title
+        self._x_label = x_label or x_column
+        self._y_label = y_label or y_column
+        self._highlight_column = highlight_column
+        self._annotation_column = annotation_column
+        self._styling = styling or {}
+        self._plot_config = config or {}
+
+        # Dynamic annotations set at render time (not cached)
+        self._dynamic_annotations: Optional[Dict[str, Any]] = None
+        self._dynamic_title: Optional[str] = None
+
+        super().__init__(
+            cache_id=cache_id,
+            data=data,
+            filters=filters,
+            filter_defaults=filter_defaults,
+            interactivity=interactivity,
+            cache_path=cache_path,
+            regenerate_cache=regenerate_cache,
+            **kwargs
+        )
+
+    def _get_cache_config(self) -> Dict[str, Any]:
+        """
+        Get configuration that affects cache validity.
+
+        Returns:
+            Dict of config values that affect preprocessing
+        """
+        return {
+            'x_column': self._x_column,
+            'y_column': self._y_column,
+            'highlight_column': self._highlight_column,
+            'annotation_column': self._annotation_column,
+        }
+
+    def _get_row_group_size(self) -> int:
+        """
+        Get optimal row group size for parquet writing.
+
+        Filtered plots use smaller row groups (10K) for better predicate
+        pushdown granularity - this allows Polars to skip row groups that
+        don't contain the filter value. Unfiltered plots use larger groups
+        (50K) since we read all data anyway.
+
+        Returns:
+            Number of rows per row group
+        """
+        if self._filters:
+            return 10_000  # Smaller groups for better filter performance
+        return 50_000  # Larger groups for unfiltered plots
+
+    def _validate_mappings(self) -> None:
+        """Validate columns exist in data schema."""
+        super()._validate_mappings()
+
+        schema = self._raw_data.collect_schema()
+        column_names = schema.names()
+
+        # Validate x and y columns exist
+        for col_name, col_label in [(self._x_column, 'x_column'),
+                                     (self._y_column, 'y_column')]:
+            if col_name not in column_names:
+                raise ValueError(
+                    f"{col_label} '{col_name}' not found in data. "
+                    f"Available columns: {column_names}"
+                )
+
+        # Validate optional columns if specified
+        if self._highlight_column and self._highlight_column not in column_names:
+            raise ValueError(
+                f"highlight_column '{self._highlight_column}' not found in data. "
+                f"Available columns: {column_names}"
+            )
+
+        if self._annotation_column and self._annotation_column not in column_names:
+            raise ValueError(
+                f"annotation_column '{self._annotation_column}' not found in data. "
+                f"Available columns: {column_names}"
+            )
+
+    def _preprocess(self) -> None:
+        """
+        Preprocess plot data.
+
+        Sorts by filter columns for efficient predicate pushdown, then
+        collects the LazyFrame for caching by base class.
+        """
+        data = self._raw_data
+
+        # Sort by filter columns for efficient predicate pushdown.
+        # This clusters identical filter values together, enabling Polars
+        # to skip row groups that don't contain the target value when
+        # filtering by selection state.
+        if self._filters:
+            sort_columns = list(self._filters.values())
+            data = data.sort(sort_columns)
+
+        # Store configuration in preprocessed data for serialization
+        self._preprocessed_data['plot_config'] = {
+            'x_column': self._x_column,
+            'y_column': self._y_column,
+            'highlight_column': self._highlight_column,
+            'annotation_column': self._annotation_column,
+        }
+
+        # Collect data for caching (filter happens at render time)
+        # Base class will serialize this to parquet
+        self._preprocessed_data['data'] = data.collect()
+
+    def _get_vue_component_name(self) -> str:
+        """Return the Vue component name."""
+        return 'PlotlyLineplotUnified'
+
+    def _get_data_key(self) -> str:
+        """Return the key used to send primary data to Vue."""
+        return 'plotData'
+
+    def _prepare_vue_data(self, state: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Prepare plot data for Vue component.
+
+        LinePlots filter based on filters mapping if provided.
+
+        Sends data as a pandas DataFrame for efficient Arrow serialization.
+        Vue parses the Arrow table and extracts column arrays for rendering.
+
+        Args:
+            state: Current selection state from StateManager
+
+        Returns:
+            Dict with plotData (pandas DataFrame) and _hash for change detection
+        """
+        # Build list of columns to select (projection pushdown for efficiency)
+        columns_to_select = [self._x_column, self._y_column]
+        if self._highlight_column:
+            columns_to_select.append(self._highlight_column)
+        if self._annotation_column:
+            columns_to_select.append(self._annotation_column)
+        # Include columns needed for interactivity (e.g., peak_id)
+        if self._interactivity:
+            for col in self._interactivity.values():
+                if col not in columns_to_select:
+                    columns_to_select.append(col)
+        # Include filter columns for filtering to work
+        if self._filters:
+            for col in self._filters.values():
+                if col not in columns_to_select:
+                    columns_to_select.append(col)
+
+        # Get cached data (DataFrame or LazyFrame)
+        data = self._preprocessed_data.get('data')
+        if data is None:
+            # Fallback to raw data if available
+            data = self._raw_data
+
+        # Ensure we have a LazyFrame for filtering
+        if isinstance(data, pl.DataFrame):
+            data = data.lazy()
+
+        # Use cached filter+collect - returns (pandas DataFrame, hash)
+        df_pandas, data_hash = filter_and_collect_cached(
+            data,
+            self._filters,
+            state,
+            columns=columns_to_select,
+            filter_defaults=self._filter_defaults,
+        )
+
+        # Determine which highlight/annotation columns to use
+        highlight_col = self._highlight_column
+        annotation_col = self._annotation_column
+
+        # Apply dynamic annotations if set
+        # Annotations are keyed by peak_id (stable identifier from interactivity column)
+        if self._dynamic_annotations and len(df_pandas) > 0:
+            num_rows = len(df_pandas)
+            highlights = [False] * num_rows
+            annotations = [''] * num_rows
+
+            # Get the interactivity column to use for lookup (e.g., 'peak_id')
+            # Use the first interactivity column as the ID column for annotation lookup
+            id_column = None
+            if self._interactivity:
+                id_column = list(self._interactivity.values())[0]
+
+            # Apply annotations by peak_id lookup
+            if id_column and id_column in df_pandas.columns:
+                peak_ids = df_pandas[id_column].tolist()
+                for row_idx, peak_id in enumerate(peak_ids):
+                    if peak_id in self._dynamic_annotations:
+                        ann_data = self._dynamic_annotations[peak_id]
+                        highlights[row_idx] = ann_data.get('highlight', False)
+                        annotations[row_idx] = ann_data.get('annotation', '')
+            else:
+                # Fallback: use row index as key (legacy behavior)
+                for idx, ann_data in self._dynamic_annotations.items():
+                    if isinstance(idx, int) and 0 <= idx < num_rows:
+                        highlights[idx] = ann_data.get('highlight', False)
+                        annotations[idx] = ann_data.get('annotation', '')
+
+            # Add dynamic columns to dataframe
+            df_pandas = df_pandas.copy()
+            df_pandas['_dynamic_highlight'] = highlights
+            df_pandas['_dynamic_annotation'] = annotations
+
+            # Update column names to use dynamic columns
+            highlight_col = '_dynamic_highlight'
+            annotation_col = '_dynamic_annotation'
+
+            # Update hash to include dynamic annotation state
+            import hashlib
+            ann_hash = hashlib.md5(str(sorted(self._dynamic_annotations.keys())).encode()).hexdigest()[:8]
+            data_hash = f"{data_hash}_{ann_hash}"
+
+        # Send as DataFrame for Arrow serialization (efficient binary transfer)
+        # Vue will parse and extract columns using the config
+        return {
+            'plotData': df_pandas,
+            '_hash': data_hash,
+            # Config tells Vue which columns map to x, y, etc.
+            '_plotConfig': {
+                'xColumn': self._x_column,
+                'yColumn': self._y_column,
+                'highlightColumn': highlight_col,
+                'annotationColumn': annotation_col,
+                'interactivityColumns': {
+                    col: col for col in (self._interactivity.values() if self._interactivity else [])
+                },
+            }
+        }
+
+    def _get_component_args(self) -> Dict[str, Any]:
+        """
+        Get component arguments to send to Vue.
+
+        Returns:
+            Dict with all plot configuration for Vue
+        """
+        # Default styling
+        default_styling = {
+            'highlightColor': '#E4572E',
+            'selectedColor': '#F3A712',
+            'unhighlightedColor': 'lightblue',
+            'highlightHiddenColor': '#1f77b4',
+            'annotationColors': {
+                'massButton': '#E4572E',
+                'selectedMassButton': '#F3A712',
+                'sequenceArrow': '#E4572E',
+                'selectedSequenceArrow': '#F3A712',
+                'background': '#f0f0f0',
+                'buttonHover': '#e0e0e0',
+            }
+        }
+
+        # Merge user styling with defaults
+        styling = {**default_styling, **self._styling}
+        if 'annotationColors' in self._styling:
+            styling['annotationColors'] = {
+                **default_styling['annotationColors'],
+                **self._styling['annotationColors']
+            }
+
+        # Use dynamic title if set, otherwise static title
+        title = self._dynamic_title if self._dynamic_title else (self._title or '')
+
+        args: Dict[str, Any] = {
+            'componentType': self._get_vue_component_name(),
+            'title': title,
+            'xLabel': self._x_label,
+            'yLabel': self._y_label,
+            'styling': styling,
+            'config': self._plot_config,
+            # Pass interactivity for click handling (sets selection on peak click)
+            'interactivity': self._interactivity,
+            # Column mappings for Arrow data parsing in Vue
+            'xColumn': self._x_column,
+            'yColumn': self._y_column,
+            'highlightColumn': self._highlight_column,
+            'annotationColumn': self._annotation_column,
+        }
+
+        # Add any extra config options
+        args.update(self._config)
+
+        return args
+
+    def with_styling(
+        self,
+        highlight_color: Optional[str] = None,
+        selected_color: Optional[str] = None,
+        unhighlighted_color: Optional[str] = None,
+    ) -> 'LinePlot':
+        """
+        Update plot styling.
+
+        Args:
+            highlight_color: Color for highlighted points
+            selected_color: Color for selected points
+            unhighlighted_color: Color for unhighlighted points
+
+        Returns:
+            Self for method chaining
+        """
+        if highlight_color:
+            self._styling['highlightColor'] = highlight_color
+        if selected_color:
+            self._styling['selectedColor'] = selected_color
+        if unhighlighted_color:
+            self._styling['unhighlightedColor'] = unhighlighted_color
+        return self
+
+    def with_annotations(
+        self,
+        background_color: Optional[str] = None,
+        button_color: Optional[str] = None,
+        selected_button_color: Optional[str] = None,
+    ) -> 'LinePlot':
+        """
+        Configure annotation styling.
+
+        Args:
+            background_color: Background color for annotation boxes
+            button_color: Color for annotation buttons
+            selected_button_color: Color for selected annotation buttons
+
+        Returns:
+            Self for method chaining
+        """
+        if 'annotationColors' not in self._styling:
+            self._styling['annotationColors'] = {}
+
+        if background_color:
+            self._styling['annotationColors']['background'] = background_color
+        if button_color:
+            self._styling['annotationColors']['massButton'] = button_color
+        if selected_button_color:
+            self._styling['annotationColors']['selectedMassButton'] = selected_button_color
+
+        return self
+
+    def set_dynamic_annotations(
+        self,
+        annotations: Optional[Dict[int, Dict[str, Any]]] = None,
+        title: Optional[str] = None,
+    ) -> 'LinePlot':
+        """
+        Set dynamic annotations to be applied at render time.
+
+        This allows updating peak annotations without recreating the component.
+        Annotations are keyed by the interactivity column value (e.g., peak_id),
+        which provides a stable identifier independent of row order.
+
+        Args:
+            annotations: Dict mapping peak IDs to annotation data.
+                Keys should match values in the first interactivity column.
+                Each entry should have:
+                - 'highlight': bool - whether to highlight this peak
+                - 'annotation': str - label text (e.g., "b3¹⁺")
+                Example: {123: {'highlight': True, 'annotation': 'b2¹⁺'}}
+            title: Optional dynamic title override
+
+        Returns:
+            Self for method chaining
+
+        Example:
+            # Compute annotations for current identification (keyed by peak_id)
+            annotations = {
+                123: {'highlight': True, 'annotation': 'b2¹⁺'},
+                456: {'highlight': True, 'annotation': 'b3¹⁺'},
+            }
+            spectrum_plot.set_dynamic_annotations(annotations, title="PEPTIDER")
+            spectrum_plot(key="plot", state_manager=sm)
+        """
+        self._dynamic_annotations = annotations
+        self._dynamic_title = title
+        return self
+
+    def clear_dynamic_annotations(self) -> 'LinePlot':
+        """
+        Clear any dynamic annotations.
+
+        Returns:
+            Self for method chaining
+        """
+        self._dynamic_annotations = None
+        self._dynamic_title = None
+        return self