WIP: Initial clustering implementation for bubbles

Author: Balletie

dnainator-core/src/main/java/nl/tudelft/dnainator/graph/Graph.java

@@ -69,12 +69,12 @@ public interface Graph extends AnnotationCollection {
 
 	/**
 	 * Return a list of nodes that belong to the same cluster as the given startId.
-	 * @param startNodes	the start nodes
+	 * @param start	the start nodes
 	 * @param end		the maximum rank of the cluster
 	 * @param threshold	the clustering threshold
 	 * @return		a list representing the cluster
 	 */
-	Map<Integer, List<Cluster>> getAllClusters(List<String> startNodes, int end, int threshold);
+	Map<Integer, List<Cluster>> getAllClusters(int start, int end, int threshold);
 
 	/**
 	 * Sets the interestingness strategy which calculates the interestingness when

dnainator-core/src/main/java/nl/tudelft/dnainator/graph/impl/Neo4jGraph.java

@@ -171,9 +171,9 @@ public int getRankFromBasePair(int base) {
 	}
 
 	@Override
-	public Map<Integer, List<Cluster>> getAllClusters(List<String> startNodes,
+	public Map<Integer, List<Cluster>> getAllClusters(int start,
 							int end, int threshold) {
-		return query(new AllClustersQuery(startNodes, end, threshold, is));
+		return query(new AllClustersQuery(start, end, threshold, is));
 	}
 
 	@Override

dnainator-core/src/main/java/nl/tudelft/dnainator/graph/impl/command/TopologicalPathExpander.java

@@ -53,22 +53,21 @@ private boolean hasUnprocessedIncoming(Node n) {
 	public Iterable<Relationship> expand(Path path,
 			BranchState<Object> noState) {
 		Node from = path.endNode();
-
 		// Propagate all unclosed bubbles and the newly created ones.
 		Set<Long> toPropagate = getSourcesToPropagate(from);
 
 		// For each unclosed bubble source, remove the current node from the endings and
 		// add outgoing nodes to the ending nodes, thereby advancing the bubble endings.
 		toPropagate.forEach(e -> advanceEnds(e, from));
-
+		// Store in this node the bubbles in which it is nested.
+		storeOuterBubbles(from, toPropagate);
 		// Create a new bubblesource, that will have its own bubble endings.
 		createBubbleSource(from, toPropagate);
 
 		// Encode the unclosed propagated bubbles on the edges.
 		from.getRelationships(RelTypes.NEXT, Direction.OUTGOING)
 			.forEach(out -> propagateSourceIDs(toPropagate, out));
 
-		// Process all outgoing edges.
 		List<Relationship> expand = new LinkedList<>();
 		for (Relationship out : from.getRelationships(RelTypes.NEXT, Direction.OUTGOING)) {
 			setNumStrainsThrough(out);
@@ -82,6 +81,13 @@ public Iterable<Relationship> expand(Path path,
 		return expand;
 	}
 
+	private void storeOuterBubbles(Node from, Set<Long> toPropagate) {
+		// Set the source id of the bubbles to which this node belongs. Excludes its own
+		// source id if it's a source.
+		from.setProperty(BubbleProperties.BUBBLE_SOURCE_IDS.name(),
+				toPropagate.stream().mapToLong(l -> l).toArray());
+	}
+
 	private Set<Long> getSourcesToPropagate(Node from) {
 		Iterable<Relationship> ins = from.getRelationships(RelTypes.NEXT, Direction.INCOMING);
 
@@ -121,8 +127,6 @@ private void createBubbleSource(Node n, Set<Long> toPropagate) {
 	}
 
 	private void propagateSourceIDs(Set<Long> propagatedUnique, Relationship out) {
-		out.setProperty(BubbleProperties.BUBBLE_SOURCE_IDS.name(),
-				propagatedUnique.stream().mapToLong(l -> l).toArray());
 		relIDtoSourceIDs.put(out.getId(), propagatedUnique);
 	}
 

dnainator-core/src/main/java/nl/tudelft/dnainator/graph/impl/query/AllClustersQuery.java

@@ -3,130 +3,172 @@
 import nl.tudelft.dnainator.annotation.Annotation;
 import nl.tudelft.dnainator.core.EnrichedSequenceNode;
 import nl.tudelft.dnainator.core.impl.Cluster;
+import nl.tudelft.dnainator.graph.impl.Neo4jScoreContainer;
 import nl.tudelft.dnainator.graph.impl.Neo4jSequenceNode;
 import nl.tudelft.dnainator.graph.impl.NodeLabels;
 import nl.tudelft.dnainator.graph.impl.RelTypes;
+import nl.tudelft.dnainator.graph.impl.properties.BubbleProperties;
 import nl.tudelft.dnainator.graph.impl.properties.SequenceProperties;
 import nl.tudelft.dnainator.graph.interestingness.InterestingnessStrategy;
 
+import org.neo4j.graphalgo.GraphAlgoFactory;
+import org.neo4j.graphalgo.PathFinder;
 import org.neo4j.graphdb.Direction;
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.graphdb.Node;
+import org.neo4j.graphdb.Path;
+import org.neo4j.graphdb.PathExpanders;
+import org.neo4j.graphdb.traversal.Evaluation;
 import org.neo4j.graphdb.traversal.TraversalDescription;
+import org.neo4j.helpers.collection.IteratorUtil;
 
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
-import java.util.PriorityQueue;
-import java.util.Queue;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 /**
  * The {@link AllClustersQuery} creates {@link Cluster}s from all nodes,
- * starting at the startNodes, and ending when the maximum specified start rank is reached.
+ * between the given ranks using the given threshold value.
  */
 public class AllClustersQuery implements Query<Map<Integer, List<Cluster>>> {
-	private Set<String> visited;
-	private List<String> startNodes;
-	private int threshold;
+	private int minRank;
 	private int maxRank;
+	private int threshold;
 	private InterestingnessStrategy is;
 
 	/**
 	 * Create a new {@link AllClustersQuery}, which will:.
-	 * - start clustering at the specified startNodes
-	 * - stop clustering when the end rank is reached
+	 * - get all clusters between the given ranks
 	 * - use the specified clustering threshold
-	 * @param startNodes	the start nodes
+	 * @param minRank	the minimum rank
 	 * @param maxRank	the maximum rank
 	 * @param threshold	the clustering threshold
 	 * @param is the interestingness strategy, which determines how the
 	 * interestingness score is calculated.
 	 */
-	public AllClustersQuery(List<String> startNodes, int maxRank, int threshold,
+	public AllClustersQuery(int minRank, int maxRank, int threshold,
 			InterestingnessStrategy is) {
-		this.startNodes = startNodes;
+		this.minRank = minRank;
 		this.maxRank = maxRank;
 		this.threshold = threshold;
-		this.visited = new HashSet<>();
 		this.is = is;
 	}
 
+	private TraversalDescription untilMaxRank(GraphDatabaseService service) {
+		return service.traversalDescription()
+				.breadthFirst()
+				.evaluator(path -> {
+					if ((int) path.endNode().getProperty(SequenceProperties.RANK.name())
+							<= maxRank) {
+						return Evaluation.INCLUDE_AND_CONTINUE;
+					} else {
+						return Evaluation.EXCLUDE_AND_PRUNE;
+					}
+				})
+				.relationships(RelTypes.NEXT, Direction.OUTGOING);
+	}
+
 	@Override
 	public Map<Integer, List<Cluster>> execute(GraphDatabaseService service) {
-		Queue<Cluster> rootClusters = new PriorityQueue<>((e1, e2) -> 
-			e1.getStartRank() - e2.getStartRank()
-		);
-		Map<Integer, List<Cluster>> result = new HashMap<Integer, List<Cluster>>();
-	
-		rootClusters.addAll(clustersFrom(service, startNodes));
-
-		// Find adjacent clusters as long as there are root clusters in the queue
-		int minrank = rootClusters.stream().mapToInt(e -> e.getStartRank()).min().orElse(0);
-		while (!rootClusters.isEmpty()) {
-			Cluster c = rootClusters.poll();
-			if (c.getStartRank() < minrank || c.getStartRank() > maxRank) {
-				continue;
+		Set<Long> bubbleSourcesToCluster = new HashSet<>();
+		Set<Long> bubbleSourcesToKeepIntact = new HashSet<>();
+		Iterable<Node> start = IteratorUtil.loop(service.findNodes(NodeLabels.NODE,
+				SequenceProperties.RANK.name(), minRank));
+		for (Node n : untilMaxRank(service).traverse(start).nodes()) {
+			if (n.hasLabel(NodeLabels.BUBBLE_SOURCE)) {
+				bubbleSourcesToCluster.add(n.getId());
+			}
+			int interestingness = is.compute(new Neo4jScoreContainer(n));
+			if (interestingness > threshold) {
+				for (long sourceID
+						: (long[]) n.getProperty(BubbleProperties.BUBBLE_SOURCE_IDS.name())) {
+					bubbleSourcesToKeepIntact.add(sourceID);
+					bubbleSourcesToCluster.remove(sourceID);
+				}
 			}
-			result.putIfAbsent(c.getStartRank(), new ArrayList<>());
-			result.get(c.getStartRank()).add(c);
-
-			c.getNodes().forEach(sn -> {
-				rootClusters.addAll(clustersFrom(service, sn.getOutgoing()));
-			});
 		}
+		return cluster(service, bubbleSourcesToCluster, bubbleSourcesToKeepIntact);
+	}
 
-		return result;
+	private Map<Integer, List<Cluster>> cluster(GraphDatabaseService service,
+			Set<Long> bubbleSourcesToCluster, Set<Long> bubbleSourcesToKeepIntact) {
+		Map<Integer, List<Cluster>> bubblesClustered = bubbleSourcesToCluster.stream()
+				.map(service::getNodeById)
+				.map(source -> collapseBubble(service, source, getSinkFromSource(source)))
+				.collect(Collectors.groupingBy(Cluster::getStartRank));
+		Stream<Map<Integer, List<Cluster>>> singletonClusters = bubbleSourcesToKeepIntact.stream()
+				.map(service::getNodeById)
+				.map(source -> getSingletonClusters(service, source, getSinkFromSource(source)));
+		return mergeMaps(Stream.concat(Stream.of(bubblesClustered), singletonClusters));
 	}
 
-	private Queue<Cluster> clustersFrom(GraphDatabaseService service, List<String> startNodes) {
-		Queue<Cluster> rootClusters = new LinkedList<Cluster>();
+	private static Node getSinkFromSource(Node source) {
+		return source.getSingleRelationship(RelTypes.BUBBLE_SOURCE_OF, Direction.OUTGOING)
+				.getEndNode();
+	}
 
-		for (String sn : startNodes) {
-			// Continue if this startNode was consumed by another cluster
-			if (visited.contains(sn)) {
-				continue;
-			}
+	private Map<Integer, List<Cluster>> getSingletonClusters(GraphDatabaseService service,
+			Node source, Node sink) {
+		int sourceRank = (int) source.getProperty(SequenceProperties.RANK.name());
+		int sinkRank = (int) sink.getProperty(SequenceProperties.RANK.name());
+		PathFinder<Path> withinBubble = pathFinderBetweenRanks(sourceRank, sinkRank);
+		return stream(withinBubble.findAllPaths(source, sink))
+				.flatMap(path -> stream(path.nodes()))
+				.distinct()
+				.map(n -> createSingletonCluster(service, n))
+				.collect(Collectors.groupingBy(Cluster::getStartRank));
+	}
 
-			// Otherwise get the cluster starting from this startNode
-			rootClusters.add(cluster(service, sn));
-		}
+	private Cluster createSingletonCluster(GraphDatabaseService service, Node n) {
+		EnrichedSequenceNode sn = new Neo4jSequenceNode(service, n);
+		return new Cluster((int) n.getProperty(SequenceProperties.RANK.name()),
+				Collections.singletonList(sn), sn.getAnnotations());
+	}
 
-		return rootClusters;
+	private Cluster collapseBubble(GraphDatabaseService service, Node source, Node sink) {
+		int sourceRank = (int) source.getProperty(SequenceProperties.RANK.name());
+		int sinkRank = (int) sink.getProperty(SequenceProperties.RANK.name());
+		int clusterRank = sourceRank + (sinkRank - sourceRank) / 2;
+		PathFinder<Path> withinBubble = pathFinderBetweenRanks(sourceRank, sinkRank);
+		// FIXME: don't collapse source and sink, keep those intact.
+		List<EnrichedSequenceNode> nodes = stream(
+				withinBubble.findAllPaths(source, sink))
+				.flatMap(path -> stream(path.nodes()))
+				.distinct()
+				.map(n -> new Neo4jSequenceNode(service, n))
+				.collect(Collectors.toList());
+		List<Annotation> annotations = nodes.stream()
+				.flatMap(e -> e.getAnnotations().stream())
+				.collect(Collectors.toList());
+		return new Cluster(clusterRank, nodes, annotations);
 	}
 
-	private Cluster cluster(GraphDatabaseService service, String start) {
-		Cluster cluster = null;
-		Node startNode = service.findNode(NodeLabels.NODE, SequenceProperties.ID.name(), start);
-		List<Node> result = new ArrayList<>();
+	private PathFinder<Path> pathFinderBetweenRanks(int minRank, int maxRank) {
+		return GraphAlgoFactory.allSimplePaths(
+				PathExpanders.forTypeAndDirection(RelTypes.NEXT, Direction.OUTGOING),
+				maxRank - minRank);
+	}
 
-		// A depth first traversal traveling along both incoming and outgoing edges.
-		TraversalDescription clusterDesc = service.traversalDescription()
-						.depthFirst()
-						.relationships(RelTypes.NEXT, Direction.BOTH)
-						.evaluator(new ClusterEvaluator(threshold, visited, is));
-		// Traverse the cluster starting from the startNode.
-		int rankStart = (int) startNode.getProperty(SequenceProperties.RANK.name());
-		for (Node end : clusterDesc.traverse(startNode).nodes()) {
-			result.add(end);
+	private Map<Integer, List<Cluster>> mergeMaps(Stream<Map<Integer, List<Cluster>>> concat) {
+		return concat.map(Map::entrySet)
+				.flatMap(Collection::stream)
+				.collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue(), (left, right) -> {
+					List<Cluster> newList = new ArrayList<>(left.size() + right.size());
+					newList.addAll(right);
+					newList.addAll(left);
+					return left;
+				}));
+	}
 
-			// Update this cluster's start rank according to the lowest node rank.
-			int endRank = (int) startNode.getProperty(SequenceProperties.RANK.name());
-			if (endRank < rankStart) {
-				rankStart = endRank;
-			}
-		}
-		// Might want to internally pass nodes.
-		List<EnrichedSequenceNode> retrieve = result.stream()
-				.map(e -> new Neo4jSequenceNode(service, e))
-				.collect(Collectors.toList());
-		List<Annotation> annotations = retrieve.stream().flatMap(e -> e.getAnnotations().stream())
-				.collect(Collectors.toList());
-		cluster = new Cluster(rankStart, retrieve, annotations);
-		return cluster;
+	private static <T> Stream<T> stream(Iterable<T> in) {
+		// Quick utility method, for converting iterables to streams.
+	    return StreamSupport.stream(in.spliterator(), false);
 	}
 }

dnainator-core/src/main/java/nl/tudelft/dnainator/graph/impl/query/ClusterEvaluator.java

@@ -81,10 +81,8 @@ public void loadChildren(Bounds bounds) {
 		Range ranks = getRange(bounds);
 
 		System.out.println("load iteration: " + ranks.getX() + " -> " + ranks.getY());
-		List<String> roots = graph.getRank(ranks.getX()).stream()
-				.map(SequenceNode::getId).collect(Collectors.toList());
 		List<Annotation> annotations = getSortedAnnotations(ranks);
-		Map<Integer, List<Cluster>> result = graph.getAllClusters(roots, ranks.getY(),
+		Map<Integer, List<Cluster>> result = graph.getAllClusters(ranks.getX(), ranks.getY(),
 				(int) (bounds.getWidth() / CLUSTER_DIVIDER));
 		clusters.clear();
 		childContent.getChildren().clear();