Simplify HaplotypeBAMWriter code. #944 (#5122)

src/main/java/org/broadinstitute/hellbender/tools/walkers/haplotypecaller/AssemblyBasedCallerUtils.java

@@ -197,7 +197,7 @@ public static Optional<HaplotypeBAMWriter> createBamWriter(final AssemblyBasedCa
                                                                final boolean createBamOutMD5,
                                                                final SAMFileHeader header) {
         return args.bamOutputPath != null ?
-                Optional.of(HaplotypeBAMWriter.create(args.bamWriterType, IOUtils.getPath(args.bamOutputPath), createBamOutIndex, createBamOutMD5, header)) :
+                Optional.of(new HaplotypeBAMWriter(args.bamWriterType, IOUtils.getPath(args.bamOutputPath), createBamOutIndex, createBamOutMD5, header)) :
                 Optional.empty();
     }
 

src/main/java/org/broadinstitute/hellbender/utils/haplotype/HaplotypeBAMDestination.java

@@ -82,4 +82,4 @@ public SAMFileHeader getBAMOutputHeader() {
         return bamOutputHeader;
     }
 
-}
+}

src/main/java/org/broadinstitute/hellbender/utils/haplotype/HaplotypeBAMWriter.java

@@ -11,16 +11,16 @@
 import org.broadinstitute.hellbender.utils.Utils;
 import org.broadinstitute.hellbender.utils.genotyper.ReadLikelihoods;
 
-import java.io.File;
 import java.util.Collection;
+import java.util.LinkedHashSet;
 import java.util.function.Function;
 import java.util.Set;
 
 /**
  * A BAMWriter that aligns reads to haplotypes and emits their best alignments to a destination.
  *
  */
-public abstract class HaplotypeBAMWriter implements AutoCloseable {
+public class HaplotypeBAMWriter implements AutoCloseable {
     /**
      * Allows us to write out unique names for our synthetic haplotype reads
      */
@@ -32,7 +32,8 @@ public abstract class HaplotypeBAMWriter implements AutoCloseable {
     private static final int bestHaplotypeMQ = 60;
     private static final int otherMQ = 0;
 
-    protected final HaplotypeBAMDestination output;
+    private final HaplotypeBAMDestination output;
+    private WriterType writerType;
     private boolean writeHaplotypes = true;
 
     /**
@@ -43,26 +44,14 @@ public enum WriterType {
          * A mode that's for method developers.  Writes out all of the possible
          * haplotypes considered, as well as reads aligned to each
          */
-        ALL_POSSIBLE_HAPLOTYPES(AllHaplotypeBAMWriter::new),
+        ALL_POSSIBLE_HAPLOTYPES,
 
         /**
          * A mode for users.  Writes out the reads aligned only to the called
          * haplotypes.  Useful to understand why the caller is calling what it is
          */
-        CALLED_HAPLOTYPES(CalledHaplotypeBAMWriter::new);
+        CALLED_HAPLOTYPES
 
-        final private Function<HaplotypeBAMDestination, HaplotypeBAMWriter> factory;
-
-        WriterType(Function<HaplotypeBAMDestination, HaplotypeBAMWriter> factory) {
-            this.factory = factory;
-        }
-
-        /**
-         * Create an instance of the  HaplotypeBAMWriter corresponding to this type.
-         */
-        public HaplotypeBAMWriter create(HaplotypeBAMDestination destination) {
-            Utils.nonNull(destination, "destination cannot be null");
-            return factory.apply(destination); }
     }
 
     /**
@@ -75,17 +64,14 @@ public HaplotypeBAMWriter create(HaplotypeBAMDestination destination) {
      * @param sourceHeader the header of the input BAMs used to make calls, must not be null.
      * @return a new HaplotypeBAMWriter
      */
-    public static HaplotypeBAMWriter create(
+    public HaplotypeBAMWriter(
             final WriterType type,
             final Path outputPath,
             final boolean createBamOutIndex,
             final boolean createBamOutMD5,
             final SAMFileHeader sourceHeader) {
-        Utils.nonNull(type, "type cannot be null");
-        Utils.nonNull(outputPath, "outputPath cannot be null");
-        Utils.nonNull(sourceHeader, "sourceHeader cannot be null");
 
-        return create(type, new SAMFileDestination(outputPath, createBamOutIndex, createBamOutMD5, sourceHeader, DEFAULT_HAPLOTYPE_READ_GROUP_ID));
+        this(type, new SAMFileDestination(outputPath, createBamOutIndex, createBamOutMD5, sourceHeader, DEFAULT_HAPLOTYPE_READ_GROUP_ID));
     }
 
     /**
@@ -97,22 +83,14 @@ public static HaplotypeBAMWriter create(
      *
      * @return a new HaplotypeBAMWriter
      */
-    public static HaplotypeBAMWriter create(final WriterType type, final HaplotypeBAMDestination destination) {
+    public HaplotypeBAMWriter(
+            final WriterType type,
+            final HaplotypeBAMDestination destination) {
+
         Utils.nonNull(type, "type cannot be null");
         Utils.nonNull(destination, "destination cannot be null");
-
-        return type.create(destination);
-    }
-
-    /**
-     * Create a new HaplotypeBAMWriter writing its output to the given destination
-     *
-     * @param output the output destination, must not be null. Note that SAM writer associated with the destination must
-     *               have its presorted bit set to false, as reads may come in out of order during writing.
-     */
-    protected HaplotypeBAMWriter(final HaplotypeBAMDestination output) {
-        Utils.nonNull(output, "output cannot be null");
-        this.output = output;
+        this.writerType = type;
+        this.output = destination;
     }
 
     /**
@@ -125,25 +103,52 @@ public void close() {
 
     /**
      * Write out a BAM representing for the haplotype caller at this site
+     * writerType (ALL_POSSIBLE_HAPLOTYPES or CALLED_HAPLOTYPES) determines inputs to writeHaplotypesAsReads
+     * Write out a BAM representing the haplotypes at this site, based on the value for writerType used when
+     * the writer was constructed (ALL_POSSIBLE_HAPLOTYPES or CALLED_HAPLOTYPES).
      *
-     * @param haplotypes a list of all possible haplotypes at this loc
-     * @param paddedReferenceLoc the span of the based reference here
-     * @param bestHaplotypes a list of the best (a subset of all) haplotypes that actually went forward into genotyping
-     * @param calledHaplotypes a list of the haplotypes that where actually called as non-reference
-     * @param readLikelihoods a map from sample -> likelihoods for each read for each of the best haplotypes
+     * @param haplotypes a list of all possible haplotypes at this loc, cannot be null
+     * @param paddedReferenceLoc the span of the based reference here, cannot be null
+     * @param bestHaplotypes a list of the best (a subset of all) haplotypes that actually went forward into genotyping, cannot be null
+     * @param calledHaplotypes a list of the haplotypes that where actually called as non-reference, cannot be null
+     * @param readLikelihoods a map from sample -> likelihoods for each read for each of the best haplotypes, cannot be null
      */
-    public abstract void writeReadsAlignedToHaplotypes(final Collection<Haplotype> haplotypes,
-                                                       final Locatable paddedReferenceLoc,
-                                                       final Collection<Haplotype> bestHaplotypes,
-                                                       final Set<Haplotype> calledHaplotypes,
-                                                       final ReadLikelihoods<Haplotype> readLikelihoods);
+    public void writeReadsAlignedToHaplotypes(final Collection<Haplotype> haplotypes,
+                                              final Locatable paddedReferenceLoc,
+                                              final Collection<Haplotype> bestHaplotypes,
+                                              final Set<Haplotype> calledHaplotypes,
+                                              final ReadLikelihoods<Haplotype> readLikelihoods) {
+
+        Utils.nonNull(haplotypes, "haplotypes cannot be null");
+        Utils.nonNull(paddedReferenceLoc, "paddedReferenceLoc cannot be null");
+        Utils.nonNull(calledHaplotypes, "calledHaplotypes cannot be null");
+        Utils.nonNull(readLikelihoods, "readLikelihoods cannot be null");
+        Utils.nonNull(bestHaplotypes, "bestHaplotypes cannot be null");
+
+        if (writerType.equals(WriterType.CALLED_HAPLOTYPES)){
+            if (calledHaplotypes.isEmpty()){
+                return;
+            }
+            writeHaplotypesAsReads(calledHaplotypes, calledHaplotypes, paddedReferenceLoc);
+
+        } else {
+            writeHaplotypesAsReads(haplotypes, new LinkedHashSet<>(bestHaplotypes), paddedReferenceLoc);
+        }
+
+        final int sampleCount = readLikelihoods.numberOfSamples();
+        for (int i = 0; i < sampleCount; i++) {
+            for (final GATKRead read : readLikelihoods.sampleReads(i)) {
+                writeReadAgainstHaplotype(read);
+            }
+        }
+    }
 
     /**
      * Write out read aligned to haplotype to the BAM file
      *
      * @param read the read we want to write aligned to the reference genome, must not be null
      */
-    protected void writeReadAgainstHaplotype(final GATKRead read) {
+    private void writeReadAgainstHaplotype(final GATKRead read) {
         Utils.nonNull(read, "read cannot be null");
         output.add(read);
     }
@@ -156,7 +161,7 @@ protected void writeReadAgainstHaplotype(final GATKRead read) {
      *                       be null
      * @param paddedReferenceLoc the genome loc of the padded reference, must not be null
      */
-    protected void writeHaplotypesAsReads(final Collection<Haplotype> haplotypes,
+    private void writeHaplotypesAsReads(final Collection<Haplotype> haplotypes,
                                           final Set<Haplotype> bestHaplotypes,
                                           final Locatable paddedReferenceLoc) {
         Utils.nonNull(haplotypes, "haplotypes cannot be null");