Skip to content

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

Open
sanjaynagi-eit wants to merge 1 commit into
mainfrom
perf/depth-numpy-streaming
Open

perf(depth): stream samtools depth into numpy arrays (6.3x faster, -133 MB)#1
sanjaynagi-eit wants to merge 1 commit into
mainfrom
perf/depth-numpy-streaming

Conversation

@sanjaynagi-eit

Copy link
Copy Markdown
Owner

Problem

get_depths_from_bam was the single largest chunk of pure-python work in a plassembler run:

  • sp.check_output captured the entire samtools depth output as one decoded python string, then .splitlines() doubled it as a list of strings.
  • It preallocated a python list of ints per base per contig ([0] * repLength).
  • collate_depths then ran statistics.mean / statistics.stdev over those lists. The statistics module works in exact rationals and is very slow at this scale.

In hybrid mode this happens twice (short and long bam) with both dicts alive simultaneously.

Change

  • Stream samtools depth via Popen and parse it with the pandas C parser in bounded 2M-row chunks, scattering into preallocated np.int32 arrays (4 bytes/base).
  • Compute the summary with np.mean / np.std / np.percentile.
  • Surface a non-zero samtools depth exit instead of silently returning zeros.

Preserving behaviour exactly

  • np.std is called with ddof=1: statistics.stdev is the sample standard deviation, numpy defaults to the population one.
  • Contigs shorter than 2 bases report NA for all four columns, which is what the old code produced via StatisticsError.
  • collate_depths still accepts plain python lists, so existing callers and tests are unaffected.

Measurements

Real ONT isolate, 2.8 Mb Flye chromosome, 29 MB sorted bam (each implementation in a fresh process):

before after
get_depths_from_bam 1.67 s 0.95 s 1.8x
collate_depths 4.47 s 0.02 s 186x
total 6.14 s 0.97 s 6.3x
peak RSS 410 MB 277 MB -133 MB

The summary dataframe is byte-identical between the two implementations. End-to-end on the same sample, peak python RSS drops 417 MB -> 280 MB.

Tests

tests/test_deterministic.py gains four cases: numpy arrays and python lists produce identical summaries, the sample-vs-population stdev distinction is pinned, sub-2-base contigs report NA, and get_depths_from_bam is checked against a hand-built bam including a contig with no alignments and an alignment-free bam.

All 73 non-slow tests pass.

get_depths_from_bam read the whole `samtools depth` output into one decoded
python string, split it into a list of strings, and scattered it into a
python list of ints per base per contig. collate_depths then ran
statistics.mean/stdev over those lists.

For a bacterial chromosome at ONT depth that is a few hundred MB of transient
python objects and several seconds of pure-python work, repeated for the short
and long bam in hybrid mode with both dicts alive at once.

Stream the depth output through the pandas C parser in bounded chunks and
scatter it into preallocated int32 arrays (4 bytes/base), then compute the
summary with numpy. np.std needs ddof=1 to match statistics.stdev, which is
the sample standard deviation; contigs shorter than 2 bases still report NA,
which the old code reached via StatisticsError.

Measured on a real 2.8 Mb ONT assembly (29 MB sorted bam):
  get_depths_from_bam  1.67s -> 0.95s
  collate_depths       4.47s -> 0.02s
  total                6.14s -> 0.97s   (6.3x)
  peak RSS            410 MB -> 277 MB
with byte-identical summary output.

An unknown contig - one present in the bam but not in contig_lengths - now
raises with a message naming it. The old dict indexing raised a bare KeyError
there, and the bam is built by mapping against exactly those contigs, so it
means the fasta and the bam disagree and should not be silently dropped.
@sanjaynagi-eit
sanjaynagi-eit force-pushed the perf/depth-numpy-streaming branch from 73d27c0 to 02f8a04 Compare August 13, 2026 22:11
@sanjaynagi-eit

Copy link
Copy Markdown
Owner Author

Updated after a self-review pass.

A contig present in the bam but absent from contig_lengths was being silently skipped. That was wrong: the bam is produced by mapping against exactly those contigs, so a mismatch means the FASTA and the bam disagree, and the original dict indexing raised KeyError there. It now raises with a message naming the offending contig instead of quietly dropping its coverage. Test added.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant