Changed array for map to cache prior arrays by number-of-alleles

broadinstitute · Jun 14, 2021 · caace67 · caace67
1 parent 3a5b468
commit caace67
Showing 1 changed file with 11 additions and 9 deletions.
diff --git a/src/main/java/org/broadinstitute/hellbender/tools/walkers/annotator/AllelePseudoDepth.java b/src/main/java/org/broadinstitute/hellbender/tools/walkers/annotator/AllelePseudoDepth.java
@@ -2,6 +2,8 @@
 
 import htsjdk.variant.variantcontext.*;
 import htsjdk.variant.vcf.*;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 import org.apache.commons.math3.linear.DefaultRealMatrixChangingVisitor;
 import org.apache.commons.math3.linear.RealMatrix;
 import org.apache.commons.math3.linear.RealMatrixChangingVisitor;
@@ -29,7 +31,6 @@ public final class AllelePseudoDepth implements GenotypeAnnotation {
     private static DecimalFormat DEPTH_FORMAT = new DecimalFormat("#.##");
     private static DecimalFormat FRACTION_FORMAT = new DecimalFormat("#.####");
 
-    public static final int INITIAL_PRIOR_CACHE_MAX_ALLELE = 10;
     public static final VCFFormatHeaderLine DEPTH_HEADER_LINE =
             new VCFFormatHeaderLine(GATKVCFConstants.PSEUDO_DEPTH_KEY, VCFHeaderLineCount.R, VCFHeaderLineType.Float, "Allele depth based on Dirichlet posterior pseudo-counts");
     public static final VCFFormatHeaderLine FRACTION_HEADER_LINE =
@@ -54,7 +55,7 @@ public final class AllelePseudoDepth implements GenotypeAnnotation {
     private static final List<String> KEYS = Arrays.asList(GATKVCFConstants.PSEUDO_DEPTH_KEY, GATKVCFConstants.PSEUDO_FRACTION_KEY);
     private static final List<VCFCompoundHeaderLine> HEADER_LINES = Arrays.asList(DEPTH_HEADER_LINE, FRACTION_HEADER_LINE);
 
-    private double[][] priorPseudoCounts = new double[INITIAL_PRIOR_CACHE_MAX_ALLELE][]; // possibly won't ever go beyond 10.
+    private Int2ObjectMap<double[]> priorPseudoCounts = new Int2ObjectOpenHashMap<>();
     @Override
     public List<VCFCompoundHeaderLine> getDescriptions() {
         return HEADER_LINES;
@@ -161,13 +162,14 @@ private double[] calculateWeights(RealMatrix lkMatrix) {
      * @return never {@code null.}
      */
     private double[] composePriorPseudoCounts(final int numberOfAlleles) {
-        if (priorPseudoCounts.length < numberOfAlleles) {
-            priorPseudoCounts = Arrays.copyOf(priorPseudoCounts, numberOfAlleles << 1);
-        }
-        double[] result = priorPseudoCounts[numberOfAlleles];
-        if (result == null) {
-            Arrays.fill(result = priorPseudoCounts[numberOfAlleles] = new double[numberOfAlleles], prior);
+        final double[] cached = priorPseudoCounts.get(numberOfAlleles);
+        if (cached == null) {
+            final double[] result = new double[numberOfAlleles];
+            Arrays.fill(result, prior);
+            priorPseudoCounts.put(numberOfAlleles, result);
+            return result;
+        } else {
+            return cached;
         }
-        return result;
     }
 }