From 3d4bf273de8e350d27b92ff8c5411a8e4485173e Mon Sep 17 00:00:00 2001
From: Mauricio Carneiro <carneiro@broadinstitute.org>
Date: Mon, 2 Jan 2012 16:17:57 -0500
Subject: [PATCH] Added getPileupForReadGroups to ReadBackPileup

   * returns a pileup for all the read groups provided.
   * saves us from multiple calls to getPileup (which is very inefficient)
---
 .../pileup/AbstractReadBackedPileup.java      | 36 +++++++++++++++++++
 .../sting/utils/pileup/ReadBackedPileup.java  |  8 +++++
 2 files changed, 44 insertions(+)

diff --git a/public/java/src/org/broadinstitute/sting/utils/pileup/AbstractReadBackedPileup.java b/public/java/src/org/broadinstitute/sting/utils/pileup/AbstractReadBackedPileup.java
index 18051ce92..586b86490 100644
--- a/public/java/src/org/broadinstitute/sting/utils/pileup/AbstractReadBackedPileup.java
+++ b/public/java/src/org/broadinstitute/sting/utils/pileup/AbstractReadBackedPileup.java
@@ -526,6 +526,42 @@ public abstract class AbstractReadBackedPileup<RBP extends AbstractReadBackedPil
         }
     }
 
+    /**
+     * Gets the pileup for a set of read groups.  Horrendously inefficient at this point.
+     * @param rgSet List of identifiers for the read groups.
+     * @return A read-backed pileup containing only the reads in the given read groups.
+     */
+    @Override
+    public RBP getPileupForReadGroups(final HashSet<String> rgSet) {
+        if(pileupElementTracker instanceof PerSamplePileupElementTracker) {
+            PerSamplePileupElementTracker<PE> tracker = (PerSamplePileupElementTracker<PE>)pileupElementTracker;
+            PerSamplePileupElementTracker<PE> filteredTracker = new PerSamplePileupElementTracker<PE>();
+
+            for(final String sample: tracker.getSamples()) {
+                PileupElementTracker<PE> perSampleElements = tracker.getElements(sample);
+                AbstractReadBackedPileup<RBP,PE> pileup = createNewPileup(loc,perSampleElements).getPileupForReadGroups(rgSet);
+                if(pileup != null)
+                    filteredTracker.addElements(sample,pileup.pileupElementTracker);
+            }
+            return filteredTracker.size()>0 ? (RBP)createNewPileup(loc,filteredTracker) : null;
+        }
+        else {
+            UnifiedPileupElementTracker<PE> filteredTracker = new UnifiedPileupElementTracker<PE>();
+            for(PE p: pileupElementTracker) {
+                GATKSAMRecord read = p.getRead();
+                if(rgSet != null && !rgSet.isEmpty()) {
+                    if(read.getReadGroup() != null && rgSet.contains(read.getReadGroup().getReadGroupId()))
+                        filteredTracker.add(p);
+                }
+                else {
+                    if(read.getReadGroup() == null || read.getReadGroup().getReadGroupId() == null)
+                        filteredTracker.add(p);
+                }
+            }
+            return filteredTracker.size()>0 ? (RBP)createNewPileup(loc,filteredTracker) : null;
+        }
+    }
+
     @Override
     public RBP getPileupForLane(String laneID) {
         if(pileupElementTracker instanceof PerSamplePileupElementTracker) {
diff --git a/public/java/src/org/broadinstitute/sting/utils/pileup/ReadBackedPileup.java b/public/java/src/org/broadinstitute/sting/utils/pileup/ReadBackedPileup.java
index 02767df7c..ccd9d509f 100644
--- a/public/java/src/org/broadinstitute/sting/utils/pileup/ReadBackedPileup.java
+++ b/public/java/src/org/broadinstitute/sting/utils/pileup/ReadBackedPileup.java
@@ -30,6 +30,7 @@ import org.broadinstitute.sting.utils.fragments.FragmentCollection;
 import org.broadinstitute.sting.utils.sam.GATKSAMRecord;
 
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 
 /**
@@ -129,6 +130,13 @@ public interface ReadBackedPileup extends Iterable<PileupElement>, HasGenomeLoca
      */
     public ReadBackedPileup getPileupForReadGroup(String readGroupId);
 
+    /**
+     * Gets all the reads associated with a given read groups.
+     * @param rgSet Set of identifiers for the read group.
+     * @return A pileup containing only the reads in the given read groups.
+     */
+    public ReadBackedPileup getPileupForReadGroups(final HashSet<String> rgSet);
+    
     /**
      * Gets all reads in a given lane id. (Lane ID is the read group
      * id stripped of the last .XX sample identifier added by the GATK).