From e2207e5cd196248e3cc7cc8f9a29c6debc22f051 Mon Sep 17 00:00:00 2001
From: ZL Deng <dawnmsg@gmail.com>
Date: Thu, 21 Apr 2022 17:42:56 +0200
Subject: [PATCH] fixed #13, update readme

---
 README.md                  |   7 +-
 ribodetector/__init__.py   |   2 +-
 ribodetector/detect.py     |  32 ----
 ribodetector/detect_cpu.py | 319 +++++++++++++++++++++++++++++++------
 4 files changed, 271 insertions(+), 89 deletions(-)

diff --git a/README.md b/README.md
index 9f0d24e..6896e70 100644
--- a/README.md
+++ b/README.md
@@ -154,13 +154,12 @@ optional arguments:
   -t THREADS, --threads THREADS
                         number of threads to use. (default: 10)
   --chunk_size CHUNK_SIZE
-                        chunk_size * threads reads to process per thread.(default: 
-                        1024)
-                        When chunk_size=1024 and threads=20, each process will load 
-                        1024 reads, in total consumming ~20G memory.
+                        chunk_size * 1024 reads to load each time.
+                        When chunk_size=1000 and threads=20, consumming ~20G memory, better to be multiples of the number of threads.
   -v, --version         show program's version number and exit
 ```
 
+**Note**: RiboDetector uses multiprocessing with shared memory, thus the memory use of a single process indicated in `htop` or `top` is actually the total memory used by RiboDector. Some job submission system like SGE mis-calculated the total memory use by adding up the memory use of all process. If you see this do not worry it will cause out of memory issue. 
 
 <!-- ### Benchmarks
 
diff --git a/ribodetector/__init__.py b/ribodetector/__init__.py
index ab69de5..e9e2983 100644
--- a/ribodetector/__init__.py
+++ b/ribodetector/__init__.py
@@ -4,5 +4,5 @@
 Accurate and rapid RiboRNA sequences Detector based on deep learning.
 """
 
-__version__ = "0.2.4"
+__version__ = "0.2.5"
 __author__ = 'ZL Deng'
diff --git a/ribodetector/detect.py b/ribodetector/detect.py
index 72ade00..3662f4e 100644
--- a/ribodetector/detect.py
+++ b/ribodetector/detect.py
@@ -163,13 +163,6 @@ def run(self):
 
             num_batches = math.ceil(num_seqs / self.batch_size)
 
-            # Output probability to files
-            # prob_out_fh = open(
-            #     self.output[0].replace('.r1.fq', '') + '.softmax.probability.txt', 'w')
-
-            # prob_out_fh.write(
-            #     '\t'.join(['read', 'r1_0', 'r1_1', 'r2_0', 'r2_1']) + '\n')
-
             data_loader = tqdm(DataLoader(paired_reads_data,
                                           num_workers=self.args.threads,
                                           pin_memory=self.has_cuda,
@@ -184,28 +177,6 @@ def run(self):
                     r2_output = self.model(r2_data.to(
                         self.device, non_blocking=self.has_cuda))
 
-                    # output the predicted probability of two classes
-                    # for read_r1, r1_probs, r2_probs in zip(r1,
-                    #                                        torch.nn.functional.softmax(
-                    #                                            r1_output, dim=1).tolist(),
-                    #                                        torch.nn.functional.softmax(r2_output, dim=1).tolist()):
-                    #     read = read_r1.split('\n')[0].lstrip(
-                    #         '@').rsplit('-', 1)[0]
-                    #     read_probs = [
-                    #         read] + list(map(str, r1_probs)) + list(map(str, r2_probs))
-                    #     prob_out_fh.write('\t'.join(read_probs) + '\n')
-
-                    # for read_r2, r2_probs in zip(r2, torch.nn.functional.softmax(r2_output, dim=1).tolist()):
-                    #     read_r2_probs = [read_r2.split(
-                    #         '\n')[0]] + list(map(str, r2_probs))
-                    #     prob_out2_fh.write('\t'.join(read_r2_probs) + '\n')
-
-                    # r1_batch_labels = torch.argmax(r1_output, dim=1).tolist()
-                    # r2_batch_labels = torch.argmax(r2_output, dim=1).tolist()
-
-                    # r1_dict, r2_dict = self.separate_paired_reads(
-                    #     r1, r1_batch_labels, r2, r2_batch_labels)
-
                     r1_dict, r2_dict = self.separate_paired_reads(
                         r1, r1_output, r2, r2_output)
                     if r1_dict[0]:
@@ -248,9 +219,6 @@ def run(self):
             norrna1_fh.close()
             norrna2_fh.close()
 
-            # close prob out file handle
-            # prob_out_fh.close()
-
         # Single end reads
         else:
             reads_data = SeqData(SeqEncoder.load_reads(*self.input))
diff --git a/ribodetector/detect_cpu.py b/ribodetector/detect_cpu.py
index a99e8cd..7c33d44 100644
--- a/ribodetector/detect_cpu.py
+++ b/ribodetector/detect_cpu.py
@@ -38,9 +38,9 @@ def __init__(self, config, args):
         self.args = args
         self.logger = config.get_logger('predict', 1)
         self.chunk_size = self.args.chunk_size
-        self.input = self.args.input
-        self.output = self.args.output
-        self.rrna = self.args.rrna
+        # self.input = self.args.input
+        # self.output = self.args.output
+        # self.rrna = self.args.rrna
         self.SENTINEL = 1  # progressbar signal
 
     def load_model(self):
@@ -75,7 +75,7 @@ def load_model(self):
                                                                             colors.OKCYAN,
                                                                             self.model_file,
                                                                             colors.ENDC))
-        
+
         so = onnxruntime.SessionOptions()
         so.intra_op_num_threads = 1
         so.inter_op_num_threads = 1
@@ -86,31 +86,12 @@ def load_model(self):
 
         self.model = onnxruntime.InferenceSession(self.model_file, so)
 
-        
-
     def run(self):
         """
         Load data and run the predictor
         """
 
         num_workers = self.args.threads
-        num_inputs = len(self.input)
-        num_rrna_outputs = None if self.rrna is None else len(self.rrna)
-        num_norrna_outputs = len(self.output)
-        if num_inputs != num_norrna_outputs or num_inputs > 2:
-            self.logger.error('{}The number of input and output sequence files is invalid!{}'.format(
-                colors.FAIL,
-                colors.ENDC))
-            raise RuntimeError(
-                "Input or output should have no more than two files and they should have the same number of files.")
-        if num_rrna_outputs != None and num_rrna_outputs != num_inputs:
-            self.logger.error('{}The number of output rRNA sequence files is invalid!{}'.format(
-                colors.FAIL,
-                colors.ENDC))
-            raise RuntimeError(
-                "Ouput rRNA should have no more than two files and they should the same number with input files.")
-
-        is_paired = (num_inputs == 2)
 
         # Manager for multiprocessing
         manager = mp.Manager()
@@ -124,15 +105,15 @@ def run(self):
         # queue for progressbar signal
         q_pbar = mp.Queue()
 
-        if is_paired:
+        if self.is_paired:
             # Load paired end read files with multiprocessing
             with mp.Pool(2) as p:
                 input_reads = p.map(SeqEncoder.load_reads, self.input)
 
             num_seqs = len(input_reads[0])
 
-            self.num_chunks = math.ceil(
-                num_seqs / (self.chunk_size))
+            self.num_batches = math.ceil(
+                num_seqs / (self.batch_size))
 
             self.logger.info('{}{}{}{} sequences loaded!'.format(
                 colors.BOLD,
@@ -182,9 +163,9 @@ def run(self):
                 p.start()
                 pool.append(p)
 
-            # Input reads chunks
-            iters = itertools.chain(Predictor.generate_paired_read_chunks(
-                input_reads, self.chunk_size), (None,) * num_workers)
+            # Input reads batches
+            iters = itertools.chain(Predictor.generate_paired_read_batches(
+                input_reads, self.batch_size), (None,) * num_workers)
             for read in iters:
                 work.put(read)
 
@@ -245,8 +226,8 @@ def run(self):
 
             num_seqs = len(input_reads)
 
-            self.num_chunks = math.ceil(
-                num_seqs / (self.chunk_size))
+            self.num_batches = math.ceil(
+                num_seqs / (self.batch_size))
 
             self.logger.info('{}{}{}{} sequences loaded!'.format(
                 colors.BOLD,
@@ -279,8 +260,8 @@ def run(self):
                 p.start()
                 pool.append(p)
 
-            iters = itertools.chain(Predictor.generate_read_chunks(
-                input_reads, self.chunk_size), (None,) * num_workers)
+            iters = itertools.chain(Predictor.generate_read_batches(
+                input_reads, self.batch_size), (None,) * num_workers)
             for read in iters:
                 work.put(read)
 
@@ -311,15 +292,248 @@ def run(self):
                 rrna_fh.close()
             norrna_fh.close()
 
+    def run_with_chunks(self):
+        """
+        Load data with chunks and run the predictor
+        """
+
+        num_workers = self.args.threads
+
+        read_chunk_size = self.batch_size * self.chunk_size
+
+        if self.is_paired:
+            if self.rrna is not None:
+                self.logger.info('Writing output rRNA sequences into file: {}{}{}'.format(
+                    colors.OKBLUE,
+                    ", ".join(self.rrna),
+                    colors.ENDC))
+                rrna1_fh = open_for_write(self.rrna[0])
+                rrna2_fh = open_for_write(self.rrna[1])
+                num_rrna = 0
+
+            self.logger.info('Writing output non-rRNA sequences into file: {}{}{}'.format(
+                colors.OKBLUE,
+                ", ".join(self.output),
+                colors.ENDC))
+
+            norrna1_fh = open_for_write(self.output[0])
+            norrna2_fh = open_for_write(self.output[1])
+
+            if self.args.ensure == 'both':
+
+                unclf1 = self.output[0] + '.unclassified.gz'
+                unclf2 = self.output[1] + '.unclassified.gz'
+                unclf1_fh = open_for_write(unclf1)
+                unclf2_fh = open_for_write(unclf2)
+                self.logger.info('Writing unclassified sequences into file: {}{}, {}{}'.format(
+                    colors.OKYELLOW,
+                    unclf1,
+                    unclf2,
+                    colors.ENDC))
+
+                num_unknown = 0
+
+            num_read = 0
+
+            # Load paired end reads with chunks
+            for chunk in SeqEncoder.get_pairedread_chunks(*self.input,
+                                                          chunk_size=read_chunk_size):
+                # Manager for multiprocessing
+                manager = mp.Manager()
+
+                # List to store prediction results
+                results = manager.list()
+                work = manager.Queue(num_workers)
+
+                pool = []
+
+                # Start the classification processes
+                for _i in range(num_workers):
+                    p = mp.Process(target=self.classify_paired_reads,
+                                   args=(work, results))
+                    p.start()
+                    pool.append(p)
+
+                # Input reads batches for each chunk
+                iters = itertools.chain(Predictor.generate_paired_read_batches(
+                    chunk, self.batch_size), (None,) * num_workers)
+                for read in iters:
+                    work.put(read)
+
+                for p in pool:
+                    p.join()
+
+                for r1_dict, r2_dict in results:
+                    # Load the prediciton results and split the input reads accordingly
+                    if r1_dict[0]:
+                        norrna1_fh.write('\n'.join(r1_dict[0]) + '\n')
+                        norrna2_fh.write('\n'.join(r2_dict[0]) + '\n')
+                    if self.rrna is not None and r1_dict[1]:
+                        rrna1_fh.write('\n'.join(r1_dict[1]) + '\n')
+                        rrna2_fh.write('\n'.join(r2_dict[1]) + '\n')
+                        num_rrna += len(r1_dict[1])
+
+                    if self.args.ensure == 'both' and r1_dict[-1]:
+                        unclf1_fh.write('\n'.join(r1_dict[-1]) + '\n')
+                        unclf2_fh.write('\n'.join(r2_dict[-1]) + '\n')
+                        num_unknown += len(r1_dict[-1])
+
+                        # del r1_data, r2_data, r1_output, r2_output, r1_batch_labels, r2_batch_labels
+                num_read += len(chunk[0])
+
+                self.logger.info('{}{}{}{} reads classified!'.format(
+                    colors.BOLD,
+                    colors.OKCYAN,
+                    num_read,
+                    colors.ENDC))
+
+            if self.rrna is not None:
+                self.logger.info('Done! Detected {}{}{}{} rRNA sequences.'.format(
+                    colors.BOLD,
+                    colors.OKCYAN,
+                    num_rrna,
+                    colors.ENDC
+                ))
+
+                rrna1_fh.close()
+                rrna2_fh.close()
+
+            if self.args.ensure == 'both':
+                self.logger.info('Discarded {}{}{}{} unclassified sequences.'.format(
+                    colors.BOLD,
+                    colors.OKCYAN,
+                    num_unknown,
+                    colors.ENDC))
+
+                unclf1_fh.close()
+                unclf2_fh.close()
+
+            norrna1_fh.close()
+            norrna2_fh.close()
+
+        else:
+            num_read = 0
+            self.logger.info('Classify paired end reads with chunk size {}{}{}'.format(
+                colors.BOLD,
+                self.chunk_size,
+                colors.ENDC))
+
+            if self.rrna is not None:
+                self.logger.info('Writing output rRNA sequences into file: {}{}{}'.format(
+                    colors.OKBLUE,
+                    ", ".join(self.rrna),
+                    colors.ENDC))
+
+                rrna_fh = open_for_write(self.rrna[0])
+                num_rrna = 0
+
+            self.logger.info('Writing output non-rRNA sequences into file: {}{}{}'.format(
+                colors.OKBLUE,
+                ", ".join(self.output),
+                colors.ENDC))
+
+            norrna_fh = open_for_write(self.output[0])
+
+            for chunk in SeqEncoder.get_seq_chunks(*self.input,
+                                                   chunk_size=read_chunk_size):
+                
+                # Manager for multiprocessing
+                manager = mp.Manager()
+
+                # List to store prediction results
+                results = manager.list()
+                work = manager.Queue(num_workers)
+
+                pool = []
+
+                for _i in range(num_workers):
+                    p = mp.Process(target=self.classify_reads,
+                                   args=(work, results))
+                    p.start()
+                    pool.append(p)
+
+                iters = itertools.chain(Predictor.generate_read_batches(
+                    chunk, self.batch_size), (None,) * num_workers)
+                for read in iters:
+                    work.put(read)
+
+                for p in pool:
+                    p.join()
+
+                for r_dict in results:
+
+                    if r_dict[0]:
+                        norrna_fh.write('\n'.join(r_dict[0]) + '\n')
+                    if self.rrna is not None and r_dict[1]:
+                        rrna_fh.write('\n'.join(r_dict[1]) + '\n')
+                        num_rrna += len(r_dict[1])
+
+                num_read += len(chunk)
+
+                self.logger.info('{}{}{}{} reads classified!'.format(
+                    colors.BOLD,
+                    colors.OKCYAN,
+                    num_read,
+                    colors.ENDC))
+
+            if self.rrna is not None:
+                self.logger.info('Done! Detected {}{}{}{} rRNA sequences'.format(
+                    colors.BOLD,
+                    colors.OKCYAN,
+                    num_rrna,
+                    colors.ENDC
+                ))
+                rrna_fh.close()
+            norrna_fh.close()
+
+
+    def detect(self):
+        """Wrapper for all steps
+
+        Raises:
+            RuntimeError: raise error if the number of input read files and output files are invalid
+        """
+        self.input = self.args.input
+        self.output = self.args.output
+        self.rrna = self.args.rrna
+
+        num_inputs = len(self.input)
+        num_rrna_outputs = None if self.rrna is None else len(self.rrna)
+        num_norrna_outputs = len(self.output)
+        if num_inputs != num_norrna_outputs or num_inputs > 2:
+            self.logger.error('{}The number of input and output sequence files is invalid!{}'.format(
+                colors.FAIL,
+                colors.ENDC))
+            raise RuntimeError(
+                "Input or output should have no more than two files and they should have the same number of files.")
+        if num_rrna_outputs != None and num_rrna_outputs != num_inputs:
+            self.logger.error('{}The number of output rRNA sequence files is invalid!{}'.format(
+                colors.FAIL,
+                colors.ENDC))
+            raise RuntimeError(
+                "Ouput rRNA should have no more than two files and they should the same number with input files.")
+
+        self.batch_size = 1024
+        self.is_paired = (num_inputs == 2)
+
+        if self.chunk_size is None:
+            self.run()
+        else:
+            self.logger.info('Classify reads with chunk size {}{}{}'.format(
+                colors.BOLD,
+                self.chunk_size,
+                colors.ENDC))
+            self.run_with_chunks()
+
     @staticmethod
-    def generate_read_chunks(reads, n):
-        """Yield successive n-sized chunks from read list."""
+    def generate_read_batches(reads, n):
+        """Yield successive n-sized batches from read list."""
         for i in range(0, len(reads), n):
             yield reads[i:i + n]
 
     @staticmethod
-    def generate_paired_read_chunks(reads, n):
-        """Yield successive n-sized chunks from paired end read list."""
+    def generate_paired_read_batches(reads, n):
+        """Yield successive n-sized batches from paired end read list."""
         r1, r2 = reads
         for i in range(0, len(r1), n):
             yield r1[i:i + n], r2[i:i + n]
@@ -393,16 +607,16 @@ def separate_paired_reads(self, r1_reads, r1_outs, r2_reads, r2_outs):
 
         return r1_dict, r2_dict
 
-    def classify_reads(self, chunk, out_list, q_pbar):
-        """Classify reads chunk
+    def classify_reads(self, batch, out_list, q_pbar=None):
+        """Classify reads batch
 
         Args:
-            chunk (list): chunk of reads
+            batch (list): batch of reads
             out_list (list(dict)): list of separated {label: read} dicts
             q_pbar (mp queue): queue for progressbar signal
         """
         while True:
-            reads = chunk.get()
+            reads = batch.get()
             if reads == None:
                 return
 
@@ -414,12 +628,12 @@ def classify_reads(self, chunk, out_list, q_pbar):
 
             out_list.append(Predictor.separate_reads(
                 reads, np.argmax(outputs[0], axis=1)))
+            if q_pbar is not None:
+                q_pbar.put(self.SENTINEL)
 
-            q_pbar.put(self.SENTINEL)
-
-    def classify_paired_reads(self, chunk, out_list, q_pbar):
+    def classify_paired_reads(self, batch, out_list, q_pbar=None):
         while True:
-            reads = chunk.get()
+            reads = batch.get()
             if reads == None:
                 return
 
@@ -440,10 +654,11 @@ def classify_paired_reads(self, chunk, out_list, q_pbar):
             out_list.append(self.separate_paired_reads(
                 r1, output_r1, r2, output_r2))
 
-            q_pbar.put(self.SENTINEL)
+            if q_pbar is not None:
+                q_pbar.put(self.SENTINEL)
 
     def listener(self, q_pbar):
-        pbar = tqdm(total=self.num_chunks)
+        pbar = tqdm(total=self.num_batches)
         for _ in iter(q_pbar.get, None):
             pbar.update()
 
@@ -507,10 +722,9 @@ def main():
     args.add_argument('-t', '--threads', default=20, type=int,
                       help='number of threads to use. (default: 10)')
 
-    args.add_argument('--chunk_size', default=1024, type=int,
-                      help='chunk_size * threads reads to process per thread.(default: 1024) \n{}.'.format(
-                          'When chunk_size=1024 and threads=20, each process will load 1024 reads, in total consumming ~20G memory'
-                      ))
+    args.add_argument('--chunk_size', default=None, type=int,
+                      help='chunk_size * 1024 reads to load each time. \n{}.'.format(
+                          'When chunk_size=1000 and threads=20, consumming ~20G memory, better to be multiples of the number of threads.'))
 
     args.add_argument('-v', '--version', action='version',
                       version='%(prog)s {version}'.format(version=__version__))
@@ -528,7 +742,8 @@ def main():
 
     seq_pred = Predictor(config, args)
     seq_pred.load_model()
-    seq_pred.run()
+
+    seq_pred.detect()
 
 
 if __name__ == '__main__':