tests for training sets

broadinstitute · Apr 1, 2015 · a10df74 · a10df74
1 parent a1e1d0d
commit a10df74
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 6 deletions.
diff --git a/src/main/java/org/broadinstitute/hellbender/tools/walkers/vqsr/VQSRTrainingSetManager.java b/src/main/java/org/broadinstitute/hellbender/tools/walkers/vqsr/VQSRTrainingSetManager.java
@@ -8,6 +8,7 @@
 import org.broadinstitute.hellbender.exceptions.GATKException;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
 
@@ -61,6 +62,10 @@ public FeatureInput<VariantContext> getFeatureInput(){
         }
     }
 
+    public void addTrainingSets( final Collection<FeatureInput<VariantContext>> trainingSets ) {
+        trainingSets.forEach(this::addTrainingSet);
+    }
+
     public void addTrainingSet( final FeatureInput<VariantContext> trainingSet ) {
         trainingSets.add(new TrainingSet(trainingSet));
     }
@@ -73,7 +78,7 @@ public boolean hasTruthSet() {
         return trainingSets.stream().anyMatch(ts -> ts.isTruth);
     }
 
-    public void parseTrainingSets( final FeatureContext tracker, final VariantContext evalVC, final VariantDatum datum, final boolean TRUST_ALL_POLYMORPHIC ) {
+    public void parseTrainingSets( final FeatureContext tracker, final VariantContext evalVC, final VariantDatum datum, final boolean trust_all_polymorphic ) {
         datum.isKnown = false;
         datum.atTruthSite = false;
         datum.atTrainingSite = false;
@@ -82,7 +87,7 @@ public void parseTrainingSets( final FeatureContext tracker, final VariantContex
 
         for( final TrainingSet trainingSet : trainingSets ) {
             for( final VariantContext trainVC : tracker.getValues(trainingSet.getFeatureInput()) ) {
-                if( isValidVariant( evalVC, trainVC, TRUST_ALL_POLYMORPHIC ) ) {
+                if( isValidVariant( evalVC, trainVC, trust_all_polymorphic ) ) {
                     datum.isKnown = datum.isKnown || trainingSet.isKnown;
                     datum.atTruthSite = datum.atTruthSite || trainingSet.isTruth;
                     datum.atTrainingSite = datum.atTrainingSite || trainingSet.isTraining;
@@ -96,9 +101,9 @@ public void parseTrainingSets( final FeatureContext tracker, final VariantContex
         }
     }
 
-    private boolean isValidVariant( final VariantContext evalVC, final VariantContext trainVC, final boolean TRUST_ALL_POLYMORPHIC) {
+    private boolean isValidVariant( final VariantContext evalVC, final VariantContext trainVC, final boolean trustAllPolymorphic) {
         return trainVC != null && trainVC.isNotFiltered() && trainVC.isVariant() && checkVariationClass( evalVC, trainVC ) &&
-                (TRUST_ALL_POLYMORPHIC || !trainVC.hasGenotypes() || trainVC.isPolymorphicInSamples());
+                (trustAllPolymorphic || !trainVC.hasGenotypes() || trainVC.isPolymorphicInSamples());
     }
 
     private static boolean checkVariationClass( final VariantContext evalVC, final VariantContext trainVC ) {

diff --git a/...java/org/broadinstitute/hellbender/tools/walkers/vqsr/VQSRTrainingSetManagerUnitTest.java b/...java/org/broadinstitute/hellbender/tools/walkers/vqsr/VQSRTrainingSetManagerUnitTest.java
@@ -1,10 +1,19 @@
 package org.broadinstitute.hellbender.tools.walkers.vqsr;
 
 import htsjdk.variant.variantcontext.VariantContext;
+import joptsimple.internal.Strings;
+import org.broadinstitute.hellbender.cmdline.Argument;
+import org.broadinstitute.hellbender.cmdline.CommandLineParser;
 import org.broadinstitute.hellbender.engine.FeatureInput;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.io.PrintStream;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.apache.commons.io.output.NullOutputStream.NULL_OUTPUT_STREAM;
+
 public class VQSRTrainingSetManagerUnitTest {
     @Test
     public void testEmpty(){
@@ -13,13 +22,91 @@ public void testEmpty(){
         Assert.assertFalse(vtsm.hasTruthSet(), "truth set");
     }
 
+
+    private static class TestClass{
+        @Argument(fullName = "resource", shortName = "resource")
+        List<FeatureInput<VariantContext>> resource;
+    }
+
+    private static void makeAndAddTrainingSets(VQSRTrainingSetManager vtsm, String... args) {
+        TestClass obj= new TestClass();
+        final String argList = Strings.join(Arrays.asList(args), " ");
+        new CommandLineParser(obj).parseArguments(new PrintStream(NULL_OUTPUT_STREAM), argList.split(" ", -1));
+        vtsm.addTrainingSets(obj.resource);
+    }
+
+
     @Test
     public void test1FeatureInput(){
         VQSRTrainingSetManager vtsm = new VQSRTrainingSetManager();
-        FeatureInput<VariantContext> fi1= new FeatureInput<>();
-        vtsm.addTrainingSet(fi1);
+        final String args = "--resource resource,known=true,prior=10.0:myFile";
+        makeAndAddTrainingSets(vtsm, args);
         Assert.assertFalse(vtsm.hasTrainingSet(), "training set");
         Assert.assertFalse(vtsm.hasTruthSet(), "truth set");
     }
 
+    @Test
+    public void testFeatureInputWithTruthSetTrainingSet(){
+        VQSRTrainingSetManager vtsm = new VQSRTrainingSetManager();
+        final String[] args= {"--resource resource,known=true,prior=10.0:myFile",
+                              "--resource resource,truth=true,training=true,prior=15.0:myFile3"};
+
+        makeAndAddTrainingSets(vtsm, args);
+
+        Assert.assertTrue(vtsm.hasTrainingSet(), "training set");
+        Assert.assertTrue(vtsm.hasTruthSet(), "truth set");
+    }
+
+    @Test
+    public void testFeatureInputWithTruthSetTrainingSet_twoFiles(){
+        VQSRTrainingSetManager vtsm = new VQSRTrainingSetManager();
+        final String[] args= {"--resource resource,known=true,prior=10.0:myFile",
+                              "--resource resource,truth=true,prior=15.0:myFile2",
+                              "--resource resource,training=true,prior=15.0:myFile3"};
+        makeAndAddTrainingSets(vtsm, args);
+
+        Assert.assertTrue(vtsm.hasTrainingSet(), "training set");
+        Assert.assertTrue(vtsm.hasTruthSet(), "truth set");
+    }
+
+    @Test
+    public void testFeatureInputWithOnlyTrainingSet(){
+        VQSRTrainingSetManager vtsm = new VQSRTrainingSetManager();
+        final String[] args= {"--resource resource,known=true,prior=10.0:myFile",
+                              "--resource resource,training=true,prior=15.0:myFile2"};
+        makeAndAddTrainingSets(vtsm, args);
+
+        Assert.assertTrue(vtsm.hasTrainingSet(), "training set");
+        Assert.assertFalse(vtsm.hasTruthSet(), "truth set");
+    }
+
+    @Test
+    public void testFeatureInputWithOnlyTruthSet(){
+        VQSRTrainingSetManager vtsm = new VQSRTrainingSetManager();
+        final String[] args= {"--resource resource,known=true,prior=10.0:myFile",
+                              "--resource resource,truth=true,prior=15.0:myFile2"};
+        makeAndAddTrainingSets(vtsm, args);
+        Assert.assertFalse(vtsm.hasTrainingSet(), "training set");
+        Assert.assertTrue(vtsm.hasTruthSet(), "truth set");
+    }
+
+    @Test
+    public void testFeatureInputConsensusSet(){
+        VQSRTrainingSetManager vtsm = new VQSRTrainingSetManager();
+        final String[] args= {"--resource resource,consensus=true,prior=10.0:myFile",
+                "--resource resource,truth=true,prior=15.0:myFile2"};
+        makeAndAddTrainingSets(vtsm, args);
+        Assert.assertFalse(vtsm.hasTrainingSet(), "training set");
+        Assert.assertTrue(vtsm.hasTruthSet(), "truth set");
+    }
+
+    @Test
+    public void testFeatureInputBadSet(){
+        VQSRTrainingSetManager vtsm = new VQSRTrainingSetManager();
+        final String[] args= {"--resource resource,bad=true,prior=10.0:myFile",
+                "--resource resource,truth=true,prior=15.0:myFile2"};
+        makeAndAddTrainingSets(vtsm, args);
+        Assert.assertFalse(vtsm.hasTrainingSet(), "training set");
+        Assert.assertTrue(vtsm.hasTruthSet(), "truth set");
+    }
 }