1 package org.broadinstitute.hellbender.tools.walkers.annotator;
2 
3 import com.google.common.collect.ImmutableMap;
4 import htsjdk.variant.variantcontext.Allele;
5 import htsjdk.variant.variantcontext.VariantContext;
6 import htsjdk.variant.vcf.VCFHeaderLineCount;
7 import htsjdk.variant.vcf.VCFHeaderLineType;
8 import htsjdk.variant.vcf.VCFInfoHeaderLine;
9 import org.broadinstitute.hellbender.engine.ReferenceContext;
10 import org.broadinstitute.hellbender.utils.QualityUtils;
11 import org.broadinstitute.hellbender.utils.Utils;
12 import org.broadinstitute.hellbender.utils.genotyper.AlleleLikelihoods;
13 import org.broadinstitute.hellbender.utils.read.GATKRead;
14 
15 import java.util.*;
16 import java.util.stream.Collectors;
17 
18 /**
19  * Apply an annotation based on aggregation data from all reads supporting each allele.
20  */
21 public abstract class PerAlleleAnnotation extends InfoFieldAnnotation{
22 
23     /**
24      * Calculate annotations for each allele based on given VariantContext and likelihoods for a given genotype's sample
25      * and add the annotations to the GenotypeBuilder.  By default annotations are only calculated for alt alleles but
26      * implementations may override the {@code includeRefAllele()} method.  See parent class docs in {@link GenotypeAnnotation}.
27      */
annotate(final ReferenceContext ref, final VariantContext vc, final AlleleLikelihoods<GATKRead, Allele> likelihoods)28     public Map<String, Object> annotate(final ReferenceContext ref,
29                          final VariantContext vc,
30                          final AlleleLikelihoods<GATKRead, Allele> likelihoods) {
31         Utils.nonNull(vc);
32         if ( likelihoods == null ) {
33             return Collections.emptyMap();
34         }
35 
36         final Map<Allele, List<Integer>> values = likelihoods.alleles().stream()
37                 .collect(Collectors.toMap(a -> a, a -> new ArrayList<>()));
38 
39         Utils.stream(likelihoods.bestAllelesBreakingTies())
40                 .filter(ba -> ba.isInformative() && isUsableRead(ba.evidence))
41                 .forEach(ba -> getValueForRead(ba.evidence, vc).ifPresent(v -> values.get(ba.allele).add(v)));
42 
43         final int[] statistics = vc.getAlleles().stream().filter(this::includeAllele).mapToInt(a -> aggregate(values.get(a))).toArray();
44         return ImmutableMap.of(getVcfKey(), statistics);
45     }
46 
includeAllele(final Allele allele)47     private boolean includeAllele(final Allele allele) {
48         return allele.isNonReference() || includeRefAllele();
49     }
50 
51     // this is false by default but implementations may wish to override
includeRefAllele()52     protected boolean includeRefAllele() { return false; }
53 
isUsableRead(final GATKRead read)54     private static boolean isUsableRead(final GATKRead read) {
55         return read.getMappingQuality() != 0 && read.getMappingQuality() != QualityUtils.MAPPING_QUALITY_UNAVAILABLE;
56     }
57 
58     @Override
getDescriptions()59     public List<VCFInfoHeaderLine> getDescriptions() {
60         return Arrays.asList(new VCFInfoHeaderLine(getVcfKey(), includeRefAllele() ? VCFHeaderLineCount.R : VCFHeaderLineCount.A, VCFHeaderLineType.Integer, getDescription()));
61     }
62 
63     @Override
getKeyNames()64     public List<String> getKeyNames() { return Arrays.asList(getVcfKey()); }
65 
getValueForRead(final GATKRead read, final VariantContext vc)66     protected abstract OptionalInt getValueForRead(final GATKRead read, final VariantContext vc);
aggregate(final List<Integer> values)67     protected abstract int aggregate(final List<Integer> values);
getVcfKey()68     protected abstract String getVcfKey();
getDescription()69     protected abstract String getDescription();
70 }
71