Error improvements

neat/models/error_models.py

@@ -10,6 +10,7 @@
 
 from Bio.Seq import Seq
 from Bio import SeqRecord
+from numpy import median
 
 from neat import variants
 
@@ -165,8 +166,9 @@ def __init__(
     def get_sequencing_errors(
             self,
             padding: int,
-            reference_segment: SeqRecord,
+            reference_segment: Seq,
             quality_scores: np.ndarray,
+            num_errors,
             rng
     ):
         """
@@ -175,8 +177,9 @@ def get_sequencing_errors(
         :param padding: this is the amount of space we have in the read for deletions.
         :param reference_segment: The section of the reference from which the read is drawn
         :param quality_scores: Array of quality scores for the read
-        :return: Modified sequence and associated quality scores
+        :param num_errors: The estimated number of errors to add.
         :param rng: random number generator.
+        :return: Modified sequence and associated quality scores
         """
 
         error_indexes = []
@@ -189,9 +192,20 @@ def get_sequencing_errors(
         if self.average_error == 0:
             return introduced_errors
         else:
-            for i in range(len(quality_scores)):
-                if rng.random() < quality_score_error_rate[quality_scores[i]]:
-                    error_indexes.append(i)
+            i = len(quality_scores)
+            while len(error_indexes) < num_errors and i > 0:
+                index = rng.choice(list(range(len(quality_scores))))
+                if rng.random() < quality_score_error_rate[quality_scores[index]]:
+                    error_indexes.append(index)
+                i -= 1
+            # Fallback: if quality scores are too high to naturally reach num_errors,
+            # force errors at positions with at-or-below-median quality scores.
+            # Using <= so that uniform quality arrays (all scores equal) always make progress.
+            median_score = median(quality_scores)
+            while len(error_indexes) < num_errors:
+                index = rng.integers(len(quality_scores))
+                if quality_scores[index] <= median_score:
+                    error_indexes.append(index)
 
         total_indel_length = 0
         # To prevent deletion collisions
@@ -205,27 +219,27 @@ def get_sequencing_errors(
 
             # This is to prevent deletion error collisions and to keep there from being too many indel errors.
             if 0 < index < self.read_length - max(
-                    self.deletion_len_model) and total_indel_length > self.read_length // 4:
+                    self.deletion_len_model) and total_indel_length <= self.read_length // 4:
                 error_type = rng.choice(a=list(self.variant_probs), p=list(self.variant_probs.values()))
 
             # Deletion error
             if error_type == Deletion:
-                deletion_length = self.get_deletion_length()
+                deletion_length = self.get_deletion_length(rng)
                 if padding - deletion_length < 0:
                     # No space in this read to add this deletion
                     continue
-                deletion_reference = reference_segment.seq[index: index + deletion_length + 1]
+                deletion_reference = reference_segment[index: index + deletion_length + 1]
                 deletion_alternate = deletion_reference[0]
                 introduced_errors.append(
                     ErrorContainer(Deletion, index, deletion_length, deletion_reference, deletion_alternate)
                 )
                 total_indel_length += deletion_length
 
-                del_blacklist.extend(list(range(index, index + deletion_length)))
+                del_blacklist.extend(list(range(index + 1, index + deletion_length + 1)))
                 padding -= deletion_length
 
             elif error_type == Insertion:
-                insertion_length = self.get_insertion_length()
+                insertion_length = self.get_insertion_length(rng)
                 insertion_reference = reference_segment[index]
                 insert_string = ''.join(rng.choice(ALLOWED_NUCL, size=insertion_length))
                 insertion_alternate = insertion_reference + insert_string

neat/read_simulator/runner.py

@@ -1,11 +1,15 @@
 """
 Runner for generate_reads task
 """
+import gzip
 import logging
+import os
+import pickle
 import shutil
 import subprocess
 import time
 import multiprocessing as mp
+from math import ceil
 
 from pathlib import Path
 
@@ -16,6 +20,7 @@
 from .utils import Options, OutputFileWriter, parse_beds, parse_input_vcf
 from ..common import validate_input_path, validate_output_path
 from .single_runner import read_simulator_single
+from ..models import SequencingErrorModel
 from ..variants import ContigVariants
 from .utils.split_inputs import main as split_main
 from .utils.stitch_outputs import main as stitch_main
@@ -74,6 +79,31 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
     reference_index = SeqIO.index(str(options.reference), "fasta")
     reference_keys_with_lens = {key: len(value) for key, value in reference_index.items()}
 
+    # We need sequencing errors to get the quality score attributes, even for the vcf
+    if options.avg_seq_error:
+        average_error = options.avg_seq_error
+    elif options.error_model:
+        error_models = pickle.load(gzip.open(options.error_model))
+        average_error = error_models["error_model1"].average_error
+        # We just need the error value
+        del error_models
+    else:
+        # Use the default value
+        average_error = 0.009228843915252066
+
+    # _LOG.debug('Sequencing error and quality score models loaded')
+    # We need to estimate how many total errors to add
+    total_reference_length = sum(reference_keys_with_lens.values())
+    # This is an estimate due to some random effects
+    number_of_bases_in_analysis = total_reference_length * options.coverage
+    # Each base called has an "average_error" chance of being an error
+    total_errors = ceil(average_error * number_of_bases_in_analysis)
+    # Normalization gives the percent value for each item with a sum of all values being 1.0
+    normalized_counts = {k: v / total_reference_length for (k, v) in reference_keys_with_lens.items()}
+    # Multiply the normalized count by the total errors. This gives the number of errors that should be
+    # introduced into the contig
+    errors_per_contig = {k: ceil(v * total_errors) for (k, v) in normalized_counts.items()}
+
     count = 0
     for contig in reference_keys_with_lens:
         count += reference_keys_with_lens[contig]
@@ -141,6 +171,14 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
     for contig in splits_files_dict:
         contig_index = contig_dict[contig]
         for ((start, length), splits_file) in splits_files_dict[contig].items():
+            block_percentage = length / reference_keys_with_lens[contig]
+            block_errors = errors_per_contig[contig] * block_percentage
+            estimated_number_of_reads = (length // options.read_len) * options.coverage
+            errors_per_read = round(block_errors / estimated_number_of_reads)
+            if errors_per_read < 1.0 and block_errors > 0:
+                # We know we need a few errors, but it's a small number total
+                if options.rng.random() < average_error:
+                    errors_per_read += 1
             current_output_dir = options.temp_dir_path / splits_file.stem
             current_output_dir.mkdir(parents=True, exist_ok=True)
             # Create local filenames based on fasta indexing scheme.
@@ -175,6 +213,7 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
                     target_regions_dict[contig],
                     discard_regions_dict[contig],
                     mutation_rate_dict[contig],
+                    errors_per_read,
                 )
                 _LOG.info(f"Completed simulating contig {contig}.")
                 # TODO Remove if not needed
@@ -196,7 +235,8 @@ def read_simulator_runner(config: str, output_dir: str, file_prefix: str):
                     thread_input_variants,
                     thread_target_regions,
                     thread_discard_regions,
-                    thread_mutation_regions
+                    thread_mutation_regions,
+                    errors_per_read,
                 ))
             thread_idx += 1
 

neat/read_simulator/single_runner.py

@@ -11,6 +11,8 @@
 import logging
 from pathlib import Path
 
+from math import ceil, floor
+
 from .utils import OutputFileWriter, \
     generate_variants, generate_reads, Options, recalibrate_mutation_regions
 from ..variants import ContigVariants
@@ -25,28 +27,28 @@ def read_simulator_single(
         thread_idx: int,
         block_start: int,
         local_options: Options,
-        bam_header: list | None,
+        bam_header: dict | None,
         contig_name: str,
         contig_index: int,
         input_variants_local: ContigVariants,
         target_regions: list,
         discard_regions: list,
         mutation_regions: list,
+        errors_per_read: int,
 ) -> tuple[int, str, ContigVariants, dict[str, Path], ]:
     """
     inputs:
     :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 bam_header: Pass in the outer bam header, which is a dictionary of the contigs plus lengths.
     :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
-    TODO I'm counting on the target and discard regions not being used. They are likely broken with multithreading.
-            Probably they make more sense after the fact now
     :param target_regions: Target regions for the run
     :param discard_regions:  discard regions for the run
     :param mutation_regions: mutation regions (unchecked) for the run
+    :param errors_per_read: How many errors this contig will receive
 
     Ideally this should work for either a file chunk or contig. We'll assume here that we're
     getting either an entire contig or a file chunk, and that no new subdivisions are needed.
@@ -113,6 +115,7 @@ def read_simulator_single(
             thread_idx,
             local_seq_record,
             seq_error_model,
+            errors_per_read,
             qual_score_model,
             fraglen_model,
             local_variants,

neat/read_simulator/utils/generate_reads.py

@@ -149,6 +149,7 @@ def generate_reads(
         thread_index: int,
         reference: SeqRecord,
         error_model: SequencingErrorModel,
+        errors_per_read: int,
         qual_model: TraditionalQualityModel,
         fraglen_model: FragmentLengthModel,
         contig_variants: ContigVariants,
@@ -166,6 +167,7 @@ def generate_reads(
     :param thread_index: Index of current thread
     :param reference: The reference segment that reads will be drawn from.
     :param error_model: The error model for this run, the forward strand
+    :param errors_per_read: Total number of errors to add to contig
     :param qual_model: The quality score model for this run, forward strand
     :param fraglen_model: The fragment length model for this run
     :param contig_variants: An object containing all input and randomly generated variants to be included.
@@ -285,9 +287,9 @@ def generate_reads(
             fastq_handle,
             options.quality_offset,
             options.produce_fastq,
+            errors_per_read,
             options.rng
         )
-
         # skip over read 2 for single ended reads.
         if options.paired_ended:
             # Padding, as above
@@ -313,7 +315,6 @@ def generate_reads(
             )
 
             read_2.mutations = find_applicable_mutations(read_2, contig_variants)
-
             if options.produce_fastq:
                 fastq_handle = ofw.files_to_write[ofw.fq2]
             else:
@@ -324,6 +325,7 @@ def generate_reads(
                 fastq_handle,
                 options.quality_offset,
                 options.produce_fastq,
+                errors_per_read,
                 options.rng
             )
             reads_to_write.append((read_1, read_2))

neat/read_simulator/utils/options.py

@@ -286,7 +286,11 @@ def check_and_log_error(keyname, value_to_check, crit1, crit2):
             if value_to_check not in crit2:
                 _LOG.error(f"Must choose one of {crit2}")
                 sys.exit(1)
-        elif isinstance(crit1, (int, float)) and isinstance(crit2, (int, float)):
+        elif not crit1 and not crit2:
+            # Nothing to check
+            pass
+        else:
+            # Must be a range
             if not (crit1 <= value_to_check <= crit2):
                 _LOG.error(f'`{keyname}` must be between {crit1} and {crit2} (input: {value_to_check}).')
                 sys.exit(1)

neat/read_simulator/utils/read.py

@@ -71,7 +71,7 @@ def __init__(self,
         self.read_sequence: Seq = Seq("")  # initialize to empty sequence
         self.errors: list[ErrorContainer] = []  # initialize
         self.mutations: dict[int, list] = {}  # initialize
-        self.quality_array: np.ndarray = np.zeros(self.run_read_length)  # this will have the correct memory length
+        self.quality_array: np.ndarray = np.zeros(self.run_read_length, dtype=int)  # this will have the correct memory length
         self.mapping_quality: int = 0  # initialize at 0
         self.read_quality_string: str = ""  # This will hold the read quality string
         self.num_ns = 0
@@ -161,7 +161,7 @@ def update_quality_array(
         # Replace the given quality score with the new one
         self.quality_array = \
             np.concatenate((self.quality_array[:location],
-                            np.array(new_quality_score),
+                            np.array(new_quality_score, dtype=int),
                             self.quality_array[location+ref_length:]))
 
     def apply_errors(self, quality_model: TraditionalQualityModel):
@@ -308,6 +308,7 @@ def finalize_read_and_write(
             fastq_handle,
             quality_offset: int,
             produce_fastq: bool,
+            num_errors: int,
             rng: Generator,
     ):
         """
@@ -319,6 +320,7 @@ def finalize_read_and_write(
         :param quality_offset: the quality offset for this run
         :param produce_fastq: If true, this will write out the temp fastqs. If false, this will only write out the tsams
             to create the bam files.
+        :param num_errors: Estimated number of errors to add to this read.
         :param rng: the random number generator for this run
         """
 
@@ -336,11 +338,12 @@ def finalize_read_and_write(
         # set the read sequence to match the reference, then modify
         self.read_sequence = deepcopy(self.reference_segment)
 
-        # Get errors for the rea and update the quality score
+        # Get errors for the read and update the quality score
         self.errors, self.padding = err_model.get_sequencing_errors(
             self.padding,
             self.reference_segment,
             self.quality_array,
+            num_errors,
             rng
         )
 
@@ -362,6 +365,8 @@ def finalize_read_and_write(
             fastq_record = f'@{self.name}\n{str(self.read_sequence)}\n+\n{self.read_quality_string}\n'
             fastq_handle.write(fastq_record)
 
+        return len(self.errors)
+
 
     def convert_masking(self, quality_model: TraditionalQualityModel):
         """

tests/test_models/test_error_and_mut_models.py

@@ -5,7 +5,6 @@
 import numpy as np
 from numpy.random import default_rng
 from Bio.Seq import Seq
-from Bio.SeqRecord import SeqRecord
 
 from neat.models.mutation_model import MutationModel
 from neat.models.error_models import SequencingErrorModel, TraditionalQualityModel
@@ -57,10 +56,10 @@ def test_traditional_quality_model_shapes_and_range():
 def test_sequencing_error_model_basic_snvs_only():
     rng = default_rng(6)
     sem = SequencingErrorModel(avg_seq_error=0.05)
-    ref = SeqRecord(Seq("ACGT" * 20), id="chr1")
+    ref = Seq("ACGT" * 20)
     quals = np.array([35] * 80, dtype=int)
     introduced, pad = sem.get_sequencing_errors(
-        padding=40, reference_segment=ref, quality_scores=quals, rng=rng
+        padding=40, reference_segment=ref, quality_scores=quals, num_errors=3, rng=rng
     )
     assert isinstance(introduced, list)
     assert pad >= 0
@@ -124,17 +123,17 @@ def test_mutation_model_snv_does_not_keep_reference_base():
 def test_sequencing_error_model_reproducible_with_seed():
     """Error placement should be deterministic given the same RNG state."""
     sem = SequencingErrorModel(avg_seq_error=0.05)
-    ref = SeqRecord(Seq("ACGT" * 30), id="chr1")
+    ref = Seq("ACGT" * 30)
     quals = np.array([30] * len(ref), dtype=int)
 
     rng1 = default_rng(9)
     rng2 = default_rng(9)
 
     introduced1, pad1 = sem.get_sequencing_errors(
-        padding=20, reference_segment=ref, quality_scores=quals, rng=rng1
+        padding=20, reference_segment=ref, quality_scores=quals, num_errors=3, rng=rng1
     )
     introduced2, pad2 = sem.get_sequencing_errors(
-        padding=20, reference_segment=ref, quality_scores=quals, rng=rng2
+        padding=20, reference_segment=ref, quality_scores=quals, num_errors=3, rng=rng2
     )
 
     assert pad1 == pad2
@@ -150,11 +149,11 @@ def test_sequencing_error_model_nonzero_error_introduces_in_bounds_errors():
     """
     rng = default_rng(10)
     sem = SequencingErrorModel(avg_seq_error=0.2)
-    ref = SeqRecord(Seq("ACGT" * 50), id="chr1")
+    ref = Seq("ACGT" * 50)
     quals = np.array([10] * len(ref), dtype=int)
 
     introduced, pad = sem.get_sequencing_errors(
-        padding=20, reference_segment=ref, quality_scores=quals, rng=rng
+        padding=20, reference_segment=ref, quality_scores=quals, num_errors=3, rng=rng
     )
 
     # At least one error is expected for these settings.