Add support for filtering by selected spots in Spot and Division Count Chart

Introduce a feature to filter spots and divisions by selected spots and persist this setting. Refactor related code to use `ProjectModel` for better integration. Update tests and UI to reflect new functionality.

Author: Stefan Hahmann

src/main/java/org/mastodon/mamut/tomancak/divisioncount/SpotAndDivisionCount.java

@@ -32,8 +32,11 @@
 import java.util.List;
 
 import org.apache.commons.lang3.tuple.Triple;
+import org.mastodon.mamut.ProjectModel;
+import org.mastodon.mamut.model.Link;
 import org.mastodon.mamut.model.Model;
 import org.mastodon.mamut.model.Spot;
+import org.mastodon.model.SelectionModel;
 import org.mastodon.util.TreeUtils;
 
 public class SpotAndDivisionCount
@@ -46,23 +49,33 @@ private SpotAndDivisionCount()
 	/**
 	 * Calculates the number of spots and divisions per time point in the given model.
 	 *
-	 * @param model The model containing the spots and edges.
+	 * @param projectModel The project model containing the spots and edges.
+	 * @param onlySelectedSpots If {@code true}, only counts spots that are selected.
 	 * @return A list of triples, where each triple contains (in that order) the timepoint, the number of spots at that timepoint, and the number of divisions at that timepoint.
 	 */
-	public static List< Triple< Integer, Integer, Integer > > getSpotAndDivisionsPerTimepoint( final Model model )
+	public static List< Triple< Integer, Integer, Integer > > getSpotAndDivisionsPerTimepoint( final ProjectModel projectModel,
+			final boolean onlySelectedSpots )
 	{
+		final Model model = projectModel.getModel();
+		final SelectionModel< Spot, Link > selectionModel = projectModel.getSelectionModel();
 		int minTimepoint = TreeUtils.getMinTimepoint( model );
 		int maxTimepoint = TreeUtils.getMaxTimepoint( model );
 		List< Triple< Integer, Integer, Integer > > timepointAndDivisions = new ArrayList<>();
 		for ( int timepoint = minTimepoint; timepoint <= maxTimepoint; timepoint++ )
 		{
+			int spots = 0;
 			int divisions = 0;
 			for ( Spot spot : model.getSpatioTemporalIndex().getSpatialIndex( timepoint ) )
 			{
+				if ( onlySelectedSpots && !selectionModel.isSelected( spot ) )
+					continue;
 				if ( spot.outgoingEdges().size() > 1 )
 					divisions++;
+				if ( onlySelectedSpots )
+					spots++;
 			}
-			int spots = model.getSpatioTemporalIndex().getSpatialIndex( timepoint ).size();
+			if ( !onlySelectedSpots )
+				spots = model.getSpatioTemporalIndex().getSpatialIndex( timepoint ).size();
 			timepointAndDivisions.add( Triple.of( timepoint, spots, divisions ) );
 		}
 		return timepointAndDivisions;

src/main/java/org/mastodon/mamut/tomancak/divisioncount/SpotAndDivisionCountChart.java

@@ -38,12 +38,14 @@
 import java.awt.geom.Rectangle2D;
 import java.util.List;
 
+import javax.swing.ButtonGroup;
 import javax.swing.JButton;
 import javax.swing.JCheckBox;
 import javax.swing.JColorChooser;
 import javax.swing.JFrame;
 import javax.swing.JLabel;
 import javax.swing.JPanel;
+import javax.swing.JRadioButton;
 import javax.swing.JSpinner;
 import javax.swing.SpinnerNumberModel;
 import javax.swing.WindowConstants;
@@ -83,6 +85,8 @@ public class SpotAndDivisionCountChart extends JFrame
 
 	private static final String DIVISION_COUNT_SLIDING_AVERAGE_WINDOW_SIZE = "divisionCountSlidingAverageWindowSize";
 
+	private static final String ONLY_SELECTED_SPOTS = "onlySelectedSpots";
+
 	private static final int SPOT_COUNT_DEFAULT_COLOR = new Color( 230, 159, 0 ).getRGB(); // Dark Orange
 
 	private static final int DIVISION_COUNT_DEFAULT_COLOR = new Color( 86, 180, 233 ).getRGB(); // Light Blue
@@ -93,9 +97,11 @@ public class SpotAndDivisionCountChart extends JFrame
 
 	private Color divisionCountColor; // Light Blue
 
-	private final static String TITLE = "Spot and Division Counts over Time";
-	private final static String SPOTS_COUNT_SERIES_NAME = "Spot Counts";
-	private final static String DIVISION_COUNT_SERIES_NAME = "Division Counts";
+	private static final String TITLE = "Spot and Division Counts over Time";
+
+	private static final String SPOTS_COUNT_SERIES_NAME = "Spot Counts";
+
+	private static final String DIVISION_COUNT_SERIES_NAME = "Division Counts";
 
 	private final XYPlot plot;
 
@@ -107,11 +113,19 @@ public class SpotAndDivisionCountChart extends JFrame
 
 	private int divisionWindowSize;
 
+	private boolean onlySelectedSpots;
+
 	private final PrefService prefs;
 
-	SpotAndDivisionCountChart( final double[] timepoints, final double[] spotCounts, final double[] divisionCounts,
-			final PrefService prefs )
+	private final ProjectModel projectModel;
+
+	private XYSeriesCollection spotCountsSeries;
+
+	private XYSeriesCollection divisionCountsSeries;
+
+	SpotAndDivisionCountChart( final ProjectModel projectModel, final PrefService prefs )
 	{
+		this.projectModel = projectModel;
 		this.prefs = prefs;
 
 		this.spotCountColor = new Color(
@@ -122,12 +136,9 @@ public class SpotAndDivisionCountChart extends JFrame
 				prefs.getInt( SpotAndDivisionCountChart.class, SPOT_COUNT_SLIDING_AVERAGE_WINDOW_SIZE, DEFAULT_SLIDING_WINDOW_SIZE );
 		this.divisionWindowSize =
 				prefs.getInt( SpotAndDivisionCountChart.class, DIVISION_COUNT_SLIDING_AVERAGE_WINDOW_SIZE, DEFAULT_SLIDING_WINDOW_SIZE );
+		this.onlySelectedSpots = prefs.getBoolean( SpotAndDivisionCountChart.class, ONLY_SELECTED_SPOTS, false );
 
-		XYSeriesCollection spotCountsSeries = createSeries(
-				timepoints, spotCounts, SPOTS_COUNT_SERIES_NAME, spotWindowSize );
-
-		XYSeriesCollection divisionCountsSeries = createSeries(
-				timepoints, divisionCounts, DIVISION_COUNT_SERIES_NAME, divisionWindowSize );
+		updateChartData();
 
 		JFreeChart chart = ChartFactory.createXYLineChart(
 				TITLE,
@@ -168,7 +179,7 @@ public class SpotAndDivisionCountChart extends JFrame
 		chartPanel.setPreferredSize( new Dimension( 800, 600 ) );
 
 		// Add color chooser and visibility controls
-		JPanel controlPanel = createControlPanel( timepoints, spotCounts, divisionCounts );
+		JPanel controlPanel = createControlPanel();
 
 		// Set up the frame layout
 		setLayout( new BorderLayout() );
@@ -184,23 +195,34 @@ public class SpotAndDivisionCountChart extends JFrame
 		repaint();
 	}
 
-	public static void show( final ProjectModel projectModel, final PrefService prefService )
+	private void updateChartData()
 	{
-		List< Triple< Integer, Integer, Integer > > divisionCounts =
-				SpotAndDivisionCount.getSpotAndDivisionsPerTimepoint( projectModel.getModel() );
-		double[] timepoints = divisionCounts.stream().mapToDouble( Triple::getLeft ).toArray();
-		double[] spots = divisionCounts.stream().mapToDouble( Triple::getMiddle ).toArray();
-		double[] divisions = divisionCounts.stream().mapToDouble( Triple::getRight ).toArray();
+		List< Triple< Integer, Integer, Integer > > counts =
+				SpotAndDivisionCount.getSpotAndDivisionsPerTimepoint( projectModel, this.onlySelectedSpots );
+		double[] timepoints = counts.stream().mapToDouble( Triple::getLeft ).toArray();
+		double[] spotCounts = counts.stream().mapToDouble( Triple::getMiddle ).toArray();
+		double[] divisionCounts = counts.stream().mapToDouble( Triple::getRight ).toArray();
+		spotCountsSeries = createSeries( timepoints, spotCounts, SPOTS_COUNT_SERIES_NAME, spotWindowSize );
+		divisionCountsSeries = createSeries( timepoints, divisionCounts, DIVISION_COUNT_SERIES_NAME, divisionWindowSize );
+
+		if ( plot != null )
+		{
+			plot.setDataset( 0, spotCountsSeries );
+			plot.setDataset( 1, divisionCountsSeries );
+		}
+	}
 
-		SpotAndDivisionCountChart chart = new SpotAndDivisionCountChart( timepoints, spots, divisions, prefService );
+	public static void show( final ProjectModel projectModel, final PrefService prefService )
+	{
+		SpotAndDivisionCountChart chart = new SpotAndDivisionCountChart( projectModel, prefService );
 		chart.setDefaultCloseOperation( WindowConstants.DISPOSE_ON_CLOSE );
 		chart.setVisible( true );
 	}
 
 	/**
 	 * Creates a control panel with color choosers and checkboxes for visibility controls.
 	 */
-	private JPanel createControlPanel( final double[] timepoints, final double[] spotCounts, final double[] divisionCounts )
+	private JPanel createControlPanel()
 	{
 		JPanel controlPanel = new JPanel( new MigLayout( "fill, wrap 5", "[grow]", "[]10[]10[]10[]10[]" ) );
 
@@ -237,7 +259,7 @@ private JPanel createControlPanel( final double[] timepoints, final double[] spo
 		spotWindowSpinner.addChangeListener( e -> {
 			spotWindowSize = ( int ) spotWindowSpinner.getValue();
 			prefs.put( SpotAndDivisionCountChart.class, SPOT_COUNT_SLIDING_AVERAGE_WINDOW_SIZE, spotWindowSize );
-			updateSlidingAverage( timepoints, spotCounts, divisionCounts );
+			updateChartData();
 		} );
 
 		// Division-related controls
@@ -273,7 +295,25 @@ private JPanel createControlPanel( final double[] timepoints, final double[] spo
 		divisionWindowSpinner.addChangeListener( e -> {
 			divisionWindowSize = ( int ) divisionWindowSpinner.getValue();
 			prefs.put( SpotAndDivisionCountChart.class, DIVISION_COUNT_SLIDING_AVERAGE_WINDOW_SIZE, divisionWindowSize );
-			updateSlidingAverage( timepoints, spotCounts, divisionCounts );
+			updateChartData();
+		} );
+
+		JRadioButton allSpots = new JRadioButton( "All spots", !onlySelectedSpots );
+		JRadioButton selectedSpots = new JRadioButton( "Only selected spots", onlySelectedSpots );
+		ButtonGroup group = new ButtonGroup();
+		group.add( allSpots );
+		group.add( selectedSpots );
+		allSpots.addActionListener( e -> {
+			onlySelectedSpots = false;
+			prefs.put( SpotAndDivisionCountChart.class, ONLY_SELECTED_SPOTS, onlySelectedSpots );
+			updateChartData();
+			repaint();
+		} );
+		selectedSpots.addActionListener( e -> {
+			this.onlySelectedSpots = true;
+			prefs.put( SpotAndDivisionCountChart.class, ONLY_SELECTED_SPOTS, onlySelectedSpots );
+			updateChartData();
+			repaint();
 		} );
 
 		// Add components to the control panel
@@ -289,6 +329,9 @@ private JPanel createControlPanel( final double[] timepoints, final double[] spo
 		controlPanel.add( new JLabel( "Window Size:" ), "align right" );
 		controlPanel.add( divisionWindowSpinner, "wmax 50" );
 
+		controlPanel.add( allSpots, "growx" );
+		controlPanel.add( selectedSpots, "span, growx" );
+
 		// Add description
 		controlPanel.add( new JLabel(
 				"<html>This windows shows the number of spots and divisions at each timepoint together with a sliding average.<br>A division is defined as a spot with more than one outgoing edge.</html>" ),
@@ -331,27 +374,23 @@ private void updateChartColors()
 	 */
 	private void updateAxisVisibility()
 	{
-		XYItemRenderer spotCountRenderer = plot.getRenderer( 0 );
-		Boolean spotCountVisible = spotCountRenderer.getSeriesVisible( 0 );
-		Boolean spotAverageVisible = spotCountRenderer.getSeriesVisible( 1 );
-		boolean isSpotCountVisible = spotCountVisible != null && spotCountVisible;
-		boolean isSpotAverageVisible = spotAverageVisible != null && spotAverageVisible;
-		boolean showSpotAxis = isSpotCountVisible || isSpotAverageVisible;
-		leftAxis.setLabel( showSpotAxis ? SPOTS_COUNT_SERIES_NAME : null );
-		leftAxis.setVisible( showSpotAxis );
-
-		XYItemRenderer divisionCountRenderer = plot.getRenderer( 1 );
-		Boolean divisionCountVisible = divisionCountRenderer.getSeriesVisible( 0 );
-		Boolean divisionAverageVisible = divisionCountRenderer.getSeriesVisible( 1 );
-		boolean isDivisionCountVisible = divisionCountVisible != null && divisionCountVisible;
-		boolean isDivisionAverageVisible = divisionAverageVisible != null && divisionAverageVisible;
-		boolean showDivisionAxis = isDivisionCountVisible || isDivisionAverageVisible;
-		rightAxis.setLabel( showDivisionAxis ? DIVISION_COUNT_SERIES_NAME : null );
-		rightAxis.setVisible( showDivisionAxis );
-
+		updateAxisVisibility( plot.getRenderer( 0 ), leftAxis, SPOTS_COUNT_SERIES_NAME );
+		updateAxisVisibility( plot.getRenderer( 1 ), rightAxis, DIVISION_COUNT_SERIES_NAME );
 		repaint();
 	}
 
+	private void updateAxisVisibility( final XYItemRenderer renderer, final NumberAxis axis,
+			final String labelIfVisible )
+	{
+		Boolean countVisible = renderer.getSeriesVisible( 0 );
+		Boolean averageVisible = renderer.getSeriesVisible( 1 );
+		boolean isCountVisible = countVisible != null && countVisible;
+		boolean isAverageVisible = averageVisible != null && averageVisible;
+		boolean showAxis = isCountVisible || isAverageVisible;
+		axis.setLabel( showAxis ? labelIfVisible : null );
+		axis.setVisible( showAxis );
+	}
+
 	private XYLineAndShapeRenderer createRenderer( final Color color, final Shape shape )
 	{
 		XYLineAndShapeRenderer renderer = new XYLineAndShapeRenderer();
@@ -406,20 +445,6 @@ private static XYSeriesCollection createSeries( final double[] xValues, final do
 		return dataset;
 	}
 
-	/**
-	 * Updates the sliding average series based on the new window sizes.
-	 */
-	private void updateSlidingAverage( double[] timepoints, double[] spotCounts, double[] divisionCounts )
-	{
-		XYSeriesCollection spotCountsSeries = createSeries(
-				timepoints, spotCounts, SPOTS_COUNT_SERIES_NAME, spotWindowSize );
-		XYSeriesCollection divisionCountsSeries = createSeries(
-				timepoints, divisionCounts, DIVISION_COUNT_SERIES_NAME, divisionWindowSize );
-
-		plot.setDataset( 0, spotCountsSeries );
-		plot.setDataset( 1, divisionCountsSeries );
-	}
-
 	private static double[] calculateSlidingAverage( double[] values, int windowSize )
 	{
 		double[] result = new double[ values.length ];

src/main/java/org/mastodon/mamut/tomancak/export/ExportDivisionCountsPerTimepointCommand.java

@@ -37,7 +37,6 @@
 
 import org.apache.commons.lang3.tuple.Triple;
 import org.mastodon.mamut.ProjectModel;
-import org.mastodon.mamut.model.Model;
 import org.mastodon.mamut.tomancak.divisioncount.SpotAndDivisionCount;
 import org.scijava.Context;
 import org.scijava.ItemVisibility;
@@ -75,7 +74,7 @@ public void run()
 	{
 		try
 		{
-			writeDivisionCountsToFile( projectModel.getModel(), saveTo, context.service( StatusService.class ) );
+			writeDivisionCountsToFile( projectModel, saveTo, context.service( StatusService.class ) );
 		}
 		catch ( IOException e )
 		{
@@ -93,7 +92,8 @@ public void run()
 	 *     <li>The first line is the header.</li>
 	 * </ul>
 	 */
-	public static void writeDivisionCountsToFile( final Model model, final File file, final StatusService statusService ) throws IOException
+	public static void writeDivisionCountsToFile( final ProjectModel projectModel, final File file, final StatusService statusService )
+			throws IOException
 	{
 		if ( file == null )
 			throw new IllegalArgumentException( "Cannot write division counts to file. Given file is null." );
@@ -102,7 +102,7 @@ public static void writeDivisionCountsToFile( final Model model, final File file
 		{
 			csvWriter.writeNext( new String[] { "timepoint", "divisions" } );
 			List< Triple< Integer, Integer, Integer > > timepointAndDivisions =
-					SpotAndDivisionCount.getSpotAndDivisionsPerTimepoint( model );
+					SpotAndDivisionCount.getSpotAndDivisionsPerTimepoint( projectModel, false );
 			for ( Triple< Integer, Integer, Integer > pair : timepointAndDivisions )
 			{
 				csvWriter.writeNext( new String[] { String.valueOf( pair.getLeft() ), String.valueOf( pair.getRight() ) }, false );

src/test/java/org/mastodon/mamut/tomancak/export/ExportDivisionCountsPerTimepointCommandTest.java

@@ -34,8 +34,14 @@
 import java.io.IOException;
 import java.nio.charset.Charset;
 
+import net.imglib2.img.Img;
+import net.imglib2.img.array.ArrayImgs;
+import net.imglib2.type.numeric.real.FloatType;
+
 import org.apache.commons.io.FileUtils;
 import org.junit.Test;
+import org.mastodon.mamut.ProjectModel;
+import org.mastodon.mamut.ProjectModelTestUtils;
 import org.mastodon.mamut.feature.branch.exampleGraph.ExampleGraph2;
 import org.mastodon.mamut.model.Model;
 import org.scijava.Context;
@@ -52,12 +58,16 @@ public void testWriteDivisionCountsToFile() throws IOException
 	{
 		try (Context context = new Context())
 		{
+			final Img< FloatType > img = ArrayImgs.floats( 1, 1, 1 );
 			Model model = new ExampleGraph2().getModel();
+			File mastodonFile = File.createTempFile( "test", ".mastodon" );
+			ProjectModel appModel = ProjectModelTestUtils.wrapAsAppModel( img, model, context, mastodonFile );
+
 			File outputFile = File.createTempFile( "divisioncounts", ".csv" );
 			outputFile.deleteOnExit();
 			StatusService service = context.service( StatusService.class );
 
-			ExportDivisionCountsPerTimepointCommand.writeDivisionCountsToFile( model, outputFile, service );
+			ExportDivisionCountsPerTimepointCommand.writeDivisionCountsToFile( appModel, outputFile, service );
 
 			String content = FileUtils.readFileToString( outputFile, Charset.defaultCharset() );
 			String expected = "\"timepoint\",\"divisions\"\n"