Merge pull request #272 from funkelab/95-visualize-by-features-other-than-area-in-the-tree-view

95 visualize-by-features-other-than-area-in-the-tree-view

Author: cmalinmayor

src/motile_tracker/data_views/_tests/test_tree_widget_utils.py

@@ -2,9 +2,11 @@
 
 import napari
 import networkx as nx
+import numpy as np
 import pandas as pd
 import pytest
 from funtracks.data_model import EdgeAttr, NodeAttr, SolutionTracks
+from funtracks.features import Feature
 
 from motile_tracker.data_views.views.tree_view.tree_widget_utils import (
     extract_sorted_tracks,
@@ -127,10 +129,14 @@ def assign_tracklet_ids(graph: nx.DiGraph) -> tuple[nx.DiGraph, list[Any], int]:
 
 def test_track_df(graph_2d):
     tracks = SolutionTracks(graph=graph_2d, ndim=3)
-    del tracks.graph.nodes[2]["area"]
-
-    assert tracks.get_area(1) == 1245
-    assert tracks.get_area(2) is None
+    for node in tracks.graph.nodes():
+        if node != 2:
+            tracks.graph.nodes[node]["custom_attr"] = node * 10
+    tracks.features["custom_attr"] = Feature(
+        feature_type="node",
+        value_type="int",
+        num_values=1,
+    )
 
     tracks.graph, _, _ = assign_tracklet_ids(tracks.graph)
 
@@ -142,6 +148,5 @@ def test_track_df(graph_2d):
 
     track_df, _ = extract_sorted_tracks(tracks, colormap)
     assert isinstance(track_df, pd.DataFrame)
-    assert track_df.loc[track_df["node_id"] == 1, "area"].values[0] == 1245
-    assert track_df.loc[track_df["node_id"] == 2, "area"].values[0] == 0
-    assert track_df["area"].notna().all()
+    assert track_df.loc[track_df["node_id"] == 1, "custom_attr"].values[0] == 10
+    assert np.isnan(track_df.loc[track_df["node_id"] == 2, "custom_attr"].values[0])

src/motile_tracker/data_views/views/tree_view/navigation_widget.py

@@ -21,24 +21,24 @@ def __init__(
         lineage_df: pd.DataFrame,
         view_direction: str,
         selected_nodes: NodeSelectionList,
-        feature: str,
+        plot_type: str,
     ):
         """Widget for controlling navigation in the tree widget
 
         Args:
             track_df (pd.DataFrame): The dataframe holding the track information
-            view_direction (str): The view direction of the tree widget. Options:
-                "vertical", "horizontal".
+            view_direction (str): The view direction of the tree widget. Options: "vertical", "horizontal".
             selected_nodes (NodeSelectionList): The list of selected nodes.
-            feature (str): The feature currently being displayed
+            plot_type (str): either 'tree' or 'feature'
         """
 
         super().__init__()
         self.track_df = track_df
         self.lineage_df = lineage_df
         self.view_direction = view_direction
         self.selected_nodes = selected_nodes
-        self.feature = feature
+        self.plot_type = plot_type
+        self.feature = None
 
         navigation_box = QGroupBox("Navigation [\u2b05 \u27a1 \u2b06 \u2b07]")
         navigation_layout = QHBoxLayout()
@@ -102,8 +102,8 @@ def move(self, direction: str) -> None:
                 next_node = self.get_next_track_node(
                     self.lineage_df, node_id, forward=False
                 )
-                # if not found, look in the whole dataframe to enable jumping to the
-                # next node outside the current tree view content
+                # if not found, look in the whole dataframe
+                # to enable jumping to the next node outside the current tree view content
                 if next_node is None:
                     next_node = self.get_next_track_node(
                         self.track_df, node_id, forward=False
@@ -123,14 +123,13 @@ def get_next_track_node(
         """Get the node at the same time point in an adjacent track.
 
         Args:
-            df (pd.DataFrame): The dataframe to be used (full track_df or subset
-                lineage_df).
+            df (pd.DataFrame): The dataframe to be used (full track_df or subset lineage_df).
             node_id (str): The current node ID to get the next from.
             forward (bool, optional): If true, pick the next track (right/down).
                 Otherwise, pick the previous track (left/up). Defaults to True.
         """
         # Determine which axis to use for finding neighbors
-        axis_label = "area" if self.feature == "area" else "x_axis_pos"
+        axis_label = self.feature if self.plot_type == "feature" else "x_axis_pos"
 
         if df.empty:
             return None
@@ -153,7 +152,6 @@ def get_next_track_node(
             )
             neighbor = neighbors.loc[closest_index_label, "node_id"]
             return neighbor
-        return None
 
     def get_predecessor(self, node_id: str) -> str | None:
         """Get the predecessor node of the given node_id
@@ -171,7 +169,6 @@ def get_predecessor(self, node_id: str) -> str | None:
         parent_row = self.track_df.loc[self.track_df["node_id"] == parent_id]
         if not parent_row.empty:
             return parent_row["node_id"].values[0]
-        return None
 
     def get_successor(self, node_id: str) -> str | None:
         """Get the successor node of the given node_id. If there are two children,
@@ -188,4 +185,3 @@ def get_successor(self, node_id: str) -> str | None:
         if not children.empty:
             child = children.to_dict("records")[0]
             return child["node_id"]
-        return None

src/motile_tracker/data_views/views/tree_view/tree_view_feature_widget.py

@@ -1,6 +1,7 @@
 from psygnal import Signal
 from qtpy.QtWidgets import (
     QButtonGroup,
+    QComboBox,
     QGroupBox,
     QHBoxLayout,
     QRadioButton,
@@ -10,53 +11,97 @@
 
 
 class TreeViewFeatureWidget(QWidget):
-    """Widget to switch between viewing all nodes versus nodes of one or more lineages
-    in the tree widget
-    """
+    """Widget to switch between viewing all nodes versus nodes of one or more lineages in the tree widget"""
 
-    change_feature = Signal(str)
+    change_plot_type = Signal(str)
 
-    def __init__(self):
+    def __init__(self, features: list[str]):
         super().__init__()
 
-        self.feature = "tree"
+        self.plot_type = "tree"
 
-        display_box = QGroupBox("Feature [W]")
+        display_box = QGroupBox("Plot [W]")
         display_layout = QHBoxLayout()
         button_group = QButtonGroup()
         self.show_tree_radio = QRadioButton("Lineage Tree")
         self.show_tree_radio.setChecked(True)
-        self.show_tree_radio.clicked.connect(lambda: self._set_feature("tree"))
-        self.show_area_radio = QRadioButton("Object size")
-        self.show_area_radio.clicked.connect(lambda: self._set_feature("area"))
+        self.show_tree_radio.clicked.connect(lambda: self._set_plot_type("tree"))
+        self.show_area_radio = QRadioButton("Feature")
+        self.show_area_radio.clicked.connect(lambda: self._set_plot_type("feature"))
         button_group.addButton(self.show_tree_radio)
         button_group.addButton(self.show_area_radio)
         display_layout.addWidget(self.show_tree_radio)
         display_layout.addWidget(self.show_area_radio)
+
+        self.feature_dropdown = QComboBox()
+        for feature in features:
+            self.feature_dropdown.addItem(feature)
+        self.feature_dropdown.currentIndexChanged.connect(self._update_feature)
+        if len(features) > 0:
+            self.current_feature = features[0]
+        else:
+            self.current_feature = None
+            self.show_area_radio.setEnabled(False)
+        display_layout.addWidget(self.feature_dropdown)
+
         display_box.setLayout(display_layout)
-        display_box.setMaximumWidth(250)
+        display_box.setMaximumWidth(400)
         display_box.setMaximumHeight(60)
 
         layout = QVBoxLayout()
         layout.addWidget(display_box)
 
         self.setLayout(layout)
 
-    def _toggle_feature_mode(self, event=None) -> None:
+    def _toggle_plot_type(self, event=None) -> None:
         """Toggle display mode"""
 
         if (
             self.show_area_radio.isEnabled
         ):  # if button is disabled, toggle is not allowed
-            if self.feature == "area":
-                self._set_feature("tree")
+            if self.plot_type == "feature":
+                self._set_plot_type("tree")
                 self.show_tree_radio.setChecked(True)
             else:
-                self._set_feature("area")
+                self._set_plot_type("feature")
                 self.show_area_radio.setChecked(True)
 
-    def _set_feature(self, mode: str):
+    def _set_plot_type(self, plot_type: str):
         """Emit signal to change the display mode"""
 
-        self.feature = mode
-        self.change_feature.emit(mode)
+        self.plot_type = plot_type
+        self.change_plot_type.emit(plot_type)
+
+    def _update_feature(self) -> None:
+        """Update the feature to be plotted if the plot_type == 'feature'"""
+
+        self.current_feature = self.feature_dropdown.currentText()
+        self.change_plot_type.emit(self.plot_type)
+
+    def get_current_feature(self):
+        """Return the current feature that is being plotted"""
+
+        return self.current_feature
+
+    def update_feature_dropdown(self, feature_list: list[str]) -> None:
+        """Update the list of features in the dropdown"""
+
+        self.feature_dropdown.currentIndexChanged.disconnect(self._update_feature)
+        self.feature_dropdown.clear()
+        self.show_area_radio.setEnabled(True)
+        for feature in feature_list:
+            self.feature_dropdown.addItem(feature)
+
+        if self.current_feature not in feature_list:
+            if len(feature_list) > 0:
+                self.current_feature = feature_list[0]
+                self.feature_dropdown.setCurrentIndex(0)
+            else:
+                self.current_feature = None
+                self.show_area_radio.setEnabled(False)
+        else:
+            self.feature_dropdown.setCurrentIndex(
+                self.feature_dropdown.findText(self.current_feature)
+            )
+
+        self.feature_dropdown.currentIndexChanged.connect(self._update_feature)