Add annotation labels for 3D plots

pyopenms_viz/_bokeh/core.py

@@ -611,10 +611,15 @@ def _add_box_boundaries(self, annotation_data):
         )
         legend_items = []
         for idx, (_, feature) in enumerate(annotation_data.iterrows()):
-            x0 = feature["leftWidth"]
-            x1 = feature["rightWidth"]
-            y0 = feature["IM_leftWidth"]
-            y1 = feature["IM_rightWidth"]
+            x0 = feature[self.annotation_x_lb]
+            x1 = feature[self.annotation_x_ub]
+            y0 = feature[self.annotation_y_lb]
+            y1 = feature[self.annotation_y_ub]
+
+            if self.annotation_colors in feature:
+                color = feature[self.annotation_colors]
+            else:
+                color = next(color_gen)
 
             # Calculate center points and dimensions
             center_x = (x0 + x1) / 2
@@ -627,13 +632,13 @@ def _add_box_boundaries(self, annotation_data):
                 y=center_y,
                 width=width,
                 height=height,
-                color=next(color_gen),
+                color=color,
                 line_dash=self.feature_config.line_type,
                 line_width=self.feature_config.line_width,
                 fill_alpha=0,
             )
-            if "name" in annotation_data.columns:
-                use_name = feature["name"]
+            if self.annotation_names in feature:
+                use_name = feature[self.annotation_names]
             else:
                 use_name = f"Feature {idx}"
             if "q_value" in annotation_data.columns:

pyopenms_viz/_config.py

@@ -126,7 +126,7 @@ def from_dict(cls, legend_dict: Dict[str, Any]) -> "LegendConfig":
 @dataclass(kw_only=True)
 class FeatureConfig:
     def default_legend_factory():
-        return LegendConfig(title="Features", loc="right", bbox_to_anchor=(1.5, 0.5))
+        return LegendConfig(title="Features", loc="left", bbox_to_anchor=(1, 0.5))
 
     colormap: str = "viridis"
     line_width: float = 1

pyopenms_viz/_core.py

@@ -2,6 +2,7 @@
 
 from abc import ABC, abstractmethod
 from typing import Any, Tuple, Literal, Union, List, Dict
+import numpy as np
 import importlib
 import types
 import re
@@ -1032,6 +1033,12 @@ def __init__(
         y_kind="spectrum",
         x_kind="chromatogram",
         annotation_data: DataFrame | None = None,
+        annotation_x_lb : str = 'leftWidth',
+        annotation_x_ub : str = 'rightWidth',
+        annotation_y_lb : str = 'IM_leftWidth',
+        annotation_y_ub : str = 'IM_rightWidth',
+        annotation_colors : str = 'color',
+        annotation_names : str = 'name',
         bin_peaks: Union[Literal["auto"], bool] = "auto",
         aggregation_method: Literal["mean", "sum", "max"] = "mean",
         num_x_bins: int = 50,
@@ -1057,7 +1064,13 @@ def __init__(
             self.annotation_data = annotation_data.copy()
         else:
             self.annotation_data = None
-
+        self.annotation_x_lb = annotation_x_lb
+        self.annotation_x_ub = annotation_x_ub
+        self.annotation_y_lb = annotation_y_lb
+        self.annotation_y_ub = annotation_y_ub
+        self.annotation_colors = annotation_colors
+        self.annotation_names = annotation_names
+
         super().__init__(data, x, y, z=z, **kwargs)
         self._check_and_aggregate_duplicates()
 
@@ -1258,6 +1271,35 @@ def _add_box_boundaries(self, annotation_data):
             None
         """
         pass
+
+    def _compute_3D_annotations(self, annotation_data, x, y, z):
+        def center_of_gravity(x, m):
+            return np.sum(x*m) / np.sum(m)
+
+        # Contains tuple of coordinates + text + color (x, y, z, t, c)
+        annotations_3d = []
+        for _, feature in annotation_data.iterrows():
+            x0 = feature[self.annotation_x_lb]
+            x1 = feature[self.annotation_x_ub]
+            y0 = feature[self.annotation_y_lb]
+            y1 = feature[self.annotation_y_ub]
+            t = feature[self.annotation_names]
+            c = feature[self.annotation_colors]
+            selected_data = self.data[
+                (self.data[x] > x0) & (self.data[x] < x1) 
+                & (self.data[y] > y0) & (self.data[y] < y1)
+            ]
+            if len(selected_data) == 0:
+                annotations_3d.append((
+                    np.mean((x0, x1)), np.mean((y0, y1)), np.mean(self.data[z]), t, c
+                ))
+            else:
+                annotations_3d.append((
+                    center_of_gravity(selected_data[x], selected_data[z]),
+                    center_of_gravity(selected_data[y], selected_data[z]),
+                    np.max(selected_data[z])*1.05, t, c
+                ))
+        return map(list, zip(*annotations_3d))
 
 
 class PlotAccessor:

pyopenms_viz/_matplotlib/core.py

@@ -225,9 +225,9 @@ def show_default(self):
         Show the plot.
         """
         if isinstance(self.fig, Axes):
-            self.fig.get_figure().tight_layout()
+            self.fig.get_figure()
         else:
-            self.superFig.tight_layout()
+            self.superFig
         plt.show()
 
 
@@ -349,27 +349,42 @@ def plot(
 
                 return ax, (legend_lines, legend_labels)
 
+    def _get_annotations():
+        pass
+
+
     def _add_annotations(
         self,
         fig,
         ann_texts: list[list[str]],
         ann_xs: list[float],
         ann_ys: list[float],
         ann_colors: list[str],
+        ann_zs: list[float] = None,
     ):
-        for text, x, y, color in zip(ann_texts, ann_xs, ann_ys, ann_colors):
+        for i, (text, x, y, color) in enumerate(zip(ann_texts, ann_xs, ann_ys, ann_colors)):
             if text is not nan and text != "" and text != "nan":
                 if is_latex_formatted(text):
+                    # Wrap the text in '$' to indicate LaTeX math mode
                     text = "\n".join([r"${}$".format(line) for line in text.split("\n")])
-                fig.annotate(
-                    text,
-                    xy=(x, y),
-                    xytext=(3, 0),
-                    textcoords="offset points",
-                    fontsize=self.annotation_font_size,
-                    color=color,
-                )
-
+                if not self.plot_3d:
+                    fig.annotate(
+                        text,
+                        xy=(x, y),
+                        xytext=(3, 0),
+                        textcoords="offset points",
+                        fontsize=self.annotation_font_size,
+                        color=color,
+                    )
+                else:
+                    fig.text(
+                        x=x, 
+                        y=y, 
+                        z=ann_zs[i], 
+                        s=text, 
+                        fontsize=self.annotation_font_size,
+                        color=color
+                    )
 
 class MATPLOTLIBScatterPlot(MATPLOTLIBPlot, ScatterPlot):
     """
@@ -664,6 +679,13 @@ def create_main_plot(self, x, y, z, class_kwargs, other_kwargs):
                 zlabel=self.zlabel,
                 **other_kwargs,
             )
+            if self.annotation_data is not None:
+                a_x, a_y, a_z, a_t, a_c = self._compute_3D_annotations(
+                    self.annotation_data, x, y, z
+                )
+                vlinePlot._add_annotations(
+                    self.fig, a_t, a_x, a_y, a_c, a_z
+                )
 
     def create_main_plot_marginals(self, x, y, z, class_kwargs, other_kwargs):
         scatterPlot = self.get_scatter_renderer(
@@ -691,16 +713,19 @@ def _add_box_boundaries(self, annotation_data):
         legend_items = []
 
         for idx, (_, feature) in enumerate(annotation_data.iterrows()):
-            x0 = feature["leftWidth"]
-            x1 = feature["rightWidth"]
-            y0 = feature["IM_leftWidth"]
-            y1 = feature["IM_rightWidth"]
+            x0 = feature[self.annotation_x_lb]
+            x1 = feature[self.annotation_x_ub]
+            y0 = feature[self.annotation_y_lb]
+            y1 = feature[self.annotation_y_ub]
 
             # Calculate center points and dimensions
             width = abs(x1 - x0)
             height = abs(y1 - y0)
 
-            color = next(color_gen)
+            if self.annotation_colors in feature:
+                color = feature[self.annotation_colors]
+            else:
+                color = next(color_gen)
             custom_lines = Rectangle(
                 (x0, y0),
                 width,
@@ -709,11 +734,12 @@ def _add_box_boundaries(self, annotation_data):
                 edgecolor=color,
                 linestyle=self.feature_config.line_type,
                 linewidth=self.feature_config.line_width,
+                zorder=5
             )
             self.fig.add_patch(custom_lines)
 
-            if "name" in annotation_data.columns:
-                use_name = feature["name"]
+            if self.annotation_names in feature:
+                use_name = feature[self.annotation_names]
             else:
                 use_name = f"Feature {idx}"
             if "q_value" in annotation_data.columns:

pyopenms_viz/_plotly/core.py

@@ -438,32 +438,52 @@ def _add_annotations(
         ann_xs: list[float],
         ann_ys: list[float],
         ann_colors: list[str],
+        ann_zs: list[float] = None,
     ):
         annotations = []
-        for text, x, y, color in zip(ann_texts, ann_xs, ann_ys, ann_colors):
+        for i, (text, x, y, color) in enumerate(zip(ann_texts, ann_xs, ann_ys, ann_colors)):
             if text is not nan and text != "" and text != "nan":
                 if is_latex_formatted(text):
                     # NOTE: Plotly uses MathJax for LaTeX rendering. Newlines are rendered as \\.
                     text = text.replace("\n", r" \\\ ")
                     text = r'${}$'.format(text)
                 else:
                     text = text.replace("\n", "<br>")
-                annotation = go.layout.Annotation(
-                    text=text,
-                    x=x,
-                    y=y,
-                    showarrow=False,
-                    xanchor="left",
-                    font=dict(
-                        family="Open Sans Mono, monospace",
-                        size=self.annotation_font_size,
-                        color=color,
-                    ),
-                )
-                annotations.append(annotation)
+                if not self.plot_3d:
+                    annotation = go.layout.Annotation(
+                        text=text,
+                        x=x,
+                        y=y,
+                        showarrow=False,
+                        xanchor="left",
+                        font=dict(
+                            family="Open Sans Mono, monospace",
+                            size=self.annotation_font_size,
+                            color=color,
+                        ),
+                    )
+                else:
+                    annotation = go.layout.scene.Annotation(
+                        text=text,
+                        x=x,
+                        y=y,
+                        z=ann_zs[i],
+                        showarrow=False,
+                        xanchor="left",
+                        font=dict(
+                            family="Open Sans Mono, monospace",
+                            size=self.annotation_font_size,
+                            color=color,
+                        ),
+                    )
 
-        for annotation in annotations:
-            fig.add_annotation(annotation)
+                annotations.append(annotation)
+        if not self.plot_3d:
+            for annotation in annotations:
+                fig.add_annotation(annotation)
+        else:
+            for annotation in annotations:
+                fig.layout.scene.annotations += (annotation,)
 
 
 class PLOTLYScatterPlot(PLOTLYPlot, ScatterPlot):
@@ -649,6 +669,13 @@ def create_main_plot(self, x, y, z, class_kwargs, other_kwargs):
                 zlabel=self.zlabel,
                 **other_kwargs,
             )
+            if self.annotation_data is not None:
+                a_x, a_y, a_z, a_t, a_c = self._compute_3D_annotations(
+                    self.annotation_data, x, y, z
+                )
+                vlinePlot._add_annotations(
+                    self.fig, a_t, a_x, a_y, a_c, a_z
+                )
 
             # TODO: Custom tooltips currently not working as expected for 3D plot, it has it's own tooltip that works out of the box, but with set x, y, z name to value
             # tooltips, custom_hover_data = self._create_tooltips({self.xlabel: x, self.ylabel: y, self.zlabel: z})
@@ -778,41 +805,52 @@ def _add_box_boundaries(self, annotation_data, **kwargs):
             colormap=self.feature_config.colormap, n=annotation_data.shape[0]
         )
         for idx, (_, feature) in enumerate(annotation_data.iterrows()):
-            x0 = feature["leftWidth"]
-            x1 = feature["rightWidth"]
-            y0 = feature["IM_leftWidth"]
-            y1 = feature["IM_rightWidth"]
+            x0 = feature[self.annotation_x_lb]
+            x1 = feature[self.annotation_x_ub]
+            y0 = feature[self.annotation_y_lb]
+            y1 = feature[self.annotation_y_ub]
 
-            color = next(color_gen)
+            if self.annotation_colors in feature:
+                color = feature[self.annotation_colors]
+            else:
+                color = next(color_gen)
 
-            if "name" in annotation_data.columns:
-                use_name = feature["name"]
+            if self.annotation_names in annotation_data.columns:
+                use_name = feature[self.annotation_names]
             else:
                 use_name = f"Feature {idx}"
             if "q_value" in annotation_data.columns:
                 legend_label = f"{use_name} (q-value: {feature['q_value']:.4f})"
             else:
                 legend_label = f"{use_name}"
+            # Plot rectangle
+            self.fig.add_shape(
+                type="rect",
+                x0=x0,
+                y0=y0,
+                x1=x1,
+                y1=y1,
+                line=dict(
+                    color=color,
+                    width=self.feature_config.line_width,
+                    dash=bokeh_line_dash_mapper(self.feature_config.line_type, "plotly"),
+                ),
+                fillcolor="rgba(0,0,0,0)",
+                opacity=0.5,
+                layer="above",
+            )
+            # Add a dummy trace for the legend
             self.fig.add_trace(
                 go.Scatter(
-                    x=[
-                        x0,
-                        x1,
-                        x1,
-                        x0,
-                        x0,
-                    ],  # Start and end at the same point to close the shape
-                    y=[y0, y0, y1, y1, y0],
+                    x=[None],
+                    y=[None],
                     mode="lines",
-                    fill="none",
-                    opacity=0.5,
                     line=dict(
                         color=color,
                         width=self.feature_config.line_width,
-                        dash=bokeh_line_dash_mapper(
-                            self.feature_config.line_type, "plotly"
-                        ),
+                        dash=bokeh_line_dash_mapper(self.feature_config.line_type, "plotly"),
                     ),
+                    showlegend=True,
                     name=legend_label,
                 )
             )