Add total time to basecaller progress.

bonito/cli/basecaller.py

@@ -91,14 +91,15 @@ def main(args):
         exit(1)
 
     if fmt.name != 'fastq':
-        groups = reader.get_read_groups(
+        groups, num_reads = reader.get_read_groups(
             args.reads_directory, args.model_directory,
             n_proc=8, recursive=args.recursive,
             read_ids=column_to_set(args.read_ids), skip=args.skip,
             cancel=process_cancel()
         )
     else:
         groups = []
+        num_reads = None
 
     reads = reader.get_reads(
         args.reads_directory, n_proc=8, recursive=args.recursive,
@@ -137,7 +138,8 @@ def main(args):
         results = align_map(aligner, results, n_thread=args.alignment_threads)
 
     writer = ResultsWriter(
-        fmt.mode, tqdm(results, desc="> calling", unit=" reads", leave=False),
+        fmt.mode, tqdm(results, desc="> calling", unit=" reads", leave=False,
+                       total=num_reads, smoothing=0, ascii=True, ncols=100),
         aligner=aligner, group_key=args.model_directory,
         ref_fn=args.reference, groups=groups,
     )

bonito/fast5.py

@@ -105,6 +105,7 @@ def get_read_groups(directory, model, read_ids=None, skip=False, n_proc=1, recur
     Get all the read meta data for a given `directory`.
     """
     groups = set()
+    num_reads = 0
     pattern = "**/*.fast5" if recursive else "*.fast5"
     fast5s = [Path(x) for x in glob(directory + "/" + pattern, recursive=True)]
     get_filtered_meta_data = partial(get_meta_data, read_ids=read_ids, skip=skip)
@@ -115,7 +116,8 @@ def get_read_groups(directory, model, read_ids=None, skip=False, n_proc=1, recur
                 desc="> preprocessing reads", unit=" fast5s", ascii=True, ncols=100
         ):
             groups.update({read.readgroup(model) for read in reads})
-        return groups
+            num_reads += len(reads)
+        return groups, num_reads
 
 
 def get_read_ids(filename, read_ids=None, skip=False):

bonito/pod5.py

@@ -91,6 +91,7 @@ def get_read_groups(directory, model, read_ids=None, skip=False, n_proc=1, recur
     Get all the read meta data for a given `directory`.
     """
     groups = set()
+    num_reads = 0
     pattern = "**/*.pod5" if recursive else "*.pod5"
     pod5_files = (Path(x) for x in glob(directory + "/" + pattern, recursive=True))
 
@@ -101,7 +102,8 @@ def get_read_groups(directory, model, read_ids=None, skip=False, n_proc=1, recur
         ):
             read = Read(read, pod5_file, meta=True)
             groups.add(read.readgroup(model))
-    return groups
+            num_reads += 1
+    return groups, num_reads
 
 
 def get_reads(directory, read_ids=None, skip=False, n_proc=1, recursive=False, cancel=None):