Extended standalone

neusomatic/python/call.py

@@ -27,7 +27,7 @@
 from dataloader import NeuSomaticDataset, matrix_transform
 from utils import get_chromosomes_order, prob2phred
 from merge_tsvs import merge_tsvs
-from defaults import VARTYPE_CLASSES
+from defaults import VARTYPE_CLASSES, NUM_ENS_FEATURES, NUM_ST_FEATURES
 
 import torch._utils
 try:
@@ -396,7 +396,7 @@ def write_vcf(vcf_records, output_vcf, chroms_order, pass_threshold, lowqual_thr
 
 def call_neusomatic(candidates_tsv, ref_file, out_dir, checkpoint, num_threads,
                     batch_size, max_load_candidates, pass_threshold, lowqual_threshold,
-                    ensemble,
+                    force_zero_ann_cols,
                     use_cuda):
     logger = logging.getLogger(call_neusomatic.__name__)
 
@@ -412,37 +412,75 @@ def call_neusomatic(candidates_tsv, ref_file, out_dir, checkpoint, num_threads,
         chroms = rf.references
 
     data_transform = matrix_transform((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-    num_channels = 119 if ensemble else 26
-    net = NeuSomaticNet(num_channels)
-    if use_cuda:
-        logger.info("GPU calling!")
-        net.cuda()
-    else:
-        logger.info("CPU calling!")
 
-    if torch.cuda.device_count() > 1:
-        logger.info("We use {} GPUs!".format(torch.cuda.device_count()))
-        net = nn.DataParallel(net)
-
-    if not os.path.exists(out_dir):
-        os.mkdir(out_dir)
     logger.info("Load pretrained model from checkpoint {}".format(checkpoint))
     pretrained_dict = torch.load(
         checkpoint, map_location=lambda storage, loc: storage)
     pretrained_state_dict = pretrained_dict["state_dict"]
     model_tag = pretrained_dict["tag"]
     logger.info("tag: {}".format(model_tag))
-
-    matrices_dir = "{}/matrices_{}".format(out_dir, model_tag)
-    if os.path.exists(matrices_dir):
-        logger.warning("Remove matrices directory: {}".format(matrices_dir))
-        shutil.rmtree(matrices_dir)
-    os.mkdir(matrices_dir)
     coverage_thr = pretrained_dict["coverage_thr"]
     if "normalize_channels" in pretrained_dict:
         normalize_channels = pretrained_dict["normalize_channels"]
     else:
         normalize_channels = False
+    if "no_seq_complexity" in pretrained_dict:
+        no_seq_complexity = pretrained_dict["no_seq_complexity"]
+    else:
+        no_seq_complexity = True
+    if "zero_ann_cols" in pretrained_dict:
+        zero_ann_cols = pretrained_dict["zero_ann_cols"]
+    else:
+        zero_ann_cols = []
+
+    if force_zero_ann_cols:
+        logger.info(
+            "Override zero_ann_cols from force_zero_ann_cols: {}".format(force_zero_ann_cols))
+        zero_ann_cols = force_zero_ann_cols
+
+    logger.info("coverage_thr: {}".format(coverage_thr))
+    logger.info("normalize_channels: {}".format(normalize_channels))
+    logger.info("no_seq_complexity: {}".format(no_seq_complexity))
+    logger.info("zero_ann_cols: {}".format(zero_ann_cols))
+
+
+    expected_ens_fields = NUM_ENS_FEATURES
+    if not no_seq_complexity:
+        expected_ens_fields += 2
+
+    logger.info("expected_ens_fields: {}".format(expected_ens_fields))
+
+    expected_st_fields = 4
+
+    logger.info("expected_st_fields: {}".format(expected_st_fields))
+
+    ensemble = False
+    for tsv in candidates_tsv:
+        with open(tsv) as i_f:
+            x = i_f.readline().strip().split()
+            if x:
+                if len(x) == expected_ens_fields + 4:
+                    ensemble = True
+                    break
+                elif len(x) == 4:
+                    break
+                else:
+                    raise Exception("Wrong number of fields in {}: {}".format(tsv, len(x)))
+
+    num_channels = expected_ens_fields + \
+        NUM_ST_FEATURES if ensemble else NUM_ST_FEATURES
+
+    logger.info("Number of channels: {}".format(num_channels))
+    net = NeuSomaticNet(num_channels)
+    if use_cuda:
+        logger.info("GPU calling!")
+        net.cuda()
+    else:
+        logger.info("CPU calling!")
+
+    if torch.cuda.device_count() > 1:
+        logger.info("We use {} GPUs!".format(torch.cuda.device_count()))
+        net = nn.DataParallel(net)
 
     model_dict = net.state_dict()
 
@@ -465,6 +503,16 @@ def call_neusomatic(candidates_tsv, ref_file, out_dir, checkpoint, num_threads,
     # 3. load the new state dict
     net.load_state_dict(pretrained_state_dict)
 
+
+    if not os.path.exists(out_dir):
+        os.mkdir(out_dir)
+    matrices_dir = "{}/matrices_{}".format(out_dir, model_tag)
+    if os.path.exists(matrices_dir):
+        logger.warning("Remove matrices directory: {}".format(matrices_dir))
+        shutil.rmtree(matrices_dir)
+    os.mkdir(matrices_dir)
+
+
     new_split_tsvs_dir = os.path.join(out_dir, "split_tsvs")
     if os.path.exists(new_split_tsvs_dir):
         logger.warning(
@@ -517,7 +565,8 @@ def call_neusomatic(candidates_tsv, ref_file, out_dir, checkpoint, num_threads,
                                          transform=data_transform, is_test=True,
                                          num_threads=num_threads,
                                          coverage_thr=coverage_thr,
-                                         normalize_channels=normalize_channels)
+                                         normalize_channels=normalize_channels,
+                                         zero_ann_cols=zero_ann_cols)
             call_loader = torch.utils.data.DataLoader(call_set,
                                                       batch_size=batch_size,
                                                       shuffle=True, pin_memory=True,
@@ -597,6 +646,11 @@ def call_neusomatic(candidates_tsv, ref_file, out_dir, checkpoint, num_threads,
     parser.add_argument('--lowqual_threshold', type=float,
                         help='SCORE for LowQual (PASS for lowqual_threshold <= score < pass_threshold)',
                         default=0.4)
+    parser.add_argument('--force_zero_ann_cols', nargs="*", type=int,
+                        help='force columns to be set to zero in the annotations. Higher priority than \
+                              --zero_ann_cols and pretrained setting.\
+                              idx starts from 5th column in candidate.tsv file',
+                        default=[])
     args = parser.parse_args()
 
     use_cuda = torch.cuda.is_available()
@@ -607,7 +661,7 @@ def call_neusomatic(candidates_tsv, ref_file, out_dir, checkpoint, num_threads,
                                      args.checkpoint,
                                      args.num_threads, args.batch_size, args.max_load_candidates,
                                      args.pass_threshold, args.lowqual_threshold,
-                                     args.ensemble,
+                                     args.force_zero_ann_cols,
                                      use_cuda)
     except Exception as e:
         logger.error(traceback.format_exc())

neusomatic/python/dataloader.py

@@ -131,6 +131,7 @@ def __init__(self, roots, max_load_candidates, transform=None,
                  num_threads=1, disable_ensemble=False, data_augmentation=False,
                  nclasses_t=4, nclasses_l=4, coverage_thr=100,
                  normalize_channels=False,
+                 zero_ann_cols=[],
                  max_opended_tsv=-1):
 
         soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
@@ -141,6 +142,7 @@ def __init__(self, roots, max_load_candidates, transform=None,
             max_opended_tsv = min(max_opended_tsv, soft)
         self.max_opended_tsv = max_opended_tsv
         self.normalize_channels = normalize_channels
+        self.zero_ann_cols = zero_ann_cols
         self.da_shift_p = 0.3
         self.da_base_p = 0.05
         self.da_rev_p = 0.1
@@ -264,6 +266,11 @@ def __getitem__(self, index):
         if self.disable_ensemble:
             anns = []
 
+        if self.zero_ann_cols and len(anns)>0:
+            anns=np.array(anns)
+            anns[self.zero_ann_cols] = 0
+            anns = anns.tolist()
+
         tag = path.split("/")[-1]
         _, _, _, _, vartype, center, length, tumor_cov, normal_cov = tag.split(
             ".")

neusomatic/python/defaults.py

@@ -1,4 +1,5 @@
-NUM_ENS_FEATURES=93
-VCF_HEADER="##fileformat=VCFv4.2"
+NUM_ENS_FEATURES = 93
+NUM_ST_FEATURES = 26
+VCF_HEADER = "##fileformat=VCFv4.2"
 TYPE_CLASS_DICT = {"DEL": 0, "INS": 1, "NONE": 2, "SNP": 3}
-VARTYPE_CLASSES = ['DEL', 'INS', 'NONE', 'SNP']
+VARTYPE_CLASSES = ['DEL', 'INS', 'NONE', 'SNP']