Add type of variant as input

genotype_results.py

@@ -26,10 +26,10 @@ def passed_variant(record):
     return record.filter is None or len(record.filter) == 0 or "PASS" in record.filter
 
 
-def readgenotypes(vcffile: str, variants: dict, dofilter: bool=False):
+def readgenotypes(vcffile: str, variants: dict, type_v, dofilter: bool=False):
     vcfin = VariantFile(vcffile)
     for r in vcfin:
-        if not dofilter or passed_variant(r):
+        if (not dofilter or passed_variant(r)) and r.sv_type.lower() == type_v:
             variants[r.id] = r
     return variants
 
@@ -65,17 +65,24 @@ def getvarsize(variant):
     return variant.stop - variant.start + 1
 
 
-def main(genotypes, predicted, filtered, output, verbose=False):
-    true_genotypes = readgenotypes(genotypes, {})
+def main(genotypes, predicted, filtered, output, type_v="del", verbose=False):
+    true_genotypes = readgenotypes(vcffile=genotypes,
+                                   variants={},
+                                   type_v=type_v)
     pred_genotypes = {}
     filtered_genotypes = {}
     with open(predicted, "r") as preds:
         for pred in preds:
             pred = pred.rstrip()
             if pred != "":
-                readgenotypes(pred, pred_genotypes)
+                readgenotypes(vcffile=pred,
+                              variants=pred_genotypes,
+                              type_v=type_v)
                 if filtered:
-                    readgenotypes(pred, filtered_genotypes, True)
+                    readgenotypes(vcffile=pred,
+                                  variants=filtered_genotypes,
+                                  type_v=type_v,
+                                  dofilter=True)
 
     true_pybed = variants_to_pybed(true_genotypes)
     pred_pybed = variants_to_pybed(pred_genotypes)
@@ -93,8 +100,9 @@ def main(genotypes, predicted, filtered, output, verbose=False):
             soft = sv[3].split('_')[0]
             if sv[3] in filtered_genotypes:
                 out.write("%s\t%s\t%d\t%3.2f\t%d\t%d\n" %
-                      (sv[3], "pass", varsize, correct / (wrong + correct), left_prec, right_prec))
-            out.write("%s\t%s\t%d\t%3.2f\t%d\t%d\n" % (sv[3], soft, varsize, correct / (wrong + correct), left_prec, right_prec))
+                          (sv[3], "pass", varsize, correct / (wrong + correct), left_prec, right_prec))
+            out.write("%s\t%s\t%d\t%3.2f\t%d\t%d\n" %
+                      (sv[3], soft, varsize, correct / (wrong + correct), left_prec, right_prec))
 
     if verbose:
         print("%d variants" % (len(true_genotypes)))
@@ -112,6 +120,7 @@ if __name__ == '__main__':
                         help='File listing VCF files containing the predicted results, with genotypes')
     parser.add_argument('--filter', action='store_const', const=True, default=False,
                         help='Predicted results has filter information')
+    parser.add_argument('-t', '--type', type=str, required=True, choices=["inv", "del"], help="Type of variant")
     parser.add_argument('-o', '--output', type=str, required=True, help="Output file")
     parser.add_argument('-v', '--verbose', action='store_const', const=True, default=False,
                         help='Verbose mode')
@@ -122,4 +131,5 @@ if __name__ == '__main__':
          predicted=args.predictions_vcf,
          filtered=args.filter,
          output=args.output,
+         type_v=args.type,
          verbose=args.verbose)