Record edge SUIDs in AbstractDataSet

Author: mikekucera

EnrichmentMapPlugin/src/main/java/org/baderlab/csplugins/enrichmentmap/model/AbstractDataSet.java

@@ -1,13 +1,16 @@
 package org.baderlab.csplugins.enrichmentmap.model;
 
 import java.text.Collator;
-import java.util.HashMap;
-import java.util.Map;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
 
 public abstract class AbstractDataSet implements Comparable<AbstractDataSet> {
 
 	private final String name;
-	private final Map<String, Long> nodeSuids = new HashMap<>();
+	private final Set<Long> nodeSuids = new HashSet<>();
+	private final Set<Long> edgeSuids = new HashSet<>();
+	
 	/** EnrichmentMap only creates nodes for these genes. */
 	private SetOfGeneSets geneSetsOfInterest = new SetOfGeneSets();
 
@@ -22,22 +25,22 @@ public String getName() {
 		return name;
 	}
 
-	public Map<String, Long> getNodeSuids() {
+	public Set<Long> getNodeSuids() {
 		synchronized (lock) {
-			return new HashMap<>(nodeSuids);
+			return Collections.unmodifiableSet(nodeSuids);
 		}
 	}
 
-	public void setNodeSuids(Map<String, Long> newValue) {
+	public void setNodeSuids(Set<Long> newValue) {
 		synchronized (lock) {
 			nodeSuids.clear();
-			nodeSuids.putAll(newValue);
+			nodeSuids.addAll(newValue);
 		}
 	}
 
-	public void addNodeSuid(String key, Long suid) {
+	public void addNodeSuid(Long suid) {
 		synchronized (lock) {
-			nodeSuids.put(key, suid);
+			nodeSuids.add(suid);
 		}
 	}
 
@@ -47,6 +50,31 @@ public void clearNodeSuids() {
 		}
 	}
 
+	public Set<Long> getEdgeSuids() {
+		synchronized (lock) {
+			return Collections.unmodifiableSet(edgeSuids);
+		}
+	}
+	
+	public void setEdgeSuids(Set<Long> newValue) {
+		synchronized (lock) {
+			edgeSuids.clear();
+			edgeSuids.addAll(newValue);
+		}
+	}
+	
+	public void addEdgeSuid(Long suid) {
+		synchronized (lock) {
+			edgeSuids.add(suid);
+		}
+	}
+	
+	public void clearEdgeSuids() {
+		synchronized (lock) {
+			edgeSuids.clear();
+		}
+	}
+	
 	public SetOfGeneSets getGeneSetsOfInterest() {
 		return geneSetsOfInterest;
 	}

EnrichmentMapPlugin/src/main/java/org/baderlab/csplugins/enrichmentmap/model/EnrichmentMap.java

@@ -13,6 +13,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import org.baderlab.csplugins.enrichmentmap.model.EMDataSet.Method;
@@ -359,43 +360,64 @@ public void setNetworkID(long networkID) {
 		this.networkID = networkID;
 	}
 
+	public static Set<Long> getNodesUnion(Collection<AbstractDataSet> dataSets) {
+		return getUnion(dataSets, AbstractDataSet::getNodeSuids);
+	}
+	
+	public static Set<Long> getNodesIntersection(Collection<? extends AbstractDataSet> dataSets) { 
+		return getIntersection(dataSets, AbstractDataSet::getNodeSuids);
+	}
+			
 	/**
-	 * Returns the node SUIDs for all the gene-sets in the given collection of DataSets.
+	 * Returns the SUIDs for all the gene-sets in the given collection of DataSets.
 	 * Each returned gene-set is contained in at least one of the given DataSets.
+	 * 
+	 * Note, this will only return distinct edges, not compound edges.
 	 */
-	public static Set<Long> getNodesUnion(Collection<AbstractDataSet> dataSets) {
+	public static Set<Long> getEdgesUnion(Collection<AbstractDataSet> dataSets) {
+		return getUnion(dataSets, AbstractDataSet::getEdgeSuids);
+	}
+	
+	/**
+	 * Returns the SUIDs for all the gene-sets in the given collection of DataSets.
+	 * Each returned gene-set is contained all of the given DataSets.
+	 * 
+	 * Note, this will only return distinct edges, not compound edges.
+	 */
+	public static Set<Long> getEdgesIntersection(Collection<? extends AbstractDataSet> dataSets) { 
+		return getIntersection(dataSets, AbstractDataSet::getEdgeSuids);
+	}
+	
+	private static Set<Long> getUnion(Collection<? extends AbstractDataSet> dataSets, Function<AbstractDataSet,Set<Long>> suidSupplier) {
 		if (dataSets.isEmpty())
 			return Collections.emptySet();
 
 		Set<Long> suids = new HashSet<>();
 
 		for (AbstractDataSet ds : dataSets) {
-			suids.addAll(ds.getNodeSuids().values());
+			suids.addAll(suidSupplier.apply(ds));
 		}
 
 		return suids;
 	}
 
-	/**
-	 * Returns the node SUIDs for all the gene-sets in the given collection of DataSets.
-	 * Each returned gene-set is contained all of the given DataSets.
-	 */
-	public static Set<Long> getNodesIntersection(Collection<EMDataSet> queryDatasets) {
-		if (queryDatasets.isEmpty())
+	private static Set<Long> getIntersection(Collection<? extends AbstractDataSet> dataSets, Function<AbstractDataSet,Set<Long>> suidSupplier) {
+		if (dataSets.isEmpty())
 			return Collections.emptySet();
 
-		Iterator<EMDataSet> iter = queryDatasets.iterator();
-		EMDataSet first = iter.next();
-		Set<Long> suids = new HashSet<>(first.getNodeSuids().values());
+		Iterator<? extends AbstractDataSet> iter = dataSets.iterator();
+		AbstractDataSet first = iter.next();
+		Set<Long> suids = new HashSet<>(suidSupplier.apply(first));
 
 		while (iter.hasNext()) {
-			EMDataSet dataset = iter.next();
-			suids.retainAll(dataset.getNodeSuids().values());
+			AbstractDataSet dataset = iter.next();
+			suids.retainAll(suidSupplier.apply(dataset));
 		}
 
 		return suids;
 	}
 
+	
 	public Set<String> getAllRankNames() {
 		Set<String> allRankNames = new HashSet<>();
 

EnrichmentMapPlugin/src/main/java/org/baderlab/csplugins/enrichmentmap/model/SimilarityKey.java

@@ -7,18 +7,17 @@ public class SimilarityKey {
 	private final String geneSet1;
 	private final String geneSet2;
 	private final String interaction;
-	// a set of zero indicates a compound edge
-	private final int set;
+	private final String name;
 
 
-	public SimilarityKey(String geneSet1, String geneSet2, String interaction, int set) {
+	public SimilarityKey(String geneSet1, String geneSet2, String interaction, String name) {
 		Objects.requireNonNull(geneSet1);
 		Objects.requireNonNull(interaction);
 		Objects.requireNonNull(geneSet2);
 		this.geneSet1 = geneSet1;
 		this.geneSet2 = geneSet2;
 		this.interaction = interaction;
-		this.set = set;
+		this.name = name;
 	}
 
 	public String getGeneSet1() {
@@ -33,27 +32,30 @@ public String getInteraction() {
 		return interaction;
 	}
 
-	public int getSet() {
-		return set;
+	public boolean isCompound() {
+		return name == null;
 	}
 
-	public boolean isCompound() {
-		return set == 0;
+	public String getName() {
+		return name;
 	}
 
 	@Override
 	public int hashCode() {
 		// add the hash codes from the genesets so that we get the same hash code regardless of the order
-		return Objects.hash(geneSet1.hashCode() + geneSet2.hashCode(), interaction, set);
+		return Objects.hash(geneSet1.hashCode() + geneSet2.hashCode(), interaction, name);
 	}
 
 	@Override
 	public boolean equals(Object o) {
 		if(!(o instanceof SimilarityKey))
 			return false;
 		SimilarityKey other = (SimilarityKey)o;
-		
-		if(set != other.set)
+		if(name != null && other.name == null)
+			return false;
+		if(name == null && other.name != null)
+			return false;
+		if(name != null && other.name != null && !name.equals(other.name))
 			return false;
 		if(!interaction.equals(other.interaction))
 			return false;
@@ -65,10 +67,10 @@ public boolean equals(Object o) {
 
 	@Override
 	public String toString() {
-		if(set == 0)
+		if(isCompound())
 			return String.format("%s (%s) %s", geneSet1, interaction, geneSet2);
 		else
-			return String.format("%s (%s_set%d) %s", geneSet1, interaction, set, geneSet2);
+			return String.format("%s (%s_set%s) %s", geneSet1, interaction, name, geneSet2);
 	}
 
 }

EnrichmentMapPlugin/src/main/java/org/baderlab/csplugins/enrichmentmap/model/io/SessionModelListener.java

@@ -3,9 +3,10 @@
 import java.awt.event.ActionEvent;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import org.baderlab.csplugins.enrichmentmap.ApplicationModule.Headless;
 import org.baderlab.csplugins.enrichmentmap.CyActivator;
@@ -17,6 +18,8 @@
 import org.baderlab.csplugins.enrichmentmap.view.heatmap.HeatMapMediator;
 import org.cytoscape.application.swing.AbstractCyAction;
 import org.cytoscape.model.CyColumn;
+import org.cytoscape.model.CyEdge;
+import org.cytoscape.model.CyIdentifiable;
 import org.cytoscape.model.CyNetwork;
 import org.cytoscape.model.CyNetworkManager;
 import org.cytoscape.model.CyNode;
@@ -160,7 +163,7 @@ private boolean restoreModelFromTables(CySession session) {
 						if(em != null) {
 							em.setServiceRegistrar(serviceRegistrar);
 							em.setNetworkID(network.getSUID());
-							updateNodeSuids(em, session);
+							updateSuids(em, session);
 							emManager.registerEnrichmentMap(em);
 							sessionHasEM = true;
 						}
@@ -171,24 +174,24 @@ private boolean restoreModelFromTables(CySession session) {
 		return sessionHasEM;
 	}
 
-	private void updateNodeSuids(EnrichmentMap map, CySession session) {
-		for (EMDataSet ds : map.getDataSetList()) {
-			Map<String, Long> oldSuids = ds.getNodeSuids();
-			Map<String, Long> newSuids = new HashMap<>();
-			
-			for (String key : oldSuids.keySet()) {
-				Long suid = oldSuids.get(key);
-				
-				if (session != null) {
-					// If we are loading from a session file then we need to re-map the ids
-					CyNode node = session.getObject(suid, CyNode.class);
-					suid = node.getSUID();
-				}
-				
-				newSuids.put(key, suid);
+	private static Set<Long> mapSuids(Set<Long> oldSuids, CySession session, Class<? extends CyIdentifiable> type) {
+		Set<Long> newSuids = new HashSet<>();
+		
+		for (Long suid : oldSuids) {
+			if (session != null) {
+				// If we are loading from a session file then we need to re-map the ids
+				CyIdentifiable obj = session.getObject(suid, type);
+				suid = obj.getSUID();
 			}
-			
-			ds.setNodeSuids(newSuids);
+			newSuids.add(suid);
+		}
+		return newSuids;
+	}
+	
+	private void updateSuids(EnrichmentMap map, CySession session) {
+		for(EMDataSet ds : map.getDataSetList()) {
+			ds.setNodeSuids(mapSuids(ds.getNodeSuids(), session, CyNode.class));
+			ds.setEdgeSuids(mapSuids(ds.getEdgeSuids(), session, CyEdge.class));
 		}
 	}
 

EnrichmentMapPlugin/src/main/java/org/baderlab/csplugins/enrichmentmap/style/EMStyleBuilder.java

@@ -374,7 +374,7 @@ private void setNodeColors(VisualStyle vs, EMStyleOptions options) {
 			CyNetwork net = netView.getModel();
 
 			for (EMSignatureDataSet sds : signatureDataSets) {
-				for (Long suid : sds.getNodeSuids().values()) {
+				for (Long suid : sds.getNodeSuids()) {
 					CyNode node = net.getNode(suid);
 
 					if (node != null) {

EnrichmentMapPlugin/src/main/java/org/baderlab/csplugins/enrichmentmap/task/ComputeSimilarityTaskParallel.java

@@ -87,9 +87,8 @@ private Map<SimilarityKey,GenesetSimilarity> startComputeSimilarities(TaskMonito
 						continue; //don't compare two identical gene sets
 
 					if(distinct) {
-						int i = 1;
 						for(EMDataSet dataset : dataSets) {
-							SimilarityKey key = new SimilarityKey(geneset1Name, geneset2Name, edgeType, i++);
+							SimilarityKey key = new SimilarityKey(geneset1Name, geneset2Name, edgeType, dataset.getName());
 
 							if(!similarities.containsKey(key)) {
 								Map<String,GeneSet> genesets = dataset.getGeneSetsOfInterest().getGeneSets();
@@ -108,7 +107,7 @@ private Map<SimilarityKey,GenesetSimilarity> startComputeSimilarities(TaskMonito
 					}
 
 					if(compound) {
-						SimilarityKey key = new SimilarityKey(geneset1Name, geneset2Name, edgeType, 0);
+						SimilarityKey key = new SimilarityKey(geneset1Name, geneset2Name, edgeType, null);
 
 						if(!similarities.containsKey(key)) {
 							Set<Integer> geneset1 = unionedGenesets.get(geneset1Name);

EnrichmentMapPlugin/src/main/java/org/baderlab/csplugins/enrichmentmap/task/CreateDiseaseSignatureTask.java

@@ -330,7 +330,7 @@ public void buildDiseaseSignature(TaskMonitor taskMonitor) {
 					continue;
 
 				boolean passedCutoff = passesCutoff(edgeName);
-				createEdge(edgeName, currentNetwork, currentView, prefix, edgeTable, nodeTable, passedCutoff);
+				createEdge(edgeName, currentNetwork, currentView, prefix, edgeTable, nodeTable, passedCutoff, sigDataSet);
 			}
 
 			widthFunctionProvider.get().setEdgeWidths(currentNetwork, prefix, taskMonitor);
@@ -395,7 +395,7 @@ private boolean createHubNode(
 		if (hubNode == null) {
 			hubNode = network.addNode();
 			taskResult.addNewNode(hubNode);
-			sigDataSet.addNodeSuid(hubName, hubNode.getSUID());
+			sigDataSet.addNodeSuid(hubNode.getSUID());
 			created = true;
 		}
 
@@ -452,7 +452,7 @@ private boolean createHubNode(
 	 * returned, if the edge had to be created it will not be returned.
 	 */
 	private void createEdge(String edgeName, CyNetwork network, CyNetworkView netView, String prefix, CyTable edgeTable,
-			CyTable nodeTable, boolean passedCutoff) {
+			CyTable nodeTable, boolean passedCutoff, EMSignatureDataSet sigDataSet) {
 		CyEdge edge = NetworkUtil.getEdgeWithValue(network, edgeTable, CyNetwork.NAME, edgeName);
 		GenesetSimilarity genesetSimilarity = geneSetSimilarities.get(edgeName);
 
@@ -467,6 +467,7 @@ private void createEdge(String edgeName, CyNetwork network, CyNetworkView netVie
 					return;
 
 				edge = network.addEdge(hubNode, geneSet, false);
+				sigDataSet.addEdgeSuid(edge.getSUID());
 				taskResult.addNewEdge(edge);
 			} else {
 				return; // edge does not exist and does not pass cutoff, do nothing