diff --git a/src/plassembler/utils/depth.py b/src/plassembler/utils/depth.py index 72fca27..ce33564 100644 --- a/src/plassembler/utils/depth.py +++ b/src/plassembler/utils/depth.py @@ -3,7 +3,6 @@ # https://github.com/rrwick/Small-plasmid-Nanopore/blob/main/scripts/get_depths.py ######################################### import os -import statistics import subprocess as sp from pathlib import Path @@ -66,22 +65,69 @@ def get_contig_circularity(fasta): return circular_status -def get_depths_from_bam(bam_file: Path, contig_lengths: pd.DataFrame): - """maps runs samtools depth on bam +# rows of `samtools depth` output parsed per chunk. Bounds the transient memory +# of the parse regardless of genome size / coverage. +DEPTH_CHUNK_ROWS = 2_000_000 + + +def get_depths_from_bam(bam_file: Path, contig_lengths: dict): + """runs samtools depth on a bam and returns per-base depth arrays + + `samtools depth` only emits positions with non-zero coverage, so each + contig starts as an all-zero array that the reported positions are scattered + into. + + The output is streamed and parsed in bounded chunks rather than being read + into one big string: for a 5 Mb chromosome at high coverage the old + `check_output` + `splitlines` + list-of-python-ints approach held several + hundred MB, where the int32 arrays below hold 4 bytes per base. + :param bam_file: Path - :param: contig_lengths: dictionary of headers and contig lengths - :return: depths: dictionary of contigs and depths + :param contig_lengths: dictionary of headers and contig lengths + :return: depths: dictionary of contigs and per-base depth arrays """ - depths = {} - for repName, repLength in contig_lengths.items(): - depths[repName] = [0] * repLength - depthCommand = ["samtools", "depth", bam_file] + depths = { + repName: np.zeros(repLength, dtype=np.int32) + for repName, repLength in contig_lengths.items() + } + + depthCommand = ["samtools", "depth", str(bam_file)] with open(os.devnull, "wb") as devNull: - depthOutput = sp.check_output(depthCommand, stderr=devNull).decode() - for line in depthOutput.splitlines(): # parse output - parts = line.strip().split("\t") - repName = parts[0] - depths[repName][int(parts[1]) - 1] = int(parts[2]) + proc = sp.Popen(depthCommand, stdout=sp.PIPE, stderr=devNull) + try: + chunks = pd.read_csv( + proc.stdout, + sep="\t", + header=None, + names=["contig", "pos", "depth"], + dtype={"contig": str, "pos": np.int64, "depth": np.int32}, + chunksize=DEPTH_CHUNK_ROWS, + ) + for chunk in chunks: + for repName, group in chunk.groupby("contig", sort=False): + if repName not in depths: + # the bam was made by mapping against exactly these + # contigs, so this means the fasta and the bam disagree. + # The old code hit a KeyError here; keep failing loudly, + # just with a message that says what went wrong. + raise KeyError( + f"{bam_file} contains alignments to contig " + f"'{repName}', which is not in the reference fasta" + ) + depths[repName][group["pos"].to_numpy() - 1] = group[ + "depth" + ].to_numpy() + except pd.errors.EmptyDataError: + # no aligned bases at all - every contig keeps its zero array + pass + finally: + if proc.stdout is not None: + proc.stdout.close() + returncode = proc.wait() + + if returncode != 0: + raise sp.CalledProcessError(returncode, depthCommand) + return depths @@ -105,16 +151,22 @@ def collate_depths(depths, shortFlag, contig_lengths): # iterate over the contigs for replicon_name, base_depths in depths.items(): replicon_length = contig_lengths[replicon_name] - try: - mean_depth = round(statistics.mean(base_depths), 2) - depth_stdev = round(statistics.stdev(base_depths), 2) + base_depths = np.asarray(base_depths) + # statistics.stdev needed >= 2 observations and statistics.mean >= 1, so + # anything shorter than 2 bases used to raise StatisticsError and report + # NA for all four columns. numpy would return nan instead, so guard. + if base_depths.size < 2: + mean_depth, depth_stdev, q25, q75 = "NA", "NA", "NA", "NA" + else: + mean_depth = round(float(np.mean(base_depths)), 2) + # ddof=1: statistics.stdev is the *sample* standard deviation, and + # numpy defaults to the population one + depth_stdev = round(float(np.std(base_depths, ddof=1)), 2) q25, q75 = np.percentile(base_depths, [25, 75]) q25, q75 = int(q25), int(q75) # save the chromosome depth if replicon_name == "chromosome": chromosome_depth = mean_depth - except statistics.StatisticsError: # if can't calculate - mean_depth, depth_stdev, q25, q75 = "NA", "NA", "NA", "NA" # append to list contig_names.append(replicon_name) contig_length.append(replicon_length) diff --git a/tests/test_deterministic.py b/tests/test_deterministic.py index 04b0db2..104184a 100644 --- a/tests/test_deterministic.py +++ b/tests/test_deterministic.py @@ -5,14 +5,18 @@ fully reproducible, so they act as true regression guards. """ +import shutil from pathlib import Path import numpy as np +import pysam +import pytest from src.plassembler.utils.depth import ( collate_depths, get_contig_circularity, get_contig_lengths, + get_depths_from_bam, ) GOLDEN_FASTA = Path("tests/test_data/golden/contigs.fasta") @@ -76,3 +80,112 @@ def test_collate_depths_zero_chromosome_depth_stays_float(): assert np.isinf(df.loc["plasmid00001", "plasmid_copy_number_short"]) # the actual CI failure: this must not raise df["plasmid_copy_number_short"].astype(float) + + +def test_collate_depths_accepts_numpy_arrays(): + """get_depths_from_bam returns numpy arrays; collate_depths must produce the + same numbers it did for the old python lists (stdev is the *sample* stdev).""" + from_lists = collate_depths( + {"chromosome": [10, 12, 8, 10], "plasmid00001": [20, 20, 20, 20]}, + "short", + {"chromosome": 4, "plasmid00001": 4}, + ) + from_arrays = collate_depths( + { + "chromosome": np.array([10, 12, 8, 10], dtype=np.int32), + "plasmid00001": np.array([20, 20, 20, 20], dtype=np.int32), + }, + "short", + {"chromosome": 4, "plasmid00001": 4}, + ) + assert from_lists.to_csv(index=False) == from_arrays.to_csv(index=False) + # sample stdev of [10, 12, 8, 10] is 1.63, the population one would be 1.41 + assert from_arrays.set_index("contig").loc["chromosome", "sd_depth_short"] == 1.63 + + +def test_collate_depths_short_contig_is_na(): + """Fewer than 2 bases has no sample stdev, so all four columns are NA - the + old code got there via statistics.StatisticsError.""" + df = collate_depths( + {"plasmid00001": np.array([7], dtype=np.int32)}, "long", {"plasmid00001": 1} + ).set_index("contig") + assert df.loc["plasmid00001", "mean_depth_long"] == "NA" + assert df.loc["plasmid00001", "sd_depth_long"] == "NA" + assert df.loc["plasmid00001", "q25_depth_long"] == "NA" + assert df.loc["plasmid00001", "q75_depth_long"] == "NA" + + +@pytest.mark.requires_tool +@pytest.mark.skipif(shutil.which("samtools") is None, reason="samtools not installed") +def test_get_depths_from_bam_exact(tmp_path): + """Per-base depths are scattered into zero-filled arrays: positions samtools + never reports stay 0, and contigs with no alignments at all stay all-zero.""" + header = { + "HD": {"VN": "1.6", "SO": "coordinate"}, + "SQ": [ + {"SN": "chromosome", "LN": 10}, + {"SN": "plasmid00001", "LN": 6}, + {"SN": "plasmid00002", "LN": 4}, # deliberately gets no reads + ], + } + bam_path = tmp_path / "test.bam" + with pysam.AlignmentFile(bam_path, "wb", header=header) as bam: + for ref_id, start, length in [(0, 2, 4), (0, 4, 4), (1, 0, 3)]: + read = pysam.AlignedSegment() + read.query_name = f"r{ref_id}_{start}" + read.query_sequence = "A" * length + read.query_qualities = pysam.qualitystring_to_array("I" * length) + read.flag = 0 + read.reference_id = ref_id + read.reference_start = start + read.mapping_quality = 60 + read.cigarstring = f"{length}M" + bam.write(read) + pysam.index(str(bam_path)) + + depths = get_depths_from_bam( + bam_path, {"chromosome": 10, "plasmid00001": 6, "plasmid00002": 4} + ) + # reads cover chromosome 3-6 and 5-8 (1-based), so 5-6 are doubly covered + np.testing.assert_array_equal( + depths["chromosome"], np.array([0, 0, 1, 1, 2, 2, 1, 1, 0, 0]) + ) + np.testing.assert_array_equal(depths["plasmid00001"], np.array([1, 1, 1, 0, 0, 0])) + np.testing.assert_array_equal(depths["plasmid00002"], np.zeros(4, dtype=np.int32)) + + +@pytest.mark.requires_tool +@pytest.mark.skipif(shutil.which("samtools") is None, reason="samtools not installed") +def test_get_depths_from_bam_empty(tmp_path): + """An alignment-free bam yields all-zero arrays rather than raising on the + empty `samtools depth` stream.""" + header = {"HD": {"VN": "1.6"}, "SQ": [{"SN": "chromosome", "LN": 5}]} + bam_path = tmp_path / "empty.bam" + with pysam.AlignmentFile(bam_path, "wb", header=header): + pass + depths = get_depths_from_bam(bam_path, {"chromosome": 5}) + np.testing.assert_array_equal(depths["chromosome"], np.zeros(5, dtype=np.int32)) + + +@pytest.mark.requires_tool +@pytest.mark.skipif(shutil.which("samtools") is None, reason="samtools not installed") +def test_get_depths_from_bam_rejects_an_unknown_contig(tmp_path): + """A contig in the bam but not in the reference means the two disagree. The + old dict-indexing raised KeyError here; keep failing rather than silently + dropping that contig's coverage.""" + header = {"HD": {"VN": "1.6"}, "SQ": [{"SN": "unexpected_contig", "LN": 8}]} + bam_path = tmp_path / "mismatch.bam" + with pysam.AlignmentFile(bam_path, "wb", header=header) as bam: + read = pysam.AlignedSegment() + read.query_name = "r1" + read.query_sequence = "ACGT" + read.query_qualities = pysam.qualitystring_to_array("IIII") + read.flag = 0 + read.reference_id = 0 + read.reference_start = 0 + read.mapping_quality = 60 + read.cigarstring = "4M" + bam.write(read) + + with pytest.raises(KeyError, match="unexpected_contig"): + get_depths_from_bam(bam_path, {"chromosome": 100})