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 index 7a28f8f8634..a6b0e64a5a8 100644 --- 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 getFeatureInput(){ } } + public void addTrainingSets( final Collection> trainingSets ) { + trainingSets.forEach(this::addTrainingSet); + } + public void addTrainingSet( final FeatureInput 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/src/test/java/org/broadinstitute/hellbender/tools/walkers/vqsr/VQSRTrainingSetManagerUnitTest.java b/src/test/java/org/broadinstitute/hellbender/tools/walkers/vqsr/VQSRTrainingSetManagerUnitTest.java index f6ed489e8bf..928b6dbcc15 100644 --- a/src/test/java/org/broadinstitute/hellbender/tools/walkers/vqsr/VQSRTrainingSetManagerUnitTest.java +++ b/src/test/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> 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 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"); + } }