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; + } }