added context menu for subcluster export

src/napari_clusters_plotter/_new_plotter_widget.py

@@ -11,7 +11,7 @@
 from qtpy import uic
 from qtpy.QtCore import Qt, Signal
 from qtpy.QtGui import QColor
-from qtpy.QtWidgets import QComboBox, QVBoxLayout, QWidget
+from qtpy.QtWidgets import QComboBox, QMenu, QVBoxLayout, QWidget
 
 from ._algorithm_widget import BaseWidget
 
@@ -64,6 +64,13 @@
             ArtistType.SCATTER
         ]
 
+        # Context menu
+        self.context_menu = QMenu(self.plotting_widget)
+        self.export_clusters = self.context_menu.addAction(
+            "Export selected cluster to new layer"
+        )
+        self.export_clusters.triggered.connect(self._on_export_clusters)
+
         # Add plot and options as widgets
         self.layout.addWidget(self.plotting_widget)
         self.layout.addWidget(self.control_widget)
@@ -85,6 +92,30 @@
         self.control_widget.bins_settings_container.setVisible(False)
         self.control_widget.log_scale_container.setVisible(False)
 
+    def contextMenuEvent(self, event):
+        self.context_menu.exec_(event.globalPos())
+
+    def _on_export_clusters(self):
+        """
+        Export the selected cluster to a new layer.
+        """
+
+        # get currently selected cluster from plotting widget
+        selected_cluster = self.plotting_widget.class_spinbox.value
+        indices = (
+            self.plotting_widget.active_artist.color_indices
+            == selected_cluster
+        )
+
+        # get the layer to export from
+        layer = self.layers[0]
+
+        export_layer = _export_cluster_to_layer(
+            layer, indices, subcluster_index=selected_cluster
+        )
+        if export_layer is not None:
+            self.viewer.add_layer(export_layer)
+
     def _setup_callbacks(self):
         """
         Set up the callbacks for the widget.
@@ -449,33 +480,90 @@
     """
     from napari.utils import DirectLabelColormap
 
-    color_mapping = {
-        napari.layers.Points: lambda _layer, _color: setattr(
-            _layer, "face_color", _color
-        ),
-        napari.layers.Vectors: lambda _layer, _color: setattr(
-            _layer, "edge_color", _color
-        ),
-        napari.layers.Surface: lambda _layer, _color: setattr(
-            _layer, "vertex_colors", _color
-        ),
-        napari.layers.Shapes: lambda _layer, _color: setattr(
-            _layer, "face_color", _color
-        ),
-        napari.layers.Labels: lambda _layer, _color: setattr(
-            _layer,
-            "colormap",
-            DirectLabelColormap(
-                color_dict={
-                    label: _color[label] for label in np.unique(_layer.data)
-                }
-            ),
-        ),
-    }
-
-    if type(layer) in color_mapping:
-        if type(layer) is napari.layers.Labels:
-            # add a color for the background at the first index
-            colors = np.insert(colors, 0, [0, 0, 0, 0], axis=0)
-        color_mapping[type(layer)](layer, colors)
-        layer.refresh()
+    if isinstance(layer, napari.layers.Points):
+        layer.face_color = colors
+
+    elif isinstance(layer, napari.layers.Vectors):
+        layer.edge_color = colors
+
+    elif isinstance(layer, napari.layers.Surface):
+        layer.vertex_colors = colors
+
+    elif isinstance(layer, napari.layers.Shapes):
+        layer.face_color = colors
+
+    elif isinstance(layer, napari.layers.Labels):
+
+        colors = np.insert(colors, 0, [0, 0, 0, 0], axis=0)
+        color_dict = dict(zip(np.unique(layer.data), colors))
+
+        # Insert default colors for labels that are not in the color_dict
+        # Relevant for non-sequential label images
+        if max(color_dict.keys()) > len(colors):
+            for i in range(1, max(color_dict.keys()) - 1):
+                color_dict[i] = [0, 0, 0, 0]
+        # Add a color for the background at the first index
+        layer.colormap = DirectLabelColormap(color_dict=color_dict)
+
+    layer.refresh()
+
+
+def _export_cluster_to_layer(layer, indices, subcluster_index: int = None):
+    """
+    Export the selected cluster to a new layer.
+
+    Parameters
+    ----------
+    layer : napari.layers.Layer
+        The layer to export the cluster from.
+
+    indices : np.ndarray
+        The indices of the cluster to export.
+
+    subcluster_index : str
+        The name of the new layer. If not provided, the name of the layer will be used.
+
+    Returns
+    -------
+    napari.layers.Layer
+        The new layer with the selected cluster.
+    """
+
+    if isinstance(layer, napari.layers.Labels):
+        LUT = np.arange(layer.data.max() + 1)
+        LUT[1:][~indices] = 0
+        new_layer = napari.layers.Labels(LUT[layer.data])
+
+    elif isinstance(layer, napari.layers.Points):
+        new_layer = napari.layers.Points(layer.data[indices])
+        new_layer.size = layer.size[indices]
+
+    elif isinstance(layer, napari.layers.Shapes):
+        new_shapes = [shape for shape, i in zip(layer.data, indices) if i]
+        new_layer = napari.layers.Shapes(new_shapes)
+
+    elif isinstance(layer, napari.layers.Surface):
+        # TODO implement surface export
+        return None
+
+    elif isinstance(layer, napari.layers.Vectors):
+        new_layer = napari.layers.Vectors(layer.data[indices])
+
+    else:
+        return None
+
+    new_layer.scale = layer.scale
+    new_layer.translate = layer.translate
+    new_layer.rotate = layer.rotate
+
+    if not subcluster_index:
+        new_layer.name = f"{layer.name} subcluster"
+    else:
+        new_layer.name = f"{layer.name} subcluster {subcluster_index}"
+
+    # copy features to new layer if available and drop cluster column
+    new_layer.features = layer.features.iloc[indices].copy()
+    if "cluster" in new_layer.features.columns:
+        new_layer.features.drop(columns=["cluster"], inplace=True)
+
+    return new_layer

src/napari_clusters_plotter/_tests/test_plotter.py

@@ -1,9 +1,137 @@
 import numpy as np
+import pandas as pd
+import pytest
+from napari.layers import Labels, Points, Shapes
+
+
+def create_points(n_samples=100, loc=5):
+
+    loc = 5
+    n_timeframes = 5
+    frame = np.arange(n_timeframes).repeat(n_samples // n_timeframes)
+    # make some random points with random features
+    points = np.random.random((n_samples, 4))
+    points2 = np.random.random((n_samples - 1, 4))
+
+    points[:, 0] = frame
+    points2[:, 0] = frame[:-1]
+
+    features = pd.DataFrame(
+        {
+            "frame": frame,
+            "feature1": np.random.normal(size=n_samples, loc=loc),
+            "feature2": np.random.normal(size=n_samples, loc=loc),
+            "feature3": np.random.normal(size=n_samples, loc=loc),
+            "feature4": np.random.normal(size=n_samples, loc=loc),
+        }
+    )
+
+    features2 = pd.DataFrame(
+        {
+            "frame": frame[:-1],
+            "feature2": np.random.normal(size=n_samples - 1, loc=-loc),
+            "feature3": np.random.normal(size=n_samples - 1, loc=-loc),
+            "feature4": np.random.normal(size=n_samples - 1, loc=-loc),
+        }
+    )
+
+    layer1 = Points(
+        points, features=features, size=0.1, blending="translucent_no_depth"
+    )
+    layer2 = Points(
+        points2,
+        features=features2,
+        size=0.1,
+        translate=(0, 0, 2),
+        blending="translucent_no_depth",
+    )
+
+    return layer1, layer2
+
+
+def create_shapes(n_samples=100):
+
+    # create 100 random anchors
+    np.random.seed(0)
+    anchors = np.random.random((n_samples, 2)) * 100
+
+    # create 100 random widths and heights
+    widths = np.random.random(n_samples) * 10
+    heights = np.random.random(n_samples) * 10
+
+    # combine into lists of corner coordinates
+    corner1 = anchors - np.c_[widths, heights] / 2
+    corner2 = anchors + np.c_[widths, -heights] / 2
+    corner3 = anchors + np.c_[widths, heights] / 2
+    corner4 = anchors + np.c_[-widths, heights] / 2
+
+    # create a list of polygons
+    polygons = np.stack([corner1, corner2, corner3, corner4], axis=1)
+
+    layer1 = Shapes(polygons[:49], shape_type="polygon")
+    layer2 = Shapes(polygons[50:], shape_type="polygon")
+    features1 = pd.DataFrame(
+        {
+            "feature1": np.random.normal(size=49),
+            "feature2": np.random.normal(size=49),
+            "feature3": np.random.normal(size=49),
+            "feature4": np.random.normal(size=49),
+        }
+    )
+
+    features2 = pd.DataFrame(
+        {
+            "feature1": np.random.normal(size=50),
+            "feature2": np.random.normal(size=50),
+            "feature3": np.random.normal(size=50),
+            "feature4": np.random.normal(size=50),
+        }
+    )
+
+    layer1.features = features1
+    layer2.features = features2
+
+    return layer1, layer2
+
+
+def create_labels(n_samples=100):
+    from skimage import data, measure
+
+    binary_image1 = data.binary_blobs(length=128, n_dim=3, volume_fraction=0.1)
+    binary_image2 = data.binary_blobs(length=128, n_dim=3, volume_fraction=0.1)
+
+    labels1 = measure.label(binary_image1)
+    labels2 = measure.label(binary_image2)
+
+    n_labels1 = len(np.unique(labels1))
+    n_labels2 = len(np.unique(labels2))
+
+    features1 = pd.DataFrame(
+        {
+            "feature1": np.random.normal(size=n_labels1),
+            "feature2": np.random.normal(size=n_labels1),
+            "feature3": np.random.normal(size=n_labels1),
+            "feature4": np.random.normal(size=n_labels1),
+        }
+    )
+
+    features2 = pd.DataFrame(
+        {
+            "feature1": np.random.normal(size=n_labels2),
+            "feature2": np.random.normal(size=n_labels2),
+            "feature3": np.random.normal(size=n_labels2),
+            "feature4": np.random.normal(size=n_labels2),
+        }
+    )
+
+    layer1 = Labels(labels1, features=features1)
+    layer2 = Labels(labels2, features=features2)
+
+    return layer1, layer2
 
 
 def create_multi_point_layer(n_samples: int = 100):
     import pandas as pd
-    from napari.layers import Points
 
     loc = 5
     n_timeframes = 5
@@ -71,7 +199,36 @@
     viewer.add_image(random_image)
     viewer.add_labels(sample_labels)
 
-    #
+
+@pytest.mark.parametrize("create_data", [create_points, create_shapes])
+def test_cluster_export(make_napari_viewer, create_data):
+    from napari_clusters_plotter import PlotterWidget
+
+    viewer = make_napari_viewer()
+    widget = PlotterWidget(viewer)
+    viewer.window.add_dock_widget(widget, area="right")
+
+    layer1, _ = create_data()
+    viewer.add_layer(layer1)
+
+    # select some random features in the plotting widget
+    n_samples = layer1.features.shape[0]
+    selected_clusters = np.zeros(n_samples, dtype=int)
+    selected_clusters[:5] = 1
+
+    widget.plotting_widget.active_artist.color_indices = selected_clusters
+    widget._on_export_clusters()
+
+    assert len(viewer.layers) == 2
+
+    if isinstance(layer1, Points):
+        assert (
+            viewer.layers[-1].data.shape[0] == n_samples - 5
+        )  # selected cluster is 0, by default.
+        assert np.array_equal(
+            viewer.layers[-1].data,
+            layer1.data[~selected_clusters.astype(bool)],
+        )
 
 
 def test_cluster_memorization(make_napari_viewer, n_samples: int = 100):