perf(delfi): ~80x faster via blacklist preloading and binary search

[DucoG]

Summary

delfi is roughly 80x faster on whole-genome inputs while producing bit-identical output. On a chr1+chr2 hg38 benchmark (4912 100kb bins, 16 worker processes), wall-clock time drops from 643s → 8s and total CPU time from 9845s → 59s.

Motivation

While running DELFI on 211 WGS samples I profiled the worker function and noticed almost all of the per-window time was spent re-reading the blacklist BED file from disk and doing per-fragment linear scans against it. Per sample the existing implementation does ~26K full re-parses of the blacklist file (one per 100kb window) and reopens the BAM and 2bit reference inside every worker call. With many samples processed in parallel on shared/NFS storage this also produces a heavy and mostly-redundant I/O load.

Changes

src/finaletoolkit/frag/_delfi.py:

1. Preload the blacklist once and index regions by contig as sorted numpy arrays. Eliminates ~26K redundant disk reads per sample.
2. Binary search the blacklist per window via np.searchsorted instead of a linear scan over every fragment.
3. Share BAM and 2bit handles per worker via the multiprocessing.Pool initializer (one open per worker, reused for every window). The handle is reused by passing the open pysam.AlignmentFile to frag_generator, which already supports it.
4. Pre-load ContigGaps into worker globals instead of pickling one into every task argument.
5. Vectorise GC counting using str.count('G') + str.count('C') in place of a Python sum over 'G' or 'C' comparisons.

tests/test_delfi.py:

• Adds test_workers_equivalence which asserts delfi(workers=1) and delfi(workers=4) produce bit-identical output on every column. This guards against any future regression that introduces parallel-only state in the worker pool.

CHANGELOG.md: documents the speedup and the new test under [Unreleased].

The public API (delfi(...) signature, output schema, output values) is unchanged.

Benchmarks

Hardware: 56-core x86_64, NFS-backed BAM, hg38, 16 workers, identical inputs across runs. Wall and CPU times are means over 5 runs each (σ < 1% in every cell).

Implementation	Wall (s)	CPU (s)	Speedup (wall)
Upstream (current main)	643	9845	1×
This PR	8.0	59.3	80×

Ablation (which optimisation contributes what)

Cumulative speedup as each optimisation is added on top of the previous:

Config	Wall (s)	CPU (s)	Speedup
baseline	643	9845	1×
+ preload blacklist	96	1378	6.7×
+ sorted-array blacklist	9.6	85	67×
+ reuse BAM/ref handle	8.5	65	76×
+ preload contig gaps	8.0	60	80×
+ vectorised GC count	8.0	59	80×

The two dominant wins are preloading the blacklist (eliminates ~26K file reopens) and switching to a sorted-array blacklist filter (eliminates O(blacklist) work per fragment). The remaining three optimisations together are worth ~1.5s wall / ~25s CPU on this input.

Correctness

• Existing tests/test_delfi.py::test_overall passes unchanged.
• New tests/test_delfi.py::test_workers_equivalence verifies workers=1 and workers=4 produce identical output on every column.
• Independently verified on a real 211-sample WGS dataset by running the upstream implementation and this PR side-by-side and diffing the output TSVs: every value of every column (contig, start, stop, arm, short, long, gc, num_frags) matches exactly.

Notes

• I considered an alternative restructuring (one task per contig with a single pysam.fetch per contig instead of per window) hoping to reduce I/O. It was actually slower in every regime (cold and warm cache) because the existing per-window fetches use the BAI index and skip gap regions, while a per-contig fetch reads the full contig. Not pursued.
• The same patterns (per-task BAM/refseq reopens, etc.) exist in _end_motifs, _frag_length, _coverage, _wps, and _multi_wps. Happy to send follow-up PRs that extract a small shared _pool helper and apply the same pattern there if there's appetite for it.

[jamesli124]

Thank you for this PR. The performance improvement is very welcome and clever. I've left a couple of notes for you in my review, the main one being that support for fragment input in tabix-indexed bed.gz was inadvertently dropped. Otherwise, we appreciate your contribution!

[jamesli124]

The new docstring says "See original docstring for full parameter
documentation", but the original docstring was deleted in this same commit. Users calling help(delfi) or reading rendered
docs will get no parameter descriptions at all. Please either restore the
parameter docs or link to the documentation site.

[jamesli124]

This unconditionally opens the input as a BAM/SAM file, but frag_generator
(which the rest of the codebase uses) also supports tabix-indexed .frag.gz/bed.gz
files by opening them as pysam.TabixFile. When given a .frag.gz/bed.gz input,
AlignmentFile raises "file does not contain alignment data" in every worker
process, causing the pool to hang silently with an empty output.

The fix is to mirror the detection already in frag_generator:

  s = str(input_file)
  if s.endswith('.sam') or s.endswith('.bam') or s.endswith('.cram'):
      _WORKER_BAM = pysam.AlignmentFile(s, 'r')
  else:
      _WORKER_BAM = pysam.TabixFile(s, 'r')

The existing tests pass because they exclusively use .bam input — this code
path is never exercised for .frag.gz.

[edawson]

@DucoG please consider using the AlignmentWrapper class, which should be compatible with all alignment file / frag.(gz/txt) inputs:

FinaleToolkit/src/finaletoolkit/io/alignment.py

Line 30 in 63d5c52

class AlignmentWrapper:

[edawson]

I also think this can be done without opening a global BAM wrapper (yikes).

1. read the BED intervals (for blocklist / passlist) to Dict[str, numpy] structure.
2. Do some set algebra to remove the blocklist intervals from the passlist intervals.
3. Chunk the remaining intervals into n_threads chunks
4. For each chunk, use the original fragmentgenerator codepath.

[jamesli124]

Small note: none of the tests test opening fragment information in .frag.gz/bed.gz files.

[edawson]

@jamesli124 one thing to consider for these repo would be setting a format convention and integrating pre-commit to run it automatically. That would prevent these little format changes from sneaking in to MRs. Happy to help if it's of interest.

[jamesli124]

@edawson Sorry, I was away when you wrote this, but I think it is a great idea. Is this something you are still interested in helping with?

[DucoG]

Thanks for the reviews, @jamesli124 and @edawson. I've reworked the PR to address all the feedback. Since main moved on quite a bit since this was opened (the io/ refactor in #171 and CRAM support in #174/#178), I rebased onto current main and re-implemented the optimization on top of the new AlignmentWrapper/ReferenceWrapper infrastructure. It's now a single commit against main.

What changed vs. the previous revision

1. Restored the full docstring (@jamesli124). The change is now based on main's current delfi docstring, so all parameter documentation is intact. The only docstring edit is to note that input_file now also accepts CRAM and tabix-indexed fragment files.

2. Fixed the dropped tabix fragment-file support (@jamesli124). The worker no longer opens input unconditionally as pysam.AlignmentFile. Instead it opens an AlignmentWrapper once per worker via the Pool initializer, which transparently handles BAM/CRAM/.frag.gz/.bed.gz — exactly as frag_generator does. So .frag.gz/.bed.gz input works again (no more silent hang on empty output).

3. Used AlignmentWrapper (@edawson). Both the alignment and the reference are now read through the project wrappers (AlignmentWrapper + ReferenceWrapper), so the "raw global pysam.AlignmentFile" is gone. GC also goes through ReferenceWrapper, so .2bit and FASTA both work and CRAM GC uses the FASTA, consistent with the rest of main.

4. Added tests for fragment-file input (@jamesli124).

• test_fragfile_input builds tabix-indexed .frag.gz (5-col FinaleDB) and .bed.gz (6-col) files from the test BAM and asserts DELFI: reads them without hanging, returns non-empty output, is invariant to worker count, and produces identical results across the two tabix layouts.
• test_workers_equivalence (kept) asserts workers=1 and workers=4 are bit-identical.

5. Avoided cosmetic churn (@edawson). The diff is restricted to functional changes — trim_coverage and the unchanged docstring lines are byte-for-byte identical to main (no stray operator/whitespace reformatting this time).

On the "no global wrapper / set-algebra" suggestion (@edawson)

I kept the per-worker handle but moved it behind AlignmentWrapper and opened it via the Pool initializer rather than as task arguments. Open pysam/py2bit handles can't be pickled into task args, so a per-worker handle created in the initializer is the standard multiprocessing.Pool pattern for sharing a non-picklable resource — the "global" is just process-local worker state, set once per worker. I didn't pursue the full passlist−blocklist set-algebra + interval-chunking rewrite: it changes the windowing/output contract and a related per-contig restructuring benchmarked slower (the per-window fetches use the BAI/CRAI/tabix index and skip gap regions). Happy to explore it as a separate change if you'd like.

Correctness

• Verified bit-identical to main: every integer column (contig, start, stop, arm, short, long, num_frags) matches exactly, and gc/ratio match to full float precision, on the chr1 6Mb test data (workers=1 and workers=4).
• Full suite passes locally: 80 passed, 27 skipped (skips are the samtools/network-gated tests). The DELFI + CRAM-DELFI equivalence tests pass.

One pre-existing observation (not changed here)

BAM and tabix input give slightly different per-window counts (~0.02% on the test data): the BAM reader fetches each window by read-alignment position while the tabix reader fetches by fragment span, so a handful of window-boundary fragments land in adjacent windows. This is independent of this PR (it's in frag_generator/AlignmentWrapper on main), which is why test_fragfile_input checks closeness to the BAM rather than exact equality. Flagging in case it's worth a separate issue.

[ravibandaru-lab]

@DucoG Thank you for the help. We have incorporated this into our refactor and have credited you accordingly in changelog and as an Github contributor.