From 5ade2d6116bc8169d94acf198c2d2f13c0bb5fab Mon Sep 17 00:00:00 2001
From: Steve Huang <shuang@broadinstitute.org>
Date: Sat, 7 Mar 2020 15:12:15 -0500
Subject: [PATCH] racon_wrapper supports gpu options

---
 scripts/racon_wrapper.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/scripts/racon_wrapper.py b/scripts/racon_wrapper.py
index 6631fee..9b7b45c 100644
--- a/scripts/racon_wrapper.py
+++ b/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,6 +117,7 @@ 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),
@@ -120,6 +125,8 @@ def run(self):
             '-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()