Got bam production up and running again

Author: joshfactorial

neat/read_simulator/runner.py

@@ -111,7 +111,7 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
 
     # Creates files and sets up objects for files that can be written to as needed.
     # Also creates headers for bam and vcf.
-    output_file_writer = OutputFileWriter(options=options, bam_header=bam_header)
+    output_file_writer = OutputFileWriter(options=options, header=bam_header)
 
     # Split file by chunk for parallel analysis or by contig for either parallel or single analysis
     _LOG.info("Splitting reference...")
@@ -140,7 +140,7 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
             # Create local filenames based on fasta indexing scheme.
             fq1 = None
             fq2 = None
-            block_reads_pickle = None
+            bam = None
             if options.produce_fastq:
                 fq1 = current_output_dir / options.fq1.name
                 # validate to double-check we don't have name collisions
@@ -150,14 +150,15 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
                     fq2 = current_output_dir / options.fq2.name
                     validate_output_path(fq2, True, False)
             if options.produce_bam:
-                block_reads_pickle = current_output_dir / "reads.p.gz"
-                validate_output_path(block_reads_pickle, True, False)
-            current_options = options.copy_with_changes(splits_file, current_output_dir, fq1, fq2, block_reads_pickle)
+                bam = current_output_dir / options.bam.name
+                validate_output_path(bam, True, False)
+            current_options = options.copy_with_changes(splits_file, current_output_dir, fq1, fq2, bam)
             if options.threads == 1:
                 idx, contig, local_variants, files_written = read_simulator_single(
                     1,
                     start,
                     current_options,
+                    bam_header,
                     contig,
                     contig_index,
                     input_variants_dict[contig],
@@ -178,6 +179,7 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
                     thread_idx,
                     start,
                     current_options,
+                    bam_header,
                     contig,
                     contig_index,
                     thread_input_variants,
@@ -208,12 +210,12 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
 
     start = time.time()
     if options.produce_bam:
-        stitch_main(output_file_writer, output_files, contig_dict, options.read_len, options.threads)
+        stitch_main(output_file_writer, output_files, options.threads)
     else:
         stitch_main(output_file_writer, output_files)
     _LOG.info(f"It took {time.time()-start} s to write the bam file")
 
-    output_file_writer.close_files()
+    output_file_writer.flush_and_close_files()
     _LOG.info(f"Read simulator complete in {time.time() - analysis_start} s")
 
 def write_final_vcf(

neat/read_simulator/single_runner.py

@@ -2,8 +2,11 @@
 Runner for read-simulator in single-ended mode
 """
 import gzip
+import os
 import pickle
-from Bio import SeqIO
+
+import pysam
+from Bio import SeqIO, bgzf
 import logging
 from pathlib import Path
 
@@ -21,6 +24,7 @@ def read_simulator_single(
         thread_idx: int,
         block_start: int,
         local_options: Options,
+        bam_header: list | None,
         contig_name: str,
         contig_index: int,
         input_variants_local: ContigVariants,
@@ -33,6 +37,7 @@ def read_simulator_single(
     :param thread_idx: index of current thread
     :param block_start: Where on the reference does this block start? For a full contig, this will be 0.
     :param local_options: options for current thread and reference chunk
+    :param bam_header: Pass in the outer bam header.
     :param contig_name: The original list of contig names.
     :param contig_index: The index of the contig which this chunk comes from
     :param input_variants_local: The input variants for this block
@@ -83,7 +88,7 @@ def read_simulator_single(
     # We'll also keep track here of what files we are producing.
     # We don't really need to write out the VCF. We should be able to store it in memory
     local_options.produce_vcf = False
-    local_output_file_writer = OutputFileWriter(options=local_options)
+    local_output_file_writer = OutputFileWriter(options=local_options, header=bam_header)
     """
     Begin Analysis
     """
@@ -100,7 +105,7 @@ def read_simulator_single(
     )
 
     if local_options.produce_fastq or local_options.produce_bam:
-        generate_reads(
+        reads_to_write = generate_reads(
             thread_idx,
             local_seq_record,
             seq_error_model,
@@ -115,11 +120,38 @@ def read_simulator_single(
             contig_index,
         )
 
-    local_output_file_writer.close_files()
+        # TODO write the per-block bam here. For some reason moving it entirely out of the loop slows it down in all cases
+        #   Hopefully this lets us take advantage of multithreading for writing the bams. Should make the final stitching
+        #   easier -- basically, that gets moved to here and there will just be a straight file dump.
+        bam_handle = bgzf.BgzfWriter(local_output_file_writer.bam, 'a', compresslevel=6)
+        for read_data in reads_to_write:
+            read1 = read_data[0]
+            read2 = read_data[1]
+            if read1:
+                local_output_file_writer.write_bam_record(
+                    read1,
+                    contig_index,
+                    bam_handle,
+                    local_options.read_len
+                )
+            if read2:
+                local_output_file_writer.write_bam_record(
+                    read2,
+                    contig_index,
+                    bam_handle,
+                    local_options.read_len
+                )
+        bam_handle.flush()
+        bam_handle.close()
+        sorted_bam = local_output_file_writer.bam.with_suffix(".sorted.bam")
+        pysam.sort("-@", str(local_options.threads), "-o", str(sorted_bam), str(local_output_file_writer.bam))
+        os.rename(str(sorted_bam), str(local_output_file_writer.bam))
+
+    local_output_file_writer.flush_and_close_files()
     file_dict = {
         "fq1": local_output_file_writer.fq1,
         "fq2": local_output_file_writer.fq2,
-        "reads": local_options.reads_pickle,
+        "bam": local_options.bam,
     }
     return (
         thread_idx,

neat/read_simulator/utils/generate_reads.py

@@ -309,12 +309,5 @@ def generate_reads(
         else:
             reads_to_write.append((read_1, None))
 
-    if options.produce_fastq:
-        # _LOG.info(f"Chunk fastq(s) written in: {(time.time() - t)/60:.2f} m")
-        pass
-    if options.produce_bam:
-        with open_output(options.reads_pickle) as reads:
-            pickle.dump(reads_to_write, reads)
-        # _LOG.info(f"Chunk bam(s) written.")
-
     _LOG.info(f"Finished sampling reads for thread {thread_index} in {(time.time() - start_time)/60:.2f} m")
+    return reads_to_write

neat/read_simulator/utils/options.py

@@ -56,7 +56,6 @@ def __init__(self,
                  fq2: Path | None = None,
                  vcf: Path | None = None,
                  bam: Path | None = None,
-                 reads_pickle: Path | None = None,
                  output_prefix: str = "neat_sim",
                  overwrite_output: bool = False,
                  rng_seed: int | None = None,
@@ -180,8 +179,6 @@ def __init__(self,
         self.fq2: Path | None = fq2
         self.vcf: Path | None = vcf
         self.bam: Path | None = bam
-        # kind of a temporary holding pen for bam files
-        self.reads_pickle: Path | None = reads_pickle
 
         # Set the rng for the run
         self.rng = self.set_random_seed()
@@ -337,7 +334,7 @@ def copy_with_changes(self,
                           current_output_dir: Path | None = None,
                           fq1: Path | None = None,
                           fq2: Path | None = None,
-                          reads_pickle: Path | None = None,
+                          bam: Path | None = None
                           ):
         return_options = deepcopy(self)
         if reference is not None:
@@ -348,9 +345,8 @@ def copy_with_changes(self,
             return_options.fq1 = fq1
         if fq2 is not None:
             return_options.fq2 = fq2
-        if reads_pickle is not None:
-            return_options.reads_pickle = reads_pickle
-            self.bam = None
+        if bam is not None:
+            return_options.bam = bam
         return return_options
 
     def set_random_seed(self) -> Generator:

neat/read_simulator/utils/output_file_writer.py

@@ -8,22 +8,24 @@
     "OutputFileWriter"
 ]
 
+import os
 import re
 from struct import pack
 import logging
+from typing import Any
 
 from Bio import bgzf
 from pathlib import Path
 
+from Bio.bgzf import BgzfWriter
+
 from .read import Read
 from .options import Options
 
 _LOG = logging.getLogger(__name__)
 
 
 # Some Constants
-# TODO make bam compression a configurable option
-BAM_COMPRESSION_LEVEL = 6
 CIGAR_PACKED = {'M': 0, 'I': 1, 'D': 2, 'N': 3, 'S': 4, 'H': 5, 'P': 6, '=': 7, 'X': 8}
 SEQ_PACKED = {'=': 0, 'A': 1, 'C': 2, 'M': 3, 'G': 4, 'R': 5, 'S': 6, 'V': 7,
               'T': 8, 'W': 9, 'Y': 10, 'H': 11, 'K': 12, 'D': 13, 'B': 14, 'N': 15}
@@ -62,16 +64,16 @@ class OutputFileWriter:
     in the various formats.
 
     :param options: Options for the current run.
-    :param bam_header: A dictionary of lengths of each contig from the reference, keyed by contig id.
+    :param header: A dictionary of lengths of each contig from the reference, keyed by contig id.
     """
     def __init__(self,
                  options: Options,
-                 bam_header: dict = None):
+                 header: dict = None):
 
         self.paired_ended = options.paired_ended
-        self.bam_header = bam_header
+        self.bam_header = header
 
-        file_handles = {}
+        file_handles: dict[Path, Any] = {}
 
         # Set up filenames based on booleans
         if options.fq1 is not None:
@@ -91,7 +93,6 @@ def __init__(self,
             vcf = None
         if options.bam is not None:
             bam = options.bam
-            file_handles[bam] = bgzf.BgzfWriter(bam, 'w', compresslevel=BAM_COMPRESSION_LEVEL)
         else:
             bam = None
 
@@ -119,26 +120,27 @@ def __init__(self,
                          f'#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tNEAT_simulated_sample\n'
             self.files_to_write[self.vcf].write(vcf_header)
 
-        if options.produce_bam and self.bam_header:
+        if options.produce_bam:
             # bam header
-            bam_handle = self.files_to_write[self.bam]
+            bam_handle = bgzf.BgzfWriter(self.bam, 'w', compresslevel=6)
             bam_handle.write("BAM\1")
             # Without a header, we can't write these as bams.
-            bam_header = "@HD\tVN:1.4\tSO:coordinate\n"
+            header = "@HD\tVN:1.4\tSO:coordinate\n"
             for item in self.bam_header:
-                bam_header += f'@SQ\tSN:{item}\tLN:{str(self.bam_header[item])}\n'
-            bam_header += "@RG\tID:NEAT\tSM:NEAT\tLB:NEAT\tPL:NEAT\n"
-            header_bytes = len(bam_header)
+                header += f'@SQ\tSN:{item}\tLN:{str(self.bam_header[item])}\n'
+            header += "@RG\tID:NEAT\tSM:NEAT\tLB:NEAT\tPL:NEAT\n"
+            header_bytes = len(header)
             num_refs = len(self.bam_header)
             bam_handle.write(pack('<i', header_bytes))
-            bam_handle.write(bam_header)
+            bam_handle.write(header)
             bam_handle.write(pack('<i', num_refs))
             # Contigs and lengths. If we can skip writing this out for intermediate files, great
             for item in self.bam_header:
                 name_length = len(item) + 1
                 bam_handle.write(pack('<i', name_length))
                 bam_handle.write(f'{item}\0')
                 bam_handle.write(pack('<i', self.bam_header[item]))
+            bam_handle.flush()
             bam_handle.close()
 
     def write_fastq_record(self, filename: Path, record: str):
@@ -148,10 +150,15 @@ def write_fastq_record(self, filename: Path, record: str):
             _LOG.error(f"Tried to write fastq record to unknown file {filename}")
             raise ValueError
 
-    def close_files(self):
+    def flush_and_close_files(self):
         for file_name in self.files_to_write:
             file_handle = self.files_to_write[file_name]
-            file_handle.close()
+            try:
+                file_handle.flush()
+                os.fsync(file_handle.fileno())
+                file_handle.close()
+            except ValueError:
+                _LOG.debug(f"file {file_name} already closed")
 
     def write_vcf_record(self, line: str):
         """
@@ -171,14 +178,15 @@ def write_bam_record(
             self,
             read: Read,
             contig_id: int,
-            bam_handle,
+            bam_handle: BgzfWriter,
             read_length: int
     ):
         """
         Takes a read object and writes it out as a bam record
 
         :param read: Read object to write to file
         :param contig_id: the index of the reference for this
+        :param bam_handle: the handle to write data to
         :param read_length: the length of the read to output
         """
         read_bin = reg2bin(read.position, read.end_point)