perf(depth): stream samtools depth into numpy arrays (6.3x faster, -133 MB)

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)

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})