Lineage Registration: dynamic spacial registration

src/main/java/org/mastodon/mamut/tomancak/TomancakPlugins.java

@@ -6,13 +6,13 @@
  * %%
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
- * 
+ *
  * 1. Redistributions of source code must retain the above copyright notice,
  *    this list of conditions and the following disclaimer.
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  *    this list of conditions and the following disclaimer in the documentation
  *    and/or other materials provided with the distribution.
- * 
+ *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
@@ -53,6 +53,7 @@
 import org.mastodon.mamut.tomancak.sort_tree.FlipDescendants;
 import org.mastodon.mamut.tomancak.sort_tree.SortTreeLeftRightDialog;
 import org.mastodon.mamut.tomancak.sort_tree.SortTreeExternInternDialog;
+import org.mastodon.mamut.tomancak.spots.AddCenterSpots;
 import org.mastodon.mamut.tomancak.spots.FilterOutSolists;
 import org.mastodon.mamut.tomancak.spots.InterpolateMissingSpots;
 import org.mastodon.ui.keymap.CommandDescriptionProvider;
@@ -77,13 +78,13 @@ public class TomancakPlugins extends AbstractContextual implements MamutPlugin
 	private static final String COMPACT_LINEAGE_VIEW = "[tomancak] show compact lineage";
 	private static final String SORT_TREE = "[tomancak] sort lineage tree";
 	private static final String SORT_TREE_EXTERN_INTERN = "[tomancak] sort lineage tree extern intern";
-
 	private static final String LABEL_SPOTS_SYSTEMATICALLY = "[tomancak] label spots systematically";
 	private static final String REMOVE_SOLISTS_SPOTS = "[tomancak] remove solists spots";
 	private static final String EXPORTS_LINEAGE_LENGTHS = "[tomancak] export lineage lengths";
 	private static final String EXPORTS_SPOTS_COUNTS = "[tomancak] export spots counts";
 	private static final String MERGE_PROJECTS = "[tomancak] merge projects";
 	private static final String TWEAK_DATASET_PATH = "[tomancak] fix project image path";
+	private static final String ADD_CENTER_SPOTS = "[tomancak] add center spot";
 
 	private static final String[] EXPORT_PHYLOXML_KEYS = { "not mapped" };
 	private static final String[] FLIP_DESCENDANTS_KEYS = { "ctrl E" };
@@ -94,13 +95,13 @@ public class TomancakPlugins extends AbstractContextual implements MamutPlugin
 	private static final String[] COMPACT_LINEAGE_VIEW_KEYS = { "not mapped" };
 	private static final String[] SORT_TREE_KEYS = { "ctrl S" };
 	private static final String[] SORT_TREE_EXTERN_INTERN_KEYS = { "not mapped" };
-
 	private static final String[] LABEL_SPOTS_SYSTEMATICALLY_KEYS = { "not mapped" };
 	private static final String[] REMOVE_SOLISTS_SPOTS_KEYS = { "not mapped" };
 	private static final String[] EXPORTS_LINEAGE_LENGTHS_KEYS = { "not mapped" };
 	private static final String[] EXPORTS_SPOTS_COUNTS_KEYS = { "not mapped" };
 	private static final String[] MERGE_PROJECTS_KEYS = { "not mapped" };
 	private static final String[] TWEAK_DATASET_PATH_KEYS = { "not mapped" };
+	private static final String[] ADD_CENTER_SPOTS_KEYS = { "not mapped" };
 
 	private static Map< String, String > menuTexts = new HashMap<>();
 
@@ -121,6 +122,7 @@ public class TomancakPlugins extends AbstractContextual implements MamutPlugin
 		menuTexts.put( EXPORTS_SPOTS_COUNTS, "Export Spots Counts" );
 		menuTexts.put( MERGE_PROJECTS, "Merge Two Projects" );
 		menuTexts.put( TWEAK_DATASET_PATH, "Fix Image Path" );
+		menuTexts.put( ADD_CENTER_SPOTS, "Add center spot" );
 	}
 
 	/*
@@ -152,6 +154,8 @@ public void getCommandDescriptions( final CommandDescriptions descriptions )
 			descriptions.add( EXPORTS_SPOTS_COUNTS, EXPORTS_SPOTS_COUNTS_KEYS, "Exports counts of spots into CSV-like files to be imported in data processors." );
 			descriptions.add( MERGE_PROJECTS, MERGE_PROJECTS_KEYS, "Merge two Mastodon projects into one." );
 			descriptions.add( TWEAK_DATASET_PATH, TWEAK_DATASET_PATH_KEYS, "Allows to insert new path to the BDV data and whether it is relative or absolute." );
+			descriptions.add( ADD_CENTER_SPOTS, ADD_CENTER_SPOTS_KEYS,
+					"On each timepoint with selected spots, add a new spot that is in the center (average position)." );
 		}
 	}
 
@@ -185,6 +189,8 @@ public void getCommandDescriptions( final CommandDescriptions descriptions )
 
 	private final AbstractNamedAction tweakDatasetPathAction;
 
+	private final AbstractNamedAction addCenterSpots;
+
 	private MamutPluginAppModel pluginAppModel;
 
 	public TomancakPlugins()
@@ -204,6 +210,7 @@ public TomancakPlugins()
 		exportSpotsCountsAction = new RunnableAction( EXPORTS_SPOTS_COUNTS, this::exportCounts );
 		mergeProjectsAction = new RunnableAction( MERGE_PROJECTS, this::mergeProjects );
 		tweakDatasetPathAction = new RunnableAction( TWEAK_DATASET_PATH, this::tweakDatasetPath );
+		addCenterSpots = new RunnableAction( ADD_CENTER_SPOTS, this::addCenterSpots );
 		updateEnabledActions();
 	}
 
@@ -226,6 +233,7 @@ public List< ViewMenuBuilder.MenuItem > getMenuItems()
 								item( LABEL_SELECTED_SPOTS ),
 								item( CHANGE_BRANCH_LABELS ),
 								item( REMOVE_SOLISTS_SPOTS ),
+								item( ADD_CENTER_SPOTS ),
 								item( INTERPOLATE_SPOTS ),
 								item( FLIP_DESCENDANTS ),
 								item( SORT_TREE ),
@@ -264,6 +272,7 @@ public void installGlobalActions( final Actions actions )
 		actions.namedAction( exportSpotsCountsAction, EXPORTS_SPOTS_COUNTS_KEYS );
 		actions.namedAction( mergeProjectsAction, MERGE_PROJECTS_KEYS );
 		actions.namedAction( tweakDatasetPathAction, TWEAK_DATASET_PATH_KEYS );
+		actions.namedAction( addCenterSpots, ADD_CENTER_SPOTS_KEYS );
 	}
 
 	private void updateEnabledActions()
@@ -283,6 +292,7 @@ private void updateEnabledActions()
 		exportSpotsCountsAction.setEnabled( appModel != null );
 		mergeProjectsAction.setEnabled( appModel != null );
 		tweakDatasetPathAction.setEnabled( appModel != null );
+		addCenterSpots.setEnabled( appModel != null );
 	}
 
 	private void exportPhyloXml()
@@ -413,4 +423,12 @@ private void labelSpotsSystematically()
 			LabelSpotsSystematicallyDialog.showDialog( pluginAppModel.getAppModel() );
 		}
 	}
+
+	private void addCenterSpots()
+	{
+		if ( pluginAppModel != null )
+		{
+			AddCenterSpots.addSpots( pluginAppModel.getAppModel() );
+		}
+	}
 }

src/main/java/org/mastodon/mamut/tomancak/lineage_registration/LineageColoring.java

@@ -8,9 +8,7 @@
 
 import org.mastodon.collection.RefCollection;
 import org.mastodon.collection.RefRefMap;
-import org.mastodon.graph.algorithm.traversal.DepthFirstSearch;
-import org.mastodon.graph.algorithm.traversal.GraphSearch;
-import org.mastodon.graph.algorithm.traversal.SearchListener;
+import org.mastodon.graph.algorithm.traversal.DepthFirstIterator;
 import org.mastodon.mamut.model.Link;
 import org.mastodon.mamut.model.Model;
 import org.mastodon.mamut.model.Spot;
@@ -27,9 +25,11 @@ public class LineageColoring
 	 * The tags of matching roots node gets the same colors assigned.
 	 * Finally, the tag is applied to all the descendants of a root node.
 	 */
-	public static void tagLineages( Model embryoA, Model embryoB )
+	public static void tagLineages( Model embryoA, int firstTimepointA, Model embryoB, int firstTimepointB )
 	{
-		RefRefMap< Spot, Spot > roots = RootsPairing.pairDividingRoots( embryoA.getGraph(), embryoB.getGraph() );
+		RefRefMap< Spot, Spot > roots = RootsPairing.pairDividingRoots(
+				embryoA.getGraph(), firstTimepointA,
+				embryoB.getGraph(), firstTimepointB );
 		List< String > labels = roots.keySet().stream().map( Spot::getLabel ).sorted().collect( Collectors.toList() );
 		Map< String, Integer > colorMap = createColorMap( labels );
 		tagLineages( colorMap, roots.keySet(), embryoA );
@@ -48,38 +48,14 @@ private static void tagLineage( Model model, Spot root, TagSetStructure.TagSet t
 	{
 		ObjTagMap< Spot, TagSetStructure.Tag > spotTags = model.getTagSetModel().getVertexTags().tags( tagSet );
 		ObjTagMap< Link, TagSetStructure.Tag > edgeTags = model.getTagSetModel().getEdgeTags().tags( tagSet );
-
-		SearchListener< Spot, Link, DepthFirstSearch< Spot, Link > > searchListener =
-				new SearchListener< Spot, Link, DepthFirstSearch< Spot, Link > >()
-				{
-					@Override
-					public void processVertexLate( Spot spot, DepthFirstSearch< Spot, Link > search )
-					{
-						// do nothing
-					}
-
-					@Override
-					public void processVertexEarly( Spot spot, DepthFirstSearch< Spot, Link > spotLinkDepthFirstSearch )
-					{
-						spotTags.set( spot, tag );
-					}
-
-					@Override
-					public void processEdge( Link link, Spot spot, Spot v1, DepthFirstSearch< Spot, Link > spotLinkDepthFirstSearch )
-					{
-						edgeTags.set( link, tag );
-					}
-
-					@Override
-					public void crossComponent( Spot spot, Spot v1, DepthFirstSearch< Spot, Link > spotLinkDepthFirstSearch )
-					{
-						// do nothing
-					}
-				};
-
-		DepthFirstSearch< Spot, Link > search = new DepthFirstSearch<>( model.getGraph(), GraphSearch.SearchDirection.DIRECTED );
-		search.setTraversalListener( searchListener );
-		search.start( root );
+		DepthFirstIterator< Spot, Link > iterator = new DepthFirstIterator<>( root, model.getGraph() );
+		while ( iterator.hasNext() )
+		{
+			Spot spot = iterator.next();
+			spotTags.set( spot, tag );
+			for ( Link edge : spot.outgoingEdges() )
+				edgeTags.set( edge, tag );
+		}
 	}
 
 	/**

src/main/java/org/mastodon/mamut/tomancak/lineage_registration/LineageRegistrationAlgorithm.java

@@ -4,8 +4,13 @@
 
 import org.mastodon.collection.RefRefMap;
 import org.mastodon.collection.ref.RefRefHashMap;
+import org.mastodon.mamut.model.Model;
 import org.mastodon.mamut.model.ModelGraph;
 import org.mastodon.mamut.model.Spot;
+import org.mastodon.mamut.tomancak.lineage_registration.spatial_registration.NotEnoughPairedRootsException;
+import org.mastodon.mamut.tomancak.lineage_registration.spatial_registration.SpatialRegistration;
+import org.mastodon.mamut.tomancak.lineage_registration.spatial_registration.SpatialRegistrationFactory;
+import org.mastodon.mamut.tomancak.lineage_registration.spatial_registration.SpatialRegistrationMethod;
 import org.mastodon.mamut.tomancak.sort_tree.SortTreeUtils;
 
 /**
@@ -17,7 +22,9 @@
  */
 public class LineageRegistrationAlgorithm
 {
-	private final AffineTransform3D transformAB;
+	private static final int TIME_OFFSET = SortTreeUtils.DIVISION_DIRECTION_TIME_OFFSET;
+
+	private final SpatialRegistration spatialRegistration;
 
 	private final ModelGraph graphA;
 
@@ -28,25 +35,63 @@ public class LineageRegistrationAlgorithm
 	 */
 	private final RefRefMap< Spot, Spot > mapAB;
 
-	public static RegisteredGraphs run( ModelGraph graphA, ModelGraph graphB,
-			RefRefMap< Spot, Spot > roots, AffineTransform3D transformAB )
+	/**
+	 * Runs the lineage registration algorithm for to given graphs. The spots before
+	 * the given timepoints are ignored.
+	 *
+	 * @return a {@link RegisteredGraphs} object that contains the two graphs and the
+	 * 	   mapping between the first spots of the branches in the two graphs.
+	 */
+	public static RegisteredGraphs run(
+			Model modelA,
+			int firstTimepointA,
+			Model modelB,
+			int firstTimepointB,
+			SpatialRegistrationMethod spatialRegistrationMethod )
+	{
+		try
+		{
+			RefRefMap< Spot, Spot > roots =
+					RootsPairing.pairDividingRoots( modelA.getGraph(), firstTimepointA, modelB.getGraph(), firstTimepointB );
+			SpatialRegistrationFactory algorithm = SpatialRegistrationMethod.getFactory( spatialRegistrationMethod );
+			SpatialRegistration spatialRegistration = algorithm.run( modelA, modelB, roots );
+			return run( modelA.getGraph(), modelB.getGraph(), roots, spatialRegistration );
+		}
+		catch ( NotEnoughPairedRootsException e )
+		{
+			throw newDetailedNotEnoughPairedRootsException( e, modelA, firstTimepointA, modelB, firstTimepointB );
+		}
+	}
+
+	private static NotEnoughPairedRootsException newDetailedNotEnoughPairedRootsException( NotEnoughPairedRootsException e, Model modelA, int firstTimepointA, Model modelB, int firstTimepointB )
 	{
-		RefRefMap< Spot, Spot > mapping = calculateMapping( graphA, graphB, roots, transformAB );
-		return new RegisteredGraphs( graphA, graphB, transformAB, mapping );
+		String message = e.getMessage() + "\n"
+				+ "\n"
+				+ RootsPairing.report( modelA.getGraph(), firstTimepointA, modelB.getGraph(), firstTimepointB )
+				+ "\n"
+				+ "Please make sure to:\n"
+				+ "  - Select timepoints at which both embryos are at a similar stage\n"
+				+ "    and have at least 3 cells that divide.\n"
+				+ "  - Name those cells by setting the label of the first spot of the cell.\n"
+				+ "    The cell names need to match between the two datasets.\n"
+				+ "  - If there are less than three dividing cells, consider using the "
+				+ "    dynamic spatial registration that is based on landmarks.\n";
+		return new NotEnoughPairedRootsException( message );
 	}
 
-	private static RefRefMap< Spot, Spot > calculateMapping( ModelGraph graphA, ModelGraph graphB,
-			RefRefMap< Spot, Spot > roots, AffineTransform3D transformAB )
+	public static RegisteredGraphs run( ModelGraph graphA, ModelGraph graphB,
+			RefRefMap< Spot, Spot > roots, SpatialRegistration spatialRegistration )
 	{
-		return new LineageRegistrationAlgorithm(
+		RefRefMap< Spot, Spot > mapping = new LineageRegistrationAlgorithm(
 				graphA, graphB,
-				roots, transformAB ).getMapping();
+				roots, spatialRegistration ).getMapping();
+		return new RegisteredGraphs( graphA, graphB, spatialRegistration, mapping );
 	}
 
 	private LineageRegistrationAlgorithm( ModelGraph graphA, ModelGraph graphB, RefRefMap< Spot, Spot > roots,
-			AffineTransform3D transformAB )
+			SpatialRegistration spatialRegistration )
 	{
-		this.transformAB = noOffsetTransform( transformAB );
+		this.spatialRegistration = spatialRegistration;
 		this.graphA = graphA;
 		this.graphB = graphB;
 		this.mapAB = new RefRefHashMap<>( graphA.vertices().getRefPool(), graphB.vertices().getRefPool() );
@@ -65,14 +110,6 @@ private LineageRegistrationAlgorithm( ModelGraph graphA, ModelGraph graphB, RefR
 		}
 	}
 
-	public static RegisteredGraphs run( ModelGraph graphA, ModelGraph graphB )
-	{
-		RefRefMap< Spot, Spot > roots = RootsPairing.pairDividingRoots( graphA, graphB );
-		AffineTransform3D transformAB = EstimateTransformation.estimateScaleRotationAndTranslation( roots );
-		RegisteredGraphs result = run( graphA, graphB, roots, transformAB );
-		return result;
-	}
-
 	private void matchTree( Spot rootA, Spot rootB )
 	{
 		mapAB.put( rootA, rootB );
@@ -88,6 +125,8 @@ private void matchTree( Spot rootA, Spot rootB )
 				return;
 			double[] directionA = SortTreeUtils.directionOfCellDevision( graphA, dividingA );
 			double[] directionB = SortTreeUtils.directionOfCellDevision( graphB, dividingB );
+			AffineTransform3D transformAB = noOffsetTransform( spatialRegistration.getTransformationAtoB(
+					dividingA.getTimepoint() + TIME_OFFSET, dividingB.getTimepoint() + TIME_OFFSET ) );
 			transformAB.apply( directionA, directionA );
 			boolean flip = SortTreeUtils.scalarProduct( directionA, directionB ) < 0;
 			matchChildTree( dividingA, dividingB, 0, flip ? 1 : 0 );
@@ -124,11 +163,14 @@ public RefRefMap< Spot, Spot > getMapping()
 		return mapAB;
 	}
 
+	// -- Helper methods --
+
 	private static AffineTransform3D noOffsetTransform( AffineTransform3D transformAB )
 	{
 		AffineTransform3D noOffsetTransform = new AffineTransform3D();
 		noOffsetTransform.set( transformAB );
 		noOffsetTransform.setTranslation( 0, 0, 0 );
 		return noOffsetTransform;
 	}
+
 }