parallelization

[iron-lion]

Dear @lilit-nersisyan

Thanks for the nice tool.

I was analyzing our Ribo-seq data using your library, but it was too slow for our analyses. The original Fivepseq processes all ~20,000 transcripts sequentially at every stage of the analysis pipeline.

This PR adds multiprocessing, resulting in a roughly 10× speedup (with 8 workers) in computation time.
I tested with our own data.

Test scenario: serial (--n-workers =1) vs. parallel (--n-workers= 8)
Note: This PR was done with assistance from Claude.

---

Pipeline Execution Order

CountPipeline.run()
│
├─ run_transcript_assembly()            [sequential]
├─ run_transcript_and_count_descriptors()  [calls generate_transcript_descriptors()]
│   └─ generate_transcript_descriptors()   [PROCESS-PARALLEL — 2 pool passes]
│       ├─ pass 1: per-transcript CDS sequence + count vector → descriptors + count_distribution_dict
│       ├─ merge count_distribution_dict → compute outlier_lower
│       └─ pass 2: downsampled read counts (uses outlier_lower)
├─ run_outliers()                       [sequential]
├─ run_start_stop_dicts()               [sequential]
├─ run_full_length_counts()             [sequential]
├─ run_count_stats()                    [sequential]
├─ run_transcript_indices()             [sequential]
│
├─ ThreadPoolExecutor (4 workers) ──────[THREAD-PARALLEL barrier]
│   ├─ run_meta_counts()
│   ├─ run_frame_counts()
│   ├─ run_codon_counts()               [calls compute_codon_pauses() → PROCESS-PARALLEL]
│   └─ run_loci_counts()
│
├─ run_aa_island_counts()               [calls compute_aa_islands() → PROCESS-PARALLEL]
├─ run_codon_stats()                    [sequential]
├─ run_queue_analysis()                 [sequential]
└─ run_out_of_frame_codon_stats()       [calls get_out_of_frame_codon_pauses() × 2 → PROCESS-PARALLEL]

---

Process-Parallel Pattern

All five multiprocessing functions follow the same three-phase pattern.

Chunk strategy

chunk_size = (transcript_count + n_workers - 1) // n_workers   # ceiling division

Transcripts are divided into n_workers equal chunks. Each chunk becomes one task_id. If transcript_count < n_workers, n_workers is clamped to transcript_count.

Phase 1 — Set globals before fork (copy-on-write)

global _shared_transcript_assembly, _shared_genome_dict
_shared_transcript_assembly = transcript_assembly   # List[Transcript], ~500 MB
_shared_genome_dict         = self.genome.genome_dict

get_context('fork') workers inherit these via the OS copy-on-write mechanism. The objects are never pickled or transferred; workers read pages only as needed. Setting them before Pool() construction is required — workers see the module state at fork time.

Phase 2 — Pool initializer: fresh BAM handle per worker

def _init_bam_worker(bam_path, three_prime, span_size, outlier_lower, fpat):
    global _worker_bam_array, _worker_span_size, _worker_outlier_lower, _worker_fpat
    af = pysam.AlignmentFile(bam_path, "rb")
    mapping = ThreePrimeMapFactory() if three_prime else FivePrimeMapFactory()
    _worker_bam_array = BAMGenomeArray(af, mapping=mapping)
    ...

def _init_codon_worker(bam_path, three_prime, outlier_lower, fpat, shm_specs):
    # same BAM setup, plus receives shared-memory specs
    ...
    _codon_worker_shm_specs = shm_specs

Each worker opens its own pysam.AlignmentFile. File descriptors cannot be shared across fork (POSIX semantics), so the parent never opens the BAM — it passes the path as a string.

Phase 3 — Shared-memory accumulator (zero-copy reduction)

# Main process allocates before pool
shape  = (n_workers, n_rows, n_cols)    # int32
nbytes = int(np.prod(shape)) * 4
shm    = SharedMemory(create=True, size=nbytes)
np.ndarray(shape, dtype=np.int32, buffer=shm.buf)[:] = 0
shm_specs['name'] = (shm.name, shape)

# Worker: accumulate locally, write once at end of chunk
local_arr = np.zeros((n_rows, n_cols), dtype=np.int32)
for idx in chunk_indices:
    ...   # fill local_arr
shm = SharedMemory(name=shm_specs['name'][0], create=False)
np.ndarray(shape, dtype=np.int32, buffer=shm.buf)[task_id] = local_arr
shm.close()

# Main process: reduce after pool.map() returns
result = np.ndarray(shape, dtype=np.int32, buffer=shm.buf).sum(axis=0).copy()
shm.close(); shm.unlink()

Each worker writes to its own [task_id] slice — no locks, no contention. The main process sums across axis=0 to get the final aggregate array.

---

Shared-Memory Lifecycle

All SharedMemory blocks follow this pattern to avoid leaks on exceptions:

shm_objects = {...}   # dict of SharedMemory instances
try:
    with ctx.Pool(...) as pool:
        pool.map(worker_fn, task_args)
finally:
    result = np.ndarray(shape, buffer=shm.buf).sum(axis=0).copy()
    for shm in shm_objects.values():
        shm.close()
        shm.unlink()

finally guarantees cleanup even if workers crash. shm.unlink() removes the POSIX shared-memory name from /dev/shm; shm.close() releases the local mapping.

---

Comparison

All 33 output files match exactly between serial (n_workers=1) and parallel (n_workers=8) runs.

Test env : 2,000 genes of Ribo-seq data

Mode	Wall time
Serial (n_workers=1)	2 487 s (41.5 min)
Parallel (n_workers=8)	191 s (3.2 min)
Speedup	13.02×