Skip to content

Vectorize compute-properties: replace per-row BioPython with array math (~10x at scale)#17

Merged
PaulNewling merged 21 commits into
mainfrom
feat/seq-properties-vectorization
Jun 11, 2026
Merged

Vectorize compute-properties: replace per-row BioPython with array math (~10x at scale)#17
PaulNewling merged 21 commits into
mainfrom
feat/seq-properties-vectorization

Conversation

@PaulNewling

@PaulNewling PaulNewling commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

What

Replaces the row-by-row BioPython computation in compute-properties with a single-threaded vectorized engine (software/src/vectorized.py): a per-residue count matrix feeds numpy/polars array math, reusing BioPython's exact constant tables (kd, DIWV, protein_weights) so the numbers are unchanged. The old software/src/properties.py is retained unchanged as the test oracle.

Why

The block's "calculation" step was the wall-clock bottleneck — not p-frame manipulation, as first suspected. The Python step iterated clones one at a time, constructing BioPython ProteinAnalysis/IsoelectricPoint objects and running a pI bisection per clone — O(N) serial Python.

Result

Baseline (pre-change per-row code) vs this PR, measured back-to-back on identical full-paired antibody input:

Clones Before After Speedup
1M ~119 s ~12 s ~10×
2M ~243 s ~25 s ~10×

~10× at production scale, where it plateaus (~120 µs/clone baseline vs ~12 µs/clone vectorized). Best case at small N (~50k, isolated) is ~14×. Single-threaded (POLARS_MAX_THREADS=1), so behavior is identical on a 2-core laptop and a 64-core server — no cpu()/core-count dependency.

Memory & resources

The vectorized engine materializes full-N count matrices, so it uses more peak memory than the old per-row code:

  • Old per-row: ~2.4 GB/million clones.
  • New, optimized: ~4.85 GB/million (first cut ~7.9 GB; reduced via int32 index arrays, a shared clean pass, and dropped dead retention — all bit-identical).

The compute step's mem() is raised 4 → 12 GiB to cover ~2M clones; cpu(1) stays (single-threaded by design).

Correctness

Output is unchanged within a quantized-equal contract (every emitted float rounded below display precision at the pipeline boundary). Verified against the retained scalar oracle:

  • Frozen whole-frame golden snapshot — 7/7 byte-equal across all modes + edge cases.
  • Per-property parity vs the oracle: bit-exact to float epsilon; charge & pI bit-identical (pI bisection: 0 straddles over 5,108 cells).
  • Pipeline-composition parity vs the oracle — catches pKa-set / include_cys / chain-pairing wiring regressions the self-generated golden can't.
  • Hypothesis invariants; cross-subprocess + row-permutation determinism.
  • 259 tests pass (the @slow perf benchmark runs separately, excluded from the CI gate where wall-clock budgets flake).

One refinement: signed-zero canonicalization (-0.0 → +0.0) at the quantization boundary — numpy's reduction order can flip the sign of a numerically-zero residual. It keeps the CID stable and fixes a pre-existing -0.000 display wart; the golden changed only on 3 verified-benign zero cells.

Upgrade impact

Upgrading from the published block recomputes this block (the engine changed) and dirties downstream consumers of gravy_*, mw_*, instability_*, aliphatic_*, aromaticity_*, and aaFraction — their bytes change (the published block stored them at full BioPython precision; the new ones are numpy-computed and rounded below display). charge_*, chargeShift_*, pi_*, eox_*, ered_* are byte-stable. The values are scientifically identical (≤1e-13); the downstream recompute is one-time and lands on the same answer. Forward-stable from here — the round-below-display contract holds the CID steady across future numpy/libm upgrades.

Live validation (in-platform, real data)

Ran the full pipeline Samples & Data → Import V(D)J Data → Sequence Properties against a backend on a real MiXCR bulk BCR export (280 IGHeavy clonotypes) — once with the published block and once with this branch, both fed from the same Import V(D)J IGHeavy output. Both produced the same mode/coverage (antibody_tcr_legacy_bulk / full_chain / IG), no errors.

Comparing every emitted column across all 280 clonotypes (order-independent hash of clonotypeKey=value, at both stored precision and raw float bytes):

  • All 12 columns identical at display precision — the vectorized engine reproduces the published block's user-visible values exactly on real data.
  • Byte-identical (raw bytes) on the 6 columns promised stable: chargeShift, charge (CDR-H3 + VH), eox, ered, pi.
  • Raw bytes differ only in the sub-display FP tail on the other 6 (gravy ×2, mw, instability, aliphatic, aromaticity) — published emits full BioPython precision (e.g. -0.30640000000000006), this branch rounds below display (-0.3064).

This confirms the Upgrade impact section empirically: the predicted byte-stable set matched byte-for-byte, the predicted recompute set differs only below display precision (values unchanged). The quantized-equal contract — previously checked only against the in-repo oracle — now also holds against the actually-published block.

Notes

  • Runtime dependency versions are unchanged; this PR also makes pyproject.toml [project.dependencies] the single source of truth and generates src/requirements.txt from it, gated by a CI requirements-sync check (matching titeseq-analysis).
  • Changeset bumps .software and .workflow (patch).
  • Design spec, plan, and decision journal: milaboratory/text#156.

Greptile Summary

This PR replaces the row-by-row BioPython computation in compute-properties with a vectorized numpy/polars engine (vectorized.py) that builds a per-residue count matrix once and derives all properties with array math, reusing BioPython's exact constant tables so numeric output is unchanged. Correctness is pinned by 259 tests including a new oracle-parity integration test that validates every emitted cell against the scalar BioPython path.

  • vectorized.py (new): single-pass cleaning via latin-1 byte gather → count matrix scatter; lockstep fixed-iteration bisection for pI (bit-identical to scalar); order-preserving dipeptide instability via np.add.at; signed-zero canonicalization at the output boundary.
  • pipeline.py: per-row loops removed; full-chain path shares a single _Cleaned object across the count substrate and the instability index; run_peptide calls build_counts + instability_index separately (two clean passes — see inline comment).
  • Memory & workflow: mem() raised 4 → 12 GiB; POLARS_MAX_THREADS=1 pinned before polars import; pyproject.toml is now the single source of truth for runtime deps with a new CI requirements-sync job.

Confidence Score: 5/5

Safe to merge. The vectorized engine is extensively tested against the scalar BioPython oracle on every emitted cell, and the single remaining suboptimal pattern (peptide path double-cleans) is a pure performance concern with no correctness impact.

The vectorized math is backed by property-level Hypothesis parity tests, a frozen whole-frame golden snapshot (7/7 byte-equal), and a new pipeline-vs-oracle integration test covering every column's pKa-set/include_cys/chain-pairing wiring. The only finding is a redundant cleaning pass in the peptide path with no effect on correctness or output values. No data-affecting bugs were found.

software/src/pipeline.py (run_peptide double-cleans sequences)

Important Files Changed

Filename Overview
software/src/vectorized.py New vectorized compute engine — single-pass cleaning, count-matrix dispatch, lockstep bisection pI, and order-preserving instability index. Logic is correct and well-tested.
software/src/pipeline.py Scalar per-row path replaced by vectorized engine. Antibody/TCR path shares _Cleaned across instability and counts; peptide path calls _clean_vectorized twice (once in build_counts, once in instability_index) — redundant at large N.
software/src/main.py Adds POLARS_MAX_THREADS=1 via os.environ.setdefault before polars is imported — correct placement and allows env override.
workflow/src/main.tpl.tengo mem() raised from 4GiB to 12GiB to accommodate full-N count matrices; cpu(1) preserved for single-threaded determinism.
software/tests/integration/test_pipeline_oracle_parity.py New integration test that cross-validates the vectorized pipeline against the scalar BioPython oracle on every emitted cell, catching pKa-set/include_cys/chain-pairing wiring regressions.
software/tests/unit/test_vectorized_instability.py Hypothesis + deterministic tests for instability index: oracle parity, floor boundary (9 vs 10 residues), invalid-row NaN, order stability.
software/tests/unit/test_vectorized_fv_pi.py Hypothesis parity + deterministic straddle test (0 straddles across 1500 VH/VL pairs) for vectorized Fv isoelectric point.
.github/workflows/python-tests.yaml Adds requirements-sync CI job (skipped on internal repo), pins uv to 0.11.16, excludes @slow perf benchmark from the gate.
Prompt To Fix All With AI
Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 1
software/src/pipeline.py:189-200
**Peptide path cleans sequences twice**

`build_counts(seqs)` (line 189) calls `_clean_vectorized` internally, and then `instability_index(seqs)` (line 200) calls `_clean_vectorized` again on the same list. This is exactly the pattern the `instability_index` docstring warns against: "Callers that also need the count substrate should clean once with `_clean_vectorized` and feed the shared `_Cleaned` to both `counts_from_cleaned` and `instability_from_cleaned`."

For a 1M-peptide load, each clean pass encodes, joins, and scatters ~100 MB of character data plus `row_of_res` construction — running it twice is a redundant transient allocation and numpy pass. The full-chain path already uses the shared-`_Cleaned` pattern; the peptide path should too.

Reviews (3): Last reviewed commit: "chore(workflow): trim compute-properties..." | Re-trigger Greptile

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request vectorizes sequence property computations by replacing per-row BioPython calculations with single-threaded NumPy and Polars array operations, significantly improving performance. It also introduces comprehensive quantization rounding across all float families to maintain determinism and absorb floating-point drift. Feedback on the changes suggests reordering the quantization prefix list to prevent incorrect prefix matching of 'chargeShift_' with 'charge_', and using safe dictionary retrieval with fallbacks when accessing chain substrates for Fv properties to avoid potential KeyErrors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread software/src/pipeline.py
Comment thread software/src/pipeline.py
…ero canonicalization, sync-guard + signed-zero tests, comment/log-token fixes
@PaulNewling

Copy link
Copy Markdown
Collaborator Author

@greptileai

Comment thread software/src/pipeline.py
@PaulNewling PaulNewling changed the title Vectorize compute-properties: replace per-row BioPython with array math (~14.5x) Vectorize compute-properties: replace per-row BioPython with array math (~14x) Jun 10, 2026
@PaulNewling PaulNewling changed the title Vectorize compute-properties: replace per-row BioPython with array math (~14x) Vectorize compute-properties: replace per-row BioPython with array math (~10x at scale) Jun 10, 2026
…rop vestigial constants, exclude perf from CI gate

- Add pipeline<->oracle composition parity test (catches pKa/include_cys/chain-pairing wiring regressions the self-generated golden can't)
- Add vectorized-engine invariant tests (invariants hold on the shipping code, not only the oracle)
- Add edge tests: Fv AND-gating with a present-but-invalid chain; zero-row full-chain+Fv; non-ASCII single-byte cleaning divergence
- Remove vestigial CID_QUANTIZE_PREFIXES/DECIMALS; tests derive decimals from the live _decimals_for
- Exclude the @slow perf benchmark from the pytest CI gate (wall-clock budget flakes on shared runners)
~4.85 GB/1M optimized footprint; 12GiB covers ~2M clones with headroom (was 16GiB).
@PaulNewling

Copy link
Copy Markdown
Collaborator Author

@greptileai

@PaulNewling PaulNewling marked this pull request as ready for review June 11, 2026 01:55
@PaulNewling PaulNewling added this pull request to the merge queue Jun 11, 2026
Merged via the queue into main with commit dc11b5b Jun 11, 2026
14 checks passed
@PaulNewling PaulNewling deleted the feat/seq-properties-vectorization branch June 11, 2026 01:55
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