Merge pull request #25 from SHuang-Broad/sh_wrapper_gpu

racon_wrapper supports gpu options

scripts/racon_wrapper.py

@@ -15,7 +15,8 @@ class RaconWrapper:
 
     def __init__(self, sequences, overlaps, target_sequences, split, subsample,
         include_unpolished, fragment_correction, window_length, quality_threshold,
-        error_threshold, match, mismatch, gap, threads):
+        error_threshold, match, mismatch, gap, threads,
+        cudaaligner_batches, cudapoa_batches, cuda_banded_alignment):
 
         self.sequences = os.path.abspath(sequences)
         self.subsampled_sequences = None
@@ -34,6 +35,9 @@ def __init__(self, sequences, overlaps, target_sequences, split, subsample,
         self.mismatch = mismatch
         self.gap = gap
         self.threads = threads
+        self.cudaaligner_batches = cudaaligner_batches
+        self.cudapoa_batches = cudapoa_batches
+        self.cuda_banded_alignment = cuda_banded_alignment
         self.work_directory = os.getcwd() + '/racon_work_directory_' + str(time.time())
 
     def __enter__(self):
@@ -113,13 +117,16 @@ def run(self):
         racon_params = [RaconWrapper.__racon]
         if (self.include_unpolished == True): racon_params.append('-u')
         if (self.fragment_correction == True): racon_params.append('-f')
+        if (self.cuda_banded_alignment == True): racon_params.append('-b')
         racon_params.extend(['-w', str(self.window_length),
             '-q', str(self.quality_threshold),
             '-e', str(self.error_threshold),
             '-m', str(self.match),
             '-x', str(self.mismatch),
             '-g', str(self.gap),
             '-t', str(self.threads),
+            '--cudaaligner-batches', str(self.cudaaligner_batches),
+            '-c', str(self.cudapoa_batches),
             self.subsampled_sequences, self.overlaps, ""])
 
         for target_sequences_part in self.split_target_sequences:
@@ -179,13 +186,17 @@ def run(self):
     parser.add_argument('-g', '--gap', default=-8, help='''gap penalty (must be
         negative)''')
     parser.add_argument('-t', '--threads', default=1, help='''number of threads''')
+    parser.add_argument('--cudaaligner-batches', default=0, help='''number of batches for CUDA accelerated alignment''')
+    parser.add_argument('-c', '--cudapoa-batches', default=0, help='''number of batches for CUDA accelerated polishing''')
+    parser.add_argument('-b', '--cuda-banded-alignment', action='store_true', help='''use banding approximation for polishing on GPU. Only applicable when -c is used.''')
 
     args = parser.parse_args()
 
     racon = RaconWrapper(args.sequences, args.overlaps, args.target_sequences,
         args.split, args.subsample, args.include_unpolished,
         args.fragment_correction, args.window_length, args.quality_threshold,
-        args.error_threshold, args.match, args.mismatch, args.gap, args.threads)
+        args.error_threshold, args.match, args.mismatch, args.gap, args.threads,
+        args.cudaaligner_batches, args.cudapoa_batches, args.cuda_banded_alignment)
 
     with racon:
         racon.run()