Improve handling of intersecting groups in SVStratify (#9185)

Author: mwalker174

src/main/java/org/broadinstitute/hellbender/tools/walkers/sv/SVStratify.java

@@ -6,6 +6,7 @@
 import htsjdk.variant.variantcontext.VariantContextBuilder;
 import htsjdk.variant.variantcontext.writer.VariantContextWriter;
 import htsjdk.variant.vcf.VCFHeader;
+import htsjdk.variant.vcf.VCFHeaderLineCount;
 import htsjdk.variant.vcf.VCFHeaderLineType;
 import htsjdk.variant.vcf.VCFInfoHeaderLine;
 import org.broadinstitute.barclay.argparser.Argument;
@@ -40,7 +41,8 @@
  *     <li>Reference track overlap</li>
  * </ul>
  * Records are annotated with their respective strata names in the {@link GATKSVVCFConstants#STRATUM_INFO_KEY} INFO
- * field. Users must provide a stratification configuration .tsv file (tab-delimited table) with the following column
+ * field. SVs that do not match any of the groups will be annotated with the {@link SVStratify#DEFAULT_STRATUM} group.
+ * Users must provide a stratification configuration .tsv file (tab-delimited table) with the following column
  * header on the first line:
  * <ol>
  *     <li>NAME</li>
@@ -78,7 +80,7 @@
  * {@link SVStratificationEngineArgumentsCollection#trackNameList} parameters. For example,
  * </p>
  * <pre>
- *     gatk GroupedSVCluster \
+ *     gatk SVStratify \
  *       --track-name RM \
  *       --track-intervals repeatmasker.bed \
  *       --track-name SD \
@@ -104,12 +106,11 @@
  * overlap is only defined by {@link SVStratificationEngineArgumentsCollection#numBreakpointOverlapsInterchrom}.
  * </p>
  *
- * <p>By default, each stratification group must be mutually exclusive, meaning that any given SV can only belong to
+ * <p>If using the --split-output option then each stratification group must be mutually exclusive, meaning that any given SV can only belong to
  * one group. An error is thrown if the tool encounters a variant that meets the criteria for more than one group.
- * This restriction can be overridden with the {@link SVStratify#ALLOW_MULTIPLE_MATCHES_LONG_NAME} argument, in which
- * case the record will be written out multiple times: once for each matching stratification group with the corresponding
- * {@link GATKSVVCFConstants#STRATUM_INFO_KEY} value. Furthermore, SVs that do not match any of the groups will be
- * annotated with the {@link SVStratify#DEFAULT_STRATUM} group.</p>
+ * This restriction can be overridden with the {@link SVStratify#ALLOW_MULTIPLE_MATCHES_LONG_NAME} argument, in which case
+ * records belonging to multiple stratification groups will be written to each corresponding file (hence possibly
+ * resulting in duplicated records).</p>
  *
  * <p>If using {@link #SPLIT_OUTPUT_LONG_NAME} then the tool generates a set of VCFs as output with each VCF containing
  * the records of each group.</p>
@@ -242,7 +243,7 @@ protected void createGroupWriter(final String name, final Path path) {
     }
 
     public static void addStratifyMetadata(final VCFHeader header) {
-        header.addMetaDataLine(new VCFInfoHeaderLine(GATKSVVCFConstants.STRATUM_INFO_KEY, 1,
+        header.addMetaDataLine(new VCFInfoHeaderLine(GATKSVVCFConstants.STRATUM_INFO_KEY, VCFHeaderLineCount.UNBOUNDED,
                 VCFHeaderLineType.String, "Stratum ID"));
     }
 
@@ -322,16 +323,23 @@ public void apply(final VariantContext variant, final ReadsContext readsContext,
         if (stratifications.isEmpty()) {
             writers.get(DEFAULT_STRATUM).add(builder.attribute(GATKSVVCFConstants.STRATUM_INFO_KEY, DEFAULT_STRATUM).make());
         } else {
-            if (!allowMultipleMatches && stratifications.size() > 1) {
-                final String matchesString = String.join(", ", stratifications.stream().map(SVStratificationEngine.Stratum::getName).collect(Collectors.toList()));
-                throw new GATKException("Record " + record.getId() + " matched multiple groups: " + matchesString + ". Bypass this error using the --" + ALLOW_MULTIPLE_MATCHES_LONG_NAME + " argument");
-            }
-            for (final SVStratificationEngine.Stratum stratum : stratifications) {
-                final VariantContextWriter writer = splitOutput ? writers.get(stratum.getName()) : writers.get(DEFAULT_STRATUM);
-                if (writer == null) {
-                    throw new GATKException("Writer not found for group: " + stratum.getName());
+            final List<String> stratumNames = new ArrayList<>(stratifications).stream().map(SVStratificationEngine.Stratum::getName).sorted().collect(Collectors.toUnmodifiableList());
+            final VariantContext outputVariant = builder.attribute(GATKSVVCFConstants.STRATUM_INFO_KEY, stratumNames).make();
+            if (splitOutput) {
+                if (!allowMultipleMatches && stratifications.size() > 1) {
+                    final String matchesString = String.join(", ", stratumNames);
+                    throw new GATKException("Record " + record.getId() + " matched multiple groups: " + matchesString + ". Bypass this error using the --" + ALLOW_MULTIPLE_MATCHES_LONG_NAME + " argument");
+                }
+                for (final SVStratificationEngine.Stratum stratum : stratifications) {
+                    final VariantContextWriter writer = writers.get(stratum.getName());
+                    if (writer == null) {
+                        throw new GATKException("Writer not found for group: " + stratum.getName());
+                    }
+                    writer.add(outputVariant);
                 }
-                writer.add(builder.attribute(GATKSVVCFConstants.STRATUM_INFO_KEY, stratum.getName()).make());
+            } else {
+                final VariantContextWriter writer = writers.get(DEFAULT_STRATUM);
+                writer.add(outputVariant);
             }
         }
     }

src/test/java/org/broadinstitute/hellbender/tools/walkers/sv/SVStratifyIntegrationTest.java

@@ -124,8 +124,42 @@ public void testBwaMeltCohortSingleOutput() {
         Assert.assertEquals(outputVcf.getRight().size(), inputVcf.getRight().size());
     }
 
-    @Test(expectedExceptions = GATKException.class)
+    @Test
     public void testBwaMeltCohortRedundant() {
+        final File outputDir = createTempDir("stratify");
+        final File outputFile = outputDir.toPath().resolve("out.vcf.gz").toFile();
+        final String inputVcfPath = getToolTestDataDir() + "bwa_melt.chr22.vcf.gz";
+        final String configFile = getToolTestDataDir() + "test_config_redundant.tsv";
+
+        final String segdupFile = getToolTestDataDir() + "hg38.SegDup.chr22.bed";
+        final String segdupName = "SD";
+        final String repeatmaskerFile = getToolTestDataDir() + "hg38.RM.chr22_subsampled.bed";
+        final String repeatmaskerName = "RM";
+
+        final ArgumentsBuilder args = new ArgumentsBuilder()
+                .addOutput(outputFile)
+                .add(SVStratificationEngineArgumentsCollection.STRATIFY_CONFIG_FILE_LONG_NAME, configFile)
+                .add(SVStratificationEngineArgumentsCollection.TRACK_NAME_FILE_LONG_NAME, segdupName)
+                .add(SVStratificationEngineArgumentsCollection.TRACK_INTERVAL_FILE_LONG_NAME, segdupFile)
+                .add(SVStratificationEngineArgumentsCollection.TRACK_NAME_FILE_LONG_NAME, repeatmaskerName)
+                .add(SVStratificationEngineArgumentsCollection.TRACK_INTERVAL_FILE_LONG_NAME, repeatmaskerFile)
+                .add(SVStratificationEngineArgumentsCollection.OVERLAP_FRACTION_LONG_NAME, 0.5)
+                .add(StandardArgumentDefinitions.SEQUENCE_DICTIONARY_NAME, GATKBaseTest.FULL_HG38_DICT)
+                .add(StandardArgumentDefinitions.VARIANT_LONG_NAME, inputVcfPath);
+
+        runCommandLine(args, SVStratify.class.getSimpleName());
+
+        final List<File> outputFiles = Lists.newArrayList(outputDir.listFiles()).stream().filter(VcfUtils::isVariantFile).collect(Collectors.toUnmodifiableList());
+        Assert.assertEquals(outputFiles.size(), 1);
+        Assert.assertEquals(outputFiles.get(0).getAbsolutePath(), outputFile.getAbsolutePath());
+        final Pair<VCFHeader, List<VariantContext>> inputVcf = VariantContextTestUtils.readEntireVCFIntoMemory(inputVcfPath);
+        final Pair<VCFHeader, List<VariantContext>> outputVcf = VariantContextTestUtils.readEntireVCFIntoMemory(outputFile.getAbsolutePath());
+        // No duplicated records
+        Assert.assertEquals(outputVcf.getRight().size(), inputVcf.getRight().size());
+    }
+
+    @Test(expectedExceptions = GATKException.class)
+    public void testBwaMeltCohortRedundantFails() {
         final File outputDir = createTempDir("stratify");
         final String inputVcfPath = getToolTestDataDir() + "bwa_melt.chr22.vcf.gz";
         final String configFile = getToolTestDataDir() + "test_config_redundant.tsv";
@@ -153,7 +187,7 @@ public void testBwaMeltCohortRedundant() {
 
     @Test
     public void testBwaMeltCohortBypassRedundant() {
-        final File outputDir = createTempDir("stratify");
+        final File outputDir = createTempDir("stratify_bypass_redundant");
         final String inputVcfPath = getToolTestDataDir() + "bwa_melt.chr22.vcf.gz";
         final String configFile = getToolTestDataDir() + "test_config_redundant.tsv";
 
@@ -177,6 +211,16 @@ public void testBwaMeltCohortBypassRedundant() {
                 .addFlag(SVStratify.ALLOW_MULTIPLE_MATCHES_LONG_NAME);
 
         runCommandLine(args, SVStratify.class.getSimpleName());
+
+        final List<File> outputFiles = Lists.newArrayList(outputDir.listFiles()).stream().filter(VcfUtils::isVariantFile).collect(Collectors.toUnmodifiableList());
+        Assert.assertEquals(outputFiles.size(), 8);
+        long totalRecords = 0;
+        for (final File outputFile : outputFiles) {
+            final Pair<VCFHeader, List<VariantContext>> inputVcf = VariantContextTestUtils.readEntireVCFIntoMemory(inputVcfPath);
+            final Pair<VCFHeader, List<VariantContext>> outputVcf = VariantContextTestUtils.readEntireVCFIntoMemory(outputFile.getAbsolutePath());
+            totalRecords += outputVcf.getRight().size();
+        }
+        Assert.assertEquals(totalRecords, 1458);
     }
 
     @Test(expectedExceptions = {GATKException.class})