Adding argument to GenotypeGVCFs to keep only RAW_GT_COUNT

src/main/java/org/broadinstitute/hellbender/cmdline/GATKPlugin/DefaultGATKVariantAnnotationArgumentCollection.java

@@ -1,6 +1,5 @@
 package org.broadinstitute.hellbender.cmdline.GATKPlugin;
 
-import com.google.common.collect.Lists;
 import org.broadinstitute.barclay.argparser.Advanced;
 import org.broadinstitute.barclay.argparser.Argument;
 import org.broadinstitute.hellbender.cmdline.StandardArgumentDefinitions;

src/main/java/org/broadinstitute/hellbender/cmdline/GATKPlugin/GATKAnnotationPluginDescriptor.java

@@ -16,7 +16,6 @@
 import org.broadinstitute.hellbender.utils.config.ConfigFactory;
 import org.broadinstitute.hellbender.utils.config.GATKConfig;
 
-import java.io.File;
 import java.lang.reflect.Modifier;
 import java.util.*;
 import java.util.stream.Collectors;
@@ -485,6 +484,21 @@ public List<Annotation> getResolvedInstances() {
         return resolvedInstances;
     }
 
+    /**
+     * Returns a map of the String to Annotations only in the resolved instances. Should only be used if
+     * getResolvedInstances() is not sufficient.
+     *
+     * @return a Map of Strings to Annotations of resolved instances
+     */
+    public Map<String, Annotation> getResolvedInstancesMap() {
+        if (resolvedInstances == null) {
+            getResolvedInstances();
+        }
+        Map<String, Annotation> resolvedInstancesMap = new HashMap<>();
+        allDiscoveredAnnotations.keySet().stream().filter(s -> resolvedInstances.contains(allDiscoveredAnnotations.get(s))).
+                forEach(s -> resolvedInstancesMap.put(s, allDiscoveredAnnotations.get(s)));
+        return(resolvedInstancesMap);
+    }
 
     /**
      * Return the class representing the instance of the plugin specified by {@code pluginName}
@@ -496,5 +510,4 @@ public List<Annotation> getResolvedInstances() {
     public Class<?> getClassForPluginHelp(final String pluginName) {
         return allDiscoveredAnnotations.containsKey(pluginName) ? allDiscoveredAnnotations.get(pluginName).getClass() : null;
     }
-
 }

src/main/java/org/broadinstitute/hellbender/tools/walkers/GenotypeGVCFs.java

@@ -9,6 +9,8 @@
 import htsjdk.variant.vcf.VCFHeaderLine;
 import org.broadinstitute.barclay.argparser.*;
 import org.broadinstitute.barclay.help.DocumentedFeature;
+import org.broadinstitute.hellbender.cmdline.GATKPlugin.GATKAnnotationPluginDescriptor;
+import org.broadinstitute.hellbender.cmdline.GATKPlugin.GATKReadFilterPluginDescriptor;
 import org.broadinstitute.hellbender.cmdline.StandardArgumentDefinitions;
 import org.broadinstitute.hellbender.cmdline.argumentcollections.DbsnpArgumentCollection;
 import org.broadinstitute.hellbender.cmdline.programgroups.ShortVariantDiscoveryProgramGroup;
@@ -28,6 +30,7 @@
 import org.broadinstitute.hellbender.tools.walkers.mutect.M2ArgumentCollection;
 import org.broadinstitute.hellbender.utils.*;
 import org.broadinstitute.hellbender.utils.variant.GATKVariantContextUtils;
+import org.broadinstitute.hellbender.tools.walkers.annotator.allelespecific.ReducibleAnnotation;
 
 import java.util.*;
 import java.util.stream.Collectors;
@@ -132,17 +135,17 @@ public final class GenotypeGVCFs extends VariantLocusWalker {
     doc = "LOD threshold to emit variant to VCF.")
     protected double tlodThreshold = 3.5;  //allow for some lower quality variants
 
-
     /**
      * Margin of error in allele fraction to consider a somatic variant homoplasmic, i.e. if there is less than a 0.1% reference allele fraction, those reads are likely errors
      */
     @Argument(fullName=CombineGVCFs.ALLELE_FRACTION_DELTA_LONG_NAME, doc = "Margin of error in allele fraction to consider a somatic variant homoplasmic")
     protected double afTolerance = 1e-3;  //based on Q30 as a "good" base quality score
 
     /**
-     * If specified, keep the combined raw annotations (e.g. AS_SB_TABLE) after genotyping.  This is applicable to Allele-Specific annotations
+     * If specified, keep all the combined raw annotations (e.g. AS_SB_TABLE) after genotyping.  This is applicable to Allele-Specific annotations. See {@link ReducibleAnnotation}
      */
-    @Argument(fullName=KEEP_COMBINED_LONG_NAME, shortName = KEEP_COMBINED_SHORT_NAME, doc = "If specified, keep the combined raw annotations")
+    @Argument(fullName=KEEP_COMBINED_LONG_NAME, shortName = KEEP_COMBINED_SHORT_NAME, doc = "If specified, keep the combined raw annotations",
+            mutex = {GenotypeGVCFsAnnotationArgumentCollection.KEEP_SPECIFIED_RAW_COMBINED_ANNOTATION_LONG_NAME})
     protected boolean keepCombined = false;
 
     @ArgumentCollection
@@ -172,6 +175,9 @@ public final class GenotypeGVCFs extends VariantLocusWalker {
     @ArgumentCollection
     private final DbsnpArgumentCollection dbsnp = new DbsnpArgumentCollection();
 
+    // @ArgumentCollection deliberately omitted since this is passed to the annotation plugin
+    final GenotypeGVCFsAnnotationArgumentCollection genotypeGVCFsAnnotationArgs = new GenotypeGVCFsAnnotationArgumentCollection();
+
     // the annotation engine
     private VariantAnnotatorEngine annotationEngine;
 
@@ -221,6 +227,16 @@ protected GenomicsDBOptions getGenomicsDBOptions() {
     @Override
     public boolean useVariantAnnotations() { return true;}
 
+    @Override
+    public List<? extends CommandLinePluginDescriptor<?>> getPluginDescriptors() {
+        GATKReadFilterPluginDescriptor readFilterDescriptor = new GATKReadFilterPluginDescriptor(getDefaultReadFilters());
+        return useVariantAnnotations()?
+                Arrays.asList(readFilterDescriptor, new GATKAnnotationPluginDescriptor(
+                        genotypeGVCFsAnnotationArgs,
+                        getDefaultVariantAnnotations(), getDefaultVariantAnnotationGroups())):
+                Collections.singletonList(readFilterDescriptor);
+    }
+
     @Override
     public List<Class<? extends Annotation>> getDefaultVariantAnnotationGroups() {
         return Arrays.asList(StandardAnnotation.class);
@@ -261,8 +277,9 @@ public void onTraversalStart() {
         intervals = hasUserSuppliedIntervals() ? intervalArgumentCollection.getIntervals(getBestAvailableSequenceDictionary()) :
                 Collections.emptyList();
 
-        Collection<Annotation>  variantAnnotations = makeVariantAnnotations();
-        annotationEngine = new VariantAnnotatorEngine(variantAnnotations, dbsnp.dbsnp, Collections.emptyList(), false, keepCombined);
+        final Collection<Annotation>  variantAnnotations = makeVariantAnnotations();
+        final Set<Annotation> annotationsToKeep = getAnnotationsToKeep();
+        annotationEngine = new VariantAnnotatorEngine(variantAnnotations, dbsnp.dbsnp, Collections.emptyList(), false, keepCombined, annotationsToKeep);
 
         merger = new ReferenceConfidenceVariantContextMerger(annotationEngine, getHeaderForVariants(), somaticInput, false, true);
 
@@ -279,6 +296,13 @@ public void onTraversalStart() {
 
     }
 
+    private Set<Annotation> getAnnotationsToKeep() {
+        final GATKAnnotationPluginDescriptor pluginDescriptor = getCommandLineParser().getPluginDescriptor(GATKAnnotationPluginDescriptor.class);
+        final List<String> annotationStringsToKeep = genotypeGVCFsAnnotationArgs.getKeepSpecifiedCombinedAnnotationNames();
+        final Map<String, Annotation> resolvedInstancesMap = pluginDescriptor.getResolvedInstancesMap();
+        return annotationStringsToKeep.stream().map(resolvedInstancesMap::get).collect(Collectors.toSet());
+    }
+
     @Override
     public void apply(final Locatable loc, List<VariantContext> variants, ReadsContext reads, ReferenceContext ref, FeatureContext features) {
 

src/main/java/org/broadinstitute/hellbender/tools/walkers/GenotypeGVCFsAnnotationArgumentCollection.java

@@ -0,0 +1,27 @@
+package org.broadinstitute.hellbender.tools.walkers;
+
+import org.broadinstitute.barclay.argparser.Argument;
+import org.broadinstitute.hellbender.cmdline.GATKPlugin.DefaultGATKVariantAnnotationArgumentCollection;
+import org.broadinstitute.hellbender.tools.walkers.annotator.allelespecific.ReducibleAnnotation;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+public class GenotypeGVCFsAnnotationArgumentCollection extends DefaultGATKVariantAnnotationArgumentCollection {
+    private static final long serialVersionUID = 1L;
+
+    public static final String KEEP_SPECIFIED_RAW_COMBINED_ANNOTATION_LONG_NAME = "keep-specific-combined-raw-annotation";
+    public static final String KEEP_SPECIFIED_RAW_COMBINED_ANNOTATION_SHORT_NAME = "keep-specific-combined";
+
+    /**
+     * Keep only the specific combined raw annotations specified. Cannot be used with --keep-combined-raw-annotations which saves all raw annotations.
+     * Duplicate values will be ignored. See {@link ReducibleAnnotation} for more information on raw annotations.
+     */
+    @Argument(fullName= KEEP_SPECIFIED_RAW_COMBINED_ANNOTATION_LONG_NAME, shortName = KEEP_SPECIFIED_RAW_COMBINED_ANNOTATION_SHORT_NAME, optional = true,
+            mutex = {GenotypeGVCFs.KEEP_COMBINED_LONG_NAME},
+            doc="Keep only the specific combined raw annotations specified (removing the other raw annotations). Duplicate values will be ignored.")
+    protected List<String> keepSpecifiedCombined = new ArrayList<>();
+
+    public List<String> getKeepSpecifiedCombinedAnnotationNames() {return Collections.unmodifiableList(keepSpecifiedCombined);}
+}

src/main/java/org/broadinstitute/hellbender/tools/walkers/annotator/VariantAnnotatorEngine.java

@@ -4,12 +4,14 @@
 import htsjdk.variant.vcf.*;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.broadinstitute.hellbender.cmdline.StandardArgumentDefinitions;
 import org.broadinstitute.hellbender.engine.FeatureContext;
 import org.broadinstitute.hellbender.engine.FeatureDataSource;
 import org.broadinstitute.hellbender.engine.FeatureInput;
 import org.broadinstitute.hellbender.engine.ReferenceContext;
 import org.broadinstitute.hellbender.exceptions.GATKException;
 import org.broadinstitute.hellbender.exceptions.UserException;
+import org.broadinstitute.hellbender.tools.walkers.GenotypeGVCFsAnnotationArgumentCollection;
 import org.broadinstitute.hellbender.tools.walkers.annotator.allelespecific.ReducibleAnnotation;
 import org.broadinstitute.hellbender.tools.walkers.annotator.allelespecific.ReducibleAnnotationData;
 import org.broadinstitute.hellbender.utils.Utils;
@@ -43,6 +45,7 @@ public final class VariantAnnotatorEngine {
     private boolean expressionAlleleConcordance;
     private final boolean useRawAnnotations;
     private final boolean keepRawCombinedAnnotations;
+    private final List<String> rawAnnotationsToKeep;
 
     private final static Logger logger = LogManager.getLogger(VariantAnnotatorEngine.class);
     private final static OneShotLogger jumboAnnotationsLogger = new OneShotLogger(VariantAnnotatorEngine.class);
@@ -59,17 +62,21 @@ public final class VariantAnnotatorEngine {
      * @param useRaw When this is set to true, the annotation engine will call {@link ReducibleAnnotation#annotateRawData(ReferenceContext, VariantContext, AlleleLikelihoods)}
 *               on annotations that extend {@link ReducibleAnnotation}, instead of {@link InfoFieldAnnotation#annotate(ReferenceContext, VariantContext, AlleleLikelihoods)},
      * @param keepCombined If true, retain the combined raw annotation values instead of removing them after finalizing
+     * @param rawAnnotationsToKeep List of raw annotations to keep even when others are removed
      */
     public VariantAnnotatorEngine(final Collection<Annotation> annotationList,
                                   final FeatureInput<VariantContext> dbSNPInput,
                                   final List<FeatureInput<VariantContext>> featureInputs,
                                   final boolean useRaw,
-                                  boolean keepCombined){
+                                  final boolean keepCombined,
+                                  final Collection<Annotation> rawAnnotationsToKeep){
         Utils.nonNull(featureInputs, "comparisonFeatureInputs is null");
         infoAnnotations = new ArrayList<>();
         genotypeAnnotations = new ArrayList<>();
         jumboInfoAnnotations = new ArrayList<>();
         jumboGenotypeAnnotations = new ArrayList<>();
+        final List<String> variantAnnotationKeys = new ArrayList<>();
+        final List<String> rawVariantAnnotationKeysToKeep = new ArrayList<>();
         for (Annotation annot : annotationList) {
             if (annot instanceof InfoFieldAnnotation) {
                 infoAnnotations.add((InfoFieldAnnotation) annot);
@@ -82,11 +89,19 @@ public VariantAnnotatorEngine(final Collection<Annotation> annotationList,
             } else {
                 throw new GATKException.ShouldNeverReachHereException("Unexpected annotation type: " + annot.getClass().getName());
             }
+            variantAnnotationKeys.addAll(((VariantAnnotation) annot).getKeyNames());
         }
         variantOverlapAnnotator = initializeOverlapAnnotator(dbSNPInput, featureInputs);
         reducibleKeys = new LinkedHashSet<>();
         useRawAnnotations = useRaw;
         keepRawCombinedAnnotations = keepCombined;
+        for (final Annotation rawAnnot : rawAnnotationsToKeep) {
+            if (!annotationList.contains(rawAnnot)) {
+                throw new UserException("Requested --" + GenotypeGVCFsAnnotationArgumentCollection.KEEP_SPECIFIED_RAW_COMBINED_ANNOTATION_LONG_NAME + ": " + rawAnnot + " is not available. Add requested annotation with --" + StandardArgumentDefinitions.ANNOTATION_LONG_NAME + ".");
+            }
+            rawVariantAnnotationKeysToKeep.addAll(((VariantAnnotation) rawAnnot).getKeyNames());
+        }
+        this.rawAnnotationsToKeep = rawVariantAnnotationKeysToKeep;
         for (InfoFieldAnnotation annot : infoAnnotations) {
             if (annot instanceof ReducibleAnnotation) {
                 for (final String rawKey : ((ReducibleAnnotation) annot).getRawKeyNames()) {
@@ -96,6 +111,14 @@ public VariantAnnotatorEngine(final Collection<Annotation> annotationList,
         }
     }
 
+    public VariantAnnotatorEngine(final Collection<Annotation> annotationList,
+                                  final FeatureInput<VariantContext> dbSNPInput,
+                                  final List<FeatureInput<VariantContext>> featureInputs,
+                                  final boolean useRaw,
+                                  boolean keepCombined){
+        this(annotationList, dbSNPInput, featureInputs, useRaw, keepCombined, Collections.emptyList());
+    }
+
     private VariantOverlapAnnotator initializeOverlapAnnotator(final FeatureInput<VariantContext> dbSNPInput, final List<FeatureInput<VariantContext>> featureInputs) {
         final Map<FeatureInput<VariantContext>, String> overlaps = new LinkedHashMap<>();
         for ( final FeatureInput<VariantContext> fi : featureInputs) {
@@ -253,6 +276,14 @@ public Map<String, Object> combineAnnotations(final List<Allele> allelesList, Ma
     public VariantContext finalizeAnnotations(VariantContext vc, VariantContext originalVC) {
         final Map<String, Object> variantAnnotations = new LinkedHashMap<>(vc.getAttributes());
 
+        //save annotations that have been requested to be kept
+        final Map<String, Object> savedRawAnnotations = new LinkedHashMap<>();
+        for(final String rawAnnot : rawAnnotationsToKeep) {
+            if (variantAnnotations.containsKey(rawAnnot)) {
+                savedRawAnnotations.put(rawAnnot, variantAnnotations.get(rawAnnot));
+            }
+        }
+
         // go through all the requested info annotationTypes
         for (final InfoFieldAnnotation annotationType : infoAnnotations) {
             if (annotationType instanceof ReducibleAnnotation) {
@@ -280,6 +311,8 @@ public VariantContext finalizeAnnotations(VariantContext vc, VariantContext orig
             variantAnnotations.remove(GATKVCFConstants.VARIANT_DEPTH_KEY);
             variantAnnotations.remove(GATKVCFConstants.RAW_GENOTYPE_COUNT_KEY);
         }
+        //add back raw annotations that have specifically been requested to keep
+        variantAnnotations.putAll(savedRawAnnotations);
 
         // generate a new annotated VC
         final VariantContextBuilder builder = new VariantContextBuilder(vc).attributes(variantAnnotations);

src/main/java/org/broadinstitute/hellbender/tools/walkers/annotator/allelespecific/ReducibleAnnotation.java

@@ -3,21 +3,21 @@
 import htsjdk.variant.variantcontext.Allele;
 import htsjdk.variant.variantcontext.VariantContext;
 import htsjdk.variant.vcf.VCFCompoundHeaderLine;
-import htsjdk.variant.vcf.VCFHeaderLine;
-import htsjdk.variant.vcf.VCFInfoHeaderLine;
 import org.broadinstitute.hellbender.engine.ReferenceContext;
 import org.broadinstitute.hellbender.tools.walkers.annotator.Annotation;
 import org.broadinstitute.hellbender.utils.genotyper.AlleleLikelihoods;
 import org.broadinstitute.hellbender.utils.read.GATKRead;
 import org.broadinstitute.hellbender.utils.variant.GATKVCFHeaderLines;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 
 /**
  * An interface for annotations that are calculated using raw data across samples, rather than the median (or median of median) of samples values
+ * The Raw annotation keeps some summary (one example might be a histogram of the raw values for each sample) of the individual sample (or allele)
+ * level annotation. As the annotations are combined across multiple samples the raw annotation continues to contain individual values while
+ * the final reduced annotation will typically be a summary statistic from these raw values.
  *
  */
 public interface ReducibleAnnotation extends Annotation {

src/main/java/org/broadinstitute/hellbender/utils/variant/GATKVCFHeaderLines.java

@@ -159,7 +159,7 @@ public static VCFFormatHeaderLine getEquivalentFormatHeaderLine(final String inf
         addInfoLine(new VCFInfoHeaderLine(INBREEDING_COEFFICIENT_KEY, 1, VCFHeaderLineType.Float, "Inbreeding coefficient as estimated from the genotype likelihoods per-sample when compared against the Hardy-Weinberg expectation"));
         addInfoLine(new VCFInfoHeaderLine(AS_INBREEDING_COEFFICIENT_KEY, VCFHeaderLineCount.A, VCFHeaderLineType.Float, "Allele-specific inbreeding coefficient as estimated from the genotype likelihoods per-sample when compared against the Hardy-Weinberg expectation"));
         addInfoLine(new VCFInfoHeaderLine(EXCESS_HET_KEY, 1, VCFHeaderLineType.Float, "Phred-scaled p-value for exact test of excess heterozygosity"));
-        addInfoLine(new VCFInfoHeaderLine(RAW_GENOTYPE_COUNT_KEY, 3, VCFHeaderLineType.Integer, "Counts of genotypes w.r.t. the reference allele: 0/0, 0/*, */*, i.e. all alts lumped together; for use in calculating excess heterozygosity"));
+        addInfoLine(new VCFInfoHeaderLine(RAW_GENOTYPE_COUNT_KEY, 3, VCFHeaderLineType.Integer, "Counts of genotypes w.r.t. the reference allele in the following order: 0/0, 0/*, */*, i.e. all alts lumped together; for use in calculating excess heterozygosity"));
         addInfoLine(new VCFInfoHeaderLine(LIKELIHOOD_RANK_SUM_KEY, 1, VCFHeaderLineType.Float, "Z-score from Wilcoxon rank sum test of Alt Vs. Ref haplotype likelihoods"));
         addInfoLine(new VCFInfoHeaderLine(MAP_QUAL_RANK_SUM_KEY, 1, VCFHeaderLineType.Float, "Z-score From Wilcoxon rank sum test of Alt vs. Ref read mapping qualities"));
         addInfoLine(new VCFInfoHeaderLine(AS_MAP_QUAL_RANK_SUM_KEY, VCFHeaderLineCount.A, VCFHeaderLineType.Float, "allele specific Z-score From Wilcoxon rank sum test of each Alt vs. Ref read mapping qualities"));

src/test/java/org/broadinstitute/hellbender/tools/walkers/GenotypeGVCFsIntegrationTest.java

@@ -43,8 +43,6 @@
 import java.security.NoSuchAlgorithmException;
 import java.util.*;
 import java.util.function.BiConsumer;
-import java.util.function.Consumer;
-import java.util.function.IntUnaryOperator;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -930,7 +928,7 @@ public void testRawGtCountAnnotation() {
         args.addReference(b37_reference_20_21)
                 .addVCF(reblockedGVCF)
                 .addOutput(output)
-                .add("keep-combined-raw-annotations", true)
+                .add(GenotypeGVCFsAnnotationArgumentCollection.KEEP_SPECIFIED_RAW_COMBINED_ANNOTATION_LONG_NAME, "RawGtCount")
                 .add("A", "RawGtCount");
         runCommandLine(args);
 
@@ -942,6 +940,7 @@ public void testRawGtCountAnnotation() {
                 Assert.assertEquals(rawGtCount.get(0), ".");
                 Assert.assertEquals(rawGtCount.get(1), "2");
                 Assert.assertEquals(rawGtCount.get(2), "0");
+                Assert.assertFalse(vc.getAttributes().containsKey(GATKVCFConstants.RAW_MAPPING_QUALITY_WITH_DEPTH_KEY));
             }
 
         }