From 907ddeca0a052c88746a9b63736827ae0575c2f0 Mon Sep 17 00:00:00 2001
From: Stefan Hahmann <stefan.hahmann@tu-dresden.de>
Date: Thu, 7 Sep 2023 14:50:46 +0200
Subject: [PATCH] Use crop start param also when finding roots

---
 .../ClusterRootNodesController.java           |  9 +++----
 .../mastodon/mamut/util/LineageTreeUtils.java | 25 +++++++++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/mastodon/mamut/clustering/ClusterRootNodesController.java b/src/main/java/org/mastodon/mamut/clustering/ClusterRootNodesController.java
index 4137b3f73..7a7844a93 100644
--- a/src/main/java/org/mastodon/mamut/clustering/ClusterRootNodesController.java
+++ b/src/main/java/org/mastodon/mamut/clustering/ClusterRootNodesController.java
@@ -1,7 +1,7 @@
 package org.mastodon.mamut.clustering;
 
 import org.apache.commons.lang3.tuple.Pair;
-import org.mastodon.graph.algorithm.RootFinder;
+import org.mastodon.collection.RefSet;
 import org.mastodon.graph.algorithm.traversal.DepthFirstIterator;
 import org.mastodon.mamut.clustering.config.ClusteringMethod;
 import org.mastodon.mamut.clustering.config.CropCriteria;
@@ -167,8 +167,7 @@ private List< BranchSpotTree > getRoots()
 	{
 		if ( !synchronizer.isUptodate() )
 			model.getBranchGraph().graphRebuilt();
-		// TODO: take into account crop start - there may be more roots if the crop start is > 0
-		Set< Spot > roots = RootFinder.getRoots( model.getGraph() );
+		RefSet< Spot > roots = LineageTreeUtils.getRoots( model.getGraph(), cropStart );
 		List< BranchSpotTree > trees = new ArrayList<>();
 		for ( Spot root : roots )
 		{
@@ -197,8 +196,8 @@ public void setInputParams( CropCriteria cropCriterion, int cropStart, int cropE
 		if ( cropCriterion.equals( CropCriteria.NUMBER_OF_CELLS ) )
 		{
 			logger.debug( "Crop criterion cells, crop start cells: {}, crop end cells: {}", cropStart, cropEnd );
-			cropStart = LineageTreeUtils.getFirstTimepointWithNSpots( model, cropStart );
-			cropEnd = LineageTreeUtils.getFirstTimepointWithNSpots( model, cropEnd );
+			this.cropStart = LineageTreeUtils.getFirstTimepointWithNSpots( model, cropStart );
+			this.cropEnd = LineageTreeUtils.getFirstTimepointWithNSpots( model, cropEnd );
 		}
 		logger.debug( "Crop criterion {}, start timepoint: {}, crop end timepoint: {}", cropCriterion, cropStart, cropEnd );
 	}
diff --git a/src/main/java/org/mastodon/mamut/util/LineageTreeUtils.java b/src/main/java/org/mastodon/mamut/util/LineageTreeUtils.java
index 3614579c8..dfca7b7a8 100644
--- a/src/main/java/org/mastodon/mamut/util/LineageTreeUtils.java
+++ b/src/main/java/org/mastodon/mamut/util/LineageTreeUtils.java
@@ -1,5 +1,7 @@
 package org.mastodon.mamut.util;
 
+import org.mastodon.collection.RefCollection;
+import org.mastodon.collection.RefCollections;
 import org.mastodon.collection.RefSet;
 import org.mastodon.graph.Edge;
 import org.mastodon.graph.Graph;
@@ -9,6 +11,7 @@
 import org.mastodon.graph.algorithm.traversal.GraphSearch;
 import org.mastodon.graph.algorithm.traversal.SearchListener;
 import org.mastodon.mamut.model.Model;
+import org.mastodon.mamut.model.ModelGraph;
 import org.mastodon.mamut.model.Spot;
 import org.mastodon.pool.PoolCollectionWrapper;
 
@@ -17,6 +20,7 @@
 import java.util.NoSuchElementException;
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
+import java.util.function.Predicate;
 
 public class LineageTreeUtils {
 
@@ -126,4 +130,25 @@ public static int getFirstTimepointWithNSpots( final Model model, final int numb
 				"No time point with at least " + numberOfSpots + " spots in the range [minTimepoint=" + minTimepoint + ", maxTimepoint="
 						+ maxTimepoint + "]." );
 	}
+
+	// Replace with new method after has been resolved https://github.com/mastodon-sc/mastodon-tomancak/issues/13
+	public static RefSet< Spot > getRoots( ModelGraph graph, int timepoint )
+	{
+		Predicate< Spot > isRoot = spot -> spot.getTimepoint() == timepoint
+				|| ( spot.incomingEdges().isEmpty() && spot.getTimepoint() > timepoint );
+		return filterSet( graph.vertices(), isRoot );
+	}
+
+	/**
+	 * Returns a new {@link RefSet} containing all elements of the given
+	 * {@link RefCollection} that satisfy the given {@link Predicate}.
+	 */
+	public static < T > RefSet< T > filterSet( RefCollection< T > values, Predicate< T > predicate )
+	{
+		RefSet< T > filtered = RefCollections.createRefSet( values );
+		for ( T t : values )
+			if ( predicate.test( t ) )
+				filtered.add( t );
+		return filtered;
+	}
 }