Show all classes when compressing data in heat map.

Refs #436

Author: mikekucera

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

@@ -1,6 +1,7 @@
 package org.baderlab.csplugins.enrichmentmap.model;
 
 import java.util.ArrayList;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Optional;
 
@@ -9,20 +10,37 @@
 public class CompressedClass implements ExpressionData {
 
 	private final ExpressionCache expressionCache;
-	private List<Pair<EMDataSet, String>> headers = new ArrayList<>();
+	private List<Pair<EMDataSet,String>> headers = new ArrayList<>();
 
 	public CompressedClass(List<EMDataSet> datasets, ExpressionCache expressionCache) {
 		this.expressionCache = expressionCache;
 
-		for (EMDataSet dataset : datasets) {
+		for(EMDataSet dataset : datasets) {
+			LinkedHashSet<String> uniquePhenos = new LinkedHashSet<>();
 			SetOfEnrichmentResults enrichments = dataset.getEnrichments();
+			
+			// Make sure the chosen classes go first
 			String pheno1 = enrichments.getPhenotype1();
+			if(pheno1 != null)
+				uniquePhenos.add(pheno1);
+			
 			String pheno2 = enrichments.getPhenotype2();
+			if(pheno2 != null)
+				uniquePhenos.add(pheno2);
 
-			if (pheno1 != null)
-				headers.add(Pair.of(dataset, pheno1));
-			if (pheno2 != null)
-				headers.add(Pair.of(dataset, pheno2));
+			// Add the rest of the classes
+			String[] phenotypes = enrichments.getPhenotypes();
+			if(phenotypes != null) {
+				for(String pheno : phenotypes) {
+					if(pheno != null) {
+						uniquePhenos.add(pheno);
+					}
+				}
+			}
+			
+			for(String pheno : uniquePhenos) {
+				headers.add(Pair.of(dataset, pheno));
+			}
 		}
 	}
 
@@ -47,22 +65,18 @@ public double getValue(int geneID, int idx, Compress compress, Transform transfo
 		String pheno = getName(idx);
 
 		String[] phenotypes = dataset.getEnrichments().getPhenotypes();
-
 		if (phenotypes == null || phenotypes.length == 0)
 			return Double.NaN;
 
 		Optional<float[]> optExpr = expressionCache.getExpressions(geneID, dataset, transform);
-
 		if (!optExpr.isPresent())
 			return Double.NaN;
 
 		float[] expressions = optExpr.get();
-
 		if (expressions.length == 0 || expressions.length != phenotypes.length)
 			return Double.NaN;
 
 		int size = 0;
-
 		for (int i = 0; i < expressions.length; i++) {
 			if (pheno.equals(phenotypes[i]))
 				size++;
@@ -80,7 +94,7 @@ public double getValue(int geneID, int idx, Compress compress, Transform transfo
 			case CLASS_MEDIAN: return GeneExpression.median(vals);
 			case CLASS_MAX:    return GeneExpression.max(vals);
 			case CLASS_MIN:    return GeneExpression.min(vals);
-			default:	           return Double.NaN;
+			default:           return Double.NaN;
 		}
 	}