Java tutorial
/* * Copyright (C) 2015 Google Inc. * * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except * in compliance with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software distributed under the License * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express * or implied. See the License for the specific language governing permissions and limitations under * the License. */ package com.google.cloud.genomics.dataflow.pipelines; import com.google.api.client.util.Strings; import com.google.api.services.genomics.Genomics; import com.google.api.services.genomics.Genomics.Annotationsets; import com.google.api.services.genomics.model.Annotation; import com.google.api.services.genomics.model.AnnotationSet; import com.google.api.services.genomics.model.BatchCreateAnnotationsRequest; import com.google.api.services.genomics.model.Position; import com.google.cloud.dataflow.sdk.Pipeline; import com.google.cloud.dataflow.sdk.options.Default; import com.google.cloud.dataflow.sdk.options.Description; import com.google.cloud.dataflow.sdk.options.PipelineOptionsFactory; import com.google.cloud.dataflow.sdk.options.Validation; import com.google.cloud.dataflow.sdk.transforms.ApproximateQuantiles; import com.google.cloud.dataflow.sdk.transforms.Combine; import com.google.cloud.dataflow.sdk.transforms.Create; import com.google.cloud.dataflow.sdk.transforms.DoFn; import com.google.cloud.dataflow.sdk.transforms.GroupByKey; import com.google.cloud.dataflow.sdk.transforms.PTransform; import com.google.cloud.dataflow.sdk.transforms.ParDo; import com.google.cloud.dataflow.sdk.transforms.SerializableFunction; import com.google.cloud.dataflow.sdk.values.KV; import com.google.cloud.dataflow.sdk.values.PCollection; import com.google.cloud.genomics.dataflow.coders.GenericJsonCoder; import com.google.cloud.genomics.dataflow.model.PosRgsMq; import com.google.cloud.genomics.dataflow.readers.ReadGroupStreamer; import com.google.cloud.genomics.dataflow.utils.GenomicsOptions; import com.google.cloud.genomics.dataflow.utils.ShardOptions; import com.google.cloud.genomics.utils.Contig; import com.google.cloud.genomics.utils.GenomicsFactory; import com.google.cloud.genomics.utils.GenomicsUtils; import com.google.cloud.genomics.utils.OfflineAuth; import com.google.cloud.genomics.utils.RetryPolicy; import com.google.cloud.genomics.utils.ShardBoundary; import com.google.cloud.genomics.utils.ShardUtils.SexChromosomeFilter; import com.google.common.collect.Lists; import com.google.genomics.v1.CigarUnit; import com.google.genomics.v1.Read; import java.io.IOException; import java.security.GeneralSecurityException; import java.util.HashMap; import java.util.List; /** * Class for calculating read depth coverage for a given data set. * * For each "bucket" in the given input references, this computes the average coverage (rounded to * six decimal places) across the bucket that 10%, 20%, 30%, etc. of the input ReadGroupsSets have * for each mapping quality of the reads (<10:Low(L), 10-29:Medium(M), >=30:High(H)) as well as * these same percentiles of ReadGroupSets for all reads regardless of mapping quality (Mapping * quality All(A)). * * There is also the option to change the number of quantiles accordingly (numQuantiles = 5 would * give you the minimum ReadGroupSet mean coverage for each and across all mapping qualities, the * 25th, 50th, and 75th percentiles, and the maximum of these values). * * See * http://googlegenomics.readthedocs.org/en/latest/use_cases/analyze_reads/calculate_coverage.html * for running instructions. */ public class CalculateCoverage { /** * Options required to run this pipeline. * * When specifying references, if you specify a nonzero starting point for any given reference, * reads that start before that reference, even if they overlap with the reference, will NOT be * counted. */ public static interface Options extends ShardOptions { @Description("A comma delimited list of the IDs of the Google Genomics ReadGroupSets this " + "pipeline is working with. Default (empty) indicates all ReadGroupSets in InputDatasetId." + " This or InputDatasetId must be set. InputDatasetId overrides ReadGroupSetIds " + "(if InputDatasetId is set, this field will be ignored). All of the referenceSetIds for" + " all ReadGroupSets in this list must be the same for the purposes of setting the" + " referenceSetId of the output AnnotationSet.") @Default.String("") String getReadGroupSetIds(); void setReadGroupSetIds(String readGroupSetId); @Description("The ID of the Google Genomics Dataset that the pipeline will get its input reads" + " from. Default (empty) means to use ReadGroupSetIds instead. This or ReadGroupSetIds" + " must be set. InputDatasetId overrides ReadGroupSetIds (if this field is set, " + "ReadGroupSetIds will be ignored). All of the referenceSetIds for all ReadGroupSets in" + " this Dataset must be the same for the purposes of setting the referenceSetId" + " of the output AnnotationSet.") @Default.String("") String getInputDatasetId(); void setInputDatasetId(String inputDatasetId); @Validation.Required @Description("The ID of the Google Genomics Dataset that the output AnnotationSet will be " + "posted to.") @Default.String("") String getOutputDatasetId(); void setOutputDatasetId(String outputDatasetId); @Description("The bucket width you would like to calculate coverage for. Default is 2048. " + "Buckets are non-overlapping. If bucket width does not divide into the reference " + "size, then the remainder will be a smaller bucket at the end of the reference.") @Default.Integer(2048) int getBucketWidth(); void setBucketWidth(int bucketWidth); @Description("The number of quantiles you would like to calculate for in the output. Default" + "is 11 (minimum value for a particular grouping, 10-90 percentiles, and the maximum" + " value). The InputDataset or list of ReadGroupSetIds specified must have an amount of " + "ReadGroupSetIds greater than or equal to the number of quantiles that you are asking " + "for.") @Default.Integer(11) int getNumQuantiles(); void setNumQuantiles(int numQuantiles); @Description("This provides the name for the AnnotationSet. Default (empty) will set the " + "name to the input References. For more information on AnnotationSets, please visit: " + "https://cloud.google.com/genomics/v1beta2/reference/annotationSets#resource") @Default.String("") String getAnnotationSetName(); void setAnnotationSetName(String name); } // Tip: Use the API explorer to test which fields to include in partial responses. // https://developers.google.com/apis-explorer/#p/genomics/v1/genomics.reads.stream?fields=alignments(alignedSequence%252Cid)&_h=2&resource=%257B%250A++%2522readGroupSetId%2522%253A+%2522CMvnhpKTFhD3he72j4KZuyc%2522%252C%250A++%2522referenceName%2522%253A+%2522chr17%2522%252C%250A++%2522start%2522%253A+%252241196311%2522%252C%250A++%2522end%2522%253A+%252241196312%2522%250A%257D& private static final String READ_FIELDS = "alignments(alignment,readGroupSetId)"; private static Options options; private static Pipeline p; private static OfflineAuth auth; public static void main(String[] args) throws GeneralSecurityException, IOException { // Register the options so that they show up via --help PipelineOptionsFactory.register(Options.class); options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class); auth = GenomicsOptions.Methods.getGenomicsAuth(options); p = Pipeline.create(options); p.getCoderRegistry().setFallbackCoderProvider(GenericJsonCoder.PROVIDER); if (options.getInputDatasetId().isEmpty() && options.getReadGroupSetIds().isEmpty()) { throw new IllegalArgumentException("InputDatasetId or ReadGroupSetIds must be specified"); } List<String> rgsIds; if (options.getInputDatasetId().isEmpty()) { rgsIds = Lists.newArrayList(options.getReadGroupSetIds().split(",")); } else { rgsIds = GenomicsUtils.getReadGroupSetIds(options.getInputDatasetId(), auth); } if (rgsIds.size() < options.getNumQuantiles()) { throw new IllegalArgumentException("Number of ReadGroupSets must be greater than or equal to" + " the number of requested quantiles."); } // Grab one ReferenceSetId to be used within the pipeline to confirm that all ReadGroupSets // are associated with the same ReferenceSet. String referenceSetId = GenomicsUtils.getReferenceSetId(rgsIds.get(0), auth); if (Strings.isNullOrEmpty(referenceSetId)) { throw new IllegalArgumentException("No ReferenceSetId associated with ReadGroupSetId " + rgsIds.get(0) + ". All ReadGroupSets in given input must have an associated ReferenceSet."); } // Create our destination AnnotationSet for the associated ReferenceSet. AnnotationSet annotationSet = createAnnotationSet(referenceSetId); PCollection<Read> reads = p.begin().apply(Create.of(rgsIds)) .apply(ParDo.of(new CheckMatchingReferenceSet(referenceSetId, auth))).apply(new ReadGroupStreamer( auth, ShardBoundary.Requirement.STRICT, READ_FIELDS, SexChromosomeFilter.INCLUDE_XY)); PCollection<KV<PosRgsMq, Double>> coverageMeans = reads.apply(new CalculateCoverageMean()); PCollection<KV<Position, KV<PosRgsMq.MappingQuality, List<Double>>>> quantiles = coverageMeans .apply(new CalculateQuantiles(options.getNumQuantiles())); PCollection<KV<Position, Iterable<KV<PosRgsMq.MappingQuality, List<Double>>>>> answer = quantiles .apply(GroupByKey.<Position, KV<PosRgsMq.MappingQuality, List<Double>>>create()); answer.apply(ParDo.of(new CreateAnnotations(annotationSet.getId(), auth, true))); p.run(); } static class CheckMatchingReferenceSet extends DoFn<String, String> { private final String referenceSetIdForAllReadGroupSets; private final OfflineAuth auth; public CheckMatchingReferenceSet(String referenceSetIdForAllReadGroupSets, OfflineAuth auth) { this.referenceSetIdForAllReadGroupSets = referenceSetIdForAllReadGroupSets; this.auth = auth; } @Override public void processElement(DoFn<String, String>.ProcessContext c) throws Exception { String readGroupSetId = c.element(); String referenceSetId = GenomicsUtils.getReferenceSetId(readGroupSetId, auth); if (Strings.isNullOrEmpty(referenceSetId)) { throw new IllegalArgumentException( "No ReferenceSetId associated with ReadGroupSetId " + readGroupSetId + ". All ReadGroupSets in given input must have an associated ReferenceSet."); } if (!referenceSetIdForAllReadGroupSets.equals(referenceSetId)) { throw new IllegalArgumentException("ReadGroupSets in given input must all be associated with" + " the same ReferenceSetId : " + referenceSetId); } c.output(readGroupSetId); } } /** * Composite PTransform that takes in Reads read in from various ReadGroupSetIds using the * Genomics API and calculates the coverage mean at every distinct PosRgsMq object for the data * input (Position, ReadgroupSetId, and Mapping Quality). */ public static class CalculateCoverageMean extends PTransform<PCollection<Read>, PCollection<KV<PosRgsMq, Double>>> { @Override public PCollection<KV<PosRgsMq, Double>> apply(PCollection<Read> input) { return input.apply(ParDo.of(new CoverageCounts())) .apply(Combine.<PosRgsMq, Long>perKey(new SumCounts())).apply(ParDo.of(new CoverageMeans())); } } static class CoverageCounts extends DoFn<Read, KV<PosRgsMq, Long>> { private static final int LOW_MQ = 10; private static final int HIGH_MQ = 30; @Override public void processElement(ProcessContext c) { if (c.element().getAlignment() != null) { //is mapped // Calculate length of read long readLength = 0; for (CigarUnit cigar : c.element().getAlignment().getCigarList()) { switch (cigar.getOperation()) { case ALIGNMENT_MATCH: case SEQUENCE_MATCH: case SEQUENCE_MISMATCH: case PAD: case DELETE: case SKIP: readLength += cigar.getOperationLength(); break; } } // Get options for getting references/bucket width Options pOptions = c.getPipelineOptions().as(Options.class); // Calculate readEnd by shifting readStart by readLength long readStart = c.element().getAlignment().getPosition().getPosition(); long readEnd = readStart + readLength; // Get starting and ending references long refStart = 0; long refEnd = Long.MAX_VALUE; if (!pOptions.isAllReferences()) { // If we aren't using all references, parse the references to get proper start and end Iterable<Contig> contigs = Contig.parseContigsFromCommandLine(pOptions.getReferences()); for (Contig ct : contigs) { if (ct.referenceName.equals(c.element().getAlignment().getPosition().getReferenceName())) { refStart = ct.start; refEnd = ct.end; } } } // Calculate the index of the first bucket this read falls into long bucket = ((readStart - refStart) / pOptions.getBucketWidth()) * pOptions.getBucketWidth() + refStart; long readCurr = readStart; if (readStart < refStart) { // If the read starts before the first bucket, start our calculations at the start of // the reference, since we know the read has to overlap there or else it wouldn't // be in our data set. bucket = refStart; readCurr = refStart; } while (readCurr < readEnd && readCurr < refEnd) { // Loop over read to generate output for each bucket long baseCount = 0; long dist = Math.min(bucket + pOptions.getBucketWidth(), readEnd) - readCurr; readCurr += dist; baseCount += dist; Position position = new Position().setPosition(bucket) .setReferenceName(c.element().getAlignment().getPosition().getReferenceName()); Integer mq = c.element().getAlignment().getMappingQuality(); if (mq == null) { mq = 0; } PosRgsMq.MappingQuality mqEnum; if (mq < LOW_MQ) { mqEnum = PosRgsMq.MappingQuality.L; } else if (mq >= LOW_MQ && mq < HIGH_MQ) { mqEnum = PosRgsMq.MappingQuality.M; } else { mqEnum = PosRgsMq.MappingQuality.H; } c.output(KV.of(new PosRgsMq(position, c.element().getReadGroupSetId(), mqEnum), baseCount)); c.output(KV.of( new PosRgsMq(position, c.element().getReadGroupSetId(), PosRgsMq.MappingQuality.A), baseCount)); bucket += pOptions.getBucketWidth(); } } } } static class SumCounts implements SerializableFunction<Iterable<Long>, Long> { @Override public Long apply(Iterable<Long> input) { long ans = 0; for (Long l : input) { ans += l; } return ans; } } static class CoverageMeans extends DoFn<KV<PosRgsMq, Long>, KV<PosRgsMq, Double>> { @Override public void processElement(ProcessContext c) { KV ans = KV.of(c.element().getKey(), (double) c.element().getValue() / (double) c.getPipelineOptions().as(Options.class).getBucketWidth()); c.output(ans); } } /** * Composite PTransform that takes in the KV representing the coverage mean for a distinct * PosRgsMq object (Position, ReadgroupsetId, and Mapping Quality) and calculates the quantiles * such that 10% of the ReadGroupSetIds at a given Position (bucket) and mapping quality * have at least this much coverage, and the same for 20%, etc., all the way up until 90%. It puts * this into a KV object that can later be used to create annotations of the results. * * The percentiles will change based on the number of quantiles specified (see class description). */ public static class CalculateQuantiles extends PTransform<PCollection<KV<PosRgsMq, Double>>, PCollection<KV<Position, KV<PosRgsMq.MappingQuality, List<Double>>>>> { private final int numQuantiles; public CalculateQuantiles(int numQuantiles) { this.numQuantiles = numQuantiles; } @Override public PCollection<KV<Position, KV<PosRgsMq.MappingQuality, List<Double>>>> apply( PCollection<KV<PosRgsMq, Double>> input) { return input.apply(ParDo.of(new RemoveRgsId())) .apply(ApproximateQuantiles.<KV<Position, PosRgsMq.MappingQuality>, Double>perKey(numQuantiles)) .apply(ParDo.of(new MoveMapping())); } } static class RemoveRgsId extends DoFn<KV<PosRgsMq, Double>, KV<KV<Position, PosRgsMq.MappingQuality>, Double>> { @Override public void processElement(ProcessContext c) { KV ans = KV.of(KV.of(c.element().getKey().getPos(), c.element().getKey().getMq()), c.element().getValue()); c.output(ans); } } static class MoveMapping extends DoFn<KV<KV<Position, PosRgsMq.MappingQuality>, List<Double>>, KV<Position, KV<PosRgsMq.MappingQuality, List<Double>>>> { @Override public void processElement(ProcessContext c) { KV ans = KV.of(c.element().getKey().getKey(), KV.of(c.element().getKey().getValue(), c.element().getValue())); c.output(ans); } } /* TODO: If a particular batchCreateAnnotations request fails more than 4 times and the bundle gets retried by the dataflow pipeline, some annotations may be written multiple times. */ static class CreateAnnotations extends DoFn<KV<Position, Iterable<KV<PosRgsMq.MappingQuality, List<Double>>>>, Annotation> { private final String asId; private final OfflineAuth auth; private final List<Annotation> currAnnotations; private final boolean write; public CreateAnnotations(String asId, OfflineAuth auth, boolean write) { this.asId = asId; this.auth = auth; this.currAnnotations = Lists.newArrayList(); this.write = write; } @Override public void processElement(ProcessContext c) throws GeneralSecurityException, IOException { Options pOptions = c.getPipelineOptions().as(Options.class); Position bucket = c.element().getKey(); Annotation a = new Annotation().setAnnotationSetId(asId).setStart(bucket.getPosition()) .setEnd(bucket.getPosition() + pOptions.getBucketWidth()) .setReferenceName(bucket.getReferenceName()).setType("GENERIC") .setInfo(new HashMap<String, List<Object>>()); for (KV<PosRgsMq.MappingQuality, List<Double>> mappingQualityKV : c.element().getValue()) { List<Object> output = Lists.newArrayList(); for (int i = 0; i < mappingQualityKV.getValue().size(); i++) { double value = Math.round(mappingQualityKV.getValue().get(i) * 1000000.0) / 1000000.0; output.add(Double.toString(value)); } a.getInfo().put(mappingQualityKV.getKey().toString(), output); } if (write) { currAnnotations.add(a); if (currAnnotations.size() == 512) { // Batch create annotations once we hit the max amount for a batch batchCreateAnnotations(); } } c.output(a); } @Override public void finishBundle(Context c) throws IOException, GeneralSecurityException { // Finish up any leftover annotations at the end. if (write && !currAnnotations.isEmpty()) { batchCreateAnnotations(); } } private void batchCreateAnnotations() throws IOException, GeneralSecurityException { Genomics genomics = GenomicsFactory.builder().build().fromOfflineAuth(auth); Genomics.Annotations.BatchCreate aRequest = genomics.annotations() .batchCreate(new BatchCreateAnnotationsRequest().setAnnotations(currAnnotations)); RetryPolicy retryP = RetryPolicy.nAttempts(4); retryP.execute(aRequest); currAnnotations.clear(); } } private static AnnotationSet createAnnotationSet(String referenceSetId) throws GeneralSecurityException, IOException { if (referenceSetId == null) { throw new IOException("ReferenceSetId was null for this readgroupset"); } // Create a new annotation set using the given datasetId, or the one that this data was from if // one was not provided AnnotationSet as = new AnnotationSet(); as.setDatasetId(options.getOutputDatasetId()); if (!"".equals(options.getAnnotationSetName())) { as.setName(options.getAnnotationSetName()); } else { if (!options.isAllReferences()) { as.setName("Read Depth for " + options.getReferences()); } else { as.setName("Read Depth for all references"); } } as.setReferenceSetId(referenceSetId); as.setType("GENERIC"); Genomics genomics = GenomicsFactory.builder().build().fromOfflineAuth(auth); Annotationsets.Create asRequest = genomics.annotationsets().create(as); AnnotationSet asWithId = asRequest.execute(); return asWithId; } }