Downsample reads with a ReservoirDownsampler in CNNScoreVariants. (#5622

) * Downsample reads with ReservoirDownsampler in CNNScoreVariants. * updated test files and reset random number generator in tests
broadinstitute · Jan 30, 2019 · 335fac0 · 335fac0
1 parent 64e2d3a
commit 335fac0
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/src/main/java/org/broadinstitute/hellbender/tools/walkers/vqsr/CNNScoreVariants.java b/src/main/java/org/broadinstitute/hellbender/tools/walkers/vqsr/CNNScoreVariants.java
@@ -18,6 +18,10 @@
 import org.broadinstitute.hellbender.engine.*;
 import org.broadinstitute.barclay.argparser.*;
 import org.broadinstitute.hellbender.engine.filters.*;
+import org.broadinstitute.hellbender.exceptions.GATKException;
+import org.broadinstitute.hellbender.utils.downsampling.ReadsDownsamplingIterator;
+import org.broadinstitute.hellbender.utils.downsampling.ReservoirDownsampler;
+import org.broadinstitute.hellbender.utils.haplotype.HaplotypeBAMWriter;
 import org.broadinstitute.hellbender.utils.io.IOUtils;
 import org.broadinstitute.hellbender.utils.io.Resource;
 import org.broadinstitute.barclay.help.DocumentedFeature;
@@ -420,7 +424,7 @@ private void transferReadsToPythonViaFifo(final VariantContext variant, final Re
         } catch (UnsupportedEncodingException e) {
             throw new GATKException("Trying to make string from reference, but unsupported encoding UTF-8.", e);
         }
-        Iterator<GATKRead> readIt = readsContext.iterator();
+        Iterator<GATKRead> readIt = new ReadsDownsamplingIterator(readsContext.iterator(), new ReservoirDownsampler(readLimit));
         if (!readIt.hasNext()) {
             logger.warn("No reads at contig:" + variant.getContig() + " site:" + String.valueOf(variant.getStart()));
         }

diff --git a/...ava/org/broadinstitute/hellbender/tools/walkers/vqsr/CNNScoreVariantsIntegrationTest.java b/...ava/org/broadinstitute/hellbender/tools/walkers/vqsr/CNNScoreVariantsIntegrationTest.java
@@ -4,6 +4,7 @@
 import org.broadinstitute.hellbender.cmdline.StandardArgumentDefinitions;
 import org.broadinstitute.hellbender.testutils.ArgumentsBuilder;
 import org.broadinstitute.hellbender.testutils.IntegrationTestSpec;
+import org.broadinstitute.hellbender.utils.Utils;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
@@ -155,6 +156,9 @@ public void testOnContigEdge() throws IOException {
      */
     @Test(groups = {"python"})
     public void testInference2dResourceModel() throws IOException {
+        // We reset the random number generator at the beginning of each test so that the random down-sampling of reads
+        // by the reservoir down-sampler does not cause slightly different scores.
+        Utils.resetRandomGenerator();
         TensorType tt = TensorType.read_tensor;
         final ArgumentsBuilder argsBuilder = new ArgumentsBuilder();
         argsBuilder.addArgument(StandardArgumentDefinitions.VARIANT_LONG_NAME, inputVCF)
@@ -177,6 +181,7 @@ public void testInference2dResourceModel() throws IOException {
      */
     @Test(groups = {"python"})
     public void testInferenceArchitecture2d() throws IOException {
+        Utils.resetRandomGenerator();
         final boolean newExpectations = false;
         TensorType tt = TensorType.read_tensor;
         final ArgumentsBuilder argsBuilder = new ArgumentsBuilder();
@@ -202,6 +207,7 @@ public void testInferenceArchitecture2d() throws IOException {
 
     @Test(groups = {"python"})
     public void testInferenceWeights2d() throws IOException {
+        Utils.resetRandomGenerator();
         TensorType tt = TensorType.read_tensor;
         final ArgumentsBuilder argsBuilder = new ArgumentsBuilder();
         argsBuilder.addArgument(StandardArgumentDefinitions.VARIANT_LONG_NAME, inputVCF)
@@ -222,6 +228,7 @@ public void testInferenceWeights2d() throws IOException {
 
     @Test(groups = {"python"})
     public void testInferenceArchitectureAndWeights2d() throws IOException {
+        Utils.resetRandomGenerator();
         TensorType tt = TensorType.read_tensor;
         final ArgumentsBuilder argsBuilder = new ArgumentsBuilder();
         argsBuilder.addArgument(StandardArgumentDefinitions.VARIANT_LONG_NAME, inputVCF)

diff --git a/src/test/resources/large/VQSR/expected/cnn_2d_chr20_subset_expected.vcf b/src/test/resources/large/VQSR/expected/cnn_2d_chr20_subset_expected.vcf