diff --git a/audit-docs/allele_model_audit.md b/audit-docs/allele_model_audit.md index aae3fd4..7421f33 100644 --- a/audit-docs/allele_model_audit.md +++ b/audit-docs/allele_model_audit.md @@ -39,7 +39,7 @@ The following Allele methods exist on every subclass: ### `_find_anchor` **Build-time only.** `_find_anchor` runs inside `Allele.__init__` -when one of the bundled-data builders in `.private/scripts/` +when one of the private bundled-data builder scripts constructs `VAllele("name", gapped_seq, length)` — the constructor calls `self._find_anchor()` unless `anchor_override=...` was passed. diff --git a/audit-docs/docs_website_audit.md b/audit-docs/docs_website_audit.md index 4f9386d..887c380 100644 --- a/audit-docs/docs_website_audit.md +++ b/audit-docs/docs_website_audit.md @@ -92,11 +92,10 @@ sibling audits. | **Designs** (per-slice / per-mechanism scoping pre-implementation) | 17 | `clonal_family_design.md`, `clonal_parent_outcome_design.md`, `clonal_plan_split_design.md`, `d_inversion_design.md`, `d_inversion_extension_design.md`, `paired_end_design.md`, `receptor_revision_design.md`, `shm_segment_rate_design.md`, `v_subregion_shm_rate_design.md`, `np_markov_base_generator_design.md`, `p_nucleotide_design.md`, `fastq_export_design.md`, `allele_usage_estimation_design.md`, `trim_distribution_estimation_design.md`, `np_length_estimation_design.md`, `np_base_model_estimation_design.md`, `p_nucleotide_length_estimation_design.md` | | **Hubs / contributor entry points** | 6 | `engine_architecture.md`, `adding_a_pass.md`, `validation_matrix.md`, `reference_cartridge.md`, `airr_record_validator.md`, `allele_model_audit.md` | -The `docs/superpowers/plans/` directory holds Claude- -session planning artifacts (e.g. -`2026-05-18-mcp-redesign-v2.md`) — **not user-facing -documentation**, but currently mixed into the same -`docs/` tree. +A private planning directory holds pre-implementation +planning notes (e.g. `2026-05-18-mcp-redesign-v2.md`) — +**not user-facing documentation**, but currently mixed +into the same `docs/` tree. The `docs/build/` directory holds Python wheel build artifacts (a side-effect of `python -m build`) — also @@ -151,7 +150,6 @@ MkDocs Material lands. - `pin_scaffold_docs_dir_carries_thirty_eight_md_audit_design_files` - `pin_scaffold_old_docs_dir_exists_as_abandoned_earlier_attempt` - `pin_scaffold_deploy_docs_workflow_targets_website_dir` -- `pin_scaffold_docs_superpowers_subdir_holds_session_artifacts_not_docs` - `pin_scaffold_docs_build_subdir_holds_wheel_artefacts_not_docs` --- @@ -559,10 +557,10 @@ expand this audit. versioning; the MkDocs migration could add it (via `mike`). Decision belongs to the framework choice. - **Search backend.** Same — depends on framework choice. -- **`docs/superpowers/plans/` cleanup.** The Claude- - session planning artefacts should probably move to a - separate `.private/` directory but that's a housekeeping - matter, not a docs structure issue. +- **Private planning-notes cleanup.** The pre-implementation + planning notes should live in a separate private directory, + not under `docs/` — a housekeeping matter, not a docs + structure issue. --- diff --git a/audit-docs/reference_cartridge_authoring_audit.md b/audit-docs/reference_cartridge_authoring_audit.md index 8daffef..5c08bf8 100644 --- a/audit-docs/reference_cartridge_authoring_audit.md +++ b/audit-docs/reference_cartridge_authoring_audit.md @@ -74,10 +74,12 @@ in three places: "Populated by `RandomDataConfigBuilder.make_from_reference`", but `RandomDataConfigBuilder` is no longer importable anywhere in `src/GenAIRR/`. -2. `.private/scripts/build_imgt_configs.py` still imports - `from GenAIRR.dataconfig.make.random import - RandomDataConfigBuilder` and is **currently broken** — - running it would `ModuleNotFoundError` at the top. +2. The IMGT build script (historically a `.private/` script that + imported the removed `RandomDataConfigBuilder`) has since been + ported to a first-class, tracked maintainer tool at + `tools/build_imgt_configs.py`, built on + `ReferenceCartridgeBuilder`. The paragraphs below document the + interim broken-stub state that the port resolved. 3. The build-cache mirror at `docs/build/lib.linux-x86_64-cpython-312/GenAIRR/dataconfig/make/` carries the historical implementation as a compile @@ -129,7 +131,7 @@ reuse them rather than reinvent them. | Reference site | Pre-slice state | Post-slice state | |---|---|---| | [`dataconfig/data_config.py:159-166`](../src/GenAIRR/dataconfig/data_config.py#L159-L166) — `DataConfig.build_report` field docstring | "Populated by `RandomDataConfigBuilder.make_from_reference`" (named a removed class) | **[Cleaned]** Docstring now references `GenAIRR.cartridge_builder.ReferenceCartridgeBuilder.build`. Pinned by `test_pin_present_build_report_docstring_now_references_new_builder`. | -| [`.private/scripts/build_imgt_configs.py`](../.private/scripts/build_imgt_configs.py) | `from GenAIRR.dataconfig.make.random import RandomDataConfigBuilder` — broken at module load with `ModuleNotFoundError` | **[Cleaned]** Top-level `raise NotImplementedError(...)` with explicit porting hint to `ReferenceCartridgeBuilder` + audit-doc reference; dead import moved into the unreachable function body for porting reference. Pinned by `test_pin_present_private_build_script_now_raises_explicit_legacy_error`. | +| [`tools/build_imgt_configs.py`](../tools/build_imgt_configs.py) | The legacy build script imported the removed `from GenAIRR.dataconfig.make.random import RandomDataConfigBuilder` | **[Ported]** Rewritten as a first-class, tracked maintainer tool using `ReferenceCartridgeBuilder` (`from_fasta → infer_identity → infer_v_subregions → build`), producing structural cartridges. Pinned by `test_pin_present_imgt_build_tool_uses_reference_cartridge_builder` + `tests/test_build_imgt_configs.py`. | | `docs/build/lib.linux-x86_64-cpython-312/GenAIRR/dataconfig/make/...` | Historical implementation of `RandomDataConfigBuilder` / `CustomDataConfigBuilder` as a compile artefact | **Unchanged.** Build-cache mirror is not on the import path; regenerated by the next wheel build. Pinned absent by `test_pin_scaffold_historical_random_builder_module_is_gone` (verifies `ModuleNotFoundError` on import). | ### 3.2 Historical shape (for design inspiration only) @@ -527,9 +529,9 @@ sub-steps: onto `cfg.build_report`). 6. **Dead-reference cleanup**: update the - `DataConfig.build_report` docstring + update - `.private/scripts/build_imgt_configs.py` to import from - the new module. + `DataConfig.build_report` docstring + port the IMGT build + script to the new module (done: now + `tools/build_imgt_configs.py`). Cost estimate: @@ -573,8 +575,8 @@ builder reuses every bundled cartridge (no producer in live source). 11. `DataConfig.build_report` docstring still names `RandomDataConfigBuilder` — a dead class. -12. `.private/scripts/build_imgt_configs.py` still imports - from the dead module path. +12. The IMGT build script imported from the dead module path + (since ported to `tools/build_imgt_configs.py`). ### `pin_absence_*` — the gaps the slice closes @@ -624,12 +626,11 @@ expand the work. - **Cartridge diff / merge tooling.** A future `CartridgeDiff` surface could compare two cartridges' build reports — separate slice. -- **Live-data download (IMGT / OGRDB / etc.).** The - historical `.private/scripts/build_imgt_configs.py` - downloads FASTA from the web. The new builder accepts file - paths or pre-parsed allele lists; download tooling is a - separate concern (the private script can stay private, - updated to import the new builder). +- **Live-data download (IMGT / OGRDB / etc.).** The IMGT + build tool (`tools/build_imgt_configs.py`) downloads FASTA + from the web. The builder itself accepts file paths or + pre-parsed allele lists; download tooling stays a separate + concern in the maintainer tool. - **GUI / web wrapper.** Out of scope. - **Auto-bundling.** The builder produces a `DataConfig` ready for pickling, but the wheel-build process that ships diff --git a/audit-docs/validation_matrix.md b/audit-docs/validation_matrix.md index 191453d..f6d07d9 100644 --- a/audit-docs/validation_matrix.md +++ b/audit-docs/validation_matrix.md @@ -103,7 +103,7 @@ required step in the contributor flow: | **Paired-end FASTQ export** (`SimulationResult.to_paired_fastq(r1_path, r2_path, *, quality="illumina", overwrite=False, **quality_kwargs)` writes two synchronized FASTQ files — one R1 record + one R2 record per AIRR record — sourced verbatim from the `r1_sequence` / `r2_sequence` projection fields. Headers use `@{sequence_id}/1` and `@{sequence_id}/2` (universally-portable suffix; no `\|`-pipe metadata that some aligners split on). R2 is **already** reverse-complemented at projection time (the `PairedEndWindowMismatch { side: R2 }` validator invariant enforces it); the writer outputs it verbatim and applies no second flip. Reuses the existing pluggable quality models (`ConstantQualityModel` / `IlluminaQualityModel` from `_qmodel.py`) consumed by the single-end `to_fastq`; each read is scored independently so R1 and R2 get their own ramp-up + tail-down shape from position 0. Refuses to clobber existing files unless `overwrite=True`; raises on non-paired records (`read_layout != "paired_end"`), empty R1/R2 windows, and length-mismatched quality arrays. Pure projection — no engine, validator, trace, manifest, or plan-signature touchpoints.) | [`fastq_export_design.md`](fastq_export_design.md) (slice shipped), [`paired_end_design.md`](paired_end_design.md) (the projection plumbing the writer reads from) | [`test_fastq_export_contract.py`](../tests/test_fastq_export_contract.py) (scaffold + present pins for the method + signature; absence pins remain on record-generator surface, manifest block, and trace addresses), [`test_to_paired_fastq.py`](../tests/test_to_paired_fastq.py) (13 spec tests: file record count, header format `@{sequence_id}/1` + `/2`, R1/R2 body equality against AIRR fields, quality length parity, constant + illumina quality models, non-paired rejection, empty R1/R2 rejection, overwrite guard with single-side pre-existence, synchronized R1/R2 sequence_id), paired-end FASTQ export smoke + decodability in [`test_release_validation.py`](../tests/test_release_validation.py) | `src/GenAIRR/result.py` (the `to_paired_fastq` method on `SimulationResult`), `src/GenAIRR/_qmodel.py` (reused unchanged — `ConstantQualityModel`, `IlluminaQualityModel`, `resolve_quality_model`, `phred_to_ascii`). **Zero Rust changes**; pure Python export layer on already-projected AIRR fields. | | **NP base models / Markov N-addition** (cartridge-authored NP-region base sampling via `NpBaseModelSpec(kind="uniform"\|"empirical_first_base"\|"markov")` on `ReferenceEmpiricalModels.np_bases`. Three concrete generators behind the new `NpBaseGenerator` trait: `UniformNpGenerator` (legacy 4-way ACGT, byte-identical signature to pre-typed-model baseline), `CategoricalNpGenerator` (position-independent weighted A/C/G/T, byte-identical to pre-Markov-slice `CategoricalBase` signature), and `MarkovBaseGenerator` (1-step previous-base-conditional with `first_base: [f64;4]` + `transitions: [[f64;4];4]` in canonical A/C/G/T from/to order). Per-position support is materialised via `base_generator.support(position, previous)`; the `GenerateNPPass` loop threads a `previous: Option` set to `Some(base)` only after the trace record is committed. **No new trace addresses** — replay reconstructs `previous` from `np.npN.bases[i-1]`. Plan signature folds the full Markov payload (5 rows × A/C/G/T canonical order) so two cartridges with identical `first_base` but different transitions fail the signature gate before any choice is consumed. Productive-only composes through the admit-mask × generator-support intersection; mid-stream `b'N'` sentinel falls back to `first_base`. **Legacy `NP_transitions` / `NP_first_bases` auto-lift remains deferred** — orphan fields stay in `_DOCUMENTED_ORPHAN_DATACONFIG_FIELDS`; manifest reports `legacy_fallback=False`.) | [`junction_n_addition_audit.md`](junction_n_addition_audit.md) (typed model + Markov shipped; legacy auto-lift deferred), [`np_markov_base_generator_design.md`](np_markov_base_generator_design.md) (Markov implementation surface table) | [`test_np_base_model_implementation.py`](../tests/test_np_base_model_implementation.py) (typed model spec + lowering + manifest + replay), [`test_np_markov_base_generator_contract.py`](../tests/test_np_markov_base_generator_contract.py) (20 post-slice pins: trait + generators + previous-base param + bridge kwarg + manifest flip + legacy orphan preservation), [`test_np_markov_base_generator_implementation.py`](../tests/test_np_markov_base_generator_implementation.py) (13 behaviour tests: lowering, deterministic A→T→G→C walk, strong-matrix convergence, replay round-trip, signature-gate mismatch, productive-only triad, byte-identical legacy signatures, manifest flip), [`test_np_markov_release.py`](../tests/test_np_markov_release.py) (release-tier: productive IGH full stack + replay round-trip + dependency invariant) | `engine_rs/src/passes/generate_np/np_base_generator.rs` (trait + `UniformNpGenerator` + `CategoricalNpGenerator` + `MarkovBaseGenerator` + 8 unit tests covering signature shapes, support-by-previous, mid-stream sentinel fallback, row-weight validation, signature divergence), `engine_rs/src/passes/generate_np.rs` (`base_generator: Box` field + `with_generator(…)` constructor + `parameter_signature` fold), `engine_rs/src/passes/generate_np/execution.rs` (`let mut previous` loop-local + post-commit update), `engine_rs/src/passes/generate_np/sampling.rs` (`previous: Option` threaded through `sample_base` + `validate_replayed_np_base` + private `SupportPairsDist` adapter feeding the generic `sample_base_with_admit_mask` / `sample_filtered_with_policy` helpers verbatim), `engine_rs/src/python/plan.rs` (`push_generate_np(..., markov_transitions=…)` kwarg + `parse_canonical_base_weights` row-completeness validator); Python bridge: [`src/GenAIRR/_dataconfig_extract.py`](../src/GenAIRR/_dataconfig_extract.py) (`_np_bases_from_models` + `_np_markov_transitions_from_models`), [`src/GenAIRR/_normalize.py`](../src/GenAIRR/_normalize.py) (`_to_immutable_byte_pair_matrix`), [`src/GenAIRR/_pipeline_ir.py`](../src/GenAIRR/_pipeline_ir.py) (`_RecombineStep.np{1,2}_markov_transitions`), [`src/GenAIRR/_compile.py`](../src/GenAIRR/_compile.py) (`_lower_recombine` thread), [`src/GenAIRR/dataconfig/data_config.py`](../src/GenAIRR/dataconfig/data_config.py) (manifest `np_base_models.supported_kinds = ["uniform", "empirical_first_base", "markov"]`, `deferred_kinds = []`) | | **P-nucleotide / palindromic addition** (cartridge-authored per-end P-nucleotide length distributions via `ReferenceEmpiricalModels.p_nucleotide_lengths`, keyed by junction side label `"V_3"` / `"D_5"` / `"D_3"` / `"J_5"`. v1 ships **lengths-only**: each `PAdditionPass` records one `Int(length)` choice at `p.{end}.length`; bases derive deterministically from `(allele, trim, orientation, length)` via `complement_base` — 3' extensions (V_3 / D_3) reverse-complement the last `length` post-trim coding bytes; 5' extensions (D_5 / J_5) reverse-complement the first `length`. P-bytes are pushed into the pool with `Nucleotide::flags::P_NUC` set; the descriptive `Region` flows through the new `SimulationEvent::PRegionAdded { end, region }` event only — **no structural region added to `sim.sequence.regions`**, so existing `regions.iter().find(|r| r.segment == ...)` projection sites stay correct (one structural region per biological segment). Pipeline ordering is post-trim: `assemble.V → p_addition.V_3 → generate_np.NP1 → [invert_d] → p_addition.D_5 → assemble.D → p_addition.D_3 → generate_np.NP2 → p_addition.J_5 → assemble.J`. **Critical IR ordering: `p_addition.D_5` runs AFTER `invert_d` commits** so D's effective_seq is read under the post-inversion orientation. AIRR projection surfaces four new int fields (`p_v_3_length` / `p_d_5_length` / `p_d_3_length` / `p_j_5_length`); the validator's `PLengthMismatch { end, reported, event_count }` issue kind catches tampered records via an event-ledger recompute. Plan signature folds each pass's per-end length distribution via `fmt_int_dist`. **Legacy `DataConfig.p_nucleotide_length_probs` is NOT auto-lifted** — orphan field stays in `_DOCUMENTED_ORPHAN_DATACONFIG_FIELDS`, manifest reports `legacy_fallback=False`. Per-base P strings (`p_v_3`, ...), aggregate `n_p_nucleotides`, and pre-trim P mode all remain explicitly out of scope per [`docs/p_nucleotide_design.md`](p_nucleotide_design.md) §15.) | [`p_nucleotide_design.md`](p_nucleotide_design.md) (audit + v1 implementation shipped; pre-trim / per-base strings / aggregate / legacy auto-lift deferred), [`junction_n_addition_audit.md`](junction_n_addition_audit.md) (junction-biology umbrella) | [`test_p_nucleotide_contract.py`](../tests/test_p_nucleotide_contract.py) (21 post-slice pins: PEnd enum / PRegionAdded variant / PLength choice address / PAdditionPass struct / cartridge plane / manifest block / four AIRR fields / PLengthMismatch validator surface / legacy-orphan boundary / IR-correct D_5 lowering position), [`test_p_nucleotide_implementation.py`](../tests/test_p_nucleotide_implementation.py) (10 behaviour tests: byte-identical baseline, VJ V_3/J_5 emission, VDJ all-four-ends emission, `P_NUC` flag emission proof, replay round-trip, productive-only triad, validator clean recompute, manifest exposure, legacy-orphan invariant, plan-signature fold + cross-replay rejection), [`test_p_nucleotide_release.py`](../tests/test_p_nucleotide_release.py) (release-tier: productive IGH full stack with all four ends + D inversion + receptor revision + SHM + corruption + `validate_records` + cache parity; replay round-trip with `invert_d(prob=1.0)`; deterministic palindrome derivation on a synthetic short allele) | `engine_rs/src/passes/p_addition.rs` (`PAdditionPass` + `effective_coding_bytes` + `derive_p_bytes` + 6 unit tests covering palindrome math at each end + zero-length no-op + declared-choice-pattern + forward/reverse orientation), `engine_rs/src/address.rs` (`PEnd { V3, D5, D3, J5 }` enum + `ChoiceAddress::PLength { end }` variant + 4 canonical address spellings + parser + `frozen_address_spellings` test extension + `P_ADDITION_{V_3,D_5,D_3,J_5}` pass-name constants), `engine_rs/src/ir/sim_event.rs` (`PRegionAdded { end: PEnd, region: Region }` variant), `engine_rs/src/ir/builder.rs` (`record_p_region(end, region)` method that broadcasts the event without mutating `sim.sequence.regions`), `engine_rs/src/airr_record/record.rs` (four new `p_*_length: i64` fields), `engine_rs/src/airr_record/builder.rs` (event-ledger walk filtered to the four `p_addition.*` pass names, summing `region.len()` per `PRegionAdded.end`), `engine_rs/src/airr_record/validate.rs` (`PLengthMismatch` variant + per-end recompute), `engine_rs/src/python/plan.rs` (`push_p_addition(end, length_pairs)` PyO3 method + `parse_p_end` helper), `engine_rs/src/python/outcome.rs` (four new `dict.set_item` calls + `PLengthMismatch` serialisation with `details.source: "event-ledger:PRegionAdded"`), `engine_rs/src/airr_record/tests/projection.rs` (2 Rust unit tests: tampered-field detection + clean-projection acceptance), `engine_rs/src/pass/support.rs` (`PassCompileFact::PLengthSupport { end, support }` variant for downstream schedule-analyser introspection), `engine_rs/src/live_call/{dirty_signal_observer,walker_observer,refresh_plan}.rs` + `engine_rs/src/ir/event_log_observer.rs` (no-op match arms documenting that `PRegionAdded` doesn't invalidate live-call state); Python: [`src/GenAIRR/reference_models.py`](../src/GenAIRR/reference_models.py) (`p_nucleotide_lengths: Dict[str, EmpiricalDistributionSpec]` field + VJ-rejects-D-end validation), [`src/GenAIRR/_dataconfig_extract.py`](../src/GenAIRR/_dataconfig_extract.py) (`_p_nucleotide_lengths_from_models` resolver + `_explicit_models` widening), [`src/GenAIRR/_pipeline_ir.py`](../src/GenAIRR/_pipeline_ir.py) (`_RecombineStep.p_{v_3,d_5,d_3,j_5}_lengths` immutable fields), [`src/GenAIRR/_compile.py`](../src/GenAIRR/_compile.py) (conditional `plan.push_p_addition(...)` insertions at the four audited boundaries with `invert_d → p_addition.D_5 → assemble.D` ordering enforced), [`src/GenAIRR/experiment.py`](../src/GenAIRR/experiment.py) (threading through `recombine()`), [`src/GenAIRR/dataconfig/data_config.py`](../src/GenAIRR/dataconfig/data_config.py) (`models.p_nucleotide_models` manifest block: `length_keys` / `legacy_p_nucleotide_length_probs_present` / `legacy_fallback=False` / `supported_ends` / `in_plan_signature=True` / `in_content_hash=False`) | -| **Reference cartridge authoring** (`ReferenceCartridgeBuilder` v1 facade + `CartridgeBuildReport` audit trail. Fluent staged builder: `from_fasta(v_fasta=..., j_fasta=..., d_fasta=..., chain_type=...)` constructor → `infer_identity(species=..., locus=..., reference_set=..., name=..., source=...)` → `infer_v_subregions()` → `with_rules(...)` → `with_models(...)` → `build()` returning a plain `DataConfig` ready for `Experiment.on(cfg)`. v1 ships structural authoring only — FASTA parsing via the existing `parse_fasta` helper; V-subregion derivation via the same `compute_v_region_boundaries` helper used by `dataconfig_to_refdata`; rules / empirical-models attachment validated at attach time. `build()` stamps the canonical `schema_sha256`, attaches the typed `CartridgeBuildReport`, runs `verify_integrity()`. Statistical estimators (`estimate_allele_usage`, `estimate_trim_distributions`, `estimate_np_length_distributions`, `estimate_np_base_model`, `estimate_p_nucleotide_lengths`, `estimate_shm_rates`) are explicitly deferred and raise `AttributeError`. The 3 historical dead-reference sites (`DataConfig.build_report` docstring + `.private/scripts/build_imgt_configs.py` + `docs/build/` mirror) cleaned up in lockstep — docstring rewritten to reference the new builder, private script raises `NotImplementedError` at module load with explicit porting hint to `ReferenceCartridgeBuilder`, build-cache mirror remains as compile artefact only and is verified absent from import path. Output is a plain `DataConfig` — no parallel cartridge object, flows through existing `dataconfig_to_refdata` bridge unchanged. Internal `_SafeAnchorMixin` + `_BuilderVAllele` / `_BuilderJAllele` subclasses degrade gracefully when the native `_native._anchor` C resolver isn't built so v1 succeeds in any environment.) | [`reference_cartridge_authoring_audit.md`](reference_cartridge_authoring_audit.md) (v1 shipped; estimators deferred — §11.2 lists explicit method names + recommended ordering), [`reference_cartridge.md`](reference_cartridge.md) (3-paths comparison: bundled vs manual vs builder + when-not-to-use guide) | [`test_reference_cartridge_authoring_contract.py`](../tests/test_reference_cartridge_authoring_contract.py) (25 post-slice pins covering: live inference-adjacent helpers reused — `parse_fasta` + `compute_v_region_boundaries` + `cartridge_manifest` JSON-cleanliness + `verify_integrity` + `RefDataConfig.{vj,vdj}` + `ReferenceEmpiricalModels` + `ReferenceRulesSpec` + `dataconfig_to_refdata` bridge + bundled-loader's 106 configs; dead-reference cleanup state — `RandomDataConfigBuilder` docstring removed + private build script raises explicit `NotImplementedError` + `GenAIRR.dataconfig.make.*` namespaces stay absent; v1 surface — module landed at `cartridge_builder` + class + dataclass at top level + `from_fasta` on builder only + `infer_*` step methods on builder; v1-deferred surfaces — no `from_airr` / no `estimate_*` at any of the three potential owner locations), [`test_reference_cartridge_authoring_implementation.py`](../tests/test_reference_cartridge_authoring_implementation.py) (15 behaviour tests: top-level imports + `from_fasta().build()` round-trip + compile through Experiment + `infer_identity()` manifest population + `infer_v_subregions()` 5-label coverage + missing-gapped-V warning-not-crash + `with_rules()`/`with_models()` attach + `build_report` pickle round-trip + `.to_dict()` JSON-cleanliness + manual `DataConfig(...)` still supported + dead-reference integration probe + v1 estimator absence + 3 error-path bonus: VJ-rejects-d_fasta, VDJ-requires-d_fasta, duplicate-allele-rejection), [`test_reference_cartridge_authoring_release.py`](../tests/test_reference_cartridge_authoring_release.py) (release-tier composition: tiny inline-FASTA VDJ cartridge built end-to-end + manifest JSON-clean + build-report JSON-clean + compile-and-recombine via `allow_curatable_refdata()` + report contains `from_fasta`/`infer_identity`/`build` stages) | `src/GenAIRR/cartridge_builder.py` (~570 lines: `CartridgeBuildReport` dataclass + `to_dict` + `ReferenceCartridgeBuilder` class + 6 stages + `from_fasta` + `_open_fasta` path/string/file-object dispatcher + `_parse_allele_name` IMGT-pipe-aware header parser + `_parse_segment_fasta` per-segment ingestion with structured rejection + `_SafeAnchorMixin` for graceful native-resolver fallback + `_BuilderVAllele` / `_BuilderJAllele` / `_BuilderDAllele` subclasses + `_resolve_chain_type` / `_resolve_species` / `_locus_label_from_chain_type` helpers), `src/GenAIRR/__init__.py` (top-level exports: `ReferenceCartridgeBuilder`, `CartridgeBuildReport`), `src/GenAIRR/dataconfig/data_config.py` (cleaned docstring on `build_report` field references new builder; defaults stay `None` for bundled + manual paths), `.private/scripts/build_imgt_configs.py` (explicit `raise NotImplementedError` at module load with porting hint); engine bridge: existing `dataconfig_to_refdata` validates builder-produced cartridges unchanged. **Zero Rust changes**; pure Python authoring layer on existing live infrastructure (FASTA parser, V-subregion derivation, typed-plane validators, manifest builder, integrity checker). | +| **Reference cartridge authoring** (`ReferenceCartridgeBuilder` v1 facade + `CartridgeBuildReport` audit trail. Fluent staged builder: `from_fasta(v_fasta=..., j_fasta=..., d_fasta=..., chain_type=...)` constructor → `infer_identity(species=..., locus=..., reference_set=..., name=..., source=...)` → `infer_v_subregions()` → `with_rules(...)` → `with_models(...)` → `build()` returning a plain `DataConfig` ready for `Experiment.on(cfg)`. v1 ships structural authoring only — FASTA parsing via the existing `parse_fasta` helper; V-subregion derivation via the same `compute_v_region_boundaries` helper used by `dataconfig_to_refdata`; rules / empirical-models attachment validated at attach time. `build()` stamps the canonical `schema_sha256`, attaches the typed `CartridgeBuildReport`, runs `verify_integrity()`. Statistical estimators (`estimate_allele_usage`, `estimate_trim_distributions`, `estimate_np_length_distributions`, `estimate_np_base_model`, `estimate_p_nucleotide_lengths`, `estimate_shm_rates`) are explicitly deferred and raise `AttributeError`. The 3 historical dead-reference sites (`DataConfig.build_report` docstring + the IMGT build script + `docs/build/` mirror) resolved in lockstep — docstring rewritten to reference the new builder, the build script ported to the first-class tool `tools/build_imgt_configs.py` on `ReferenceCartridgeBuilder`, build-cache mirror remains as compile artefact only and is verified absent from import path. Output is a plain `DataConfig` — no parallel cartridge object, flows through existing `dataconfig_to_refdata` bridge unchanged. Internal `_SafeAnchorMixin` + `_BuilderVAllele` / `_BuilderJAllele` subclasses degrade gracefully when the native `_native._anchor` C resolver isn't built so v1 succeeds in any environment.) | [`reference_cartridge_authoring_audit.md`](reference_cartridge_authoring_audit.md) (v1 shipped; estimators deferred — §11.2 lists explicit method names + recommended ordering), [`reference_cartridge.md`](reference_cartridge.md) (3-paths comparison: bundled vs manual vs builder + when-not-to-use guide) | [`test_reference_cartridge_authoring_contract.py`](../tests/test_reference_cartridge_authoring_contract.py) (25 post-slice pins covering: live inference-adjacent helpers reused — `parse_fasta` + `compute_v_region_boundaries` + `cartridge_manifest` JSON-cleanliness + `verify_integrity` + `RefDataConfig.{vj,vdj}` + `ReferenceEmpiricalModels` + `ReferenceRulesSpec` + `dataconfig_to_refdata` bridge + bundled-loader's 106 configs; dead-reference cleanup state — `RandomDataConfigBuilder` docstring removed + the IMGT build tool uses `ReferenceCartridgeBuilder` + `GenAIRR.dataconfig.make.*` namespaces stay absent; v1 surface — module landed at `cartridge_builder` + class + dataclass at top level + `from_fasta` on builder only + `infer_*` step methods on builder; v1-deferred surfaces — no `from_airr` / no `estimate_*` at any of the three potential owner locations), [`test_reference_cartridge_authoring_implementation.py`](../tests/test_reference_cartridge_authoring_implementation.py) (15 behaviour tests: top-level imports + `from_fasta().build()` round-trip + compile through Experiment + `infer_identity()` manifest population + `infer_v_subregions()` 5-label coverage + missing-gapped-V warning-not-crash + `with_rules()`/`with_models()` attach + `build_report` pickle round-trip + `.to_dict()` JSON-cleanliness + manual `DataConfig(...)` still supported + dead-reference integration probe + v1 estimator absence + 3 error-path bonus: VJ-rejects-d_fasta, VDJ-requires-d_fasta, duplicate-allele-rejection), [`test_reference_cartridge_authoring_release.py`](../tests/test_reference_cartridge_authoring_release.py) (release-tier composition: tiny inline-FASTA VDJ cartridge built end-to-end + manifest JSON-clean + build-report JSON-clean + compile-and-recombine via `allow_curatable_refdata()` + report contains `from_fasta`/`infer_identity`/`build` stages) | `src/GenAIRR/cartridge_builder.py` (~570 lines: `CartridgeBuildReport` dataclass + `to_dict` + `ReferenceCartridgeBuilder` class + 6 stages + `from_fasta` + `_open_fasta` path/string/file-object dispatcher + `_parse_allele_name` IMGT-pipe-aware header parser + `_parse_segment_fasta` per-segment ingestion with structured rejection + `_SafeAnchorMixin` for graceful native-resolver fallback + `_BuilderVAllele` / `_BuilderJAllele` / `_BuilderDAllele` subclasses + `_resolve_chain_type` / `_resolve_species` / `_locus_label_from_chain_type` helpers), `src/GenAIRR/__init__.py` (top-level exports: `ReferenceCartridgeBuilder`, `CartridgeBuildReport`), `src/GenAIRR/dataconfig/data_config.py` (cleaned docstring on `build_report` field references new builder; defaults stay `None` for bundled + manual paths), `tools/build_imgt_configs.py` (first-class maintainer tool on `ReferenceCartridgeBuilder`); engine bridge: existing `dataconfig_to_refdata` validates builder-produced cartridges unchanged. **Zero Rust changes**; pure Python authoring layer on existing live infrastructure (FASTA parser, V-subregion derivation, typed-plane validators, manifest builder, integrity checker). | | **Allele usage estimation** (`ReferenceCartridgeBuilder.estimate_allele_usage(rearrangements, *, min_count=1.0, ambiguous="fractional", replace=True)` reads AIRR `v_call` / `d_call` / `j_call` from `list[dict]` / path / open text handle and writes per-segment weights into the typed `ReferenceEmpiricalModels.allele_usage: Optional[AlleleUsageSpec]` plane — `AlleleUsageSpec` carries `v` / `d` / `j` `Dict[str, float]` fields, validates non-empty allele names + finite positive weights + VJ-D-rejection, and lowers through `_dataconfig_extract._allele_usage_from_models` into `Experiment.recombine` with precedence kwarg > cartridge plane > uniform (no Rust changes; reuses the existing `_RecombineStep.weights_*` → `plan.push_sample_allele(weights=...)` → `AllelePoolDist::from_weights` engine surface). Three ambiguity policies: `"fractional"` (splits 1.0 / k credit across the k tie-set entries), `"truth_first"` (uses the first call only — mirrors `_mcp_summary.first_call`), `"reject"` (drops ambiguous rows into `report.rejected` with `reason="ambiguous_rejected"`). Unknown alleles, VJ-D-warning (one-time per stage), VDJ-missing-D (`reason="missing_d_call_on_vdj"`), and below-`min_count` drops all surface in the canonical `{stage, inputs, inferred, warnings}` stage entry on `CartridgeBuildReport`. Per-segment weights normalised to sum to 1.0. Idempotent via `previously_estimated` + `replace` kwarg (mirrors `infer_v_subregions` discipline). Manifest `models.allele_usage` block surfaces `available` / `segments` / `nonempty_segments` / `legacy_gene_use_dict_present` / `legacy_fallback=False` / `in_plan_signature=False` (inherited soft gap 1 — same boundary as the per-experiment `v_allele_weights` kwarg) / `source`. **`gene_use_dict` stays orphan**: the estimator does NOT auto-lift the legacy dict — same boundary the Markov / P-nucleotide slices respected for their legacy fields.) | [`allele_usage_estimation_design.md`](allele_usage_estimation_design.md) (slice shipped; plan-signature soft-gap tightening + `gene_use_dict` auto-lift remain deferred), [`reference_cartridge_authoring_audit.md`](reference_cartridge_authoring_audit.md) §11.2 (estimator ordering — 1 of 6 shipped) | [`test_allele_usage_estimation_contract.py`](../tests/test_allele_usage_estimation_contract.py) (25 post-slice pins: scaffold + present-state surfaces + soft-gap pin held + `gene_use_dict` orphan-preservation pin + the 5 absence pins flipped to present), [`test_allele_usage_estimation_implementation.py`](../tests/test_allele_usage_estimation_implementation.py) (15 behaviour tests: `AlleleUsageSpec` validation + explicit typed plane biases recombination + explicit recombine kwarg overrides plane + `gene_use_dict` orphan-no-auto-lift + fractional splits tie-set credit + truth_first first-call-only + reject drops to `report.rejected` + unknown alleles in rejected + VJ-ignores-D + VDJ-missing-D + `min_count` drops low-support + stage entry shape + manifest block + pickle round-trip + contract-pin-flip integration probe), [`test_allele_usage_estimation_release.py`](../tests/test_allele_usage_estimation_release.py) (release-tier: tiny inline-FASTA VDJ cartridge built via the v1 facade + `estimate_allele_usage` on a synthetic AIRR record set + recombination shows the estimated bias + `manifest["models"]["allele_usage"]["available"]=True` + report carries the `estimate_allele_usage` stage + explicit `recombine(v_allele_weights=...)` override still wins) | **Zero Rust changes**; pure Python authoring layer on the existing engine surface. Reuses `engine_rs/src/dist/allele_pool.rs::AllelePoolDist::from_weights` + `engine_rs/src/python/plan.rs::push_sample_allele(weights=...)` (already wired end-to-end pre-slice — confirmed by `pin_scaffold_*` group in the contract file); Python: [`src/GenAIRR/reference_models.py`](../src/GenAIRR/reference_models.py) (`AlleleUsageSpec` dataclass + chain-type-aware validate + `ReferenceEmpiricalModels.allele_usage` field), [`src/GenAIRR/_dataconfig_extract.py`](../src/GenAIRR/_dataconfig_extract.py) (`_allele_usage_from_models` resolver + `_explicit_models` widening), [`src/GenAIRR/experiment.py`](../src/GenAIRR/experiment.py) (precedence shim before `_resolve_allele_weights`), [`src/GenAIRR/cartridge_builder.py`](../src/GenAIRR/cartridge_builder.py) (`estimate_allele_usage` method + `_split_tie_set` + `_allele_names_for_segment` + `_load_rearrangements` helpers), [`src/GenAIRR/dataconfig/data_config.py`](../src/GenAIRR/dataconfig/data_config.py) (`_allele_usage_manifest_block` + manifest insertion), [`src/GenAIRR/__init__.py`](../src/GenAIRR/__init__.py) (top-level `AlleleUsageSpec` export). | | **Allele-call ambiguity & disambiguation** | [`allele_call_audit.md`](allele_call_audit.md) | [`test_allele_call_provenance.py`](../tests/test_allele_call_provenance.py) (21 tests) | `engine_rs/src/live_call/tests/` (7 modules: walker, caller, state, bitset, reference_index, …) | | **Junction-call fields** | [`junction_call_audit.md`](junction_call_audit.md) | [`test_junction_call_provenance.py`](../tests/test_junction_call_provenance.py) (25 tests) | `engine_rs/src/junction.rs::tests` (12 unit tests), `engine_rs/src/airr_record/tests/anchors.rs` | diff --git a/engine_rs/src/address.rs b/engine_rs/src/address.rs index 55b189f..8c14235 100644 --- a/engine_rs/src/address.rs +++ b/engine_rs/src/address.rs @@ -941,632 +941,4 @@ fn parse_choice_address_pattern(address: &str) -> Option { } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn segment_helpers_emit_existing_recombination_addresses() { - assert_eq!(sample_allele_vdj(Segment::V), "sample_allele.v"); - assert_eq!(sample_allele_vdj(Segment::D), "sample_allele.d"); - assert_eq!(sample_allele_vdj(Segment::J), "sample_allele.j"); - assert_eq!(trim_vdj(Segment::V, TrimEnd::Five), "trim.v_5"); - assert_eq!(trim_vdj(Segment::J, TrimEnd::Three), "trim.j_3"); - assert_eq!(assemble_vdj(Segment::D), "assemble.d"); - assert_eq!(np_length_region(Segment::Np1), "np.np1.length"); - } - - #[test] - fn parsers_accept_existing_indexed_addresses() { - assert_eq!( - parse_indexed("mutate.s5f.base[2]", MUTATE_S5F_BASE_PREFIX), - Some(2) - ); - assert_eq!( - ChoiceAddress::parse("np.np1.bases[17]"), - Some(ChoiceAddress::NpBase { - segment: NpSegment::Np1, - index: 17, - }) - ); - assert_eq!( - ChoiceAddress::parse("np.np2.bases[4]"), - Some(ChoiceAddress::NpBase { - segment: NpSegment::Np2, - index: 4, - }) - ); - assert_eq!(ChoiceAddress::parse("np.np2.bases[x]"), None); - assert_eq!( - ChoiceAddress::parse("np.np1.length"), - Some(ChoiceAddress::NpLength(NpSegment::Np1)) - ); - } - - #[test] - fn typed_choice_addresses_round_trip_persisted_strings() { - let cases = [ - (ChoiceAddress::SampleAllele(VdjSegment::V), SAMPLE_ALLELE_V), - (ChoiceAddress::SampleAllele(VdjSegment::D), SAMPLE_ALLELE_D), - (ChoiceAddress::SampleAllele(VdjSegment::J), SAMPLE_ALLELE_J), - ( - ChoiceAddress::Trim { - segment: VdjSegment::V, - end: TrimEnd::Three, - }, - TRIM_V_3, - ), - ( - ChoiceAddress::Trim { - segment: VdjSegment::D, - end: TrimEnd::Five, - }, - TRIM_D_5, - ), - ( - ChoiceAddress::Trim { - segment: VdjSegment::J, - end: TrimEnd::Five, - }, - TRIM_J_5, - ), - (ChoiceAddress::NpLength(NpSegment::Np1), NP1_LENGTH), - ( - ChoiceAddress::NpBase { - segment: NpSegment::Np2, - index: 17, - }, - "np.np2.bases[17]", - ), - (ChoiceAddress::MutateUniformCount, MUTATE_UNIFORM_COUNT), - ( - ChoiceAddress::MutateUniformSite(2), - "mutate.uniform.site[2]", - ), - ( - ChoiceAddress::MutateUniformBase(3), - "mutate.uniform.base[3]", - ), - (ChoiceAddress::MutateS5fCount, MUTATE_S5F_COUNT), - (ChoiceAddress::MutateS5fSite(4), "mutate.s5f.site[4]"), - (ChoiceAddress::MutateS5fBase(5), "mutate.s5f.base[5]"), - (ChoiceAddress::CorruptPcrCount, CORRUPT_PCR_COUNT), - ( - ChoiceAddress::CorruptPcrSite(6), - "corrupt.pcr.error_site[6]", - ), - ( - ChoiceAddress::CorruptPcrBase(7), - "corrupt.pcr.error_base[7]", - ), - (ChoiceAddress::CorruptQualityCount, CORRUPT_QUALITY_COUNT), - ( - ChoiceAddress::CorruptQualitySite(8), - "corrupt.quality.error_site[8]", - ), - ( - ChoiceAddress::CorruptQualityBase(9), - "corrupt.quality.error_base[9]", - ), - ( - ChoiceAddress::CorruptContaminantApplied, - CORRUPT_CONTAMINANT_APPLIED, - ), - ( - ChoiceAddress::CorruptContaminantBase(10), - "corrupt.contaminant.bases[10]", - ), - (ChoiceAddress::CorruptIndelCount, CORRUPT_INDEL_COUNT), - ( - ChoiceAddress::CorruptIndelKind(11), - "corrupt.indel.kind[11]", - ), - ( - ChoiceAddress::CorruptIndelSite(12), - "corrupt.indel.site[12]", - ), - ( - ChoiceAddress::CorruptIndelBase(13), - "corrupt.indel.base[13]", - ), - (ChoiceAddress::CorruptNsCount, CORRUPT_NS_COUNT), - (ChoiceAddress::CorruptNsSite(14), "corrupt.ns.site[14]"), - ( - ChoiceAddress::CorruptEndLoss(PrimeEnd::Five), - CORRUPT_END_LOSS_5, - ), - ( - ChoiceAddress::CorruptEndLoss(PrimeEnd::Three), - CORRUPT_END_LOSS_3, - ), - ( - ChoiceAddress::CorruptRevCompApplied, - CORRUPT_REV_COMP_APPLIED, - ), - ( - ChoiceAddress::SampleAlleleDInverted, - SAMPLE_ALLELE_D_INVERTED, - ), - ( - ChoiceAddress::ReceptorRevisionApplied, - RECEPTOR_REVISION_APPLIED, - ), - ( - ChoiceAddress::ReceptorRevisionVAllele, - RECEPTOR_REVISION_V_ALLELE, - ), - ( - ChoiceAddress::ReceptorRevisionVTrim3, - RECEPTOR_REVISION_V_TRIM_3, - ), - (ChoiceAddress::PairedEndR1Length, PAIRED_END_R1_LENGTH), - (ChoiceAddress::PairedEndR2Length, PAIRED_END_R2_LENGTH), - ( - ChoiceAddress::PairedEndInsertSize, - PAIRED_END_INSERT_SIZE, - ), - ]; - - for (typed, raw) in cases { - assert_eq!(typed.to_string(), raw); - assert_eq!(ChoiceAddress::parse(raw), Some(typed)); - assert_eq!(raw.parse::().unwrap(), typed); - } - } - - #[test] - fn typed_choice_address_rejects_patterns_and_unknown_strings() { - for raw in [ - "mutate.s5f.site[0..n]", - "np.np1.bases[x]", - "assemble.v", - "sample_allele.np1", - "trim.np1_3", - "custom.choice", - "", - ] { - assert_eq!(ChoiceAddress::parse(raw), None, "raw={raw:?}"); - assert!(raw.parse::().is_err(), "raw={raw:?}"); - } - } - - #[test] - fn typed_segment_conversions_reject_wrong_segment_family() { - assert_eq!(VdjSegment::try_from(Segment::V), Ok(VdjSegment::V)); - assert_eq!(VdjSegment::try_from(Segment::Np1), Err(())); - assert_eq!(NpSegment::try_from(Segment::Np2), Ok(NpSegment::Np2)); - assert_eq!(NpSegment::try_from(Segment::J), Err(())); - } - - #[test] - fn typed_choice_address_patterns_round_trip_report_strings() { - let cases = [ - ( - ChoiceAddressPattern::SampleAllele(VdjSegment::V), - SAMPLE_ALLELE_V, - ), - ( - ChoiceAddressPattern::SampleAllele(VdjSegment::D), - SAMPLE_ALLELE_D, - ), - ( - ChoiceAddressPattern::SampleAllele(VdjSegment::J), - SAMPLE_ALLELE_J, - ), - ( - ChoiceAddressPattern::Trim { - segment: VdjSegment::V, - end: TrimEnd::Five, - }, - TRIM_V_5, - ), - ( - ChoiceAddressPattern::Trim { - segment: VdjSegment::D, - end: TrimEnd::Three, - }, - TRIM_D_3, - ), - ( - ChoiceAddressPattern::Trim { - segment: VdjSegment::J, - end: TrimEnd::Five, - }, - TRIM_J_5, - ), - (ChoiceAddressPattern::NpLength(NpSegment::Np1), NP1_LENGTH), - (ChoiceAddressPattern::NpLength(NpSegment::Np2), NP2_LENGTH), - ( - ChoiceAddressPattern::NpBase(NpSegment::Np1), - NP1_BASES_PATTERN, - ), - ( - ChoiceAddressPattern::NpBase(NpSegment::Np2), - NP2_BASES_PATTERN, - ), - ( - ChoiceAddressPattern::MutateUniformCount, - MUTATE_UNIFORM_COUNT, - ), - ( - ChoiceAddressPattern::MutateUniformSite, - MUTATE_UNIFORM_SITE_PATTERN, - ), - ( - ChoiceAddressPattern::MutateUniformBase, - MUTATE_UNIFORM_BASE_PATTERN, - ), - (ChoiceAddressPattern::MutateS5fCount, MUTATE_S5F_COUNT), - (ChoiceAddressPattern::MutateS5fSite, MUTATE_S5F_SITE_PATTERN), - (ChoiceAddressPattern::MutateS5fBase, MUTATE_S5F_BASE_PATTERN), - (ChoiceAddressPattern::CorruptPcrCount, CORRUPT_PCR_COUNT), - ( - ChoiceAddressPattern::CorruptPcrSite, - CORRUPT_PCR_SITE_PATTERN, - ), - ( - ChoiceAddressPattern::CorruptPcrBase, - CORRUPT_PCR_BASE_PATTERN, - ), - ( - ChoiceAddressPattern::CorruptQualityCount, - CORRUPT_QUALITY_COUNT, - ), - ( - ChoiceAddressPattern::CorruptQualitySite, - CORRUPT_QUALITY_SITE_PATTERN, - ), - ( - ChoiceAddressPattern::CorruptQualityBase, - CORRUPT_QUALITY_BASE_PATTERN, - ), - ( - ChoiceAddressPattern::CorruptContaminantApplied, - CORRUPT_CONTAMINANT_APPLIED, - ), - ( - ChoiceAddressPattern::CorruptContaminantBase, - CORRUPT_CONTAMINANT_BASES_PATTERN, - ), - (ChoiceAddressPattern::CorruptIndelCount, CORRUPT_INDEL_COUNT), - ( - ChoiceAddressPattern::CorruptIndelKind, - CORRUPT_INDEL_KIND_PATTERN, - ), - ( - ChoiceAddressPattern::CorruptIndelSite, - CORRUPT_INDEL_SITE_PATTERN, - ), - ( - ChoiceAddressPattern::CorruptIndelBase, - CORRUPT_INDEL_BASE_PATTERN, - ), - (ChoiceAddressPattern::CorruptNsCount, CORRUPT_NS_COUNT), - (ChoiceAddressPattern::CorruptNsSite, CORRUPT_NS_SITE_PATTERN), - ( - ChoiceAddressPattern::CorruptEndLoss(PrimeEnd::Five), - CORRUPT_END_LOSS_5, - ), - ( - ChoiceAddressPattern::CorruptEndLoss(PrimeEnd::Three), - CORRUPT_END_LOSS_3, - ), - ( - ChoiceAddressPattern::CorruptRevCompApplied, - CORRUPT_REV_COMP_APPLIED, - ), - ( - ChoiceAddressPattern::SampleAlleleDInverted, - SAMPLE_ALLELE_D_INVERTED, - ), - ( - ChoiceAddressPattern::ReceptorRevisionApplied, - RECEPTOR_REVISION_APPLIED, - ), - ( - ChoiceAddressPattern::ReceptorRevisionVAllele, - RECEPTOR_REVISION_V_ALLELE, - ), - ( - ChoiceAddressPattern::ReceptorRevisionVTrim3, - RECEPTOR_REVISION_V_TRIM_3, - ), - ( - ChoiceAddressPattern::PairedEndR1Length, - PAIRED_END_R1_LENGTH, - ), - ( - ChoiceAddressPattern::PairedEndR2Length, - PAIRED_END_R2_LENGTH, - ), - ( - ChoiceAddressPattern::PairedEndInsertSize, - PAIRED_END_INSERT_SIZE, - ), - ]; - - for (typed, raw) in cases { - assert_eq!(typed.to_string(), raw); - assert_eq!(ChoiceAddressPattern::parse(raw), Some(typed)); - assert_eq!(raw.parse::().unwrap(), typed); - } - } - - #[test] - fn typed_choice_address_pattern_rejects_concrete_indexes_and_unknown_strings() { - for raw in [ - "mutate.s5f.site[2]", - "np.np1.bases[0]", - "corrupt.indel.kind[7]", - "mutate.s5f.site[x]", - "assemble.v", - "custom.choice", - "", - ] { - assert_eq!(ChoiceAddressPattern::parse(raw), None, "raw={raw:?}"); - assert!(raw.parse::().is_err(), "raw={raw:?}"); - } - } - - // ── Frozen address vocabulary (compile-fence) ───────────────── - // - // Every persisted trace file carries an `address_schema_version`. - // Bumping the constant is the explicit signal that the on-disk - // vocabulary has changed; this test makes accidental drift loud. - // - // For each `ChoiceAddress` variant we pin: - // (a) the exact `Display` string (the on-disk spelling), and - // (b) that `Display → parse → Display` round-trips to the same - // string (the bidirectional contract). - // - // A code change that touches the Display or parse paths must - // either preserve the pinned strings or bump - // `ADDRESS_SCHEMA_VERSION` and re-pin the new spellings here. - - fn assert_pinned(addr: ChoiceAddress, expected: &str) { - let s = addr.to_string(); - assert_eq!( - s, expected, - "ChoiceAddress::Display drift detected; if intentional, bump ADDRESS_SCHEMA_VERSION", - ); - let parsed = ChoiceAddress::parse(expected) - .unwrap_or_else(|| panic!("frozen address string fails to parse: {expected:?}")); - assert_eq!( - parsed, addr, - "ChoiceAddress::parse drift; round-trip would break replay of committed traces", - ); - let round = parsed.to_string(); - assert_eq!(round, expected, "Display ∘ parse must equal Display"); - } - - #[test] - fn frozen_address_spellings_for_choice_address_schema_v1() { - assert_eq!( - ADDRESS_SCHEMA_VERSION, 1, - "if you bumped ADDRESS_SCHEMA_VERSION, also re-pin this test", - ); - - // SampleAllele.{v,d,j} - assert_pinned( - ChoiceAddress::SampleAllele(VdjSegment::V), - "sample_allele.v", - ); - assert_pinned( - ChoiceAddress::SampleAllele(VdjSegment::D), - "sample_allele.d", - ); - assert_pinned( - ChoiceAddress::SampleAllele(VdjSegment::J), - "sample_allele.j", - ); - - // Trim.{segment}_{end} - for (seg, seg_str) in [ - (VdjSegment::V, "v"), - (VdjSegment::D, "d"), - (VdjSegment::J, "j"), - ] { - for (end, end_str) in [(TrimEnd::Five, "5"), (TrimEnd::Three, "3")] { - let expected = format!("trim.{seg_str}_{end_str}"); - assert_pinned(ChoiceAddress::Trim { segment: seg, end }, &expected); - } - } - - // NpLength.{np1,np2} - assert_pinned( - ChoiceAddress::NpLength(NpSegment::Np1), - "np.np1.length", - ); - assert_pinned( - ChoiceAddress::NpLength(NpSegment::Np2), - "np.np2.length", - ); - - // NpBase[i] — pin a few representative indices. - for i in [0u32, 1, 5, 42] { - assert_pinned( - ChoiceAddress::NpBase { - segment: NpSegment::Np1, - index: i, - }, - &format!("np.np1.bases[{i}]"), - ); - assert_pinned( - ChoiceAddress::NpBase { - segment: NpSegment::Np2, - index: i, - }, - &format!("np.np2.bases[{i}]"), - ); - } - - // Mutate kernels. - assert_pinned(ChoiceAddress::MutateUniformCount, "mutate.uniform.count"); - for i in [0u32, 7, 99] { - assert_pinned( - ChoiceAddress::MutateUniformSite(i), - &format!("mutate.uniform.site[{i}]"), - ); - assert_pinned( - ChoiceAddress::MutateUniformBase(i), - &format!("mutate.uniform.base[{i}]"), - ); - } - assert_pinned(ChoiceAddress::MutateS5fCount, "mutate.s5f.count"); - for i in [0u32, 3] { - assert_pinned( - ChoiceAddress::MutateS5fSite(i), - &format!("mutate.s5f.site[{i}]"), - ); - assert_pinned( - ChoiceAddress::MutateS5fBase(i), - &format!("mutate.s5f.base[{i}]"), - ); - } - - // Corruption passes — PCR, quality, contaminant, ns, end_loss, - // rev_comp, indel. - assert_pinned(ChoiceAddress::CorruptPcrCount, "corrupt.pcr.count"); - for i in [0u32, 2] { - assert_pinned( - ChoiceAddress::CorruptPcrSite(i), - &format!("corrupt.pcr.error_site[{i}]"), - ); - assert_pinned( - ChoiceAddress::CorruptPcrBase(i), - &format!("corrupt.pcr.error_base[{i}]"), - ); - } - assert_pinned(ChoiceAddress::CorruptQualityCount, "corrupt.quality.count"); - for i in [0u32, 2] { - assert_pinned( - ChoiceAddress::CorruptQualitySite(i), - &format!("corrupt.quality.error_site[{i}]"), - ); - assert_pinned( - ChoiceAddress::CorruptQualityBase(i), - &format!("corrupt.quality.error_base[{i}]"), - ); - } - assert_pinned( - ChoiceAddress::CorruptContaminantApplied, - "corrupt.contaminant.applied", - ); - for i in [0u32, 5] { - assert_pinned( - ChoiceAddress::CorruptContaminantBase(i), - &format!("corrupt.contaminant.bases[{i}]"), - ); - } - assert_pinned(ChoiceAddress::CorruptIndelCount, "corrupt.indel.count"); - for i in [0u32, 4] { - assert_pinned( - ChoiceAddress::CorruptIndelKind(i), - &format!("corrupt.indel.kind[{i}]"), - ); - assert_pinned( - ChoiceAddress::CorruptIndelSite(i), - &format!("corrupt.indel.site[{i}]"), - ); - assert_pinned( - ChoiceAddress::CorruptIndelBase(i), - &format!("corrupt.indel.base[{i}]"), - ); - } - assert_pinned(ChoiceAddress::CorruptNsCount, "corrupt.ns.count"); - for i in [0u32, 3] { - assert_pinned( - ChoiceAddress::CorruptNsSite(i), - &format!("corrupt.ns.site[{i}]"), - ); - } - assert_pinned( - ChoiceAddress::CorruptEndLoss(PrimeEnd::Five), - "corrupt.end_loss.5", - ); - assert_pinned( - ChoiceAddress::CorruptEndLoss(PrimeEnd::Three), - "corrupt.end_loss.3", - ); - assert_pinned( - ChoiceAddress::CorruptRevCompApplied, - "corrupt.rev_comp.applied", - ); - - // D inversion (Slice C). The on-disk spelling sits under the - // existing `sample_allele.d` namespace so the persisted - // address-schema-version stays at v1 — no migration burden - // for traces emitted by pre-Slice-C engines that simply - // never wrote this address. - assert_pinned( - ChoiceAddress::SampleAlleleDInverted, - "sample_allele.d.inverted", - ); - - // Receptor revision (Slice C of the receptor-revision roadmap). - // New top-level `receptor_revision.*` namespace; additive - // under the v1 vocabulary policy — old traces don't reference - // these strings, new traces parse on engines that postdate - // this slice. No ADDRESS_SCHEMA_VERSION bump required. - assert_pinned( - ChoiceAddress::ReceptorRevisionApplied, - "receptor_revision.applied", - ); - assert_pinned( - ChoiceAddress::ReceptorRevisionVAllele, - "receptor_revision.v_allele", - ); - assert_pinned( - ChoiceAddress::ReceptorRevisionVTrim3, - "receptor_revision.v_trim_3", - ); - - // Paired-end / read layout (Slice C of the paired-end - // roadmap). New top-level `paired_end.*` namespace; same - // additive policy as receptor revision — old traces don't - // reference these strings, no ADDRESS_SCHEMA_VERSION bump. - assert_pinned( - ChoiceAddress::PairedEndR1Length, - "paired_end.r1_length", - ); - assert_pinned( - ChoiceAddress::PairedEndR2Length, - "paired_end.r2_length", - ); - assert_pinned( - ChoiceAddress::PairedEndInsertSize, - "paired_end.insert_size", - ); - - // P-nucleotide length per end (palindromic addition - // slice). New top-level `p.*` namespace; same additive - // policy as receptor revision / paired-end — old traces - // don't reference these strings, no - // ADDRESS_SCHEMA_VERSION bump. - assert_pinned(ChoiceAddress::PLength { end: PEnd::V3 }, "p.v_3.length"); - assert_pinned(ChoiceAddress::PLength { end: PEnd::D5 }, "p.d_5.length"); - assert_pinned(ChoiceAddress::PLength { end: PEnd::D3 }, "p.d_3.length"); - assert_pinned(ChoiceAddress::PLength { end: PEnd::J5 }, "p.j_5.length"); - - // Phased genotype (genotype-modeling PR1). New top-level - // `sample_haplotype` + `sample_gene.*` + `sample_allele_in_slot.*` - // namespaces; same additive policy as receptor revision / - // paired-end — old traces don't reference these strings, no - // ADDRESS_SCHEMA_VERSION bump. - assert_pinned(ChoiceAddress::SampleHaplotype, "sample_haplotype"); - assert_pinned(ChoiceAddress::SampleGene(VdjSegment::V), "sample_gene.v"); - assert_pinned(ChoiceAddress::SampleGene(VdjSegment::D), "sample_gene.d"); - assert_pinned(ChoiceAddress::SampleGene(VdjSegment::J), "sample_gene.j"); - assert_pinned( - ChoiceAddress::SampleAlleleInSlot(VdjSegment::V), - "sample_allele_in_slot.v", - ); - assert_pinned( - ChoiceAddress::SampleAlleleInSlot(VdjSegment::D), - "sample_allele_in_slot.d", - ); - assert_pinned( - ChoiceAddress::SampleAlleleInSlot(VdjSegment::J), - "sample_allele_in_slot.j", - ); - } -} +mod tests; diff --git a/engine_rs/src/address/tests.rs b/engine_rs/src/address/tests.rs new file mode 100644 index 0000000..76f14f2 --- /dev/null +++ b/engine_rs/src/address/tests.rs @@ -0,0 +1,627 @@ + use super::*; + + #[test] + fn segment_helpers_emit_existing_recombination_addresses() { + assert_eq!(sample_allele_vdj(Segment::V), "sample_allele.v"); + assert_eq!(sample_allele_vdj(Segment::D), "sample_allele.d"); + assert_eq!(sample_allele_vdj(Segment::J), "sample_allele.j"); + assert_eq!(trim_vdj(Segment::V, TrimEnd::Five), "trim.v_5"); + assert_eq!(trim_vdj(Segment::J, TrimEnd::Three), "trim.j_3"); + assert_eq!(assemble_vdj(Segment::D), "assemble.d"); + assert_eq!(np_length_region(Segment::Np1), "np.np1.length"); + } + + #[test] + fn parsers_accept_existing_indexed_addresses() { + assert_eq!( + parse_indexed("mutate.s5f.base[2]", MUTATE_S5F_BASE_PREFIX), + Some(2) + ); + assert_eq!( + ChoiceAddress::parse("np.np1.bases[17]"), + Some(ChoiceAddress::NpBase { + segment: NpSegment::Np1, + index: 17, + }) + ); + assert_eq!( + ChoiceAddress::parse("np.np2.bases[4]"), + Some(ChoiceAddress::NpBase { + segment: NpSegment::Np2, + index: 4, + }) + ); + assert_eq!(ChoiceAddress::parse("np.np2.bases[x]"), None); + assert_eq!( + ChoiceAddress::parse("np.np1.length"), + Some(ChoiceAddress::NpLength(NpSegment::Np1)) + ); + } + + #[test] + fn typed_choice_addresses_round_trip_persisted_strings() { + let cases = [ + (ChoiceAddress::SampleAllele(VdjSegment::V), SAMPLE_ALLELE_V), + (ChoiceAddress::SampleAllele(VdjSegment::D), SAMPLE_ALLELE_D), + (ChoiceAddress::SampleAllele(VdjSegment::J), SAMPLE_ALLELE_J), + ( + ChoiceAddress::Trim { + segment: VdjSegment::V, + end: TrimEnd::Three, + }, + TRIM_V_3, + ), + ( + ChoiceAddress::Trim { + segment: VdjSegment::D, + end: TrimEnd::Five, + }, + TRIM_D_5, + ), + ( + ChoiceAddress::Trim { + segment: VdjSegment::J, + end: TrimEnd::Five, + }, + TRIM_J_5, + ), + (ChoiceAddress::NpLength(NpSegment::Np1), NP1_LENGTH), + ( + ChoiceAddress::NpBase { + segment: NpSegment::Np2, + index: 17, + }, + "np.np2.bases[17]", + ), + (ChoiceAddress::MutateUniformCount, MUTATE_UNIFORM_COUNT), + ( + ChoiceAddress::MutateUniformSite(2), + "mutate.uniform.site[2]", + ), + ( + ChoiceAddress::MutateUniformBase(3), + "mutate.uniform.base[3]", + ), + (ChoiceAddress::MutateS5fCount, MUTATE_S5F_COUNT), + (ChoiceAddress::MutateS5fSite(4), "mutate.s5f.site[4]"), + (ChoiceAddress::MutateS5fBase(5), "mutate.s5f.base[5]"), + (ChoiceAddress::CorruptPcrCount, CORRUPT_PCR_COUNT), + ( + ChoiceAddress::CorruptPcrSite(6), + "corrupt.pcr.error_site[6]", + ), + ( + ChoiceAddress::CorruptPcrBase(7), + "corrupt.pcr.error_base[7]", + ), + (ChoiceAddress::CorruptQualityCount, CORRUPT_QUALITY_COUNT), + ( + ChoiceAddress::CorruptQualitySite(8), + "corrupt.quality.error_site[8]", + ), + ( + ChoiceAddress::CorruptQualityBase(9), + "corrupt.quality.error_base[9]", + ), + ( + ChoiceAddress::CorruptContaminantApplied, + CORRUPT_CONTAMINANT_APPLIED, + ), + ( + ChoiceAddress::CorruptContaminantBase(10), + "corrupt.contaminant.bases[10]", + ), + (ChoiceAddress::CorruptIndelCount, CORRUPT_INDEL_COUNT), + ( + ChoiceAddress::CorruptIndelKind(11), + "corrupt.indel.kind[11]", + ), + ( + ChoiceAddress::CorruptIndelSite(12), + "corrupt.indel.site[12]", + ), + ( + ChoiceAddress::CorruptIndelBase(13), + "corrupt.indel.base[13]", + ), + (ChoiceAddress::CorruptNsCount, CORRUPT_NS_COUNT), + (ChoiceAddress::CorruptNsSite(14), "corrupt.ns.site[14]"), + ( + ChoiceAddress::CorruptEndLoss(PrimeEnd::Five), + CORRUPT_END_LOSS_5, + ), + ( + ChoiceAddress::CorruptEndLoss(PrimeEnd::Three), + CORRUPT_END_LOSS_3, + ), + ( + ChoiceAddress::CorruptRevCompApplied, + CORRUPT_REV_COMP_APPLIED, + ), + ( + ChoiceAddress::SampleAlleleDInverted, + SAMPLE_ALLELE_D_INVERTED, + ), + ( + ChoiceAddress::ReceptorRevisionApplied, + RECEPTOR_REVISION_APPLIED, + ), + ( + ChoiceAddress::ReceptorRevisionVAllele, + RECEPTOR_REVISION_V_ALLELE, + ), + ( + ChoiceAddress::ReceptorRevisionVTrim3, + RECEPTOR_REVISION_V_TRIM_3, + ), + (ChoiceAddress::PairedEndR1Length, PAIRED_END_R1_LENGTH), + (ChoiceAddress::PairedEndR2Length, PAIRED_END_R2_LENGTH), + ( + ChoiceAddress::PairedEndInsertSize, + PAIRED_END_INSERT_SIZE, + ), + ]; + + for (typed, raw) in cases { + assert_eq!(typed.to_string(), raw); + assert_eq!(ChoiceAddress::parse(raw), Some(typed)); + assert_eq!(raw.parse::().unwrap(), typed); + } + } + + #[test] + fn typed_choice_address_rejects_patterns_and_unknown_strings() { + for raw in [ + "mutate.s5f.site[0..n]", + "np.np1.bases[x]", + "assemble.v", + "sample_allele.np1", + "trim.np1_3", + "custom.choice", + "", + ] { + assert_eq!(ChoiceAddress::parse(raw), None, "raw={raw:?}"); + assert!(raw.parse::().is_err(), "raw={raw:?}"); + } + } + + #[test] + fn typed_segment_conversions_reject_wrong_segment_family() { + assert_eq!(VdjSegment::try_from(Segment::V), Ok(VdjSegment::V)); + assert_eq!(VdjSegment::try_from(Segment::Np1), Err(())); + assert_eq!(NpSegment::try_from(Segment::Np2), Ok(NpSegment::Np2)); + assert_eq!(NpSegment::try_from(Segment::J), Err(())); + } + + #[test] + fn typed_choice_address_patterns_round_trip_report_strings() { + let cases = [ + ( + ChoiceAddressPattern::SampleAllele(VdjSegment::V), + SAMPLE_ALLELE_V, + ), + ( + ChoiceAddressPattern::SampleAllele(VdjSegment::D), + SAMPLE_ALLELE_D, + ), + ( + ChoiceAddressPattern::SampleAllele(VdjSegment::J), + SAMPLE_ALLELE_J, + ), + ( + ChoiceAddressPattern::Trim { + segment: VdjSegment::V, + end: TrimEnd::Five, + }, + TRIM_V_5, + ), + ( + ChoiceAddressPattern::Trim { + segment: VdjSegment::D, + end: TrimEnd::Three, + }, + TRIM_D_3, + ), + ( + ChoiceAddressPattern::Trim { + segment: VdjSegment::J, + end: TrimEnd::Five, + }, + TRIM_J_5, + ), + (ChoiceAddressPattern::NpLength(NpSegment::Np1), NP1_LENGTH), + (ChoiceAddressPattern::NpLength(NpSegment::Np2), NP2_LENGTH), + ( + ChoiceAddressPattern::NpBase(NpSegment::Np1), + NP1_BASES_PATTERN, + ), + ( + ChoiceAddressPattern::NpBase(NpSegment::Np2), + NP2_BASES_PATTERN, + ), + ( + ChoiceAddressPattern::MutateUniformCount, + MUTATE_UNIFORM_COUNT, + ), + ( + ChoiceAddressPattern::MutateUniformSite, + MUTATE_UNIFORM_SITE_PATTERN, + ), + ( + ChoiceAddressPattern::MutateUniformBase, + MUTATE_UNIFORM_BASE_PATTERN, + ), + (ChoiceAddressPattern::MutateS5fCount, MUTATE_S5F_COUNT), + (ChoiceAddressPattern::MutateS5fSite, MUTATE_S5F_SITE_PATTERN), + (ChoiceAddressPattern::MutateS5fBase, MUTATE_S5F_BASE_PATTERN), + (ChoiceAddressPattern::CorruptPcrCount, CORRUPT_PCR_COUNT), + ( + ChoiceAddressPattern::CorruptPcrSite, + CORRUPT_PCR_SITE_PATTERN, + ), + ( + ChoiceAddressPattern::CorruptPcrBase, + CORRUPT_PCR_BASE_PATTERN, + ), + ( + ChoiceAddressPattern::CorruptQualityCount, + CORRUPT_QUALITY_COUNT, + ), + ( + ChoiceAddressPattern::CorruptQualitySite, + CORRUPT_QUALITY_SITE_PATTERN, + ), + ( + ChoiceAddressPattern::CorruptQualityBase, + CORRUPT_QUALITY_BASE_PATTERN, + ), + ( + ChoiceAddressPattern::CorruptContaminantApplied, + CORRUPT_CONTAMINANT_APPLIED, + ), + ( + ChoiceAddressPattern::CorruptContaminantBase, + CORRUPT_CONTAMINANT_BASES_PATTERN, + ), + (ChoiceAddressPattern::CorruptIndelCount, CORRUPT_INDEL_COUNT), + ( + ChoiceAddressPattern::CorruptIndelKind, + CORRUPT_INDEL_KIND_PATTERN, + ), + ( + ChoiceAddressPattern::CorruptIndelSite, + CORRUPT_INDEL_SITE_PATTERN, + ), + ( + ChoiceAddressPattern::CorruptIndelBase, + CORRUPT_INDEL_BASE_PATTERN, + ), + (ChoiceAddressPattern::CorruptNsCount, CORRUPT_NS_COUNT), + (ChoiceAddressPattern::CorruptNsSite, CORRUPT_NS_SITE_PATTERN), + ( + ChoiceAddressPattern::CorruptEndLoss(PrimeEnd::Five), + CORRUPT_END_LOSS_5, + ), + ( + ChoiceAddressPattern::CorruptEndLoss(PrimeEnd::Three), + CORRUPT_END_LOSS_3, + ), + ( + ChoiceAddressPattern::CorruptRevCompApplied, + CORRUPT_REV_COMP_APPLIED, + ), + ( + ChoiceAddressPattern::SampleAlleleDInverted, + SAMPLE_ALLELE_D_INVERTED, + ), + ( + ChoiceAddressPattern::ReceptorRevisionApplied, + RECEPTOR_REVISION_APPLIED, + ), + ( + ChoiceAddressPattern::ReceptorRevisionVAllele, + RECEPTOR_REVISION_V_ALLELE, + ), + ( + ChoiceAddressPattern::ReceptorRevisionVTrim3, + RECEPTOR_REVISION_V_TRIM_3, + ), + ( + ChoiceAddressPattern::PairedEndR1Length, + PAIRED_END_R1_LENGTH, + ), + ( + ChoiceAddressPattern::PairedEndR2Length, + PAIRED_END_R2_LENGTH, + ), + ( + ChoiceAddressPattern::PairedEndInsertSize, + PAIRED_END_INSERT_SIZE, + ), + ]; + + for (typed, raw) in cases { + assert_eq!(typed.to_string(), raw); + assert_eq!(ChoiceAddressPattern::parse(raw), Some(typed)); + assert_eq!(raw.parse::().unwrap(), typed); + } + } + + #[test] + fn typed_choice_address_pattern_rejects_concrete_indexes_and_unknown_strings() { + for raw in [ + "mutate.s5f.site[2]", + "np.np1.bases[0]", + "corrupt.indel.kind[7]", + "mutate.s5f.site[x]", + "assemble.v", + "custom.choice", + "", + ] { + assert_eq!(ChoiceAddressPattern::parse(raw), None, "raw={raw:?}"); + assert!(raw.parse::().is_err(), "raw={raw:?}"); + } + } + + // ── Frozen address vocabulary (compile-fence) ───────────────── + // + // Every persisted trace file carries an `address_schema_version`. + // Bumping the constant is the explicit signal that the on-disk + // vocabulary has changed; this test makes accidental drift loud. + // + // For each `ChoiceAddress` variant we pin: + // (a) the exact `Display` string (the on-disk spelling), and + // (b) that `Display → parse → Display` round-trips to the same + // string (the bidirectional contract). + // + // A code change that touches the Display or parse paths must + // either preserve the pinned strings or bump + // `ADDRESS_SCHEMA_VERSION` and re-pin the new spellings here. + + fn assert_pinned(addr: ChoiceAddress, expected: &str) { + let s = addr.to_string(); + assert_eq!( + s, expected, + "ChoiceAddress::Display drift detected; if intentional, bump ADDRESS_SCHEMA_VERSION", + ); + let parsed = ChoiceAddress::parse(expected) + .unwrap_or_else(|| panic!("frozen address string fails to parse: {expected:?}")); + assert_eq!( + parsed, addr, + "ChoiceAddress::parse drift; round-trip would break replay of committed traces", + ); + let round = parsed.to_string(); + assert_eq!(round, expected, "Display ∘ parse must equal Display"); + } + + #[test] + fn frozen_address_spellings_for_choice_address_schema_v1() { + assert_eq!( + ADDRESS_SCHEMA_VERSION, 1, + "if you bumped ADDRESS_SCHEMA_VERSION, also re-pin this test", + ); + + // SampleAllele.{v,d,j} + assert_pinned( + ChoiceAddress::SampleAllele(VdjSegment::V), + "sample_allele.v", + ); + assert_pinned( + ChoiceAddress::SampleAllele(VdjSegment::D), + "sample_allele.d", + ); + assert_pinned( + ChoiceAddress::SampleAllele(VdjSegment::J), + "sample_allele.j", + ); + + // Trim.{segment}_{end} + for (seg, seg_str) in [ + (VdjSegment::V, "v"), + (VdjSegment::D, "d"), + (VdjSegment::J, "j"), + ] { + for (end, end_str) in [(TrimEnd::Five, "5"), (TrimEnd::Three, "3")] { + let expected = format!("trim.{seg_str}_{end_str}"); + assert_pinned(ChoiceAddress::Trim { segment: seg, end }, &expected); + } + } + + // NpLength.{np1,np2} + assert_pinned( + ChoiceAddress::NpLength(NpSegment::Np1), + "np.np1.length", + ); + assert_pinned( + ChoiceAddress::NpLength(NpSegment::Np2), + "np.np2.length", + ); + + // NpBase[i] — pin a few representative indices. + for i in [0u32, 1, 5, 42] { + assert_pinned( + ChoiceAddress::NpBase { + segment: NpSegment::Np1, + index: i, + }, + &format!("np.np1.bases[{i}]"), + ); + assert_pinned( + ChoiceAddress::NpBase { + segment: NpSegment::Np2, + index: i, + }, + &format!("np.np2.bases[{i}]"), + ); + } + + // Mutate kernels. + assert_pinned(ChoiceAddress::MutateUniformCount, "mutate.uniform.count"); + for i in [0u32, 7, 99] { + assert_pinned( + ChoiceAddress::MutateUniformSite(i), + &format!("mutate.uniform.site[{i}]"), + ); + assert_pinned( + ChoiceAddress::MutateUniformBase(i), + &format!("mutate.uniform.base[{i}]"), + ); + } + assert_pinned(ChoiceAddress::MutateS5fCount, "mutate.s5f.count"); + for i in [0u32, 3] { + assert_pinned( + ChoiceAddress::MutateS5fSite(i), + &format!("mutate.s5f.site[{i}]"), + ); + assert_pinned( + ChoiceAddress::MutateS5fBase(i), + &format!("mutate.s5f.base[{i}]"), + ); + } + + // Corruption passes — PCR, quality, contaminant, ns, end_loss, + // rev_comp, indel. + assert_pinned(ChoiceAddress::CorruptPcrCount, "corrupt.pcr.count"); + for i in [0u32, 2] { + assert_pinned( + ChoiceAddress::CorruptPcrSite(i), + &format!("corrupt.pcr.error_site[{i}]"), + ); + assert_pinned( + ChoiceAddress::CorruptPcrBase(i), + &format!("corrupt.pcr.error_base[{i}]"), + ); + } + assert_pinned(ChoiceAddress::CorruptQualityCount, "corrupt.quality.count"); + for i in [0u32, 2] { + assert_pinned( + ChoiceAddress::CorruptQualitySite(i), + &format!("corrupt.quality.error_site[{i}]"), + ); + assert_pinned( + ChoiceAddress::CorruptQualityBase(i), + &format!("corrupt.quality.error_base[{i}]"), + ); + } + assert_pinned( + ChoiceAddress::CorruptContaminantApplied, + "corrupt.contaminant.applied", + ); + for i in [0u32, 5] { + assert_pinned( + ChoiceAddress::CorruptContaminantBase(i), + &format!("corrupt.contaminant.bases[{i}]"), + ); + } + assert_pinned(ChoiceAddress::CorruptIndelCount, "corrupt.indel.count"); + for i in [0u32, 4] { + assert_pinned( + ChoiceAddress::CorruptIndelKind(i), + &format!("corrupt.indel.kind[{i}]"), + ); + assert_pinned( + ChoiceAddress::CorruptIndelSite(i), + &format!("corrupt.indel.site[{i}]"), + ); + assert_pinned( + ChoiceAddress::CorruptIndelBase(i), + &format!("corrupt.indel.base[{i}]"), + ); + } + assert_pinned(ChoiceAddress::CorruptNsCount, "corrupt.ns.count"); + for i in [0u32, 3] { + assert_pinned( + ChoiceAddress::CorruptNsSite(i), + &format!("corrupt.ns.site[{i}]"), + ); + } + assert_pinned( + ChoiceAddress::CorruptEndLoss(PrimeEnd::Five), + "corrupt.end_loss.5", + ); + assert_pinned( + ChoiceAddress::CorruptEndLoss(PrimeEnd::Three), + "corrupt.end_loss.3", + ); + assert_pinned( + ChoiceAddress::CorruptRevCompApplied, + "corrupt.rev_comp.applied", + ); + + // D inversion (Slice C). The on-disk spelling sits under the + // existing `sample_allele.d` namespace so the persisted + // address-schema-version stays at v1 — no migration burden + // for traces emitted by pre-Slice-C engines that simply + // never wrote this address. + assert_pinned( + ChoiceAddress::SampleAlleleDInverted, + "sample_allele.d.inverted", + ); + + // Receptor revision (Slice C of the receptor-revision roadmap). + // New top-level `receptor_revision.*` namespace; additive + // under the v1 vocabulary policy — old traces don't reference + // these strings, new traces parse on engines that postdate + // this slice. No ADDRESS_SCHEMA_VERSION bump required. + assert_pinned( + ChoiceAddress::ReceptorRevisionApplied, + "receptor_revision.applied", + ); + assert_pinned( + ChoiceAddress::ReceptorRevisionVAllele, + "receptor_revision.v_allele", + ); + assert_pinned( + ChoiceAddress::ReceptorRevisionVTrim3, + "receptor_revision.v_trim_3", + ); + + // Paired-end / read layout (Slice C of the paired-end + // roadmap). New top-level `paired_end.*` namespace; same + // additive policy as receptor revision — old traces don't + // reference these strings, no ADDRESS_SCHEMA_VERSION bump. + assert_pinned( + ChoiceAddress::PairedEndR1Length, + "paired_end.r1_length", + ); + assert_pinned( + ChoiceAddress::PairedEndR2Length, + "paired_end.r2_length", + ); + assert_pinned( + ChoiceAddress::PairedEndInsertSize, + "paired_end.insert_size", + ); + + // P-nucleotide length per end (palindromic addition + // slice). New top-level `p.*` namespace; same additive + // policy as receptor revision / paired-end — old traces + // don't reference these strings, no + // ADDRESS_SCHEMA_VERSION bump. + assert_pinned(ChoiceAddress::PLength { end: PEnd::V3 }, "p.v_3.length"); + assert_pinned(ChoiceAddress::PLength { end: PEnd::D5 }, "p.d_5.length"); + assert_pinned(ChoiceAddress::PLength { end: PEnd::D3 }, "p.d_3.length"); + assert_pinned(ChoiceAddress::PLength { end: PEnd::J5 }, "p.j_5.length"); + + // Phased genotype (genotype-modeling PR1). New top-level + // `sample_haplotype` + `sample_gene.*` + `sample_allele_in_slot.*` + // namespaces; same additive policy as receptor revision / + // paired-end — old traces don't reference these strings, no + // ADDRESS_SCHEMA_VERSION bump. + assert_pinned(ChoiceAddress::SampleHaplotype, "sample_haplotype"); + assert_pinned(ChoiceAddress::SampleGene(VdjSegment::V), "sample_gene.v"); + assert_pinned(ChoiceAddress::SampleGene(VdjSegment::D), "sample_gene.d"); + assert_pinned(ChoiceAddress::SampleGene(VdjSegment::J), "sample_gene.j"); + assert_pinned( + ChoiceAddress::SampleAlleleInSlot(VdjSegment::V), + "sample_allele_in_slot.v", + ); + assert_pinned( + ChoiceAddress::SampleAlleleInSlot(VdjSegment::D), + "sample_allele_in_slot.d", + ); + assert_pinned( + ChoiceAddress::SampleAlleleInSlot(VdjSegment::J), + "sample_allele_in_slot.j", + ); + } diff --git a/engine_rs/src/airr_record/validate.rs b/engine_rs/src/airr_record/validate.rs index c0a6ab9..dfee9ea 100644 --- a/engine_rs/src/airr_record/validate.rs +++ b/engine_rs/src/airr_record/validate.rs @@ -9,19 +9,19 @@ //! See [`docs/airr_record_validator.md`](../../../docs/airr_record_validator.md) //! for the full check catalogue and the design discussion. -use crate::address::{ChoiceAddress, PrimeEnd, VdjSegment}; -use crate::ir::{Segment, SimulationEvent}; -use crate::live_call::scoring::{ - allele_pool_for_segment, score_alleles_with_extensions, tie_set_ids_at_max_score, -}; +use crate::address::PrimeEnd; +use crate::ir::Segment; use crate::pass::Outcome; use crate::refdata::RefDataConfig; -use crate::trace::ChoiceValue; -use super::junction::{anchor_amino_acid_preserved, anchor_pool_position, junction_has_stop}; -use super::projection::lookup_allele; use super::record::AirrRecord; -use super::sequence::pool_bases; + +mod allele_oracle; +mod counters; +mod junction; +mod paired_end; +mod region_invariants; +mod structural; /// One reported issue. Variants are stable identifiers; downstream /// CI / dashboards can pattern-match on them. @@ -345,1128 +345,12 @@ pub fn validate_airr_record( let mut issues = Vec::new(); let sim = outcome.final_simulation(); - check_structural(record, sim, &mut issues); - check_counters(record, outcome, refdata, &mut issues); - check_junction(record, sim, refdata, &mut issues); - check_allele_oracle(record, outcome, refdata, &mut issues); - check_region_and_hypothesis_invariants(sim, &mut issues); - check_paired_end_defaults(record, &mut issues); + structural::check_structural(record, sim, &mut issues); + counters::check_counters(record, outcome, refdata, &mut issues); + junction::check_junction(record, sim, refdata, &mut issues); + allele_oracle::check_allele_oracle(record, outcome, refdata, &mut issues); + region_invariants::check_region_and_hypothesis_invariants(sim, &mut issues); + paired_end::check_paired_end_defaults(record, &mut issues); issues } - -// ────────────────────────────────────────────────────────────────── -// C6: Paired-end / read-layout invariants -// -// Slice A: enforces the no-layout default invariant (`read_layout -// == ""` ⇒ every paired-end field is at its default). -// -// Slice B: extends the dispatch with per-layout geometry checks. -// When `read_layout == "paired_end"`, the validator re-derives -// R1/R2 windows from the rules in `docs/paired_end_design.md` §8 -// and surfaces any divergence as one of the four reserved variants -// (`ReadWindowOutOfBounds`, `ReadSequenceMismatch`, -// `ReadInsertSizeMismatch`). Unknown non-empty layouts surface as -// `ReadLayoutMismatch`. `"single_end"` is documented as reserved -// (§2.2) and treated as a no-op for now. -// ────────────────────────────────────────────────────────────────── - -fn check_paired_end_defaults( - record: &AirrRecord, - issues: &mut Vec, -) { - match record.read_layout.as_str() { - // Slice A: no-layout default invariant. - "" => check_paired_end_default_values(record, issues), - // Slice B: geometry against the projection rules. - "paired_end" => check_paired_end_geometry(record, issues), - // Reserved (per §2.2). A future slice may attach a - // single-read geometry check; today we don't validate - // anything beyond accepting the layout string. - "single_end" => {} - // Anything else is an unsupported layout value. Surface - // the structured mismatch so a typo (`"pair_end"`) or a - // refactor that introduced a new layout without wiring - // it through the validator fails closed. - _ => issues.push(RecordValidationIssue::ReadLayoutMismatch { - reported: record.read_layout.clone(), - expected: r#"one of: "", "paired_end", "single_end""#.to_string(), - }), - } -} - -fn check_paired_end_default_values( - record: &AirrRecord, - issues: &mut Vec, -) { - if !record.r1_sequence.is_empty() { - issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { - field: PairedEndField::R1Sequence, - }); - } - if !record.r2_sequence.is_empty() { - issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { - field: PairedEndField::R2Sequence, - }); - } - if record.r1_start.is_some() { - issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { - field: PairedEndField::R1Start, - }); - } - if record.r1_end.is_some() { - issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { - field: PairedEndField::R1End, - }); - } - if record.r2_start.is_some() { - issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { - field: PairedEndField::R2Start, - }); - } - if record.r2_end.is_some() { - issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { - field: PairedEndField::R2End, - }); - } - if record.insert_size != 0 { - issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { - field: PairedEndField::InsertSize, - }); - } -} - -fn check_paired_end_geometry( - record: &AirrRecord, - issues: &mut Vec, -) { - let seq_len = record.sequence_length; - let seq = &record.sequence; - - // R1 window: bounds + byte equality. - match resolve_window(record.r1_start, record.r1_end, seq_len) { - WindowResolution::Valid { start, end } => { - let expected = seq[start as usize..end as usize].to_string(); - if record.r1_sequence != expected { - issues.push(RecordValidationIssue::ReadSequenceMismatch { - side: PairedEndRead::R1, - reported: record.r1_sequence.clone(), - expected, - }); - } - } - WindowResolution::OutOfBounds { start, end } => { - issues.push(RecordValidationIssue::ReadWindowOutOfBounds { - side: PairedEndRead::R1, - start, - end, - sequence_length: seq_len, - }); - } - } - - // R2 window: bounds + reverse-complement byte equality + insert - // size consistency. The audit pins - // `insert_size == r2_end` (§8) — the only insert-size-mismatch - // surface today. - match resolve_window(record.r2_start, record.r2_end, seq_len) { - WindowResolution::Valid { start, end } => { - let r2_inner = &seq[start as usize..end as usize]; - let expected = super::sequence::reverse_complement(r2_inner); - if record.r2_sequence != expected { - issues.push(RecordValidationIssue::ReadSequenceMismatch { - side: PairedEndRead::R2, - reported: record.r2_sequence.clone(), - expected, - }); - } - if record.insert_size != end { - issues.push(RecordValidationIssue::ReadInsertSizeMismatch { - reported: record.insert_size, - expected: end, - }); - } - } - WindowResolution::OutOfBounds { start, end } => { - issues.push(RecordValidationIssue::ReadWindowOutOfBounds { - side: PairedEndRead::R2, - start, - end, - sequence_length: seq_len, - }); - // With R2 bounds unresolved we can't recompute the - // expected insert size; skip the insert-size check - // rather than fabricate a sentinel. The window - // out-of-bounds issue is the actionable signal. - } - } -} - -/// Result of resolving a paired-end window `(start, end)` against -/// the projected `sequence_length`. `OutOfBounds` collapses missing -/// coords (`None`) and out-of-range coords (negative, swapped, -/// past the end) into one variant whose `start`/`end` fields use -/// the sentinel `-1` for missing values — surfaces cleanly in the -/// structured-issue payload without inventing a new variant. -enum WindowResolution { - Valid { start: i64, end: i64 }, - OutOfBounds { start: i64, end: i64 }, -} - -fn resolve_window( - start: Option, - end: Option, - sequence_length: i64, -) -> WindowResolution { - let start_val = start.unwrap_or(-1); - let end_val = end.unwrap_or(-1); - match (start, end) { - (Some(s), Some(e)) if s >= 0 && e >= s && e <= sequence_length => { - WindowResolution::Valid { start: s, end: e } - } - _ => WindowResolution::OutOfBounds { - start: start_val, - end: end_val, - }, - } -} - -// ────────────────────────────────────────────────────────────────── -// C1: Structural record invariants -// ────────────────────────────────────────────────────────────────── - -fn check_structural( - record: &AirrRecord, - sim: &crate::ir::Simulation, - issues: &mut Vec, -) { - // sequence_length matches actual sequence string bytes. - let actual_bytes = record.sequence.len(); - if record.sequence_length as usize != actual_bytes { - issues.push(RecordValidationIssue::SequenceLengthMismatch { - reported: record.sequence_length, - actual_bytes, - }); - } - - // sequence content matches the pool's bases (case-folded). - // Compare case-insensitively because sequencing errors lowercase - // their substitutions, and the AIRR record may carry rev-comp. - let pool_seq = String::from_utf8(pool_bases(sim)).unwrap_or_default(); - if !record.rev_comp && !record.sequence.eq_ignore_ascii_case(&pool_seq) { - let prefix_len = record.sequence.len().min(pool_seq.len()).min(40); - issues.push(RecordValidationIssue::SequenceContentMismatch { - reported_prefix: record.sequence.chars().take(prefix_len).collect(), - actual_prefix: pool_seq.chars().take(prefix_len).collect(), - }); - } - - // Coordinate ordering for V/D/J segments and germline ranges. - check_coord_pair( - Segment::V, - record.v_sequence_start, - record.v_sequence_end, - false, - issues, - ); - check_coord_pair( - Segment::V, - record.v_germline_start, - record.v_germline_end, - true, - issues, - ); - check_coord_pair( - Segment::D, - record.d_sequence_start, - record.d_sequence_end, - false, - issues, - ); - check_coord_pair( - Segment::D, - record.d_germline_start, - record.d_germline_end, - true, - issues, - ); - check_coord_pair( - Segment::J, - record.j_sequence_start, - record.j_sequence_end, - false, - issues, - ); - check_coord_pair( - Segment::J, - record.j_germline_start, - record.j_germline_end, - true, - issues, - ); - - // CIGAR span sanity: M+I ops must equal the sequence-side span. - for (seg, start, end, cigar) in [ - ( - Segment::V, - record.v_sequence_start, - record.v_sequence_end, - &record.v_cigar, - ), - ( - Segment::D, - record.d_sequence_start, - record.d_sequence_end, - &record.d_cigar, - ), - ( - Segment::J, - record.j_sequence_start, - record.j_sequence_end, - &record.j_cigar, - ), - ] { - if cigar.is_empty() { - continue; - } - match parse_cigar(cigar) { - Ok(ops) => { - let query_span: usize = ops - .iter() - .filter(|(_, op)| *op == b'M' || *op == b'I') - .map(|(n, _)| *n) - .sum(); - if let (Some(s), Some(e)) = (start, end) { - let sequence_span = (e - s).max(0) as usize; - if query_span != sequence_span { - issues.push(RecordValidationIssue::CigarSpanMismatch { - segment: seg, - cigar_query_span: query_span, - sequence_span, - }); - } - } - } - Err(reason) => issues.push(RecordValidationIssue::CigarReadsInvalid { - segment: seg, - reason, - }), - } - } -} - -fn check_coord_pair( - segment: Segment, - start: Option, - end: Option, - is_germline: bool, - issues: &mut Vec, -) { - if let (Some(s), Some(e)) = (start, end) { - // Half-open: end > start is required (or both 0 for empty). - if s < 0 || e < s { - let issue = if is_germline { - RecordValidationIssue::GermlineCoordinatesOutOfOrder { - segment, - start: s, - end: e, - } - } else { - RecordValidationIssue::SegmentCoordinatesOutOfOrder { - segment, - start: s, - end: e, - } - }; - issues.push(issue); - } - } -} - -fn parse_cigar(cigar: &str) -> Result, String> { - let mut ops = Vec::new(); - let mut digits = String::new(); - for ch in cigar.chars() { - if ch.is_ascii_digit() { - digits.push(ch); - } else { - let n: usize = digits - .parse() - .map_err(|_| format!("non-numeric op length in CIGAR {cigar:?}"))?; - if !matches!(ch as u8, b'M' | b'I' | b'D' | b'S' | b'N' | b'P' | b'X' | b'=') { - return Err(format!("unrecognized CIGAR op {ch:?} in {cigar:?}")); - } - ops.push((n, ch as u8)); - digits.clear(); - } - } - if !digits.is_empty() { - return Err(format!("trailing digits in CIGAR {cigar:?}")); - } - Ok(ops) -} - -// ────────────────────────────────────────────────────────────────── -// C2: Counter provenance -// ────────────────────────────────────────────────────────────────── - -fn check_counters( - record: &AirrRecord, - outcome: &Outcome, - refdata: &RefDataConfig, - issues: &mut Vec, -) { - let sim = outcome.final_simulation(); - - // n_mutations comes from sim.mutation_count (set by S5F / Uniform - // at seal). Direct equality. - if record.n_mutations != sim.mutation_count as i64 { - issues.push(RecordValidationIssue::NMutationsMismatch { - reported: record.n_mutations, - sim_count: sim.mutation_count as i64, - }); - } - - // n_pcr_errors / n_quality_errors from trace. - let trace_int = |addr: ChoiceAddress| -> i64 { - outcome - .trace - .find_choice(addr) - .and_then(|rec| match rec.value { - ChoiceValue::Int(n) => Some(n), - _ => None, - }) - .unwrap_or(0) - }; - - let pcr_attempts = trace_int(ChoiceAddress::CorruptPcrCount); - if record.n_pcr_errors != pcr_attempts { - issues.push(RecordValidationIssue::NPcrErrorsMismatch { - reported: record.n_pcr_errors, - trace_count: pcr_attempts, - }); - } - - let quality_attempts = trace_int(ChoiceAddress::CorruptQualityCount); - if record.n_quality_errors != quality_attempts { - issues.push(RecordValidationIssue::NQualityErrorsMismatch { - reported: record.n_quality_errors, - trace_count: quality_attempts, - }); - } - - // n_indels and per-segment indel counters from event ledger of the - // corrupt.indel pass only (per audit §6.1 / §6.2). - let mut total = 0i64; - let mut by_segment = [0i64; Segment::COUNT]; - for er in outcome.events() { - if er.pass_name != crate::address::CORRUPT_INDEL { - continue; - } - for ev in &er.simulation_events { - let segment = match ev { - SimulationEvent::IndelInserted { segment, .. } => *segment, - SimulationEvent::IndelDeleted { segment, .. } => *segment, - _ => continue, - }; - total += 1; - by_segment[segment as usize] += 1; - } - } - if record.n_indels != total { - issues.push(RecordValidationIssue::NIndelsMismatch { - reported: record.n_indels, - event_count: total, - }); - } - for (seg, reported) in [ - (Segment::V, record.n_v_indels), - (Segment::D, record.n_d_indels), - (Segment::J, record.n_j_indels), - ] { - let expected = by_segment[seg as usize]; - if reported != expected { - issues.push(RecordValidationIssue::NSegmentIndelsMismatch { - segment: seg, - reported, - event_count: expected, - }); - } - } - - // Per-segment SHM counters from the event ledger of the - // mutate.{uniform,s5f} passes only. Mirrors the indel walk - // above but counts `BaseChanged` events instead of indel - // events, and rolls NP1+NP2 into the single NP bucket. The - // four per-bucket checks fire `N*MutationsMismatch` on - // disagreement; the sum-invariant cross-check fires - // `MutationCountSumMismatch`. - // - // Same walk also produces the V-subregion partition (Slice — - // V-Subregion Mutation Counters): each V `BaseChanged.germline_pos` - // is matched against the assigned V allele's `subregions` table - // from scratch, independent of the projection's bucketing. Six - // per-bucket mismatches plus a cross-field sum invariant - // (`VSubregionMutationCountSumMismatch`) fire on disagreement. - let mut shm_by_segment = [0i64; Segment::COUNT]; - let mut shm_by_subregion = [0i64; 5]; - let mut shm_v_unannotated = 0i64; - let v_subregions_for_validate: Option<&[crate::refdata::VSubregion]> = outcome - .final_simulation() - .assignments - .get(Segment::V) - .and_then(|inst| refdata.v_pool.get(inst.allele_id)) - .map(|allele| allele.subregions.as_slice()); - for er in outcome.events() { - if er.pass_name != crate::address::MUTATE_UNIFORM - && er.pass_name != crate::address::MUTATE_S5F - { - continue; - } - for ev in &er.simulation_events { - let SimulationEvent::BaseChanged { - segment, - germline_pos, - .. - } = ev - else { - continue; - }; - shm_by_segment[*segment as usize] += 1; - if *segment == Segment::V { - let label = v_subregions_for_validate.and_then(|subs| { - germline_pos.and_then(|pos| { - subs.iter() - .find(|s| s.start <= pos && pos < s.end) - .map(|s| s.label) - }) - }); - match label { - Some(crate::refdata::VSubregionLabel::Fwr1) => { - shm_by_subregion[0] += 1 - } - Some(crate::refdata::VSubregionLabel::Cdr1) => { - shm_by_subregion[1] += 1 - } - Some(crate::refdata::VSubregionLabel::Fwr2) => { - shm_by_subregion[2] += 1 - } - Some(crate::refdata::VSubregionLabel::Cdr2) => { - shm_by_subregion[3] += 1 - } - Some(crate::refdata::VSubregionLabel::Fwr3) => { - shm_by_subregion[4] += 1 - } - None => shm_v_unannotated += 1, - } - } - } - } - let expected_v = shm_by_segment[Segment::V as usize]; - let expected_d = shm_by_segment[Segment::D as usize]; - let expected_j = shm_by_segment[Segment::J as usize]; - let expected_np = shm_by_segment[Segment::Np1 as usize] - + shm_by_segment[Segment::Np2 as usize]; - if record.n_v_mutations != expected_v { - issues.push(RecordValidationIssue::NVMutationsMismatch { - reported: record.n_v_mutations, - event_count: expected_v, - }); - } - if record.n_d_mutations != expected_d { - issues.push(RecordValidationIssue::NDMutationsMismatch { - reported: record.n_d_mutations, - event_count: expected_d, - }); - } - if record.n_j_mutations != expected_j { - issues.push(RecordValidationIssue::NJMutationsMismatch { - reported: record.n_j_mutations, - event_count: expected_j, - }); - } - if record.n_np_mutations != expected_np { - issues.push(RecordValidationIssue::NNpMutationsMismatch { - reported: record.n_np_mutations, - event_count: expected_np, - }); - } - // Sum-invariant: the four per-bucket fields must add up to - // ``n_mutations``. Validates the consistency of any consumer- - // supplied record dict; the engine-projected record satisfies - // it by construction. - let sum_of_buckets = record - .n_v_mutations - .saturating_add(record.n_d_mutations) - .saturating_add(record.n_j_mutations) - .saturating_add(record.n_np_mutations); - if record.n_mutations != sum_of_buckets { - issues.push(RecordValidationIssue::MutationCountSumMismatch { - reported_total: record.n_mutations, - sum_of_buckets, - }); - } - // V-subregion partition mismatch checks. Each per-bucket - // mismatch fires `NMutationsMismatch`; the sum - // invariant fires `VSubregionMutationCountSumMismatch`. - if record.n_fwr1_mutations != shm_by_subregion[0] { - issues.push(RecordValidationIssue::NFwr1MutationsMismatch { - reported: record.n_fwr1_mutations, - event_count: shm_by_subregion[0], - }); - } - if record.n_cdr1_mutations != shm_by_subregion[1] { - issues.push(RecordValidationIssue::NCdr1MutationsMismatch { - reported: record.n_cdr1_mutations, - event_count: shm_by_subregion[1], - }); - } - if record.n_fwr2_mutations != shm_by_subregion[2] { - issues.push(RecordValidationIssue::NFwr2MutationsMismatch { - reported: record.n_fwr2_mutations, - event_count: shm_by_subregion[2], - }); - } - if record.n_cdr2_mutations != shm_by_subregion[3] { - issues.push(RecordValidationIssue::NCdr2MutationsMismatch { - reported: record.n_cdr2_mutations, - event_count: shm_by_subregion[3], - }); - } - if record.n_fwr3_mutations != shm_by_subregion[4] { - issues.push(RecordValidationIssue::NFwr3MutationsMismatch { - reported: record.n_fwr3_mutations, - event_count: shm_by_subregion[4], - }); - } - if record.n_v_unannotated_mutations != shm_v_unannotated { - issues.push(RecordValidationIssue::NVUnannotatedMutationsMismatch { - reported: record.n_v_unannotated_mutations, - event_count: shm_v_unannotated, - }); - } - let sum_of_subregion_buckets = record - .n_fwr1_mutations - .saturating_add(record.n_cdr1_mutations) - .saturating_add(record.n_fwr2_mutations) - .saturating_add(record.n_cdr2_mutations) - .saturating_add(record.n_fwr3_mutations) - .saturating_add(record.n_v_unannotated_mutations); - if record.n_v_mutations != sum_of_subregion_buckets { - issues.push( - RecordValidationIssue::VSubregionMutationCountSumMismatch { - reported_v_total: record.n_v_mutations, - sum_of_subregion_buckets, - }, - ); - } - - // Per-end P-nucleotide length counters (Slice — - // P-nucleotide v1). Independent event-ledger recompute: - // walk `PRegionAdded { end, region }` events from the - // matching `p_addition.*` passes and sum `region.len()` - // per end. Catches downstream consumers that tamper with - // the four `p_*_length` fields (record edits, fork- - // patched builders, deserialised dicts). - let mut p_v_3_recompute = 0i64; - let mut p_d_5_recompute = 0i64; - let mut p_d_3_recompute = 0i64; - let mut p_j_5_recompute = 0i64; - for ev_record in outcome.events() { - let is_p_addition = ev_record.pass_name == crate::address::P_ADDITION_V_3 - || ev_record.pass_name == crate::address::P_ADDITION_D_5 - || ev_record.pass_name == crate::address::P_ADDITION_D_3 - || ev_record.pass_name == crate::address::P_ADDITION_J_5; - if !is_p_addition { - continue; - } - for ev in &ev_record.simulation_events { - if let crate::ir::SimulationEvent::PRegionAdded { end, region } = ev { - let len = region.len() as i64; - match end { - crate::address::PEnd::V3 => p_v_3_recompute += len, - crate::address::PEnd::D5 => p_d_5_recompute += len, - crate::address::PEnd::D3 => p_d_3_recompute += len, - crate::address::PEnd::J5 => p_j_5_recompute += len, - } - } - } - } - if record.p_v_3_length != p_v_3_recompute { - issues.push(RecordValidationIssue::PLengthMismatch { - end: crate::address::PEnd::V3, - reported: record.p_v_3_length, - event_count: p_v_3_recompute, - }); - } - if record.p_d_5_length != p_d_5_recompute { - issues.push(RecordValidationIssue::PLengthMismatch { - end: crate::address::PEnd::D5, - reported: record.p_d_5_length, - event_count: p_d_5_recompute, - }); - } - if record.p_d_3_length != p_d_3_recompute { - issues.push(RecordValidationIssue::PLengthMismatch { - end: crate::address::PEnd::D3, - reported: record.p_d_3_length, - event_count: p_d_3_recompute, - }); - } - if record.p_j_5_length != p_j_5_recompute { - issues.push(RecordValidationIssue::PLengthMismatch { - end: crate::address::PEnd::J5, - reported: record.p_j_5_length, - event_count: p_j_5_recompute, - }); - } - - // End-loss lengths from trace. - let el5 = trace_int(ChoiceAddress::CorruptEndLoss(PrimeEnd::Five)); - if record.end_loss_5_length != el5 { - issues.push(RecordValidationIssue::EndLossLengthMismatch { - side: PrimeEnd::Five, - reported: record.end_loss_5_length, - trace_count: el5, - }); - } - let el3 = trace_int(ChoiceAddress::CorruptEndLoss(PrimeEnd::Three)); - if record.end_loss_3_length != el3 { - issues.push(RecordValidationIssue::EndLossLengthMismatch { - side: PrimeEnd::Three, - reported: record.end_loss_3_length, - trace_count: el3, - }); - } - - // D inversion provenance (Slice E). Expected value reads from - // the simulation's final D assignment; defaults to `false` when - // D is absent (VJ chains) — matching the builder's `unwrap_or`. - let expected_inverted = sim - .assignments - .get(Segment::D) - .map(|inst| inst.orientation.is_reverse()) - .unwrap_or(false); - if record.d_inverted != expected_inverted { - issues.push(RecordValidationIssue::DInvertedMismatch { - reported: record.d_inverted, - expected: expected_inverted, - }); - } - - // Receptor revision provenance — IR-sourced (Bug D fix). - // Originally trace-sourced; the descendant trace omits pre-fork - // choices, which made every clonal descendant's projection - // disagree with the parent's actual revision state. The - // assignments slot persists across the parent→descendant - // boundary, so reading from it produces identical behaviour - // for non-clonal and clonal pipelines. - let v_inst = sim.assignments.get(Segment::V); - let expected_applied = v_inst - .map(|inst| inst.receptor_revision_original_id.is_some()) - .unwrap_or(false); - if record.receptor_revision_applied != expected_applied { - issues.push(RecordValidationIssue::ReceptorRevisionAppliedMismatch { - reported: record.receptor_revision_applied, - expected: expected_applied, - }); - } - - let expected_original_v_call = if expected_applied { - original_v_name_from_assignment( - v_inst.expect("expected_applied implies v_inst is Some"), - refdata, - ) - } else { - String::new() - }; - if record.original_v_call != expected_original_v_call { - issues.push(RecordValidationIssue::OriginalVCallMismatch { - reported: record.original_v_call.clone(), - expected: expected_original_v_call, - }); - } -} - -/// IR-sourced counterpart of the (now-removed) trace-based helper. -/// Resolves the pre-revision V allele's refdata name from the -/// persistent provenance slot the receptor-revision pass installs. -/// Mirrors `original_v_call_from_assignment` in -/// `airr_record::builder` so the validator stays self-contained. -fn original_v_name_from_assignment( - v_inst: &crate::assignment::AlleleInstance, - refdata: &RefDataConfig, -) -> String { - let Some(original_id) = v_inst.receptor_revision_original_id else { - return String::new(); - }; - refdata - .get(Segment::V, original_id) - .map(|a| a.name.clone()) - .unwrap_or_default() -} - -// ────────────────────────────────────────────────────────────────── -// C3: Junction truth -// ────────────────────────────────────────────────────────────────── - -fn check_junction( - record: &AirrRecord, - sim: &crate::ir::Simulation, - refdata: &RefDataConfig, - issues: &mut Vec, -) { - // Junction recomputation only meaningful for non-rev-comp records; - // rev-comp flips coordinates and re-translates AA after the - // junction is sliced. Skip when record is reverse-complemented — - // the §5 audit covers that path with dedicated tests. - if record.rev_comp { - return; - } - - // Re-derive the junction window the same way builder.rs does: - // germline_pos scan over V/J regions, not offset arithmetic. - // (The contract-side `crate::junction::compute_junction` uses - // offsets and diverges under indels/end-loss; the AIRR builder - // uses the scan, so the validator must match the builder.) - let v_region = sim - .sequence - .regions - .iter() - .find(|r| r.segment == Segment::V); - let j_region = sim - .sequence - .regions - .iter() - .find(|r| r.segment == Segment::J); - let v_id = sim.assignments.get(Segment::V).map(|i| i.allele_id); - let j_id = sim.assignments.get(Segment::J).map(|i| i.allele_id); - let v_anchor = lookup_allele(refdata, Segment::V, v_id).and_then(|a| a.anchor); - let j_anchor = lookup_allele(refdata, Segment::J, j_id).and_then(|a| a.anchor); - - let (Some(vr), Some(jr), Some(va), Some(ja)) = (v_region, j_region, v_anchor, j_anchor) else { - // Junction undefined; builder leaves fields as defaults. - return; - }; - let (Some(vap), Some(jap)) = ( - anchor_pool_position(sim, vr, va as u32), - anchor_pool_position(sim, jr, ja as u32), - ) else { - return; - }; - let v_anchor_in_pool = vap as i64; - let j_anchor_in_pool = jap as i64; - if j_anchor_in_pool + 3 <= v_anchor_in_pool { - return; - } - let jstart = v_anchor_in_pool; - let jend = j_anchor_in_pool + 3; - let seq_len = record.sequence_length; - let safe_start = jstart.clamp(0, seq_len) as usize; - let safe_end = jend.clamp(0, seq_len) as usize; - let recomputed_content: String = if safe_end > safe_start { - record.sequence[safe_start..safe_end].to_string() - } else { - String::new() - }; - let recomputed_len = recomputed_content.len() as u32; - - let reported_len = record.junction_length; - if reported_len != Some(recomputed_len as i64) { - issues.push(RecordValidationIssue::JunctionLengthMismatch { - reported: reported_len, - recomputed: recomputed_len, - }); - } - - if !record.junction.eq_ignore_ascii_case(&recomputed_content) { - issues.push(RecordValidationIssue::JunctionContentMismatch { - reported: record.junction.clone(), - recomputed: recomputed_content.clone(), - }); - } - - // Frame. - let recomputed_in_frame = recomputed_len % 3 == 0; - if record.vj_in_frame != Some(recomputed_in_frame) { - issues.push(RecordValidationIssue::VjInFrameMismatch { - reported: record.vj_in_frame, - recomputed: recomputed_in_frame, - }); - } - - // Stop codon (only meaningful when in-frame). - let recomputed_stop = recomputed_in_frame && junction_has_stop(&record.junction); - if record.stop_codon != Some(recomputed_stop) { - issues.push(RecordValidationIssue::StopCodonMismatch { - reported: record.stop_codon, - recomputed: recomputed_stop, - }); - } - - // Productive triad: in-frame ∧ no stop ∧ V/J anchor amino acids preserved. - let v_region = sim - .sequence - .regions - .iter() - .find(|r| r.segment == Segment::V); - let j_region = sim - .sequence - .regions - .iter() - .find(|r| r.segment == Segment::J); - let v_anchor_ok = v_region - .map(|r| { - anchor_amino_acid_preserved( - sim, - refdata, - Segment::V, - r, - sim.assignments.get(Segment::V).map(|i| i.allele_id), - record.v_trim_5, - ) - }) - .unwrap_or(true); - let j_anchor_ok = j_region - .map(|r| { - anchor_amino_acid_preserved( - sim, - refdata, - Segment::J, - r, - sim.assignments.get(Segment::J).map(|i| i.allele_id), - record.j_trim_5, - ) - }) - .unwrap_or(true); - - let (recomputed_productive, reason) = if !recomputed_in_frame { - (false, ProductiveDecidedBy::OutOfFrame) - } else if recomputed_stop { - (false, ProductiveDecidedBy::JunctionStopCodon) - } else if !v_anchor_ok { - (false, ProductiveDecidedBy::VAnchorAaChanged) - } else if !j_anchor_ok { - (false, ProductiveDecidedBy::JAnchorAaChanged) - } else { - (true, ProductiveDecidedBy::InFrameAndAnchorsPreserved) - }; - if record.productive != Some(recomputed_productive) { - issues.push(RecordValidationIssue::ProductiveMismatch { - reported: record.productive, - recomputed: recomputed_productive, - reason, - }); - } -} - -// ────────────────────────────────────────────────────────────────── -// C4: Allele-call oracle -// -// Independent reimplementation of the walker's max-match-count -// tie-set selection. We walk the segment's region bytes, count -// matches per allele, and pick the alleles at the max score. -// Then we re-derive the projected CSV order (truth first if in -// tie-set, otherwise ascending allele id). -// ────────────────────────────────────────────────────────────────── - -fn check_allele_oracle( - record: &AirrRecord, - outcome: &Outcome, - refdata: &RefDataConfig, - issues: &mut Vec, -) { - let sim = outcome.final_simulation(); - if record.rev_comp { - // Rev-comp flips the sequence; the allele call was computed - // pre-flip. Skip the oracle here; the rev-comp audit covers - // it with dedicated tests. - return; - } - for (segment, vdj, reported_call) in [ - (Segment::V, VdjSegment::V, &record.v_call), - (Segment::D, VdjSegment::D, &record.d_call), - (Segment::J, VdjSegment::J, &record.j_call), - ] { - let _ = vdj; - oracle_check_segment(sim, refdata, segment, reported_call, issues); - } -} - -fn oracle_check_segment( - sim: &crate::ir::Simulation, - refdata: &RefDataConfig, - segment: Segment, - reported_call: &str, - issues: &mut Vec, -) { - let Some(region) = sim - .sequence - .regions - .iter() - .find(|r| r.segment == segment) - else { - return; - }; - let Some(allele_pool) = allele_pool_for_segment(refdata, segment) else { - return; - }; - if allele_pool.is_empty() { - return; - } - let assignment = sim.assignments.get(segment); - let truth_id = assignment.map(|a| a.allele_id); - let trim_5_cap = assignment.map(|a| a.trim_5 as u32).unwrap_or(0); - let trim_3_cap = assignment.map(|a| a.trim_3 as u32).unwrap_or(0); - // Orientation drives the per-byte comparison rule via the - // shared `matches_observed_with_orientation` primitive: under - // `ReverseComplement` the observed byte is pre-complemented - // before matching the allele's germline byte at the same - // `germline_pos`. Defaults to `Forward` when the segment is - // unassigned. See `scoring::observed_in_germline_orientation` - // for the rationale. - let orientation = assignment - .map(|a| a.orientation) - .unwrap_or(crate::assignment::SegmentOrientation::Forward); - - // Independent rescore via the shared scoring kernel: structural - // region + NP-region extensions under the assigned allele's trim - // caps. Mirrors `live_call::walker::call_from_region` so the - // oracle and the walker agree on the tie-set under arbitrary - // trim. - let scores = score_alleles_with_extensions( - sim, - segment, - allele_pool, - region, - trim_5_cap, - trim_3_cap, - orientation, - ); - let tied_ids = tie_set_ids_at_max_score(&scores); - if tied_ids.is_empty() { - return; // No germline evidence; oracle abstains. - } - let tied_indices: Vec = tied_ids.iter().map(|id| id.as_usize()).collect(); - - // Expected CSV order: truth allele first when in tie-set, - // otherwise ascending by allele id (already sorted by the - // kernel's iteration order). - let truth_idx = truth_id - .map(|id| id.as_usize()) - .filter(|&i| i < allele_pool.len()); - - let mut expected_order = tied_indices.clone(); - if let Some(t) = truth_idx { - if let Some(pos) = expected_order.iter().position(|&i| i == t) { - let truth_first = expected_order.remove(pos); - expected_order.insert(0, truth_first); - } - } - - let expected_names: Vec = expected_order - .iter() - .map(|&i| allele_pool[i].name.clone()) - .collect(); - let reported_names: Vec = reported_call - .split(',') - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .collect(); - - // Tie-set equality (order-insensitive). - let mut reported_sorted = reported_names.clone(); - reported_sorted.sort(); - let mut expected_sorted = expected_names.clone(); - expected_sorted.sort(); - if reported_sorted != expected_sorted { - issues.push(RecordValidationIssue::AlleleCallTieSetMismatch { - segment, - reported: reported_names.clone(), - recomputed: expected_names.clone(), - }); - return; // Order check is meaningless if the sets differ. - } - - // Order check: first element must match expected_order's first. - if let (Some(reported_first), Some(expected_first)) = - (reported_names.first(), expected_names.first()) - { - if reported_first != expected_first { - let reason = if truth_idx.is_some() - && expected_order - .first() - .map(|&i| i == truth_idx.unwrap()) - .unwrap_or(false) - { - AlleleOrderReason::TruthFirstIfInTieSet - } else { - AlleleOrderReason::AscendingAlleleIdOtherwise - }; - issues.push(RecordValidationIssue::AlleleCallOrderMismatch { - segment, - reported_first: reported_first.clone(), - expected_first: expected_first.clone(), - reason, - }); - } - } -} - -// Match semantics live in crate::live_call::scoring; this module -// uses them via score_alleles_in_region / tie_set_ids_at_max_score. - -// ────────────────────────────────────────────────────────────────── -// C5: Region / live-call structural invariants -// -// Per the §5 architectural audit: -// - Each Segment (V/D/J) appears at most once in -// sim.sequence.regions. Live-call and AIRR projection pick the -// "latest" and "first" respectively; identical-by-construction -// today, drift-able if invariant breaks. -// - SegmentLiveCall.hypotheses has length <= 1 in production runs. -// Projection silently uses hypotheses[0]; multi-hypothesis would -// be lossy. -// ────────────────────────────────────────────────────────────────── - -fn check_region_and_hypothesis_invariants( - sim: &crate::ir::Simulation, - issues: &mut Vec, -) { - for seg in [Segment::V, Segment::D, Segment::J] { - let count = sim - .sequence - .regions - .iter() - .filter(|r| r.segment == seg) - .count(); - if count > 1 { - issues.push(RecordValidationIssue::MultipleRegionsForSegment { - segment: seg, - count, - }); - } - } - - for seg in [Segment::V, Segment::D, Segment::J] { - if let Some(call) = sim.segment_calls.get(seg) { - if call.hypotheses.len() > 1 { - issues.push(RecordValidationIssue::MultipleHypothesesInLiveCall { - segment: seg, - count: call.hypotheses.len(), - }); - } - } - } -} - -// ────────────────────────────────────────────────────────────────── -// Tests -// ────────────────────────────────────────────────────────────────── - -#[cfg(test)] -mod tests { - //! Unit tests live alongside the integration tests in - //! `engine_rs/src/airr_record/tests/validate.rs`. Keeping them - //! out of this file makes the module easier to skim. -} diff --git a/engine_rs/src/airr_record/validate/allele_oracle.rs b/engine_rs/src/airr_record/validate/allele_oracle.rs new file mode 100644 index 0000000..596eb6d --- /dev/null +++ b/engine_rs/src/airr_record/validate/allele_oracle.rs @@ -0,0 +1,161 @@ +//! C4: Allele-call oracle. +//! +//! Independent reimplementation of the walker's max-match-count +//! tie-set selection. We walk the segment's region bytes, count +//! matches per allele, and pick the alleles at the max score. +//! Then we re-derive the projected CSV order (truth first if in +//! tie-set, otherwise ascending allele id). + +use super::*; + +use crate::address::VdjSegment; +use crate::ir::Segment; +use crate::live_call::scoring::{ + allele_pool_for_segment, score_alleles_with_extensions, tie_set_ids_at_max_score, +}; + +pub(super) fn check_allele_oracle( + record: &AirrRecord, + outcome: &Outcome, + refdata: &RefDataConfig, + issues: &mut Vec, +) { + let sim = outcome.final_simulation(); + if record.rev_comp { + // Rev-comp flips the sequence; the allele call was computed + // pre-flip. Skip the oracle here; the rev-comp audit covers + // it with dedicated tests. + return; + } + for (segment, vdj, reported_call) in [ + (Segment::V, VdjSegment::V, &record.v_call), + (Segment::D, VdjSegment::D, &record.d_call), + (Segment::J, VdjSegment::J, &record.j_call), + ] { + let _ = vdj; + oracle_check_segment(sim, refdata, segment, reported_call, issues); + } +} + +fn oracle_check_segment( + sim: &crate::ir::Simulation, + refdata: &RefDataConfig, + segment: Segment, + reported_call: &str, + issues: &mut Vec, +) { + let Some(region) = sim + .sequence + .regions + .iter() + .find(|r| r.segment == segment) + else { + return; + }; + let Some(allele_pool) = allele_pool_for_segment(refdata, segment) else { + return; + }; + if allele_pool.is_empty() { + return; + } + let assignment = sim.assignments.get(segment); + let truth_id = assignment.map(|a| a.allele_id); + let trim_5_cap = assignment.map(|a| a.trim_5 as u32).unwrap_or(0); + let trim_3_cap = assignment.map(|a| a.trim_3 as u32).unwrap_or(0); + // Orientation drives the per-byte comparison rule via the + // shared `matches_observed_with_orientation` primitive: under + // `ReverseComplement` the observed byte is pre-complemented + // before matching the allele's germline byte at the same + // `germline_pos`. Defaults to `Forward` when the segment is + // unassigned. See `scoring::observed_in_germline_orientation` + // for the rationale. + let orientation = assignment + .map(|a| a.orientation) + .unwrap_or(crate::assignment::SegmentOrientation::Forward); + + // Independent rescore via the shared scoring kernel: structural + // region + NP-region extensions under the assigned allele's trim + // caps. Mirrors `live_call::walker::call_from_region` so the + // oracle and the walker agree on the tie-set under arbitrary + // trim. + let scores = score_alleles_with_extensions( + sim, + segment, + allele_pool, + region, + trim_5_cap, + trim_3_cap, + orientation, + ); + let tied_ids = tie_set_ids_at_max_score(&scores); + if tied_ids.is_empty() { + return; // No germline evidence; oracle abstains. + } + let tied_indices: Vec = tied_ids.iter().map(|id| id.as_usize()).collect(); + + // Expected CSV order: truth allele first when in tie-set, + // otherwise ascending by allele id (already sorted by the + // kernel's iteration order). + let truth_idx = truth_id + .map(|id| id.as_usize()) + .filter(|&i| i < allele_pool.len()); + + let mut expected_order = tied_indices.clone(); + if let Some(t) = truth_idx { + if let Some(pos) = expected_order.iter().position(|&i| i == t) { + let truth_first = expected_order.remove(pos); + expected_order.insert(0, truth_first); + } + } + + let expected_names: Vec = expected_order + .iter() + .map(|&i| allele_pool[i].name.clone()) + .collect(); + let reported_names: Vec = reported_call + .split(',') + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .collect(); + + // Tie-set equality (order-insensitive). + let mut reported_sorted = reported_names.clone(); + reported_sorted.sort(); + let mut expected_sorted = expected_names.clone(); + expected_sorted.sort(); + if reported_sorted != expected_sorted { + issues.push(RecordValidationIssue::AlleleCallTieSetMismatch { + segment, + reported: reported_names.clone(), + recomputed: expected_names.clone(), + }); + return; // Order check is meaningless if the sets differ. + } + + // Order check: first element must match expected_order's first. + if let (Some(reported_first), Some(expected_first)) = + (reported_names.first(), expected_names.first()) + { + if reported_first != expected_first { + let reason = if truth_idx.is_some() + && expected_order + .first() + .map(|&i| i == truth_idx.unwrap()) + .unwrap_or(false) + { + AlleleOrderReason::TruthFirstIfInTieSet + } else { + AlleleOrderReason::AscendingAlleleIdOtherwise + }; + issues.push(RecordValidationIssue::AlleleCallOrderMismatch { + segment, + reported_first: reported_first.clone(), + expected_first: expected_first.clone(), + reason, + }); + } + } +} + +// Match semantics live in crate::live_call::scoring; this module +// uses them via score_alleles_in_region / tie_set_ids_at_max_score. diff --git a/engine_rs/src/airr_record/validate/counters.rs b/engine_rs/src/airr_record/validate/counters.rs new file mode 100644 index 0000000..d13190e --- /dev/null +++ b/engine_rs/src/airr_record/validate/counters.rs @@ -0,0 +1,403 @@ +//! C2: Counter provenance. + +use super::*; + +use crate::address::{ChoiceAddress, PrimeEnd}; +use crate::ir::{Segment, SimulationEvent}; +use crate::trace::ChoiceValue; + +pub(super) fn check_counters( + record: &AirrRecord, + outcome: &Outcome, + refdata: &RefDataConfig, + issues: &mut Vec, +) { + let sim = outcome.final_simulation(); + + // n_mutations comes from sim.mutation_count (set by S5F / Uniform + // at seal). Direct equality. + if record.n_mutations != sim.mutation_count as i64 { + issues.push(RecordValidationIssue::NMutationsMismatch { + reported: record.n_mutations, + sim_count: sim.mutation_count as i64, + }); + } + + // n_pcr_errors / n_quality_errors from trace. + let trace_int = |addr: ChoiceAddress| -> i64 { + outcome + .trace + .find_choice(addr) + .and_then(|rec| match rec.value { + ChoiceValue::Int(n) => Some(n), + _ => None, + }) + .unwrap_or(0) + }; + + let pcr_attempts = trace_int(ChoiceAddress::CorruptPcrCount); + if record.n_pcr_errors != pcr_attempts { + issues.push(RecordValidationIssue::NPcrErrorsMismatch { + reported: record.n_pcr_errors, + trace_count: pcr_attempts, + }); + } + + let quality_attempts = trace_int(ChoiceAddress::CorruptQualityCount); + if record.n_quality_errors != quality_attempts { + issues.push(RecordValidationIssue::NQualityErrorsMismatch { + reported: record.n_quality_errors, + trace_count: quality_attempts, + }); + } + + // n_indels and per-segment indel counters from event ledger of the + // corrupt.indel pass only (per audit §6.1 / §6.2). + let mut total = 0i64; + let mut by_segment = [0i64; Segment::COUNT]; + for er in outcome.events() { + if er.pass_name != crate::address::CORRUPT_INDEL { + continue; + } + for ev in &er.simulation_events { + let segment = match ev { + SimulationEvent::IndelInserted { segment, .. } => *segment, + SimulationEvent::IndelDeleted { segment, .. } => *segment, + _ => continue, + }; + total += 1; + by_segment[segment as usize] += 1; + } + } + if record.n_indels != total { + issues.push(RecordValidationIssue::NIndelsMismatch { + reported: record.n_indels, + event_count: total, + }); + } + for (seg, reported) in [ + (Segment::V, record.n_v_indels), + (Segment::D, record.n_d_indels), + (Segment::J, record.n_j_indels), + ] { + let expected = by_segment[seg as usize]; + if reported != expected { + issues.push(RecordValidationIssue::NSegmentIndelsMismatch { + segment: seg, + reported, + event_count: expected, + }); + } + } + + // Per-segment SHM counters from the event ledger of the + // mutate.{uniform,s5f} passes only. Mirrors the indel walk + // above but counts `BaseChanged` events instead of indel + // events, and rolls NP1+NP2 into the single NP bucket. The + // four per-bucket checks fire `N*MutationsMismatch` on + // disagreement; the sum-invariant cross-check fires + // `MutationCountSumMismatch`. + // + // Same walk also produces the V-subregion partition (Slice — + // V-Subregion Mutation Counters): each V `BaseChanged.germline_pos` + // is matched against the assigned V allele's `subregions` table + // from scratch, independent of the projection's bucketing. Six + // per-bucket mismatches plus a cross-field sum invariant + // (`VSubregionMutationCountSumMismatch`) fire on disagreement. + let mut shm_by_segment = [0i64; Segment::COUNT]; + let mut shm_by_subregion = [0i64; 5]; + let mut shm_v_unannotated = 0i64; + let v_subregions_for_validate: Option<&[crate::refdata::VSubregion]> = outcome + .final_simulation() + .assignments + .get(Segment::V) + .and_then(|inst| refdata.v_pool.get(inst.allele_id)) + .map(|allele| allele.subregions.as_slice()); + for er in outcome.events() { + if er.pass_name != crate::address::MUTATE_UNIFORM + && er.pass_name != crate::address::MUTATE_S5F + { + continue; + } + for ev in &er.simulation_events { + let SimulationEvent::BaseChanged { + segment, + germline_pos, + .. + } = ev + else { + continue; + }; + shm_by_segment[*segment as usize] += 1; + if *segment == Segment::V { + let label = v_subregions_for_validate.and_then(|subs| { + germline_pos.and_then(|pos| { + subs.iter() + .find(|s| s.start <= pos && pos < s.end) + .map(|s| s.label) + }) + }); + match label { + Some(crate::refdata::VSubregionLabel::Fwr1) => { + shm_by_subregion[0] += 1 + } + Some(crate::refdata::VSubregionLabel::Cdr1) => { + shm_by_subregion[1] += 1 + } + Some(crate::refdata::VSubregionLabel::Fwr2) => { + shm_by_subregion[2] += 1 + } + Some(crate::refdata::VSubregionLabel::Cdr2) => { + shm_by_subregion[3] += 1 + } + Some(crate::refdata::VSubregionLabel::Fwr3) => { + shm_by_subregion[4] += 1 + } + None => shm_v_unannotated += 1, + } + } + } + } + let expected_v = shm_by_segment[Segment::V as usize]; + let expected_d = shm_by_segment[Segment::D as usize]; + let expected_j = shm_by_segment[Segment::J as usize]; + let expected_np = shm_by_segment[Segment::Np1 as usize] + + shm_by_segment[Segment::Np2 as usize]; + if record.n_v_mutations != expected_v { + issues.push(RecordValidationIssue::NVMutationsMismatch { + reported: record.n_v_mutations, + event_count: expected_v, + }); + } + if record.n_d_mutations != expected_d { + issues.push(RecordValidationIssue::NDMutationsMismatch { + reported: record.n_d_mutations, + event_count: expected_d, + }); + } + if record.n_j_mutations != expected_j { + issues.push(RecordValidationIssue::NJMutationsMismatch { + reported: record.n_j_mutations, + event_count: expected_j, + }); + } + if record.n_np_mutations != expected_np { + issues.push(RecordValidationIssue::NNpMutationsMismatch { + reported: record.n_np_mutations, + event_count: expected_np, + }); + } + // Sum-invariant: the four per-bucket fields must add up to + // ``n_mutations``. Validates the consistency of any consumer- + // supplied record dict; the engine-projected record satisfies + // it by construction. + let sum_of_buckets = record + .n_v_mutations + .saturating_add(record.n_d_mutations) + .saturating_add(record.n_j_mutations) + .saturating_add(record.n_np_mutations); + if record.n_mutations != sum_of_buckets { + issues.push(RecordValidationIssue::MutationCountSumMismatch { + reported_total: record.n_mutations, + sum_of_buckets, + }); + } + // V-subregion partition mismatch checks. Each per-bucket + // mismatch fires `NMutationsMismatch`; the sum + // invariant fires `VSubregionMutationCountSumMismatch`. + if record.n_fwr1_mutations != shm_by_subregion[0] { + issues.push(RecordValidationIssue::NFwr1MutationsMismatch { + reported: record.n_fwr1_mutations, + event_count: shm_by_subregion[0], + }); + } + if record.n_cdr1_mutations != shm_by_subregion[1] { + issues.push(RecordValidationIssue::NCdr1MutationsMismatch { + reported: record.n_cdr1_mutations, + event_count: shm_by_subregion[1], + }); + } + if record.n_fwr2_mutations != shm_by_subregion[2] { + issues.push(RecordValidationIssue::NFwr2MutationsMismatch { + reported: record.n_fwr2_mutations, + event_count: shm_by_subregion[2], + }); + } + if record.n_cdr2_mutations != shm_by_subregion[3] { + issues.push(RecordValidationIssue::NCdr2MutationsMismatch { + reported: record.n_cdr2_mutations, + event_count: shm_by_subregion[3], + }); + } + if record.n_fwr3_mutations != shm_by_subregion[4] { + issues.push(RecordValidationIssue::NFwr3MutationsMismatch { + reported: record.n_fwr3_mutations, + event_count: shm_by_subregion[4], + }); + } + if record.n_v_unannotated_mutations != shm_v_unannotated { + issues.push(RecordValidationIssue::NVUnannotatedMutationsMismatch { + reported: record.n_v_unannotated_mutations, + event_count: shm_v_unannotated, + }); + } + let sum_of_subregion_buckets = record + .n_fwr1_mutations + .saturating_add(record.n_cdr1_mutations) + .saturating_add(record.n_fwr2_mutations) + .saturating_add(record.n_cdr2_mutations) + .saturating_add(record.n_fwr3_mutations) + .saturating_add(record.n_v_unannotated_mutations); + if record.n_v_mutations != sum_of_subregion_buckets { + issues.push( + RecordValidationIssue::VSubregionMutationCountSumMismatch { + reported_v_total: record.n_v_mutations, + sum_of_subregion_buckets, + }, + ); + } + + // Per-end P-nucleotide length counters (Slice — + // P-nucleotide v1). Independent event-ledger recompute: + // walk `PRegionAdded { end, region }` events from the + // matching `p_addition.*` passes and sum `region.len()` + // per end. Catches downstream consumers that tamper with + // the four `p_*_length` fields (record edits, fork- + // patched builders, deserialised dicts). + let mut p_v_3_recompute = 0i64; + let mut p_d_5_recompute = 0i64; + let mut p_d_3_recompute = 0i64; + let mut p_j_5_recompute = 0i64; + for ev_record in outcome.events() { + let is_p_addition = ev_record.pass_name == crate::address::P_ADDITION_V_3 + || ev_record.pass_name == crate::address::P_ADDITION_D_5 + || ev_record.pass_name == crate::address::P_ADDITION_D_3 + || ev_record.pass_name == crate::address::P_ADDITION_J_5; + if !is_p_addition { + continue; + } + for ev in &ev_record.simulation_events { + if let crate::ir::SimulationEvent::PRegionAdded { end, region } = ev { + let len = region.len() as i64; + match end { + crate::address::PEnd::V3 => p_v_3_recompute += len, + crate::address::PEnd::D5 => p_d_5_recompute += len, + crate::address::PEnd::D3 => p_d_3_recompute += len, + crate::address::PEnd::J5 => p_j_5_recompute += len, + } + } + } + } + if record.p_v_3_length != p_v_3_recompute { + issues.push(RecordValidationIssue::PLengthMismatch { + end: crate::address::PEnd::V3, + reported: record.p_v_3_length, + event_count: p_v_3_recompute, + }); + } + if record.p_d_5_length != p_d_5_recompute { + issues.push(RecordValidationIssue::PLengthMismatch { + end: crate::address::PEnd::D5, + reported: record.p_d_5_length, + event_count: p_d_5_recompute, + }); + } + if record.p_d_3_length != p_d_3_recompute { + issues.push(RecordValidationIssue::PLengthMismatch { + end: crate::address::PEnd::D3, + reported: record.p_d_3_length, + event_count: p_d_3_recompute, + }); + } + if record.p_j_5_length != p_j_5_recompute { + issues.push(RecordValidationIssue::PLengthMismatch { + end: crate::address::PEnd::J5, + reported: record.p_j_5_length, + event_count: p_j_5_recompute, + }); + } + + // End-loss lengths from trace. + let el5 = trace_int(ChoiceAddress::CorruptEndLoss(PrimeEnd::Five)); + if record.end_loss_5_length != el5 { + issues.push(RecordValidationIssue::EndLossLengthMismatch { + side: PrimeEnd::Five, + reported: record.end_loss_5_length, + trace_count: el5, + }); + } + let el3 = trace_int(ChoiceAddress::CorruptEndLoss(PrimeEnd::Three)); + if record.end_loss_3_length != el3 { + issues.push(RecordValidationIssue::EndLossLengthMismatch { + side: PrimeEnd::Three, + reported: record.end_loss_3_length, + trace_count: el3, + }); + } + + // D inversion provenance (Slice E). Expected value reads from + // the simulation's final D assignment; defaults to `false` when + // D is absent (VJ chains) — matching the builder's `unwrap_or`. + let expected_inverted = sim + .assignments + .get(Segment::D) + .map(|inst| inst.orientation.is_reverse()) + .unwrap_or(false); + if record.d_inverted != expected_inverted { + issues.push(RecordValidationIssue::DInvertedMismatch { + reported: record.d_inverted, + expected: expected_inverted, + }); + } + + // Receptor revision provenance — IR-sourced (Bug D fix). + // Originally trace-sourced; the descendant trace omits pre-fork + // choices, which made every clonal descendant's projection + // disagree with the parent's actual revision state. The + // assignments slot persists across the parent→descendant + // boundary, so reading from it produces identical behaviour + // for non-clonal and clonal pipelines. + let v_inst = sim.assignments.get(Segment::V); + let expected_applied = v_inst + .map(|inst| inst.receptor_revision_original_id.is_some()) + .unwrap_or(false); + if record.receptor_revision_applied != expected_applied { + issues.push(RecordValidationIssue::ReceptorRevisionAppliedMismatch { + reported: record.receptor_revision_applied, + expected: expected_applied, + }); + } + + let expected_original_v_call = if expected_applied { + original_v_name_from_assignment( + v_inst.expect("expected_applied implies v_inst is Some"), + refdata, + ) + } else { + String::new() + }; + if record.original_v_call != expected_original_v_call { + issues.push(RecordValidationIssue::OriginalVCallMismatch { + reported: record.original_v_call.clone(), + expected: expected_original_v_call, + }); + } +} + +/// IR-sourced counterpart of the (now-removed) trace-based helper. +/// Resolves the pre-revision V allele's refdata name from the +/// persistent provenance slot the receptor-revision pass installs. +/// Mirrors `original_v_call_from_assignment` in +/// `airr_record::builder` so the validator stays self-contained. +fn original_v_name_from_assignment( + v_inst: &crate::assignment::AlleleInstance, + refdata: &RefDataConfig, +) -> String { + let Some(original_id) = v_inst.receptor_revision_original_id else { + return String::new(); + }; + refdata + .get(Segment::V, original_id) + .map(|a| a.name.clone()) + .unwrap_or_default() +} diff --git a/engine_rs/src/airr_record/validate/junction.rs b/engine_rs/src/airr_record/validate/junction.rs new file mode 100644 index 0000000..9f8385e --- /dev/null +++ b/engine_rs/src/airr_record/validate/junction.rs @@ -0,0 +1,157 @@ +//! C3: Junction truth. + +use super::*; + +use super::super::junction::{anchor_amino_acid_preserved, anchor_pool_position, junction_has_stop}; +use super::super::projection::lookup_allele; +use crate::ir::Segment; + +pub(super) fn check_junction( + record: &AirrRecord, + sim: &crate::ir::Simulation, + refdata: &RefDataConfig, + issues: &mut Vec, +) { + // Junction recomputation only meaningful for non-rev-comp records; + // rev-comp flips coordinates and re-translates AA after the + // junction is sliced. Skip when record is reverse-complemented — + // the §5 audit covers that path with dedicated tests. + if record.rev_comp { + return; + } + + // Re-derive the junction window the same way builder.rs does: + // germline_pos scan over V/J regions, not offset arithmetic. + // (The contract-side `crate::junction::compute_junction` uses + // offsets and diverges under indels/end-loss; the AIRR builder + // uses the scan, so the validator must match the builder.) + let v_region = sim + .sequence + .regions + .iter() + .find(|r| r.segment == Segment::V); + let j_region = sim + .sequence + .regions + .iter() + .find(|r| r.segment == Segment::J); + let v_id = sim.assignments.get(Segment::V).map(|i| i.allele_id); + let j_id = sim.assignments.get(Segment::J).map(|i| i.allele_id); + let v_anchor = lookup_allele(refdata, Segment::V, v_id).and_then(|a| a.anchor); + let j_anchor = lookup_allele(refdata, Segment::J, j_id).and_then(|a| a.anchor); + + let (Some(vr), Some(jr), Some(va), Some(ja)) = (v_region, j_region, v_anchor, j_anchor) else { + // Junction undefined; builder leaves fields as defaults. + return; + }; + let (Some(vap), Some(jap)) = ( + anchor_pool_position(sim, vr, va as u32), + anchor_pool_position(sim, jr, ja as u32), + ) else { + return; + }; + let v_anchor_in_pool = vap as i64; + let j_anchor_in_pool = jap as i64; + if j_anchor_in_pool + 3 <= v_anchor_in_pool { + return; + } + let jstart = v_anchor_in_pool; + let jend = j_anchor_in_pool + 3; + let seq_len = record.sequence_length; + let safe_start = jstart.clamp(0, seq_len) as usize; + let safe_end = jend.clamp(0, seq_len) as usize; + let recomputed_content: String = if safe_end > safe_start { + record.sequence[safe_start..safe_end].to_string() + } else { + String::new() + }; + let recomputed_len = recomputed_content.len() as u32; + + let reported_len = record.junction_length; + if reported_len != Some(recomputed_len as i64) { + issues.push(RecordValidationIssue::JunctionLengthMismatch { + reported: reported_len, + recomputed: recomputed_len, + }); + } + + if !record.junction.eq_ignore_ascii_case(&recomputed_content) { + issues.push(RecordValidationIssue::JunctionContentMismatch { + reported: record.junction.clone(), + recomputed: recomputed_content.clone(), + }); + } + + // Frame. + let recomputed_in_frame = recomputed_len % 3 == 0; + if record.vj_in_frame != Some(recomputed_in_frame) { + issues.push(RecordValidationIssue::VjInFrameMismatch { + reported: record.vj_in_frame, + recomputed: recomputed_in_frame, + }); + } + + // Stop codon (only meaningful when in-frame). + let recomputed_stop = recomputed_in_frame && junction_has_stop(&record.junction); + if record.stop_codon != Some(recomputed_stop) { + issues.push(RecordValidationIssue::StopCodonMismatch { + reported: record.stop_codon, + recomputed: recomputed_stop, + }); + } + + // Productive triad: in-frame ∧ no stop ∧ V/J anchor amino acids preserved. + let v_region = sim + .sequence + .regions + .iter() + .find(|r| r.segment == Segment::V); + let j_region = sim + .sequence + .regions + .iter() + .find(|r| r.segment == Segment::J); + let v_anchor_ok = v_region + .map(|r| { + anchor_amino_acid_preserved( + sim, + refdata, + Segment::V, + r, + sim.assignments.get(Segment::V).map(|i| i.allele_id), + record.v_trim_5, + ) + }) + .unwrap_or(true); + let j_anchor_ok = j_region + .map(|r| { + anchor_amino_acid_preserved( + sim, + refdata, + Segment::J, + r, + sim.assignments.get(Segment::J).map(|i| i.allele_id), + record.j_trim_5, + ) + }) + .unwrap_or(true); + + let (recomputed_productive, reason) = if !recomputed_in_frame { + (false, ProductiveDecidedBy::OutOfFrame) + } else if recomputed_stop { + (false, ProductiveDecidedBy::JunctionStopCodon) + } else if !v_anchor_ok { + (false, ProductiveDecidedBy::VAnchorAaChanged) + } else if !j_anchor_ok { + (false, ProductiveDecidedBy::JAnchorAaChanged) + } else { + (true, ProductiveDecidedBy::InFrameAndAnchorsPreserved) + }; + if record.productive != Some(recomputed_productive) { + issues.push(RecordValidationIssue::ProductiveMismatch { + reported: record.productive, + recomputed: recomputed_productive, + reason, + }); + } +} diff --git a/engine_rs/src/airr_record/validate/paired_end.rs b/engine_rs/src/airr_record/validate/paired_end.rs new file mode 100644 index 0000000..7a7e851 --- /dev/null +++ b/engine_rs/src/airr_record/validate/paired_end.rs @@ -0,0 +1,175 @@ +//! C6: Paired-end / read-layout invariants. +//! +//! Slice A: enforces the no-layout default invariant (`read_layout +//! == ""` ⇒ every paired-end field is at its default). +//! +//! Slice B: extends the dispatch with per-layout geometry checks. +//! When `read_layout == "paired_end"`, the validator re-derives +//! R1/R2 windows from the rules in `docs/paired_end_design.md` §8 +//! and surfaces any divergence as one of the four reserved variants +//! (`ReadWindowOutOfBounds`, `ReadSequenceMismatch`, +//! `ReadInsertSizeMismatch`). Unknown non-empty layouts surface as +//! `ReadLayoutMismatch`. `"single_end"` is documented as reserved +//! (§2.2) and treated as a no-op for now. + +use super::*; + +pub(super) fn check_paired_end_defaults( + record: &AirrRecord, + issues: &mut Vec, +) { + match record.read_layout.as_str() { + // Slice A: no-layout default invariant. + "" => check_paired_end_default_values(record, issues), + // Slice B: geometry against the projection rules. + "paired_end" => check_paired_end_geometry(record, issues), + // Reserved (per §2.2). A future slice may attach a + // single-read geometry check; today we don't validate + // anything beyond accepting the layout string. + "single_end" => {} + // Anything else is an unsupported layout value. Surface + // the structured mismatch so a typo (`"pair_end"`) or a + // refactor that introduced a new layout without wiring + // it through the validator fails closed. + _ => issues.push(RecordValidationIssue::ReadLayoutMismatch { + reported: record.read_layout.clone(), + expected: r#"one of: "", "paired_end", "single_end""#.to_string(), + }), + } +} + +fn check_paired_end_default_values( + record: &AirrRecord, + issues: &mut Vec, +) { + if !record.r1_sequence.is_empty() { + issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { + field: PairedEndField::R1Sequence, + }); + } + if !record.r2_sequence.is_empty() { + issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { + field: PairedEndField::R2Sequence, + }); + } + if record.r1_start.is_some() { + issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { + field: PairedEndField::R1Start, + }); + } + if record.r1_end.is_some() { + issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { + field: PairedEndField::R1End, + }); + } + if record.r2_start.is_some() { + issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { + field: PairedEndField::R2Start, + }); + } + if record.r2_end.is_some() { + issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { + field: PairedEndField::R2End, + }); + } + if record.insert_size != 0 { + issues.push(RecordValidationIssue::PairedEndFieldWithoutLayout { + field: PairedEndField::InsertSize, + }); + } +} + +fn check_paired_end_geometry( + record: &AirrRecord, + issues: &mut Vec, +) { + let seq_len = record.sequence_length; + let seq = &record.sequence; + + // R1 window: bounds + byte equality. + match resolve_window(record.r1_start, record.r1_end, seq_len) { + WindowResolution::Valid { start, end } => { + let expected = seq[start as usize..end as usize].to_string(); + if record.r1_sequence != expected { + issues.push(RecordValidationIssue::ReadSequenceMismatch { + side: PairedEndRead::R1, + reported: record.r1_sequence.clone(), + expected, + }); + } + } + WindowResolution::OutOfBounds { start, end } => { + issues.push(RecordValidationIssue::ReadWindowOutOfBounds { + side: PairedEndRead::R1, + start, + end, + sequence_length: seq_len, + }); + } + } + + // R2 window: bounds + reverse-complement byte equality + insert + // size consistency. The audit pins + // `insert_size == r2_end` (§8) — the only insert-size-mismatch + // surface today. + match resolve_window(record.r2_start, record.r2_end, seq_len) { + WindowResolution::Valid { start, end } => { + let r2_inner = &seq[start as usize..end as usize]; + let expected = super::super::sequence::reverse_complement(r2_inner); + if record.r2_sequence != expected { + issues.push(RecordValidationIssue::ReadSequenceMismatch { + side: PairedEndRead::R2, + reported: record.r2_sequence.clone(), + expected, + }); + } + if record.insert_size != end { + issues.push(RecordValidationIssue::ReadInsertSizeMismatch { + reported: record.insert_size, + expected: end, + }); + } + } + WindowResolution::OutOfBounds { start, end } => { + issues.push(RecordValidationIssue::ReadWindowOutOfBounds { + side: PairedEndRead::R2, + start, + end, + sequence_length: seq_len, + }); + // With R2 bounds unresolved we can't recompute the + // expected insert size; skip the insert-size check + // rather than fabricate a sentinel. The window + // out-of-bounds issue is the actionable signal. + } + } +} + +/// Result of resolving a paired-end window `(start, end)` against +/// the projected `sequence_length`. `OutOfBounds` collapses missing +/// coords (`None`) and out-of-range coords (negative, swapped, +/// past the end) into one variant whose `start`/`end` fields use +/// the sentinel `-1` for missing values — surfaces cleanly in the +/// structured-issue payload without inventing a new variant. +enum WindowResolution { + Valid { start: i64, end: i64 }, + OutOfBounds { start: i64, end: i64 }, +} + +fn resolve_window( + start: Option, + end: Option, + sequence_length: i64, +) -> WindowResolution { + let start_val = start.unwrap_or(-1); + let end_val = end.unwrap_or(-1); + match (start, end) { + (Some(s), Some(e)) if s >= 0 && e >= s && e <= sequence_length => { + WindowResolution::Valid { start: s, end: e } + } + _ => WindowResolution::OutOfBounds { + start: start_val, + end: end_val, + }, + } +} diff --git a/engine_rs/src/airr_record/validate/region_invariants.rs b/engine_rs/src/airr_record/validate/region_invariants.rs new file mode 100644 index 0000000..5342108 --- /dev/null +++ b/engine_rs/src/airr_record/validate/region_invariants.rs @@ -0,0 +1,45 @@ +//! C5: Region / live-call structural invariants. +//! +//! Per the §5 architectural audit: +//! - Each Segment (V/D/J) appears at most once in +//! sim.sequence.regions. Live-call and AIRR projection pick the +//! "latest" and "first" respectively; identical-by-construction +//! today, drift-able if invariant breaks. +//! - SegmentLiveCall.hypotheses has length <= 1 in production runs. +//! Projection silently uses hypotheses[0]; multi-hypothesis would +//! be lossy. + +use super::*; + +use crate::ir::Segment; + +pub(super) fn check_region_and_hypothesis_invariants( + sim: &crate::ir::Simulation, + issues: &mut Vec, +) { + for seg in [Segment::V, Segment::D, Segment::J] { + let count = sim + .sequence + .regions + .iter() + .filter(|r| r.segment == seg) + .count(); + if count > 1 { + issues.push(RecordValidationIssue::MultipleRegionsForSegment { + segment: seg, + count, + }); + } + } + + for seg in [Segment::V, Segment::D, Segment::J] { + if let Some(call) = sim.segment_calls.get(seg) { + if call.hypotheses.len() > 1 { + issues.push(RecordValidationIssue::MultipleHypothesesInLiveCall { + segment: seg, + count: call.hypotheses.len(), + }); + } + } + } +} diff --git a/engine_rs/src/airr_record/validate/structural.rs b/engine_rs/src/airr_record/validate/structural.rs new file mode 100644 index 0000000..ec73d20 --- /dev/null +++ b/engine_rs/src/airr_record/validate/structural.rs @@ -0,0 +1,177 @@ +//! C1: Structural record invariants. + +use super::*; + +use super::super::sequence::pool_bases; +use crate::ir::Segment; + +pub(super) fn check_structural( + record: &AirrRecord, + sim: &crate::ir::Simulation, + issues: &mut Vec, +) { + // sequence_length matches actual sequence string bytes. + let actual_bytes = record.sequence.len(); + if record.sequence_length as usize != actual_bytes { + issues.push(RecordValidationIssue::SequenceLengthMismatch { + reported: record.sequence_length, + actual_bytes, + }); + } + + // sequence content matches the pool's bases (case-folded). + // Compare case-insensitively because sequencing errors lowercase + // their substitutions, and the AIRR record may carry rev-comp. + let pool_seq = String::from_utf8(pool_bases(sim)).unwrap_or_default(); + if !record.rev_comp && !record.sequence.eq_ignore_ascii_case(&pool_seq) { + let prefix_len = record.sequence.len().min(pool_seq.len()).min(40); + issues.push(RecordValidationIssue::SequenceContentMismatch { + reported_prefix: record.sequence.chars().take(prefix_len).collect(), + actual_prefix: pool_seq.chars().take(prefix_len).collect(), + }); + } + + // Coordinate ordering for V/D/J segments and germline ranges. + check_coord_pair( + Segment::V, + record.v_sequence_start, + record.v_sequence_end, + false, + issues, + ); + check_coord_pair( + Segment::V, + record.v_germline_start, + record.v_germline_end, + true, + issues, + ); + check_coord_pair( + Segment::D, + record.d_sequence_start, + record.d_sequence_end, + false, + issues, + ); + check_coord_pair( + Segment::D, + record.d_germline_start, + record.d_germline_end, + true, + issues, + ); + check_coord_pair( + Segment::J, + record.j_sequence_start, + record.j_sequence_end, + false, + issues, + ); + check_coord_pair( + Segment::J, + record.j_germline_start, + record.j_germline_end, + true, + issues, + ); + + // CIGAR span sanity: M+I ops must equal the sequence-side span. + for (seg, start, end, cigar) in [ + ( + Segment::V, + record.v_sequence_start, + record.v_sequence_end, + &record.v_cigar, + ), + ( + Segment::D, + record.d_sequence_start, + record.d_sequence_end, + &record.d_cigar, + ), + ( + Segment::J, + record.j_sequence_start, + record.j_sequence_end, + &record.j_cigar, + ), + ] { + if cigar.is_empty() { + continue; + } + match parse_cigar(cigar) { + Ok(ops) => { + let query_span: usize = ops + .iter() + .filter(|(_, op)| *op == b'M' || *op == b'I') + .map(|(n, _)| *n) + .sum(); + if let (Some(s), Some(e)) = (start, end) { + let sequence_span = (e - s).max(0) as usize; + if query_span != sequence_span { + issues.push(RecordValidationIssue::CigarSpanMismatch { + segment: seg, + cigar_query_span: query_span, + sequence_span, + }); + } + } + } + Err(reason) => issues.push(RecordValidationIssue::CigarReadsInvalid { + segment: seg, + reason, + }), + } + } +} + +fn check_coord_pair( + segment: Segment, + start: Option, + end: Option, + is_germline: bool, + issues: &mut Vec, +) { + if let (Some(s), Some(e)) = (start, end) { + // Half-open: end > start is required (or both 0 for empty). + if s < 0 || e < s { + let issue = if is_germline { + RecordValidationIssue::GermlineCoordinatesOutOfOrder { + segment, + start: s, + end: e, + } + } else { + RecordValidationIssue::SegmentCoordinatesOutOfOrder { + segment, + start: s, + end: e, + } + }; + issues.push(issue); + } + } +} + +fn parse_cigar(cigar: &str) -> Result, String> { + let mut ops = Vec::new(); + let mut digits = String::new(); + for ch in cigar.chars() { + if ch.is_ascii_digit() { + digits.push(ch); + } else { + let n: usize = digits + .parse() + .map_err(|_| format!("non-numeric op length in CIGAR {cigar:?}"))?; + if !matches!(ch as u8, b'M' | b'I' | b'D' | b'S' | b'N' | b'P' | b'X' | b'=') { + return Err(format!("unrecognized CIGAR op {ch:?} in {cigar:?}")); + } + ops.push((n, ch as u8)); + digits.clear(); + } + } + if !digits.is_empty() { + return Err(format!("trailing digits in CIGAR {cigar:?}")); + } + Ok(ops) +} diff --git a/engine_rs/src/lib.rs b/engine_rs/src/lib.rs index 6d85e2a..2b660fc 100644 --- a/engine_rs/src/lib.rs +++ b/engine_rs/src/lib.rs @@ -1,7 +1,6 @@ //! GenAIRR engine — Rust kernel. //! -//! This crate implements the simulation architecture described in -//! `.private/engine_v6_living_design_2026-05-05.md`. +//! This crate implements the GenAIRR simulation architecture. //! //! # Contributor-facing architecture //! diff --git a/engine_rs/src/passes/mod.rs b/engine_rs/src/passes/mod.rs index 3b90d4b..1ae0a49 100644 --- a/engine_rs/src/passes/mod.rs +++ b/engine_rs/src/passes/mod.rs @@ -58,1357 +58,4 @@ pub use trim::TrimPass; // ────────────────────────────────────────────────────────────────── #[cfg(test)] -mod tests { - use super::*; - use crate::assignment::TrimEnd; - use crate::contract::{ - productive, ChoiceContext, Contract, ContractSet, ContractViolation, - ProductiveJunctionFrame, - }; - use crate::dist::{ - AllelePoolDist, Distribution, EmpiricalLengthDist, FilteredSampleError, UniformBase, - }; - use crate::ir::{NucHandle, Segment, Simulation}; - use crate::pass::testing::PassRuntime; - use crate::pass::{Pass, PassPlan}; - use crate::passes::sample_allele::test_support::make_test_pool; - use crate::s5f::{S5FKernel, S5F_NUM_CONTEXTS, S5F_SUBSTITUTION_LEN}; - use crate::trace::ChoiceValue; - - fn fixed_count_dist() -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])) - } - - fn test_s5f_kernel() -> S5FKernel { - S5FKernel::new( - vec![1.0; S5F_NUM_CONTEXTS], - vec![0.25; S5F_SUBSTITUTION_LEN], - ) - } - - fn assert_typed_patterns_project_to_declared_choices(pass: &dyn Pass) { - let projected: Vec = pass - .declared_choice_patterns() - .into_iter() - .map(String::from) - .collect(); - assert_eq!( - projected, - pass.declared_choices(), - "{} typed declared-choice patterns must project to the stable string surface", - pass.name() - ); - } - - #[test] - fn built_in_pass_declared_choice_patterns_match_string_surface() { - let v_pool = make_test_pool(2, Segment::V); - let d_pool = make_test_pool(2, Segment::D); - let j_pool = make_test_pool(2, Segment::J); - - let passes: Vec> = vec![ - Box::new(EchoPass::new(b'A', 0, Segment::V)), - Box::new(AssembleSegmentPass::new(Segment::V)), - Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&v_pool)), - )), - Box::new(SampleAllelePass::new( - Segment::D, - Box::new(AllelePoolDist::uniform(&d_pool)), - )), - Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&j_pool)), - )), - Box::new(TrimPass::new(Segment::V, TrimEnd::Five, fixed_count_dist())), - Box::new(TrimPass::new( - Segment::D, - TrimEnd::Three, - fixed_count_dist(), - )), - Box::new(TrimPass::new(Segment::J, TrimEnd::Five, fixed_count_dist())), - Box::new(GenerateNPPass::new( - Segment::Np1, - fixed_count_dist(), - Box::new(UniformBase), - )), - Box::new(GenerateNPPass::new( - Segment::Np2, - fixed_count_dist(), - Box::new(UniformBase), - )), - Box::new(UniformMutationPass::new( - fixed_count_dist(), - Box::new(UniformBase), - )), - Box::new(S5FMutationPass::new(test_s5f_kernel(), fixed_count_dist())), - Box::new(PCRErrorPass::new(fixed_count_dist(), Box::new(UniformBase))), - Box::new(QualityErrorPass::new( - fixed_count_dist(), - Box::new(UniformBase), - )), - Box::new(ContaminantPass::new(0.5, Box::new(UniformBase))), - Box::new(IndelPass::new( - fixed_count_dist(), - 0.5, - Box::new(UniformBase), - )), - Box::new(NCorruptionPass::new(fixed_count_dist())), - Box::new(EndLossPass::new(LossEnd::Five, fixed_count_dist())), - Box::new(EndLossPass::new(LossEnd::Three, fixed_count_dist())), - Box::new(RevCompPass::new(0.5)), - ]; - - for pass in passes { - assert_typed_patterns_project_to_declared_choices(pass.as_ref()); - } - } - - /// Build a synthetic VJ refdata: V "AAACCCGGG" (9bp, anchor 6), - /// J "TTTAAA" (6bp, anchor 0). Junction = V_anchor_to_end (3bp) - /// + NP1 + J_anchor_to_W3 (3bp). For productive frame: - /// (3 + NP1 + 3) % 3 == 0 → NP1 % 3 == 0. - fn make_vj_refdata_for_filter() -> crate::refdata::RefDataConfig { - let mut cfg = crate::refdata::RefDataConfig::empty(crate::refdata::ChainType::Vj); - let _ = cfg.v_pool.push(crate::refdata::Allele { - name: "v_test*01".into(), - gene: "v_test".into(), - seq: b"AAACCCGGG".to_vec(), - segment: Segment::V, - anchor: Some(6), - functional_status: None, - subregions: Vec::new(), - }); - let _ = cfg.j_pool.push(crate::refdata::Allele { - name: "j_test*01".into(), - gene: "j_test".into(), - seq: b"TTTAAA".to_vec(), - segment: Segment::J, - anchor: Some(0), - functional_status: None, - subregions: Vec::new(), - }); - cfg - } - - #[test] - fn productive_admits_filters_np1_for_vj_chain() { - let cfg = make_vj_refdata_for_filter(); - let dist = EmpiricalLengthDist::from_pairs((0..7).map(|i| (i, 1.0)).collect::>()); - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(dist), - Box::new(UniformBase), - ))); - - let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); - - for seed in 0..20u64 { - let outcome = PassRuntime::execute_with_context( - &plan, - Simulation::new(), - seed, - Some(&cfg), - Some(&contracts), - ); - let np1_len = match outcome.trace.find("np.np1.length").unwrap().value { - ChoiceValue::Int(n) => n, - _ => panic!("wrong variant"), - }; - assert!( - np1_len % 3 == 0, - "seed {} produced NP1 length {} (not divisible by 3)", - seed, - np1_len - ); - } - } - - #[test] - fn productive_admits_unconstrained_without_contracts() { - let cfg = make_vj_refdata_for_filter(); - let dist = EmpiricalLengthDist::from_pairs((0..7).map(|i| (i, 1.0)).collect::>()); - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(dist), - Box::new(UniformBase), - ))); - - // Without contracts: any value in [0, 6] is allowed. - let mut seen_out_of_frame = false; - for seed in 0..50u64 { - let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), seed, &cfg); - let np1_len = match outcome.trace.find("np.np1.length").unwrap().value { - ChoiceValue::Int(n) => n, - _ => panic!("wrong variant"), - }; - assert!(np1_len >= 0 && np1_len <= 6); - if np1_len % 3 != 0 { - seen_out_of_frame = true; - } - } - assert!( - seen_out_of_frame, - "Without contracts, expected at least one out-of-frame NP1 sample" - ); - } - - #[test] - fn productive_admits_empty_filter_returns_explicit_zero_length() { - // v3.0 rule: when natural ∩ admissible is empty under - // active contracts, permissive mode must NOT fall back - // to an unconstrained draw from the natural support - // (that would re-introduce reject-after-propose). The - // architectural no-op for a length sampler is 0 — same - // shape as `TrimPass::sample_trim`'s empty-support - // fallback. (Pre-v3.0 the pass fell through to - // `length_dist.sample(rng)` and produced one of - // {1,2,4,5}, all of which violate frame in this fixture.) - let cfg = make_vj_refdata_for_filter(); - let dist = EmpiricalLengthDist::from_pairs(vec![(1, 1.0), (2, 1.0), (4, 1.0), (5, 1.0)]); - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(dist), - Box::new(UniformBase), - ))); - - let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); - - let outcome = PassRuntime::execute_with_context( - &plan, - Simulation::new(), - 0, - Some(&cfg), - Some(&contracts), - ); - let np1_len = match outcome.trace.find("np.np1.length").unwrap().value { - ChoiceValue::Int(n) => n, - _ => panic!("wrong variant"), - }; - assert_eq!(np1_len, 0); - } - - #[derive(Clone, Debug)] - struct UnenumerableLengthDist; - - impl Distribution for UnenumerableLengthDist { - type Output = i64; - - fn sample(&self, _rng: &mut crate::rng::Rng) -> i64 { - 0 - } - } - - #[test] - fn productive_strict_errors_when_length_filter_empty() { - let cfg = make_vj_refdata_for_filter(); - let dist = EmpiricalLengthDist::from_pairs(vec![(1, 1.0), (2, 1.0), (4, 1.0), (5, 1.0)]); - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(dist), - Box::new(UniformBase), - ))); - - let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); - let err = PassRuntime::execute_strict_with_context( - &plan, - Simulation::new(), - 0, - Some(&cfg), - Some(&contracts), - ) - .unwrap_err(); - - assert_eq!(err.pass_name(), "generate_np.np1"); - assert_eq!(err.address(), "np.np1.length"); - assert_eq!( - err.constraint_reason(), - Some(FilteredSampleError::EmptyAdmissibleSupport) - ); - assert!(err.to_string().contains("no admissible candidates")); - } - - #[test] - fn productive_strict_errors_when_length_support_unavailable() { - let cfg = make_vj_refdata_for_filter(); - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(UnenumerableLengthDist), - Box::new(UniformBase), - ))); - - let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); - let err = PassRuntime::execute_strict_with_context( - &plan, - Simulation::new(), - 0, - Some(&cfg), - Some(&contracts), - ) - .unwrap_err(); - - assert_eq!(err.pass_name(), "generate_np.np1"); - assert_eq!(err.address(), "np.np1.length"); - assert_eq!( - err.constraint_reason(), - Some(FilteredSampleError::SupportUnavailable) - ); - } - - struct RejectNpBases; - - impl Contract for RejectNpBases { - fn name(&self) -> &str { - "reject_np_bases" - } - - fn verify( - &self, - _sim: &Simulation, - _refdata: Option<&crate::refdata::RefDataConfig>, - ) -> Result<(), ContractViolation> { - Ok(()) - } - - // Post-flip the trait's primary surface is `admits_typed`, - // so this test contract dispatches on `ChoiceAddress::NpBase` - // directly rather than prefix-matching the legacy string. - fn admits_typed( - &self, - _sim: &Simulation, - _refdata: Option<&crate::refdata::RefDataConfig>, - context: ChoiceContext<'_>, - _candidate: &ChoiceValue, - ) -> Result<(), ContractViolation> { - if matches!( - context.address, - Some(crate::address::ChoiceAddress::NpBase { - segment: crate::address::NpSegment::Np1, - .. - }) - ) { - return Err(ContractViolation::new(self.name(), "rejected by test")); - } - Ok(()) - } - } - - #[test] - fn productive_admits_empty_base_filter_writes_n_sentinel_in_permissive() { - // v3.0 rule (mirror of `sample_length`'s explicit-zero - // fallback): when the bundle leaves no admissible base - // for an NP slot, permissive mode writes the IUPAC `N` - // sentinel rather than falling back to an unconstrained - // draw. `N` translates to amino acid `X`, which is not - // a stop and not bound by anchor preservation (NP slots - // are between V/J anchor codons), so it satisfies the - // productive bundle's full admit check. - // - // The matching strict test below (`productive_strict_errors_…`) - // pins that strict mode surfaces the same empty support - // as `EmptyAdmissibleSupport` instead. - let mut plan = PassPlan::new(); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])), - Box::new(UniformBase), - ))); - let contracts = ContractSet::new().with(Box::new(RejectNpBases)); - - let outcome = - PassRuntime::execute_with_context(&plan, Simulation::new(), 0, None, Some(&contracts)); - - // The trace's recorded base for NP slot 0 is `b'N'`. - assert_eq!( - outcome.trace.find("np.np1.bases[0]").unwrap().value, - ChoiceValue::Base(b'N') - ); - // The pool reflects the same sentinel. - let final_sim = outcome.final_simulation(); - assert_eq!(final_sim.pool.len(), 1); - assert_eq!(final_sim.pool.get(NucHandle::new(0)).unwrap().base, b'N'); - } - - #[test] - fn productive_strict_errors_when_base_filter_empty() { - let mut plan = PassPlan::new(); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])), - Box::new(UniformBase), - ))); - let contracts = ContractSet::new().with(Box::new(RejectNpBases)); - - let err = PassRuntime::execute_strict_with_context( - &plan, - Simulation::new(), - 0, - None, - Some(&contracts), - ) - .unwrap_err(); - - assert_eq!(err.pass_name(), "generate_np.np1"); - assert_eq!(err.address(), "np.np1.bases[0]"); - assert_eq!( - err.constraint_reason(), - Some(FilteredSampleError::EmptyAdmissibleSupport) - ); - } - - #[test] - fn productive_strict_succeeds_when_admissible_length_exists() { - let cfg = make_vj_refdata_for_filter(); - let dist = EmpiricalLengthDist::from_pairs((0..7).map(|i| (i, 1.0)).collect::>()); - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(dist), - Box::new(UniformBase), - ))); - - let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); - let outcome = PassRuntime::execute_strict_with_context( - &plan, - Simulation::new(), - 0, - Some(&cfg), - Some(&contracts), - ) - .unwrap(); - let np1_len = match outcome.trace.find("np.np1.length").unwrap().value { - ChoiceValue::Int(n) => n, - _ => panic!("wrong variant"), - }; - - assert_eq!(np1_len % 3, 0); - } - - #[test] - fn productive_admits_makes_full_pipeline_in_frame_for_vj() { - use crate::junction::compute_junction; - - let cfg = make_vj_refdata_for_filter(); - let dist = EmpiricalLengthDist::from_pairs((0..7).map(|i| (i, 1.0)).collect::>()); - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(dist), - Box::new(UniformBase), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); - - let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); - - for seed in 0..30u64 { - let outcome = PassRuntime::execute_with_context( - &plan, - Simulation::new(), - seed, - Some(&cfg), - Some(&contracts), - ); - let junction = compute_junction(outcome.final_simulation(), &cfg) - .expect("junction should be defined"); - assert!( - junction.is_in_frame(), - "seed {} produced out-of-frame junction (length {})", - seed, - junction.length - ); - } - } - - #[test] - fn productive_full_bundle_in_frame_and_admits_dispatch_works() { - use crate::junction::compute_junction; - - let cfg = make_vj_refdata_for_filter(); - let dist = EmpiricalLengthDist::from_pairs((0..10).map(|i| (i, 1.0)).collect::>()); - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - Box::new(dist), - Box::new(UniformBase), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); - - let contracts = productive(); - - for seed in 0..20u64 { - let outcome = PassRuntime::execute_with_context( - &plan, - Simulation::new(), - seed, - Some(&cfg), - Some(&contracts), - ); - let junction = compute_junction(outcome.final_simulation(), &cfg) - .expect("junction should be defined"); - assert!(junction.is_in_frame()); - } - } - - #[test] - fn sample_allele_pass_replay_in_mixed_plan_with_echo_passes() { - // Mixed plan: SampleAllele + Echo (transform) interleaved. - // The trace records only the sampling choices; the IR - // accumulates both the assignment and the echoed nucleotides. - let pool = make_test_pool(3, Segment::V); - let mut plan = PassPlan::new(); - plan.push(Box::new(EchoPass::new(b'A', 0, Segment::V))); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&pool)), - ))); - plan.push(Box::new(EchoPass::new(b'C', 1, Segment::V))); - - let oa = PassRuntime::execute(&plan, Simulation::new(), 0xfeed); - let ob = PassRuntime::execute(&plan, Simulation::new(), 0xfeed); - - assert_eq!( - oa.final_simulation().pool.len(), - ob.final_simulation().pool.len() - ); - assert_eq!( - oa.final_simulation() - .assignments - .get(Segment::V) - .copied() - .unwrap() - .allele_id, - ob.final_simulation() - .assignments - .get(Segment::V) - .copied() - .unwrap() - .allele_id - ); - // Trace contains exactly one entry — only the sampling pass records. - assert_eq!(oa.trace.len(), 1); - assert_eq!(oa.trace.choices()[0].address, "sample_allele.v"); - } - - /// A representative mixed plan: alternating EchoPass (no RNG) - /// and SampleBasePass (RNG-consuming). Used by the cross-cutting - /// replay-determinism test below. - fn mixed_plan() -> PassPlan { - use crate::ir::flag; - let mut plan = PassPlan::new(); - for i in 0..8 { - plan.push(Box::new(EchoPass::new(b'A', i as u16, Segment::V))); - plan.push(Box::new(SampleBasePass::new( - format!("np.np1.bases[{}]", i), - Box::new(UniformBase), - Segment::Np1, - flag::N_NUC, - ))); - } - plan - } - - #[test] - fn replay_determinism_holds_under_repeated_runs() { - // Run the same plan with the same seed five times. Every - // trace and every final IR must be identical to the first. - let baseline = PassRuntime::execute(&mixed_plan(), Simulation::new(), 0xfeed_face); - for _ in 0..5 { - let other = PassRuntime::execute(&mixed_plan(), Simulation::new(), 0xfeed_face); - assert_eq!(other.trace.len(), baseline.trace.len()); - for (a, b) in baseline - .trace - .choices() - .iter() - .zip(other.trace.choices().iter()) - { - assert_eq!(a.address, b.address); - assert_eq!(a.value, b.value); - } - assert_eq!( - other.final_simulation().pool.len(), - baseline.final_simulation().pool.len() - ); - } - } - - // ────────────────────────────────────────────────────────────── - // Per-pass `EventRecord.simulation_events` outcome contract - // - // Every committed pass now carries two complementary surfaces: - // - `trace_span` → choices the pass consumed/emitted - // - `simulation_events` → state consequences of those choices - // - // These tests pin which events each pass produces, indexed off - // the committed `outcome.events` ledger — i.e. observability - // from the *outcome* surface, not the test-only PassContext - // hook. - // ────────────────────────────────────────────────────────────── - - /// Find exactly-one event matching `pat` inside `events`. Fails - /// the test if zero or more-than-one match. Returns the match. - fn find_one( - events: &[crate::ir::SimulationEvent], - pat: F, - ) -> &crate::ir::SimulationEvent - where - F: Fn(&crate::ir::SimulationEvent) -> bool, - { - let matches: Vec<&crate::ir::SimulationEvent> = - events.iter().filter(|e| pat(e)).collect(); - assert_eq!( - matches.len(), - 1, - "expected exactly one matching event, got {} from {:?}", - matches.len(), - events - ); - matches[0] - } - - #[test] - fn outcome_event_records_carry_assignment_trim_and_region_events() { - // Plan: sample V → trim V 5' → assemble V → generate NP1 → - // assemble J. Each committed pass's `EventRecord` should - // carry the consequence event it actually fired. - let cfg = make_vj_refdata_for_filter(); - let trim_dist = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(0, 1.0)])) - }; - let np_dist = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(3, 1.0)])) - }; - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(TrimPass::new( - Segment::V, - TrimEnd::Five, - trim_dist(), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - np_dist(), - Box::new(UniformBase), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); - - let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 42, &cfg); - - // Sanity: pass count matches plan length. - assert_eq!(outcome.events.len(), plan.len()); - - // Pass 0 — SampleAllele V → one AssignmentChanged for V. - let ev0 = &outcome.events[0]; - assert_eq!(ev0.pass_name, "sample_allele.v"); - let assign = find_one(&ev0.simulation_events, |e| { - matches!( - e, - crate::ir::SimulationEvent::AssignmentChanged { - segment: Segment::V, - .. - } - ) - }); - match assign { - crate::ir::SimulationEvent::AssignmentChanged { segment, old, .. } => { - assert_eq!(*segment, Segment::V); - assert!(old.is_none(), "first-install AssignmentChanged has old=None"); - } - _ => unreachable!(), - } - - // Pass 2 — TrimPass V 5' → one TrimChanged for (V, Five). - let ev2 = &outcome.events[2]; - assert_eq!(ev2.pass_name, "trim.v_5"); - let trim_event = find_one(&ev2.simulation_events, |e| { - matches!( - e, - crate::ir::SimulationEvent::TrimChanged { - segment: Segment::V, - end: TrimEnd::Five, - .. - } - ) - }); - match trim_event { - crate::ir::SimulationEvent::TrimChanged { new, .. } => { - assert_eq!(*new, 0, "trim distribution fixed at 0"); - } - _ => unreachable!(), - } - - // Pass 3 — AssembleSegment V → one RegionAdded for V. - let ev3 = &outcome.events[3]; - assert_eq!(ev3.pass_name, "assemble.v"); - let v_region = find_one(&ev3.simulation_events, |e| { - matches!( - e, - crate::ir::SimulationEvent::RegionAdded { region } if region.segment == Segment::V - ) - }); - match v_region { - crate::ir::SimulationEvent::RegionAdded { region } => { - // V allele is 9 bases long with trim_5=0 → region [0, 9). - assert_eq!(region.start, NucHandle::new(0)); - assert_eq!(region.end, NucHandle::new(9)); - } - _ => unreachable!(), - } - - // Pass 4 — GenerateNP Np1 → one RegionAdded for Np1. - let ev4 = &outcome.events[4]; - assert_eq!(ev4.pass_name, "generate_np.np1"); - let np_region = find_one(&ev4.simulation_events, |e| { - matches!( - e, - crate::ir::SimulationEvent::RegionAdded { region } if region.segment == Segment::Np1 - ) - }); - match np_region { - crate::ir::SimulationEvent::RegionAdded { region } => { - // NP1 starts where V ended (handle 9), length 3. - assert_eq!(region.start, NucHandle::new(9)); - assert_eq!(region.end, NucHandle::new(12)); - } - _ => unreachable!(), - } - - // Pass 5 — AssembleSegment J → one RegionAdded for J. - let ev5 = &outcome.events[5]; - assert_eq!(ev5.pass_name, "assemble.j"); - let _ = find_one(&ev5.simulation_events, |e| { - matches!( - e, - crate::ir::SimulationEvent::RegionAdded { region } if region.segment == Segment::J - ) - }); - } - - #[test] - fn outcome_event_records_carry_base_changed_and_mutation_count_for_mutation_pass() { - use crate::passes::UniformMutationPass; - // Plan: sample V → assemble V → uniform-mutation 1. - // The mutation pass's EventRecord must carry exactly one - // `BaseChanged` (the single applied substitution) and - // exactly one `MutationCountChanged` (the commit-time count - // bump). - let cfg = make_vj_refdata_for_filter(); - let mutation_count_dist = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])) - }; - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(UniformMutationPass::new( - mutation_count_dist(), - Box::new(UniformBase), - ))); - - let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 42, &cfg); - - // The mutation pass is the third committed pass. - let mut_event = outcome - .events - .iter() - .find(|e| e.pass_name == "mutate.uniform") - .expect("uniform mutation pass committed an event record"); - - let _ = find_one(&mut_event.simulation_events, |e| { - matches!(e, crate::ir::SimulationEvent::BaseChanged { .. }) - }); - let mc = find_one(&mut_event.simulation_events, |e| { - matches!(e, crate::ir::SimulationEvent::MutationCountChanged { .. }) - }); - match mc { - crate::ir::SimulationEvent::MutationCountChanged { - old, - new, - delta, - } => { - assert_eq!(*old, 0); - assert_eq!(*new, 1); - assert_eq!(*delta, 1); - } - _ => unreachable!(), - } - - // Final mutation_count carried on the sealed sim matches the - // event's `new` field — pin the outcome-ledger / sim - // consistency contract. - assert_eq!(outcome.final_simulation().mutation_count, 1); - } - - // ────────────────────────────────────────────────────────────── - // Event-coverage audit + invariant tests - // - // Three coordinated tests that lock down the consequence-event - // ledger before any derived-state lifecycle refactor: - // - // 1. Cross-pass coverage — every meaningful pass emits the - // consequence variants it owns. - // 2. Event/state consistency — outcome counts reconstructed - // from `simulation_events` match the final `Simulation`. - // 3. Trace/event alignment — `TraceSpan`s are contiguous, - // cover the full trace, and replay equality holds. - // ────────────────────────────────────────────────────────────── - - /// Find every event matching `pat` in a slice. Used in the - /// audit tests where we care about "at least one of X" and - /// counts, not exact-one-of-X. - fn count_matching(events: &[crate::ir::SimulationEvent], pat: F) -> usize - where - F: Fn(&crate::ir::SimulationEvent) -> bool, - { - events.iter().filter(|e| pat(e)).count() - } - - /// Build the heavy-stack VJ plan used by the coverage audit. - /// Every category the user listed (sample/trim/assemble/np/ - /// mutation/pcr/quality/ncorrupt/indel/rev-comp) is represented - /// by exactly one pass. Distributions are pinned to fixed - /// counts so the audit is deterministic across RNG paths. - fn build_coverage_audit_plan( - cfg: &crate::refdata::RefDataConfig, - ) -> PassPlan { - use crate::passes::{ - IndelPass, NCorruptionPass, PCRErrorPass, QualityErrorPass, RevCompPass, - UniformMutationPass, - }; - let count_one = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])) - }; - let count_zero = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(0, 1.0)])) - }; - let np_three = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(3, 1.0)])) - }; - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(TrimPass::new( - Segment::V, - TrimEnd::Five, - count_zero(), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - np_three(), - Box::new(UniformBase), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); - plan.push(Box::new(UniformMutationPass::new( - count_one(), - Box::new(UniformBase), - ))); - plan.push(Box::new(PCRErrorPass::new(count_one(), Box::new(UniformBase)))); - plan.push(Box::new(QualityErrorPass::new( - count_one(), - Box::new(UniformBase), - ))); - plan.push(Box::new(NCorruptionPass::new(count_one()))); - plan.push(Box::new(IndelPass::new( - count_one(), - 0.5, - Box::new(UniformBase), - ))); - plan.push(Box::new(RevCompPass::new(1.0))); - plan - } - - #[test] - fn cross_pass_event_coverage_emits_expected_consequence_types() { - // Drive a heavy-stack VJ plan covering every event-emitting - // pass category. Assert each pass's `EventRecord.simulation_events` - // carries at least one event of the expected variant — this - // is the load-bearing "no holes in the ledger" check before - // we lean on it for lifecycle decisions. - use crate::ir::SimulationEvent; - - let cfg = make_vj_refdata_for_filter(); - let plan = build_coverage_audit_plan(&cfg); - let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 42, &cfg); - - // Convenience: locate a committed pass's event record by - // its `pass_name`. Panics if missing because the plan is - // hard-coded — any missing pass is a regression. - let find = |name: &str| { - outcome - .events - .iter() - .find(|e| e.pass_name == name) - .unwrap_or_else(|| panic!("missing event record for {name}")) - }; - - // Sample allele: AssignmentChanged for V and J. - assert!( - count_matching(&find("sample_allele.v").simulation_events, |e| matches!( - e, - SimulationEvent::AssignmentChanged { segment: Segment::V, .. } - )) >= 1 - ); - assert!( - count_matching(&find("sample_allele.j").simulation_events, |e| matches!( - e, - SimulationEvent::AssignmentChanged { segment: Segment::J, .. } - )) >= 1 - ); - - // Trim: TrimChanged. - assert!( - count_matching(&find("trim.v_5").simulation_events, |e| matches!( - e, - SimulationEvent::TrimChanged { segment: Segment::V, .. } - )) >= 1 - ); - - // Assemble: many BasePushed + one RegionAdded per segment. - for (name, seg) in [("assemble.v", Segment::V), ("assemble.j", Segment::J)] { - let ev = find(name); - assert!( - count_matching(&ev.simulation_events, |e| matches!(e, SimulationEvent::BasePushed { .. })) > 0, - "{name} should emit BasePushed events" - ); - assert!( - count_matching(&ev.simulation_events, |e| matches!( - e, - SimulationEvent::RegionAdded { region } if region.segment == seg - )) == 1, - "{name} should emit exactly one RegionAdded for its segment" - ); - } - - // NP: BasePushed (length 3) + one RegionAdded for Np1. - let np = find("generate_np.np1"); - assert!( - count_matching(&np.simulation_events, |e| matches!(e, SimulationEvent::BasePushed { .. })) >= 1 - ); - assert!( - count_matching(&np.simulation_events, |e| matches!( - e, - SimulationEvent::RegionAdded { region } if region.segment == Segment::Np1 - )) == 1 - ); - - // Uniform mutation: BaseChanged + MutationCountChanged. - let mu = find("mutate.uniform"); - assert!( - count_matching(&mu.simulation_events, |e| matches!(e, SimulationEvent::BaseChanged { .. })) >= 1 - ); - assert!( - count_matching(&mu.simulation_events, |e| matches!( - e, - SimulationEvent::MutationCountChanged { .. } - )) == 1 - ); - - // PCR / quality / ncorrupt: BaseChanged (no mutation-count - // bump — these are sequencing artifacts, not biological - // mutations, and don't tag `add_to_mutation_count`). - for name in ["corrupt.pcr", "corrupt.quality", "corrupt.ns"] { - let ev = find(name); - assert!( - count_matching(&ev.simulation_events, |e| matches!(e, SimulationEvent::BaseChanged { .. })) >= 1, - "{name} should emit at least one BaseChanged" - ); - assert!( - count_matching(&ev.simulation_events, |e| matches!( - e, - SimulationEvent::MutationCountChanged { .. } - )) == 0, - "{name} is a sequencing-artifact pass and must not bump n_mutations" - ); - } - - // Indel: at least one of IndelInserted / IndelDeleted. - let indel = find("corrupt.indel"); - let indels = count_matching(&indel.simulation_events, |e| { - matches!( - e, - SimulationEvent::IndelInserted { .. } | SimulationEvent::IndelDeleted { .. } - ) - }); - assert!(indels >= 1, "corrupt.indel should emit at least one indel event"); - - // Rev-comp: one ReverseComplementFlagRecorded with applied=true - // (apply_prob = 1.0 in the audit plan). - let rc = find("corrupt.rev_comp"); - assert_eq!( - count_matching(&rc.simulation_events, |e| matches!( - e, - SimulationEvent::ReverseComplementFlagRecorded { applied: true } - )), - 1 - ); - } - - #[test] - fn outcome_events_are_consistent_with_final_simulation_state() { - // Reconstruct three coarse counts purely from the - // `simulation_events` stream and assert they match the - // sealed final `Simulation`. This is the load-bearing - // self-consistency proof: a future derived-state observer - // can trust the ledger as the source of truth. - use crate::ir::SimulationEvent; - - let cfg = make_vj_refdata_for_filter(); - let count_one = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])) - }; - let np_three = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(3, 1.0)])) - }; - - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - np_three(), - Box::new(UniformBase), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); - plan.push(Box::new(crate::passes::UniformMutationPass::new( - count_one(), - Box::new(UniformBase), - ))); - - let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 0, &cfg); - let final_sim = outcome.final_simulation(); - - // ── Invariant 1: regions ──────────────────────────────── - // Indel-free plan → final region count equals the total - // `RegionAdded` events across all passes. - let region_adds: usize = outcome - .events - .iter() - .flat_map(|e| &e.simulation_events) - .filter(|e| matches!(e, SimulationEvent::RegionAdded { .. })) - .count(); - assert_eq!( - region_adds, - final_sim.sequence.regions.len(), - "sum of RegionAdded events must equal final region count for indel-free plan" - ); - - // ── Invariant 2: assignments ──────────────────────────── - // The last `AssignmentChanged` per segment must equal the - // segment's final `AlleleInstance` in `final_sim.assignments`. - use crate::assignment::AlleleInstance; - let mut last_by_segment: std::collections::HashMap = - std::collections::HashMap::new(); - for ev in outcome.events.iter().flat_map(|e| &e.simulation_events) { - if let SimulationEvent::AssignmentChanged { segment, new, .. } = ev { - last_by_segment.insert(*segment, *new); - } - } - for (seg, last_assignment) in &last_by_segment { - assert_eq!( - final_sim.assignments.get(*seg).copied(), - Some(*last_assignment), - "last AssignmentChanged for {:?} must equal final assignment", - seg - ); - } - // V and J should both have been assigned. - assert!(last_by_segment.contains_key(&Segment::V)); - assert!(last_by_segment.contains_key(&Segment::J)); - - // ── Invariant 3: mutation_count ───────────────────────── - // The last `MutationCountChanged.new` across the run must - // equal `final_sim.mutation_count`. (Plan has exactly one - // count-bumping pass; this still holds with multiple.) - let final_mc = outcome - .events - .iter() - .flat_map(|e| &e.simulation_events) - .filter_map(|e| match e { - SimulationEvent::MutationCountChanged { new, .. } => Some(*new), - _ => None, - }) - .last() - .expect("plan has one mutation pass"); - assert_eq!(final_mc, final_sim.mutation_count); - } - - #[test] - fn trace_spans_partition_outcome_trace_and_replay_preserves_trace() { - // Two-part alignment proof: - // - // (a) Every committed pass yields one `EventRecord` whose - // `trace_span` is contiguous with its neighbours and - // whose union covers the entire outcome trace. No - // gaps, no overlaps. - // - // (b) Replaying the captured trace through `CompiledSimulator` - // produces an outcome whose trace matches the original - // record-for-record — proving event capture is - // transparent w.r.t. the trace, the only persisted - // artifact. - use crate::compiled::ExecutionPolicy; - - let cfg = make_vj_refdata_for_filter(); - let np_three = || -> Box> { - Box::new(EmpiricalLengthDist::from_pairs(vec![(3, 1.0)])) - }; - let mut plan = PassPlan::new(); - plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - np_three(), - Box::new(UniformBase), - ))); - plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); - - // ── (a) Trace-span partition ───────────────────────────── - let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 7, &cfg); - assert_eq!( - outcome.events.len(), - plan.len(), - "every committed pass must contribute exactly one EventRecord" - ); - assert_eq!(outcome.events[0].trace_span.start, 0); - for i in 1..outcome.events.len() { - assert_eq!( - outcome.events[i - 1].trace_span.end, - outcome.events[i].trace_span.start, - "trace spans must be contiguous (gap or overlap at boundary {})", - i - ); - } - assert_eq!( - outcome - .events - .last() - .map(|e| e.trace_span.end) - .unwrap_or(0), - outcome.trace.len(), - "the union of all trace spans must cover the full outcome trace" - ); - - // ── (b) Replay equality ────────────────────────────────── - // Compile the same plan + refdata into an owned simulator, - // replay the captured trace, and assert the replayed - // trace matches the original record-for-record. Event - // capture happens on both runs (the compiled path always - // supplies a sink) but neither side persists events — so a - // matching trace is sufficient proof that capture is - // transparent. - // - // `OwnedCompiledSimulator::compile` takes ownership of the - // plan + refdata, so build a fresh equivalent plan for the - // replay leg. (The borrowing `CompiledSimulator` doesn't - // expose `replay_from_trace_records`.) - let mut replay_plan = PassPlan::new(); - replay_plan.push(Box::new(SampleAllelePass::new( - Segment::V, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ))); - replay_plan.push(Box::new(SampleAllelePass::new( - Segment::J, - Box::new(AllelePoolDist::uniform(&cfg.j_pool)), - ))); - replay_plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); - replay_plan.push(Box::new(GenerateNPPass::new( - Segment::Np1, - np_three(), - Box::new(UniformBase), - ))); - replay_plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); - let compiled = crate::compiled::OwnedCompiledSimulator::compile_with_options( - replay_plan, - Some(cfg.clone()), - None, - ExecutionPolicy::Permissive, - crate::compiled::CompileOptions::skip_refdata_validation(), - ) - .expect("audit plan should compile"); - let original_records: Vec<_> = outcome.trace.choices().to_vec(); - let replayed = compiled - .replay_from_trace_records( - &original_records, - /* seed unused for migrated sites */ 0, - ExecutionPolicy::Strict, - ) - .expect("replay should succeed against the same plan"); - let replayed_records: Vec<_> = replayed.trace.choices().to_vec(); - assert_eq!( - original_records.len(), - replayed_records.len(), - "replay must produce a trace of the same length" - ); - for (i, (orig, repl)) in original_records.iter().zip(replayed_records.iter()).enumerate() { - assert_eq!(orig.address, repl.address, "address divergence at record {}", i); - assert_eq!(orig.value, repl.value, "value divergence at record {} (addr={})", i, orig.address); - } - } - - // ────────────────────────────────────────────────────────────── - // Compile-effect / event-emission policy conformance - // - // Built-in mutating passes that declare a `PassCompileEffect` - // are expected to emit at least one corresponding - // `SimulationEvent` when they actually mutate state. This test - // exercises the heavy-stack VJ plan and asserts every - // `EventRecord` passes the policy check in - // [`crate::event::check_event_emission_consistency`]. - // - // It is the regression net for the "new pass mutates Simulation - // directly via sim.with_* and forgets to emit events" anti- - // pattern. The runtime live-call refresh trusts events; a - // silent zero-event mutating pass would skip the refresh - // (this is also what the divergence tests in - // `compiled::tests::live_call_edits` demonstrate by construction). - // ────────────────────────────────────────────────────────────── - - #[test] - fn builtin_passes_emit_events_consistent_with_declared_compile_effects() { - let cfg = make_vj_refdata_for_filter(); - let plan = build_coverage_audit_plan(&cfg); - let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 42, &cfg); - - // Apply the policy to every committed pass. The plan's - // distributions are pinned to deterministic non-zero - // counts so the EditBases / StructuralIndel zero- - // exemption never fires here — meaningful coverage. - let mut violations: Vec = Vec::new(); - for record in &outcome.events { - if let Err(err) = crate::event::check_event_emission_consistency(record) { - violations.push(format!( - "{}: declared {:?} but {}", - err.pass_name, err.compile_effect, err.reason - )); - } - } - assert!( - violations.is_empty(), - "built-in passes must emit events matching their compile-effect declarations.\n\ - Violations:\n {}\n\n\ - Likely cause: the pass mutates `Simulation` via `sim.with_*` directly \ - instead of routing through `SimulationBuilder` — runtime derived-state \ - refresh trusts events, so a silent mutator gets no refresh.", - violations.join("\n ") - ); - } -} +mod tests; diff --git a/engine_rs/src/passes/receptor_revision.rs b/engine_rs/src/passes/receptor_revision.rs index 4b71b94..4c39d31 100644 --- a/engine_rs/src/passes/receptor_revision.rs +++ b/engine_rs/src/passes/receptor_revision.rs @@ -871,770 +871,4 @@ impl Pass for ReceptorRevisionPass { } #[cfg(test)] -mod tests { - use super::*; - use crate::assignment::AlleleInstance; - use crate::contract::ContractSet; - use crate::dist::AllelePoolDist; - use crate::ir::{NucHandle, Region, Segment, SimulationEvent}; - use crate::pass::testing::PassRuntime; - use crate::pass::{PassError, PassPlan}; - use crate::refdata::{Allele, AllelePool, ChainType, RefDataConfig}; - use crate::replay::TraceCursor; - use crate::rng::Rng; - use crate::trace::Trace; - - fn allele(name: &str, seq: &[u8]) -> Allele { - Allele { - name: name.to_string(), - gene: name.split('*').next().unwrap_or(name).to_string(), - seq: seq.to_vec(), - segment: Segment::V, - anchor: None, - functional_status: None, - subregions: Vec::new(), - } - } - - /// Reference data with two V alleles whose retained 6-byte - /// prefix (under 0 trim) differs. The original assembled V is - /// `AAAAAA` matching allele V0 exactly; replacement candidate - /// V1 carries `GGGGGGCC` — length 8, retains 6 bytes at - /// `trim_3 = 2`. - fn two_v_refdata() -> (RefDataConfig, AlleleId, AlleleId) { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - let v0 = cfg.v_pool.push(allele("V1*01", b"AAAAAA")); - let v1 = cfg.v_pool.push(allele("V1*02", b"GGGGGGCC")); - (cfg, v0, v1) - } - - /// Build a sim with the V slot already assigned to V0 and a - /// V region of length 6 (`AAAAAA`). Doubles for tests as the - /// post-recombination starting state. - fn sim_v_assembled(v0: AlleleId) -> Simulation { - let mut sim = Simulation::new(); - for (i, &b) in b"AAAAAA".iter().enumerate() { - let (next, _) = sim.with_nucleotide_pushed(Nucleotide::germline( - b, - i as u16, - Segment::V, - )); - sim = next; - } - sim.with_allele_assigned(Segment::V, AlleleInstance::new(v0)) - .with_region_added(Region::new( - Segment::V, - NucHandle::new(0), - NucHandle::new(6), - )) - } - - fn v_only_pool(cfg: &RefDataConfig) -> AllelePool { - cfg.v_pool.clone() - } - - fn run_pass( - prob: f64, - seed: u64, - cfg: RefDataConfig, - sim: Simulation, - ) -> (Trace, Simulation) { - let mut plan = PassPlan::new(); - let pool = v_only_pool(&cfg); - plan.push(Box::new(ReceptorRevisionPass::new( - prob, - Box::new(AllelePoolDist::uniform(&pool)), - ))); - let outcome = PassRuntime::execute_with_refdata(&plan, sim, seed, &cfg); - let final_sim = outcome.final_simulation().clone(); - (outcome.trace, final_sim) - } - - fn run_with_ctx( - pass: &ReceptorRevisionPass, - cfg: &RefDataConfig, - contracts: Option<&ContractSet>, - initial: Simulation, - cursor: Option<&mut TraceCursor>, - event_sink: Option<&mut Vec>, - ) -> Result<(Trace, Simulation), PassError> { - let mut trace = Trace::new(); - let mut rng = Rng::new(0xc0ff_ee); - let mut ctx = PassContext { - trace: &mut trace, - rng: &mut rng, - pass_index: 0, - refdata: Some(cfg), - contracts, - feasibility: None, - reference_index: None, - replay_cursor: cursor, - event_log_sink: event_sink, - }; - let next = pass.execute_checked(&initial, &mut ctx)?; - Ok((trace, next)) - } - - // ── Construction guards ────────────────────────────────────── - - #[test] - #[should_panic(expected = "prob must be in [0.0, 1.0]")] - fn receptor_revision_rejects_out_of_range_prob() { - let (cfg, _, _) = two_v_refdata(); - let _ = ReceptorRevisionPass::new( - 1.5, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ); - } - - #[test] - #[should_panic(expected = "prob must be in [0.0, 1.0]")] - fn receptor_revision_rejects_negative_prob() { - let (cfg, _, _) = two_v_refdata(); - let _ = ReceptorRevisionPass::new( - -0.5, - Box::new(AllelePoolDist::uniform(&cfg.v_pool)), - ); - } - - // ── genotype-aware candidate selection ────────────────────── - - fn constraint(per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> GenotypeVConstraint { - GenotypeVConstraint { per_hap, same_haplotype: same } - } - - fn geno_sim(v0: AlleleId) -> Simulation { - sim_v_assembled(v0) - .with_allele_assigned(Segment::V, AlleleInstance::new(v0).with_haplotype(0)) - } - - fn geno_pass(cfg: &RefDataConfig, per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> ReceptorRevisionPass { - ReceptorRevisionPass::new(1.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) - .with_genotype_constraint(constraint(per_hap, same)) - } - - #[test] - fn genotype_same_haplotype_excludes_current_and_restricts_to_chromosome() { - let (cfg, v0, v1) = two_v_refdata(); - // hap0 carries {V0, V1}; hap1 carries {V0}. Current is V0 on hap0. - let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); - let (trace, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap(); - assert_eq!(trace.find("receptor_revision.applied").unwrap().value, ChoiceValue::Bool(true)); - // exclude-current: the only eligible alternate on hap0 is V1 - assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); - assert_eq!(after.assignments.get(Segment::V).unwrap().receptor_revision_original_id, Some(v0)); - assert_eq!(after.assignments.get(Segment::V).unwrap().haplotype, Some(0)); - } - - fn run_permissive(pass: &ReceptorRevisionPass, cfg: &RefDataConfig, initial: Simulation) -> (Trace, Simulation) { - let mut trace = Trace::new(); - let mut rng = Rng::new(0xc0ff_ee); - let mut ctx = PassContext { - trace: &mut trace, - rng: &mut rng, - pass_index: 0, - refdata: Some(cfg), - contracts: None, - feasibility: None, - reference_index: None, - replay_cursor: None, - event_log_sink: None, - }; - let next = pass.execute(&initial, &mut ctx); - (trace, next) - } - - #[test] - fn genotype_no_eligible_alternate_permissive_applied_false() { - let (cfg, v0, _v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); - let (trace, after) = run_permissive(&pass, &cfg, geno_sim(v0)); - assert_eq!(trace.find("receptor_revision.applied").unwrap().value, ChoiceValue::Bool(false)); - assert!(trace.find("receptor_revision.v_allele").is_none()); - assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v0); - } - - #[test] - fn genotype_no_eligible_alternate_strict_errors() { - let (cfg, v0, _v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); - let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap_err(); - assert!(matches!(err, PassError::ConstraintSampling { .. }), "got {err:?}"); - } - - #[test] - fn genotype_both_haplotypes_admits_other_chromosome_allele() { - let (cfg, v0, v1) = two_v_refdata(); - // V1 carried only on hap1; same_haplotype=false aggregates both. - let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v1, 1.0)]], false); - let (_t, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap(); - assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); - // haplotype provenance preserved as the original rearrangement chromosome (0) - assert_eq!(after.assignments.get(Segment::V).unwrap().haplotype, Some(0)); - } - - #[test] - fn genotype_missing_haplotype_stamp_errors() { - let (cfg, v0, v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); - // sim_v_assembled assigns V0 WITHOUT a haplotype stamp. - let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), None, None).unwrap_err(); - assert!(matches!(err, PassError::InvalidPlanState { .. }), "got {err:?}"); - } - - #[test] - fn genotype_strict_post_event_contract_rejection_errors() { - use crate::contract::{Contract, ContractViolation}; - - // A contract whose verify() always rejects the post-event state. - struct RejectAll; - impl Contract for RejectAll { - fn name(&self) -> &str { - "reject_all_test" - } - fn verify( - &self, - _sim: &Simulation, - _refdata: Option<&RefDataConfig>, - ) -> Result<(), ContractViolation> { - Err(ContractViolation::new(self.name(), "rejected by test")) - } - } - - let (cfg, v0, v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); - let contracts = ContractSet::new().with(Box::new(RejectAll)); - // Strict mode (execute_checked) with a rejecting contract must surface a - // contract violation — parity with the non-genotype path. - let err = - run_with_ctx(&pass, &cfg, Some(&contracts), geno_sim(v0), None, None).unwrap_err(); - assert!(matches!(err, PassError::ContractViolation { .. }), "got {err:?}"); - } - - #[test] - fn genotype_replay_strict_post_event_contract_rejection_errors() { - use crate::contract::{Contract, ContractViolation}; - - struct RejectAll; - impl Contract for RejectAll { - fn name(&self) -> &str { - "reject_all_test" - } - fn verify( - &self, - _sim: &Simulation, - _refdata: Option<&RefDataConfig>, - ) -> Result<(), ContractViolation> { - Err(ContractViolation::new(self.name(), "rejected by test")) - } - } - - let (cfg, v0, v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); - let contracts = ContractSet::new().with(Box::new(RejectAll)); - // A valid applied=true replay (v1, trim 2) must STILL be rejected by the - // post-event contract in strict mode — replay parity with the fresh path. - let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); - let err = run_with_ctx(&pass, &cfg, Some(&contracts), geno_sim(v0), Some(&mut cursor), None) - .unwrap_err(); - assert!(matches!(err, PassError::ContractViolation { .. }), "got {err:?}"); - } - - // ── genotype-aware replay validation ──────────────────────── - - fn replay_records(applied: bool, allele: Option, trim: Option) -> Vec { - let mut t = Trace::new(); - t.record_choice(address::ChoiceAddress::ReceptorRevisionApplied, ChoiceValue::Bool(applied)); - if let Some(a) = allele { - t.record_choice(address::ChoiceAddress::ReceptorRevisionVAllele, ChoiceValue::AlleleId(a)); - } - if let Some(tr) = trim { - t.record_choice(address::ChoiceAddress::ReceptorRevisionVTrim3, ChoiceValue::Int(tr)); - } - t.choices().to_vec() - } - - #[test] - fn replay_genotype_valid_replacement_reproduces() { - let (cfg, v0, v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); - let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); - let (_t, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap(); - assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); - assert!(cursor.is_drained()); - } - - #[test] - fn replay_genotype_allele_not_carried_on_haplotype_errors() { - let (cfg, v0, v1) = two_v_refdata(); - // hap0 carries only V0; recorded V1 is not carried on the drawn chromosome. - let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0), (v1, 1.0)]], true); - let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); - let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); - assert!(matches!(err, PassError::InvalidDistributionOutput { .. }), "got {err:?}"); - } - - #[test] - fn replay_genotype_unresolvable_allele_id_reports_missing_allele() { - let (cfg, v0, _v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); - // out-of-range recorded v_allele -> missing_allele (not "not_carried"), - // matching the non-genotype replay diagnostic contract. - let mut cursor = TraceCursor::from_owned(replay_records(true, Some(999), Some(0))); - let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); - match err { - PassError::MissingAllele { allele_id, .. } => assert_eq!(allele_id, 999), - other => panic!("expected MissingAllele, got {other:?}"), - } - } - - #[test] - fn replay_genotype_equals_current_allele_errors() { - let (cfg, v0, v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); - let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v0.index()), Some(0))); - let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); - assert!(matches!(err, PassError::InvalidDistributionOutput { .. }), "got {err:?}"); - } - - #[test] - fn replay_genotype_trim_length_mismatch_errors() { - let (cfg, v0, v1) = two_v_refdata(); - let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); - // V1 len 8, old 6 -> trim must be 2; record 3 (retains 5) => mismatch. - let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(3))); - let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); - assert!(matches!(err, PassError::InvalidPlanState { .. }), "got {err:?}"); - } - - #[test] - fn genotype_signature_differs_by_candidate_set_and_same_haplotype() { - let (cfg, v0, v1) = two_v_refdata(); - let no_geno = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) - .parameter_signature(); - let g_true = geno_pass_prob(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true).parameter_signature(); - let g_false = geno_pass_prob(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], false).parameter_signature(); - let g_other = geno_pass_prob(&cfg, [vec![(v0, 2.0)], vec![(v1, 1.0)]], true).parameter_signature(); - assert_ne!(g_true, no_geno); - assert_ne!(g_true, g_false); - assert_ne!(g_true, g_other); - } - - fn geno_pass_prob(cfg: &RefDataConfig, per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> ReceptorRevisionPass { - ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) - .with_genotype_constraint(constraint(per_hap, same)) - } - - // ── prob=0: no replacement ────────────────────────────────── - - #[test] - fn prob_zero_records_applied_false_and_no_mutation_events() { - let (cfg, v0, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - let mut events: Vec = Vec::new(); - let (trace, after) = run_with_ctx( - &pass, - &cfg, - None, - sim_v_assembled(v0), - None, - Some(&mut events), - ) - .unwrap(); - - // Applied=false recorded. - let rec = trace - .find("receptor_revision.applied") - .expect("applied Bool must be recorded even when prob = 0"); - assert_eq!(rec.value, ChoiceValue::Bool(false)); - // No allele/trim records. - assert!(trace.find("receptor_revision.v_allele").is_none()); - assert!(trace.find("receptor_revision.v_trim_3").is_none()); - - // No state-changing events. - assert!(events.iter().all(|e| !matches!( - e, - SimulationEvent::AssignmentChanged { .. } - | SimulationEvent::TrimChanged { .. } - | SimulationEvent::SegmentReplaced { .. } - ))); - - // V assignment + pool bytes unchanged. - assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v0); - let bases: Vec = after.pool.as_slice().iter().map(|n| n.base).collect(); - assert_eq!(&bases, b"AAAAAA"); - } - - // ── prob=1: replacement ───────────────────────────────────── - - #[test] - fn prob_one_records_three_choices_and_emits_three_events() { - let (cfg, v0, v1) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(1.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - let mut events: Vec = Vec::new(); - let (trace, after) = run_with_ctx( - &pass, - &cfg, - None, - sim_v_assembled(v0), - None, - Some(&mut events), - ) - .unwrap(); - - // applied = true - assert_eq!( - trace.find("receptor_revision.applied").unwrap().value, - ChoiceValue::Bool(true), - ); - // The only length-eligible candidate at old_v_len=6 is V1 - // (length 8). V0 is length 6 too — same length is eligible - // (trim_3=0), so RNG decides. Either way, the recorded - // pair must satisfy `len - trim_3 == 6`. - let recorded_id = match trace.find("receptor_revision.v_allele").unwrap().value { - ChoiceValue::AlleleId(id) => id, - _ => panic!("expected AlleleId record"), - }; - let recorded_trim = match trace.find("receptor_revision.v_trim_3").unwrap().value { - ChoiceValue::Int(t) => t, - _ => panic!("expected Int record"), - }; - let recorded_allele = cfg - .get(Segment::V, AlleleId::new(recorded_id)) - .expect("recorded allele must resolve"); - assert_eq!( - recorded_allele.len() as i64 - recorded_trim, - 6, - "retained length must equal old V region length" - ); - let _unused = v1; - - // Exactly one of each of the three replacement events. - let assignments = events - .iter() - .filter(|e| matches!(e, SimulationEvent::AssignmentChanged { segment: Segment::V, .. })) - .count(); - let trims = events - .iter() - .filter(|e| matches!(e, SimulationEvent::TrimChanged { segment: Segment::V, end: TrimEnd::Three, .. })) - .count(); - let replaces = events - .iter() - .filter(|e| matches!(e, SimulationEvent::SegmentReplaced { segment: Segment::V, .. })) - .count(); - assert_eq!(assignments, 1); - assert_eq!(trims, 1); - assert_eq!(replaces, 1); - - // The committed V assignment matches the recorded id. - assert_eq!( - after.assignments.get(Segment::V).unwrap().allele_id.index(), - recorded_id, - ); - // Pool length unchanged (same-length constraint). - assert_eq!(after.pool.len(), 6); - // The replacement bytes equal the recorded allele's - // 6-byte prefix. - let bases: Vec = after.pool.as_slice().iter().map(|n| n.base).collect(); - assert_eq!(bases, recorded_allele.seq[..6].to_vec()); - } - - // ── Replay ────────────────────────────────────────────────── - - #[test] - fn replay_applied_true_reproduces_replacement_without_rng() { - let (cfg, v0, v1) = two_v_refdata(); - // prob=0.0 would normally fire applied=false; the trace - // overrides via cursor. Pins "trace is the source of truth". - let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - - let mut input = Trace::new(); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionApplied, - ChoiceValue::Bool(true), - ); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionVAllele, - ChoiceValue::AlleleId(v1.index()), - ); - // V1 length 8, old V length 6 → trim_3 = 2. - input.record_choice( - address::ChoiceAddress::ReceptorRevisionVTrim3, - ChoiceValue::Int(2), - ); - let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); - - let (trace, after) = - run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None).unwrap(); - - assert_eq!( - trace.find("receptor_revision.applied").unwrap().value, - ChoiceValue::Bool(true), - ); - assert_eq!( - after.assignments.get(Segment::V).unwrap().allele_id, - v1, - ); - assert_eq!( - after.assignments.get(Segment::V).unwrap().trim_3, - 2, - ); - // V1's retained 6-byte prefix is "GGGGGG". - let bases: Vec = after.pool.as_slice().iter().map(|n| n.base).collect(); - assert_eq!(&bases, b"GGGGGG"); - assert!(cursor.is_drained()); - } - - #[test] - fn replay_applied_false_consumes_only_bool_record() { - let (cfg, v0, _) = two_v_refdata(); - // prob=1.0 would normally fire applied=true; trace says - // false → no allele/trim consumption. - let pass = ReceptorRevisionPass::new(1.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - - let mut input = Trace::new(); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionApplied, - ChoiceValue::Bool(false), - ); - let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); - - let (trace, after) = - run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None).unwrap(); - - assert_eq!( - trace.find("receptor_revision.applied").unwrap().value, - ChoiceValue::Bool(false), - ); - assert!(trace.find("receptor_revision.v_allele").is_none()); - // Unchanged. - assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v0); - assert!(cursor.is_drained()); - } - - #[test] - fn replay_missing_allele_record_after_applied_true_errors() { - let (cfg, v0, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - - let mut input = Trace::new(); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionApplied, - ChoiceValue::Bool(true), - ); - // Missing allele + trim records. - let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); - - let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None) - .unwrap_err(); - match err { - PassError::Replay { pass_name, reason } => { - assert_eq!(pass_name, "receptor_revision"); - let msg = format!("{reason}"); - assert!( - msg.contains("receptor_revision.v_allele") || msg.contains("exhausted"), - "expected replay error to mention v_allele or exhaustion, got: {msg}", - ); - } - other => panic!("expected PassError::Replay, got {other:?}"), - } - } - - #[test] - fn replay_unknown_allele_id_errors() { - let (cfg, v0, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - - let mut input = Trace::new(); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionApplied, - ChoiceValue::Bool(true), - ); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionVAllele, - ChoiceValue::AlleleId(999), - ); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionVTrim3, - ChoiceValue::Int(0), - ); - let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); - - let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None) - .unwrap_err(); - match err { - PassError::MissingAllele { pass_name, segment, allele_id } => { - assert_eq!(pass_name, "receptor_revision"); - assert_eq!(segment, Segment::V); - assert_eq!(allele_id, 999); - } - other => panic!("expected PassError::MissingAllele, got {other:?}"), - } - } - - #[test] - fn replay_trim_causing_length_mismatch_errors() { - let (cfg, v0, v1) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - - let mut input = Trace::new(); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionApplied, - ChoiceValue::Bool(true), - ); - input.record_choice( - address::ChoiceAddress::ReceptorRevisionVAllele, - ChoiceValue::AlleleId(v1.index()), - ); - // V1 has length 8, old V region is 6. A trim_3 of 3 would - // retain 5 bytes — mismatch. - input.record_choice( - address::ChoiceAddress::ReceptorRevisionVTrim3, - ChoiceValue::Int(3), - ); - let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); - - let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None) - .unwrap_err(); - match err { - PassError::InvalidPlanState { pass_name, reason } => { - assert_eq!(pass_name, "receptor_revision"); - assert!(reason.contains("length mismatch"), "got: {reason}"); - } - other => panic!("expected PassError::InvalidPlanState, got {other:?}"), - } - } - - // ── Plan-state guards ─────────────────────────────────────── - - #[test] - fn missing_v_assignment_errors_in_checked_path() { - let (cfg, _, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - let err = run_with_ctx(&pass, &cfg, None, Simulation::new(), None, None).unwrap_err(); - match err { - PassError::MissingAssignment { pass_name, segment } => { - assert_eq!(pass_name, "receptor_revision"); - assert_eq!(segment, Segment::V); - } - other => panic!("expected MissingAssignment, got {other:?}"), - } - } - - #[test] - fn missing_refdata_errors_in_checked_path() { - // Build a sim with V assigned but execute without refdata. - let (cfg, v0, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - - let mut trace = Trace::new(); - let mut rng = Rng::new(0); - let mut ctx = PassContext { - trace: &mut trace, - rng: &mut rng, - pass_index: 0, - refdata: None, - contracts: None, - feasibility: None, - reference_index: None, - replay_cursor: None, - event_log_sink: None, - }; - let err = pass - .execute_checked(&sim_v_assembled(v0), &mut ctx) - .unwrap_err(); - match err { - PassError::MissingRefData { pass_name } => { - assert_eq!(pass_name, "receptor_revision"); - } - other => panic!("expected MissingRefData, got {other:?}"), - } - } - - #[test] - fn no_v_region_errors_in_checked_path() { - let (cfg, v0, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - // Assignment without a region — assembly never ran. - let sim = Simulation::new().with_allele_assigned(Segment::V, AlleleInstance::new(v0)); - let err = run_with_ctx(&pass, &cfg, None, sim, None, None).unwrap_err(); - match err { - PassError::InvalidPlanState { pass_name, reason } => { - assert_eq!(pass_name, "receptor_revision"); - assert!(reason.contains("no V region"), "got: {reason}"); - } - other => panic!("expected InvalidPlanState, got {other:?}"), - } - } - - // ── Pass metadata ─────────────────────────────────────────── - - #[test] - fn declares_three_choice_patterns_in_order() { - let (cfg, _, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - assert_eq!( - pass.declared_choice_patterns(), - vec![ - address::ChoiceAddressPattern::ReceptorRevisionApplied, - address::ChoiceAddressPattern::ReceptorRevisionVAllele, - address::ChoiceAddressPattern::ReceptorRevisionVTrim3, - ], - ); - } - - #[test] - fn declares_refdata_and_v_assignment_requirements() { - let (cfg, _, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - assert_eq!( - pass.requirements(), - vec![ - PassRequirement::RefData, - PassRequirement::AlleleAssignment(Segment::V), - ], - ); - } - - #[test] - fn declares_no_compile_effects() { - // Pinning the empty declaration — receptor revision is - // event-driven, the schedule analyser must not treat it - // as initial recombination. - let (cfg, _, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - assert!(pass.effects().is_empty()); - } - - #[test] - fn pass_name_is_stable() { - // The frozen pass name flows through `pass_plan_signature` - // — a rename here breaks every existing trace. - let (cfg, _, _) = two_v_refdata(); - let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); - assert_eq!(pass.name(), "receptor_revision"); - } - - // ── Full plan round-trip (V assembled → replacement) ──────── - - #[test] - fn full_runtime_round_trip_with_prob_one() { - let (cfg, v0, _) = two_v_refdata(); - let (trace, after) = run_pass(1.0, 0, cfg.clone(), sim_v_assembled(v0)); - - // Three records present. - assert!(trace.find("receptor_revision.applied").is_some()); - assert!(trace.find("receptor_revision.v_allele").is_some()); - assert!(trace.find("receptor_revision.v_trim_3").is_some()); - - // V region length unchanged (same-length constraint). - let v_region = after - .sequence - .regions - .iter() - .find(|r| r.segment == Segment::V) - .expect("V region must remain after replacement"); - assert_eq!(v_region.len(), 6); - assert_eq!(after.pool.len(), 6); - } -} +mod tests; diff --git a/engine_rs/src/passes/receptor_revision/tests.rs b/engine_rs/src/passes/receptor_revision/tests.rs new file mode 100644 index 0000000..a61eed7 --- /dev/null +++ b/engine_rs/src/passes/receptor_revision/tests.rs @@ -0,0 +1,765 @@ + use super::*; + use crate::assignment::AlleleInstance; + use crate::contract::ContractSet; + use crate::dist::AllelePoolDist; + use crate::ir::{NucHandle, Region, Segment, SimulationEvent}; + use crate::pass::testing::PassRuntime; + use crate::pass::{PassError, PassPlan}; + use crate::refdata::{Allele, AllelePool, ChainType, RefDataConfig}; + use crate::replay::TraceCursor; + use crate::rng::Rng; + use crate::trace::Trace; + + fn allele(name: &str, seq: &[u8]) -> Allele { + Allele { + name: name.to_string(), + gene: name.split('*').next().unwrap_or(name).to_string(), + seq: seq.to_vec(), + segment: Segment::V, + anchor: None, + functional_status: None, + subregions: Vec::new(), + } + } + + /// Reference data with two V alleles whose retained 6-byte + /// prefix (under 0 trim) differs. The original assembled V is + /// `AAAAAA` matching allele V0 exactly; replacement candidate + /// V1 carries `GGGGGGCC` — length 8, retains 6 bytes at + /// `trim_3 = 2`. + fn two_v_refdata() -> (RefDataConfig, AlleleId, AlleleId) { + let mut cfg = RefDataConfig::empty(ChainType::Vdj); + let v0 = cfg.v_pool.push(allele("V1*01", b"AAAAAA")); + let v1 = cfg.v_pool.push(allele("V1*02", b"GGGGGGCC")); + (cfg, v0, v1) + } + + /// Build a sim with the V slot already assigned to V0 and a + /// V region of length 6 (`AAAAAA`). Doubles for tests as the + /// post-recombination starting state. + fn sim_v_assembled(v0: AlleleId) -> Simulation { + let mut sim = Simulation::new(); + for (i, &b) in b"AAAAAA".iter().enumerate() { + let (next, _) = sim.with_nucleotide_pushed(Nucleotide::germline( + b, + i as u16, + Segment::V, + )); + sim = next; + } + sim.with_allele_assigned(Segment::V, AlleleInstance::new(v0)) + .with_region_added(Region::new( + Segment::V, + NucHandle::new(0), + NucHandle::new(6), + )) + } + + fn v_only_pool(cfg: &RefDataConfig) -> AllelePool { + cfg.v_pool.clone() + } + + fn run_pass( + prob: f64, + seed: u64, + cfg: RefDataConfig, + sim: Simulation, + ) -> (Trace, Simulation) { + let mut plan = PassPlan::new(); + let pool = v_only_pool(&cfg); + plan.push(Box::new(ReceptorRevisionPass::new( + prob, + Box::new(AllelePoolDist::uniform(&pool)), + ))); + let outcome = PassRuntime::execute_with_refdata(&plan, sim, seed, &cfg); + let final_sim = outcome.final_simulation().clone(); + (outcome.trace, final_sim) + } + + fn run_with_ctx( + pass: &ReceptorRevisionPass, + cfg: &RefDataConfig, + contracts: Option<&ContractSet>, + initial: Simulation, + cursor: Option<&mut TraceCursor>, + event_sink: Option<&mut Vec>, + ) -> Result<(Trace, Simulation), PassError> { + let mut trace = Trace::new(); + let mut rng = Rng::new(0xc0ff_ee); + let mut ctx = PassContext { + trace: &mut trace, + rng: &mut rng, + pass_index: 0, + refdata: Some(cfg), + contracts, + feasibility: None, + reference_index: None, + replay_cursor: cursor, + event_log_sink: event_sink, + }; + let next = pass.execute_checked(&initial, &mut ctx)?; + Ok((trace, next)) + } + + // ── Construction guards ────────────────────────────────────── + + #[test] + #[should_panic(expected = "prob must be in [0.0, 1.0]")] + fn receptor_revision_rejects_out_of_range_prob() { + let (cfg, _, _) = two_v_refdata(); + let _ = ReceptorRevisionPass::new( + 1.5, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ); + } + + #[test] + #[should_panic(expected = "prob must be in [0.0, 1.0]")] + fn receptor_revision_rejects_negative_prob() { + let (cfg, _, _) = two_v_refdata(); + let _ = ReceptorRevisionPass::new( + -0.5, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ); + } + + // ── genotype-aware candidate selection ────────────────────── + + fn constraint(per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> GenotypeVConstraint { + GenotypeVConstraint { per_hap, same_haplotype: same } + } + + fn geno_sim(v0: AlleleId) -> Simulation { + sim_v_assembled(v0) + .with_allele_assigned(Segment::V, AlleleInstance::new(v0).with_haplotype(0)) + } + + fn geno_pass(cfg: &RefDataConfig, per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> ReceptorRevisionPass { + ReceptorRevisionPass::new(1.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) + .with_genotype_constraint(constraint(per_hap, same)) + } + + #[test] + fn genotype_same_haplotype_excludes_current_and_restricts_to_chromosome() { + let (cfg, v0, v1) = two_v_refdata(); + // hap0 carries {V0, V1}; hap1 carries {V0}. Current is V0 on hap0. + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let (trace, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap(); + assert_eq!(trace.find("receptor_revision.applied").unwrap().value, ChoiceValue::Bool(true)); + // exclude-current: the only eligible alternate on hap0 is V1 + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); + assert_eq!(after.assignments.get(Segment::V).unwrap().receptor_revision_original_id, Some(v0)); + assert_eq!(after.assignments.get(Segment::V).unwrap().haplotype, Some(0)); + } + + fn run_permissive(pass: &ReceptorRevisionPass, cfg: &RefDataConfig, initial: Simulation) -> (Trace, Simulation) { + let mut trace = Trace::new(); + let mut rng = Rng::new(0xc0ff_ee); + let mut ctx = PassContext { + trace: &mut trace, + rng: &mut rng, + pass_index: 0, + refdata: Some(cfg), + contracts: None, + feasibility: None, + reference_index: None, + replay_cursor: None, + event_log_sink: None, + }; + let next = pass.execute(&initial, &mut ctx); + (trace, next) + } + + #[test] + fn genotype_no_eligible_alternate_permissive_applied_false() { + let (cfg, v0, _v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); + let (trace, after) = run_permissive(&pass, &cfg, geno_sim(v0)); + assert_eq!(trace.find("receptor_revision.applied").unwrap().value, ChoiceValue::Bool(false)); + assert!(trace.find("receptor_revision.v_allele").is_none()); + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v0); + } + + #[test] + fn genotype_no_eligible_alternate_strict_errors() { + let (cfg, v0, _v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap_err(); + assert!(matches!(err, PassError::ConstraintSampling { .. }), "got {err:?}"); + } + + #[test] + fn genotype_both_haplotypes_admits_other_chromosome_allele() { + let (cfg, v0, v1) = two_v_refdata(); + // V1 carried only on hap1; same_haplotype=false aggregates both. + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v1, 1.0)]], false); + let (_t, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), None, None).unwrap(); + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); + // haplotype provenance preserved as the original rearrangement chromosome (0) + assert_eq!(after.assignments.get(Segment::V).unwrap().haplotype, Some(0)); + } + + #[test] + fn genotype_missing_haplotype_stamp_errors() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + // sim_v_assembled assigns V0 WITHOUT a haplotype stamp. + let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), None, None).unwrap_err(); + assert!(matches!(err, PassError::InvalidPlanState { .. }), "got {err:?}"); + } + + #[test] + fn genotype_strict_post_event_contract_rejection_errors() { + use crate::contract::{Contract, ContractViolation}; + + // A contract whose verify() always rejects the post-event state. + struct RejectAll; + impl Contract for RejectAll { + fn name(&self) -> &str { + "reject_all_test" + } + fn verify( + &self, + _sim: &Simulation, + _refdata: Option<&RefDataConfig>, + ) -> Result<(), ContractViolation> { + Err(ContractViolation::new(self.name(), "rejected by test")) + } + } + + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let contracts = ContractSet::new().with(Box::new(RejectAll)); + // Strict mode (execute_checked) with a rejecting contract must surface a + // contract violation — parity with the non-genotype path. + let err = + run_with_ctx(&pass, &cfg, Some(&contracts), geno_sim(v0), None, None).unwrap_err(); + assert!(matches!(err, PassError::ContractViolation { .. }), "got {err:?}"); + } + + #[test] + fn genotype_replay_strict_post_event_contract_rejection_errors() { + use crate::contract::{Contract, ContractViolation}; + + struct RejectAll; + impl Contract for RejectAll { + fn name(&self) -> &str { + "reject_all_test" + } + fn verify( + &self, + _sim: &Simulation, + _refdata: Option<&RefDataConfig>, + ) -> Result<(), ContractViolation> { + Err(ContractViolation::new(self.name(), "rejected by test")) + } + } + + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let contracts = ContractSet::new().with(Box::new(RejectAll)); + // A valid applied=true replay (v1, trim 2) must STILL be rejected by the + // post-event contract in strict mode — replay parity with the fresh path. + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); + let err = run_with_ctx(&pass, &cfg, Some(&contracts), geno_sim(v0), Some(&mut cursor), None) + .unwrap_err(); + assert!(matches!(err, PassError::ContractViolation { .. }), "got {err:?}"); + } + + // ── genotype-aware replay validation ──────────────────────── + + fn replay_records(applied: bool, allele: Option, trim: Option) -> Vec { + let mut t = Trace::new(); + t.record_choice(address::ChoiceAddress::ReceptorRevisionApplied, ChoiceValue::Bool(applied)); + if let Some(a) = allele { + t.record_choice(address::ChoiceAddress::ReceptorRevisionVAllele, ChoiceValue::AlleleId(a)); + } + if let Some(tr) = trim { + t.record_choice(address::ChoiceAddress::ReceptorRevisionVTrim3, ChoiceValue::Int(tr)); + } + t.choices().to_vec() + } + + #[test] + fn replay_genotype_valid_replacement_reproduces() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); + let (_t, after) = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap(); + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v1); + assert!(cursor.is_drained()); + } + + #[test] + fn replay_genotype_allele_not_carried_on_haplotype_errors() { + let (cfg, v0, v1) = two_v_refdata(); + // hap0 carries only V0; recorded V1 is not carried on the drawn chromosome. + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0), (v1, 1.0)]], true); + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(2))); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); + assert!(matches!(err, PassError::InvalidDistributionOutput { .. }), "got {err:?}"); + } + + #[test] + fn replay_genotype_unresolvable_allele_id_reports_missing_allele() { + let (cfg, v0, _v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0)], vec![(v0, 1.0)]], true); + // out-of-range recorded v_allele -> missing_allele (not "not_carried"), + // matching the non-genotype replay diagnostic contract. + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(999), Some(0))); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); + match err { + PassError::MissingAllele { allele_id, .. } => assert_eq!(allele_id, 999), + other => panic!("expected MissingAllele, got {other:?}"), + } + } + + #[test] + fn replay_genotype_equals_current_allele_errors() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v0.index()), Some(0))); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); + assert!(matches!(err, PassError::InvalidDistributionOutput { .. }), "got {err:?}"); + } + + #[test] + fn replay_genotype_trim_length_mismatch_errors() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = geno_pass(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true); + // V1 len 8, old 6 -> trim must be 2; record 3 (retains 5) => mismatch. + let mut cursor = TraceCursor::from_owned(replay_records(true, Some(v1.index()), Some(3))); + let err = run_with_ctx(&pass, &cfg, None, geno_sim(v0), Some(&mut cursor), None).unwrap_err(); + assert!(matches!(err, PassError::InvalidPlanState { .. }), "got {err:?}"); + } + + #[test] + fn genotype_signature_differs_by_candidate_set_and_same_haplotype() { + let (cfg, v0, v1) = two_v_refdata(); + let no_geno = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) + .parameter_signature(); + let g_true = geno_pass_prob(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], true).parameter_signature(); + let g_false = geno_pass_prob(&cfg, [vec![(v0, 1.0), (v1, 1.0)], vec![(v0, 1.0)]], false).parameter_signature(); + let g_other = geno_pass_prob(&cfg, [vec![(v0, 2.0)], vec![(v1, 1.0)]], true).parameter_signature(); + assert_ne!(g_true, no_geno); + assert_ne!(g_true, g_false); + assert_ne!(g_true, g_other); + } + + fn geno_pass_prob(cfg: &RefDataConfig, per_hap: [Vec<(AlleleId, f64)>; 2], same: bool) -> ReceptorRevisionPass { + ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))) + .with_genotype_constraint(constraint(per_hap, same)) + } + + // ── prob=0: no replacement ────────────────────────────────── + + #[test] + fn prob_zero_records_applied_false_and_no_mutation_events() { + let (cfg, v0, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + let mut events: Vec = Vec::new(); + let (trace, after) = run_with_ctx( + &pass, + &cfg, + None, + sim_v_assembled(v0), + None, + Some(&mut events), + ) + .unwrap(); + + // Applied=false recorded. + let rec = trace + .find("receptor_revision.applied") + .expect("applied Bool must be recorded even when prob = 0"); + assert_eq!(rec.value, ChoiceValue::Bool(false)); + // No allele/trim records. + assert!(trace.find("receptor_revision.v_allele").is_none()); + assert!(trace.find("receptor_revision.v_trim_3").is_none()); + + // No state-changing events. + assert!(events.iter().all(|e| !matches!( + e, + SimulationEvent::AssignmentChanged { .. } + | SimulationEvent::TrimChanged { .. } + | SimulationEvent::SegmentReplaced { .. } + ))); + + // V assignment + pool bytes unchanged. + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v0); + let bases: Vec = after.pool.as_slice().iter().map(|n| n.base).collect(); + assert_eq!(&bases, b"AAAAAA"); + } + + // ── prob=1: replacement ───────────────────────────────────── + + #[test] + fn prob_one_records_three_choices_and_emits_three_events() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(1.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + let mut events: Vec = Vec::new(); + let (trace, after) = run_with_ctx( + &pass, + &cfg, + None, + sim_v_assembled(v0), + None, + Some(&mut events), + ) + .unwrap(); + + // applied = true + assert_eq!( + trace.find("receptor_revision.applied").unwrap().value, + ChoiceValue::Bool(true), + ); + // The only length-eligible candidate at old_v_len=6 is V1 + // (length 8). V0 is length 6 too — same length is eligible + // (trim_3=0), so RNG decides. Either way, the recorded + // pair must satisfy `len - trim_3 == 6`. + let recorded_id = match trace.find("receptor_revision.v_allele").unwrap().value { + ChoiceValue::AlleleId(id) => id, + _ => panic!("expected AlleleId record"), + }; + let recorded_trim = match trace.find("receptor_revision.v_trim_3").unwrap().value { + ChoiceValue::Int(t) => t, + _ => panic!("expected Int record"), + }; + let recorded_allele = cfg + .get(Segment::V, AlleleId::new(recorded_id)) + .expect("recorded allele must resolve"); + assert_eq!( + recorded_allele.len() as i64 - recorded_trim, + 6, + "retained length must equal old V region length" + ); + let _unused = v1; + + // Exactly one of each of the three replacement events. + let assignments = events + .iter() + .filter(|e| matches!(e, SimulationEvent::AssignmentChanged { segment: Segment::V, .. })) + .count(); + let trims = events + .iter() + .filter(|e| matches!(e, SimulationEvent::TrimChanged { segment: Segment::V, end: TrimEnd::Three, .. })) + .count(); + let replaces = events + .iter() + .filter(|e| matches!(e, SimulationEvent::SegmentReplaced { segment: Segment::V, .. })) + .count(); + assert_eq!(assignments, 1); + assert_eq!(trims, 1); + assert_eq!(replaces, 1); + + // The committed V assignment matches the recorded id. + assert_eq!( + after.assignments.get(Segment::V).unwrap().allele_id.index(), + recorded_id, + ); + // Pool length unchanged (same-length constraint). + assert_eq!(after.pool.len(), 6); + // The replacement bytes equal the recorded allele's + // 6-byte prefix. + let bases: Vec = after.pool.as_slice().iter().map(|n| n.base).collect(); + assert_eq!(bases, recorded_allele.seq[..6].to_vec()); + } + + // ── Replay ────────────────────────────────────────────────── + + #[test] + fn replay_applied_true_reproduces_replacement_without_rng() { + let (cfg, v0, v1) = two_v_refdata(); + // prob=0.0 would normally fire applied=false; the trace + // overrides via cursor. Pins "trace is the source of truth". + let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + + let mut input = Trace::new(); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionApplied, + ChoiceValue::Bool(true), + ); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionVAllele, + ChoiceValue::AlleleId(v1.index()), + ); + // V1 length 8, old V length 6 → trim_3 = 2. + input.record_choice( + address::ChoiceAddress::ReceptorRevisionVTrim3, + ChoiceValue::Int(2), + ); + let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); + + let (trace, after) = + run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None).unwrap(); + + assert_eq!( + trace.find("receptor_revision.applied").unwrap().value, + ChoiceValue::Bool(true), + ); + assert_eq!( + after.assignments.get(Segment::V).unwrap().allele_id, + v1, + ); + assert_eq!( + after.assignments.get(Segment::V).unwrap().trim_3, + 2, + ); + // V1's retained 6-byte prefix is "GGGGGG". + let bases: Vec = after.pool.as_slice().iter().map(|n| n.base).collect(); + assert_eq!(&bases, b"GGGGGG"); + assert!(cursor.is_drained()); + } + + #[test] + fn replay_applied_false_consumes_only_bool_record() { + let (cfg, v0, _) = two_v_refdata(); + // prob=1.0 would normally fire applied=true; trace says + // false → no allele/trim consumption. + let pass = ReceptorRevisionPass::new(1.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + + let mut input = Trace::new(); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionApplied, + ChoiceValue::Bool(false), + ); + let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); + + let (trace, after) = + run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None).unwrap(); + + assert_eq!( + trace.find("receptor_revision.applied").unwrap().value, + ChoiceValue::Bool(false), + ); + assert!(trace.find("receptor_revision.v_allele").is_none()); + // Unchanged. + assert_eq!(after.assignments.get(Segment::V).unwrap().allele_id, v0); + assert!(cursor.is_drained()); + } + + #[test] + fn replay_missing_allele_record_after_applied_true_errors() { + let (cfg, v0, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + + let mut input = Trace::new(); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionApplied, + ChoiceValue::Bool(true), + ); + // Missing allele + trim records. + let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); + + let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None) + .unwrap_err(); + match err { + PassError::Replay { pass_name, reason } => { + assert_eq!(pass_name, "receptor_revision"); + let msg = format!("{reason}"); + assert!( + msg.contains("receptor_revision.v_allele") || msg.contains("exhausted"), + "expected replay error to mention v_allele or exhaustion, got: {msg}", + ); + } + other => panic!("expected PassError::Replay, got {other:?}"), + } + } + + #[test] + fn replay_unknown_allele_id_errors() { + let (cfg, v0, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + + let mut input = Trace::new(); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionApplied, + ChoiceValue::Bool(true), + ); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionVAllele, + ChoiceValue::AlleleId(999), + ); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionVTrim3, + ChoiceValue::Int(0), + ); + let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); + + let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None) + .unwrap_err(); + match err { + PassError::MissingAllele { pass_name, segment, allele_id } => { + assert_eq!(pass_name, "receptor_revision"); + assert_eq!(segment, Segment::V); + assert_eq!(allele_id, 999); + } + other => panic!("expected PassError::MissingAllele, got {other:?}"), + } + } + + #[test] + fn replay_trim_causing_length_mismatch_errors() { + let (cfg, v0, v1) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.0, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + + let mut input = Trace::new(); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionApplied, + ChoiceValue::Bool(true), + ); + input.record_choice( + address::ChoiceAddress::ReceptorRevisionVAllele, + ChoiceValue::AlleleId(v1.index()), + ); + // V1 has length 8, old V region is 6. A trim_3 of 3 would + // retain 5 bytes — mismatch. + input.record_choice( + address::ChoiceAddress::ReceptorRevisionVTrim3, + ChoiceValue::Int(3), + ); + let mut cursor = TraceCursor::from_owned(input.choices().to_vec()); + + let err = run_with_ctx(&pass, &cfg, None, sim_v_assembled(v0), Some(&mut cursor), None) + .unwrap_err(); + match err { + PassError::InvalidPlanState { pass_name, reason } => { + assert_eq!(pass_name, "receptor_revision"); + assert!(reason.contains("length mismatch"), "got: {reason}"); + } + other => panic!("expected PassError::InvalidPlanState, got {other:?}"), + } + } + + // ── Plan-state guards ─────────────────────────────────────── + + #[test] + fn missing_v_assignment_errors_in_checked_path() { + let (cfg, _, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + let err = run_with_ctx(&pass, &cfg, None, Simulation::new(), None, None).unwrap_err(); + match err { + PassError::MissingAssignment { pass_name, segment } => { + assert_eq!(pass_name, "receptor_revision"); + assert_eq!(segment, Segment::V); + } + other => panic!("expected MissingAssignment, got {other:?}"), + } + } + + #[test] + fn missing_refdata_errors_in_checked_path() { + // Build a sim with V assigned but execute without refdata. + let (cfg, v0, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + + let mut trace = Trace::new(); + let mut rng = Rng::new(0); + let mut ctx = PassContext { + trace: &mut trace, + rng: &mut rng, + pass_index: 0, + refdata: None, + contracts: None, + feasibility: None, + reference_index: None, + replay_cursor: None, + event_log_sink: None, + }; + let err = pass + .execute_checked(&sim_v_assembled(v0), &mut ctx) + .unwrap_err(); + match err { + PassError::MissingRefData { pass_name } => { + assert_eq!(pass_name, "receptor_revision"); + } + other => panic!("expected MissingRefData, got {other:?}"), + } + } + + #[test] + fn no_v_region_errors_in_checked_path() { + let (cfg, v0, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + // Assignment without a region — assembly never ran. + let sim = Simulation::new().with_allele_assigned(Segment::V, AlleleInstance::new(v0)); + let err = run_with_ctx(&pass, &cfg, None, sim, None, None).unwrap_err(); + match err { + PassError::InvalidPlanState { pass_name, reason } => { + assert_eq!(pass_name, "receptor_revision"); + assert!(reason.contains("no V region"), "got: {reason}"); + } + other => panic!("expected InvalidPlanState, got {other:?}"), + } + } + + // ── Pass metadata ─────────────────────────────────────────── + + #[test] + fn declares_three_choice_patterns_in_order() { + let (cfg, _, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + assert_eq!( + pass.declared_choice_patterns(), + vec![ + address::ChoiceAddressPattern::ReceptorRevisionApplied, + address::ChoiceAddressPattern::ReceptorRevisionVAllele, + address::ChoiceAddressPattern::ReceptorRevisionVTrim3, + ], + ); + } + + #[test] + fn declares_refdata_and_v_assignment_requirements() { + let (cfg, _, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + assert_eq!( + pass.requirements(), + vec![ + PassRequirement::RefData, + PassRequirement::AlleleAssignment(Segment::V), + ], + ); + } + + #[test] + fn declares_no_compile_effects() { + // Pinning the empty declaration — receptor revision is + // event-driven, the schedule analyser must not treat it + // as initial recombination. + let (cfg, _, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + assert!(pass.effects().is_empty()); + } + + #[test] + fn pass_name_is_stable() { + // The frozen pass name flows through `pass_plan_signature` + // — a rename here breaks every existing trace. + let (cfg, _, _) = two_v_refdata(); + let pass = ReceptorRevisionPass::new(0.5, Box::new(AllelePoolDist::uniform(&cfg.v_pool))); + assert_eq!(pass.name(), "receptor_revision"); + } + + // ── Full plan round-trip (V assembled → replacement) ──────── + + #[test] + fn full_runtime_round_trip_with_prob_one() { + let (cfg, v0, _) = two_v_refdata(); + let (trace, after) = run_pass(1.0, 0, cfg.clone(), sim_v_assembled(v0)); + + // Three records present. + assert!(trace.find("receptor_revision.applied").is_some()); + assert!(trace.find("receptor_revision.v_allele").is_some()); + assert!(trace.find("receptor_revision.v_trim_3").is_some()); + + // V region length unchanged (same-length constraint). + let v_region = after + .sequence + .regions + .iter() + .find(|r| r.segment == Segment::V) + .expect("V region must remain after replacement"); + assert_eq!(v_region.len(), 6); + assert_eq!(after.pool.len(), 6); + } diff --git a/engine_rs/src/passes/tests.rs b/engine_rs/src/passes/tests.rs new file mode 100644 index 0000000..a069cd0 --- /dev/null +++ b/engine_rs/src/passes/tests.rs @@ -0,0 +1,1352 @@ + use super::*; + use crate::assignment::TrimEnd; + use crate::contract::{ + productive, ChoiceContext, Contract, ContractSet, ContractViolation, + ProductiveJunctionFrame, + }; + use crate::dist::{ + AllelePoolDist, Distribution, EmpiricalLengthDist, FilteredSampleError, UniformBase, + }; + use crate::ir::{NucHandle, Segment, Simulation}; + use crate::pass::testing::PassRuntime; + use crate::pass::{Pass, PassPlan}; + use crate::passes::sample_allele::test_support::make_test_pool; + use crate::s5f::{S5FKernel, S5F_NUM_CONTEXTS, S5F_SUBSTITUTION_LEN}; + use crate::trace::ChoiceValue; + + fn fixed_count_dist() -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])) + } + + fn test_s5f_kernel() -> S5FKernel { + S5FKernel::new( + vec![1.0; S5F_NUM_CONTEXTS], + vec![0.25; S5F_SUBSTITUTION_LEN], + ) + } + + fn assert_typed_patterns_project_to_declared_choices(pass: &dyn Pass) { + let projected: Vec = pass + .declared_choice_patterns() + .into_iter() + .map(String::from) + .collect(); + assert_eq!( + projected, + pass.declared_choices(), + "{} typed declared-choice patterns must project to the stable string surface", + pass.name() + ); + } + + #[test] + fn built_in_pass_declared_choice_patterns_match_string_surface() { + let v_pool = make_test_pool(2, Segment::V); + let d_pool = make_test_pool(2, Segment::D); + let j_pool = make_test_pool(2, Segment::J); + + let passes: Vec> = vec![ + Box::new(EchoPass::new(b'A', 0, Segment::V)), + Box::new(AssembleSegmentPass::new(Segment::V)), + Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&v_pool)), + )), + Box::new(SampleAllelePass::new( + Segment::D, + Box::new(AllelePoolDist::uniform(&d_pool)), + )), + Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&j_pool)), + )), + Box::new(TrimPass::new(Segment::V, TrimEnd::Five, fixed_count_dist())), + Box::new(TrimPass::new( + Segment::D, + TrimEnd::Three, + fixed_count_dist(), + )), + Box::new(TrimPass::new(Segment::J, TrimEnd::Five, fixed_count_dist())), + Box::new(GenerateNPPass::new( + Segment::Np1, + fixed_count_dist(), + Box::new(UniformBase), + )), + Box::new(GenerateNPPass::new( + Segment::Np2, + fixed_count_dist(), + Box::new(UniformBase), + )), + Box::new(UniformMutationPass::new( + fixed_count_dist(), + Box::new(UniformBase), + )), + Box::new(S5FMutationPass::new(test_s5f_kernel(), fixed_count_dist())), + Box::new(PCRErrorPass::new(fixed_count_dist(), Box::new(UniformBase))), + Box::new(QualityErrorPass::new( + fixed_count_dist(), + Box::new(UniformBase), + )), + Box::new(ContaminantPass::new(0.5, Box::new(UniformBase))), + Box::new(IndelPass::new( + fixed_count_dist(), + 0.5, + Box::new(UniformBase), + )), + Box::new(NCorruptionPass::new(fixed_count_dist())), + Box::new(EndLossPass::new(LossEnd::Five, fixed_count_dist())), + Box::new(EndLossPass::new(LossEnd::Three, fixed_count_dist())), + Box::new(RevCompPass::new(0.5)), + ]; + + for pass in passes { + assert_typed_patterns_project_to_declared_choices(pass.as_ref()); + } + } + + /// Build a synthetic VJ refdata: V "AAACCCGGG" (9bp, anchor 6), + /// J "TTTAAA" (6bp, anchor 0). Junction = V_anchor_to_end (3bp) + /// + NP1 + J_anchor_to_W3 (3bp). For productive frame: + /// (3 + NP1 + 3) % 3 == 0 → NP1 % 3 == 0. + fn make_vj_refdata_for_filter() -> crate::refdata::RefDataConfig { + let mut cfg = crate::refdata::RefDataConfig::empty(crate::refdata::ChainType::Vj); + let _ = cfg.v_pool.push(crate::refdata::Allele { + name: "v_test*01".into(), + gene: "v_test".into(), + seq: b"AAACCCGGG".to_vec(), + segment: Segment::V, + anchor: Some(6), + functional_status: None, + subregions: Vec::new(), + }); + let _ = cfg.j_pool.push(crate::refdata::Allele { + name: "j_test*01".into(), + gene: "j_test".into(), + seq: b"TTTAAA".to_vec(), + segment: Segment::J, + anchor: Some(0), + functional_status: None, + subregions: Vec::new(), + }); + cfg + } + + #[test] + fn productive_admits_filters_np1_for_vj_chain() { + let cfg = make_vj_refdata_for_filter(); + let dist = EmpiricalLengthDist::from_pairs((0..7).map(|i| (i, 1.0)).collect::>()); + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(dist), + Box::new(UniformBase), + ))); + + let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); + + for seed in 0..20u64 { + let outcome = PassRuntime::execute_with_context( + &plan, + Simulation::new(), + seed, + Some(&cfg), + Some(&contracts), + ); + let np1_len = match outcome.trace.find("np.np1.length").unwrap().value { + ChoiceValue::Int(n) => n, + _ => panic!("wrong variant"), + }; + assert!( + np1_len % 3 == 0, + "seed {} produced NP1 length {} (not divisible by 3)", + seed, + np1_len + ); + } + } + + #[test] + fn productive_admits_unconstrained_without_contracts() { + let cfg = make_vj_refdata_for_filter(); + let dist = EmpiricalLengthDist::from_pairs((0..7).map(|i| (i, 1.0)).collect::>()); + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(dist), + Box::new(UniformBase), + ))); + + // Without contracts: any value in [0, 6] is allowed. + let mut seen_out_of_frame = false; + for seed in 0..50u64 { + let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), seed, &cfg); + let np1_len = match outcome.trace.find("np.np1.length").unwrap().value { + ChoiceValue::Int(n) => n, + _ => panic!("wrong variant"), + }; + assert!(np1_len >= 0 && np1_len <= 6); + if np1_len % 3 != 0 { + seen_out_of_frame = true; + } + } + assert!( + seen_out_of_frame, + "Without contracts, expected at least one out-of-frame NP1 sample" + ); + } + + #[test] + fn productive_admits_empty_filter_returns_explicit_zero_length() { + // v3.0 rule: when natural ∩ admissible is empty under + // active contracts, permissive mode must NOT fall back + // to an unconstrained draw from the natural support + // (that would re-introduce reject-after-propose). The + // architectural no-op for a length sampler is 0 — same + // shape as `TrimPass::sample_trim`'s empty-support + // fallback. (Pre-v3.0 the pass fell through to + // `length_dist.sample(rng)` and produced one of + // {1,2,4,5}, all of which violate frame in this fixture.) + let cfg = make_vj_refdata_for_filter(); + let dist = EmpiricalLengthDist::from_pairs(vec![(1, 1.0), (2, 1.0), (4, 1.0), (5, 1.0)]); + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(dist), + Box::new(UniformBase), + ))); + + let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); + + let outcome = PassRuntime::execute_with_context( + &plan, + Simulation::new(), + 0, + Some(&cfg), + Some(&contracts), + ); + let np1_len = match outcome.trace.find("np.np1.length").unwrap().value { + ChoiceValue::Int(n) => n, + _ => panic!("wrong variant"), + }; + assert_eq!(np1_len, 0); + } + + #[derive(Clone, Debug)] + struct UnenumerableLengthDist; + + impl Distribution for UnenumerableLengthDist { + type Output = i64; + + fn sample(&self, _rng: &mut crate::rng::Rng) -> i64 { + 0 + } + } + + #[test] + fn productive_strict_errors_when_length_filter_empty() { + let cfg = make_vj_refdata_for_filter(); + let dist = EmpiricalLengthDist::from_pairs(vec![(1, 1.0), (2, 1.0), (4, 1.0), (5, 1.0)]); + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(dist), + Box::new(UniformBase), + ))); + + let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); + let err = PassRuntime::execute_strict_with_context( + &plan, + Simulation::new(), + 0, + Some(&cfg), + Some(&contracts), + ) + .unwrap_err(); + + assert_eq!(err.pass_name(), "generate_np.np1"); + assert_eq!(err.address(), "np.np1.length"); + assert_eq!( + err.constraint_reason(), + Some(FilteredSampleError::EmptyAdmissibleSupport) + ); + assert!(err.to_string().contains("no admissible candidates")); + } + + #[test] + fn productive_strict_errors_when_length_support_unavailable() { + let cfg = make_vj_refdata_for_filter(); + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(UnenumerableLengthDist), + Box::new(UniformBase), + ))); + + let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); + let err = PassRuntime::execute_strict_with_context( + &plan, + Simulation::new(), + 0, + Some(&cfg), + Some(&contracts), + ) + .unwrap_err(); + + assert_eq!(err.pass_name(), "generate_np.np1"); + assert_eq!(err.address(), "np.np1.length"); + assert_eq!( + err.constraint_reason(), + Some(FilteredSampleError::SupportUnavailable) + ); + } + + struct RejectNpBases; + + impl Contract for RejectNpBases { + fn name(&self) -> &str { + "reject_np_bases" + } + + fn verify( + &self, + _sim: &Simulation, + _refdata: Option<&crate::refdata::RefDataConfig>, + ) -> Result<(), ContractViolation> { + Ok(()) + } + + // Post-flip the trait's primary surface is `admits_typed`, + // so this test contract dispatches on `ChoiceAddress::NpBase` + // directly rather than prefix-matching the legacy string. + fn admits_typed( + &self, + _sim: &Simulation, + _refdata: Option<&crate::refdata::RefDataConfig>, + context: ChoiceContext<'_>, + _candidate: &ChoiceValue, + ) -> Result<(), ContractViolation> { + if matches!( + context.address, + Some(crate::address::ChoiceAddress::NpBase { + segment: crate::address::NpSegment::Np1, + .. + }) + ) { + return Err(ContractViolation::new(self.name(), "rejected by test")); + } + Ok(()) + } + } + + #[test] + fn productive_admits_empty_base_filter_writes_n_sentinel_in_permissive() { + // v3.0 rule (mirror of `sample_length`'s explicit-zero + // fallback): when the bundle leaves no admissible base + // for an NP slot, permissive mode writes the IUPAC `N` + // sentinel rather than falling back to an unconstrained + // draw. `N` translates to amino acid `X`, which is not + // a stop and not bound by anchor preservation (NP slots + // are between V/J anchor codons), so it satisfies the + // productive bundle's full admit check. + // + // The matching strict test below (`productive_strict_errors_…`) + // pins that strict mode surfaces the same empty support + // as `EmptyAdmissibleSupport` instead. + let mut plan = PassPlan::new(); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])), + Box::new(UniformBase), + ))); + let contracts = ContractSet::new().with(Box::new(RejectNpBases)); + + let outcome = + PassRuntime::execute_with_context(&plan, Simulation::new(), 0, None, Some(&contracts)); + + // The trace's recorded base for NP slot 0 is `b'N'`. + assert_eq!( + outcome.trace.find("np.np1.bases[0]").unwrap().value, + ChoiceValue::Base(b'N') + ); + // The pool reflects the same sentinel. + let final_sim = outcome.final_simulation(); + assert_eq!(final_sim.pool.len(), 1); + assert_eq!(final_sim.pool.get(NucHandle::new(0)).unwrap().base, b'N'); + } + + #[test] + fn productive_strict_errors_when_base_filter_empty() { + let mut plan = PassPlan::new(); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])), + Box::new(UniformBase), + ))); + let contracts = ContractSet::new().with(Box::new(RejectNpBases)); + + let err = PassRuntime::execute_strict_with_context( + &plan, + Simulation::new(), + 0, + None, + Some(&contracts), + ) + .unwrap_err(); + + assert_eq!(err.pass_name(), "generate_np.np1"); + assert_eq!(err.address(), "np.np1.bases[0]"); + assert_eq!( + err.constraint_reason(), + Some(FilteredSampleError::EmptyAdmissibleSupport) + ); + } + + #[test] + fn productive_strict_succeeds_when_admissible_length_exists() { + let cfg = make_vj_refdata_for_filter(); + let dist = EmpiricalLengthDist::from_pairs((0..7).map(|i| (i, 1.0)).collect::>()); + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(dist), + Box::new(UniformBase), + ))); + + let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); + let outcome = PassRuntime::execute_strict_with_context( + &plan, + Simulation::new(), + 0, + Some(&cfg), + Some(&contracts), + ) + .unwrap(); + let np1_len = match outcome.trace.find("np.np1.length").unwrap().value { + ChoiceValue::Int(n) => n, + _ => panic!("wrong variant"), + }; + + assert_eq!(np1_len % 3, 0); + } + + #[test] + fn productive_admits_makes_full_pipeline_in_frame_for_vj() { + use crate::junction::compute_junction; + + let cfg = make_vj_refdata_for_filter(); + let dist = EmpiricalLengthDist::from_pairs((0..7).map(|i| (i, 1.0)).collect::>()); + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(dist), + Box::new(UniformBase), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); + + let contracts = ContractSet::new().with(Box::new(ProductiveJunctionFrame::new())); + + for seed in 0..30u64 { + let outcome = PassRuntime::execute_with_context( + &plan, + Simulation::new(), + seed, + Some(&cfg), + Some(&contracts), + ); + let junction = compute_junction(outcome.final_simulation(), &cfg) + .expect("junction should be defined"); + assert!( + junction.is_in_frame(), + "seed {} produced out-of-frame junction (length {})", + seed, + junction.length + ); + } + } + + #[test] + fn productive_full_bundle_in_frame_and_admits_dispatch_works() { + use crate::junction::compute_junction; + + let cfg = make_vj_refdata_for_filter(); + let dist = EmpiricalLengthDist::from_pairs((0..10).map(|i| (i, 1.0)).collect::>()); + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + Box::new(dist), + Box::new(UniformBase), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); + + let contracts = productive(); + + for seed in 0..20u64 { + let outcome = PassRuntime::execute_with_context( + &plan, + Simulation::new(), + seed, + Some(&cfg), + Some(&contracts), + ); + let junction = compute_junction(outcome.final_simulation(), &cfg) + .expect("junction should be defined"); + assert!(junction.is_in_frame()); + } + } + + #[test] + fn sample_allele_pass_replay_in_mixed_plan_with_echo_passes() { + // Mixed plan: SampleAllele + Echo (transform) interleaved. + // The trace records only the sampling choices; the IR + // accumulates both the assignment and the echoed nucleotides. + let pool = make_test_pool(3, Segment::V); + let mut plan = PassPlan::new(); + plan.push(Box::new(EchoPass::new(b'A', 0, Segment::V))); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&pool)), + ))); + plan.push(Box::new(EchoPass::new(b'C', 1, Segment::V))); + + let oa = PassRuntime::execute(&plan, Simulation::new(), 0xfeed); + let ob = PassRuntime::execute(&plan, Simulation::new(), 0xfeed); + + assert_eq!( + oa.final_simulation().pool.len(), + ob.final_simulation().pool.len() + ); + assert_eq!( + oa.final_simulation() + .assignments + .get(Segment::V) + .copied() + .unwrap() + .allele_id, + ob.final_simulation() + .assignments + .get(Segment::V) + .copied() + .unwrap() + .allele_id + ); + // Trace contains exactly one entry — only the sampling pass records. + assert_eq!(oa.trace.len(), 1); + assert_eq!(oa.trace.choices()[0].address, "sample_allele.v"); + } + + /// A representative mixed plan: alternating EchoPass (no RNG) + /// and SampleBasePass (RNG-consuming). Used by the cross-cutting + /// replay-determinism test below. + fn mixed_plan() -> PassPlan { + use crate::ir::flag; + let mut plan = PassPlan::new(); + for i in 0..8 { + plan.push(Box::new(EchoPass::new(b'A', i as u16, Segment::V))); + plan.push(Box::new(SampleBasePass::new( + format!("np.np1.bases[{}]", i), + Box::new(UniformBase), + Segment::Np1, + flag::N_NUC, + ))); + } + plan + } + + #[test] + fn replay_determinism_holds_under_repeated_runs() { + // Run the same plan with the same seed five times. Every + // trace and every final IR must be identical to the first. + let baseline = PassRuntime::execute(&mixed_plan(), Simulation::new(), 0xfeed_face); + for _ in 0..5 { + let other = PassRuntime::execute(&mixed_plan(), Simulation::new(), 0xfeed_face); + assert_eq!(other.trace.len(), baseline.trace.len()); + for (a, b) in baseline + .trace + .choices() + .iter() + .zip(other.trace.choices().iter()) + { + assert_eq!(a.address, b.address); + assert_eq!(a.value, b.value); + } + assert_eq!( + other.final_simulation().pool.len(), + baseline.final_simulation().pool.len() + ); + } + } + + // ────────────────────────────────────────────────────────────── + // Per-pass `EventRecord.simulation_events` outcome contract + // + // Every committed pass now carries two complementary surfaces: + // - `trace_span` → choices the pass consumed/emitted + // - `simulation_events` → state consequences of those choices + // + // These tests pin which events each pass produces, indexed off + // the committed `outcome.events` ledger — i.e. observability + // from the *outcome* surface, not the test-only PassContext + // hook. + // ────────────────────────────────────────────────────────────── + + /// Find exactly-one event matching `pat` inside `events`. Fails + /// the test if zero or more-than-one match. Returns the match. + fn find_one( + events: &[crate::ir::SimulationEvent], + pat: F, + ) -> &crate::ir::SimulationEvent + where + F: Fn(&crate::ir::SimulationEvent) -> bool, + { + let matches: Vec<&crate::ir::SimulationEvent> = + events.iter().filter(|e| pat(e)).collect(); + assert_eq!( + matches.len(), + 1, + "expected exactly one matching event, got {} from {:?}", + matches.len(), + events + ); + matches[0] + } + + #[test] + fn outcome_event_records_carry_assignment_trim_and_region_events() { + // Plan: sample V → trim V 5' → assemble V → generate NP1 → + // assemble J. Each committed pass's `EventRecord` should + // carry the consequence event it actually fired. + let cfg = make_vj_refdata_for_filter(); + let trim_dist = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(0, 1.0)])) + }; + let np_dist = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(3, 1.0)])) + }; + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(TrimPass::new( + Segment::V, + TrimEnd::Five, + trim_dist(), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + np_dist(), + Box::new(UniformBase), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); + + let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 42, &cfg); + + // Sanity: pass count matches plan length. + assert_eq!(outcome.events.len(), plan.len()); + + // Pass 0 — SampleAllele V → one AssignmentChanged for V. + let ev0 = &outcome.events[0]; + assert_eq!(ev0.pass_name, "sample_allele.v"); + let assign = find_one(&ev0.simulation_events, |e| { + matches!( + e, + crate::ir::SimulationEvent::AssignmentChanged { + segment: Segment::V, + .. + } + ) + }); + match assign { + crate::ir::SimulationEvent::AssignmentChanged { segment, old, .. } => { + assert_eq!(*segment, Segment::V); + assert!(old.is_none(), "first-install AssignmentChanged has old=None"); + } + _ => unreachable!(), + } + + // Pass 2 — TrimPass V 5' → one TrimChanged for (V, Five). + let ev2 = &outcome.events[2]; + assert_eq!(ev2.pass_name, "trim.v_5"); + let trim_event = find_one(&ev2.simulation_events, |e| { + matches!( + e, + crate::ir::SimulationEvent::TrimChanged { + segment: Segment::V, + end: TrimEnd::Five, + .. + } + ) + }); + match trim_event { + crate::ir::SimulationEvent::TrimChanged { new, .. } => { + assert_eq!(*new, 0, "trim distribution fixed at 0"); + } + _ => unreachable!(), + } + + // Pass 3 — AssembleSegment V → one RegionAdded for V. + let ev3 = &outcome.events[3]; + assert_eq!(ev3.pass_name, "assemble.v"); + let v_region = find_one(&ev3.simulation_events, |e| { + matches!( + e, + crate::ir::SimulationEvent::RegionAdded { region } if region.segment == Segment::V + ) + }); + match v_region { + crate::ir::SimulationEvent::RegionAdded { region } => { + // V allele is 9 bases long with trim_5=0 → region [0, 9). + assert_eq!(region.start, NucHandle::new(0)); + assert_eq!(region.end, NucHandle::new(9)); + } + _ => unreachable!(), + } + + // Pass 4 — GenerateNP Np1 → one RegionAdded for Np1. + let ev4 = &outcome.events[4]; + assert_eq!(ev4.pass_name, "generate_np.np1"); + let np_region = find_one(&ev4.simulation_events, |e| { + matches!( + e, + crate::ir::SimulationEvent::RegionAdded { region } if region.segment == Segment::Np1 + ) + }); + match np_region { + crate::ir::SimulationEvent::RegionAdded { region } => { + // NP1 starts where V ended (handle 9), length 3. + assert_eq!(region.start, NucHandle::new(9)); + assert_eq!(region.end, NucHandle::new(12)); + } + _ => unreachable!(), + } + + // Pass 5 — AssembleSegment J → one RegionAdded for J. + let ev5 = &outcome.events[5]; + assert_eq!(ev5.pass_name, "assemble.j"); + let _ = find_one(&ev5.simulation_events, |e| { + matches!( + e, + crate::ir::SimulationEvent::RegionAdded { region } if region.segment == Segment::J + ) + }); + } + + #[test] + fn outcome_event_records_carry_base_changed_and_mutation_count_for_mutation_pass() { + use crate::passes::UniformMutationPass; + // Plan: sample V → assemble V → uniform-mutation 1. + // The mutation pass's EventRecord must carry exactly one + // `BaseChanged` (the single applied substitution) and + // exactly one `MutationCountChanged` (the commit-time count + // bump). + let cfg = make_vj_refdata_for_filter(); + let mutation_count_dist = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])) + }; + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(UniformMutationPass::new( + mutation_count_dist(), + Box::new(UniformBase), + ))); + + let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 42, &cfg); + + // The mutation pass is the third committed pass. + let mut_event = outcome + .events + .iter() + .find(|e| e.pass_name == "mutate.uniform") + .expect("uniform mutation pass committed an event record"); + + let _ = find_one(&mut_event.simulation_events, |e| { + matches!(e, crate::ir::SimulationEvent::BaseChanged { .. }) + }); + let mc = find_one(&mut_event.simulation_events, |e| { + matches!(e, crate::ir::SimulationEvent::MutationCountChanged { .. }) + }); + match mc { + crate::ir::SimulationEvent::MutationCountChanged { + old, + new, + delta, + } => { + assert_eq!(*old, 0); + assert_eq!(*new, 1); + assert_eq!(*delta, 1); + } + _ => unreachable!(), + } + + // Final mutation_count carried on the sealed sim matches the + // event's `new` field — pin the outcome-ledger / sim + // consistency contract. + assert_eq!(outcome.final_simulation().mutation_count, 1); + } + + // ────────────────────────────────────────────────────────────── + // Event-coverage audit + invariant tests + // + // Three coordinated tests that lock down the consequence-event + // ledger before any derived-state lifecycle refactor: + // + // 1. Cross-pass coverage — every meaningful pass emits the + // consequence variants it owns. + // 2. Event/state consistency — outcome counts reconstructed + // from `simulation_events` match the final `Simulation`. + // 3. Trace/event alignment — `TraceSpan`s are contiguous, + // cover the full trace, and replay equality holds. + // ────────────────────────────────────────────────────────────── + + /// Find every event matching `pat` in a slice. Used in the + /// audit tests where we care about "at least one of X" and + /// counts, not exact-one-of-X. + fn count_matching(events: &[crate::ir::SimulationEvent], pat: F) -> usize + where + F: Fn(&crate::ir::SimulationEvent) -> bool, + { + events.iter().filter(|e| pat(e)).count() + } + + /// Build the heavy-stack VJ plan used by the coverage audit. + /// Every category the user listed (sample/trim/assemble/np/ + /// mutation/pcr/quality/ncorrupt/indel/rev-comp) is represented + /// by exactly one pass. Distributions are pinned to fixed + /// counts so the audit is deterministic across RNG paths. + fn build_coverage_audit_plan( + cfg: &crate::refdata::RefDataConfig, + ) -> PassPlan { + use crate::passes::{ + IndelPass, NCorruptionPass, PCRErrorPass, QualityErrorPass, RevCompPass, + UniformMutationPass, + }; + let count_one = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])) + }; + let count_zero = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(0, 1.0)])) + }; + let np_three = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(3, 1.0)])) + }; + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(TrimPass::new( + Segment::V, + TrimEnd::Five, + count_zero(), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + np_three(), + Box::new(UniformBase), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); + plan.push(Box::new(UniformMutationPass::new( + count_one(), + Box::new(UniformBase), + ))); + plan.push(Box::new(PCRErrorPass::new(count_one(), Box::new(UniformBase)))); + plan.push(Box::new(QualityErrorPass::new( + count_one(), + Box::new(UniformBase), + ))); + plan.push(Box::new(NCorruptionPass::new(count_one()))); + plan.push(Box::new(IndelPass::new( + count_one(), + 0.5, + Box::new(UniformBase), + ))); + plan.push(Box::new(RevCompPass::new(1.0))); + plan + } + + #[test] + fn cross_pass_event_coverage_emits_expected_consequence_types() { + // Drive a heavy-stack VJ plan covering every event-emitting + // pass category. Assert each pass's `EventRecord.simulation_events` + // carries at least one event of the expected variant — this + // is the load-bearing "no holes in the ledger" check before + // we lean on it for lifecycle decisions. + use crate::ir::SimulationEvent; + + let cfg = make_vj_refdata_for_filter(); + let plan = build_coverage_audit_plan(&cfg); + let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 42, &cfg); + + // Convenience: locate a committed pass's event record by + // its `pass_name`. Panics if missing because the plan is + // hard-coded — any missing pass is a regression. + let find = |name: &str| { + outcome + .events + .iter() + .find(|e| e.pass_name == name) + .unwrap_or_else(|| panic!("missing event record for {name}")) + }; + + // Sample allele: AssignmentChanged for V and J. + assert!( + count_matching(&find("sample_allele.v").simulation_events, |e| matches!( + e, + SimulationEvent::AssignmentChanged { segment: Segment::V, .. } + )) >= 1 + ); + assert!( + count_matching(&find("sample_allele.j").simulation_events, |e| matches!( + e, + SimulationEvent::AssignmentChanged { segment: Segment::J, .. } + )) >= 1 + ); + + // Trim: TrimChanged. + assert!( + count_matching(&find("trim.v_5").simulation_events, |e| matches!( + e, + SimulationEvent::TrimChanged { segment: Segment::V, .. } + )) >= 1 + ); + + // Assemble: many BasePushed + one RegionAdded per segment. + for (name, seg) in [("assemble.v", Segment::V), ("assemble.j", Segment::J)] { + let ev = find(name); + assert!( + count_matching(&ev.simulation_events, |e| matches!(e, SimulationEvent::BasePushed { .. })) > 0, + "{name} should emit BasePushed events" + ); + assert!( + count_matching(&ev.simulation_events, |e| matches!( + e, + SimulationEvent::RegionAdded { region } if region.segment == seg + )) == 1, + "{name} should emit exactly one RegionAdded for its segment" + ); + } + + // NP: BasePushed (length 3) + one RegionAdded for Np1. + let np = find("generate_np.np1"); + assert!( + count_matching(&np.simulation_events, |e| matches!(e, SimulationEvent::BasePushed { .. })) >= 1 + ); + assert!( + count_matching(&np.simulation_events, |e| matches!( + e, + SimulationEvent::RegionAdded { region } if region.segment == Segment::Np1 + )) == 1 + ); + + // Uniform mutation: BaseChanged + MutationCountChanged. + let mu = find("mutate.uniform"); + assert!( + count_matching(&mu.simulation_events, |e| matches!(e, SimulationEvent::BaseChanged { .. })) >= 1 + ); + assert!( + count_matching(&mu.simulation_events, |e| matches!( + e, + SimulationEvent::MutationCountChanged { .. } + )) == 1 + ); + + // PCR / quality / ncorrupt: BaseChanged (no mutation-count + // bump — these are sequencing artifacts, not biological + // mutations, and don't tag `add_to_mutation_count`). + for name in ["corrupt.pcr", "corrupt.quality", "corrupt.ns"] { + let ev = find(name); + assert!( + count_matching(&ev.simulation_events, |e| matches!(e, SimulationEvent::BaseChanged { .. })) >= 1, + "{name} should emit at least one BaseChanged" + ); + assert!( + count_matching(&ev.simulation_events, |e| matches!( + e, + SimulationEvent::MutationCountChanged { .. } + )) == 0, + "{name} is a sequencing-artifact pass and must not bump n_mutations" + ); + } + + // Indel: at least one of IndelInserted / IndelDeleted. + let indel = find("corrupt.indel"); + let indels = count_matching(&indel.simulation_events, |e| { + matches!( + e, + SimulationEvent::IndelInserted { .. } | SimulationEvent::IndelDeleted { .. } + ) + }); + assert!(indels >= 1, "corrupt.indel should emit at least one indel event"); + + // Rev-comp: one ReverseComplementFlagRecorded with applied=true + // (apply_prob = 1.0 in the audit plan). + let rc = find("corrupt.rev_comp"); + assert_eq!( + count_matching(&rc.simulation_events, |e| matches!( + e, + SimulationEvent::ReverseComplementFlagRecorded { applied: true } + )), + 1 + ); + } + + #[test] + fn outcome_events_are_consistent_with_final_simulation_state() { + // Reconstruct three coarse counts purely from the + // `simulation_events` stream and assert they match the + // sealed final `Simulation`. This is the load-bearing + // self-consistency proof: a future derived-state observer + // can trust the ledger as the source of truth. + use crate::ir::SimulationEvent; + + let cfg = make_vj_refdata_for_filter(); + let count_one = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(1, 1.0)])) + }; + let np_three = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(3, 1.0)])) + }; + + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + np_three(), + Box::new(UniformBase), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); + plan.push(Box::new(crate::passes::UniformMutationPass::new( + count_one(), + Box::new(UniformBase), + ))); + + let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 0, &cfg); + let final_sim = outcome.final_simulation(); + + // ── Invariant 1: regions ──────────────────────────────── + // Indel-free plan → final region count equals the total + // `RegionAdded` events across all passes. + let region_adds: usize = outcome + .events + .iter() + .flat_map(|e| &e.simulation_events) + .filter(|e| matches!(e, SimulationEvent::RegionAdded { .. })) + .count(); + assert_eq!( + region_adds, + final_sim.sequence.regions.len(), + "sum of RegionAdded events must equal final region count for indel-free plan" + ); + + // ── Invariant 2: assignments ──────────────────────────── + // The last `AssignmentChanged` per segment must equal the + // segment's final `AlleleInstance` in `final_sim.assignments`. + use crate::assignment::AlleleInstance; + let mut last_by_segment: std::collections::HashMap = + std::collections::HashMap::new(); + for ev in outcome.events.iter().flat_map(|e| &e.simulation_events) { + if let SimulationEvent::AssignmentChanged { segment, new, .. } = ev { + last_by_segment.insert(*segment, *new); + } + } + for (seg, last_assignment) in &last_by_segment { + assert_eq!( + final_sim.assignments.get(*seg).copied(), + Some(*last_assignment), + "last AssignmentChanged for {:?} must equal final assignment", + seg + ); + } + // V and J should both have been assigned. + assert!(last_by_segment.contains_key(&Segment::V)); + assert!(last_by_segment.contains_key(&Segment::J)); + + // ── Invariant 3: mutation_count ───────────────────────── + // The last `MutationCountChanged.new` across the run must + // equal `final_sim.mutation_count`. (Plan has exactly one + // count-bumping pass; this still holds with multiple.) + let final_mc = outcome + .events + .iter() + .flat_map(|e| &e.simulation_events) + .filter_map(|e| match e { + SimulationEvent::MutationCountChanged { new, .. } => Some(*new), + _ => None, + }) + .last() + .expect("plan has one mutation pass"); + assert_eq!(final_mc, final_sim.mutation_count); + } + + #[test] + fn trace_spans_partition_outcome_trace_and_replay_preserves_trace() { + // Two-part alignment proof: + // + // (a) Every committed pass yields one `EventRecord` whose + // `trace_span` is contiguous with its neighbours and + // whose union covers the entire outcome trace. No + // gaps, no overlaps. + // + // (b) Replaying the captured trace through `CompiledSimulator` + // produces an outcome whose trace matches the original + // record-for-record — proving event capture is + // transparent w.r.t. the trace, the only persisted + // artifact. + use crate::compiled::ExecutionPolicy; + + let cfg = make_vj_refdata_for_filter(); + let np_three = || -> Box> { + Box::new(EmpiricalLengthDist::from_pairs(vec![(3, 1.0)])) + }; + let mut plan = PassPlan::new(); + plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + np_three(), + Box::new(UniformBase), + ))); + plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); + + // ── (a) Trace-span partition ───────────────────────────── + let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 7, &cfg); + assert_eq!( + outcome.events.len(), + plan.len(), + "every committed pass must contribute exactly one EventRecord" + ); + assert_eq!(outcome.events[0].trace_span.start, 0); + for i in 1..outcome.events.len() { + assert_eq!( + outcome.events[i - 1].trace_span.end, + outcome.events[i].trace_span.start, + "trace spans must be contiguous (gap or overlap at boundary {})", + i + ); + } + assert_eq!( + outcome + .events + .last() + .map(|e| e.trace_span.end) + .unwrap_or(0), + outcome.trace.len(), + "the union of all trace spans must cover the full outcome trace" + ); + + // ── (b) Replay equality ────────────────────────────────── + // Compile the same plan + refdata into an owned simulator, + // replay the captured trace, and assert the replayed + // trace matches the original record-for-record. Event + // capture happens on both runs (the compiled path always + // supplies a sink) but neither side persists events — so a + // matching trace is sufficient proof that capture is + // transparent. + // + // `OwnedCompiledSimulator::compile` takes ownership of the + // plan + refdata, so build a fresh equivalent plan for the + // replay leg. (The borrowing `CompiledSimulator` doesn't + // expose `replay_from_trace_records`.) + let mut replay_plan = PassPlan::new(); + replay_plan.push(Box::new(SampleAllelePass::new( + Segment::V, + Box::new(AllelePoolDist::uniform(&cfg.v_pool)), + ))); + replay_plan.push(Box::new(SampleAllelePass::new( + Segment::J, + Box::new(AllelePoolDist::uniform(&cfg.j_pool)), + ))); + replay_plan.push(Box::new(AssembleSegmentPass::new(Segment::V))); + replay_plan.push(Box::new(GenerateNPPass::new( + Segment::Np1, + np_three(), + Box::new(UniformBase), + ))); + replay_plan.push(Box::new(AssembleSegmentPass::new(Segment::J))); + let compiled = crate::compiled::OwnedCompiledSimulator::compile_with_options( + replay_plan, + Some(cfg.clone()), + None, + ExecutionPolicy::Permissive, + crate::compiled::CompileOptions::skip_refdata_validation(), + ) + .expect("audit plan should compile"); + let original_records: Vec<_> = outcome.trace.choices().to_vec(); + let replayed = compiled + .replay_from_trace_records( + &original_records, + /* seed unused for migrated sites */ 0, + ExecutionPolicy::Strict, + ) + .expect("replay should succeed against the same plan"); + let replayed_records: Vec<_> = replayed.trace.choices().to_vec(); + assert_eq!( + original_records.len(), + replayed_records.len(), + "replay must produce a trace of the same length" + ); + for (i, (orig, repl)) in original_records.iter().zip(replayed_records.iter()).enumerate() { + assert_eq!(orig.address, repl.address, "address divergence at record {}", i); + assert_eq!(orig.value, repl.value, "value divergence at record {} (addr={})", i, orig.address); + } + } + + // ────────────────────────────────────────────────────────────── + // Compile-effect / event-emission policy conformance + // + // Built-in mutating passes that declare a `PassCompileEffect` + // are expected to emit at least one corresponding + // `SimulationEvent` when they actually mutate state. This test + // exercises the heavy-stack VJ plan and asserts every + // `EventRecord` passes the policy check in + // [`crate::event::check_event_emission_consistency`]. + // + // It is the regression net for the "new pass mutates Simulation + // directly via sim.with_* and forgets to emit events" anti- + // pattern. The runtime live-call refresh trusts events; a + // silent zero-event mutating pass would skip the refresh + // (this is also what the divergence tests in + // `compiled::tests::live_call_edits` demonstrate by construction). + // ────────────────────────────────────────────────────────────── + + #[test] + fn builtin_passes_emit_events_consistent_with_declared_compile_effects() { + let cfg = make_vj_refdata_for_filter(); + let plan = build_coverage_audit_plan(&cfg); + let outcome = PassRuntime::execute_with_refdata(&plan, Simulation::new(), 42, &cfg); + + // Apply the policy to every committed pass. The plan's + // distributions are pinned to deterministic non-zero + // counts so the EditBases / StructuralIndel zero- + // exemption never fires here — meaningful coverage. + let mut violations: Vec = Vec::new(); + for record in &outcome.events { + if let Err(err) = crate::event::check_event_emission_consistency(record) { + violations.push(format!( + "{}: declared {:?} but {}", + err.pass_name, err.compile_effect, err.reason + )); + } + } + assert!( + violations.is_empty(), + "built-in passes must emit events matching their compile-effect declarations.\n\ + Violations:\n {}\n\n\ + Likely cause: the pass mutates `Simulation` via `sim.with_*` directly \ + instead of routing through `SimulationBuilder` — runtime derived-state \ + refresh trusts events, so a silent mutator gets no refresh.", + violations.join("\n ") + ); + } diff --git a/engine_rs/src/refdata.rs b/engine_rs/src/refdata.rs index cb97b75..0157893 100644 --- a/engine_rs/src/refdata.rs +++ b/engine_rs/src/refdata.rs @@ -716,593 +716,7 @@ fn tag_identity_source(identity: &mut ReferenceIdentity, tag: &str) { // ────────────────────────────────────────────────────────────────── #[cfg(test)] -mod tests { - use super::*; - use std::mem::size_of; - - /// Helper: a tiny synthetic V allele for tests. - fn make_v(name: &str, gene: &str, seq: &[u8], anchor: Option) -> Allele { - Allele { - name: name.to_string(), - gene: gene.to_string(), - seq: seq.to_vec(), - segment: Segment::V, - anchor, - functional_status: None, - subregions: Vec::new(), - } - } - - #[test] - fn allele_id_is_zero_cost_newtype() { - assert_eq!(size_of::(), size_of::()); - } - - #[test] - fn allele_id_round_trip() { - let id = AlleleId::new(7); - assert_eq!(id.index(), 7); - assert_eq!(id.as_usize(), 7); - } - - #[test] - fn chain_type_has_d_distinction() { - assert!(!ChainType::Vj.has_d()); - assert!(ChainType::Vdj.has_d()); - } - - #[test] - fn allele_basic_accessors() { - let a = make_v("IGHV1-2*01", "IGHV1-2", b"ACGTACGT", Some(3)); - assert_eq!(a.len(), 8); - assert!(a.has_anchor()); - assert_eq!(a.anchor, Some(3)); - assert_eq!(a.segment, Segment::V); - } - - #[test] - fn allele_anchorless_round_trip() { - let a = make_v("IGHV-pseudo*01", "IGHV-pseudo", b"ACGT", None); - assert!(!a.has_anchor()); - assert_eq!(a.anchor, None); - } - - #[test] - fn allele_pool_starts_empty() { - let p = AllelePool::new(); - assert_eq!(p.len(), 0); - assert!(p.is_empty()); - assert!(p.get(AlleleId::new(0)).is_none()); - assert!(p.find_by_name("nonexistent").is_none()); - } - - #[test] - fn allele_pool_push_returns_sequential_ids() { - let mut p = AllelePool::new(); - let id0 = p.push(make_v("a*01", "a", b"AA", Some(0))); - let id1 = p.push(make_v("b*01", "b", b"CC", Some(0))); - let id2 = p.push(make_v("c*01", "c", b"GG", None)); - - assert_eq!(id0.index(), 0); - assert_eq!(id1.index(), 1); - assert_eq!(id2.index(), 2); - assert_eq!(p.len(), 3); - } - - #[test] - fn allele_pool_get_returns_stored_allele() { - let mut p = AllelePool::new(); - let id = p.push(make_v("x*01", "x", b"ATGC", Some(2))); - let got = p.get(id).expect("just pushed"); - assert_eq!(got.name, "x*01"); - assert_eq!(got.seq, b"ATGC"); - assert_eq!(got.anchor, Some(2)); - } - - #[test] - fn allele_pool_get_out_of_bounds_returns_none() { - let p = AllelePool::new(); - assert!(p.get(AlleleId::new(99)).is_none()); - } - - #[test] - fn allele_pool_iter_yields_id_allele_pairs_in_order() { - let mut p = AllelePool::new(); - let _ = p.push(make_v("a*01", "a", b"AA", None)); - let _ = p.push(make_v("b*01", "b", b"CC", None)); - let _ = p.push(make_v("c*01", "c", b"GG", None)); - - let collected: Vec<(u32, String)> = p - .iter() - .map(|(id, a)| (id.index(), a.name.clone())) - .collect(); - assert_eq!( - collected, - vec![ - (0, "a*01".to_string()), - (1, "b*01".to_string()), - (2, "c*01".to_string()), - ] - ); - } - - #[test] - fn allele_pool_find_by_name_locates_exact_match() { - let mut p = AllelePool::new(); - let _ = p.push(make_v("IGHV1-2*01", "IGHV1-2", b"AA", None)); - let target_id = p.push(make_v("IGHV1-2*02", "IGHV1-2", b"AC", None)); - let _ = p.push(make_v("IGHV3-23*01", "IGHV3-23", b"GG", None)); - - let (id, allele) = p.find_by_name("IGHV1-2*02").expect("name should exist"); - assert_eq!(id, target_id); - assert_eq!(allele.seq, b"AC"); - - // Partial / similar names should not match. - assert!(p.find_by_name("IGHV1-2").is_none()); - assert!(p.find_by_name("IGHV1-2*03").is_none()); - } - - #[test] - fn ref_data_config_empty_for_chain_type() { - let cfg = RefDataConfig::empty(ChainType::Vdj); - assert_eq!(cfg.chain_type, ChainType::Vdj); - assert!(cfg.v_pool.is_empty()); - assert!(cfg.d_pool.is_empty()); - assert!(cfg.j_pool.is_empty()); - assert!(cfg.c_pool.is_empty()); - } - - #[test] - fn ref_data_config_pool_for_segment_routes_correctly() { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - let _ = cfg.v_pool.push(make_v("v*01", "v", b"AA", None)); - let _ = cfg.d_pool.push(Allele { - name: "d*01".into(), - gene: "d".into(), - seq: b"GG".to_vec(), - segment: Segment::D, - anchor: None, - functional_status: None, - subregions: Vec::new(), - }); - let _ = cfg.j_pool.push(Allele { - name: "j*01".into(), - gene: "j".into(), - seq: b"TT".to_vec(), - segment: Segment::J, - anchor: Some(0), - functional_status: None, - subregions: Vec::new(), - }); - - assert_eq!(cfg.pool_for(Segment::V).unwrap().len(), 1); - assert_eq!(cfg.pool_for(Segment::D).unwrap().len(), 1); - assert_eq!(cfg.pool_for(Segment::J).unwrap().len(), 1); - assert!(cfg.pool_for(Segment::Np1).is_none()); - assert!(cfg.pool_for(Segment::Np2).is_none()); - } - - #[test] - fn ref_data_config_get_resolves_segment_id_pair() { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - let v_id = cfg.v_pool.push(make_v("v*01", "v", b"AAAT", Some(1))); - - let v = cfg.get(Segment::V, v_id).expect("v*01 should resolve"); - assert_eq!(v.name, "v*01"); - assert_eq!(v.anchor, Some(1)); - - // Wrong segment -> None. - assert!(cfg.get(Segment::J, v_id).is_none()); - // NP segment -> None defensively. - assert!(cfg.get(Segment::Np1, v_id).is_none()); - } - - // ── Curation policy ─────────────────────────────────────────── - - fn allele_with_anchor(name: &str, seq: &[u8], anchor: Option, segment: Segment) -> Allele { - Allele { - name: name.to_string(), - gene: name.split('*').next().unwrap_or(name).to_string(), - seq: seq.to_vec(), - segment, - anchor, - functional_status: None, - subregions: Vec::new(), - } - } - - fn vj_cfg_for_curation() -> RefDataConfig { - // Two V alleles: one Cys-anchored, one Gly-anchored. - // Two J alleles: one W-anchored, one anchorless. - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg - .v_pool - .push(allele_with_anchor("MYV-good*01", b"TGTAAACCC", Some(0), Segment::V)); - let _ = cfg - .v_pool - .push(allele_with_anchor("MYV-gly*01", b"GGGAAACCC", Some(0), Segment::V)); - let _ = cfg - .j_pool - .push(allele_with_anchor("MYJ-good*01", b"TGGAAACCC", Some(0), Segment::J)); - let _ = cfg - .j_pool - .push(allele_with_anchor("MYJ-orphan*01", b"GGG", None, Segment::J)); - cfg - } - - #[test] - fn curation_raw_is_identity() { - let cfg = vj_cfg_for_curation(); - let curated = cfg.curated(RefDataCurationPolicy::Raw); - assert_eq!(curated.v_pool.len(), 2); - assert_eq!(curated.j_pool.len(), 2); - // Raw policy does NOT touch identity. - assert_eq!(curated.identity.source, None); - } - - #[test] - fn curation_functional_anchors_only_filters_v_and_j() { - let cfg = vj_cfg_for_curation(); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - assert_eq!(curated.v_pool.len(), 1); - assert_eq!(curated.v_pool.iter().next().unwrap().1.name, "MYV-good*01"); - assert_eq!(curated.j_pool.len(), 1); - assert_eq!(curated.j_pool.iter().next().unwrap().1.name, "MYJ-good*01"); - } - - #[test] - fn curation_tags_identity_source() { - let cfg = vj_cfg_for_curation(); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - assert_eq!( - curated.identity.source.as_deref(), - Some("curated:functional_anchors_only"), - ); - } - - #[test] - fn curation_extends_existing_identity_source() { - let mut cfg = vj_cfg_for_curation(); - cfg.identity.source = Some("DataConfig".to_string()); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - assert_eq!( - curated.identity.source.as_deref(), - Some("DataConfig|curated:functional_anchors_only"), - ); - } - - #[test] - fn curation_preserves_other_identity_fields() { - let mut cfg = vj_cfg_for_curation(); - cfg.identity.species = Some("Human".into()); - cfg.identity.locus = Some("IGK".into()); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - assert_eq!(curated.identity.species.as_deref(), Some("Human")); - assert_eq!(curated.identity.locus.as_deref(), Some("IGK")); - } - - #[test] - fn curation_does_not_silently_fix_duplicate_names() { - // Add a duplicate V allele AFTER the good one. Both are - // Cys-anchored — they survive curation, then validation - // surfaces the duplicate as a Fatal issue. Curation never - // removes a structurally-corrupt entry to make a catalogue - // look healthier. - let mut cfg = vj_cfg_for_curation(); - let _ = cfg.v_pool.push(allele_with_anchor( - "MYV-good*01", - b"TGTAAACCC", - Some(0), - Segment::V, - )); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - let issues = curated.validate(); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::DuplicateAlleleName { .. } - ))); - } - - #[test] - fn curation_does_not_silently_fix_invalid_bytes() { - // Add an allele with a gap byte ('.') and a valid Cys anchor. - // The invalid byte is Fatal and must still surface post-curation. - let mut cfg = vj_cfg_for_curation(); - let _ = cfg.v_pool.push(allele_with_anchor( - "MYV-badbyte*01", - b"TGT.AAACC", - Some(0), - Segment::V, - )); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - assert!(curated.validate().iter().any(|i| matches!( - i, - RefDataValidationIssue::InvalidAlleleByte { .. } - ))); - } - - #[test] - fn curation_to_empty_pool_surfaces_empty_required_pool() { - // Every V allele is non-Cys — curation empties V. - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg - .v_pool - .push(allele_with_anchor("MYV-gly*01", b"GGGAAACCC", Some(0), Segment::V)); - let _ = cfg - .j_pool - .push(allele_with_anchor("MYJ-good*01", b"TGGAAACCC", Some(0), Segment::J)); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - assert_eq!(curated.v_pool.len(), 0); - let issues = curated.validate(); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::EmptyRequiredPool { segment: Segment::V } - ))); - } - - #[test] - fn curation_d_pool_passes_through_unchanged() { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - let _ = cfg - .v_pool - .push(allele_with_anchor("MYV*01", b"TGTAAACCC", Some(0), Segment::V)); - let _ = cfg.d_pool.push(allele_with_anchor( - "MYD*01", - b"GGGCCCAAA", - None, - Segment::D, - )); - let _ = cfg - .j_pool - .push(allele_with_anchor("MYJ*01", b"TGGAAACCC", Some(0), Segment::J)); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - // Anchorless D survives — D pool is never filtered. - assert_eq!(curated.d_pool.len(), 1); - } - - #[test] - fn curation_respects_custom_j_anchor_rule() { - // Rule expects Y at J anchor. TGG (W) and TTC (F) are dropped; - // TAT (Y) is kept. - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.j_anchor.expected_amino_acids = vec!['Y']; - let _ = cfg - .v_pool - .push(allele_with_anchor("MYV*01", b"TGTAAACCC", Some(0), Segment::V)); - let _ = cfg - .j_pool - .push(allele_with_anchor("MYJ-w*01", b"TGGAAACCC", Some(0), Segment::J)); - let _ = cfg - .j_pool - .push(allele_with_anchor("MYJ-y*01", b"TATAAACCC", Some(0), Segment::J)); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - assert_eq!(curated.j_pool.len(), 1); - assert_eq!(curated.j_pool.iter().next().unwrap().1.name, "MYJ-y*01"); - } - - #[test] - fn curation_anchor_required_false_keeps_anchorless_alleles() { - // If the rule says anchor is optional, anchorless alleles - // are kept under FunctionalAnchorsOnly. - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.j_anchor.required = false; - let _ = cfg - .v_pool - .push(allele_with_anchor("MYV*01", b"TGTAAACCC", Some(0), Segment::V)); - let _ = cfg - .j_pool - .push(allele_with_anchor("MYJ-orphan*01", b"GGG", None, Segment::J)); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); - assert_eq!(curated.j_pool.len(), 1); - } - - // ── Functional-status curation ──────────────────────────────── - - fn allele_with_status( - name: &str, - seq: &[u8], - segment: Segment, - anchor: Option, - status: Option, - ) -> Allele { - Allele { - name: name.to_string(), - gene: name.split('*').next().unwrap_or(name).to_string(), - seq: seq.to_vec(), - segment, - anchor, - functional_status: status, - subregions: Vec::new(), - } - } - - fn mixed_status_cfg() -> RefDataConfig { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - let _ = cfg.v_pool.push(allele_with_status( - "v-f*01", b"TGTAAACCC", Segment::V, Some(0), - Some(FunctionalStatus::Functional), - )); - let _ = cfg.v_pool.push(allele_with_status( - "v-o*01", b"TGTAAACCC", Segment::V, Some(0), - Some(FunctionalStatus::Orf), - )); - let _ = cfg.v_pool.push(allele_with_status( - "v-p*01", b"TGTAAACCC", Segment::V, Some(0), - Some(FunctionalStatus::Pseudogene), - )); - let _ = cfg.v_pool.push(allele_with_status( - "v-na*01", b"TGTAAACCC", Segment::V, Some(0), None, - )); - let _ = cfg.d_pool.push(allele_with_status( - "d-f*01", b"GGG", Segment::D, None, - Some(FunctionalStatus::Functional), - )); - let _ = cfg.d_pool.push(allele_with_status( - "d-p*01", b"GGG", Segment::D, None, - Some(FunctionalStatus::Pseudogene), - )); - let _ = cfg.j_pool.push(allele_with_status( - "j-f*01", b"TGGAAACCC", Segment::J, Some(0), - Some(FunctionalStatus::Functional), - )); - cfg - } - - #[test] - fn curation_functional_status_filters_v_d_j() { - let cfg = mixed_status_cfg(); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalStatus { - allowed: vec![FunctionalStatus::Functional], - keep_unannotated: false, - }); - let v_names: Vec<&str> = - curated.v_pool.iter().map(|(_, a)| a.name.as_str()).collect(); - let d_names: Vec<&str> = - curated.d_pool.iter().map(|(_, a)| a.name.as_str()).collect(); - let j_names: Vec<&str> = - curated.j_pool.iter().map(|(_, a)| a.name.as_str()).collect(); - assert_eq!(v_names, vec!["v-f*01"]); - assert_eq!(d_names, vec!["d-f*01"]); - assert_eq!(j_names, vec!["j-f*01"]); - } - - #[test] - fn curation_functional_status_keeps_unannotated_when_flagged() { - let cfg = mixed_status_cfg(); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalStatus { - allowed: vec![FunctionalStatus::Functional], - keep_unannotated: true, - }); - let v_names: Vec<&str> = - curated.v_pool.iter().map(|(_, a)| a.name.as_str()).collect(); - assert_eq!(v_names, vec!["v-f*01", "v-na*01"]); - } - - #[test] - fn curation_functional_status_tag_is_canonical() { - let p = RefDataCurationPolicy::FunctionalStatus { - allowed: vec![FunctionalStatus::Orf, FunctionalStatus::Functional], - keep_unannotated: false, - }; - // allowed list is sorted into canonical lexical order; the - // policy tag is the source of identity-source provenance, so - // two policies producing identical curated catalogues must - // produce identical tags regardless of input order. - assert_eq!( - p.tag(), - "functional_status:functional,orf|keep_unannotated=false", - ); - } - - #[test] - fn curation_functional_status_dedupes_allowed_in_tag() { - let p = RefDataCurationPolicy::FunctionalStatus { - allowed: vec![ - FunctionalStatus::Functional, - FunctionalStatus::Functional, - ], - keep_unannotated: true, - }; - assert_eq!(p.tag(), "functional_status:functional|keep_unannotated=true"); - } - - #[test] - fn curation_functional_status_empty_v_surfaces_empty_required_pool() { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg.v_pool.push(allele_with_status( - "v-p*01", b"TGTAAACCC", Segment::V, Some(0), - Some(FunctionalStatus::Pseudogene), - )); - let _ = cfg.j_pool.push(allele_with_status( - "j-f*01", b"TGGAAACCC", Segment::J, Some(0), - Some(FunctionalStatus::Functional), - )); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalStatus { - allowed: vec![FunctionalStatus::Functional], - keep_unannotated: false, - }); - assert!(curated.v_pool.is_empty()); - let issues = curated.validate(); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::EmptyRequiredPool { segment: Segment::V } - ))); - } - - #[test] - fn curation_functional_status_tags_identity_source() { - let cfg = mixed_status_cfg(); - let curated = cfg.curated(RefDataCurationPolicy::FunctionalStatus { - allowed: vec![FunctionalStatus::Functional], - keep_unannotated: true, - }); - let src = curated.identity.source.as_deref().unwrap_or(""); - assert!(src.contains("curated:functional_status:functional|keep_unannotated=true")); - } - - #[test] - fn ref_data_config_supports_vj_chain_with_empty_d_pool() { - // VJ chains: d_pool is conventionally empty. Construction - // does not enforce this; assembly (C.8) handles VJ vs VDJ - // explicitly via chain_type.has_d(). - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg.v_pool.push(make_v("v*01", "v", b"AA", Some(0))); - let _ = cfg.j_pool.push(Allele { - name: "j*01".into(), - gene: "j".into(), - seq: b"TT".to_vec(), - segment: Segment::J, - anchor: Some(0), - functional_status: None, - subregions: Vec::new(), - }); - - assert!(!cfg.chain_type.has_d()); - assert!(cfg.d_pool.is_empty()); - assert_eq!(cfg.v_pool.len(), 1); - assert_eq!(cfg.j_pool.len(), 1); - } -} +mod tests; #[cfg(test)] -mod gene_index_tests { - use super::*; - - fn pool_with(genes: &[(&str, &str)]) -> AllelePool { - // genes: (allele_name, gene_name) - let mut p = AllelePool::new(); - for (name, gene) in genes { - let _ = p.push(Allele { - name: (*name).to_string(), - gene: (*gene).to_string(), - seq: vec![b'A'; 10], - segment: Segment::V, - anchor: Some(3), - functional_status: None, - subregions: Vec::new(), - }); - } - p - } - - #[test] - fn gene_index_groups_alleles_by_gene_in_first_appearance_order() { - let pool = pool_with(&[ - ("IGHV1-2*01", "IGHV1-2"), - ("IGHV1-2*02", "IGHV1-2"), - ("IGHV3-23*01", "IGHV3-23"), - ]); - let idx = GeneIndex::build(&pool); - - assert_eq!(idx.len(), 2); - let g12 = idx.gene_id("IGHV1-2").unwrap(); - let g323 = idx.gene_id("IGHV3-23").unwrap(); - assert_eq!(g12.index(), 0); // first appearance - assert_eq!(g323.index(), 1); - assert_eq!(idx.gene_name(g12), "IGHV1-2"); - assert_eq!(idx.alleles_of(g12), &[AlleleId::new(0), AlleleId::new(1)]); - assert_eq!(idx.alleles_of(g323), &[AlleleId::new(2)]); - assert_eq!(idx.gene_of(AlleleId::new(1)), g12); - assert!(idx.gene_id("nope").is_none()); - } -} +mod gene_index_tests; diff --git a/engine_rs/src/refdata/allele.rs b/engine_rs/src/refdata/allele.rs deleted file mode 100644 index baa23b7..0000000 --- a/engine_rs/src/refdata/allele.rs +++ /dev/null @@ -1,40 +0,0 @@ -use crate::ir::Segment; - -/// One germline allele in a reference data set. -/// -/// Immutable once constructed. Per-simulation state (which allele was -/// sampled, what trim was applied, what ambiguity set the post-trim -/// retained bases project to) lives on `AlleleInstance`. -/// -/// **Field discipline:** -/// - `seq` is uppercase (`b'A'`, `b'C'`, `b'G'`, `b'T'`). Mixed case -/// is not allowed — case carries semantic meaning later in the -/// simulation (lowercase = SHM-mutated, etc.) and should not leak -/// into reference data. -/// - `anchor` is `Some(pos)` when the allele has a conserved codon -/// (Cys for V, W/F for J) at position `pos` (0-indexed within -/// `seq`). `None` for anchorless / partial alleles. The -/// `AnchorPreserved` contract reads this field. -/// - `name` is the canonical allele identifier (e.g., -/// `"IGHV1-2*01"`). `gene` is the truncation to the gene level -/// (e.g., `"IGHV1-2"`). -#[derive(Clone, Debug)] -pub struct Allele { - pub name: String, - pub gene: String, - pub seq: Vec, - pub segment: Segment, - pub anchor: Option, -} - -impl Allele { - /// Length of the allele sequence in nucleotides. - pub fn len(&self) -> u32 { - self.seq.len() as u32 - } - - /// Whether the allele has a known anchor position. - pub fn has_anchor(&self) -> bool { - self.anchor.is_some() - } -} diff --git a/engine_rs/src/refdata/chain.rs b/engine_rs/src/refdata/chain.rs deleted file mode 100644 index ebde24a..0000000 --- a/engine_rs/src/refdata/chain.rs +++ /dev/null @@ -1,23 +0,0 @@ -/// The biological chain configuration. -/// -/// `Vj` chains (light: kappa, lambda, TCR alpha/gamma) recombine V -/// and J only — there is no D segment, NP1 spans the V→J interval, -/// and NP2 / D pools are unused. -/// -/// `Vdj` chains (heavy: IGH, TCR beta/delta) recombine V, D, and J -/// with NP1 between V and D and NP2 between D and J. -/// -/// Other axes (allele frequency models, isotype constants, etc.) -/// are independent of `ChainType` and live elsewhere in the config. -#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)] -pub enum ChainType { - Vj, - Vdj, -} - -impl ChainType { - /// Whether this chain has a D segment (and therefore an NP2 region). - pub const fn has_d(self) -> bool { - matches!(self, ChainType::Vdj) - } -} diff --git a/engine_rs/src/refdata/config.rs b/engine_rs/src/refdata/config.rs deleted file mode 100644 index 74bd478..0000000 --- a/engine_rs/src/refdata/config.rs +++ /dev/null @@ -1,62 +0,0 @@ -use super::{Allele, AlleleId, AllelePool, ChainType}; -use crate::ir::Segment; - -/// Top-level immutable reference configuration for one simulation -/// target (e.g., human IGH). -/// -/// In production, this is loaded from a `.v6dat` file (D10) at -/// `Experiment.compile()` time. In tests, build directly via -/// `RefDataConfig::builder()` or by constructing the fields -/// manually. -/// -/// **Invariant:** for `chain_type == ChainType::Vj`, the `d_pool` is -/// expected to be empty. Construction does not enforce this — callers -/// of the builder API in tests are expected to honor it; the -/// assembly pass in C.8 will explicitly skip D for VJ chains. -#[derive(Clone, Debug)] -pub struct RefDataConfig { - pub chain_type: ChainType, - pub v_pool: AllelePool, - pub d_pool: AllelePool, - pub j_pool: AllelePool, - pub c_pool: AllelePool, -} - -impl RefDataConfig { - /// Empty config for the given chain type. Use the builder / - /// direct field assignment to populate the pools. - pub fn empty(chain_type: ChainType) -> Self { - Self { - chain_type, - v_pool: AllelePool::new(), - d_pool: AllelePool::new(), - j_pool: AllelePool::new(), - c_pool: AllelePool::new(), - } - } - - /// Resolve an allele by its segment role and id. Returns `None` - /// if the segment isn't in this config's pools (e.g., asking for - /// a D allele on a VJ chain) or if the id is out of bounds. - pub fn get(&self, segment: Segment, id: AlleleId) -> Option<&Allele> { - match segment { - Segment::V => self.v_pool.get(id), - Segment::D => self.d_pool.get(id), - Segment::J => self.j_pool.get(id), - // NP regions don't have alleles; return None defensively - // rather than panicking — a caller asking for an NP - // allele is misusing the API. - Segment::Np1 | Segment::Np2 => None, - } - } - - /// Pool for the given segment, or `None` for NP segments. - pub fn pool_for(&self, segment: Segment) -> Option<&AllelePool> { - match segment { - Segment::V => Some(&self.v_pool), - Segment::D => Some(&self.d_pool), - Segment::J => Some(&self.j_pool), - Segment::Np1 | Segment::Np2 => None, - } - } -} diff --git a/engine_rs/src/refdata/gene_index_tests.rs b/engine_rs/src/refdata/gene_index_tests.rs new file mode 100644 index 0000000..13de78a --- /dev/null +++ b/engine_rs/src/refdata/gene_index_tests.rs @@ -0,0 +1,39 @@ + use super::*; + + fn pool_with(genes: &[(&str, &str)]) -> AllelePool { + // genes: (allele_name, gene_name) + let mut p = AllelePool::new(); + for (name, gene) in genes { + let _ = p.push(Allele { + name: (*name).to_string(), + gene: (*gene).to_string(), + seq: vec![b'A'; 10], + segment: Segment::V, + anchor: Some(3), + functional_status: None, + subregions: Vec::new(), + }); + } + p + } + + #[test] + fn gene_index_groups_alleles_by_gene_in_first_appearance_order() { + let pool = pool_with(&[ + ("IGHV1-2*01", "IGHV1-2"), + ("IGHV1-2*02", "IGHV1-2"), + ("IGHV3-23*01", "IGHV3-23"), + ]); + let idx = GeneIndex::build(&pool); + + assert_eq!(idx.len(), 2); + let g12 = idx.gene_id("IGHV1-2").unwrap(); + let g323 = idx.gene_id("IGHV3-23").unwrap(); + assert_eq!(g12.index(), 0); // first appearance + assert_eq!(g323.index(), 1); + assert_eq!(idx.gene_name(g12), "IGHV1-2"); + assert_eq!(idx.alleles_of(g12), &[AlleleId::new(0), AlleleId::new(1)]); + assert_eq!(idx.alleles_of(g323), &[AlleleId::new(2)]); + assert_eq!(idx.gene_of(AlleleId::new(1)), g12); + assert!(idx.gene_id("nope").is_none()); + } diff --git a/engine_rs/src/refdata/id.rs b/engine_rs/src/refdata/id.rs deleted file mode 100644 index f8f0c3c..0000000 --- a/engine_rs/src/refdata/id.rs +++ /dev/null @@ -1,20 +0,0 @@ -/// Index into an `AllelePool`. Stable for the lifetime of the -/// `RefDataConfig` it was issued for. -/// -/// Distinct type from `NucHandle` and `RegionHandle` so the compiler -/// catches handle-confusion at call sites — this is the same -/// discipline as the IR handles in §3 of the design doc. -#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)] -pub struct AlleleId(u32); - -impl AlleleId { - pub const fn new(idx: u32) -> Self { - Self(idx) - } - pub const fn index(self) -> u32 { - self.0 - } - pub const fn as_usize(self) -> usize { - self.0 as usize - } -} diff --git a/engine_rs/src/refdata/pool.rs b/engine_rs/src/refdata/pool.rs deleted file mode 100644 index bbb861e..0000000 --- a/engine_rs/src/refdata/pool.rs +++ /dev/null @@ -1,75 +0,0 @@ -use super::{Allele, AlleleId}; - -/// All alleles available for one segment role (V, D, J, or C). -/// -/// Indexed by `AlleleId`. The pool is the reference data; *which* -/// allele a particular simulation sampled is recorded by an -/// `AlleleInstance` referring to a stable `AlleleId`. -#[derive(Clone, Debug, Default)] -pub struct AllelePool { - alleles: Vec, -} - -impl AllelePool { - pub fn new() -> Self { - Self::default() - } - - /// Construct from an existing `Vec`. Caller is responsible - /// for ensuring all alleles have the right segment role for the - /// pool being built. - pub fn from_vec(alleles: Vec) -> Self { - Self { alleles } - } - - /// Number of alleles in the pool. - pub fn len(&self) -> usize { - self.alleles.len() - } - - /// Whether the pool contains zero alleles. - pub fn is_empty(&self) -> bool { - self.alleles.is_empty() - } - - /// Append an allele. Returns the issued `AlleleId` so callers can - /// reference the just-added allele in subsequent setup. This is - /// the construction-time API; reference data is sealed before any - /// simulation runs against it. - #[must_use = "AllelePool::push returns the issued AlleleId; use \ - `_ = pool.push(...)` if it isn't needed"] - pub fn push(&mut self, allele: Allele) -> AlleleId { - let id = AlleleId::new(self.alleles.len() as u32); - self.alleles.push(allele); - id - } - - /// Look up an allele by id. Returns `None` if `id` is out of - /// bounds — defensive to allow `RefDataConfig` to expose - /// fallible lookup at the boundary even though correct callers - /// always have valid handles. - pub fn get(&self, id: AlleleId) -> Option<&Allele> { - self.alleles.get(id.as_usize()) - } - - /// Read-only slice of all alleles in pool order. Iterated by - /// `AllelePoolDist` to build cumulative frequency tables. - pub fn as_slice(&self) -> &[Allele] { - &self.alleles - } - - /// Iterator over `(AlleleId, &Allele)` pairs in pool order. - pub fn iter(&self) -> impl Iterator { - self.alleles - .iter() - .enumerate() - .map(|(i, a)| (AlleleId::new(i as u32), a)) - } - - /// Find the first allele whose name matches exactly. O(N). - /// Used in tests and tooling; production sampling goes through - /// `AllelePoolDist` (C.3) and never name-resolves here. - pub fn find_by_name(&self, name: &str) -> Option<(AlleleId, &Allele)> { - self.iter().find(|(_, a)| a.name == name) - } -} diff --git a/engine_rs/src/refdata/tests.rs b/engine_rs/src/refdata/tests.rs index 5c7b349..7736f84 100644 --- a/engine_rs/src/refdata/tests.rs +++ b/engine_rs/src/refdata/tests.rs @@ -1,203 +1,545 @@ -use super::*; -use crate::ir::Segment; -use std::mem::size_of; - -/// Helper: a tiny synthetic V allele for tests. -fn make_v(name: &str, gene: &str, seq: &[u8], anchor: Option) -> Allele { - Allele { - name: name.to_string(), - gene: gene.to_string(), - seq: seq.to_vec(), - segment: Segment::V, - anchor, - functional_status: None, - subregions: Vec::new(), - } -} - -#[test] -fn allele_id_is_zero_cost_newtype() { - assert_eq!(size_of::(), size_of::()); -} - -#[test] -fn allele_id_round_trip() { - let id = AlleleId::new(7); - assert_eq!(id.index(), 7); - assert_eq!(id.as_usize(), 7); -} - -#[test] -fn chain_type_has_d_distinction() { - assert!(!ChainType::Vj.has_d()); - assert!(ChainType::Vdj.has_d()); -} - -#[test] -fn allele_basic_accessors() { - let a = make_v("IGHV1-2*01", "IGHV1-2", b"ACGTACGT", Some(3)); - assert_eq!(a.len(), 8); - assert!(a.has_anchor()); - assert_eq!(a.anchor, Some(3)); - assert_eq!(a.segment, Segment::V); -} - -#[test] -fn allele_anchorless_round_trip() { - let a = make_v("IGHV-pseudo*01", "IGHV-pseudo", b"ACGT", None); - assert!(!a.has_anchor()); - assert_eq!(a.anchor, None); -} - -#[test] -fn allele_pool_starts_empty() { - let p = AllelePool::new(); - assert_eq!(p.len(), 0); - assert!(p.is_empty()); - assert!(p.get(AlleleId::new(0)).is_none()); - assert!(p.find_by_name("nonexistent").is_none()); -} - -#[test] -fn allele_pool_push_returns_sequential_ids() { - let mut p = AllelePool::new(); - let id0 = p.push(make_v("a*01", "a", b"AA", Some(0))); - let id1 = p.push(make_v("b*01", "b", b"CC", Some(0))); - let id2 = p.push(make_v("c*01", "c", b"GG", None)); - - assert_eq!(id0.index(), 0); - assert_eq!(id1.index(), 1); - assert_eq!(id2.index(), 2); - assert_eq!(p.len(), 3); -} - -#[test] -fn allele_pool_get_returns_stored_allele() { - let mut p = AllelePool::new(); - let id = p.push(make_v("x*01", "x", b"ATGC", Some(2))); - let got = p.get(id).expect("just pushed"); - assert_eq!(got.name, "x*01"); - assert_eq!(got.seq, b"ATGC"); - assert_eq!(got.anchor, Some(2)); -} - -#[test] -fn allele_pool_get_out_of_bounds_returns_none() { - let p = AllelePool::new(); - assert!(p.get(AlleleId::new(99)).is_none()); -} - -#[test] -fn allele_pool_iter_yields_id_allele_pairs_in_order() { - let mut p = AllelePool::new(); - let _ = p.push(make_v("a*01", "a", b"AA", None)); - let _ = p.push(make_v("b*01", "b", b"CC", None)); - let _ = p.push(make_v("c*01", "c", b"GG", None)); - - let collected: Vec<(u32, String)> = p - .iter() - .map(|(id, a)| (id.index(), a.name.clone())) - .collect(); - assert_eq!( - collected, - vec![ - (0, "a*01".to_string()), - (1, "b*01".to_string()), - (2, "c*01".to_string()), - ] - ); -} - -#[test] -fn allele_pool_find_by_name_locates_exact_match() { - let mut p = AllelePool::new(); - let _ = p.push(make_v("IGHV1-2*01", "IGHV1-2", b"AA", None)); - let target_id = p.push(make_v("IGHV1-2*02", "IGHV1-2", b"AC", None)); - let _ = p.push(make_v("IGHV3-23*01", "IGHV3-23", b"GG", None)); - - let (id, allele) = p.find_by_name("IGHV1-2*02").expect("name should exist"); - assert_eq!(id, target_id); - assert_eq!(allele.seq, b"AC"); - - // Partial / similar names should not match. - assert!(p.find_by_name("IGHV1-2").is_none()); - assert!(p.find_by_name("IGHV1-2*03").is_none()); -} - -#[test] -fn ref_data_config_empty_for_chain_type() { - let cfg = RefDataConfig::empty(ChainType::Vdj); - assert_eq!(cfg.chain_type, ChainType::Vdj); - assert!(cfg.v_pool.is_empty()); - assert!(cfg.d_pool.is_empty()); - assert!(cfg.j_pool.is_empty()); - assert!(cfg.c_pool.is_empty()); -} - -#[test] -fn ref_data_config_pool_for_segment_routes_correctly() { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - let _ = cfg.v_pool.push(make_v("v*01", "v", b"AA", None)); - let _ = cfg.d_pool.push(Allele { - name: "d*01".into(), - gene: "d".into(), - seq: b"GG".to_vec(), - segment: Segment::D, - anchor: None, - functional_status: None, - subregions: Vec::new(), - }); - let _ = cfg.j_pool.push(Allele { - name: "j*01".into(), - gene: "j".into(), - seq: b"TT".to_vec(), - segment: Segment::J, - anchor: Some(0), - functional_status: None, - subregions: Vec::new(), - }); - - assert_eq!(cfg.pool_for(Segment::V).unwrap().len(), 1); - assert_eq!(cfg.pool_for(Segment::D).unwrap().len(), 1); - assert_eq!(cfg.pool_for(Segment::J).unwrap().len(), 1); - assert!(cfg.pool_for(Segment::Np1).is_none()); - assert!(cfg.pool_for(Segment::Np2).is_none()); -} - -#[test] -fn ref_data_config_get_resolves_segment_id_pair() { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - let v_id = cfg.v_pool.push(make_v("v*01", "v", b"AAAT", Some(1))); - - let v = cfg.get(Segment::V, v_id).expect("v*01 should resolve"); - assert_eq!(v.name, "v*01"); - assert_eq!(v.anchor, Some(1)); - - // Wrong segment -> None. - assert!(cfg.get(Segment::J, v_id).is_none()); - // NP segment -> None defensively. - assert!(cfg.get(Segment::Np1, v_id).is_none()); -} - -#[test] -fn ref_data_config_supports_vj_chain_with_empty_d_pool() { - // VJ chains: d_pool is conventionally empty. Construction - // does not enforce this; assembly (C.8) handles VJ vs VDJ - // explicitly via chain_type.has_d(). - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg.v_pool.push(make_v("v*01", "v", b"AA", Some(0))); - let _ = cfg.j_pool.push(Allele { - name: "j*01".into(), - gene: "j".into(), - seq: b"TT".to_vec(), - segment: Segment::J, - anchor: Some(0), - functional_status: None, - subregions: Vec::new(), - }); - - assert!(!cfg.chain_type.has_d()); - assert!(cfg.d_pool.is_empty()); - assert_eq!(cfg.v_pool.len(), 1); - assert_eq!(cfg.j_pool.len(), 1); -} + use super::*; + use std::mem::size_of; + + /// Helper: a tiny synthetic V allele for tests. + fn make_v(name: &str, gene: &str, seq: &[u8], anchor: Option) -> Allele { + Allele { + name: name.to_string(), + gene: gene.to_string(), + seq: seq.to_vec(), + segment: Segment::V, + anchor, + functional_status: None, + subregions: Vec::new(), + } + } + + #[test] + fn allele_id_is_zero_cost_newtype() { + assert_eq!(size_of::(), size_of::()); + } + + #[test] + fn allele_id_round_trip() { + let id = AlleleId::new(7); + assert_eq!(id.index(), 7); + assert_eq!(id.as_usize(), 7); + } + + #[test] + fn chain_type_has_d_distinction() { + assert!(!ChainType::Vj.has_d()); + assert!(ChainType::Vdj.has_d()); + } + + #[test] + fn allele_basic_accessors() { + let a = make_v("IGHV1-2*01", "IGHV1-2", b"ACGTACGT", Some(3)); + assert_eq!(a.len(), 8); + assert!(a.has_anchor()); + assert_eq!(a.anchor, Some(3)); + assert_eq!(a.segment, Segment::V); + } + + #[test] + fn allele_anchorless_round_trip() { + let a = make_v("IGHV-pseudo*01", "IGHV-pseudo", b"ACGT", None); + assert!(!a.has_anchor()); + assert_eq!(a.anchor, None); + } + + #[test] + fn allele_pool_starts_empty() { + let p = AllelePool::new(); + assert_eq!(p.len(), 0); + assert!(p.is_empty()); + assert!(p.get(AlleleId::new(0)).is_none()); + assert!(p.find_by_name("nonexistent").is_none()); + } + + #[test] + fn allele_pool_push_returns_sequential_ids() { + let mut p = AllelePool::new(); + let id0 = p.push(make_v("a*01", "a", b"AA", Some(0))); + let id1 = p.push(make_v("b*01", "b", b"CC", Some(0))); + let id2 = p.push(make_v("c*01", "c", b"GG", None)); + + assert_eq!(id0.index(), 0); + assert_eq!(id1.index(), 1); + assert_eq!(id2.index(), 2); + assert_eq!(p.len(), 3); + } + + #[test] + fn allele_pool_get_returns_stored_allele() { + let mut p = AllelePool::new(); + let id = p.push(make_v("x*01", "x", b"ATGC", Some(2))); + let got = p.get(id).expect("just pushed"); + assert_eq!(got.name, "x*01"); + assert_eq!(got.seq, b"ATGC"); + assert_eq!(got.anchor, Some(2)); + } + + #[test] + fn allele_pool_get_out_of_bounds_returns_none() { + let p = AllelePool::new(); + assert!(p.get(AlleleId::new(99)).is_none()); + } + + #[test] + fn allele_pool_iter_yields_id_allele_pairs_in_order() { + let mut p = AllelePool::new(); + let _ = p.push(make_v("a*01", "a", b"AA", None)); + let _ = p.push(make_v("b*01", "b", b"CC", None)); + let _ = p.push(make_v("c*01", "c", b"GG", None)); + + let collected: Vec<(u32, String)> = p + .iter() + .map(|(id, a)| (id.index(), a.name.clone())) + .collect(); + assert_eq!( + collected, + vec![ + (0, "a*01".to_string()), + (1, "b*01".to_string()), + (2, "c*01".to_string()), + ] + ); + } + + #[test] + fn allele_pool_find_by_name_locates_exact_match() { + let mut p = AllelePool::new(); + let _ = p.push(make_v("IGHV1-2*01", "IGHV1-2", b"AA", None)); + let target_id = p.push(make_v("IGHV1-2*02", "IGHV1-2", b"AC", None)); + let _ = p.push(make_v("IGHV3-23*01", "IGHV3-23", b"GG", None)); + + let (id, allele) = p.find_by_name("IGHV1-2*02").expect("name should exist"); + assert_eq!(id, target_id); + assert_eq!(allele.seq, b"AC"); + + // Partial / similar names should not match. + assert!(p.find_by_name("IGHV1-2").is_none()); + assert!(p.find_by_name("IGHV1-2*03").is_none()); + } + + #[test] + fn ref_data_config_empty_for_chain_type() { + let cfg = RefDataConfig::empty(ChainType::Vdj); + assert_eq!(cfg.chain_type, ChainType::Vdj); + assert!(cfg.v_pool.is_empty()); + assert!(cfg.d_pool.is_empty()); + assert!(cfg.j_pool.is_empty()); + assert!(cfg.c_pool.is_empty()); + } + + #[test] + fn ref_data_config_pool_for_segment_routes_correctly() { + let mut cfg = RefDataConfig::empty(ChainType::Vdj); + let _ = cfg.v_pool.push(make_v("v*01", "v", b"AA", None)); + let _ = cfg.d_pool.push(Allele { + name: "d*01".into(), + gene: "d".into(), + seq: b"GG".to_vec(), + segment: Segment::D, + anchor: None, + functional_status: None, + subregions: Vec::new(), + }); + let _ = cfg.j_pool.push(Allele { + name: "j*01".into(), + gene: "j".into(), + seq: b"TT".to_vec(), + segment: Segment::J, + anchor: Some(0), + functional_status: None, + subregions: Vec::new(), + }); + + assert_eq!(cfg.pool_for(Segment::V).unwrap().len(), 1); + assert_eq!(cfg.pool_for(Segment::D).unwrap().len(), 1); + assert_eq!(cfg.pool_for(Segment::J).unwrap().len(), 1); + assert!(cfg.pool_for(Segment::Np1).is_none()); + assert!(cfg.pool_for(Segment::Np2).is_none()); + } + + #[test] + fn ref_data_config_get_resolves_segment_id_pair() { + let mut cfg = RefDataConfig::empty(ChainType::Vdj); + let v_id = cfg.v_pool.push(make_v("v*01", "v", b"AAAT", Some(1))); + + let v = cfg.get(Segment::V, v_id).expect("v*01 should resolve"); + assert_eq!(v.name, "v*01"); + assert_eq!(v.anchor, Some(1)); + + // Wrong segment -> None. + assert!(cfg.get(Segment::J, v_id).is_none()); + // NP segment -> None defensively. + assert!(cfg.get(Segment::Np1, v_id).is_none()); + } + + // ── Curation policy ─────────────────────────────────────────── + + fn allele_with_anchor(name: &str, seq: &[u8], anchor: Option, segment: Segment) -> Allele { + Allele { + name: name.to_string(), + gene: name.split('*').next().unwrap_or(name).to_string(), + seq: seq.to_vec(), + segment, + anchor, + functional_status: None, + subregions: Vec::new(), + } + } + + fn vj_cfg_for_curation() -> RefDataConfig { + // Two V alleles: one Cys-anchored, one Gly-anchored. + // Two J alleles: one W-anchored, one anchorless. + let mut cfg = RefDataConfig::empty(ChainType::Vj); + let _ = cfg + .v_pool + .push(allele_with_anchor("MYV-good*01", b"TGTAAACCC", Some(0), Segment::V)); + let _ = cfg + .v_pool + .push(allele_with_anchor("MYV-gly*01", b"GGGAAACCC", Some(0), Segment::V)); + let _ = cfg + .j_pool + .push(allele_with_anchor("MYJ-good*01", b"TGGAAACCC", Some(0), Segment::J)); + let _ = cfg + .j_pool + .push(allele_with_anchor("MYJ-orphan*01", b"GGG", None, Segment::J)); + cfg + } + + #[test] + fn curation_raw_is_identity() { + let cfg = vj_cfg_for_curation(); + let curated = cfg.curated(RefDataCurationPolicy::Raw); + assert_eq!(curated.v_pool.len(), 2); + assert_eq!(curated.j_pool.len(), 2); + // Raw policy does NOT touch identity. + assert_eq!(curated.identity.source, None); + } + + #[test] + fn curation_functional_anchors_only_filters_v_and_j() { + let cfg = vj_cfg_for_curation(); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + assert_eq!(curated.v_pool.len(), 1); + assert_eq!(curated.v_pool.iter().next().unwrap().1.name, "MYV-good*01"); + assert_eq!(curated.j_pool.len(), 1); + assert_eq!(curated.j_pool.iter().next().unwrap().1.name, "MYJ-good*01"); + } + + #[test] + fn curation_tags_identity_source() { + let cfg = vj_cfg_for_curation(); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + assert_eq!( + curated.identity.source.as_deref(), + Some("curated:functional_anchors_only"), + ); + } + + #[test] + fn curation_extends_existing_identity_source() { + let mut cfg = vj_cfg_for_curation(); + cfg.identity.source = Some("DataConfig".to_string()); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + assert_eq!( + curated.identity.source.as_deref(), + Some("DataConfig|curated:functional_anchors_only"), + ); + } + + #[test] + fn curation_preserves_other_identity_fields() { + let mut cfg = vj_cfg_for_curation(); + cfg.identity.species = Some("Human".into()); + cfg.identity.locus = Some("IGK".into()); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + assert_eq!(curated.identity.species.as_deref(), Some("Human")); + assert_eq!(curated.identity.locus.as_deref(), Some("IGK")); + } + + #[test] + fn curation_does_not_silently_fix_duplicate_names() { + // Add a duplicate V allele AFTER the good one. Both are + // Cys-anchored — they survive curation, then validation + // surfaces the duplicate as a Fatal issue. Curation never + // removes a structurally-corrupt entry to make a catalogue + // look healthier. + let mut cfg = vj_cfg_for_curation(); + let _ = cfg.v_pool.push(allele_with_anchor( + "MYV-good*01", + b"TGTAAACCC", + Some(0), + Segment::V, + )); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + let issues = curated.validate(); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::DuplicateAlleleName { .. } + ))); + } + + #[test] + fn curation_does_not_silently_fix_invalid_bytes() { + // Add an allele with a gap byte ('.') and a valid Cys anchor. + // The invalid byte is Fatal and must still surface post-curation. + let mut cfg = vj_cfg_for_curation(); + let _ = cfg.v_pool.push(allele_with_anchor( + "MYV-badbyte*01", + b"TGT.AAACC", + Some(0), + Segment::V, + )); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + assert!(curated.validate().iter().any(|i| matches!( + i, + RefDataValidationIssue::InvalidAlleleByte { .. } + ))); + } + + #[test] + fn curation_to_empty_pool_surfaces_empty_required_pool() { + // Every V allele is non-Cys — curation empties V. + let mut cfg = RefDataConfig::empty(ChainType::Vj); + let _ = cfg + .v_pool + .push(allele_with_anchor("MYV-gly*01", b"GGGAAACCC", Some(0), Segment::V)); + let _ = cfg + .j_pool + .push(allele_with_anchor("MYJ-good*01", b"TGGAAACCC", Some(0), Segment::J)); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + assert_eq!(curated.v_pool.len(), 0); + let issues = curated.validate(); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::EmptyRequiredPool { segment: Segment::V } + ))); + } + + #[test] + fn curation_d_pool_passes_through_unchanged() { + let mut cfg = RefDataConfig::empty(ChainType::Vdj); + let _ = cfg + .v_pool + .push(allele_with_anchor("MYV*01", b"TGTAAACCC", Some(0), Segment::V)); + let _ = cfg.d_pool.push(allele_with_anchor( + "MYD*01", + b"GGGCCCAAA", + None, + Segment::D, + )); + let _ = cfg + .j_pool + .push(allele_with_anchor("MYJ*01", b"TGGAAACCC", Some(0), Segment::J)); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + // Anchorless D survives — D pool is never filtered. + assert_eq!(curated.d_pool.len(), 1); + } + + #[test] + fn curation_respects_custom_j_anchor_rule() { + // Rule expects Y at J anchor. TGG (W) and TTC (F) are dropped; + // TAT (Y) is kept. + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.j_anchor.expected_amino_acids = vec!['Y']; + let _ = cfg + .v_pool + .push(allele_with_anchor("MYV*01", b"TGTAAACCC", Some(0), Segment::V)); + let _ = cfg + .j_pool + .push(allele_with_anchor("MYJ-w*01", b"TGGAAACCC", Some(0), Segment::J)); + let _ = cfg + .j_pool + .push(allele_with_anchor("MYJ-y*01", b"TATAAACCC", Some(0), Segment::J)); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + assert_eq!(curated.j_pool.len(), 1); + assert_eq!(curated.j_pool.iter().next().unwrap().1.name, "MYJ-y*01"); + } + + #[test] + fn curation_anchor_required_false_keeps_anchorless_alleles() { + // If the rule says anchor is optional, anchorless alleles + // are kept under FunctionalAnchorsOnly. + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.j_anchor.required = false; + let _ = cfg + .v_pool + .push(allele_with_anchor("MYV*01", b"TGTAAACCC", Some(0), Segment::V)); + let _ = cfg + .j_pool + .push(allele_with_anchor("MYJ-orphan*01", b"GGG", None, Segment::J)); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalAnchorsOnly); + assert_eq!(curated.j_pool.len(), 1); + } + + // ── Functional-status curation ──────────────────────────────── + + fn allele_with_status( + name: &str, + seq: &[u8], + segment: Segment, + anchor: Option, + status: Option, + ) -> Allele { + Allele { + name: name.to_string(), + gene: name.split('*').next().unwrap_or(name).to_string(), + seq: seq.to_vec(), + segment, + anchor, + functional_status: status, + subregions: Vec::new(), + } + } + + fn mixed_status_cfg() -> RefDataConfig { + let mut cfg = RefDataConfig::empty(ChainType::Vdj); + let _ = cfg.v_pool.push(allele_with_status( + "v-f*01", b"TGTAAACCC", Segment::V, Some(0), + Some(FunctionalStatus::Functional), + )); + let _ = cfg.v_pool.push(allele_with_status( + "v-o*01", b"TGTAAACCC", Segment::V, Some(0), + Some(FunctionalStatus::Orf), + )); + let _ = cfg.v_pool.push(allele_with_status( + "v-p*01", b"TGTAAACCC", Segment::V, Some(0), + Some(FunctionalStatus::Pseudogene), + )); + let _ = cfg.v_pool.push(allele_with_status( + "v-na*01", b"TGTAAACCC", Segment::V, Some(0), None, + )); + let _ = cfg.d_pool.push(allele_with_status( + "d-f*01", b"GGG", Segment::D, None, + Some(FunctionalStatus::Functional), + )); + let _ = cfg.d_pool.push(allele_with_status( + "d-p*01", b"GGG", Segment::D, None, + Some(FunctionalStatus::Pseudogene), + )); + let _ = cfg.j_pool.push(allele_with_status( + "j-f*01", b"TGGAAACCC", Segment::J, Some(0), + Some(FunctionalStatus::Functional), + )); + cfg + } + + #[test] + fn curation_functional_status_filters_v_d_j() { + let cfg = mixed_status_cfg(); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalStatus { + allowed: vec![FunctionalStatus::Functional], + keep_unannotated: false, + }); + let v_names: Vec<&str> = + curated.v_pool.iter().map(|(_, a)| a.name.as_str()).collect(); + let d_names: Vec<&str> = + curated.d_pool.iter().map(|(_, a)| a.name.as_str()).collect(); + let j_names: Vec<&str> = + curated.j_pool.iter().map(|(_, a)| a.name.as_str()).collect(); + assert_eq!(v_names, vec!["v-f*01"]); + assert_eq!(d_names, vec!["d-f*01"]); + assert_eq!(j_names, vec!["j-f*01"]); + } + + #[test] + fn curation_functional_status_keeps_unannotated_when_flagged() { + let cfg = mixed_status_cfg(); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalStatus { + allowed: vec![FunctionalStatus::Functional], + keep_unannotated: true, + }); + let v_names: Vec<&str> = + curated.v_pool.iter().map(|(_, a)| a.name.as_str()).collect(); + assert_eq!(v_names, vec!["v-f*01", "v-na*01"]); + } + + #[test] + fn curation_functional_status_tag_is_canonical() { + let p = RefDataCurationPolicy::FunctionalStatus { + allowed: vec![FunctionalStatus::Orf, FunctionalStatus::Functional], + keep_unannotated: false, + }; + // allowed list is sorted into canonical lexical order; the + // policy tag is the source of identity-source provenance, so + // two policies producing identical curated catalogues must + // produce identical tags regardless of input order. + assert_eq!( + p.tag(), + "functional_status:functional,orf|keep_unannotated=false", + ); + } + + #[test] + fn curation_functional_status_dedupes_allowed_in_tag() { + let p = RefDataCurationPolicy::FunctionalStatus { + allowed: vec![ + FunctionalStatus::Functional, + FunctionalStatus::Functional, + ], + keep_unannotated: true, + }; + assert_eq!(p.tag(), "functional_status:functional|keep_unannotated=true"); + } + + #[test] + fn curation_functional_status_empty_v_surfaces_empty_required_pool() { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + let _ = cfg.v_pool.push(allele_with_status( + "v-p*01", b"TGTAAACCC", Segment::V, Some(0), + Some(FunctionalStatus::Pseudogene), + )); + let _ = cfg.j_pool.push(allele_with_status( + "j-f*01", b"TGGAAACCC", Segment::J, Some(0), + Some(FunctionalStatus::Functional), + )); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalStatus { + allowed: vec![FunctionalStatus::Functional], + keep_unannotated: false, + }); + assert!(curated.v_pool.is_empty()); + let issues = curated.validate(); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::EmptyRequiredPool { segment: Segment::V } + ))); + } + + #[test] + fn curation_functional_status_tags_identity_source() { + let cfg = mixed_status_cfg(); + let curated = cfg.curated(RefDataCurationPolicy::FunctionalStatus { + allowed: vec![FunctionalStatus::Functional], + keep_unannotated: true, + }); + let src = curated.identity.source.as_deref().unwrap_or(""); + assert!(src.contains("curated:functional_status:functional|keep_unannotated=true")); + } + + #[test] + fn ref_data_config_supports_vj_chain_with_empty_d_pool() { + // VJ chains: d_pool is conventionally empty. Construction + // does not enforce this; assembly (C.8) handles VJ vs VDJ + // explicitly via chain_type.has_d(). + let mut cfg = RefDataConfig::empty(ChainType::Vj); + let _ = cfg.v_pool.push(make_v("v*01", "v", b"AA", Some(0))); + let _ = cfg.j_pool.push(Allele { + name: "j*01".into(), + gene: "j".into(), + seq: b"TT".to_vec(), + segment: Segment::J, + anchor: Some(0), + functional_status: None, + subregions: Vec::new(), + }); + + assert!(!cfg.chain_type.has_d()); + assert!(cfg.d_pool.is_empty()); + assert_eq!(cfg.v_pool.len(), 1); + assert_eq!(cfg.j_pool.len(), 1); + } diff --git a/engine_rs/src/refdata/validation.rs b/engine_rs/src/refdata/validation.rs index 7f05ef7..ff5c0ce 100644 --- a/engine_rs/src/refdata/validation.rs +++ b/engine_rs/src/refdata/validation.rs @@ -672,815 +672,4 @@ fn _chain_type_imported(ct: ChainType) -> bool { // ────────────────────────────────────────────────────────────────── #[cfg(test)] -mod tests { - use super::*; - use crate::refdata::{Allele, AlleleId, ChainType, RefDataConfig}; - - fn v_allele(name: &str, seq: &[u8], anchor: Option) -> Allele { - Allele { - name: name.to_string(), - gene: name.split('*').next().unwrap_or(name).to_string(), - seq: seq.to_vec(), - segment: Segment::V, - anchor, - functional_status: None, - subregions: Vec::new(), - } - } - fn d_allele(name: &str, seq: &[u8]) -> Allele { - Allele { - name: name.to_string(), - gene: name.split('*').next().unwrap_or(name).to_string(), - seq: seq.to_vec(), - segment: Segment::D, - anchor: None, - functional_status: None, - subregions: Vec::new(), - } - } - fn j_allele(name: &str, seq: &[u8], anchor: Option) -> Allele { - Allele { - name: name.to_string(), - gene: name.split('*').next().unwrap_or(name).to_string(), - seq: seq.to_vec(), - segment: Segment::J, - anchor, - functional_status: None, - subregions: Vec::new(), - } - } - - /// Build a minimal valid VDJ refdata so each test only varies - /// the one field under examination. - fn minimal_vdj() -> RefDataConfig { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - // V allele: anchor codon TGT (Cys) at position 0. - let _ = cfg.v_pool.push(v_allele("IGHV1-1*01", b"TGTAAACCC", Some(0))); - // D allele: no anchor. - let _ = cfg.d_pool.push(d_allele("IGHD1-1*01", b"GGGCCCAAA")); - // J allele: anchor codon TGG (Trp = W) at position 0. - let _ = cfg.j_pool.push(j_allele("IGHJ1*01", b"TGGAAACCC", Some(0))); - cfg - } - - fn minimal_vj() -> RefDataConfig { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg.v_pool.push(v_allele("IGKV1-1*01", b"TGTAAACCC", Some(0))); - // IGK J anchor codon TTC (Phe = F). - let _ = cfg.j_pool.push(j_allele("IGKJ1*01", b"TTCAAACCC", Some(0))); - cfg - } - - #[test] - fn minimal_vdj_validates_clean() { - assert!(minimal_vdj().validate().is_empty()); - assert!(minimal_vdj().validate_strict().is_ok()); - } - - #[test] - fn minimal_vj_validates_clean() { - assert!(minimal_vj().validate().is_empty()); - } - - #[test] - fn empty_v_pool_fails_vj() { - let mut cfg = minimal_vj(); - cfg.v_pool = super::AllelePool::new(); - let issues = cfg.validate(); - assert!(issues.contains(&RefDataValidationIssue::EmptyRequiredPool { - segment: Segment::V, - })); - } - - #[test] - fn empty_v_pool_fails_vdj() { - let mut cfg = minimal_vdj(); - cfg.v_pool = super::AllelePool::new(); - let issues = cfg.validate(); - assert!(issues.contains(&RefDataValidationIssue::EmptyRequiredPool { - segment: Segment::V, - })); - } - - #[test] - fn empty_j_pool_fails_vj() { - let mut cfg = minimal_vj(); - cfg.j_pool = super::AllelePool::new(); - let issues = cfg.validate(); - assert!(issues.contains(&RefDataValidationIssue::EmptyRequiredPool { - segment: Segment::J, - })); - } - - #[test] - fn empty_d_pool_allowed_vj() { - // VJ chains have empty D pools by construction. This must - // not flag an EmptyRequiredPool for D. - let cfg = minimal_vj(); - let issues = cfg.validate(); - assert!(!issues - .iter() - .any(|i| matches!(i, RefDataValidationIssue::EmptyRequiredPool { segment } if *segment == Segment::D))); - } - - #[test] - fn empty_d_pool_rejected_vdj() { - let mut cfg = minimal_vdj(); - cfg.d_pool = super::AllelePool::new(); - let issues = cfg.validate(); - assert!(issues.contains(&RefDataValidationIssue::EmptyRequiredPool { - segment: Segment::D, - })); - } - - #[test] - fn duplicate_v_allele_names_fail() { - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV1-1*01", b"TGTAAA", Some(0))); - let issues = cfg.validate(); - assert!(issues.contains(&RefDataValidationIssue::DuplicateAlleleName { - segment: Segment::V, - name: "IGHV1-1*01".to_string(), - })); - } - - #[test] - fn duplicate_j_allele_names_fail() { - let mut cfg = minimal_vdj(); - let _ = cfg.j_pool.push(j_allele("IGHJ1*01", b"TGGAAA", Some(0))); - let issues = cfg.validate(); - assert!(issues.contains(&RefDataValidationIssue::DuplicateAlleleName { - segment: Segment::J, - name: "IGHJ1*01".to_string(), - })); - } - - #[test] - fn invalid_byte_in_v_seq_fails_with_exact_position() { - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV2*01", b"TGT.AAA", Some(0))); - let issues = cfg.validate(); - let bad_v_id = AlleleId::new(1); - let expected = RefDataValidationIssue::InvalidAlleleByte { - segment: Segment::V, - allele_id: bad_v_id, - pos: 3, - byte: b'.', - }; - assert!( - issues.contains(&expected), - "expected invalid-byte issue at pos 3, got: {issues:?}" - ); - } - - #[test] - fn lowercase_bases_accepted() { - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV3*01", b"tgtaaaccc", Some(0))); - // Lowercase A/C/G/T/N is allowed. The anchor codon `tgt` still - // translates to Cys, so V anchor check also passes. - let issues = cfg.validate(); - assert!( - issues.is_empty(), - "lowercase canonical bases should validate, got: {issues:?}" - ); - } - - #[test] - fn iupac_ambiguity_other_than_n_fails() { - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV4*01", b"TGTRAACCC", Some(0))); - let issues = cfg.validate(); - let v_id = AlleleId::new(1); - assert!(issues.contains(&RefDataValidationIssue::InvalidAlleleByte { - segment: Segment::V, - allele_id: v_id, - pos: 3, - byte: b'R', - })); - } - - #[test] - fn v_anchor_out_of_bounds_fails() { - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV5*01", b"TGT", Some(5))); - let issues = cfg.validate(); - let v_id = AlleleId::new(1); - assert!(issues.contains(&RefDataValidationIssue::AnchorOutOfBounds { - segment: Segment::V, - allele_id: v_id, - anchor: 5, - len: 3, - })); - } - - #[test] - fn v_anchor_at_boundary_with_no_full_codon_fails() { - // seq.len() == 5, anchor == 3 → codon would span [3, 6) but - // len is 5 → out of bounds. - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV6*01", b"AAATG", Some(3))); - let issues = cfg.validate(); - let v_id = AlleleId::new(1); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::AnchorOutOfBounds { segment, allele_id, .. } - if *segment == Segment::V && *allele_id == v_id - ))); - } - - #[test] - fn v_anchor_non_cys_fails() { - let mut cfg = minimal_vdj(); - // GGG = Gly, not Cys. - let _ = cfg.v_pool.push(v_allele("IGHV7*01", b"GGGAAACCC", Some(0))); - let issues = cfg.validate(); - let v_id = AlleleId::new(1); - assert!(issues.contains(&RefDataValidationIssue::VAnchorNotCys { - allele_id: v_id, - codon: [b'G', b'G', b'G'], - aa: 'G', - severity: RefDataIssueSeverity::Curatable, - })); - } - - #[test] - fn v_anchor_tgc_is_cys_and_passes() { - // TGC is also Cys (degenerate; both TGT and TGC code for C). - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV8*01", b"TGCAAACCC", Some(0))); - let issues = cfg.validate(); - assert!( - issues.is_empty(), - "TGC anchor codon should pass V Cys check, got: {issues:?}" - ); - } - - #[test] - fn j_anchor_unexpected_aa_under_igh_rule_fails() { - // The catalogue claims IGH (rules.j_anchor.expected = ['W']); - // a J allele with anchor codon TTC (Phe = F) is flagged. - // Allele names are no longer used as the source of truth for - // expected J anchors — the rule on the config is. - let mut cfg = minimal_vdj(); - cfg.rules = ReferenceRules::for_locus("IGH"); - let _ = cfg.j_pool.push(j_allele("IGHJ2*01", b"TTCAAACCC", Some(0))); - let issues = cfg.validate(); - let j_id = AlleleId::new(1); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::JAnchorUnexpectedAa { allele_id, aa, .. } - if *allele_id == j_id && *aa == 'F' - ))); - } - - #[test] - fn j_anchor_accepts_f_under_igl_rule() { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules = ReferenceRules::for_locus("IGL"); - let _ = cfg.v_pool.push(v_allele("IGLV1*01", b"TGTAAA", Some(0))); - let _ = cfg.j_pool.push(j_allele("IGLJ1*01", b"TTCAAACCC", Some(0))); - assert!(cfg.validate().is_empty()); - } - - #[test] - fn j_anchor_w_under_igl_rule_is_flagged() { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules = ReferenceRules::for_locus("IGL"); - let _ = cfg.v_pool.push(v_allele("IGLV1*01", b"TGTAAA", Some(0))); - // IGL rule expects F — a W-coded J anchor is flagged. - let _ = cfg.j_pool.push(j_allele("IGLJ2*01", b"TGGAAACCC", Some(0))); - let issues = cfg.validate(); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::JAnchorUnexpectedAa { aa, .. } if *aa == 'W' - ))); - } - - #[test] - fn default_rule_accepts_w_or_f_for_unconfigured_locus() { - // `RefDataConfig::empty` initialises with the lenient default - // (J expects W or F). Any allele-name pattern that the loader - // doesn't recognise lands here. - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); - let _ = cfg.j_pool.push(j_allele("MYJ1*01", b"TGGAAACCC", Some(0))); - assert!( - cfg.validate().is_empty(), - "default-locus J anchor 'W' should be accepted" - ); - - let mut cfg2 = RefDataConfig::empty(ChainType::Vj); - let _ = cfg2.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); - let _ = cfg2.j_pool.push(j_allele("MYJ2*01", b"TTCAAACCC", Some(0))); - assert!( - cfg2.validate().is_empty(), - "default-locus J anchor 'F' should be accepted" - ); - } - - #[test] - fn missing_anchor_on_v_is_reported() { - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV9*01", b"AAACCCGGG", None)); - let issues = cfg.validate(); - let v_id = AlleleId::new(1); - assert!(issues.contains(&RefDataValidationIssue::MissingAnchor { - segment: Segment::V, - allele_id: v_id, - severity: RefDataIssueSeverity::Curatable, - })); - } - - #[test] - fn missing_anchor_on_j_is_reported() { - let mut cfg = minimal_vdj(); - let _ = cfg.j_pool.push(j_allele("IGHJ-orphan*01", b"GGGAAA", None)); - let issues = cfg.validate(); - let j_id = AlleleId::new(1); - assert!(issues.contains(&RefDataValidationIssue::MissingAnchor { - segment: Segment::J, - allele_id: j_id, - severity: RefDataIssueSeverity::Curatable, - })); - } - - #[test] - fn d_anchor_absence_is_not_reported() { - // The default minimal_vdj D allele has no anchor; this must - // not flag a MissingAnchor issue (D anchors are not part of - // the contract). - let cfg = minimal_vdj(); - let issues = cfg.validate(); - assert!(!issues - .iter() - .any(|i| matches!(i, RefDataValidationIssue::MissingAnchor { segment, .. } if *segment == Segment::D))); - } - - #[test] - fn validate_strict_packages_errors() { - let mut cfg = minimal_vdj(); - let _ = cfg.v_pool.push(v_allele("IGHV1-1*01", b"TGTAAA", Some(0))); - let err = cfg.validate_strict().expect_err("duplicate must fail strict"); - assert_eq!(err.issues.len(), 1); - assert!(format!("{err}").contains("duplicate allele name")); - } - - #[test] - fn validate_strict_collects_all_issues_not_just_first() { - let mut cfg = RefDataConfig::empty(ChainType::Vdj); - // Multiple problems at once: empty V, empty J, empty D. - let err = cfg.validate_strict().expect_err("all empty should fail"); - let segs: Vec<_> = err - .issues - .iter() - .filter_map(|i| match i { - RefDataValidationIssue::EmptyRequiredPool { segment } => Some(*segment), - _ => None, - }) - .collect(); - assert!(segs.contains(&Segment::V)); - assert!(segs.contains(&Segment::J)); - assert!(segs.contains(&Segment::D)); - - // Plug V, leave D and J empty: still has two issues. - let _ = cfg.v_pool.push(v_allele("IGHV1*01", b"TGTAAA", Some(0))); - let err2 = cfg.validate_strict().expect_err("two still missing"); - assert_eq!(err2.issues.len(), 2); - } - - // ── Severity classification + mode-aware gating ─────────────── - - /// VJ refdata with only curatable issues (Gly V anchor + no J - /// anchor). Useful to verify mode toggling without bleed-over - /// from Fatal issues. - fn vj_curatable_only() -> RefDataConfig { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg.v_pool.push(v_allele("IGKV1-1*01", b"AAACCCGGG", Some(6))); - let _ = cfg.j_pool.push(Allele { - name: "IGKJ-orphan*01".into(), - gene: "IGKJ-orphan".into(), - seq: b"TTCAAACCC".to_vec(), - segment: Segment::J, - anchor: None, - functional_status: None, - subregions: Vec::new(), - }); - cfg - } - - #[test] - fn validate_with_strict_mode_rejects_curatable_issues() { - let cfg = vj_curatable_only(); - let err = cfg - .validate_with_mode(RefDataValidationMode::Strict) - .expect_err("strict must reject curatable issues"); - let (fatal, curatable) = err.severity_counts(); - assert_eq!(fatal, 0); - assert!(curatable >= 2); - assert_eq!(err.mode, RefDataValidationMode::Strict); - } - - #[test] - fn validate_with_allow_curatable_accepts_curatable_only() { - let cfg = vj_curatable_only(); - cfg.validate_with_mode(RefDataValidationMode::AllowCuratable) - .expect("allow_curatable must accept curatable-only fixtures"); - } - - #[test] - fn validate_with_allow_curatable_still_rejects_fatal() { - let mut cfg = vj_curatable_only(); - // Add an invalid byte allele (Fatal). - let _ = cfg.v_pool.push(v_allele("IGKV2*01", b"TGT.AAACC", Some(0))); - let err = cfg - .validate_with_mode(RefDataValidationMode::AllowCuratable) - .expect_err("allow_curatable must still reject Fatal issues"); - let (fatal, curatable) = err.severity_counts(); - assert!(fatal >= 1, "expected ≥1 fatal, got: {:?}", err.issues); - assert!(curatable >= 1, "expected ≥1 curatable preserved, got: {:?}", err.issues); - assert_eq!(err.mode, RefDataValidationMode::AllowCuratable); - } - - #[test] - fn display_remediation_hint_when_only_curatable_issues_remain() { - let cfg = vj_curatable_only(); - let err = cfg.validate_strict().expect_err("must fail"); - let msg = format!("{err}"); - assert!( - msg.contains("allow_curatable_refdata"), - "msg must suggest allow_curatable_refdata when only curatable: {msg}" - ); - assert!( - msg.contains("pseudogene/ORF"), - "msg must mention pseudogene/ORF: {msg}" - ); - } - - #[test] - fn display_omits_remediation_hint_when_fatal_present() { - let mut cfg = vj_curatable_only(); - let _ = cfg.v_pool.push(v_allele("IGKV2*01", b"TGT.AAACC", Some(0))); - let err = cfg.validate_strict().expect_err("must fail"); - let msg = format!("{err}"); - // Fatal issues exist — opting into allow_curatable_refdata - // wouldn't help, so the hint must NOT appear. - assert!( - !msg.contains("allow_curatable_refdata"), - "msg must NOT suggest allow_curatable_refdata when Fatal issues remain: {msg}" - ); - } - - #[test] - fn each_issue_variant_carries_documented_severity() { - use RefDataIssueSeverity::*; - let aid = AlleleId::new(0); - // Fatal set. - assert_eq!( - RefDataValidationIssue::EmptyRequiredPool { segment: Segment::V }.severity(), - Fatal - ); - assert_eq!( - RefDataValidationIssue::DuplicateAlleleName { - segment: Segment::J, - name: "x".into() - } - .severity(), - Fatal - ); - assert_eq!( - RefDataValidationIssue::InvalidAlleleByte { - segment: Segment::V, - allele_id: aid, - pos: 0, - byte: b'.' - } - .severity(), - Fatal - ); - assert_eq!( - RefDataValidationIssue::AnchorOutOfBounds { - segment: Segment::V, - allele_id: aid, - anchor: 99, - len: 5 - } - .severity(), - Fatal - ); - // Rule-controlled set — severity is whatever the construction - // site stamped on the issue (`severity()` just returns it). - assert_eq!( - RefDataValidationIssue::VAnchorNotCys { - allele_id: aid, - codon: [b'G', b'G', b'G'], - aa: 'G', - severity: Curatable, - } - .severity(), - Curatable - ); - assert_eq!( - RefDataValidationIssue::JAnchorUnexpectedAa { - allele_id: aid, - codon: [b'T', b'T', b'A'], - aa: 'L', - expected: vec!['W'], - severity: Curatable, - } - .severity(), - Curatable - ); - assert_eq!( - RefDataValidationIssue::MissingAnchor { - segment: Segment::V, - allele_id: aid, - severity: Curatable, - } - .severity(), - Curatable - ); - // The rule-controlled variants honour the severity carried - // on the issue itself — a Fatal anchor mismatch (custom rule) - // reports as Fatal. - assert_eq!( - RefDataValidationIssue::MissingAnchor { - segment: Segment::V, - allele_id: aid, - severity: Fatal, - } - .severity(), - Fatal - ); - } - - // ── ReferenceRules v1 — configurable anchor + alphabet ────── - - #[test] - fn default_rules_match_legacy_behavior() { - // V → ['C'], J → ['W', 'F'], alphabet ACGTN, all Curatable. - let rules = ReferenceRules::default(); - assert_eq!(rules.v_anchor.expected_amino_acids, vec!['C']); - assert_eq!(rules.j_anchor.expected_amino_acids, vec!['W', 'F']); - assert!(rules.v_anchor.required); - assert!(rules.j_anchor.required); - for &b in b"acgtnACGTN" { - assert!(rules.alphabet.is_allowed(b), "should allow {}", b as char); - } - assert!(!rules.alphabet.is_allowed(b'.')); - assert!(!rules.alphabet.is_allowed(b'R')); - } - - #[test] - fn for_locus_returns_locus_appropriate_j_rule() { - assert_eq!( - ReferenceRules::for_locus("IGH").j_anchor.expected_amino_acids, - vec!['W'] - ); - assert_eq!( - ReferenceRules::for_locus("IGK").j_anchor.expected_amino_acids, - vec!['F'] - ); - assert_eq!( - ReferenceRules::for_locus("IGL").j_anchor.expected_amino_acids, - vec!['F'] - ); - for tr in ["TRA", "TRB", "TRG", "TRD"] { - assert_eq!( - ReferenceRules::for_locus(tr).j_anchor.expected_amino_acids, - vec!['F'], - "{tr} J should expect F" - ); - } - // Unknown prefix → default lenient set. - assert_eq!( - ReferenceRules::for_locus("XXX").j_anchor.expected_amino_acids, - vec!['W', 'F'] - ); - } - - #[test] - fn custom_j_rule_y_accepts_tat_and_tac() { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.j_anchor.expected_amino_acids = vec!['Y']; - let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); - // TAT → Y. - let _ = cfg.j_pool.push(j_allele("MYJ1*01", b"TATAAACCC", Some(0))); - assert!(cfg.validate().is_empty()); - // TAC → Y. - let mut cfg2 = cfg.clone(); - cfg2.j_pool = crate::refdata::AllelePool::new(); - let _ = cfg2.j_pool.push(j_allele("MYJ2*01", b"TACAAACCC", Some(0))); - assert!(cfg2.validate().is_empty()); - } - - #[test] - fn custom_j_rule_y_rejects_tgg() { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.j_anchor.expected_amino_acids = vec!['Y']; - let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); - // TGG → W. Under a Y-only J rule, this is flagged. - let _ = cfg.j_pool.push(j_allele("MYJ-bad*01", b"TGGAAACCC", Some(0))); - let issues = cfg.validate(); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::JAnchorUnexpectedAa { aa, .. } if *aa == 'W' - ))); - } - - #[test] - fn anchor_required_false_suppresses_missing_anchor_issue() { - // Build a config where the J rule says anchor is optional. - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.j_anchor.required = false; - let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); - // J allele with no anchor → would normally emit MissingAnchor. - let _ = cfg.j_pool.push(j_allele("MYJ-orphan*01", b"GGG", None)); - assert!( - cfg.validate().is_empty(), - "rule.required=false must suppress MissingAnchor" - ); - } - - #[test] - fn missing_severity_follows_rule() { - // Bump missing_severity to Fatal — the resulting issue's - // severity() reflects that, and AllowCuratable mode still - // rejects. - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.j_anchor.missing_severity = RefDataIssueSeverity::Fatal; - let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); - let _ = cfg.j_pool.push(j_allele("MYJ*01", b"TGGAAA", None)); - let issues = cfg.validate(); - let missing = issues - .iter() - .find(|i| matches!(i, RefDataValidationIssue::MissingAnchor { .. })) - .expect("must emit MissingAnchor"); - assert_eq!(missing.severity(), RefDataIssueSeverity::Fatal); - - // AllowCuratable still rejects. - cfg.validate_with_mode(RefDataValidationMode::AllowCuratable) - .expect_err("Fatal missing-anchor must not be opt-outable"); - } - - #[test] - fn mismatch_severity_follows_rule() { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.v_anchor.mismatch_severity = RefDataIssueSeverity::Fatal; - // GGG (Gly) at V anchor. - let _ = cfg.v_pool.push(v_allele("MYV-bad*01", b"GGGAAACCC", Some(0))); - let _ = cfg.j_pool.push(j_allele("MYJ*01", b"TGGAAA", Some(0))); - let issues = cfg.validate(); - let mismatch = issues - .iter() - .find(|i| matches!(i, RefDataValidationIssue::VAnchorNotCys { .. })) - .expect("must emit VAnchorNotCys"); - assert_eq!(mismatch.severity(), RefDataIssueSeverity::Fatal); - - // AllowCuratable still rejects this catalogue. - cfg.validate_with_mode(RefDataValidationMode::AllowCuratable) - .expect_err("Fatal V anchor mismatch must not be opt-outable"); - } - - #[test] - fn invalid_byte_remains_fatal_regardless_of_rule() { - // Even with an extremely permissive anchor rule, an invalid - // sequence byte stays Fatal — structural problems aren't - // rule-controlled. - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.v_anchor.expected_amino_acids = vec!['A', 'C', 'D', 'E', 'F']; - cfg.rules.j_anchor.expected_amino_acids = vec!['A', 'C', 'D', 'E', 'F']; - let _ = cfg.v_pool.push(v_allele("MYV-bad*01", b"TGT.AAACC", Some(0))); - let _ = cfg.j_pool.push(j_allele("MYJ*01", b"TGGAAA", Some(0))); - let issues = cfg.validate(); - let bad = issues - .iter() - .find(|i| matches!(i, RefDataValidationIssue::InvalidAlleleByte { .. })) - .expect("must emit InvalidAlleleByte"); - assert_eq!(bad.severity(), RefDataIssueSeverity::Fatal); - } - - #[test] - fn custom_alphabet_extends_allowed_set() { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - cfg.rules.alphabet = ReferenceAlphabet { - allowed: vec![b'A', b'C', b'G', b'T', b'N', b'R'], // R = puRine ambig - }; - let _ = cfg.v_pool.push(v_allele("MYV*01", b"TGTRAA", Some(0))); - let _ = cfg.j_pool.push(j_allele("MYJ*01", b"TGGAAA", Some(0))); - let issues = cfg.validate(); - // R is now in-alphabet → no InvalidAlleleByte issue. - assert!( - !issues - .iter() - .any(|i| matches!(i, RefDataValidationIssue::InvalidAlleleByte { .. })), - "R should be allowed under the extended alphabet; got: {issues:?}" - ); - } - - #[test] - fn reference_alphabet_is_case_insensitive() { - let alphabet = ReferenceAlphabet::default(); - assert!(alphabet.is_allowed(b'a')); - assert!(alphabet.is_allowed(b'A')); - assert!(alphabet.is_allowed(b'n')); - assert!(!alphabet.is_allowed(b'.')); - } - - // ── Identity ↔ chain-type mismatch (Reference Identity slice) ── - - fn vdj_cfg_with_identity_locus(locus: &str) -> RefDataConfig { - let mut cfg = minimal_vdj(); - cfg.identity.locus = Some(locus.to_string()); - cfg - } - - fn vj_cfg_with_identity_locus(locus: &str) -> RefDataConfig { - let mut cfg = RefDataConfig::empty(ChainType::Vj); - let _ = cfg.v_pool.push(v_allele("IGKV1-1*01", b"TGTAAACCC", Some(0))); - let _ = cfg.j_pool.push(j_allele("IGKJ1*01", b"TTCAAACCC", Some(0))); - cfg.identity.locus = Some(locus.to_string()); - cfg - } - - #[test] - fn vj_cartridge_with_igh_locus_is_flagged() { - // VJ topology declaring an IGH (VDJ) locus is structurally - // wrong — recombination shape and locus disagree. - let cfg = vj_cfg_with_identity_locus("IGH"); - let issues = cfg.validate(); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::LocusChainTypeMismatch { locus, .. } if locus == "IGH" - ))); - } - - #[test] - fn vj_cartridge_with_igk_locus_passes() { - // Locus matches topology — no LocusChainTypeMismatch. - let cfg = vj_cfg_with_identity_locus("IGK"); - let issues = cfg.validate(); - assert!(!issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::LocusChainTypeMismatch { .. } - ))); - } - - #[test] - fn vdj_cartridge_with_igk_locus_is_flagged() { - let cfg = vdj_cfg_with_identity_locus("IGK"); - assert!(cfg.validate().iter().any(|i| matches!( - i, - RefDataValidationIssue::LocusChainTypeMismatch { locus, .. } if locus == "IGK" - ))); - } - - #[test] - fn locus_chain_type_mismatch_is_fatal() { - let cfg = vj_cfg_with_identity_locus("IGH"); - let issue = cfg - .validate() - .into_iter() - .find(|i| matches!(i, RefDataValidationIssue::LocusChainTypeMismatch { .. })) - .expect("must emit LocusChainTypeMismatch"); - assert_eq!(issue.severity(), RefDataIssueSeverity::Fatal); - // AllowCuratable mode cannot opt this out. - cfg.validate_with_mode(RefDataValidationMode::AllowCuratable) - .expect_err("Fatal mismatch must not be opt-outable"); - } - - #[test] - fn unknown_locus_does_not_fail_validation() { - let cfg = vj_cfg_with_identity_locus("XYZ"); - assert!(!cfg.validate().iter().any(|i| matches!( - i, - RefDataValidationIssue::LocusChainTypeMismatch { .. } - ))); - } - - #[test] - fn empty_identity_does_not_fail_validation() { - // No locus declared → no mismatch issue regardless of chain. - let cfg_vj = RefDataConfig::empty(ChainType::Vj); - assert!(!cfg_vj.validate().iter().any(|i| matches!( - i, - RefDataValidationIssue::LocusChainTypeMismatch { .. } - ))); - let cfg_vdj = RefDataConfig::empty(ChainType::Vdj); - assert!(!cfg_vdj.validate().iter().any(|i| matches!( - i, - RefDataValidationIssue::LocusChainTypeMismatch { .. } - ))); - } - - #[test] - fn locus_is_case_insensitive() { - // Lowercase / mixed-case locus still flags correctly. - let cfg = vj_cfg_with_identity_locus("igh"); - let issues = cfg.validate(); - assert!(issues.iter().any(|i| matches!( - i, - RefDataValidationIssue::LocusChainTypeMismatch { locus, .. } if locus == "IGH" - ))); - } -} +mod tests; diff --git a/engine_rs/src/refdata/validation/tests.rs b/engine_rs/src/refdata/validation/tests.rs new file mode 100644 index 0000000..4a32163 --- /dev/null +++ b/engine_rs/src/refdata/validation/tests.rs @@ -0,0 +1,810 @@ + use super::*; + use crate::refdata::{Allele, AlleleId, ChainType, RefDataConfig}; + + fn v_allele(name: &str, seq: &[u8], anchor: Option) -> Allele { + Allele { + name: name.to_string(), + gene: name.split('*').next().unwrap_or(name).to_string(), + seq: seq.to_vec(), + segment: Segment::V, + anchor, + functional_status: None, + subregions: Vec::new(), + } + } + fn d_allele(name: &str, seq: &[u8]) -> Allele { + Allele { + name: name.to_string(), + gene: name.split('*').next().unwrap_or(name).to_string(), + seq: seq.to_vec(), + segment: Segment::D, + anchor: None, + functional_status: None, + subregions: Vec::new(), + } + } + fn j_allele(name: &str, seq: &[u8], anchor: Option) -> Allele { + Allele { + name: name.to_string(), + gene: name.split('*').next().unwrap_or(name).to_string(), + seq: seq.to_vec(), + segment: Segment::J, + anchor, + functional_status: None, + subregions: Vec::new(), + } + } + + /// Build a minimal valid VDJ refdata so each test only varies + /// the one field under examination. + fn minimal_vdj() -> RefDataConfig { + let mut cfg = RefDataConfig::empty(ChainType::Vdj); + // V allele: anchor codon TGT (Cys) at position 0. + let _ = cfg.v_pool.push(v_allele("IGHV1-1*01", b"TGTAAACCC", Some(0))); + // D allele: no anchor. + let _ = cfg.d_pool.push(d_allele("IGHD1-1*01", b"GGGCCCAAA")); + // J allele: anchor codon TGG (Trp = W) at position 0. + let _ = cfg.j_pool.push(j_allele("IGHJ1*01", b"TGGAAACCC", Some(0))); + cfg + } + + fn minimal_vj() -> RefDataConfig { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + let _ = cfg.v_pool.push(v_allele("IGKV1-1*01", b"TGTAAACCC", Some(0))); + // IGK J anchor codon TTC (Phe = F). + let _ = cfg.j_pool.push(j_allele("IGKJ1*01", b"TTCAAACCC", Some(0))); + cfg + } + + #[test] + fn minimal_vdj_validates_clean() { + assert!(minimal_vdj().validate().is_empty()); + assert!(minimal_vdj().validate_strict().is_ok()); + } + + #[test] + fn minimal_vj_validates_clean() { + assert!(minimal_vj().validate().is_empty()); + } + + #[test] + fn empty_v_pool_fails_vj() { + let mut cfg = minimal_vj(); + cfg.v_pool = super::AllelePool::new(); + let issues = cfg.validate(); + assert!(issues.contains(&RefDataValidationIssue::EmptyRequiredPool { + segment: Segment::V, + })); + } + + #[test] + fn empty_v_pool_fails_vdj() { + let mut cfg = minimal_vdj(); + cfg.v_pool = super::AllelePool::new(); + let issues = cfg.validate(); + assert!(issues.contains(&RefDataValidationIssue::EmptyRequiredPool { + segment: Segment::V, + })); + } + + #[test] + fn empty_j_pool_fails_vj() { + let mut cfg = minimal_vj(); + cfg.j_pool = super::AllelePool::new(); + let issues = cfg.validate(); + assert!(issues.contains(&RefDataValidationIssue::EmptyRequiredPool { + segment: Segment::J, + })); + } + + #[test] + fn empty_d_pool_allowed_vj() { + // VJ chains have empty D pools by construction. This must + // not flag an EmptyRequiredPool for D. + let cfg = minimal_vj(); + let issues = cfg.validate(); + assert!(!issues + .iter() + .any(|i| matches!(i, RefDataValidationIssue::EmptyRequiredPool { segment } if *segment == Segment::D))); + } + + #[test] + fn empty_d_pool_rejected_vdj() { + let mut cfg = minimal_vdj(); + cfg.d_pool = super::AllelePool::new(); + let issues = cfg.validate(); + assert!(issues.contains(&RefDataValidationIssue::EmptyRequiredPool { + segment: Segment::D, + })); + } + + #[test] + fn duplicate_v_allele_names_fail() { + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV1-1*01", b"TGTAAA", Some(0))); + let issues = cfg.validate(); + assert!(issues.contains(&RefDataValidationIssue::DuplicateAlleleName { + segment: Segment::V, + name: "IGHV1-1*01".to_string(), + })); + } + + #[test] + fn duplicate_j_allele_names_fail() { + let mut cfg = minimal_vdj(); + let _ = cfg.j_pool.push(j_allele("IGHJ1*01", b"TGGAAA", Some(0))); + let issues = cfg.validate(); + assert!(issues.contains(&RefDataValidationIssue::DuplicateAlleleName { + segment: Segment::J, + name: "IGHJ1*01".to_string(), + })); + } + + #[test] + fn invalid_byte_in_v_seq_fails_with_exact_position() { + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV2*01", b"TGT.AAA", Some(0))); + let issues = cfg.validate(); + let bad_v_id = AlleleId::new(1); + let expected = RefDataValidationIssue::InvalidAlleleByte { + segment: Segment::V, + allele_id: bad_v_id, + pos: 3, + byte: b'.', + }; + assert!( + issues.contains(&expected), + "expected invalid-byte issue at pos 3, got: {issues:?}" + ); + } + + #[test] + fn lowercase_bases_accepted() { + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV3*01", b"tgtaaaccc", Some(0))); + // Lowercase A/C/G/T/N is allowed. The anchor codon `tgt` still + // translates to Cys, so V anchor check also passes. + let issues = cfg.validate(); + assert!( + issues.is_empty(), + "lowercase canonical bases should validate, got: {issues:?}" + ); + } + + #[test] + fn iupac_ambiguity_other_than_n_fails() { + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV4*01", b"TGTRAACCC", Some(0))); + let issues = cfg.validate(); + let v_id = AlleleId::new(1); + assert!(issues.contains(&RefDataValidationIssue::InvalidAlleleByte { + segment: Segment::V, + allele_id: v_id, + pos: 3, + byte: b'R', + })); + } + + #[test] + fn v_anchor_out_of_bounds_fails() { + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV5*01", b"TGT", Some(5))); + let issues = cfg.validate(); + let v_id = AlleleId::new(1); + assert!(issues.contains(&RefDataValidationIssue::AnchorOutOfBounds { + segment: Segment::V, + allele_id: v_id, + anchor: 5, + len: 3, + })); + } + + #[test] + fn v_anchor_at_boundary_with_no_full_codon_fails() { + // seq.len() == 5, anchor == 3 → codon would span [3, 6) but + // len is 5 → out of bounds. + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV6*01", b"AAATG", Some(3))); + let issues = cfg.validate(); + let v_id = AlleleId::new(1); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::AnchorOutOfBounds { segment, allele_id, .. } + if *segment == Segment::V && *allele_id == v_id + ))); + } + + #[test] + fn v_anchor_non_cys_fails() { + let mut cfg = minimal_vdj(); + // GGG = Gly, not Cys. + let _ = cfg.v_pool.push(v_allele("IGHV7*01", b"GGGAAACCC", Some(0))); + let issues = cfg.validate(); + let v_id = AlleleId::new(1); + assert!(issues.contains(&RefDataValidationIssue::VAnchorNotCys { + allele_id: v_id, + codon: [b'G', b'G', b'G'], + aa: 'G', + severity: RefDataIssueSeverity::Curatable, + })); + } + + #[test] + fn v_anchor_tgc_is_cys_and_passes() { + // TGC is also Cys (degenerate; both TGT and TGC code for C). + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV8*01", b"TGCAAACCC", Some(0))); + let issues = cfg.validate(); + assert!( + issues.is_empty(), + "TGC anchor codon should pass V Cys check, got: {issues:?}" + ); + } + + #[test] + fn j_anchor_unexpected_aa_under_igh_rule_fails() { + // The catalogue claims IGH (rules.j_anchor.expected = ['W']); + // a J allele with anchor codon TTC (Phe = F) is flagged. + // Allele names are no longer used as the source of truth for + // expected J anchors — the rule on the config is. + let mut cfg = minimal_vdj(); + cfg.rules = ReferenceRules::for_locus("IGH"); + let _ = cfg.j_pool.push(j_allele("IGHJ2*01", b"TTCAAACCC", Some(0))); + let issues = cfg.validate(); + let j_id = AlleleId::new(1); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::JAnchorUnexpectedAa { allele_id, aa, .. } + if *allele_id == j_id && *aa == 'F' + ))); + } + + #[test] + fn j_anchor_accepts_f_under_igl_rule() { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules = ReferenceRules::for_locus("IGL"); + let _ = cfg.v_pool.push(v_allele("IGLV1*01", b"TGTAAA", Some(0))); + let _ = cfg.j_pool.push(j_allele("IGLJ1*01", b"TTCAAACCC", Some(0))); + assert!(cfg.validate().is_empty()); + } + + #[test] + fn j_anchor_w_under_igl_rule_is_flagged() { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules = ReferenceRules::for_locus("IGL"); + let _ = cfg.v_pool.push(v_allele("IGLV1*01", b"TGTAAA", Some(0))); + // IGL rule expects F — a W-coded J anchor is flagged. + let _ = cfg.j_pool.push(j_allele("IGLJ2*01", b"TGGAAACCC", Some(0))); + let issues = cfg.validate(); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::JAnchorUnexpectedAa { aa, .. } if *aa == 'W' + ))); + } + + #[test] + fn default_rule_accepts_w_or_f_for_unconfigured_locus() { + // `RefDataConfig::empty` initialises with the lenient default + // (J expects W or F). Any allele-name pattern that the loader + // doesn't recognise lands here. + let mut cfg = RefDataConfig::empty(ChainType::Vj); + let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); + let _ = cfg.j_pool.push(j_allele("MYJ1*01", b"TGGAAACCC", Some(0))); + assert!( + cfg.validate().is_empty(), + "default-locus J anchor 'W' should be accepted" + ); + + let mut cfg2 = RefDataConfig::empty(ChainType::Vj); + let _ = cfg2.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); + let _ = cfg2.j_pool.push(j_allele("MYJ2*01", b"TTCAAACCC", Some(0))); + assert!( + cfg2.validate().is_empty(), + "default-locus J anchor 'F' should be accepted" + ); + } + + #[test] + fn missing_anchor_on_v_is_reported() { + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV9*01", b"AAACCCGGG", None)); + let issues = cfg.validate(); + let v_id = AlleleId::new(1); + assert!(issues.contains(&RefDataValidationIssue::MissingAnchor { + segment: Segment::V, + allele_id: v_id, + severity: RefDataIssueSeverity::Curatable, + })); + } + + #[test] + fn missing_anchor_on_j_is_reported() { + let mut cfg = minimal_vdj(); + let _ = cfg.j_pool.push(j_allele("IGHJ-orphan*01", b"GGGAAA", None)); + let issues = cfg.validate(); + let j_id = AlleleId::new(1); + assert!(issues.contains(&RefDataValidationIssue::MissingAnchor { + segment: Segment::J, + allele_id: j_id, + severity: RefDataIssueSeverity::Curatable, + })); + } + + #[test] + fn d_anchor_absence_is_not_reported() { + // The default minimal_vdj D allele has no anchor; this must + // not flag a MissingAnchor issue (D anchors are not part of + // the contract). + let cfg = minimal_vdj(); + let issues = cfg.validate(); + assert!(!issues + .iter() + .any(|i| matches!(i, RefDataValidationIssue::MissingAnchor { segment, .. } if *segment == Segment::D))); + } + + #[test] + fn validate_strict_packages_errors() { + let mut cfg = minimal_vdj(); + let _ = cfg.v_pool.push(v_allele("IGHV1-1*01", b"TGTAAA", Some(0))); + let err = cfg.validate_strict().expect_err("duplicate must fail strict"); + assert_eq!(err.issues.len(), 1); + assert!(format!("{err}").contains("duplicate allele name")); + } + + #[test] + fn validate_strict_collects_all_issues_not_just_first() { + let mut cfg = RefDataConfig::empty(ChainType::Vdj); + // Multiple problems at once: empty V, empty J, empty D. + let err = cfg.validate_strict().expect_err("all empty should fail"); + let segs: Vec<_> = err + .issues + .iter() + .filter_map(|i| match i { + RefDataValidationIssue::EmptyRequiredPool { segment } => Some(*segment), + _ => None, + }) + .collect(); + assert!(segs.contains(&Segment::V)); + assert!(segs.contains(&Segment::J)); + assert!(segs.contains(&Segment::D)); + + // Plug V, leave D and J empty: still has two issues. + let _ = cfg.v_pool.push(v_allele("IGHV1*01", b"TGTAAA", Some(0))); + let err2 = cfg.validate_strict().expect_err("two still missing"); + assert_eq!(err2.issues.len(), 2); + } + + // ── Severity classification + mode-aware gating ─────────────── + + /// VJ refdata with only curatable issues (Gly V anchor + no J + /// anchor). Useful to verify mode toggling without bleed-over + /// from Fatal issues. + fn vj_curatable_only() -> RefDataConfig { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + let _ = cfg.v_pool.push(v_allele("IGKV1-1*01", b"AAACCCGGG", Some(6))); + let _ = cfg.j_pool.push(Allele { + name: "IGKJ-orphan*01".into(), + gene: "IGKJ-orphan".into(), + seq: b"TTCAAACCC".to_vec(), + segment: Segment::J, + anchor: None, + functional_status: None, + subregions: Vec::new(), + }); + cfg + } + + #[test] + fn validate_with_strict_mode_rejects_curatable_issues() { + let cfg = vj_curatable_only(); + let err = cfg + .validate_with_mode(RefDataValidationMode::Strict) + .expect_err("strict must reject curatable issues"); + let (fatal, curatable) = err.severity_counts(); + assert_eq!(fatal, 0); + assert!(curatable >= 2); + assert_eq!(err.mode, RefDataValidationMode::Strict); + } + + #[test] + fn validate_with_allow_curatable_accepts_curatable_only() { + let cfg = vj_curatable_only(); + cfg.validate_with_mode(RefDataValidationMode::AllowCuratable) + .expect("allow_curatable must accept curatable-only fixtures"); + } + + #[test] + fn validate_with_allow_curatable_still_rejects_fatal() { + let mut cfg = vj_curatable_only(); + // Add an invalid byte allele (Fatal). + let _ = cfg.v_pool.push(v_allele("IGKV2*01", b"TGT.AAACC", Some(0))); + let err = cfg + .validate_with_mode(RefDataValidationMode::AllowCuratable) + .expect_err("allow_curatable must still reject Fatal issues"); + let (fatal, curatable) = err.severity_counts(); + assert!(fatal >= 1, "expected ≥1 fatal, got: {:?}", err.issues); + assert!(curatable >= 1, "expected ≥1 curatable preserved, got: {:?}", err.issues); + assert_eq!(err.mode, RefDataValidationMode::AllowCuratable); + } + + #[test] + fn display_remediation_hint_when_only_curatable_issues_remain() { + let cfg = vj_curatable_only(); + let err = cfg.validate_strict().expect_err("must fail"); + let msg = format!("{err}"); + assert!( + msg.contains("allow_curatable_refdata"), + "msg must suggest allow_curatable_refdata when only curatable: {msg}" + ); + assert!( + msg.contains("pseudogene/ORF"), + "msg must mention pseudogene/ORF: {msg}" + ); + } + + #[test] + fn display_omits_remediation_hint_when_fatal_present() { + let mut cfg = vj_curatable_only(); + let _ = cfg.v_pool.push(v_allele("IGKV2*01", b"TGT.AAACC", Some(0))); + let err = cfg.validate_strict().expect_err("must fail"); + let msg = format!("{err}"); + // Fatal issues exist — opting into allow_curatable_refdata + // wouldn't help, so the hint must NOT appear. + assert!( + !msg.contains("allow_curatable_refdata"), + "msg must NOT suggest allow_curatable_refdata when Fatal issues remain: {msg}" + ); + } + + #[test] + fn each_issue_variant_carries_documented_severity() { + use RefDataIssueSeverity::*; + let aid = AlleleId::new(0); + // Fatal set. + assert_eq!( + RefDataValidationIssue::EmptyRequiredPool { segment: Segment::V }.severity(), + Fatal + ); + assert_eq!( + RefDataValidationIssue::DuplicateAlleleName { + segment: Segment::J, + name: "x".into() + } + .severity(), + Fatal + ); + assert_eq!( + RefDataValidationIssue::InvalidAlleleByte { + segment: Segment::V, + allele_id: aid, + pos: 0, + byte: b'.' + } + .severity(), + Fatal + ); + assert_eq!( + RefDataValidationIssue::AnchorOutOfBounds { + segment: Segment::V, + allele_id: aid, + anchor: 99, + len: 5 + } + .severity(), + Fatal + ); + // Rule-controlled set — severity is whatever the construction + // site stamped on the issue (`severity()` just returns it). + assert_eq!( + RefDataValidationIssue::VAnchorNotCys { + allele_id: aid, + codon: [b'G', b'G', b'G'], + aa: 'G', + severity: Curatable, + } + .severity(), + Curatable + ); + assert_eq!( + RefDataValidationIssue::JAnchorUnexpectedAa { + allele_id: aid, + codon: [b'T', b'T', b'A'], + aa: 'L', + expected: vec!['W'], + severity: Curatable, + } + .severity(), + Curatable + ); + assert_eq!( + RefDataValidationIssue::MissingAnchor { + segment: Segment::V, + allele_id: aid, + severity: Curatable, + } + .severity(), + Curatable + ); + // The rule-controlled variants honour the severity carried + // on the issue itself — a Fatal anchor mismatch (custom rule) + // reports as Fatal. + assert_eq!( + RefDataValidationIssue::MissingAnchor { + segment: Segment::V, + allele_id: aid, + severity: Fatal, + } + .severity(), + Fatal + ); + } + + // ── ReferenceRules v1 — configurable anchor + alphabet ────── + + #[test] + fn default_rules_match_legacy_behavior() { + // V → ['C'], J → ['W', 'F'], alphabet ACGTN, all Curatable. + let rules = ReferenceRules::default(); + assert_eq!(rules.v_anchor.expected_amino_acids, vec!['C']); + assert_eq!(rules.j_anchor.expected_amino_acids, vec!['W', 'F']); + assert!(rules.v_anchor.required); + assert!(rules.j_anchor.required); + for &b in b"acgtnACGTN" { + assert!(rules.alphabet.is_allowed(b), "should allow {}", b as char); + } + assert!(!rules.alphabet.is_allowed(b'.')); + assert!(!rules.alphabet.is_allowed(b'R')); + } + + #[test] + fn for_locus_returns_locus_appropriate_j_rule() { + assert_eq!( + ReferenceRules::for_locus("IGH").j_anchor.expected_amino_acids, + vec!['W'] + ); + assert_eq!( + ReferenceRules::for_locus("IGK").j_anchor.expected_amino_acids, + vec!['F'] + ); + assert_eq!( + ReferenceRules::for_locus("IGL").j_anchor.expected_amino_acids, + vec!['F'] + ); + for tr in ["TRA", "TRB", "TRG", "TRD"] { + assert_eq!( + ReferenceRules::for_locus(tr).j_anchor.expected_amino_acids, + vec!['F'], + "{tr} J should expect F" + ); + } + // Unknown prefix → default lenient set. + assert_eq!( + ReferenceRules::for_locus("XXX").j_anchor.expected_amino_acids, + vec!['W', 'F'] + ); + } + + #[test] + fn custom_j_rule_y_accepts_tat_and_tac() { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.j_anchor.expected_amino_acids = vec!['Y']; + let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); + // TAT → Y. + let _ = cfg.j_pool.push(j_allele("MYJ1*01", b"TATAAACCC", Some(0))); + assert!(cfg.validate().is_empty()); + // TAC → Y. + let mut cfg2 = cfg.clone(); + cfg2.j_pool = crate::refdata::AllelePool::new(); + let _ = cfg2.j_pool.push(j_allele("MYJ2*01", b"TACAAACCC", Some(0))); + assert!(cfg2.validate().is_empty()); + } + + #[test] + fn custom_j_rule_y_rejects_tgg() { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.j_anchor.expected_amino_acids = vec!['Y']; + let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); + // TGG → W. Under a Y-only J rule, this is flagged. + let _ = cfg.j_pool.push(j_allele("MYJ-bad*01", b"TGGAAACCC", Some(0))); + let issues = cfg.validate(); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::JAnchorUnexpectedAa { aa, .. } if *aa == 'W' + ))); + } + + #[test] + fn anchor_required_false_suppresses_missing_anchor_issue() { + // Build a config where the J rule says anchor is optional. + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.j_anchor.required = false; + let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); + // J allele with no anchor → would normally emit MissingAnchor. + let _ = cfg.j_pool.push(j_allele("MYJ-orphan*01", b"GGG", None)); + assert!( + cfg.validate().is_empty(), + "rule.required=false must suppress MissingAnchor" + ); + } + + #[test] + fn missing_severity_follows_rule() { + // Bump missing_severity to Fatal — the resulting issue's + // severity() reflects that, and AllowCuratable mode still + // rejects. + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.j_anchor.missing_severity = RefDataIssueSeverity::Fatal; + let _ = cfg.v_pool.push(v_allele("MYV1*01", b"TGTAAA", Some(0))); + let _ = cfg.j_pool.push(j_allele("MYJ*01", b"TGGAAA", None)); + let issues = cfg.validate(); + let missing = issues + .iter() + .find(|i| matches!(i, RefDataValidationIssue::MissingAnchor { .. })) + .expect("must emit MissingAnchor"); + assert_eq!(missing.severity(), RefDataIssueSeverity::Fatal); + + // AllowCuratable still rejects. + cfg.validate_with_mode(RefDataValidationMode::AllowCuratable) + .expect_err("Fatal missing-anchor must not be opt-outable"); + } + + #[test] + fn mismatch_severity_follows_rule() { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.v_anchor.mismatch_severity = RefDataIssueSeverity::Fatal; + // GGG (Gly) at V anchor. + let _ = cfg.v_pool.push(v_allele("MYV-bad*01", b"GGGAAACCC", Some(0))); + let _ = cfg.j_pool.push(j_allele("MYJ*01", b"TGGAAA", Some(0))); + let issues = cfg.validate(); + let mismatch = issues + .iter() + .find(|i| matches!(i, RefDataValidationIssue::VAnchorNotCys { .. })) + .expect("must emit VAnchorNotCys"); + assert_eq!(mismatch.severity(), RefDataIssueSeverity::Fatal); + + // AllowCuratable still rejects this catalogue. + cfg.validate_with_mode(RefDataValidationMode::AllowCuratable) + .expect_err("Fatal V anchor mismatch must not be opt-outable"); + } + + #[test] + fn invalid_byte_remains_fatal_regardless_of_rule() { + // Even with an extremely permissive anchor rule, an invalid + // sequence byte stays Fatal — structural problems aren't + // rule-controlled. + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.v_anchor.expected_amino_acids = vec!['A', 'C', 'D', 'E', 'F']; + cfg.rules.j_anchor.expected_amino_acids = vec!['A', 'C', 'D', 'E', 'F']; + let _ = cfg.v_pool.push(v_allele("MYV-bad*01", b"TGT.AAACC", Some(0))); + let _ = cfg.j_pool.push(j_allele("MYJ*01", b"TGGAAA", Some(0))); + let issues = cfg.validate(); + let bad = issues + .iter() + .find(|i| matches!(i, RefDataValidationIssue::InvalidAlleleByte { .. })) + .expect("must emit InvalidAlleleByte"); + assert_eq!(bad.severity(), RefDataIssueSeverity::Fatal); + } + + #[test] + fn custom_alphabet_extends_allowed_set() { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + cfg.rules.alphabet = ReferenceAlphabet { + allowed: vec![b'A', b'C', b'G', b'T', b'N', b'R'], // R = puRine ambig + }; + let _ = cfg.v_pool.push(v_allele("MYV*01", b"TGTRAA", Some(0))); + let _ = cfg.j_pool.push(j_allele("MYJ*01", b"TGGAAA", Some(0))); + let issues = cfg.validate(); + // R is now in-alphabet → no InvalidAlleleByte issue. + assert!( + !issues + .iter() + .any(|i| matches!(i, RefDataValidationIssue::InvalidAlleleByte { .. })), + "R should be allowed under the extended alphabet; got: {issues:?}" + ); + } + + #[test] + fn reference_alphabet_is_case_insensitive() { + let alphabet = ReferenceAlphabet::default(); + assert!(alphabet.is_allowed(b'a')); + assert!(alphabet.is_allowed(b'A')); + assert!(alphabet.is_allowed(b'n')); + assert!(!alphabet.is_allowed(b'.')); + } + + // ── Identity ↔ chain-type mismatch (Reference Identity slice) ── + + fn vdj_cfg_with_identity_locus(locus: &str) -> RefDataConfig { + let mut cfg = minimal_vdj(); + cfg.identity.locus = Some(locus.to_string()); + cfg + } + + fn vj_cfg_with_identity_locus(locus: &str) -> RefDataConfig { + let mut cfg = RefDataConfig::empty(ChainType::Vj); + let _ = cfg.v_pool.push(v_allele("IGKV1-1*01", b"TGTAAACCC", Some(0))); + let _ = cfg.j_pool.push(j_allele("IGKJ1*01", b"TTCAAACCC", Some(0))); + cfg.identity.locus = Some(locus.to_string()); + cfg + } + + #[test] + fn vj_cartridge_with_igh_locus_is_flagged() { + // VJ topology declaring an IGH (VDJ) locus is structurally + // wrong — recombination shape and locus disagree. + let cfg = vj_cfg_with_identity_locus("IGH"); + let issues = cfg.validate(); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::LocusChainTypeMismatch { locus, .. } if locus == "IGH" + ))); + } + + #[test] + fn vj_cartridge_with_igk_locus_passes() { + // Locus matches topology — no LocusChainTypeMismatch. + let cfg = vj_cfg_with_identity_locus("IGK"); + let issues = cfg.validate(); + assert!(!issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::LocusChainTypeMismatch { .. } + ))); + } + + #[test] + fn vdj_cartridge_with_igk_locus_is_flagged() { + let cfg = vdj_cfg_with_identity_locus("IGK"); + assert!(cfg.validate().iter().any(|i| matches!( + i, + RefDataValidationIssue::LocusChainTypeMismatch { locus, .. } if locus == "IGK" + ))); + } + + #[test] + fn locus_chain_type_mismatch_is_fatal() { + let cfg = vj_cfg_with_identity_locus("IGH"); + let issue = cfg + .validate() + .into_iter() + .find(|i| matches!(i, RefDataValidationIssue::LocusChainTypeMismatch { .. })) + .expect("must emit LocusChainTypeMismatch"); + assert_eq!(issue.severity(), RefDataIssueSeverity::Fatal); + // AllowCuratable mode cannot opt this out. + cfg.validate_with_mode(RefDataValidationMode::AllowCuratable) + .expect_err("Fatal mismatch must not be opt-outable"); + } + + #[test] + fn unknown_locus_does_not_fail_validation() { + let cfg = vj_cfg_with_identity_locus("XYZ"); + assert!(!cfg.validate().iter().any(|i| matches!( + i, + RefDataValidationIssue::LocusChainTypeMismatch { .. } + ))); + } + + #[test] + fn empty_identity_does_not_fail_validation() { + // No locus declared → no mismatch issue regardless of chain. + let cfg_vj = RefDataConfig::empty(ChainType::Vj); + assert!(!cfg_vj.validate().iter().any(|i| matches!( + i, + RefDataValidationIssue::LocusChainTypeMismatch { .. } + ))); + let cfg_vdj = RefDataConfig::empty(ChainType::Vdj); + assert!(!cfg_vdj.validate().iter().any(|i| matches!( + i, + RefDataValidationIssue::LocusChainTypeMismatch { .. } + ))); + } + + #[test] + fn locus_is_case_insensitive() { + // Lowercase / mixed-case locus still flags correctly. + let cfg = vj_cfg_with_identity_locus("igh"); + let issues = cfg.validate(); + assert!(issues.iter().any(|i| matches!( + i, + RefDataValidationIssue::LocusChainTypeMismatch { locus, .. } if locus == "IGH" + ))); + } diff --git a/pyproject.toml b/pyproject.toml index 2402a86..4edfb5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,9 +21,12 @@ readme = "README.md" requires-python = ">=3.9" license = { text = "GPL-3.0-or-later" } authors = [ - { name = "Thomas Konstantinovsky", email = "thomaskon90@gmail.com" }, + { name = "Thomas Konstantinovsky" }, { name = "Ayelet Peres" }, ] +maintainers = [ + { name = "Thomas Konstantinovsky", email = "thomaskon90@gmail.com" }, +] keywords = ["immunogenetics", "sequence simulation", "bioinformatics", "alignment benchmarking"] classifiers = [ "Development Status :: 5 - Production/Stable", diff --git a/site_docs/concepts/reference-cartridge.md b/site_docs/concepts/reference-cartridge.md index 9fd17b5..6500788 100644 --- a/site_docs/concepts/reference-cartridge.md +++ b/site_docs/concepts/reference-cartridge.md @@ -166,6 +166,27 @@ the legacy nested-dict fallback (`cfg.NP_lengths` / `cfg.trim_dicts`), then a uniform placeholder. So a cartridge can ship typed defaults that users override per-experiment. +!!! warning "Which bundled cartridges carry *real* data-derived distributions" + Only the human **IGH**, **IGK**, **IGL**, and **TCRB** cartridges + ship with empirical distributions fitted from **real repertoire + data** — trim lengths, NP-region lengths, the NP base / Markov + model, and allele usage. + + Every **other** bundled species/locus cartridge carries a real + germline **allele catalogue** (from IMGT), but its data-derived + parameters are **placeholders — uniformly / randomly initialised, + not estimated from real data**. This is deliberate: a uniform prior + gives even coverage of the scenarios a repertoire can express + rather than baking in the biases of one dataset. Don't read those + parameters as ground-truth statistics for the species. + + Need empirically-grounded parameters for one of them? Fit your own: + estimate the distributions from real AIRR data with the + [cartridge estimators](../guides/estimate-cartridge-models.md) and + attach them to `cfg.reference_models`, or set the legacy dicts + (`cfg.trim_dicts`, `cfg.NP_lengths`, …) on the `DataConfig` + directly. + ## Build a cartridge from FASTA For new cartridges with an audit trail, use `ReferenceCartridgeBuilder`. diff --git a/site_docs/getting-started/quick-start.md b/site_docs/getting-started/quick-start.md index 3a2ebc1..3b3051e 100644 --- a/site_docs/getting-started/quick-start.md +++ b/site_docs/getting-started/quick-start.md @@ -57,7 +57,7 @@ the compiled plan first). This means the DSL composes cleanly: | Step | Effect | |---|---| -| `Experiment.on("human_igh")` | Bind to the bundled human IGH cartridge. Other shortcuts: `"human_igk"`, `"human_igl"`, `"mouse_igh"`, `"human_tcrb"`. Pass a `DataConfig` instead of a string for a custom cartridge. | +| `Experiment.on("human_igh")` | Bind to the bundled human IGH cartridge. Other shortcuts: `"human_igk"`, `"human_igl"`, `"mouse_igh"`, `"human_tcrb"`. Pass a `DataConfig` instead of a string for a custom cartridge. (Only human IGH/IGK/IGL/TCRB carry real data-derived distributions — [other species use uniform placeholders](../concepts/reference-cartridge.md#empirical-models).) | | `.recombine()` | Append a V(D)J recombination pass - sample alleles, trim, fill NP1/NP2, assemble. Defaults to the cartridge's empirical models. | | `.productive_only()` | Constraint-aware: the engine samples only choices that produce a productive sequence (in-frame junction, no stop codons, anchors preserved). No retry loops. | | `.run_records(n=1000, seed=42)` | Compile the plan, run 1,000 seeded draws, project each into an AIRR-format record. Same seed → byte-identical output across runs and platforms. | diff --git a/src/GenAIRR/_cartridge_estimators/__init__.py b/src/GenAIRR/_cartridge_estimators/__init__.py new file mode 100644 index 0000000..d951ff2 --- /dev/null +++ b/src/GenAIRR/_cartridge_estimators/__init__.py @@ -0,0 +1,20 @@ +"""Cartridge model estimators, split one module per estimator (behavior- +preserving). ReferenceCartridgeBuilder inherits the aggregate mixin.""" +from .allele_usage import _AlleleUsageEstimatorMixin +from .trim import _TrimEstimatorMixin +from .np_lengths import _NpLengthEstimatorMixin +from .np_base_model import _NpBaseModelEstimatorMixin +from .p_nucleotide_lengths import _PNucleotideLengthEstimatorMixin + + +class _CartridgeEstimators( + _AlleleUsageEstimatorMixin, + _TrimEstimatorMixin, + _NpLengthEstimatorMixin, + _NpBaseModelEstimatorMixin, + _PNucleotideLengthEstimatorMixin, +): + __slots__ = () + + +__all__ = ["_CartridgeEstimators"] diff --git a/src/GenAIRR/_cartridge_estimators/_common.py b/src/GenAIRR/_cartridge_estimators/_common.py new file mode 100644 index 0000000..ea82201 --- /dev/null +++ b/src/GenAIRR/_cartridge_estimators/_common.py @@ -0,0 +1,70 @@ +"""Shared module-level helpers for the cartridge estimators (behavior- +preserving). Moved verbatim from the original single-file module.""" +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, List + +_NP_CANONICAL_BASES = ("A", "C", "G", "T") + + +def _split_tie_set(raw: str) -> List[str]: + """Parse a comma-separated AIRR tie-set call into a clean list. + + Empty / whitespace-only entries drop out. The returned list + preserves insertion order — important for the + ``ambiguous="truth_first"`` policy which keeps only the first + entry.""" + return [token.strip() for token in raw.split(",") if token.strip()] + + +def _allele_names_for_segment( + buckets: Dict[str, List[Any]], +) -> "set[str]": + """Flatten a builder's per-gene allele bucket dict into a set + of canonical allele names. Used by + :meth:`ReferenceCartridgeBuilder.estimate_allele_usage` to + decide which AIRR-record allele names are "known" to the + cartridge under construction.""" + out: "set[str]" = set() + for allele_list in buckets.values(): + for allele in allele_list: + name = getattr(allele, "name", None) + if isinstance(name, str) and name: + out.add(name) + return out + + +def _load_rearrangements( + source: Any, +) -> "tuple[List[Dict[str, Any]], str]": + """Normalise a ``rearrangements`` argument into a + ``(records, source_label)`` tuple. + + Accepted inputs: + + - ``list[dict]`` — used verbatim; source label + ``"records:N"`` carrying the row count. + - path-like (``str`` / ``Path``) — parsed via + :class:`csv.DictReader` with ``delimiter='\\t'`` (AIRR-C TSV + convention); source label is the path. + - open text file handle — parsed via :class:`csv.DictReader`; + source label is ``"file:N"``. + + Raises :class:`TypeError` for unrecognised shapes.""" + import csv + + if isinstance(source, list): + return list(source), f"records:{len(source)}" + if hasattr(source, "read"): + rows = list(csv.DictReader(source, delimiter="\t")) + return rows, f"file:{len(rows)}" + if isinstance(source, (str, Path)): + path = Path(source) + with open(path, "r", newline="") as fh: + rows = list(csv.DictReader(fh, delimiter="\t")) + return rows, str(path) + raise TypeError( + f"rearrangements must be a list[dict], path, or open text " + f"file, got {type(source).__name__}" + ) diff --git a/src/GenAIRR/_cartridge_estimators/allele_usage.py b/src/GenAIRR/_cartridge_estimators/allele_usage.py new file mode 100644 index 0000000..03cc09b --- /dev/null +++ b/src/GenAIRR/_cartridge_estimators/allele_usage.py @@ -0,0 +1,287 @@ +"""Allele-usage estimator sub-mixin (behavior-preserving). The estimator +method is moved verbatim from the original single-file module.""" +from __future__ import annotations + +from typing import Any, Dict, List + +from ._common import ( + _allele_names_for_segment, + _load_rearrangements, + _split_tie_set, +) +from ..reference_models import ( + AlleleUsageSpec, + ReferenceEmpiricalModels, +) + + +class _AlleleUsageEstimatorMixin: + __slots__ = () + + def estimate_allele_usage( + self, + rearrangements: Any, + *, + min_count: float = 1.0, + ambiguous: str = "fractional", + replace: bool = True, + ) -> "ReferenceCartridgeBuilder": + """Estimate per-segment allele-usage weights from observed + AIRR rearrangement records. + + ``rearrangements`` accepts: + + - a list of dicts (each row is one AIRR record), + - a path-like (filesystem path to an AIRR TSV — parsed via + ``csv.DictReader`` with tab delimiter), or + - an open text file handle pointing at AIRR TSV. + + ``ambiguous`` selects the tie-set policy (column values + like ``"IGHV1*01,IGHV2*01"``): + + - ``"fractional"`` (default): split one row's credit + ``1.0`` evenly across all known alleles in the tie set. + Unknown allele names in the tie set are excluded; if + NO names in the tie set are known to the cartridge, the + row is recorded as ``unknown_allele`` and skipped. + - ``"truth_first"``: credit only the first comma-separated + allele in the tie set. Matches the existing + :func:`GenAIRR._mcp_summary` convention. + - ``"reject"``: drop ambiguous (multi-call) rows entirely + and record them in ``report.rejected``. + + ``min_count`` drops alleles whose final per-segment count + is strictly below the threshold. Pre-normalisation; the + per-segment weights remaining after the drop are then + renormalised to sum to ``1.0``. + + ``replace`` (default ``True``) controls idempotency. When + ``True``, calling :meth:`estimate_allele_usage` twice + overwrites the previous spec and writes ``replaced=True`` + on the new stage entry. When ``False``, the second call + raises :class:`ValueError`. + + Updates ``self._reference_models.allele_usage`` so a + downstream :meth:`build` carries the estimated spec into + the cartridge. + """ + if ambiguous not in ("fractional", "truth_first", "reject"): + raise ValueError( + f"ambiguous must be one of 'fractional' / 'truth_first' / " + f"'reject', got {ambiguous!r}" + ) + previously_estimated = any( + entry.get("stage") == "estimate_allele_usage" + for entry in self._report.stages + ) + if previously_estimated and not replace: + raise ValueError( + "estimate_allele_usage already ran; pass replace=True to " + "overwrite the previous spec" + ) + + records, source_label = _load_rearrangements(rearrangements) + chain_has_d = self._chain_type.has_d + v_pool = _allele_names_for_segment(self._v_alleles) + d_pool = _allele_names_for_segment(self._d_alleles) + j_pool = _allele_names_for_segment(self._j_alleles) + + v_counts: Dict[str, float] = {} + d_counts: Dict[str, float] = {} + j_counts: Dict[str, float] = {} + skipped = { + "missing_required_column": 0, + "unknown_allele": {"V": 0, "D": 0, "J": 0}, + "missing_d_call_on_vdj": 0, + "ambiguous_rejected": 0, + } + rejected_entries: List[Dict[str, Any]] = [] + warnings: List[str] = [] + d_on_vj_warned = False + + for row_idx, row in enumerate(records): + v_raw = (row.get("v_call") or "").strip() + j_raw = (row.get("j_call") or "").strip() + d_raw = (row.get("d_call") or "").strip() + + # Required column check. + if not v_raw or not j_raw: + skipped["missing_required_column"] += 1 + rejected_entries.append( + { + "stage": "estimate_allele_usage", + "row_index": row_idx, + "reason": "missing_required_column", + } + ) + continue + + # VDJ requires d_call; VJ ignores any d_call. + if chain_has_d and not d_raw: + skipped["missing_d_call_on_vdj"] += 1 + rejected_entries.append( + { + "stage": "estimate_allele_usage", + "row_index": row_idx, + "reason": "missing_d_call_on_vdj", + } + ) + continue + if not chain_has_d and d_raw: + if not d_on_vj_warned: + warnings.append( + "d_call column present on a VJ cartridge — D " + "contribution ignored for every row" + ) + d_on_vj_warned = True + d_raw = "" # ignore D contribution silently after warning + + v_tie = _split_tie_set(v_raw) + j_tie = _split_tie_set(j_raw) + d_tie = _split_tie_set(d_raw) if d_raw else [] + + if ambiguous == "reject": + if len(v_tie) > 1 or len(j_tie) > 1 or (chain_has_d and len(d_tie) > 1): + skipped["ambiguous_rejected"] += 1 + rejected_entries.append( + { + "stage": "estimate_allele_usage", + "row_index": row_idx, + "reason": "ambiguous_rejected", + } + ) + continue + + # truth_first collapses the tie set to its first entry. + if ambiguous == "truth_first": + v_tie = v_tie[:1] + j_tie = j_tie[:1] + d_tie = d_tie[:1] if d_tie else [] + + # Resolve known alleles in each tie set. + v_known = [n for n in v_tie if n in v_pool] + j_known = [n for n in j_tie if n in j_pool] + d_known = ( + [n for n in d_tie if n in d_pool] if d_tie else [] + ) + + # Unknown-allele bookkeeping. We record one rejection + # entry per unknown allele per segment per row so the + # report names the actual unknown names. + for n in v_tie: + if n not in v_pool: + skipped["unknown_allele"]["V"] += 1 + rejected_entries.append( + { + "stage": "estimate_allele_usage", + "row_index": row_idx, + "segment": "V", + "allele_name": n, + "reason": "unknown_allele", + } + ) + for n in j_tie: + if n not in j_pool: + skipped["unknown_allele"]["J"] += 1 + rejected_entries.append( + { + "stage": "estimate_allele_usage", + "row_index": row_idx, + "segment": "J", + "allele_name": n, + "reason": "unknown_allele", + } + ) + if d_tie: + for n in d_tie: + if n not in d_pool: + skipped["unknown_allele"]["D"] += 1 + rejected_entries.append( + { + "stage": "estimate_allele_usage", + "row_index": row_idx, + "segment": "D", + "allele_name": n, + "reason": "unknown_allele", + } + ) + + # Fractional credit (the default and the truth_first + # collapsed path; both use the same accumulation + # logic now that truth_first reduced the tie set). + if v_known: + share = 1.0 / len(v_known) + for n in v_known: + v_counts[n] = v_counts.get(n, 0.0) + share + if j_known: + share = 1.0 / len(j_known) + for n in j_known: + j_counts[n] = j_counts.get(n, 0.0) + share + if d_known: + share = 1.0 / len(d_known) + for n in d_known: + d_counts[n] = d_counts.get(n, 0.0) + share + + # min_count filter + normalisation per segment. + below_min: Dict[str, int] = {"V": 0, "D": 0, "J": 0} + + def _filter_and_normalise( + counts: Dict[str, float], segment_label: str + ) -> Dict[str, float]: + kept = {n: c for n, c in counts.items() if c >= min_count} + dropped = len(counts) - len(kept) + below_min[segment_label] = dropped + total = sum(kept.values()) + if total <= 0.0: + return {} + return {n: w / total for n, w in kept.items()} + + v_weights = _filter_and_normalise(v_counts, "V") + d_weights = _filter_and_normalise(d_counts, "D") + j_weights = _filter_and_normalise(j_counts, "J") + + if below_min["V"] or below_min["D"] or below_min["J"]: + warnings.append( + f"alleles below min_count={min_count} dropped — " + f"V={below_min['V']}, D={below_min['D']}, J={below_min['J']}" + ) + + spec = AlleleUsageSpec(v=v_weights, d=d_weights, j=j_weights) + chain_label = "vdj" if chain_has_d else "vj" + spec.validate(chain_type=chain_label, name="allele_usage") + + # Attach to the existing reference_models (creating it if + # absent). The cartridge built downstream carries the + # spec automatically. + if self._reference_models is None: + self._reference_models = ReferenceEmpiricalModels() + self._reference_models = ReferenceEmpiricalModels( + np_lengths=self._reference_models.np_lengths, + trims=self._reference_models.trims, + np_bases=self._reference_models.np_bases, + p_nucleotide_lengths=self._reference_models.p_nucleotide_lengths, + allele_usage=spec, + ) + + self._report.stages.append( + { + "stage": "estimate_allele_usage", + "inputs": { + "record_count": len(records), + "ambiguous": ambiguous, + "min_count": float(min_count), + "source": source_label, + "replaced": previously_estimated, + }, + "inferred": { + "V": v_weights, + "D": d_weights, + "J": j_weights, + "skipped": skipped, + "below_min_count": below_min, + }, + "warnings": warnings, + } + ) + self._report.rejected.extend(rejected_entries) + return self diff --git a/src/GenAIRR/_cartridge_estimators/np_base_model.py b/src/GenAIRR/_cartridge_estimators/np_base_model.py new file mode 100644 index 0000000..d947aba --- /dev/null +++ b/src/GenAIRR/_cartridge_estimators/np_base_model.py @@ -0,0 +1,377 @@ +"""NP base-model estimator sub-mixin (behavior-preserving). The estimator +method is moved verbatim from the original single-file module.""" +from __future__ import annotations + +from typing import Any, Dict, List + +from ._common import _NP_CANONICAL_BASES, _load_rearrangements +from ..reference_models import ( + NP_KEYS, + NpBaseModelSpec, + ReferenceEmpiricalModels, +) + + +class _NpBaseModelEstimatorMixin: + __slots__ = () + + def estimate_np_base_model( + self, + rearrangements: Any, + *, + kind: str = "markov", + min_count: int = 1, + pseudocount: float = 0.0, + replace: bool = True, + ) -> "ReferenceCartridgeBuilder": + """Estimate per-key NP base sampling model from + observed AIRR rearrangement records. + + Writes ``NpBaseModelSpec`` instances into + ``self._reference_models.np_bases`` keyed by + ``"NP1"`` (always) and (on VDJ cartridges) + ``"NP2"``. VJ cartridges produce only the ``NP1`` + key. + + ``kind`` selects the model family: + + - ``"empirical_first_base"``: estimates a single + categorical over A/C/G/T from the **full base + composition** of every observed NP string (every + position, not just position 0). The engine + samples every NP position independently from + this distribution — the name is preserved for + API stability with existing cartridges; the + biologically correct estimate is the full-base + composition. + - ``"markov"`` (default): estimates a first-base + row from position 0 of each NP string plus a + 4×4 transition matrix from every observed + (prev, next) pair. + + AIRR columns consumed (per audit §1.2): + + - ``np1`` (always). + - ``np2`` (VDJ only — VJ rows with a non-empty + ``np2`` raise a one-time warning and the + contribution is skipped). + + Columns deliberately ignored: ``junction`` (audit + §1.2), ``p_v_3_length`` / ``p_d_5_length`` / + ``p_d_3_length`` / ``p_j_5_length`` (separate + P-nucleotide biology), ``np1_length`` / + ``np2_length`` (length-only — owned by + :meth:`estimate_np_length_distributions`). + + Per-row validation is **field-local**: a malformed + or missing value in one NP column drops only that + key's contribution. Per-row structured entries + land in ``report.rejected`` with reasons + ``"missing_required_column"`` (empty / missing + string) or ``"noncanonical_base"`` (any character + outside ``{A,C,G,T}`` after uppercasing). + + ``min_count`` (int, default 1) drops first-base + categories whose observed count is strictly below + the threshold before normalisation. For ``markov``, + ``min_count`` applies to the **first-base row + only**: dropping transition cells could leave a + from-base row with no positive weights, which the + :class:`NpBaseModelSpec` validator rejects. v1 + keeps transition rows intact and surfaces the + first-base drops in ``below_min_count.first_base``. + + ``pseudocount`` (float, default 0.0) adds a uniform + prior: + + - ``empirical_first_base``: added to every A/C/G/T + base category before normalisation. + - ``markov``: added to every A/C/G/T first-base + category AND to every cell of the 4×4 transition + matrix. + + ``replace`` (default ``True``) controls idempotency. + When ``False`` AND a prior typed-plane + ``np_bases`` is already attached, the call raises + :class:`ValueError` before consuming any records. + """ + if kind not in ("empirical_first_base", "markov"): + raise ValueError( + f"kind must be one of 'empirical_first_base' / " + f"'markov', got {kind!r}" + ) + if isinstance(min_count, bool) or not isinstance(min_count, int): + raise ValueError( + f"min_count must be an int, got {min_count!r}" + ) + if min_count < 0: + raise ValueError( + f"min_count must be non-negative, got {min_count!r}" + ) + if isinstance(pseudocount, bool) or not isinstance( + pseudocount, (int, float) + ): + raise ValueError( + f"pseudocount must be a non-negative number, got {pseudocount!r}" + ) + if pseudocount < 0.0: + raise ValueError( + f"pseudocount must be non-negative, got {pseudocount!r}" + ) + + if not replace: + existing = ( + self._reference_models.np_bases + if self._reference_models is not None + else None + ) + if existing: + raise ValueError( + "np_bases model already attached to this cartridge; " + "pass replace=True to overwrite" + ) + previously_estimated = any( + entry.get("stage") == "estimate_np_base_model" + for entry in self._report.stages + ) + + records, source_label = _load_rearrangements(rearrangements) + chain_has_d = self._chain_type.has_d + active_keys = ("NP1", "NP2") if chain_has_d else ("NP1",) + + key_to_column = {"NP1": "np1", "NP2": "np2"} + bases = _NP_CANONICAL_BASES # ("A","C","G","T") + + # Per-key tallies. first_base[key] is a dict over A/C/G/T; + # transitions[key] is a 4×4 dict-of-dict over A/C/G/T from→to. + first_base_counts: Dict[str, Dict[str, int]] = { + key: {b: 0 for b in bases} for key in active_keys + } + transition_counts: Dict[str, Dict[str, Dict[str, int]]] = { + key: {b: {t: 0 for t in bases} for b in bases} + for key in active_keys + } + skipped = { + "missing_required_column": { + key_to_column[k]: 0 for k in active_keys + }, + "noncanonical_base": { + key_to_column[k]: 0 for k in active_keys + }, + } + # VJ chains track non-empty `np2` strings as a dropped + # column with a single one-time warning. + dropped_columns: Dict[str, int] = {} + if not chain_has_d: + dropped_columns["np2"] = 0 + rejected_entries: List[Dict[str, Any]] = [] + warnings: List[str] = [] + vj_np2_warned = False + + for row_idx, row in enumerate(records): + # Surface non-empty `np2` on VJ as a dropped column + # with one warning across the dataset. + if not chain_has_d: + raw_np2 = row.get("np2") + if isinstance(raw_np2, str) and raw_np2.strip(): + dropped_columns["np2"] += 1 + if not vj_np2_warned: + warnings.append( + "VJ cartridge: np2 column non-empty in " + "records; contribution ignored (no NP2 " + "region on a VJ chain)" + ) + vj_np2_warned = True + + # Field-local validation per active key. + for key in active_keys: + col = key_to_column[key] + raw = row.get(col) + if raw is None or not isinstance(raw, str) or not raw.strip(): + skipped["missing_required_column"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_np_base_model", + "row_index": row_idx, + "column": col, + "key": key, + "reason": "missing_required_column", + } + ) + continue + upper = raw.strip().upper() + non_canonical = sorted(set(upper) - set(bases)) + if non_canonical: + skipped["noncanonical_base"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_np_base_model", + "row_index": row_idx, + "column": col, + "key": key, + "unknown_chars": non_canonical, + "reason": "noncanonical_base", + } + ) + continue + # Tally bases. + # - For `empirical_first_base` the engine samples + # every NP position independently from the same + # distribution, so the biologically correct + # estimate is the FULL base composition (every + # position). + # - For `markov` the first-base row models position + # 0 only (the seed of the chain); position 1+ + # are modelled by transitions. + if kind == "empirical_first_base": + for b in upper: + first_base_counts[key][b] += 1 + else: + first_base_counts[key][upper[0]] += 1 + for i in range(len(upper) - 1): + prev_b = upper[i] + next_b = upper[i + 1] + transition_counts[key][prev_b][next_b] += 1 + + # min_count + pseudocount + per-kind spec construction. + below_min: Dict[str, Dict[str, int]] = { + key: {"first_base": 0} for key in active_keys + } + inferred: Dict[str, Dict[str, Any]] = { + key: {} for key in active_keys + } + + new_np_bases: Dict[str, NpBaseModelSpec] = {} + from_base_errors: List[str] = [] + + for key in active_keys: + fb_counts = first_base_counts[key] + # Apply min_count filter to first-base row (drop bases + # whose count is strictly below threshold). + fb_kept = {b: c for b, c in fb_counts.items() if c >= min_count} + below_min[key]["first_base"] = len(fb_counts) - len(fb_kept) + # Apply pseudocount AFTER filter — observed-only support. + if pseudocount > 0.0: + for b in bases: + if b in fb_kept: + fb_kept[b] = fb_kept[b] + pseudocount + else: + fb_kept[b] = pseudocount + fb_total = float(sum(fb_kept.values())) + if fb_total <= 0.0: + # Nothing observed AND no pseudocount → skip this key. + inferred[key] = {"first_base": None, "transitions": None} + continue + first_base_weights = { + b: fb_kept[b] / fb_total + for b in sorted(fb_kept.keys()) + if fb_kept[b] > 0 + } + + spec_payload: Dict[str, Any] = { + "first_base": first_base_weights, + } + if kind == "markov": + # Apply pseudocount to every cell of every transition row. + tr = transition_counts[key] + tr_weights: Dict[str, Dict[str, float]] = {} + for from_b in bases: + row_counts = {t: float(tr[from_b][t]) for t in bases} + if pseudocount > 0.0: + for t in bases: + row_counts[t] += pseudocount + row_total = sum(row_counts.values()) + if row_total <= 0.0: + # Pseudocount=0 and from-base never observed: surface a + # tagged error so the caller knows to pass pseudocount. + from_base_errors.append( + f"{key}: transition row for from-base " + f"{from_b!r} has no observed transitions " + f"and pseudocount=0; pass pseudocount > 0 " + f"or supply more data" + ) + continue + tr_weights[from_b] = { + t: row_counts[t] / row_total + for t in bases + if row_counts[t] > 0 + } + if from_base_errors: + # Bail before constructing a partial spec. + raise ValueError("; ".join(from_base_errors)) + spec_payload["transitions"] = tr_weights + else: + spec_payload["transitions"] = None + + # Construct + validate the spec. + spec = NpBaseModelSpec( + kind=kind, + first_base=spec_payload["first_base"], + transitions=spec_payload["transitions"], + ) + spec.validate(name=f"np_bases[{key}]") + new_np_bases[key] = spec + inferred[key] = { + "first_base": dict(spec_payload["first_base"]), + "transitions": ( + {k: dict(v) for k, v in spec_payload["transitions"].items()} + if spec_payload["transitions"] is not None + else None + ), + } + + if any(below_min[k]["first_base"] for k in active_keys): + warnings.append( + f"first-base values below min_count={min_count} dropped — " + + ", ".join( + f"{k}={below_min[k]['first_base']}" + for k in active_keys + ) + ) + + # Attach to existing reference_models (creating if absent); + # other typed planes preserved. + if self._reference_models is None: + self._reference_models = ReferenceEmpiricalModels() + self._reference_models = ReferenceEmpiricalModels( + np_lengths=self._reference_models.np_lengths, + trims=self._reference_models.trims, + np_bases=new_np_bases, + p_nucleotide_lengths=self._reference_models.p_nucleotide_lengths, + allele_usage=self._reference_models.allele_usage, + ) + chain_label = "vdj" if chain_has_d else "vj" + self._reference_models.validate(chain_type=chain_label) + + # NP2 entry surfaces in inferred even on VJ (as empty dict) + # for shape stability — same discipline as the NP-length + # estimator's empty `D_5` / `D_3` entries. + np1_inferred = inferred.get("NP1", {"first_base": None, "transitions": None}) + np2_inferred = inferred.get("NP2", {"first_base": None, "transitions": None}) + + self._report.stages.append( + { + "stage": "estimate_np_base_model", + "inputs": { + "record_count": len(records), + "kind": kind, + "min_count": int(min_count), + "pseudocount": float(pseudocount), + "source": source_label, + "replaced": previously_estimated, + }, + "inferred": { + "NP1": np1_inferred, + "NP2": np2_inferred, + "skipped": skipped, + "below_min_count": { + k: below_min.get(k, {"first_base": 0}) + for k in NP_KEYS + }, + "dropped_columns": dropped_columns, + }, + "warnings": warnings, + } + ) + self._report.rejected.extend(rejected_entries) + return self diff --git a/src/GenAIRR/_cartridge_estimators/np_lengths.py b/src/GenAIRR/_cartridge_estimators/np_lengths.py new file mode 100644 index 0000000..3bc5d40 --- /dev/null +++ b/src/GenAIRR/_cartridge_estimators/np_lengths.py @@ -0,0 +1,277 @@ +"""NP-length distribution estimator sub-mixin (behavior-preserving). The +estimator method is moved verbatim from the original single-file module.""" +from __future__ import annotations + +from typing import Any, Dict, List + +from ._common import _load_rearrangements +from ..reference_models import ( + EmpiricalDistributionSpec, + NP_KEYS, + ReferenceEmpiricalModels, +) + + +class _NpLengthEstimatorMixin: + __slots__ = () + + def estimate_np_length_distributions( + self, + rearrangements: Any, + *, + min_count: int = 1, + pseudocount: float = 0.0, + replace: bool = True, + ) -> "ReferenceCartridgeBuilder": + """Estimate per-key NP length distributions from + observed AIRR rearrangement records. + + Writes ``EmpiricalDistributionSpec`` instances into + ``self._reference_models.np_lengths`` keyed by + ``"NP1"`` and (on VDJ cartridges) ``"NP2"`` (see + :data:`GenAIRR.reference_models.NP_KEYS`). VJ + cartridges produce only the ``NP1`` key; the + cartridge has no NP2 region. The boundary is + enforced at estimation time because the typed-plane + validator does NOT chain-type-reject ``NP2`` on + VJ at attach time (see audit §2.2 of + ``docs/np_length_estimation_design.md``). + + AIRR columns consumed (per audit §1.2): + + - ``np1_length`` (always). + - ``np2_length`` (VDJ only — VJ rows with a + non-zero ``np2_length`` raise a one-time warning + and the contribution is skipped). + + Columns deliberately ignored: ``np1`` / ``np2`` + (sequence-derived length is sensitive to + post-claim reabsorption — see audit §5.3), + ``p_v_3_length`` / ``p_d_5_length`` / + ``p_d_3_length`` / ``p_j_5_length`` (separate + P-nucleotide biology — see + ``docs/p_nucleotide_design.md``), + and ``junction_length`` (aggregate arithmetic too + fragile across simulators — audit §5.2). + + Per-row validation is **field-local**: a malformed + or missing value in one NP column drops only that + key's contribution; the row's other column still + feeds. Per-row structured entries land in + ``report.rejected`` with reasons + ``"missing_required_column"`` / + ``"malformed_length_value"`` / + ``"negative_length_value"``, each carrying the + AIRR column name. + + ``min_count`` (int, default 1) drops length values + whose observed integer count is strictly below + the threshold before normalisation. + + ``pseudocount`` (float, default 0.0) adds a uniform + prior to every **observed** length value before + normalisation (no support expansion). Applied + AFTER the ``min_count`` filter. + + ``replace`` (default ``True``) controls idempotency. + When ``False`` AND a prior typed-plane + ``np_lengths`` is attached to + ``self._reference_models``, the call raises + :class:`ValueError` before consuming any records. + """ + if isinstance(min_count, bool) or not isinstance(min_count, int): + raise ValueError( + f"min_count must be an int, got {min_count!r}" + ) + if min_count < 0: + raise ValueError( + f"min_count must be non-negative, got {min_count!r}" + ) + if isinstance(pseudocount, bool) or not isinstance( + pseudocount, (int, float) + ): + raise ValueError( + f"pseudocount must be a non-negative number, got {pseudocount!r}" + ) + if pseudocount < 0.0: + raise ValueError( + f"pseudocount must be non-negative, got {pseudocount!r}" + ) + + if not replace: + existing = ( + self._reference_models.np_lengths + if self._reference_models is not None + else None + ) + if existing: + raise ValueError( + "np_length distributions already attached to this " + "cartridge; pass replace=True to overwrite" + ) + previously_estimated = any( + entry.get("stage") == "estimate_np_length_distributions" + for entry in self._report.stages + ) + + records, source_label = _load_rearrangements(rearrangements) + chain_has_d = self._chain_type.has_d + active_keys = NP_KEYS if chain_has_d else ("NP1",) + + key_to_column = {"NP1": "np1_length", "NP2": "np2_length"} + + counters: Dict[str, Dict[int, int]] = {k: {} for k in active_keys} + skipped = { + "missing_required_column": { + key_to_column[k]: 0 for k in active_keys + }, + "malformed_length_value": { + key_to_column[k]: 0 for k in active_keys + }, + "negative_length_value": { + key_to_column[k]: 0 for k in active_keys + }, + } + # VJ chains track `np2_length` contributions as a dropped + # column with a single one-time warning. + dropped_columns: Dict[str, int] = {} + if not chain_has_d: + dropped_columns["np2_length"] = 0 + rejected_entries: List[Dict[str, Any]] = [] + warnings: List[str] = [] + vj_np2_warned = False + + for row_idx, row in enumerate(records): + # On VJ, surface non-zero `np2_length` as a dropped + # column with one warning across the dataset. + if not chain_has_d: + raw_np2 = row.get("np2_length") + if raw_np2 not in (None, "", "0"): + try: + if int(str(raw_np2).strip()) != 0: + dropped_columns["np2_length"] += 1 + if not vj_np2_warned: + warnings.append( + "VJ cartridge: np2_length column " + "present in records; contribution " + "ignored (no NP2 region on a VJ " + "chain)" + ) + vj_np2_warned = True + except (TypeError, ValueError): + pass + + # Field-local validation per active key. + for key in active_keys: + col = key_to_column[key] + raw = row.get(col) + if raw is None or str(raw).strip() == "": + skipped["missing_required_column"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_np_length_distributions", + "row_index": row_idx, + "column": col, + "key": key, + "reason": "missing_required_column", + } + ) + continue + try: + value = int(str(raw).strip()) + except (TypeError, ValueError): + skipped["malformed_length_value"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_np_length_distributions", + "row_index": row_idx, + "column": col, + "key": key, + "value": raw, + "reason": "malformed_length_value", + } + ) + continue + if value < 0: + skipped["negative_length_value"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_np_length_distributions", + "row_index": row_idx, + "column": col, + "key": key, + "value": value, + "reason": "negative_length_value", + } + ) + continue + counters[key][value] = counters[key].get(value, 0) + 1 + + # min_count + pseudocount + per-key normalisation. + below_min: Dict[str, int] = {k: 0 for k in active_keys} + inferred_pairs: Dict[str, List[tuple]] = {k: [] for k in active_keys} + + for key in active_keys: + raw_counts = counters[key] + kept = {v: c for v, c in raw_counts.items() if c >= min_count} + below_min[key] = len(raw_counts) - len(kept) + if pseudocount > 0.0: + kept = {v: (c + pseudocount) for v, c in kept.items()} + total = float(sum(kept.values())) + if total <= 0.0 or not kept: + inferred_pairs[key] = [] + continue + normalised = [(v, kept[v] / total) for v in sorted(kept.keys())] + inferred_pairs[key] = normalised + + if any(below_min[k] for k in active_keys): + warnings.append( + f"NP-length values below min_count={min_count} dropped — " + + ", ".join(f"{k}={below_min[k]}" for k in active_keys) + ) + + # Build per-key EmpiricalDistributionSpec instances. + new_np_lengths: Dict[str, EmpiricalDistributionSpec] = {} + for key, pairs in inferred_pairs.items(): + if not pairs: + continue + spec = EmpiricalDistributionSpec(pairs) + spec.validate(name=f"np_lengths[{key}]") + new_np_lengths[key] = spec + + # Attach to existing reference_models (creating if absent). + # Other typed planes are preserved. + if self._reference_models is None: + self._reference_models = ReferenceEmpiricalModels() + self._reference_models = ReferenceEmpiricalModels( + np_lengths=new_np_lengths, + trims=self._reference_models.trims, + np_bases=self._reference_models.np_bases, + p_nucleotide_lengths=self._reference_models.p_nucleotide_lengths, + allele_usage=self._reference_models.allele_usage, + ) + chain_label = "vdj" if chain_has_d else "vj" + self._reference_models.validate(chain_type=chain_label) + + self._report.stages.append( + { + "stage": "estimate_np_length_distributions", + "inputs": { + "record_count": len(records), + "min_count": int(min_count), + "pseudocount": float(pseudocount), + "source": source_label, + "replaced": previously_estimated, + }, + "inferred": { + "NP1": inferred_pairs.get("NP1", []), + "NP2": inferred_pairs.get("NP2", []), + "skipped": skipped, + "below_min_count": {k: below_min.get(k, 0) for k in NP_KEYS}, + "dropped_columns": dropped_columns, + }, + "warnings": warnings, + } + ) + self._report.rejected.extend(rejected_entries) + return self diff --git a/src/GenAIRR/_cartridge_estimators/p_nucleotide_lengths.py b/src/GenAIRR/_cartridge_estimators/p_nucleotide_lengths.py new file mode 100644 index 0000000..ad0896e --- /dev/null +++ b/src/GenAIRR/_cartridge_estimators/p_nucleotide_lengths.py @@ -0,0 +1,333 @@ +"""P-nucleotide length estimator sub-mixin (behavior-preserving). The +estimator method is moved verbatim from the original single-file module.""" +from __future__ import annotations + +from typing import Any, Dict, List + +from ._common import _load_rearrangements +from ..reference_models import ( + EmpiricalDistributionSpec, + P_NUCLEOTIDE_END_KEYS, + P_NUCLEOTIDE_END_KEYS_VJ, + ReferenceEmpiricalModels, +) + + +class _PNucleotideLengthEstimatorMixin: + __slots__ = () + + def estimate_p_nucleotide_lengths( + self, + rearrangements: Any, + *, + min_count: int = 1, + pseudocount: float = 0.0, + replace: bool = True, + ) -> "ReferenceCartridgeBuilder": + """Estimate per-end P-nucleotide length distributions + from observed AIRR rearrangement records. + + Writes ``EmpiricalDistributionSpec`` instances into + ``self._reference_models.p_nucleotide_lengths`` + keyed by ``"V_3"`` and ``"J_5"`` (always) plus + ``"D_5"`` and ``"D_3"`` (VDJ only). VJ cartridges + produce only the V_3 and J_5 keys; the typed-plane + validator chain-type-rejects D-end keys on VJ at + attach time. + + **Provenance warning.** This estimator requires + AIRR-like records that ALREADY carry GenAIRR's + P-length fields (``p_v_3_length`` / + ``p_d_5_length`` / ``p_d_3_length`` / + ``p_j_5_length``). It does NOT infer P-lengths + from generic AIRR junction sequences, NP strings, + or trim arithmetic — that inference problem is + out of scope for v1. External AIRR tools (IgBLAST, + MiXCR, …) do not model P-nucleotide additions, so + their output will either omit these columns + entirely (rejection storm) or populate them as + zero (degenerate ``[(0, 1.0)]`` distribution). + The estimator emits a stage-level warning per key + when ≥ 95% of contributing rows reported zero — + diagnostic of P-naïve input. See + ``docs/p_nucleotide_length_estimation_design.md`` + §5 for the realistic-input enumeration. + + AIRR columns consumed: + + - ``p_v_3_length`` (always). + - ``p_j_5_length`` (always). + - ``p_d_5_length`` (VDJ only — VJ rows with a + non-zero value raise a one-time warning per + column and the contribution is skipped). + - ``p_d_3_length`` (same). + + Columns deliberately ignored: every other AIRR + column. v1 does NOT derive P-lengths from + junction-arithmetic, NP strings, trim fields, or + end-loss fields. + + Per-row validation is **field-local**: a malformed + or missing value in one P-length column drops only + that key's contribution. Per-row structured + entries land in ``report.rejected`` with reasons + ``"missing_required_column"`` / + ``"malformed_length_value"`` / + ``"negative_length_value"``. + + ``min_count`` (int, default 1) drops length values + whose observed count is strictly below the + threshold before normalisation. + + ``pseudocount`` (float, default 0.0) adds a uniform + prior to every **observed** length value before + normalisation (no support expansion). Applied + AFTER the ``min_count`` filter. + + ``replace`` (default ``True``) controls idempotency. + When ``False`` AND a prior typed-plane + ``p_nucleotide_lengths`` is already attached, + the call raises :class:`ValueError` before + consuming any records. + """ + if isinstance(min_count, bool) or not isinstance(min_count, int): + raise ValueError( + f"min_count must be an int, got {min_count!r}" + ) + if min_count < 0: + raise ValueError( + f"min_count must be non-negative, got {min_count!r}" + ) + if isinstance(pseudocount, bool) or not isinstance( + pseudocount, (int, float) + ): + raise ValueError( + f"pseudocount must be a non-negative number, got {pseudocount!r}" + ) + if pseudocount < 0.0: + raise ValueError( + f"pseudocount must be non-negative, got {pseudocount!r}" + ) + + if not replace: + existing = ( + self._reference_models.p_nucleotide_lengths + if self._reference_models is not None + else None + ) + if existing: + raise ValueError( + "p_nucleotide_lengths model already attached to " + "this cartridge; pass replace=True to overwrite" + ) + previously_estimated = any( + entry.get("stage") == "estimate_p_nucleotide_lengths" + for entry in self._report.stages + ) + + records, source_label = _load_rearrangements(rearrangements) + chain_has_d = self._chain_type.has_d + active_keys = ( + P_NUCLEOTIDE_END_KEYS if chain_has_d + else P_NUCLEOTIDE_END_KEYS_VJ + ) + + key_to_column = { + "V_3": "p_v_3_length", + "D_5": "p_d_5_length", + "D_3": "p_d_3_length", + "J_5": "p_j_5_length", + } + vj_ignored_columns = ("p_d_5_length", "p_d_3_length") + + counters: Dict[str, Dict[int, int]] = {k: {} for k in active_keys} + contributing_counts: Dict[str, int] = {k: 0 for k in active_keys} + zero_counts: Dict[str, int] = {k: 0 for k in active_keys} + skipped = { + "missing_required_column": { + key_to_column[k]: 0 for k in active_keys + }, + "malformed_length_value": { + key_to_column[k]: 0 for k in active_keys + }, + "negative_length_value": { + key_to_column[k]: 0 for k in active_keys + }, + } + # VJ chains track nonzero `p_d_*_length` columns as + # dropped columns with one warning per column. + dropped_columns: Dict[str, int] = {} + if not chain_has_d: + for col in vj_ignored_columns: + dropped_columns[col] = 0 + rejected_entries: List[Dict[str, Any]] = [] + warnings: List[str] = [] + vj_ignored_warned = {col: False for col in vj_ignored_columns} + + for row_idx, row in enumerate(records): + # VJ: surface nonzero D-end columns as dropped + warn once per column. + if not chain_has_d: + for col in vj_ignored_columns: + raw = row.get(col) + if raw in (None, "", "0"): + continue + try: + value = int(str(raw).strip()) + except (TypeError, ValueError): + continue + if value != 0: + dropped_columns[col] += 1 + if not vj_ignored_warned[col]: + warnings.append( + f"VJ cartridge: {col} non-zero in " + f"input — contribution ignored (no " + f"D segment on a VJ chain)" + ) + vj_ignored_warned[col] = True + + # Field-local validation per active key. + for key in active_keys: + col = key_to_column[key] + raw = row.get(col) + if raw is None or str(raw).strip() == "": + skipped["missing_required_column"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_p_nucleotide_lengths", + "row_index": row_idx, + "column": col, + "key": key, + "reason": "missing_required_column", + } + ) + continue + try: + value = int(str(raw).strip()) + except (TypeError, ValueError): + skipped["malformed_length_value"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_p_nucleotide_lengths", + "row_index": row_idx, + "column": col, + "key": key, + "value": raw, + "reason": "malformed_length_value", + } + ) + continue + if value < 0: + skipped["negative_length_value"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_p_nucleotide_lengths", + "row_index": row_idx, + "column": col, + "key": key, + "value": value, + "reason": "negative_length_value", + } + ) + continue + counters[key][value] = counters[key].get(value, 0) + 1 + contributing_counts[key] += 1 + if value == 0: + zero_counts[key] += 1 + + # min_count + pseudocount + per-key normalisation. + below_min: Dict[str, int] = {k: 0 for k in active_keys} + inferred_pairs: Dict[str, List[tuple]] = {k: [] for k in active_keys} + zero_fraction: Dict[str, float] = {} + + for key in active_keys: + raw_counts = counters[key] + kept = {v: c for v, c in raw_counts.items() if c >= min_count} + below_min[key] = len(raw_counts) - len(kept) + if pseudocount > 0.0: + kept = {v: (c + pseudocount) for v, c in kept.items()} + total = float(sum(kept.values())) + if total <= 0.0 or not kept: + inferred_pairs[key] = [] + else: + inferred_pairs[key] = [ + (v, kept[v] / total) for v in sorted(kept.keys()) + ] + # zero_fraction tracks observed-row provenance before + # normalisation — useful diagnostic for P-naïve input. + n_contrib = contributing_counts[key] + if n_contrib > 0: + zero_fraction[key] = zero_counts[key] / n_contrib + else: + zero_fraction[key] = 0.0 + + if any(below_min[k] for k in active_keys): + warnings.append( + f"P-nucleotide length values below min_count={min_count} " + f"dropped — " + + ", ".join(f"{k}={below_min[k]}" for k in active_keys) + ) + + # Per-key provenance auto-warning (audit §5.2). + for key in active_keys: + if ( + contributing_counts[key] > 0 + and zero_fraction[key] >= 0.95 + ): + warnings.append( + f"p_nucleotide_lengths[{key}] is >=95% zero; " + f"input may be P-naive or lack P annotations" + ) + + # Build per-key EmpiricalDistributionSpec instances. + new_p_lengths: Dict[str, EmpiricalDistributionSpec] = {} + for key, pairs in inferred_pairs.items(): + if not pairs: + continue + spec = EmpiricalDistributionSpec(pairs) + spec.validate(name=f"p_nucleotide_lengths[{key}]") + new_p_lengths[key] = spec + + # Attach to existing reference_models (creating if absent); + # other typed planes preserved. + if self._reference_models is None: + self._reference_models = ReferenceEmpiricalModels() + self._reference_models = ReferenceEmpiricalModels( + np_lengths=self._reference_models.np_lengths, + trims=self._reference_models.trims, + np_bases=self._reference_models.np_bases, + p_nucleotide_lengths=new_p_lengths, + allele_usage=self._reference_models.allele_usage, + ) + chain_label = "vdj" if chain_has_d else "vj" + self._reference_models.validate(chain_type=chain_label) + + self._report.stages.append( + { + "stage": "estimate_p_nucleotide_lengths", + "inputs": { + "record_count": len(records), + "min_count": int(min_count), + "pseudocount": float(pseudocount), + "source": source_label, + "replaced": previously_estimated, + }, + "inferred": { + "V_3": inferred_pairs.get("V_3", []), + "D_5": inferred_pairs.get("D_5", []), + "D_3": inferred_pairs.get("D_3", []), + "J_5": inferred_pairs.get("J_5", []), + "skipped": skipped, + "below_min_count": { + k: below_min.get(k, 0) for k in P_NUCLEOTIDE_END_KEYS + }, + "zero_fraction": { + k: zero_fraction.get(k, 0.0) + for k in P_NUCLEOTIDE_END_KEYS + }, + "dropped_columns": dropped_columns, + }, + "warnings": warnings, + } + ) + self._report.rejected.extend(rejected_entries) + return self diff --git a/src/GenAIRR/_cartridge_estimators/trim.py b/src/GenAIRR/_cartridge_estimators/trim.py new file mode 100644 index 0000000..94e32cb --- /dev/null +++ b/src/GenAIRR/_cartridge_estimators/trim.py @@ -0,0 +1,309 @@ +"""Trim-distribution estimator sub-mixin (behavior-preserving). The +estimator method is moved verbatim from the original single-file module.""" +from __future__ import annotations + +from typing import Any, Dict, List + +from ._common import _load_rearrangements +from ..reference_models import ( + EmpiricalDistributionSpec, + ReferenceEmpiricalModels, + TRIM_KEYS, + TRIM_KEYS_VJ, +) + + +class _TrimEstimatorMixin: + __slots__ = () + + def estimate_trim_distributions( + self, + rearrangements: Any, + *, + min_count: int = 1, + pseudocount: float = 0.0, + replace: bool = True, + ) -> "ReferenceCartridgeBuilder": + """Estimate per-key trim distributions from observed + AIRR rearrangement records. + + Writes ``EmpiricalDistributionSpec`` instances into + ``self._reference_models.trims`` keyed by ``"V_3"`` / + ``"D_5"`` / ``"D_3"`` / ``"J_5"`` (see + :data:`GenAIRR.reference_models.TRIM_KEYS`). VJ cartridges + only produce the ``V_3`` and ``J_5`` keys (see + :data:`GenAIRR.reference_models.TRIM_KEYS_VJ`); the D-end + keys are skipped silently because the cartridge has no D + segment to trim. + + ``rearrangements`` accepts the same shapes as + :meth:`estimate_allele_usage`: + + - a list of dicts (each row is one AIRR record), + - a path-like to an AIRR TSV (parsed via ``csv.DictReader`` + with tab delimiter), or + - an open text file handle pointing at AIRR TSV. + + AIRR columns consumed (per audit §1.3): + + - **VJ:** ``v_trim_3``, ``j_trim_5``. + - **VDJ:** ``v_trim_3``, ``d_trim_5``, ``d_trim_3``, ``j_trim_5``. + + Columns deliberately ignored: ``v_trim_5`` / ``j_trim_3`` + (no engine pass — hard-zero in projection), and the + observation-stage ``end_loss_5_length`` / ``end_loss_3_length`` + (separate corruption surface — see + ``docs/primer_trim_end_loss_audit.md``). + + Per-row validation is **field-local**: a malformed or + missing value in one trim column drops that field's + contribution only — the row still feeds its other + well-formed fields. Per-row structured entries land in + ``report.rejected`` with reasons + ``"missing_required_column"`` / ``"malformed_trim_value"`` + / ``"negative_trim_value"``, each carrying the AIRR + column name. + + ``min_count`` (int, default 1) drops trim values whose + observed integer count is strictly below the threshold + before normalisation. + + ``pseudocount`` (float, default 0.0) adds a uniform prior + to every **observed** trim value before normalisation + (no support expansion; values never observed stay + unobserved). Applied AFTER the ``min_count`` filter so + the filter looks at raw observations. + + ``replace`` (default ``True``) controls idempotency. When + ``True``, calling this method twice overwrites the + previous specs and writes ``replaced=True`` on the new + stage entry. When ``False`` AND a prior typed-plane + ``trims`` is already attached to ``self._reference_models``, + the call raises :class:`ValueError` before consuming any + records. + """ + if isinstance(min_count, bool) or not isinstance(min_count, int): + raise ValueError( + f"min_count must be an int, got {min_count!r}" + ) + if min_count < 0: + raise ValueError( + f"min_count must be non-negative, got {min_count!r}" + ) + if isinstance(pseudocount, bool) or not isinstance( + pseudocount, (int, float) + ): + raise ValueError( + f"pseudocount must be a non-negative number, got {pseudocount!r}" + ) + if pseudocount < 0.0: + raise ValueError( + f"pseudocount must be non-negative, got {pseudocount!r}" + ) + + if not replace: + existing_trims = ( + self._reference_models.trims + if self._reference_models is not None + else None + ) + if existing_trims: + raise ValueError( + "trim distributions already attached to this cartridge; " + "pass replace=True to overwrite" + ) + previously_estimated = any( + entry.get("stage") == "estimate_trim_distributions" + for entry in self._report.stages + ) + + records, source_label = _load_rearrangements(rearrangements) + chain_has_d = self._chain_type.has_d + active_keys = TRIM_KEYS if chain_has_d else TRIM_KEYS_VJ + + # Map plane key → AIRR column. + key_to_column = { + "V_3": "v_trim_3", + "D_5": "d_trim_5", + "D_3": "d_trim_3", + "J_5": "j_trim_5", + } + # AIRR columns the estimator deliberately does NOT consume. + ignored_columns = ("v_trim_5", "j_trim_3") + + counters: Dict[str, Dict[int, int]] = {k: {} for k in active_keys} + skipped = { + "missing_required_column": {col: 0 for col in + (key_to_column[k] for k in active_keys)}, + "malformed_trim_value": {col: 0 for col in + (key_to_column[k] for k in active_keys)}, + "negative_trim_value": {col: 0 for col in + (key_to_column[k] for k in active_keys)}, + } + dropped_columns: Dict[str, int] = {col: 0 for col in ignored_columns} + rejected_entries: List[Dict[str, Any]] = [] + warnings: List[str] = [] + ignored_warned = {col: False for col in ignored_columns} + vj_d_warned = False + + for row_idx, row in enumerate(records): + # Surface non-zero `v_trim_5` / `j_trim_3` columns: track + # the count but never consume them. One warning per + # column across the dataset. + for col in ignored_columns: + raw = row.get(col) + if raw not in (None, "", "0"): + try: + if int(str(raw).strip()) != 0: + dropped_columns[col] += 1 + if not ignored_warned[col]: + warnings.append( + f"{col} column non-zero in input — " + f"contribution dropped (no V_5 / J_3 " + f"trim pass in the engine)" + ) + ignored_warned[col] = True + except (TypeError, ValueError): + # Malformed in an unused column is uninteresting. + pass + + # Warn once per dataset if a VJ cartridge sees populated + # D-trim columns in the input. Same boundary as + # `estimate_allele_usage`'s D-call ignore. + if not chain_has_d: + for col in ("d_trim_5", "d_trim_3"): + raw = row.get(col) + if raw not in (None, "", "0") and not vj_d_warned: + try: + if int(str(raw).strip()) != 0: + warnings.append( + "VJ cartridge: d_trim_5 / d_trim_3 " + "columns present in records; " + "contribution ignored" + ) + vj_d_warned = True + break + except (TypeError, ValueError): + pass + + # Field-local validation: each key/column independently. + for key in active_keys: + col = key_to_column[key] + raw = row.get(col) + if raw is None or str(raw).strip() == "": + skipped["missing_required_column"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_trim_distributions", + "row_index": row_idx, + "column": col, + "key": key, + "reason": "missing_required_column", + } + ) + continue + try: + value = int(str(raw).strip()) + except (TypeError, ValueError): + skipped["malformed_trim_value"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_trim_distributions", + "row_index": row_idx, + "column": col, + "key": key, + "value": raw, + "reason": "malformed_trim_value", + } + ) + continue + if value < 0: + skipped["negative_trim_value"][col] += 1 + rejected_entries.append( + { + "stage": "estimate_trim_distributions", + "row_index": row_idx, + "column": col, + "key": key, + "value": value, + "reason": "negative_trim_value", + } + ) + continue + counters[key][value] = counters[key].get(value, 0) + 1 + + # min_count + pseudocount + per-key normalisation. + below_min: Dict[str, int] = {k: 0 for k in active_keys} + inferred_pairs: Dict[str, List[tuple]] = {k: [] for k in active_keys} + + for key in active_keys: + raw_counts = counters[key] + # Drop values strictly below min_count. + kept = {v: c for v, c in raw_counts.items() if c >= min_count} + below_min[key] = len(raw_counts) - len(kept) + # Apply pseudocount to observed values only. + if pseudocount > 0.0: + kept = {v: (c + pseudocount) for v, c in kept.items()} + # Normalise to sum 1.0 if anything survived. + total = float(sum(kept.values())) + if total <= 0.0 or not kept: + inferred_pairs[key] = [] + continue + normalised = [(v, kept[v] / total) for v in sorted(kept.keys())] + inferred_pairs[key] = normalised + + if any(below_min[k] for k in active_keys): + warnings.append( + f"trim values below min_count={min_count} dropped — " + + ", ".join(f"{k}={below_min[k]}" for k in active_keys) + ) + + # Build the per-key EmpiricalDistributionSpec instances. + new_trims: Dict[str, EmpiricalDistributionSpec] = {} + for key, pairs in inferred_pairs.items(): + if not pairs: + continue + spec = EmpiricalDistributionSpec(pairs) + spec.validate(name=f"trims[{key}]") + new_trims[key] = spec + + # Attach to the existing reference_models (creating it if + # absent). Other typed planes are preserved. + if self._reference_models is None: + self._reference_models = ReferenceEmpiricalModels() + self._reference_models = ReferenceEmpiricalModels( + np_lengths=self._reference_models.np_lengths, + trims=new_trims, + np_bases=self._reference_models.np_bases, + p_nucleotide_lengths=self._reference_models.p_nucleotide_lengths, + allele_usage=self._reference_models.allele_usage, + ) + # Validate the full container under the cartridge's chain + # type so D-on-VJ etc. raise at attach time. + chain_label = "vdj" if chain_has_d else "vj" + self._reference_models.validate(chain_type=chain_label) + + self._report.stages.append( + { + "stage": "estimate_trim_distributions", + "inputs": { + "record_count": len(records), + "min_count": int(min_count), + "pseudocount": float(pseudocount), + "source": source_label, + "replaced": previously_estimated, + }, + "inferred": { + "V_3": inferred_pairs.get("V_3", []), + "D_5": inferred_pairs.get("D_5", []), + "D_3": inferred_pairs.get("D_3", []), + "J_5": inferred_pairs.get("J_5", []), + "skipped": skipped, + "below_min_count": {k: below_min.get(k, 0) for k in TRIM_KEYS}, + "dropped_columns": dropped_columns, + }, + "warnings": warnings, + } + ) + self._report.rejected.extend(rejected_entries) + return self diff --git a/src/GenAIRR/_compiled.py b/src/GenAIRR/_compiled.py deleted file mode 100644 index e6684e0..0000000 --- a/src/GenAIRR/_compiled.py +++ /dev/null @@ -1,906 +0,0 @@ -"""Runtime wrappers around :class:`GenAIRR._engine.CompiledSimulator`. - -:class:`CompiledExperiment` is the frozen, executable form of -:class:`GenAIRR.experiment.Experiment`. :class:`CompiledClonalExperiment` -wraps the two-stage parent-then-descendants flow produced by -``expand_clones(...)``. - -Both classes are intentionally thin: they hold a reference to a -compiled engine simulator and a few pieces of source context -(``steps``, ``dataconfig``, ``metadata``) so ``describe()`` and AIRR -record export can render a faithful narrative. The execution loops -delegate straight into the engine; the wrappers exist for ergonomics -and to keep the public surface stable across engine refactors. -""" -from __future__ import annotations - -from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, TYPE_CHECKING - -from GenAIRR import _engine # private Rust extension submodule - -from ._describe import ( - _describe_clonal_fork_step, - _describe_experiment_header, - _describe_step_sequence, - _format_active_contracts, -) -from ._pipeline_ir import _ClonalForkStep, _LineageForkStep, _RepertoireForkStep - -if TYPE_CHECKING: - from .dataconfig import DataConfig - from .result import SimulationResult - - -class CompiledExperiment: - """A frozen ``Experiment`` ready for execution. - - Holds the owning :class:`GenAIRR._engine.CompiledSimulator` and the - refdata it was built against. Contracts are captured at compile - time; ``run()`` only accepts execution parameters. - """ - - __slots__ = ( - "_simulator", - "_refdata", - "_steps", - "_dataconfig", - "_metadata", - "_genotype", - ) - - def __init__( - self, - simulator: "_engine.CompiledSimulator", - refdata: "_engine.RefDataConfig", - steps: Sequence[Any] = (), - dataconfig: Optional["DataConfig"] = None, - metadata: Optional[Dict[str, Any]] = None, - genotype: Optional[Any] = None, - ) -> None: - self._simulator = simulator - self._refdata = refdata - # Source steps stashed for `describe()`. The compiled simulator - # itself can render pass names but not biology — keeping the - # builder steps around is the cheapest way to give a faithful - # narrative back. - self._steps: Tuple[Any, ...] = tuple(steps) - self._dataconfig = dataconfig - self._metadata = dict(metadata) if metadata else {} - # Attached single-subject genotype (or None). When set, - # run_records stamps subject_id + haplotype provenance. - self._genotype = genotype - - @property - def simulator(self) -> "_engine.CompiledSimulator": - """The owning Rust compiled simulator.""" - return self._simulator - - @property - def pass_plan(self) -> Tuple[str, ...]: - """Read-only pass-name summary of the compiled pipeline.""" - return tuple(self._simulator.pass_names()) - - @property - def pass_names(self) -> Tuple[str, ...]: - """Stable names of the compiled pass sequence.""" - return self.pass_plan - - @property - def active_contracts(self) -> Tuple[str, ...]: - """Stable names of the contract bundle captured at compile time.""" - return tuple(self._simulator.active_contracts()) - - @property - def refdata(self) -> "_engine.RefDataConfig": - """The :class:`GenAIRR._engine.RefDataConfig` the plan was built against.""" - return self._refdata - - def describe(self) -> str: - """Render a biology-style narrative of the compiled experiment. - - Equivalent to ``Experiment.describe()`` but additionally - surfaces compile-time constraints (e.g. ``productive_only``) - attached via ``compile(respect=...)``. See - :meth:`Experiment.describe` for the output shape. - """ - header = _describe_experiment_header(self._refdata, self._dataconfig) - if not self._steps: - body = [" (no steps recorded)"] - else: - body = _describe_step_sequence(self._steps, self._refdata.chain_type) - lines = [header, *body] - contracts_line = _format_active_contracts(self.active_contracts) - if contracts_line: - lines.append(f" Constraints: {contracts_line}") - if self._metadata: - stamps = ", ".join(f"{k}={v!r}" for k, v in self._metadata.items()) - lines.append(f" Metadata stamped on every record: {stamps}") - return "\n".join(lines) - - def run( - self, - *, - n: int = 1, - seed: int = 0, - strict: bool = False, - ) -> List["_engine.Outcome"]: - """Run the compiled simulator ``n`` times — **fresh sampling**. - - Each iteration uses ``seed + i`` as the per-run seed so - consecutive batches stitch together by offsetting ``seed``. - - ``strict`` controls the failure mode when a pass's - contract-narrowed candidate set is empty *at sample time*: - - - ``False`` (default, **permissive**) — apply the pass's - declared ``EmptySupport`` policy. The pass writes a - documented sentinel value to the trace and continues. - Common sentinels: indel ``site = -1`` NoOp, NP length ``0``, - NP base ``N``, trim ``0``; SHM substitution skips the slot - (no trace record). - - ``True`` (**strict**) — raise - :class:`GenAIRR._engine.StrictSamplingError`. The exception's - ``args`` are a 3-tuple ``(pass_name, address, reason)``; - ``reason`` is one of ``"support_unavailable"``, - ``"empty_admissible_support"``, or - ``"invalid_filtered_support"``. - - **Compile-time precondition failures are separate.** If a - sampling distribution is *statically* incompatible with the - active contracts (e.g. every NP1 length in the distribution - violates frame divisibility), :meth:`Experiment.compile` - raises :class:`ValueError` *before* this method runs. - ``ValueError`` and ``StrictSamplingError`` have **no shared - base class** — catching only one will miss the other. See - ``docs/productive_failure_mode_audit.md`` §6.1. - - **Strict semantics apply only to fresh sampling.** Trace - replay via :meth:`CompiledExperiment.replay_from_trace_file` - consumes the recorded values verbatim and does NOT - re-evaluate contract admissibility — a permissive sentinel - trace replays cleanly even under ``strict=True``. See that - method's docstring. - - Raises ``ValueError`` for ``n < 1``. - """ - if n < 1: - raise ValueError(f"n must be at least 1, got {n}") - return self._simulator.run_batch(n, seed, strict=strict) - - def run_records( - self, - *, - n: int = 1, - seed: int = 0, - strict: bool = False, - expose_provenance: bool = False, - validate_records: bool = False, - ) -> "SimulationResult": - """Run the compiled simulator ``n`` times and return the batch as - a :class:`SimulationResult` ready for ``.to_csv`` / - ``.to_fasta`` / ``.to_dataframe`` export. - - Same arguments as :meth:`run`. ``expose_provenance=True`` - appends `truth_v_call/d_call/j_call` columns reflecting the - originally-sampled allele names. - - ``validate_records=True`` runs - :meth:`SimulationResult.validate_records` on the freshly - built batch and raises - :class:`GenAIRR._validation.RecordValidationFailedError` - (a :class:`RuntimeError` subclass) when any record fails the - postcondition validator. Default ``False`` keeps this method - zero-overhead. - """ - from .result import SimulationResult - - outcomes = self.run(n=n, seed=seed, strict=strict) - result = SimulationResult.from_outcomes( - outcomes, self._refdata, expose_provenance=expose_provenance - ) - if self._genotype is not None: - self._stamp_genotype_provenance(outcomes, result) - if validate_records: - from ._validation import _raise_on_validation_failure - - _raise_on_validation_failure(result.validate_records(self._refdata)) - return result - - def _stamp_genotype_provenance(self, outcomes, result) -> None: - """Stamp per-record ``subject_id`` + ``haplotype`` (the chromosome - the rearrangement drew from) and expose the genotype on the - result. Used only when a genotype is attached.""" - subject = self._genotype.subject_id - for outcome, rec in zip(outcomes, result._records): - rec["subject_id"] = subject - hap = outcome.trace().find("sample_haplotype") - rec["haplotype"] = hap.value if hap is not None else None - result._genotypes = [self._genotype] - - def stream( - self, - *, - n: Optional[int] = None, - seed: int = 0, - strict: bool = False, - ) -> Iterator["_engine.Outcome"]: - """Lazily yield :class:`GenAIRR._engine.Outcome` objects one at - a time, without materialising the full batch in memory. - - Useful for large simulations where holding ``n`` outcomes - would be wasteful — typical pattern is - - >>> for outcome in compiled.stream(n=1_000_000, seed=0): - ... write_to_disk(outcome) - - ``n=None`` (the default) yields outcomes indefinitely; the - caller is expected to stop with ``itertools.islice``, - ``break``, or similar. ``n=N`` yields exactly ``N`` outcomes - with seeds ``seed`` … ``seed + N - 1``. - - ``strict`` behaves as in :meth:`run`. - - Raises ``ValueError`` when ``n`` is set to a value below 1. - """ - if n is not None and n < 1: - raise ValueError(f"n must be at least 1, got {n}") - i = 0 - while n is None or i < n: - yield self._simulator.run(seed + i, strict=strict) - i += 1 - - def stream_records( - self, - *, - n: Optional[int] = None, - seed: int = 0, - strict: bool = False, - id_prefix: str = "seq", - ) -> Iterator[Dict[str, Any]]: - """Lazily yield AIRR-format record dicts (one per outcome). - - Same shape as the records inside a :class:`SimulationResult`, - but yielded one at a time so callers can write each record to - disk without retaining the prior ones. Pairs naturally with - :func:`csv.DictWriter` for streaming TSV/CSV output. - - Each record's ``sequence_id`` is set to - ``f"{id_prefix}{i}"`` so streamed batches have unique - AIRR-style identifiers without buffering. - """ - from ._airr_record import outcome_to_airr_record - - for i, outcome in enumerate( - self.stream(n=n, seed=seed, strict=strict) - ): - yield outcome_to_airr_record( - outcome, self._refdata, sequence_id=f"{id_prefix}{i}" - ) - - def __repr__(self) -> str: - return ( - f"" - ) - - -class CompiledClonalExperiment: - """A compiled experiment with a clonal-fork structure. - - Wraps two :class:`GenAIRR._engine.CompiledSimulator`s — the - pre-fork plan (run once per clone, typically the recombine - step) and the post-fork plan (run once per descendant inside - the clone, typically mutate / corrupt_*). - - :meth:`run_records` orchestrates the parent → descendants loop - and tags every record with a ``clone_id`` integer so downstream - clonotype-clustering tools can be benchmarked against the true - clonal structure. - """ - - __slots__ = ( - "_pre", - "_post", - "_fork", - "_refdata", - "_pre_steps", - "_post_steps", - "_dataconfig", - "_metadata", - ) - - def __init__( - self, - pre_simulator: "_engine.CompiledSimulator", - post_simulator: "_engine.CompiledSimulator", - fork: "_ClonalForkStep", - refdata: "_engine.RefDataConfig", - pre_steps: Sequence[Any] = (), - post_steps: Sequence[Any] = (), - dataconfig: Optional["DataConfig"] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> None: - self._pre = pre_simulator - self._post = post_simulator - self._fork = fork - self._refdata = refdata - self._pre_steps: Tuple[Any, ...] = tuple(pre_steps) - self._post_steps: Tuple[Any, ...] = tuple(post_steps) - self._dataconfig = dataconfig - self._metadata = dict(metadata) if metadata else {} - - @property - def n_clones(self) -> int: - return self._fork.n_clones - - @property - def size(self) -> int: - return self._fork.size - - @property - def total_records(self) -> int: - """Number of records produced per :meth:`run_records` call.""" - return self._fork.n_clones * self._fork.size - - @property - def refdata(self) -> "_engine.RefDataConfig": - return self._refdata - - def describe(self) -> str: - """Render a biology-style narrative of the compiled clonal - experiment, with an explicit divider at the fork. See - :meth:`Experiment.describe` for the basic shape.""" - header = _describe_experiment_header(self._refdata, self._dataconfig) - lines = [header] - # pre-fork section (per-clone) - pre_lines = _describe_step_sequence( - self._pre_steps, self._refdata.chain_type, start_index=1 - ) - lines.extend(pre_lines) - # the fork itself - lines.append(f" ── {_describe_clonal_fork_step(self._fork)} ──") - lines.append( - " (steps above run once per clone; " - "steps below run once per descendant)" - ) - # post-fork section (per-descendant) - post_start = sum( - 1 for s in self._pre_steps if not isinstance(s, _ClonalForkStep) - ) + 1 - post_lines = _describe_step_sequence( - self._post_steps, self._refdata.chain_type, start_index=post_start - ) - lines.extend(post_lines) - if self._metadata: - stamps = ", ".join(f"{k}={v!r}" for k, v in self._metadata.items()) - lines.append(f" Metadata stamped on every record: {stamps}") - return "\n".join(lines) - - def run( - self, - *, - n: Optional[int] = None, - seed: int = 0, - strict: bool = False, - ) -> List["_engine.Outcome"]: - """Run all clonal descendants and return their outcomes in - clone-major order (clone 0's descendants 0..size-1, clone 1's - descendants 0..size-1, …). - - ``n`` is optional: when omitted the runtime expands - ``n_clones * size`` outcomes. Passing ``n`` is allowed only - when ``n == n_clones * size`` (otherwise raises). - """ - total = self.total_records - if n is not None and n != total: - raise ValueError( - f"clonal pipeline produces n_clones * size = " - f"{self._fork.n_clones} * {self._fork.size} = {total} " - f"records; passing n={n} is inconsistent. Drop the n " - f"argument or pass n={total}." - ) - - outcomes: List["_engine.Outcome"] = [] - for clone_idx in range(self._fork.n_clones): - clone_seed = int(seed) + clone_idx * 1_000_000 - parent = self._pre.run(seed=clone_seed, strict=strict) - parent_sim = parent.final_simulation() - for desc_idx in range(self._fork.size): - desc_seed = clone_seed + 1 + desc_idx - desc = self._post.run_from( - parent_sim, desc_seed, strict=strict - ) - outcomes.append(desc) - return outcomes - - def run_records( - self, - *, - n: Optional[int] = None, - seed: int = 0, - strict: bool = False, - expose_provenance: bool = False, - validate_records: bool = False, - ) -> "SimulationResult": - """Same as :meth:`run` but returns a :class:`SimulationResult` - with each record dict carrying an integer ``clone_id`` field - in ``[0, n_clones)`` plus a ``parent_id`` integer indexing - into :attr:`SimulationResult.parents`. ``expose_provenance=True`` - also appends `truth_v_call` / `truth_d_call` / `truth_j_call` - columns from the originally-sampled allele names. - - The returned :class:`SimulationResult` carries the per-clone - parent ``Outcome`` objects on its ``.parents`` attribute. - Each parent holds the pre-fork addressed-choice trace, the - pre-fork event ledger, and the post-recombination IR — useful - for replay, lineage analysis, and the upcoming parent-aware - family validator. The flat ``.outcomes`` list continues to - carry only the descendant outcomes (one per record); - parents are exposed separately so the per-record list stays - the same shape clonal consumers already know. - - ``validate_records=True`` runs - :meth:`SimulationResult.validate_records` on the freshly - built batch and raises - :class:`GenAIRR._validation.RecordValidationFailedError` - on any postcondition failure. After the per-record gate - passes, this also runs - :meth:`SimulationResult.validate_families` and raises the - sibling - :class:`GenAIRR._validation.FamilyValidationFailedError` - if any clonal-family invariant is violated. The two - gates report separately so users can tell projection bugs - from family-consistency bugs. Default ``False`` keeps this - method zero-overhead. - """ - from ._airr_record import outcome_to_airr_record - from .result import SimulationResult, _inject_truth_columns - - total = self.total_records - if n is not None and n != total: - raise ValueError( - f"clonal pipeline produces n_clones * size = " - f"{self._fork.n_clones} * {self._fork.size} = {total} " - f"records; passing n={n} is inconsistent. Drop the n " - f"argument or pass n={total}." - ) - - records: List[Dict[str, Any]] = [] - outcomes: List["_engine.Outcome"] = [] - # Slice 2: retain parent outcomes — one per clone. The - # orchestration loop used to drop them after extracting - # ``final_simulation()``. We now keep them so the returned - # :class:`SimulationResult` can expose ``.parents`` for - # replay / lineage tooling. The parent ``Outcome`` itself - # is not copied onto each descendant; we hold a single - # reference per clone in the ``parents`` list. - parents: List["_engine.Outcome"] = [] - for clone_idx in range(self._fork.n_clones): - clone_seed = int(seed) + clone_idx * 1_000_000 - parent = self._pre.run(seed=clone_seed, strict=strict) - parents.append(parent) - parent_sim = parent.final_simulation() - for desc_idx in range(self._fork.size): - desc_seed = clone_seed + 1 + desc_idx - desc = self._post.run_from( - parent_sim, desc_seed, strict=strict - ) - outcomes.append(desc) - rec = outcome_to_airr_record( - desc, - self._refdata, - sequence_id=f"clone{clone_idx}_desc{desc_idx}", - ) - rec["clone_id"] = clone_idx - # ``parent_id`` is the descendant's index into - # ``result.parents``. Today clones are dense and - # zero-based, so ``parent_id == clone_id`` by - # construction — we stamp both because they carry - # distinct semantics: ``clone_id`` is the family - # identity (the existing Slice 0 contract); - # ``parent_id`` is the addressing scheme into the - # parent-outcome list (the new Slice 2 contract). - # Keeping them separate now means a future slice - # that introduces sparse / non-zero-based family - # ids doesn't have to retrofit both. - rec["parent_id"] = clone_idx - if expose_provenance: - _inject_truth_columns(desc, self._refdata, rec) - records.append(rec) - result = SimulationResult(records, outcomes=outcomes, parents=parents) - if validate_records: - from ._validation import ( - _raise_on_family_validation_failure, - _raise_on_validation_failure, - ) - - _raise_on_validation_failure(result.validate_records(self._refdata)) - _raise_on_family_validation_failure(result.validate_families()) - return result - - def __repr__(self) -> str: - return ( - f"" - ) - - -class CompiledLineageExperiment: - """A compiled experiment that grows BCR affinity-maturation lineage trees. - - Wraps a pre-fork :class:`GenAIRR._engine.CompiledSimulator` (founder - recombination) and a :class:`~GenAIRR._pipeline_ir._LineageForkStep` - (lineage parameters). :meth:`run_records` grows one lineage tree per - clone via the Rust ``simulate_family_outcomes`` kernel and returns a - :class:`~GenAIRR.result.SimulationResultWithLineages` whose records are - per-observed-node AIRR dicts with lineage metadata. - """ - - __slots__ = ( - "_pre", - "_step", - "_refdata", - "_post", - "_post_steps", - "_dataconfig", - "_metadata", - ) - - def __init__( - self, - pre_simulator: "_engine.CompiledSimulator", - step: "_LineageForkStep", - refdata: "_engine.RefDataConfig", - post_simulator: Optional["_engine.CompiledSimulator"] = None, - post_steps: Sequence[Any] = (), - dataconfig: Optional["DataConfig"] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> None: - self._pre = pre_simulator - self._step = step - self._refdata = refdata - # Per-observed-cell library-prep / sequencing corruption plan. - # ``None`` keeps the pristine-read path (no artefacts applied). - self._post = post_simulator - self._post_steps: Tuple[Any, ...] = tuple(post_steps) - self._dataconfig = dataconfig - self._metadata = dict(metadata) if metadata else {} - - @property - def refdata(self) -> "_engine.RefDataConfig": - return self._refdata - - def run_records( - self, - *, - seed: int = 0, - strict: bool = False, - expose_provenance: bool = False, - validate_records: bool = False, - ) -> "SimulationResultWithLineages": - """Grow lineage trees and return per-observed-node AIRR records. - - Each clone is seeded at ``seed + clone_idx * 1_000_000`` so - independent clones are reproducible and non-overlapping. - - The returned :class:`SimulationResultWithLineages` carries the - per-record ``Outcome`` that produced each AIRR record on its - ``.outcomes`` attribute (index-aligned with ``.records``). On the - pristine-read path these are the per-observed-node SHM - ``Outcome`` objects (``fam.observed_outcomes()``); on the - corruption path they are the per-node merged - recombination+SHM+artefact ``Outcome`` objects. Since each node - ``Outcome`` is self-consistent, ``result.validate_records(refdata)`` - passes. - - ``expose_provenance=True`` appends ``truth_v_call`` / - ``truth_d_call`` / ``truth_j_call`` columns to each record from - the founder allele assignments carried on the per-record - ``Outcome``. - - ``validate_records=True`` runs - :meth:`SimulationResult.validate_records` (per-record - postcondition) and :meth:`SimulationResult.validate_families` - (clonal-family consistency by ``clone_id``) on the freshly built - batch, raising the matching validation error on any failure. - """ - from GenAIRR import _engine as _eng - from ._airr_record import outcome_to_airr_record - from ._s5f_loader import load_builtin_s5f_kernel - from .result import SimulationResultWithLineages, _inject_truth_columns - - step = self._step - mutability, substitution = load_builtin_s5f_kernel(step.s5f_model) - - records: List[Dict[str, Any]] = [] - # Per-record source ``Outcome`` objects, index-aligned with - # ``records`` so ``result.validate_records`` can re-derive each - # record from the engine state that produced it. - outcomes: List["_engine.Outcome"] = [] - lineage_trees = [] - - # Bound on founder-survival retries. Sampling draws from the LIVING - # final-generation population, so an extinct founder (drew 0 offspring) - # yields zero observed cells. With ``allow_extinction=False`` (default) - # we condition each clone on survival by re-growing the family with a - # fresh deterministic sub-seed until it survives or we exhaust the - # bound. The sub-seed offset (a prime times the attempt index) keeps the - # retries deterministic — same top-level seed => same result. - _MAX_SURVIVAL_ATTEMPTS = 50 - _SUBSEED_STRIDE = 7_919 # prime; keeps retry sub-seeds well-separated - allow_extinction = getattr(step, "allow_extinction", False) - - for clone_idx in range(step.n_clones): - clone_seed = int(seed) + clone_idx * 1_000_000 - fam = None - tree = None - for attempt in range(_MAX_SURVIVAL_ATTEMPTS): - family_seed = clone_seed + attempt * _SUBSEED_STRIDE - # _pre.run() returns a single Outcome (not a list) — mirror - # CompiledClonalExperiment which calls self._pre.run(seed=...). - founder = self._pre.run(seed=family_seed, strict=strict) - candidate = _eng.simulate_family_outcomes( - founder, - self._refdata, - mutability, - substitution, - step.rate, - step.lambda_base, - 0.0, # lambda_mut: positional slot 6 (inert; hardcoded) - step.max_generations, - step.n_max, - step.n_sample, - family_seed, - selection_strength=step.selection_strength, - beta=step.beta, - target_aa=step.target_aa, - mature_substitutions=step.mature_substitutions, - ) - survived = len(candidate.observed_outcomes()) > 0 - if survived or allow_extinction: - fam = candidate - tree = candidate.tree() - break - # extinct + survival required => retry with next sub-seed - else: - # Exhausted the retry bound without a surviving family. - raise ValueError( - f"clonal_lineage: clone {clone_idx} went extinct on every " - f"one of {_MAX_SURVIVAL_ATTEMPTS} survival attempts (each " - "founder drew 0 offspring). Increase lambda_base (offspring " - "Poisson mean) and/or max_generations so families reliably " - "survive, or pass allow_extinction=True to accept extinct " - "clones (producing fewer than n_clones families)." - ) - - if fam is None: - # allow_extinction=True and the final attempt was extinct: - # skip this clone (it contributes no observed cells/records). - continue - - lineage_trees.append(tree) - if len(fam.observed_outcomes()) == 0: - # allow_extinction=True and this family is extinct: keep its - # (empty) tree for export parity but emit no records. - continue - - if self._post is None: - # Pristine-read path: project each observed node's - # synthesized SHM Outcome straight to an AIRR record. - # ``observed_outcomes()`` is index-aligned with both the - # observed-node list and ``airr_records()``. - observed_nodes = [n for n in tree.nodes() if n.observed] - recs = fam.airr_records(self._refdata) - observed_outcomes = fam.observed_outcomes() - for node, rec, node_outcome in zip( - observed_nodes, recs, observed_outcomes - ): - self._stamp_lineage_metadata(rec, clone_idx, node) - if expose_provenance and node_outcome is not None: - _inject_truth_columns(node_outcome, self._refdata, rec) - records.append(rec) - outcomes.append(node_outcome) - continue - - # Corruption path: per observed node, run the library-prep / - # sequencing corruption plan FROM the node's post-SHM - # simulation, then merge the founder-recombination + SHM - # provenance with the corruption trace / events so the AIRR - # record reports trims, v/d/j, SHM counts AND artefact - # counters (n_quality_errors, n_pcr_errors, n_indels, …). - node_outcomes = fam.node_outcomes() - for node, base_outcome in zip(tree.nodes(), node_outcomes): - if base_outcome is None: - continue - node_seed = clone_seed + 1 + node.id - corruption_outcome = self._post.run_from( - base_outcome.final_simulation(), node_seed, strict=strict - ) - merged = _eng.merge_lineage_corruption( - base_outcome, corruption_outcome - ) - rec = outcome_to_airr_record( - merged, - self._refdata, - sequence_id=f"clone{clone_idx}_node{node.id}", - ) - self._stamp_lineage_metadata(rec, clone_idx, node) - if expose_provenance: - _inject_truth_columns(merged, self._refdata, rec) - records.append(rec) - outcomes.append(merged) - - result = SimulationResultWithLineages( - records, outcomes=outcomes, lineage_trees=lineage_trees - ) - if validate_records: - from ._validation import ( - _raise_on_family_validation_failure, - _raise_on_validation_failure, - ) - - _raise_on_validation_failure(result.validate_records(self._refdata)) - _raise_on_family_validation_failure(result.validate_families()) - return result - - @staticmethod - def _stamp_lineage_metadata(rec: Dict[str, Any], clone_idx: int, node) -> None: - """Stamp clone + lineage-node provenance onto an AIRR record.""" - rec["clone_id"] = clone_idx - rec["lineage_node_id"] = node.id - rec["lineage_parent_id"] = ( - node.parent_id if node.parent_id is not None else -1 - ) - rec["lineage_generation"] = node.generation - rec["lineage_abundance"] = node.abundance - # AIRR-standard abundance field that abundance-aware tools (Change-O / - # SCOPer / dowser) read. Mirror lineage_abundance so the genotype- - # collapsed observed cell carries its represented count both ways. - rec["duplicate_count"] = node.abundance - rec["lineage_affinity"] = node.affinity - rec["sequence_id"] = f"clone{clone_idx}_node{node.id}" - - def __repr__(self) -> str: - return ( - f"" - ) - - -class CompiledRepertoireExperiment: - """A compiled non-tree clonal-repertoire experiment. - - Wraps a pre-fork :class:`GenAIRR._engine.CompiledSimulator` (the - founder recombination, run once per clone) and an optional - post-fork simulator (the per-read library-prep / sequencing passes). - Per clone a size is drawn from a heavy-tailed distribution via - ``_engine.sample_clone_sizes``; that many reads are emitted through - the post-fork passes and identical reads are genotype-collapsed into - AIRR records carrying a standard ``duplicate_count``. - - When there are no post-fork passes every read of a clone is - identical, so the clone collapses to a single record whose - ``duplicate_count`` equals the drawn size (the no-corruption - shortcut). Every record carries an integer ``clone_id``. - """ - - __slots__ = ( - "_pre", - "_post", - "_step", - "_refdata", - "_dataconfig", - "_metadata", - ) - - def __init__( - self, - pre_simulator: "_engine.CompiledSimulator", - post_simulator: Optional["_engine.CompiledSimulator"], - step: "_RepertoireForkStep", - refdata: "_engine.RefDataConfig", - dataconfig: Optional["DataConfig"] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> None: - self._pre = pre_simulator - self._post = post_simulator - self._step = step - self._refdata = refdata - self._dataconfig = dataconfig - self._metadata = dict(metadata) if metadata else {} - - @property - def refdata(self) -> "_engine.RefDataConfig": - return self._refdata - - def run_records( - self, - *, - seed: int = 0, - strict: bool = False, - validate_records: bool = False, - expose_provenance: bool = False, - ) -> "SimulationResult": - """Draw per-clone sizes, emit reads through the post-fork passes, - collapse identical reads, and return a :class:`SimulationResult` - whose record dicts carry ``clone_id`` and ``duplicate_count``. - """ - import GenAIRR._engine as _engine - - from ._airr_record import outcome_to_airr_record - from .result import SimulationResult, _inject_truth_columns - - step = self._step - sizes = _engine.sample_clone_sizes( - step.n_clones, - int(seed), - kind=step.size_distribution, - exponent=step.exponent, - mu=step.mu, - sigma=step.sigma, - max_size=step.max_size, - unexpanded_fraction=step.unexpanded_fraction, - ) - records: List[Dict[str, Any]] = [] - outcomes: List["_engine.Outcome"] = [] - for clone_idx, size in enumerate(sizes): - clone_seed = int(seed) + clone_idx * 1_000_000 - founder = self._pre.run(seed=clone_seed, strict=strict) - founder_sim = founder.final_simulation() - if self._post is None: - # No post-fork passes: all ``size`` copies are - # identical -> one record with duplicate_count = size. - rec = outcome_to_airr_record( - founder, - self._refdata, - sequence_id=f"clone{clone_idx}_read0", - ) - rec["clone_id"] = clone_idx - rec["duplicate_count"] = int(size) - if expose_provenance: - _inject_truth_columns(founder, self._refdata, rec) - records.append(rec) - outcomes.append(founder) - continue - # Post-fork passes present: simulate ``size`` reads and - # collapse by emitted sequence. - by_seq: Dict[str, list] = {} - for read_idx in range(int(size)): - desc = self._post.run_from( - founder_sim, clone_seed + 1 + read_idx, strict=strict - ) - rec = outcome_to_airr_record( - desc, - self._refdata, - sequence_id=f"clone{clone_idx}_read{read_idx}", - ) - key = rec["sequence"] - if key in by_seq: - by_seq[key][0] += 1 - else: - by_seq[key] = [1, desc, rec] - for count, desc, rec in by_seq.values(): - rec["clone_id"] = clone_idx - rec["duplicate_count"] = count - if expose_provenance: - _inject_truth_columns(desc, self._refdata, rec) - records.append(rec) - outcomes.append(desc) - result = SimulationResult(records, outcomes=outcomes) - if validate_records: - from ._validation import ( - _raise_on_family_validation_failure, - _raise_on_validation_failure, - ) - - _raise_on_validation_failure(result.validate_records(self._refdata)) - _raise_on_family_validation_failure(result.validate_families()) - return result - - def __repr__(self) -> str: - return ( - f"" - ) diff --git a/src/GenAIRR/_compiled/__init__.py b/src/GenAIRR/_compiled/__init__.py new file mode 100644 index 0000000..e3aca76 --- /dev/null +++ b/src/GenAIRR/_compiled/__init__.py @@ -0,0 +1,12 @@ +"""Compiled/runtime experiment wrappers, one module per runtime shape.""" +from .plain import CompiledExperiment +from .clonal import CompiledClonalExperiment +from .lineage import CompiledLineageExperiment +from .repertoire import CompiledRepertoireExperiment + +__all__ = [ + "CompiledExperiment", + "CompiledClonalExperiment", + "CompiledLineageExperiment", + "CompiledRepertoireExperiment", +] diff --git a/src/GenAIRR/_compiled/clonal.py b/src/GenAIRR/_compiled/clonal.py new file mode 100644 index 0000000..a14f8f1 --- /dev/null +++ b/src/GenAIRR/_compiled/clonal.py @@ -0,0 +1,267 @@ +"""Runtime wrapper for the clonal-fork compiled experiment. + +:class:`CompiledClonalExperiment` wraps the two-stage +parent-then-descendants flow produced by ``expand_clones(...)``. It +holds two compiled engine simulators (pre-fork and post-fork) plus +source context so ``describe()`` and AIRR record export can render a +faithful narrative. The execution loop delegates straight into the +engine; the wrapper exists for ergonomics and to keep the public +surface stable across engine refactors. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING + +from GenAIRR import _engine # private Rust extension submodule + +from .._describe import ( + _describe_clonal_fork_step, + _describe_experiment_header, + _describe_step_sequence, +) +from .._pipeline_ir import _ClonalForkStep + +if TYPE_CHECKING: + from ..dataconfig import DataConfig + from ..result import SimulationResult + + +class CompiledClonalExperiment: + """A compiled experiment with a clonal-fork structure. + + Wraps two :class:`GenAIRR._engine.CompiledSimulator`s — the + pre-fork plan (run once per clone, typically the recombine + step) and the post-fork plan (run once per descendant inside + the clone, typically mutate / corrupt_*). + + :meth:`run_records` orchestrates the parent → descendants loop + and tags every record with a ``clone_id`` integer so downstream + clonotype-clustering tools can be benchmarked against the true + clonal structure. + """ + + __slots__ = ( + "_pre", + "_post", + "_fork", + "_refdata", + "_pre_steps", + "_post_steps", + "_dataconfig", + "_metadata", + ) + + def __init__( + self, + pre_simulator: "_engine.CompiledSimulator", + post_simulator: "_engine.CompiledSimulator", + fork: "_ClonalForkStep", + refdata: "_engine.RefDataConfig", + pre_steps: Sequence[Any] = (), + post_steps: Sequence[Any] = (), + dataconfig: Optional["DataConfig"] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: + self._pre = pre_simulator + self._post = post_simulator + self._fork = fork + self._refdata = refdata + self._pre_steps: Tuple[Any, ...] = tuple(pre_steps) + self._post_steps: Tuple[Any, ...] = tuple(post_steps) + self._dataconfig = dataconfig + self._metadata = dict(metadata) if metadata else {} + + @property + def n_clones(self) -> int: + return self._fork.n_clones + + @property + def size(self) -> int: + return self._fork.size + + @property + def total_records(self) -> int: + """Number of records produced per :meth:`run_records` call.""" + return self._fork.n_clones * self._fork.size + + @property + def refdata(self) -> "_engine.RefDataConfig": + return self._refdata + + def describe(self) -> str: + """Render a biology-style narrative of the compiled clonal + experiment, with an explicit divider at the fork. See + :meth:`Experiment.describe` for the basic shape.""" + header = _describe_experiment_header(self._refdata, self._dataconfig) + lines = [header] + # pre-fork section (per-clone) + pre_lines = _describe_step_sequence( + self._pre_steps, self._refdata.chain_type, start_index=1 + ) + lines.extend(pre_lines) + # the fork itself + lines.append(f" ── {_describe_clonal_fork_step(self._fork)} ──") + lines.append( + " (steps above run once per clone; " + "steps below run once per descendant)" + ) + # post-fork section (per-descendant) + post_start = sum( + 1 for s in self._pre_steps if not isinstance(s, _ClonalForkStep) + ) + 1 + post_lines = _describe_step_sequence( + self._post_steps, self._refdata.chain_type, start_index=post_start + ) + lines.extend(post_lines) + if self._metadata: + stamps = ", ".join(f"{k}={v!r}" for k, v in self._metadata.items()) + lines.append(f" Metadata stamped on every record: {stamps}") + return "\n".join(lines) + + def run( + self, + *, + n: Optional[int] = None, + seed: int = 0, + strict: bool = False, + ) -> List["_engine.Outcome"]: + """Run all clonal descendants and return their outcomes in + clone-major order (clone 0's descendants 0..size-1, clone 1's + descendants 0..size-1, …). + + ``n`` is optional: when omitted the runtime expands + ``n_clones * size`` outcomes. Passing ``n`` is allowed only + when ``n == n_clones * size`` (otherwise raises). + """ + total = self.total_records + if n is not None and n != total: + raise ValueError( + f"clonal pipeline produces n_clones * size = " + f"{self._fork.n_clones} * {self._fork.size} = {total} " + f"records; passing n={n} is inconsistent. Drop the n " + f"argument or pass n={total}." + ) + + outcomes: List["_engine.Outcome"] = [] + for clone_idx in range(self._fork.n_clones): + clone_seed = int(seed) + clone_idx * 1_000_000 + parent = self._pre.run(seed=clone_seed, strict=strict) + parent_sim = parent.final_simulation() + for desc_idx in range(self._fork.size): + desc_seed = clone_seed + 1 + desc_idx + desc = self._post.run_from( + parent_sim, desc_seed, strict=strict + ) + outcomes.append(desc) + return outcomes + + def run_records( + self, + *, + n: Optional[int] = None, + seed: int = 0, + strict: bool = False, + expose_provenance: bool = False, + validate_records: bool = False, + ) -> "SimulationResult": + """Same as :meth:`run` but returns a :class:`SimulationResult` + with each record dict carrying an integer ``clone_id`` field + in ``[0, n_clones)`` plus a ``parent_id`` integer indexing + into :attr:`SimulationResult.parents`. ``expose_provenance=True`` + also appends `truth_v_call` / `truth_d_call` / `truth_j_call` + columns from the originally-sampled allele names. + + The returned :class:`SimulationResult` carries the per-clone + parent ``Outcome`` objects on its ``.parents`` attribute. + Each parent holds the pre-fork addressed-choice trace, the + pre-fork event ledger, and the post-recombination IR — useful + for replay, lineage analysis, and the upcoming parent-aware + family validator. The flat ``.outcomes`` list continues to + carry only the descendant outcomes (one per record); + parents are exposed separately so the per-record list stays + the same shape clonal consumers already know. + + ``validate_records=True`` runs + :meth:`SimulationResult.validate_records` on the freshly + built batch and raises + :class:`GenAIRR._validation.RecordValidationFailedError` + on any postcondition failure. After the per-record gate + passes, this also runs + :meth:`SimulationResult.validate_families` and raises the + sibling + :class:`GenAIRR._validation.FamilyValidationFailedError` + if any clonal-family invariant is violated. The two + gates report separately so users can tell projection bugs + from family-consistency bugs. Default ``False`` keeps this + method zero-overhead. + """ + from .._airr_record import outcome_to_airr_record + from ..result import SimulationResult, _inject_truth_columns + + total = self.total_records + if n is not None and n != total: + raise ValueError( + f"clonal pipeline produces n_clones * size = " + f"{self._fork.n_clones} * {self._fork.size} = {total} " + f"records; passing n={n} is inconsistent. Drop the n " + f"argument or pass n={total}." + ) + + records: List[Dict[str, Any]] = [] + outcomes: List["_engine.Outcome"] = [] + # Slice 2: retain parent outcomes — one per clone. The + # orchestration loop used to drop them after extracting + # ``final_simulation()``. We now keep them so the returned + # :class:`SimulationResult` can expose ``.parents`` for + # replay / lineage tooling. The parent ``Outcome`` itself + # is not copied onto each descendant; we hold a single + # reference per clone in the ``parents`` list. + parents: List["_engine.Outcome"] = [] + for clone_idx in range(self._fork.n_clones): + clone_seed = int(seed) + clone_idx * 1_000_000 + parent = self._pre.run(seed=clone_seed, strict=strict) + parents.append(parent) + parent_sim = parent.final_simulation() + for desc_idx in range(self._fork.size): + desc_seed = clone_seed + 1 + desc_idx + desc = self._post.run_from( + parent_sim, desc_seed, strict=strict + ) + outcomes.append(desc) + rec = outcome_to_airr_record( + desc, + self._refdata, + sequence_id=f"clone{clone_idx}_desc{desc_idx}", + ) + rec["clone_id"] = clone_idx + # ``parent_id`` is the descendant's index into + # ``result.parents``. Today clones are dense and + # zero-based, so ``parent_id == clone_id`` by + # construction — we stamp both because they carry + # distinct semantics: ``clone_id`` is the family + # identity (the existing Slice 0 contract); + # ``parent_id`` is the addressing scheme into the + # parent-outcome list (the new Slice 2 contract). + # Keeping them separate now means a future slice + # that introduces sparse / non-zero-based family + # ids doesn't have to retrofit both. + rec["parent_id"] = clone_idx + if expose_provenance: + _inject_truth_columns(desc, self._refdata, rec) + records.append(rec) + result = SimulationResult(records, outcomes=outcomes, parents=parents) + if validate_records: + from .._validation import ( + _raise_on_family_validation_failure, + _raise_on_validation_failure, + ) + + _raise_on_validation_failure(result.validate_records(self._refdata)) + _raise_on_family_validation_failure(result.validate_families()) + return result + + def __repr__(self) -> str: + return ( + f"" + ) diff --git a/src/GenAIRR/_compiled/lineage.py b/src/GenAIRR/_compiled/lineage.py new file mode 100644 index 0000000..4c71b7f --- /dev/null +++ b/src/GenAIRR/_compiled/lineage.py @@ -0,0 +1,263 @@ +"""Runtime wrapper for the BCR affinity-maturation lineage experiment. + +:class:`CompiledLineageExperiment` grows one lineage tree per clone via +the Rust ``simulate_family_outcomes`` kernel and returns per-observed-node +AIRR records with lineage metadata. It is intentionally thin: it holds a +reference to the compiled engine simulators and source context so record +export can render a faithful narrative. The execution loop delegates +straight into the engine; the wrapper exists for ergonomics and to keep +the public surface stable across engine refactors. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING + +from GenAIRR import _engine # private Rust extension submodule + +from .._pipeline_ir import _LineageForkStep + +if TYPE_CHECKING: + from ..dataconfig import DataConfig + from ..result import SimulationResultWithLineages + + +class CompiledLineageExperiment: + """A compiled experiment that grows BCR affinity-maturation lineage trees. + + Wraps a pre-fork :class:`GenAIRR._engine.CompiledSimulator` (founder + recombination) and a :class:`~GenAIRR._pipeline_ir._LineageForkStep` + (lineage parameters). :meth:`run_records` grows one lineage tree per + clone via the Rust ``simulate_family_outcomes`` kernel and returns a + :class:`~GenAIRR.result.SimulationResultWithLineages` whose records are + per-observed-node AIRR dicts with lineage metadata. + """ + + __slots__ = ( + "_pre", + "_step", + "_refdata", + "_post", + "_post_steps", + "_dataconfig", + "_metadata", + ) + + def __init__( + self, + pre_simulator: "_engine.CompiledSimulator", + step: "_LineageForkStep", + refdata: "_engine.RefDataConfig", + post_simulator: Optional["_engine.CompiledSimulator"] = None, + post_steps: Sequence[Any] = (), + dataconfig: Optional["DataConfig"] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: + self._pre = pre_simulator + self._step = step + self._refdata = refdata + # Per-observed-cell library-prep / sequencing corruption plan. + # ``None`` keeps the pristine-read path (no artefacts applied). + self._post = post_simulator + self._post_steps: Tuple[Any, ...] = tuple(post_steps) + self._dataconfig = dataconfig + self._metadata = dict(metadata) if metadata else {} + + @property + def refdata(self) -> "_engine.RefDataConfig": + return self._refdata + + def run_records( + self, + *, + seed: int = 0, + strict: bool = False, + expose_provenance: bool = False, + validate_records: bool = False, + ) -> "SimulationResultWithLineages": + """Grow lineage trees and return per-observed-node AIRR records. + + Each clone is seeded at ``seed + clone_idx * 1_000_000`` so + independent clones are reproducible and non-overlapping. + + The returned :class:`SimulationResultWithLineages` carries the + per-record ``Outcome`` that produced each AIRR record on its + ``.outcomes`` attribute (index-aligned with ``.records``). On the + pristine-read path these are the per-observed-node SHM + ``Outcome`` objects (``fam.observed_outcomes()``); on the + corruption path they are the per-node merged + recombination+SHM+artefact ``Outcome`` objects. Since each node + ``Outcome`` is self-consistent, ``result.validate_records(refdata)`` + passes. + + ``expose_provenance=True`` appends ``truth_v_call`` / + ``truth_d_call`` / ``truth_j_call`` columns to each record from + the founder allele assignments carried on the per-record + ``Outcome``. + + ``validate_records=True`` runs + :meth:`SimulationResult.validate_records` (per-record + postcondition) and :meth:`SimulationResult.validate_families` + (clonal-family consistency by ``clone_id``) on the freshly built + batch, raising the matching validation error on any failure. + """ + from GenAIRR import _engine as _eng + from .._airr_record import outcome_to_airr_record + from .._s5f_loader import load_builtin_s5f_kernel + from ..result import SimulationResultWithLineages, _inject_truth_columns + + step = self._step + mutability, substitution = load_builtin_s5f_kernel(step.s5f_model) + + records: List[Dict[str, Any]] = [] + # Per-record source ``Outcome`` objects, index-aligned with + # ``records`` so ``result.validate_records`` can re-derive each + # record from the engine state that produced it. + outcomes: List["_engine.Outcome"] = [] + lineage_trees = [] + + # Bound on founder-survival retries. Sampling draws from the LIVING + # final-generation population, so an extinct founder (drew 0 offspring) + # yields zero observed cells. With ``allow_extinction=False`` (default) + # we condition each clone on survival by re-growing the family with a + # fresh deterministic sub-seed until it survives or we exhaust the + # bound. The sub-seed offset (a prime times the attempt index) keeps the + # retries deterministic — same top-level seed => same result. + _MAX_SURVIVAL_ATTEMPTS = 50 + _SUBSEED_STRIDE = 7_919 # prime; keeps retry sub-seeds well-separated + allow_extinction = getattr(step, "allow_extinction", False) + + for clone_idx in range(step.n_clones): + clone_seed = int(seed) + clone_idx * 1_000_000 + fam = None + tree = None + for attempt in range(_MAX_SURVIVAL_ATTEMPTS): + family_seed = clone_seed + attempt * _SUBSEED_STRIDE + # _pre.run() returns a single Outcome (not a list) — mirror + # CompiledClonalExperiment which calls self._pre.run(seed=...). + founder = self._pre.run(seed=family_seed, strict=strict) + candidate = _eng.simulate_family_outcomes( + founder, + self._refdata, + mutability, + substitution, + step.rate, + step.lambda_base, + 0.0, # lambda_mut: positional slot 6 (inert; hardcoded) + step.max_generations, + step.n_max, + step.n_sample, + family_seed, + selection_strength=step.selection_strength, + beta=step.beta, + target_aa=step.target_aa, + mature_substitutions=step.mature_substitutions, + ) + survived = len(candidate.observed_outcomes()) > 0 + if survived or allow_extinction: + fam = candidate + tree = candidate.tree() + break + # extinct + survival required => retry with next sub-seed + else: + # Exhausted the retry bound without a surviving family. + raise ValueError( + f"clonal_lineage: clone {clone_idx} went extinct on every " + f"one of {_MAX_SURVIVAL_ATTEMPTS} survival attempts (each " + "founder drew 0 offspring). Increase lambda_base (offspring " + "Poisson mean) and/or max_generations so families reliably " + "survive, or pass allow_extinction=True to accept extinct " + "clones (producing fewer than n_clones families)." + ) + + if fam is None: + # allow_extinction=True and the final attempt was extinct: + # skip this clone (it contributes no observed cells/records). + continue + + lineage_trees.append(tree) + if len(fam.observed_outcomes()) == 0: + # allow_extinction=True and this family is extinct: keep its + # (empty) tree for export parity but emit no records. + continue + + if self._post is None: + # Pristine-read path: project each observed node's + # synthesized SHM Outcome straight to an AIRR record. + # ``observed_outcomes()`` is index-aligned with both the + # observed-node list and ``airr_records()``. + observed_nodes = [n for n in tree.nodes() if n.observed] + recs = fam.airr_records(self._refdata) + observed_outcomes = fam.observed_outcomes() + for node, rec, node_outcome in zip( + observed_nodes, recs, observed_outcomes + ): + self._stamp_lineage_metadata(rec, clone_idx, node) + if expose_provenance and node_outcome is not None: + _inject_truth_columns(node_outcome, self._refdata, rec) + records.append(rec) + outcomes.append(node_outcome) + continue + + # Corruption path: per observed node, run the library-prep / + # sequencing corruption plan FROM the node's post-SHM + # simulation, then merge the founder-recombination + SHM + # provenance with the corruption trace / events so the AIRR + # record reports trims, v/d/j, SHM counts AND artefact + # counters (n_quality_errors, n_pcr_errors, n_indels, …). + node_outcomes = fam.node_outcomes() + for node, base_outcome in zip(tree.nodes(), node_outcomes): + if base_outcome is None: + continue + node_seed = clone_seed + 1 + node.id + corruption_outcome = self._post.run_from( + base_outcome.final_simulation(), node_seed, strict=strict + ) + merged = _eng.merge_lineage_corruption( + base_outcome, corruption_outcome + ) + rec = outcome_to_airr_record( + merged, + self._refdata, + sequence_id=f"clone{clone_idx}_node{node.id}", + ) + self._stamp_lineage_metadata(rec, clone_idx, node) + if expose_provenance: + _inject_truth_columns(merged, self._refdata, rec) + records.append(rec) + outcomes.append(merged) + + result = SimulationResultWithLineages( + records, outcomes=outcomes, lineage_trees=lineage_trees + ) + if validate_records: + from .._validation import ( + _raise_on_family_validation_failure, + _raise_on_validation_failure, + ) + + _raise_on_validation_failure(result.validate_records(self._refdata)) + _raise_on_family_validation_failure(result.validate_families()) + return result + + @staticmethod + def _stamp_lineage_metadata(rec: Dict[str, Any], clone_idx: int, node) -> None: + """Stamp clone + lineage-node provenance onto an AIRR record.""" + rec["clone_id"] = clone_idx + rec["lineage_node_id"] = node.id + rec["lineage_parent_id"] = ( + node.parent_id if node.parent_id is not None else -1 + ) + rec["lineage_generation"] = node.generation + rec["lineage_abundance"] = node.abundance + # AIRR-standard abundance field that abundance-aware tools (Change-O / + # SCOPer / dowser) read. Mirror lineage_abundance so the genotype- + # collapsed observed cell carries its represented count both ways. + rec["duplicate_count"] = node.abundance + rec["lineage_affinity"] = node.affinity + rec["sequence_id"] = f"clone{clone_idx}_node{node.id}" + + def __repr__(self) -> str: + return ( + f"" + ) diff --git a/src/GenAIRR/_compiled/plain.py b/src/GenAIRR/_compiled/plain.py new file mode 100644 index 0000000..f7023f3 --- /dev/null +++ b/src/GenAIRR/_compiled/plain.py @@ -0,0 +1,280 @@ +"""Runtime wrapper for the plain (non-forked) compiled experiment. + +:class:`CompiledExperiment` is the frozen, executable form of +:class:`GenAIRR.experiment.Experiment`. It is intentionally thin: it +holds a reference to a compiled engine simulator and a few pieces of +source context (``steps``, ``dataconfig``, ``metadata``) so +``describe()`` and AIRR record export can render a faithful narrative. +The execution loops delegate straight into the engine; the wrapper +exists for ergonomics and to keep the public surface stable across +engine refactors. +""" +from __future__ import annotations + +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, TYPE_CHECKING + +from GenAIRR import _engine # private Rust extension submodule + +from .._describe import ( + _describe_experiment_header, + _describe_step_sequence, + _format_active_contracts, +) + +if TYPE_CHECKING: + from ..dataconfig import DataConfig + from ..result import SimulationResult + + +class CompiledExperiment: + """A frozen ``Experiment`` ready for execution. + + Holds the owning :class:`GenAIRR._engine.CompiledSimulator` and the + refdata it was built against. Contracts are captured at compile + time; ``run()`` only accepts execution parameters. + """ + + __slots__ = ( + "_simulator", + "_refdata", + "_steps", + "_dataconfig", + "_metadata", + "_genotype", + ) + + def __init__( + self, + simulator: "_engine.CompiledSimulator", + refdata: "_engine.RefDataConfig", + steps: Sequence[Any] = (), + dataconfig: Optional["DataConfig"] = None, + metadata: Optional[Dict[str, Any]] = None, + genotype: Optional[Any] = None, + ) -> None: + self._simulator = simulator + self._refdata = refdata + # Source steps stashed for `describe()`. The compiled simulator + # itself can render pass names but not biology — keeping the + # builder steps around is the cheapest way to give a faithful + # narrative back. + self._steps: Tuple[Any, ...] = tuple(steps) + self._dataconfig = dataconfig + self._metadata = dict(metadata) if metadata else {} + # Attached single-subject genotype (or None). When set, + # run_records stamps subject_id + haplotype provenance. + self._genotype = genotype + + @property + def simulator(self) -> "_engine.CompiledSimulator": + """The owning Rust compiled simulator.""" + return self._simulator + + @property + def pass_plan(self) -> Tuple[str, ...]: + """Read-only pass-name summary of the compiled pipeline.""" + return tuple(self._simulator.pass_names()) + + @property + def pass_names(self) -> Tuple[str, ...]: + """Stable names of the compiled pass sequence.""" + return self.pass_plan + + @property + def active_contracts(self) -> Tuple[str, ...]: + """Stable names of the contract bundle captured at compile time.""" + return tuple(self._simulator.active_contracts()) + + @property + def refdata(self) -> "_engine.RefDataConfig": + """The :class:`GenAIRR._engine.RefDataConfig` the plan was built against.""" + return self._refdata + + def describe(self) -> str: + """Render a biology-style narrative of the compiled experiment. + + Equivalent to ``Experiment.describe()`` but additionally + surfaces compile-time constraints (e.g. ``productive_only``) + attached via ``compile(respect=...)``. See + :meth:`Experiment.describe` for the output shape. + """ + header = _describe_experiment_header(self._refdata, self._dataconfig) + if not self._steps: + body = [" (no steps recorded)"] + else: + body = _describe_step_sequence(self._steps, self._refdata.chain_type) + lines = [header, *body] + contracts_line = _format_active_contracts(self.active_contracts) + if contracts_line: + lines.append(f" Constraints: {contracts_line}") + if self._metadata: + stamps = ", ".join(f"{k}={v!r}" for k, v in self._metadata.items()) + lines.append(f" Metadata stamped on every record: {stamps}") + return "\n".join(lines) + + def run( + self, + *, + n: int = 1, + seed: int = 0, + strict: bool = False, + ) -> List["_engine.Outcome"]: + """Run the compiled simulator ``n`` times — **fresh sampling**. + + Each iteration uses ``seed + i`` as the per-run seed so + consecutive batches stitch together by offsetting ``seed``. + + ``strict`` controls the failure mode when a pass's + contract-narrowed candidate set is empty *at sample time*: + + - ``False`` (default, **permissive**) — apply the pass's + declared ``EmptySupport`` policy. The pass writes a + documented sentinel value to the trace and continues. + Common sentinels: indel ``site = -1`` NoOp, NP length ``0``, + NP base ``N``, trim ``0``; SHM substitution skips the slot + (no trace record). + - ``True`` (**strict**) — raise + :class:`GenAIRR._engine.StrictSamplingError`. The exception's + ``args`` are a 3-tuple ``(pass_name, address, reason)``; + ``reason`` is one of ``"support_unavailable"``, + ``"empty_admissible_support"``, or + ``"invalid_filtered_support"``. + + **Compile-time precondition failures are separate.** If a + sampling distribution is *statically* incompatible with the + active contracts (e.g. every NP1 length in the distribution + violates frame divisibility), :meth:`Experiment.compile` + raises :class:`ValueError` *before* this method runs. + ``ValueError`` and ``StrictSamplingError`` have **no shared + base class** — catching only one will miss the other. See + ``docs/productive_failure_mode_audit.md`` §6.1. + + **Strict semantics apply only to fresh sampling.** Trace + replay via :meth:`CompiledExperiment.replay_from_trace_file` + consumes the recorded values verbatim and does NOT + re-evaluate contract admissibility — a permissive sentinel + trace replays cleanly even under ``strict=True``. See that + method's docstring. + + Raises ``ValueError`` for ``n < 1``. + """ + if n < 1: + raise ValueError(f"n must be at least 1, got {n}") + return self._simulator.run_batch(n, seed, strict=strict) + + def run_records( + self, + *, + n: int = 1, + seed: int = 0, + strict: bool = False, + expose_provenance: bool = False, + validate_records: bool = False, + ) -> "SimulationResult": + """Run the compiled simulator ``n`` times and return the batch as + a :class:`SimulationResult` ready for ``.to_csv`` / + ``.to_fasta`` / ``.to_dataframe`` export. + + Same arguments as :meth:`run`. ``expose_provenance=True`` + appends `truth_v_call/d_call/j_call` columns reflecting the + originally-sampled allele names. + + ``validate_records=True`` runs + :meth:`SimulationResult.validate_records` on the freshly + built batch and raises + :class:`GenAIRR._validation.RecordValidationFailedError` + (a :class:`RuntimeError` subclass) when any record fails the + postcondition validator. Default ``False`` keeps this method + zero-overhead. + """ + from ..result import SimulationResult + + outcomes = self.run(n=n, seed=seed, strict=strict) + result = SimulationResult.from_outcomes( + outcomes, self._refdata, expose_provenance=expose_provenance + ) + if self._genotype is not None: + self._stamp_genotype_provenance(outcomes, result) + if validate_records: + from .._validation import _raise_on_validation_failure + + _raise_on_validation_failure(result.validate_records(self._refdata)) + return result + + def _stamp_genotype_provenance(self, outcomes, result) -> None: + """Stamp per-record ``subject_id`` + ``haplotype`` (the chromosome + the rearrangement drew from) and expose the genotype on the + result. Used only when a genotype is attached.""" + subject = self._genotype.subject_id + for outcome, rec in zip(outcomes, result._records): + rec["subject_id"] = subject + hap = outcome.trace().find("sample_haplotype") + rec["haplotype"] = hap.value if hap is not None else None + result._genotypes = [self._genotype] + + def stream( + self, + *, + n: Optional[int] = None, + seed: int = 0, + strict: bool = False, + ) -> Iterator["_engine.Outcome"]: + """Lazily yield :class:`GenAIRR._engine.Outcome` objects one at + a time, without materialising the full batch in memory. + + Useful for large simulations where holding ``n`` outcomes + would be wasteful — typical pattern is + + >>> for outcome in compiled.stream(n=1_000_000, seed=0): + ... write_to_disk(outcome) + + ``n=None`` (the default) yields outcomes indefinitely; the + caller is expected to stop with ``itertools.islice``, + ``break``, or similar. ``n=N`` yields exactly ``N`` outcomes + with seeds ``seed`` … ``seed + N - 1``. + + ``strict`` behaves as in :meth:`run`. + + Raises ``ValueError`` when ``n`` is set to a value below 1. + """ + if n is not None and n < 1: + raise ValueError(f"n must be at least 1, got {n}") + i = 0 + while n is None or i < n: + yield self._simulator.run(seed + i, strict=strict) + i += 1 + + def stream_records( + self, + *, + n: Optional[int] = None, + seed: int = 0, + strict: bool = False, + id_prefix: str = "seq", + ) -> Iterator[Dict[str, Any]]: + """Lazily yield AIRR-format record dicts (one per outcome). + + Same shape as the records inside a :class:`SimulationResult`, + but yielded one at a time so callers can write each record to + disk without retaining the prior ones. Pairs naturally with + :func:`csv.DictWriter` for streaming TSV/CSV output. + + Each record's ``sequence_id`` is set to + ``f"{id_prefix}{i}"`` so streamed batches have unique + AIRR-style identifiers without buffering. + """ + from .._airr_record import outcome_to_airr_record + + for i, outcome in enumerate( + self.stream(n=n, seed=seed, strict=strict) + ): + yield outcome_to_airr_record( + outcome, self._refdata, sequence_id=f"{id_prefix}{i}" + ) + + def __repr__(self) -> str: + return ( + f"" + ) diff --git a/src/GenAIRR/_compiled/repertoire.py b/src/GenAIRR/_compiled/repertoire.py new file mode 100644 index 0000000..f5f84a4 --- /dev/null +++ b/src/GenAIRR/_compiled/repertoire.py @@ -0,0 +1,159 @@ +"""Runtime wrapper for the non-tree clonal-repertoire experiment. + +:class:`CompiledRepertoireExperiment` draws a heavy-tailed size per clone, +emits that many reads through the post-fork passes, and genotype-collapses +identical reads into AIRR records carrying a standard ``duplicate_count``. +It is intentionally thin: it holds references to the compiled engine +simulators and source context. The execution loop delegates straight into +the engine; the wrapper exists for ergonomics and to keep the public +surface stable across engine refactors. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional, TYPE_CHECKING + +from GenAIRR import _engine # private Rust extension submodule + +from .._pipeline_ir import _RepertoireForkStep + +if TYPE_CHECKING: + from ..dataconfig import DataConfig + from ..result import SimulationResult + + +class CompiledRepertoireExperiment: + """A compiled non-tree clonal-repertoire experiment. + + Wraps a pre-fork :class:`GenAIRR._engine.CompiledSimulator` (the + founder recombination, run once per clone) and an optional + post-fork simulator (the per-read library-prep / sequencing passes). + Per clone a size is drawn from a heavy-tailed distribution via + ``_engine.sample_clone_sizes``; that many reads are emitted through + the post-fork passes and identical reads are genotype-collapsed into + AIRR records carrying a standard ``duplicate_count``. + + When there are no post-fork passes every read of a clone is + identical, so the clone collapses to a single record whose + ``duplicate_count`` equals the drawn size (the no-corruption + shortcut). Every record carries an integer ``clone_id``. + """ + + __slots__ = ( + "_pre", + "_post", + "_step", + "_refdata", + "_dataconfig", + "_metadata", + ) + + def __init__( + self, + pre_simulator: "_engine.CompiledSimulator", + post_simulator: Optional["_engine.CompiledSimulator"], + step: "_RepertoireForkStep", + refdata: "_engine.RefDataConfig", + dataconfig: Optional["DataConfig"] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: + self._pre = pre_simulator + self._post = post_simulator + self._step = step + self._refdata = refdata + self._dataconfig = dataconfig + self._metadata = dict(metadata) if metadata else {} + + @property + def refdata(self) -> "_engine.RefDataConfig": + return self._refdata + + def run_records( + self, + *, + seed: int = 0, + strict: bool = False, + validate_records: bool = False, + expose_provenance: bool = False, + ) -> "SimulationResult": + """Draw per-clone sizes, emit reads through the post-fork passes, + collapse identical reads, and return a :class:`SimulationResult` + whose record dicts carry ``clone_id`` and ``duplicate_count``. + """ + import GenAIRR._engine as _engine + + from .._airr_record import outcome_to_airr_record + from ..result import SimulationResult, _inject_truth_columns + + step = self._step + sizes = _engine.sample_clone_sizes( + step.n_clones, + int(seed), + kind=step.size_distribution, + exponent=step.exponent, + mu=step.mu, + sigma=step.sigma, + max_size=step.max_size, + unexpanded_fraction=step.unexpanded_fraction, + ) + records: List[Dict[str, Any]] = [] + outcomes: List["_engine.Outcome"] = [] + for clone_idx, size in enumerate(sizes): + clone_seed = int(seed) + clone_idx * 1_000_000 + founder = self._pre.run(seed=clone_seed, strict=strict) + founder_sim = founder.final_simulation() + if self._post is None: + # No post-fork passes: all ``size`` copies are + # identical -> one record with duplicate_count = size. + rec = outcome_to_airr_record( + founder, + self._refdata, + sequence_id=f"clone{clone_idx}_read0", + ) + rec["clone_id"] = clone_idx + rec["duplicate_count"] = int(size) + if expose_provenance: + _inject_truth_columns(founder, self._refdata, rec) + records.append(rec) + outcomes.append(founder) + continue + # Post-fork passes present: simulate ``size`` reads and + # collapse by emitted sequence. + by_seq: Dict[str, list] = {} + for read_idx in range(int(size)): + desc = self._post.run_from( + founder_sim, clone_seed + 1 + read_idx, strict=strict + ) + rec = outcome_to_airr_record( + desc, + self._refdata, + sequence_id=f"clone{clone_idx}_read{read_idx}", + ) + key = rec["sequence"] + if key in by_seq: + by_seq[key][0] += 1 + else: + by_seq[key] = [1, desc, rec] + for count, desc, rec in by_seq.values(): + rec["clone_id"] = clone_idx + rec["duplicate_count"] = count + if expose_provenance: + _inject_truth_columns(desc, self._refdata, rec) + records.append(rec) + outcomes.append(desc) + result = SimulationResult(records, outcomes=outcomes) + if validate_records: + from .._validation import ( + _raise_on_family_validation_failure, + _raise_on_validation_failure, + ) + + _raise_on_validation_failure(result.validate_records(self._refdata)) + _raise_on_family_validation_failure(result.validate_families()) + return result + + def __repr__(self) -> str: + return ( + f"" + ) diff --git a/src/GenAIRR/_experiment/__init__.py b/src/GenAIRR/_experiment/__init__.py new file mode 100644 index 0000000..c06d551 --- /dev/null +++ b/src/GenAIRR/_experiment/__init__.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from .clonal import _ClonalMixin +from .compile import _CompileMixin +from .constraints import _ConstraintsMixin +from .corruption import _CorruptionMixin +from .genotype_alleles import _GenotypeAllelesMixin +from .introspection import _IntrospectionMixin +from .mutation import _MutationMixin +from .recombination import _RecombinationMixin +from .refdata_controls import _RefdataControlsMixin +from .run import _RunMixin + +__all__ = [ + "_ClonalMixin", + "_CompileMixin", + "_ConstraintsMixin", + "_CorruptionMixin", + "_GenotypeAllelesMixin", + "_IntrospectionMixin", + "_MutationMixin", + "_RecombinationMixin", + "_RefdataControlsMixin", + "_RunMixin", +] diff --git a/src/GenAIRR/_experiment/_common.py b/src/GenAIRR/_experiment/_common.py new file mode 100644 index 0000000..52e82c8 --- /dev/null +++ b/src/GenAIRR/_experiment/_common.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from typing import Iterable, Union + + +# Sentinel for `using(...)` arguments that the caller did not pass at +# all. We can't use ``None`` for "unchanged" because ``None`` already +# means "clear the lock." +class _Unset: + __slots__ = () + + def __repr__(self) -> str: + return "" + + +_UNSET: _Unset = _Unset() + + +# Inputs accepted by :meth:`Experiment.restrict_alleles` per segment. ``None`` +# clears any prior lock; a single name is a one-allele lock; an +# iterable of names is a multi-allele subset; ``_UNSET`` (the default) +# leaves the existing lock unchanged. +_LockInput = Union[str, Iterable[str], None, _Unset] diff --git a/src/GenAIRR/_experiment/clonal.py b/src/GenAIRR/_experiment/clonal.py new file mode 100644 index 0000000..352a7b9 --- /dev/null +++ b/src/GenAIRR/_experiment/clonal.py @@ -0,0 +1,490 @@ +from __future__ import annotations + +import warnings +from typing import Optional + +from .._pipeline_ir import ( + _ClonalForkStep, + _LineageForkStep, + _RepertoireForkStep, +) +from .._step_validation import _descendant_phase_step_classifier + + +class _ClonalMixin: + __slots__ = () + + def expand_clones( + self, + *, + n_clones: int, + per_clone: int, + ) -> "Experiment": + """Expand the pipeline into clonal lineages. + + Marks the per-clone / per-descendant boundary in the chain: + steps appended *before* this call run **once per clone** — + typically just :meth:`recombine`, which establishes the + parent V/D/J + trim + NP + assembled IR for the clonal + family. Steps appended *after* this call run **once per + read inside the family** — typically :meth:`mutate` and the + library-prep / sequencing-stage steps, which introduce + per-read divergence within the clone. + + Concrete shape:: + + exp = (Experiment.on("human_igh") + .recombine() + .expand_clones(n_clones=10, per_clone=20) + .mutate(rate=0.05) + .pcr_amplify(count=2)) + result = exp.run_records(seed=0) + # 10 clones × 20 descendants = 200 records. + # Each record carries a ``clone_id`` integer in [0, 10). + + ``n`` can be omitted from :meth:`run_records` for a clonal + experiment — the runtime expands ``n_clones * per_clone`` + records automatically. Passing ``n`` is allowed only when + ``n == n_clones * per_clone``. + + .. deprecated:: + Use :meth:`clonal_lineage` for BCR affinity-maturation + trees, or :meth:`clonal_repertoire` for TCR / flat-BCR + abundance repertoires with clone-size distributions. + ``expand_clones`` remains supported for fixed-size flat + star expansion. + + Constraints: + - Both ``n_clones`` and ``per_clone`` must be positive ints. + - At most one expansion per pipeline; calling this method + twice raises ``ValueError``. + + Implementation note: the runtime forks the parent's IR + (final ``Simulation`` after the pre-fork plan) into + descendants by running the post-fork plan from that IR + with distinct seeds. Within a clone, every descendant + shares the same recombination provenance (V allele, trim, + NP bases) and only diverges through the post-fork passes. + """ + warnings.warn( + "Experiment.expand_clones() is deprecated. Use " + "Experiment.clonal_lineage() for BCR affinity-maturation trees " + "(internal SHM, lineage_trees, no per_clone) or " + "Experiment.clonal_repertoire() for TCR / flat-BCR abundance " + "repertoires with clone-size distributions and duplicate_count. " + "Neither is a drop-in replacement for fixed per_clone output; " + "expand_clones() remains supported for legacy fixed-size flat " + "star expansion.", + DeprecationWarning, + stacklevel=2, + ) + if not isinstance(n_clones, int) or isinstance(n_clones, bool) or n_clones < 1: + raise ValueError( + f"n_clones must be a positive int, got {n_clones!r}" + ) + if not isinstance(per_clone, int) or isinstance(per_clone, bool) or per_clone < 1: + raise ValueError(f"per_clone must be a positive int, got {per_clone!r}") + if self._has_clonal_fork(): + raise ValueError( + "expand_clones() / clonal_lineage() / clonal_repertoire() " + "can only be called once per pipeline" + ) + # Unified descendant-phase ordering guard. Scan the appended + # step list for any step that came from a descendant-phase + # DSL method (see ``_descendant_phase_step_classifier``). + # Pre-fork placement of these steps either silently misreports + # the AIRR field (trace-sourced fields don't survive the + # parent→descendant boundary — Bugs C / E / F) or collapses + # descendant diversity (the pass runs once on the parent IR + # and every clone member inherits an identical effect). + # Either failure mode is a clonal-semantics violation, so + # reject at the DSL boundary with a uniform message that + # names the offending step and the fix. + for step in self._steps: + offending_method = _descendant_phase_step_classifier(step) + if offending_method is None: + continue + if offending_method == "mutate": + detail = ( + "SHM is descendant-specific in GenAIRR's current " + "clonal model" + ) + else: + detail = ( + "it is descendant-specific and must be sampled " + "independently for each clone member" + ) + raise ValueError( + f"{offending_method} must be called after " + f"expand_clones(); {detail}. Move " + f"{offending_method}(...) after expand_clones(...)." + ) + self._steps.append(_ClonalForkStep(n_clones=n_clones, size=per_clone)) + return self + + def clonal_repertoire( + self, + *, + n_clones: int, + size_distribution: str = "power_law", + exponent: float = 2.0, + mu: float = 1.0, + sigma: float = 1.0, + max_size: int = 1000, + unexpanded_fraction: float = 0.0, + ) -> "Experiment": + """Expand the pipeline into a non-tree clonal **repertoire**. + + This is the non-tree clonal model (contrast with + :meth:`clonal_lineage`, which grows real affinity-maturation + lineage *trees*). It generalizes the deprecated + :meth:`expand_clones`: instead of a fixed ``per_clone`` size, + each clone draws a size from a heavy-tailed distribution (with + an unexpanded-singleton fraction). That many reads pass through + the post-fork library-prep / sequencing passes, and identical + reads are genotype-collapsed into AIRR records carrying a + standard ``duplicate_count`` that records abundance. + + Like :meth:`expand_clones`, this call marks the per-clone / + per-read boundary: steps appended *before* it run **once per + clone** (typically just :meth:`recombine`, establishing the + clonal V/D/J + trim + NP backbone); steps appended *after* it + run **once per read** (the library-prep / sequencing passes + that introduce per-read divergence). + + Concrete shape:: + + exp = (Experiment.on("human_igh") + .recombine() + .clonal_repertoire(n_clones=200, max_size=500, + unexpanded_fraction=0.3) + .sequencing_errors(rate=0.005)) + result = exp.run_records(seed=0) + # Each record carries a `clone_id` and a `duplicate_count`. + + Parameters: + - ``n_clones`` — number of clones (positive int). + - ``size_distribution`` — ``"power_law"`` or ``"lognormal"``. + - ``exponent`` — power-law exponent (``> 0``; used when + ``size_distribution="power_law"``). + - ``mu`` / ``sigma`` — lognormal parameters (``sigma >= 0``; + used when ``size_distribution="lognormal"``). + - ``max_size`` — clamps the largest clone. Because the total + number of reads simulated is roughly the **sum** of the drawn + sizes when post-fork passes are present, keep ``max_size`` + modest to bound runtime. + - ``unexpanded_fraction`` — fraction of clones forced to size 1 + (unexpanded singletons), in ``[0, 1]``. + + TCR works out of the box (no SHM): the reads diverge only + through the post-fork sequencing passes. Note that ``mutate`` + after this call is rejected on TCR by :meth:`mutate`'s own TCR + guard; on BCR an optional post-fork ``mutate`` adds flat SHM. + + **No-corruption shortcut:** a clone with no post-fork passes + emits identical copies, so it collapses to a single record whose + ``duplicate_count`` equals the drawn size. + + Constraints: + - At most one fork per pipeline — calling this when an + :meth:`expand_clones`, :meth:`clonal_lineage`, or + :meth:`clonal_repertoire` fork is already present raises + ``ValueError``. + - The same descendant-phase ordering guard as + :meth:`expand_clones` applies: a descendant-phase step (e.g. + ``.mutate()``) appended *before* this call is rejected; + :meth:`recombine` is fine. + """ + if not isinstance(n_clones, int) or isinstance(n_clones, bool) or n_clones < 1: + raise ValueError( + f"n_clones must be a positive int, got {n_clones!r}" + ) + if size_distribution not in ("power_law", "lognormal"): + raise ValueError( + "size_distribution must be 'power_law' or 'lognormal', got " + f"{size_distribution!r}" + ) + if not (isinstance(exponent, (int, float)) and exponent > 0): + raise ValueError(f"exponent must be > 0, got {exponent!r}") + if not (isinstance(sigma, (int, float)) and sigma >= 0): + raise ValueError(f"sigma must be >= 0, got {sigma!r}") + if not isinstance(max_size, int) or isinstance(max_size, bool) or max_size < 1: + raise ValueError(f"max_size must be a positive int, got {max_size!r}") + if not ( + isinstance(unexpanded_fraction, (int, float)) + and 0.0 <= unexpanded_fraction <= 1.0 + ): + raise ValueError( + "unexpanded_fraction must be in [0, 1], got " + f"{unexpanded_fraction!r}" + ) + if any( + isinstance(s, (_ClonalForkStep, _RepertoireForkStep, _LineageForkStep)) + for s in self._steps + ): + raise ValueError( + "clonal_repertoire() / expand_clones() / clonal_lineage() " + "can only be called once per pipeline" + ) + # Same descendant-phase ordering guard as expand_clones: a + # descendant-phase step (mutate / corruption / paired_end) + # appended before the fork is rejected. + for step in self._steps: + offending_method = _descendant_phase_step_classifier(step) + if offending_method is None: + continue + if offending_method == "mutate": + detail = ( + "SHM is descendant-specific in GenAIRR's current " + "clonal model" + ) + else: + detail = ( + "it is descendant-specific and must be sampled " + "independently for each read" + ) + raise ValueError( + f"{offending_method} must be called after " + f"clonal_repertoire(); {detail}. Move " + f"{offending_method}(...) after clonal_repertoire(...)." + ) + self._steps.append( + _RepertoireForkStep( + n_clones=n_clones, + size_distribution=size_distribution, + exponent=float(exponent), + mu=float(mu), + sigma=float(sigma), + max_size=max_size, + unexpanded_fraction=float(unexpanded_fraction), + ) + ) + return self + + def clonal_lineage( + self, + *, + n_clones: int, + max_generations: int = 10, + n_max: int = 1000, + n_sample: int = 50, + rate: float = 0.05, + lambda_base: float = 1.5, + selection_strength: float = 0.0, + beta: float = 1.0, + target_aa: Optional[str] = None, + mature_substitutions: int = 5, + s5f_model: str = "hh_s5f", + allow_extinction: bool = False, + ) -> "Experiment": + """Grow BCR lineage trees (neutral by default; set ``selection_strength > 0`` + and optionally ``target_aa`` to enable affinity maturation). + + Each clone gets its own lineage tree produced by the Rust + ``simulate_family_outcomes`` kernel. The returned + :class:`~GenAIRR.result.SimulationResultWithLineages` carries: + + - ``.records`` — one AIRR dict per *observed* (genotype-collapsed) cell, + tagged with ``clone_id``, ``lineage_node_id``, ``lineage_parent_id``, + ``lineage_generation``, ``lineage_abundance``, and + ``lineage_affinity``. Because identical genotypes are collapsed before + sampling, the number of records per clone is ≤ ``n_sample``; the + ``lineage_abundance`` field accounts for how many sampled cells were + represented by each observed record. Mutation counts (``n_mutations``, + ``n_v_mutations``, …) are pool-derived and self-consistent. + - ``.lineage_trees`` — one :class:`~GenAIRR._engine.LineageTree` + per clone for ground-truth export (Newick, FASTA, node table TSV). + + Parameters + ---------- + n_clones: + Number of independent clonal lineages to grow. + max_generations: + Maximum depth of the lineage tree (≤ 1000). + n_max: + Per-generation LIVING-population carrying capacity: the live + population per generation is capped at this (the tree can contain + more total nodes across generations). It is NOT a hard cap on the + total number of cells per clone. + n_sample: + Number of cells to sample as observed leaves. Records returned + per clone are ≤ ``n_sample`` because identical genotypes are + collapsed (duplicates are counted in ``lineage_abundance``). + rate: + Per-base SHM rate for within-lineage mutations. + lambda_base: + Poisson mean for offspring count at affinity 0. + selection_strength: + Selection pressure; ``0.0`` = neutral drift (fitness is 1.0 for + every cell). This disables selection but does not force + ``lineage_affinity`` to 0 when a target sequence is supplied. + Set ``> 0`` to enable affinity maturation; combine with + ``target_aa`` for a fixed sequence target. + beta: + Scaling factor for the affinity term in ``exp(−beta·distance)``. + target_aa: + Amino-acid sequence of the full receptor used to compute + per-cell affinity via a BLOSUM62-weighted distance (compared + position-wise against the cell's translated receptor; only + the overlapping prefix is scored when lengths differ). Must be + a non-empty string of standard amino-acid letters + (``ACDEFGHIKLMNPQRSTVWY``). When ``None``, an auto target is + generated from the founder by applying ``mature_substitutions`` + random residue changes whenever selection is enabled. In fully + neutral mode (``selection_strength=0`` and ``target_aa=None``), + no affinity model is built and ``lineage_affinity`` is 0. + mature_substitutions: + Number of amino-acid substitutions used to build the auto + target (when ``target_aa`` is ``None``). + s5f_model: + Bundled S5F kernel name for within-lineage mutation context + (``"hh_s5f"``, ``"hkl_s5f"``, …). + allow_extinction: + Sampling draws from the LIVING final-generation population, so a + founder that draws 0 offspring goes extinct and yields zero + observed cells/records. With ``allow_extinction=False`` (default) + each requested clone is conditioned on survival: an extinct family + is retried with a fresh deterministic sub-seed (up to a bounded + number of attempts) so you reliably get ``n_clones`` families. With + ``allow_extinction=True`` extinction is accepted and the extinct + clone is skipped, producing fewer families than ``n_clones``. + + **BCR-only guard:** ``clonal_lineage`` applies S5F somatic + hypermutation, which is a B-cell process. Calling it on a TCR-configured + experiment raises ``ValueError`` (immunoglobulin / BCR loci only). TCR + clone-size primitives exist in the engine but are not yet exposed as a + DSL workflow. + """ + import math + import warnings + + # --- BCR-only guard (mirror mutate()'s TCR rejection) --- + # clonal_lineage applies S5F somatic hypermutation, a B-cell process, + # so it must reject TCR loci. ``_is_tcr_refdata`` inspects the first V + # allele name prefix (TR* => TCR, IG* => BCR) on the already-bound + # refdata, exactly as mutate() does. Firing here, at call time and + # before compile(), guarantees the clear BCR-only message instead of a + # downstream cartridge / compile error. + if self._is_tcr_refdata(): + locus = self._refdata.v_allele(0).name if self._refdata.v_pool_size() else "?" + raise ValueError( + "clonal_lineage models B-cell somatic hypermutation and " + "supports immunoglobulin (BCR) loci only; the locus " + f"'{locus}' is a TCR locus. (TCR clone-size simulation is not " + "yet exposed in the DSL.)" + ) + # --- allow_extinction --- + if not isinstance(allow_extinction, bool): + raise ValueError( + f"allow_extinction must be a bool, got {allow_extinction!r}" + ) + + # --- n_clones --- + if isinstance(n_clones, bool) or not isinstance(n_clones, int) or n_clones < 1: + raise ValueError(f"n_clones must be a positive int, got {n_clones!r}") + # --- max_generations --- + if ( + isinstance(max_generations, bool) + or not isinstance(max_generations, int) + or max_generations < 1 + or max_generations > 1000 + ): + raise ValueError( + f"max_generations must be a positive int <= 1000, got {max_generations!r}" + ) + # --- n_max --- + if isinstance(n_max, bool) or not isinstance(n_max, int) or n_max < 1: + raise ValueError(f"n_max must be a positive int, got {n_max!r}") + # --- n_sample --- + if isinstance(n_sample, bool) or not isinstance(n_sample, int) or n_sample < 1: + raise ValueError(f"n_sample must be a positive int, got {n_sample!r}") + # --- rate --- + if not isinstance(rate, (int, float)) or not math.isfinite(rate) or rate < 0 or rate > 1: + raise ValueError(f"rate must be a float in [0, 1], got {rate!r}") + # --- lambda_base --- + if not isinstance(lambda_base, (int, float)) or not math.isfinite(lambda_base) or lambda_base < 0: + raise ValueError(f"lambda_base must be a finite non-negative float, got {lambda_base!r}") + # --- beta --- + if not isinstance(beta, (int, float)) or not math.isfinite(beta) or beta < 0: + raise ValueError(f"beta must be a finite non-negative float, got {beta!r}") + # --- selection_strength --- + if not isinstance(selection_strength, (int, float)) or not math.isfinite(selection_strength) or selection_strength < 0: + raise ValueError( + f"selection_strength must be a finite non-negative float, got {selection_strength!r}" + ) + # --- mature_substitutions --- + if ( + isinstance(mature_substitutions, bool) + or not isinstance(mature_substitutions, int) + or mature_substitutions < 0 + ): + raise ValueError( + f"mature_substitutions must be a non-negative int, got {mature_substitutions!r}" + ) + # --- target_aa --- + _VALID_AA = set("ACDEFGHIKLMNPQRSTVWY") + if target_aa is not None: + if not isinstance(target_aa, str) or len(target_aa) == 0: + raise ValueError( + "target_aa must be a non-empty amino-acid string " + "(letters from ACDEFGHIKLMNPQRSTVWY)" + ) + target_aa = target_aa.upper() + invalid = set(target_aa) - _VALID_AA + if invalid: + raise ValueError( + f"target_aa contains invalid characters {sorted(invalid)!r}; " + "only standard amino-acid letters (ACDEFGHIKLMNPQRSTVWY) are allowed" + ) + if len(target_aa) < 30: + warnings.warn( + f"target_aa has length {len(target_aa)}, which is shorter than a typical " + "receptor sequence (~300+ aa). If this is an epitope sequence rather than " + "the full receptor, affinity scoring will be based only on the overlapping " + "prefix — consider supplying the full translated receptor instead.", + UserWarning, + stacklevel=2, + ) + # --- s5f_model (validate at call time, not at run time) --- + from GenAIRR._s5f_loader import _BUILTIN_S5F_MODELS + _s5f_key = s5f_model.lower().strip() + if _s5f_key not in _BUILTIN_S5F_MODELS: + avail = ", ".join(f'"{k}"' for k in sorted(_BUILTIN_S5F_MODELS)) + raise ValueError( + f"Unknown s5f_model {s5f_model!r}. Available: {avail}" + ) + # --- reject duplicate fork steps --- + for s in self._steps: + if isinstance(s, (_ClonalForkStep, _RepertoireForkStep, _LineageForkStep)): + raise ValueError( + "clonal_lineage() / expand_clones() / clonal_repertoire() " + "can only be called once per pipeline" + ) + # --- descendant-phase guard (same as expand_clones) --- + for step in self._steps: + offending_method = _descendant_phase_step_classifier(step) + if offending_method is None: + continue + raise ValueError( + f"{offending_method} must be called after " + f"clonal_lineage(); it is descendant-specific and must be sampled " + f"independently for each clone member. Move " + f"{offending_method}(...) after clonal_lineage(...)." + ) + self._steps.append( + _LineageForkStep( + n_clones=n_clones, + max_generations=max_generations, + n_max=n_max, + n_sample=n_sample, + rate=rate, + lambda_base=lambda_base, + selection_strength=selection_strength, + beta=beta, + target_aa=target_aa, + mature_substitutions=mature_substitutions, + s5f_model=s5f_model, + allow_extinction=allow_extinction, + ) + ) + return self diff --git a/src/GenAIRR/_experiment/compile.py b/src/GenAIRR/_experiment/compile.py new file mode 100644 index 0000000..cd5c5c5 --- /dev/null +++ b/src/GenAIRR/_experiment/compile.py @@ -0,0 +1,394 @@ +from __future__ import annotations + +import warnings +from typing import Optional + +from GenAIRR import _engine + +from .._compiled import ( + CompiledClonalExperiment, + CompiledExperiment, + CompiledLineageExperiment, + CompiledRepertoireExperiment, +) +from .._lowering import ( + _extract_invert_d_prob, + _extract_paired_end_step, + _extract_receptor_revision_prob, + _lower_paired_end, + _lower_recombine, + lower_step, +) +from .._pipeline_ir import ( + _ClonalForkStep, + _CorruptStep, + _LineageForkStep, + _MutateStep, + _PairedEndStep, + _RecombineStep, + _RepertoireForkStep, +) +from .._refdata_resolver import dataconfig_to_refdata + + +class _CompileMixin: + __slots__ = () + + def compile(self, *, allow_curatable_refdata: Optional[bool] = None): + """Compile the recorded steps into a reusable + :class:`CompiledExperiment` (or :class:`CompiledClonalExperiment` + when the pipeline contains a :meth:`expand_clones` + fork). + + Idempotent: calling ``compile()`` twice produces two distinct + compiled instances with structurally-equal simulators. + + Constraints declared via :meth:`productive_only` (or future + bundle methods) are baked into the compiled simulator at this + step; they're not runtime knobs. To run without constraints, + omit the constraint methods from the chain. + + ``allow_curatable_refdata`` selects the refdata validation + mode. ``None`` (default) inherits the instance flag set by + :meth:`allow_curatable_refdata`; an explicit ``True`` / + ``False`` overrides per-call. ``False`` runs the gate in + strict mode — every issue rejects compile with a + :class:`ValueError`. ``True`` runs the lenient mode — Fatal + issues (empty pool, duplicates, invalid byte, anchor out of + bounds) still reject, but Curatable issues (pseudogene-shape + anchor anomalies) pass. + """ + if allow_curatable_refdata is None: + allow_curatable_refdata = self._allow_curatable_refdata + + # Receptor revision with a phased genotype is now haplotype-aware: + # the lowering builds a genotype-aware ReceptorRevisionPass that + # restricts the replacement V to carried alleles on the drawn + # rearrangement chromosome (see _lower_recombine). No rejection here. + + # Genotype provenance (subject_id / haplotype / result.genotypes) + # is only threaded through the plain compiled path, not the + # clonal/lineage/repertoire forked classes. Reject the + # combination rather than silently dropping provenance (review + # #9); genotype + clonal cohorts are a planned follow-on. + if self._genotype is not None and self._has_clonal_fork(): + raise ValueError( + "with_genotype() is not supported together with expand_clones() / " + "clonal_lineage() / clonal_repertoire() in this release" + ) + from dataclasses import replace as _replace + + # On raw RefDataConfig with default-on trim, warn at compile + # time if the user didn't explicitly call .trim() to either + # disable or replace the no-op default. By compile() time we + # know the final shape of the recombine step. + if self._dataconfig is None: + for step in self._steps: + if not isinstance(step, _RecombineStep): + continue + has_trim_data = any( + (step.trim_v_3, step.trim_d_5, step.trim_d_3, step.trim_j_5) + ) + if has_trim_data or step.trim_overridden: + # Either trim is set up (DataConfig path or + # .trim(v_3=...)) or the user explicitly called + # .trim() — both silence the warning. + continue + warnings.warn( + "Experiment bound to a raw RefDataConfig has no " + "trim distributions; exonuclease trim is a no-op. " + "Call .trim(enabled=False) to silence this warning, " + "or supply custom distributions via " + ".trim(v_3=..., d_5=..., ...).", + UserWarning, + stacklevel=2, + ) + break # one warning per compile() call is enough + + contracts = self._build_contracts() + any_lock = any(self._locks[seg] is not None for seg in ("V", "D", "J")) + + # if a `_ClonalForkStep` is present, split the step list + # at it and compile two simulators (pre-fork = per-clone, + # post-fork = per-descendant). + fork_idx = next( + (i for i, s in enumerate(self._steps) if isinstance(s, _ClonalForkStep)), + None, + ) + if fork_idx is not None: + fork_step: _ClonalForkStep = self._steps[fork_idx] + pre_steps = self._steps[:fork_idx] + post_steps = self._steps[fork_idx + 1 :] + pre_simulator = self._build_simulator( + pre_steps, + contracts, + any_lock, + replace_fn=_replace, + allow_curatable_refdata=allow_curatable_refdata, + ) + # The post-fork plan inherits the parent's V/D/J/NP + # backbone (recombination already happened in the + # pre-fork half). The recombination-time precondition + # facts (np.np1.length residues, anchor trim supports) + # aren't produced here. v2 used to drop the contract + # bundle entirely on this side to avoid the compile-time + # precondition failure — but doing so left every + # post-fork mutation / corruption pass unfiltered, which + # was the dominant cause of non-productive output under + # `productive_only()`. We now pass the same bundle + # through; the engine analyzer (in + # `compiled/analyze.rs::validate_contract_preconditions`) + # skips the productive-frame check when no recombination + # facts are present in the plan, and the runtime + # admit_with_context / admits_post_event paths still + # enforce no-stop-codon-in-junction and anchor-preserved + # for every substitution and indel in the post-fork pipeline. + post_simulator = self._build_simulator( + post_steps, + contracts, + any_lock=False, + replace_fn=_replace, + allow_curatable_refdata=allow_curatable_refdata, + ) + return CompiledClonalExperiment( + pre_simulator, + post_simulator, + fork_step, + self._refdata, + pre_steps=tuple(pre_steps), + post_steps=tuple(post_steps), + dataconfig=self._dataconfig, + metadata=self._metadata, + ) + + # if a `_RepertoireForkStep` is present, split the step list + # at it and compile the pre-fork (per-clone) simulator plus an + # optional post-fork (per-read) simulator, mirroring the + # `_ClonalForkStep` branch. Per-clone sizes are drawn at run + # time from the heavy-tailed distribution; identical reads are + # collapsed into `duplicate_count`-carrying records. + repertoire_idx = next( + ( + i + for i, s in enumerate(self._steps) + if isinstance(s, _RepertoireForkStep) + ), + None, + ) + if repertoire_idx is not None: + repertoire_step: _RepertoireForkStep = self._steps[repertoire_idx] + pre_steps = self._steps[:repertoire_idx] + post_steps = self._steps[repertoire_idx + 1 :] + pre_simulator = self._build_simulator( + pre_steps, + contracts, + any_lock, + replace_fn=_replace, + allow_curatable_refdata=allow_curatable_refdata, + ) + # post_steps may be empty — that's the pure-copy case + # (each clone collapses to one record with + # duplicate_count = size). When present, the post-fork + # plan inherits the parent's V/D/J/NP backbone, so + # any_lock=False and no recombination facts on this side + # (same as the clonal branch). + post_simulator = None + if post_steps: + post_simulator = self._build_simulator( + post_steps, + contracts, + any_lock=False, + replace_fn=_replace, + allow_curatable_refdata=allow_curatable_refdata, + ) + return CompiledRepertoireExperiment( + pre_simulator, + post_simulator, + repertoire_step, + self._refdata, + dataconfig=self._dataconfig, + metadata=self._metadata, + ) + + # if a `_LineageForkStep` is present, compile a + # CompiledLineageExperiment: pre-fork steps (recombine) become + # the founder simulator; post-fork steps must be empty (the + # lineage engine handles mutation internally). + lineage_idx = next( + (i for i, s in enumerate(self._steps) if isinstance(s, _LineageForkStep)), + None, + ) + if lineage_idx is not None: + lineage_step: _LineageForkStep = self._steps[lineage_idx] + pre_steps = self._steps[:lineage_idx] + post_steps = self._steps[lineage_idx + 1:] + # Steps after clonal_lineage() are per-observed-cell + # library-prep / sequencing artefact passes (the same + # post-fork set expand_clones() allows). SHM is internal + # to the lineage engine, so .mutate() is rejected; the + # paired-end read layout is not yet wired through the + # per-cell corruption merge, so reject it for now too. + for s in post_steps: + if isinstance(s, _MutateStep): + raise ValueError( + "SHM is internal to clonal_lineage; do not add " + ".mutate() after it. Set the within-lineage SHM rate " + "via clonal_lineage(rate=...)." + ) + if isinstance(s, _PairedEndStep): + raise ValueError( + "paired_end not yet supported with clonal_lineage; " + "apply per-read library-prep passes (sequencing_errors, " + "pcr_amplify, polymerase_indels, end_loss_*, " + "ambiguous_base_calls, random_strand_orientation) instead." + ) + if not isinstance(s, _CorruptStep): + raise ValueError( + "Only per-read library-prep / sequencing artefact passes " + "may follow clonal_lineage(); got " + f"{type(s).__name__}." + ) + pre_simulator = self._build_simulator( + pre_steps, + contracts, + any_lock, + replace_fn=_replace, + allow_curatable_refdata=allow_curatable_refdata, + ) + # Build the per-cell corruption simulator from the + # post-fork steps, mirroring the clonal branch's + # post_simulator (no recombination facts on this side, so + # any_lock=False and the analyzer skips the productive + # precondition check). + post_simulator = None + if post_steps: + post_simulator = self._build_simulator( + post_steps, + contracts, + any_lock=False, + replace_fn=_replace, + allow_curatable_refdata=allow_curatable_refdata, + ) + return CompiledLineageExperiment( + pre_simulator, + lineage_step, + self._refdata, + post_simulator=post_simulator, + post_steps=tuple(post_steps), + dataconfig=self._dataconfig, + metadata=self._metadata, + ) + + # When the attached genotype defines novel/private alleles, compile + # against an *effective* reference = base catalogue + injected novel + # alleles, so they become real pool entries the engine samples, + # assembles, and reports like any allele. No genotype, or a genotype + # without novel alleles, uses the base refdata unchanged. + effective_refdata = self._refdata + if self._genotype is not None and self._genotype.has_novel(): + if self._dataconfig is None: + raise ValueError( + "genotype with novel alleles requires a DataConfig-backed " + "experiment (Experiment.on(dataconfig), not a raw RefDataConfig)" + ) + effective_refdata = dataconfig_to_refdata( + self._genotype.effective_dataconfig() + ) + + simulator = self._build_simulator( + self._steps, + contracts, + any_lock, + replace_fn=_replace, + allow_curatable_refdata=allow_curatable_refdata, + refdata=effective_refdata, + ) + return CompiledExperiment( + simulator, + effective_refdata, + steps=tuple(self._steps), + dataconfig=self._dataconfig, + metadata=self._metadata, + genotype=self._genotype, + ) + + def _build_simulator( + self, + steps, + contracts, + any_lock: bool, + *, + replace_fn, + allow_curatable_refdata: bool = False, + refdata=None, + ): + """Compile a list of steps into a `GenAIRR._engine.CompiledSimulator`. + Lifted out of `compile()` so the clonal-fork branch can build + two simulators from sub-step-lists with a shared body. + + ``refdata`` overrides ``self._refdata`` — used when a genotype with + novel alleles compiles against an *effective* reference (base + + injected private alleles).""" + refdata = refdata if refdata is not None else self._refdata + plan = _engine.PassPlan() + # Pull the (at-most-one) `_InvertDStep` out of the step + # sequence and thread its probability into the recombine + # lowering directly. Inlining the InvertDPass push between + # `push_generate_np("NP1", ...)` and `push_assemble("D")` + # (see `_lower_recombine`) keeps the canonical V-NP1-D-NP2-J + # pool layout intact; a separate `_InvertDStep` lowering + # combined with `Schedule::before(invert_d, assemble.d)` + # would re-promote the pass past `assemble.j`, swapping D + # and J in the pool. See `_extract_invert_d_prob` for the + # full rationale. + invert_d_prob, steps = _extract_invert_d_prob(steps) + # Same inline-into-recombine-lowering pattern as + # _extract_invert_d_prob: pulling the receptor-revision step + # out here lets the lowering push it at the exact slot + # between `assemble.j` and any subsequent mutate/corrupt + # passes. A standalone lower path would either need a new + # schedule edge or place the pass at the end of the plan + # (after corruption), both of which break the design doc §2 + # ordering. + receptor_revision_prob, receptor_revision_same_haplotype, steps = ( + _extract_receptor_revision_prob(steps) + ) + # Pull out the (at-most-one) paired-end step too. It must + # land at the END of the plan, not inline with recombine — + # see `_extract_paired_end_step` for the rationale on + # placement vs. trace order. + paired_end_step, steps = _extract_paired_end_step(steps) + for step in steps: + # Inject any allele-locks set via ``.restrict_alleles(...)`` into the + # recombine step at compile time. Other step types ignore + # locks. + if any_lock and isinstance(step, _RecombineStep): + step = replace_fn( + step, + locks_v=self._locks["V"], + locks_d=self._locks["D"], + locks_j=self._locks["J"], + ) + if isinstance(step, _RecombineStep): + _lower_recombine( + step, + plan, + refdata, + invert_d_prob=invert_d_prob, + receptor_revision_prob=receptor_revision_prob, + receptor_revision_same_haplotype=receptor_revision_same_haplotype, + genotype=self._genotype, + ) + else: + lower_step(step, plan, refdata) + # Paired-end is sequencing-stage / readout-stage: lower + # it AFTER every biology + corruption pass so the trace + # records land last. See `_extract_paired_end_step` for + # the full rationale. + if paired_end_step is not None: + _lower_paired_end(paired_end_step, plan) + return plan.compile( + refdata=refdata, + respect=contracts, + allow_curatable_refdata=allow_curatable_refdata, + ) diff --git a/src/GenAIRR/_experiment/constraints.py b/src/GenAIRR/_experiment/constraints.py new file mode 100644 index 0000000..a759dc5 --- /dev/null +++ b/src/GenAIRR/_experiment/constraints.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from typing import Any + + +class _ConstraintsMixin: + __slots__ = () + + def with_metadata(self, **fields: Any) -> "Experiment": + """Attach sample-level metadata to every AIRR record. + + Standard AIRR Repertoire fields like ``sample_id``, + ``donor``, ``repertoire_id``, and ``cell_id`` are commonly + used by multi-sample analysis tools — Change-O, scirpy, + immcantation pipelines all expect them populated. Custom + keys are also accepted and pass through unchanged. + + Subsequent ``with_metadata()`` calls **merge** with the + prior set (per-key replacement). Pass ``key=None`` to clear + a single key. Values are stored as-is and serialised via + the standard CSV/TSV writers; non-string values are + converted with ``str()`` on output. + + Example:: + + (Experiment.on("human_igh") + .with_metadata(sample_id="P1", donor="D001", + repertoire_id="R-001-IGH") + .recombine().mutate(count=8)) + """ + for key, value in fields.items(): + if not isinstance(key, str): + raise TypeError( + f"with_metadata: keys must be strings, got " + f"{type(key).__name__}" + ) + if value is None: + self._metadata.pop(key, None) + else: + self._metadata[key] = value + return self + + def productive_only(self) -> "Experiment": + """Require every emitted record to be a productive sequence. + + Attaches the canonical productive-sequence contract bundle to + the experiment: junction frame in-register, no stop codons in + the junction, and V/J anchor amino acids preserved. The bundle + is enforced during recombination, mutation, and corruption + passes by narrowing each pass's action support before + sampling. If a constrained support is empty, permissive + execution uses the pass's explicit no-op / sentinel behavior; + use ``strict=True`` at run time to raise instead. + + **Failure surfaces** (see + ``docs/productive_failure_mode_audit.md`` for the full matrix): + + - *Compile-time precondition* — when a sampling distribution + is statically impossible under the bundle (e.g. every NP1 + length violates frame), :meth:`compile` raises + ``ValueError`` regardless of the ``strict`` flag. + - *Runtime fresh strict* (``run(..., strict=True)``) — when + dynamic state makes a sampler's admissible support empty, + raises :class:`GenAIRR._engine.StrictSamplingError` with + structured ``(pass_name, address, reason)`` args. + - *Runtime fresh permissive* (``strict=False``, default) — + the pass records its declared sentinel (indel ``site=-1``, + NP length ``0``, NP base ``N``, trim ``0``) or skips the + slot; the record continues. + - *Trace replay* (``replay_from_trace_file``) — consumes + recorded values verbatim, does not re-evaluate + admissibility. A permissive-sentinel trace replays cleanly + even with ``strict=True``. + + **Order-independent.** This method is a constraint declaration, + not a pipeline step — it can be called anywhere in the chain + and the result is identical. Convention is to place it last + (right before ``run_records()``) so the constraint reads as a + post-hoc requirement on the emitted records. + + TCR refdata accepts the call but raises ``ValueError`` at + :meth:`compile` time because TCRs don't have somatic + hypermutation and the productive bundle's anchor checks + assume BCR semantics. Catch this early at the builder if it + matters to you. + + Example:: + + result = ( + Experiment.on("human_igh") + .recombine() + .mutate(count=(5, 15)) + .productive_only() + .run_records(seed=42) + ) + """ + if "productive" not in self._contracts: + self._contracts.append("productive") + return self diff --git a/src/GenAIRR/_experiment/corruption.py b/src/GenAIRR/_experiment/corruption.py new file mode 100644 index 0000000..3ca5514 --- /dev/null +++ b/src/GenAIRR/_experiment/corruption.py @@ -0,0 +1,359 @@ +from __future__ import annotations + +from typing import Iterable, Optional, Tuple, Union + +from .._normalize import _normalize_count +from .._pipeline_ir import ( + _CORRUPT_KIND_3PRIME_LOSS, + _CORRUPT_KIND_5PRIME_LOSS, + _CORRUPT_KIND_CONTAMINANT, + _CORRUPT_KIND_INDEL, + _CORRUPT_KIND_NS, + _CORRUPT_KIND_PCR, + _CORRUPT_KIND_QUALITY, + _CORRUPT_KIND_REV_COMP, + _CorruptStep, +) + + +class _CorruptionMixin: + __slots__ = () + + def pcr_amplify( + self, + *, + count: Optional[Union[int, Tuple[int, int], Iterable[Tuple[int, float]]]] = None, + rate: Optional[float] = None, + ) -> "Experiment": + """Append a PCR-error step modelling substitution errors + introduced during PCR amplification. + + **Specify intensity with exactly one of ``rate`` or ``count``.** + + ``rate`` is the per-base PCR error probability (e.g. + ``1e-4`` per base per cycle, depending on polymerase + fidelity). At execute time the engine draws + ``count ~ Poisson(rate × pool_len)`` against each record's + current sequence length — matching how PCR error is + universally reported in the literature. This is the + canonical biology-default form. + + ``count`` is the legacy explicit count distribution + (fixed int, ``(low, high)`` range, or empirical + ``(count, weight)`` list). Useful for benchmark scripts. + + Each error samples a uniform position in the assembled pool + and replaces the base with a uniform A/C/G/T draw. Trace + addresses: ``corrupt.pcr.{count, error_site[i], error_base[i]}``. + + Passing both ``count`` and ``rate`` raises ``ValueError``. + Passing neither raises ``ValueError``. + """ + return self._append_count_or_rate_step( + kind=_CORRUPT_KIND_PCR, + label="pcr_amplify", + count=count, + rate=rate, + ) + + def sequencing_errors( + self, + *, + count: Optional[Union[int, Tuple[int, int], Iterable[Tuple[int, float]]]] = None, + rate: Optional[float] = None, + ) -> "Experiment": + """Append a sequencing-quality-error step modelling base-call + errors during sequencing readout. + + **Specify intensity with exactly one of ``rate`` or ``count``.** + + ``rate`` is the per-base sequencing error probability (e.g. + ``1e-3`` for a Q30 base, ``1e-2`` for Q20). Drawn as + ``count ~ Poisson(rate × pool_len)`` per record — the + canonical Phred-quality framing in immunoseq. + + ``count`` is the legacy explicit count distribution. + + Same shape as :meth:`pcr_amplify` but each substitution + writes the destination base in **lowercase** to mark the + position as corrupted (the sequencing-error convention). + """ + return self._append_count_or_rate_step( + kind=_CORRUPT_KIND_QUALITY, + label="sequencing_errors", + count=count, + rate=rate, + ) + + def _append_count_or_rate_step( + self, + *, + kind: str, + label: str, + count: Optional[Union[int, Tuple[int, int], Iterable[Tuple[int, float]]]], + rate: Optional[float], + ) -> "Experiment": + """Shared validator + step appender for the corruption methods + that accept both ``count`` and ``rate`` (PCR errors, sequencing + errors). See :meth:`mutate` for the same shape on the + mutation side.""" + if count is not None and rate is not None: + raise ValueError( + f"{label}(): pass exactly one of `rate` or `count`, not both. " + f"`rate` is the canonical biology default (e.g. rate=1e-4 " + f"for per-base PCR error); `count` is the explicit per-record " + f"count for benchmark / deterministic-count workflows." + ) + if count is None and rate is None: + raise ValueError( + f"{label}(): pass exactly one of `rate` or `count`." + ) + if rate is not None: + if not isinstance(rate, (int, float)) or isinstance(rate, bool): + raise TypeError( + f"rate must be a finite float in [0.0, 1.0], got " + f"{type(rate).__name__}" + ) + if not (0.0 <= float(rate) <= 1.0): + raise ValueError( + f"rate must be in [0.0, 1.0] (got {rate!r}); rate is a " + f"per-base error probability, not an absolute count." + ) + self._steps.append( + _CorruptStep( + kind=kind, + rate=float(rate), + ) + ) + return self + pairs = _normalize_count(count) + self._steps.append( + _CorruptStep( + kind=kind, + count_pairs=pairs, + ) + ) + return self + + def polymerase_indels( + self, + *, + count: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], + insertion_prob: float = 0.5, + ) -> "Experiment": + """Append a polymerase-slippage indel step (PCR-stage artifact). + + Models the small insertions / deletions that arise from + polymerase slippage during PCR amplification — single-base + indels at low rate, typically <1 % per read at standard + polymerase fidelity. Larger structural indels are rare and + not modelled here. + + ``count`` is the per-simulation distribution over the total + number of indel events. Each event independently chooses + insertion (probability ``insertion_prob``) vs. deletion. + Insertions sample a uniform A/C/G/T base. Indels change + pool length and shift downstream region ranges accordingly. + + Raises ``ValueError`` if ``insertion_prob`` is outside + ``[0.0, 1.0]`` or non-finite. + """ + if ( + insertion_prob != insertion_prob # NaN check + or not (0.0 <= insertion_prob <= 1.0) + ): + raise ValueError( + f"insertion_prob must be a finite number in [0.0, 1.0], " + f"got {insertion_prob}" + ) + pairs = _normalize_count(count) + self._steps.append( + _CorruptStep( + kind=_CORRUPT_KIND_INDEL, + count_pairs=pairs, + insertion_prob=float(insertion_prob), + apply_prob=0.0, + ) + ) + return self + + def end_loss_5prime( + self, + *, + length: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], + ) -> "Experiment": + """Append a 5'-end loss corruption step. + + Drops bases from the start of the assembled sequence. The + underlying engine pass is `EndLossPass` — this models + **observation-stage** read-end / primer-region loss (the + AIRR record exposes it as ``end_loss_5_length``, the trace + address is ``corrupt.end_loss.5``). It is distinct from + recombination-stage trimming (`v_trim_5` etc.), which writes + allele-instance metadata rather than deleting pool bytes. + + ``length`` accepts the same shapes as ``count`` on other + corrupt ops: + + - ``length=10`` — strip exactly 10 bases. + - ``length=(0, 20)`` — strip a uniform integer in ``[0, 20]``. + - ``length=[(5, 1.0), (10, 1.0), (15, 1.0)]`` — empirical + (length, weight) distribution. + + The actual loss is clamped to the pool length (so a sample + larger than the sequence drops the whole pool). The loss is + permanent for downstream passes — subsequent corruption + operates on the shorter pool. + + :meth:`primer_trim_5prime` is a backwards-compatible alias. + """ + pairs = _normalize_count(length) + self._steps.append( + _CorruptStep( + kind=_CORRUPT_KIND_5PRIME_LOSS, + count_pairs=pairs, + insertion_prob=0.0, + apply_prob=0.0, + ) + ) + return self + + def end_loss_3prime( + self, + *, + length: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], + ) -> "Experiment": + """Append a 3'-end loss corruption step. + + Drops bases from the end of the assembled sequence. Same + observation-stage semantics as :meth:`end_loss_5prime` + (trace address ``corrupt.end_loss.3``, AIRR field + ``end_loss_3_length``). Same ``length`` shapes. + + :meth:`primer_trim_3prime` is a backwards-compatible alias. + """ + pairs = _normalize_count(length) + self._steps.append( + _CorruptStep( + kind=_CORRUPT_KIND_3PRIME_LOSS, + count_pairs=pairs, + insertion_prob=0.0, + apply_prob=0.0, + ) + ) + return self + + def primer_trim_5prime( + self, + *, + length: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], + ) -> "Experiment": + """Backwards-compatible alias for :meth:`end_loss_5prime`. + + The DSL name `primer_trim_5prime` predates the engine's + cleaner vocabulary — the same step is now exposed as + :meth:`end_loss_5prime`. Behaviour is identical (same trace + address ``corrupt.end_loss.5``, same AIRR field + ``end_loss_5_length``, byte-identical records for the same + seed). New code should prefer the `end_loss_*` form; the + alias stays so existing scripts keep working. + """ + return self.end_loss_5prime(length=length) + + def primer_trim_3prime( + self, + *, + length: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], + ) -> "Experiment": + """Backwards-compatible alias for :meth:`end_loss_3prime`. + + See :meth:`primer_trim_5prime` for the rationale. + """ + return self.end_loss_3prime(length=length) + + def ambiguous_base_calls( + self, + *, + count: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], + ) -> "Experiment": + """Append an N-substitution corruption step. + + With the per-simulation count drawn from ``count``, replace + that many uniform-random pool positions with the ambiguous + base ``N``. Models the low-quality positions that real + sequencers emit when the base caller cannot commit. + + ``count`` accepts the same shapes as other corruption ops: + - ``count=10`` — fixed. + - ``count=(0, 20)`` — uniform integer range. + - ``count=[(5, 1.0), (10, 1.0)]`` — empirical (count, weight). + + Sampled sites are with-replacement, so ``count`` is the + upper bound on resulting `N` count (collisions reduce it). + """ + pairs = _normalize_count(count) + self._steps.append( + _CorruptStep( + kind=_CORRUPT_KIND_NS, + count_pairs=pairs, + insertion_prob=0.0, + apply_prob=0.0, + ) + ) + return self + + def random_strand_orientation(self, *, prob: float = 0.5) -> "Experiment": + """Append a reverse-complement corruption step. + + With probability ``prob`` the AIRR record-builder reverse- + complements the ``sequence``, ``np1``, ``np2``, and + ``junction`` fields and flips the corresponding pool-position + coords (``*_sequence_start/end``, ``junction_start/end``). + Alignment / germline / CIGAR / identity fields stay in + forward orientation per the AIRR Rearrangement spec. + ``prob=0.0`` is a no-op (coin flip recorded but never fires); + ``prob=1.0`` always flips. The biological model is the ~50% + of antisense reads in real immune-seq libraries. + + Raises ``ValueError`` if ``prob`` is outside ``[0.0, 1.0]`` + or non-finite. + """ + if prob != prob or not (0.0 <= prob <= 1.0): + raise ValueError( + f"prob must be a finite number in [0.0, 1.0], got {prob}" + ) + self._steps.append( + _CorruptStep( + kind=_CORRUPT_KIND_REV_COMP, + count_pairs=(), + insertion_prob=0.0, + apply_prob=float(prob), + ) + ) + return self + + def contaminate(self, *, prob: float) -> "Experiment": + """Append a contaminant-replacement corruption step. + + With probability ``prob`` the entire assembled pool is + overwritten with uniform A/C/G/T bases. Models primer + dimers, bacterial DNA, or any non-receptor sequence in the + library. ``prob=0.0`` is a no-op (coin flip recorded but + never fires); ``prob=1.0`` always contaminates. + + Raises ``ValueError`` if ``prob`` is outside ``[0.0, 1.0]`` + or non-finite. + """ + if prob != prob or not (0.0 <= prob <= 1.0): + raise ValueError( + f"prob must be a finite number in [0.0, 1.0], got {prob}" + ) + self._steps.append( + _CorruptStep( + kind=_CORRUPT_KIND_CONTAMINANT, + count_pairs=(), + insertion_prob=0.0, + apply_prob=float(prob), + ) + ) + return self diff --git a/src/GenAIRR/_experiment/genotype_alleles.py b/src/GenAIRR/_experiment/genotype_alleles.py new file mode 100644 index 0000000..170895e --- /dev/null +++ b/src/GenAIRR/_experiment/genotype_alleles.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +from typing import Dict, Iterable, List, Optional, Tuple, Union + +from ._common import _UNSET + + +class _GenotypeAllelesMixin: + __slots__ = () + + def with_genotype(self, genotype) -> "Experiment": + """Attach a single-subject diploid genotype. + + With a genotype attached, V(D)J recombination becomes + haplotype-phased: V, D and J of each rearrangement are drawn from + a single chromosome, honouring the genotype's allele + presence/absence, zygosity, and copy-number/deletion. With no + genotype, the flat (uniform / usage-weighted) path runs unchanged. + + Mutually exclusive with :meth:`restrict_alleles` and the + ``recombine(*_allele_weights=...)`` kwargs — the genotype owns + allele presence and within-gene expression. + + Raises ``ValueError`` if the genotype was built against a + different cartridge (content-hash mismatch), or if allele locks / + explicit allele weights are already set. + """ + live_hash = self._refdata.content_hash() + if genotype._source_hash != live_hash: + raise ValueError( + "genotype was built against a different cartridge (content hash " + f"{genotype._source_hash!r} != experiment {live_hash!r})" + ) + if any(v is not None for v in self._locks.values()): + raise ValueError( + "with_genotype() and restrict_alleles() are mutually exclusive" + ) + if self._user_allele_weights_set: + raise ValueError( + "with_genotype() and recombine(*_allele_weights=...) are mutually " + "exclusive: the genotype owns allele expression" + ) + # Snapshot the (mutable) builder so later edits to ``genotype`` + # cannot desync the compiled engine genotype from + # ``result.genotypes`` (review #8). + self._genotype = genotype._snapshot() + return self + + def restrict_alleles( + self, + *, + v: "_LockInput" = _UNSET, + d: "_LockInput" = _UNSET, + j: "_LockInput" = _UNSET, + ) -> "Experiment": + """Narrow allele sampling to a named subset, per segment. + + Sampling stays uniform — this restricts the *support* of the + sampling distribution to the listed allele names, not pins a + single allele. If you supply one name, the result is + effectively deterministic (uniform over one element); for any + list of N names, the recombination pass samples uniformly + over those N alleles. + + Each kwarg accepts: + - a single allele name string (e.g. ``"IGHV1-2*02"``), + - a list / tuple of allele names (sample uniformly among them), + - ``None`` to clear a previously-set restriction for that segment, + - omitted (default) — the existing restriction for that segment + is left unchanged. + + The restrictions are applied to the next ``recombine()`` step's + ``push_sample_allele`` calls at compile time. Calling + ``restrict_alleles()`` more than once overlays the new values + onto the previous restrictions (per-segment), so + ``.restrict_alleles(v="A").restrict_alleles(d="B")`` restricts + both V and D. + + Raises: + - ``ValueError`` if an allele name doesn't exist in the + configured refdata, or if a restriction is set for ``D`` on + a VJ chain. + - ``TypeError`` if an unexpected input shape is passed. + """ + if self._genotype is not None: + raise ValueError( + "restrict_alleles() and with_genotype() are mutually exclusive: " + "a genotype already owns allele presence and within-gene expression" + ) + for segment, value in (("V", v), ("D", d), ("J", j)): + if value is _UNSET: + continue + if value is None: + self._locks[segment] = None + continue + ids = self._resolve_lock(segment, value) + self._locks[segment] = ids + return self + + def _resolve_lock( + self, + segment: str, + value: Union[str, Iterable[str]], + ) -> Tuple[int, ...]: + """Resolve allele-name(s) → tuple of allele IDs against this + experiment's refdata. Raises on unknown names, on D locks + for VJ chains, or on non-string inputs.""" + if segment == "D" and self._refdata.chain_type != "vdj": + raise ValueError( + f"cannot lock D alleles on a {self._refdata.chain_type!r} chain" + ) + + if isinstance(value, str): + names: Tuple[str, ...] = (value,) + else: + try: + names = tuple(value) + except TypeError as exc: + raise TypeError( + f"restrict_alleles(): {segment} lock must be a name string or an " + f"iterable of name strings; got {type(value).__name__}" + ) from exc + if not all(isinstance(n, str) for n in names): + raise TypeError( + f"restrict_alleles(): {segment} lock entries must all be strings" + ) + if not names: + raise ValueError( + f"restrict_alleles(): {segment} lock list must be non-empty " + "(pass None to clear instead)" + ) + + index = self._allele_name_index(segment) + ids: List[int] = [] + seen: set = set() + for name in names: + if name not in index: + raise ValueError( + f"restrict_alleles(): no {segment} allele named {name!r} in refdata " + "(check spelling against `list_alleles` or the loaded config)" + ) + allele_id = index[name] + if allele_id in seen: + raise ValueError( + f"restrict_alleles(): duplicate {segment} allele {name!r} in lock list" + ) + seen.add(allele_id) + ids.append(allele_id) + return tuple(ids) + + def _allele_name_index(self, segment: str) -> Dict[str, int]: + """Build a name → allele_id map for the given segment by + scanning the refdata pool. Cheap enough to do per-call given + typical pool sizes (≤ a few hundred) and the rarity of + ``.restrict_alleles()``.""" + if segment == "V": + n = self._refdata.v_pool_size() + getter = self._refdata.v_allele + elif segment == "D": + n = self._refdata.d_pool_size() + getter = self._refdata.d_allele + elif segment == "J": + n = self._refdata.j_pool_size() + getter = self._refdata.j_allele + else: # pragma: no cover — guarded above. + raise ValueError(f"unsupported segment {segment!r}") + return {getter(i).name: i for i in range(n)} + + def _resolve_allele_weights( + self, + segment: str, + user_weights: Optional[Dict[str, float]], + ) -> Optional[Tuple[float, ...]]: + """Build a dense pool-aligned weight vector from the + user-supplied ``{name: weight}`` dict. Listed alleles get the + supplied weight; everything else gets ``1.0``. Returns + ``None`` when no weights were supplied (preserving the + upstream uniform-default behavior). + + Raises ``ValueError`` for unknown allele names, non-positive + weights, or D weights on a VJ chain. + """ + if user_weights is None: + return None + if not user_weights: + raise ValueError( + f"recombine: {segment.lower()}_allele_weights must contain " + "at least one (name, weight) entry" + ) + if segment == "D" and self._refdata.chain_type != "vdj": + raise ValueError( + f"recombine: cannot weight D alleles on a " + f"{self._refdata.chain_type!r} chain" + ) + + index = self._allele_name_index(segment) + if not index: + return None # No alleles for this segment (e.g. VJ + D). + + # Dense vector indexed by allele_id; default 1.0. + pool_size = max(index.values()) + 1 + weights: List[float] = [1.0] * pool_size + for name, w in user_weights.items(): + if not isinstance(name, str): + raise TypeError( + f"recombine: {segment.lower()}_allele_weights keys must " + f"be allele-name strings, got {type(name).__name__}" + ) + if isinstance(w, bool) or not isinstance(w, (int, float)): + raise TypeError( + f"recombine: {segment.lower()}_allele_weights['{name}'] " + f"must be numeric, got {type(w).__name__}" + ) + wf = float(w) + if not (wf > 0.0 and wf == wf and wf != float("inf")): + raise ValueError( + f"recombine: {segment.lower()}_allele_weights['{name}'] " + f"must be a finite positive number, got {wf}" + ) + if name not in index: + raise ValueError( + f"recombine: no {segment} allele named {name!r} in " + "refdata (check spelling against `list_alleles` or the " + "loaded config)" + ) + weights[index[name]] = wf + return tuple(weights) diff --git a/src/GenAIRR/_experiment/introspection.py b/src/GenAIRR/_experiment/introspection.py new file mode 100644 index 0000000..f639a71 --- /dev/null +++ b/src/GenAIRR/_experiment/introspection.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from typing import Any, List + +from .._describe import ( + _describe_experiment_header, + _describe_step_sequence, + _format_declared_contracts, +) +from .._pipeline_ir import _RecombineStep + + +class _IntrospectionMixin: + __slots__ = () + + def describe(self) -> str: + """Render a biology-style narrative of this experiment. + + The output is one line per step, prefixed with its position + in the chain. Locks, allele weights, NP-length distributions, + SHM kernels, and per-corruption rates are all surfaced. Use + this to sanity-check what a fluent chain actually encodes — + if it doesn't read like an immunology protocol, the chain is + too murky. + + Returns a multi-line string ending without a trailing newline. + Safe to ``print(exp.describe())``. + + Example:: + + >>> print(Experiment.on("human_igh").recombine().mutate(count=(5, 15)).describe()) + Experiment on human_igh (vdj, DataConfig) + 1. V(D)J recombination: sample V/D/J alleles; empirical exonuclease trim (V3', D5', D3', J5'); insert NP1 (0–11 weighted bases) and NP2 (0–11 weighted bases) + 2. Somatic hypermutation (S5F context model, human heavy-chain (HH_S5F)): 5–15 mutations/record + """ + header = _describe_experiment_header(self._refdata, self._dataconfig) + if not self._steps and not self._contracts: + return header + "\n (no steps appended yet)" + + lines = [header] + resolved = self._steps_with_locks_resolved() + body_lines = _describe_step_sequence(resolved, self._refdata.chain_type) + lines.extend(body_lines) + contracts_line = _format_declared_contracts(self._contracts) + if contracts_line: + lines.append(f" Constraints: {contracts_line}") + if self._metadata: + stamps = ", ".join(f"{k}={v!r}" for k, v in self._metadata.items()) + lines.append(f" Metadata stamped on every record: {stamps}") + return "\n".join(lines) + + def _steps_with_locks_resolved(self) -> List[Any]: + """Return a copy of ``self._steps`` with per-segment allele + locks (from :meth:`restrict_alleles`) injected into the first + :class:`_RecombineStep`. Compile-time injection lives in + :meth:`_build_simulator`; ``describe()`` needs the same + substitution to render lock info correctly without + side-effecting the live step list.""" + from dataclasses import replace as _replace + + any_lock = any(self._locks[seg] is not None for seg in ("V", "D", "J")) + if not any_lock: + return list(self._steps) + out: List[Any] = [] + injected = False + for step in self._steps: + if not injected and isinstance(step, _RecombineStep): + step = _replace( + step, + locks_v=self._locks["V"], + locks_d=self._locks["D"], + locks_j=self._locks["J"], + ) + injected = True + out.append(step) + return out + + def __repr__(self) -> str: + return f"" diff --git a/src/GenAIRR/_experiment/mutation.py b/src/GenAIRR/_experiment/mutation.py new file mode 100644 index 0000000..b95847e --- /dev/null +++ b/src/GenAIRR/_experiment/mutation.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +from typing import Dict, Iterable, Optional, Tuple, Union + +from .._normalize import _normalize_count +from .._pipeline_ir import ( + _ClonalForkStep, + _LineageForkStep, + _MutateStep, + _RepertoireForkStep, +) +from .._step_validation import ( + _DEFAULT_V_SUBREGION_RATES, + _validate_segment_rates, + _validate_v_subregion_rates, +) + + +class _MutationMixin: + __slots__ = () + + def mutate( + self, + *, + model: str = "s5f", + count: Optional[Union[int, Tuple[int, int], Iterable[Tuple[int, float]]]] = None, + rate: Optional[float] = None, + s5f_model: str = "hh_s5f", + segment_rates: Optional[Dict[str, float]] = None, + v_subregion_rates: Optional[Dict[str, float]] = None, + ) -> "Experiment": + """Append a somatic-hypermutation step. + + ``model`` selects the mutation kernel: + - ``"s5f"`` (default) — context-dependent SHM via the bundled + S5F kernel named in ``s5f_model``. Available kernels: + ``"hh_s5f"``, ``"hh_s5f_60"``, ``"hh_s5f_opposite"``, + ``"hkl_s5f"``. + - ``"uniform"`` — position-independent SHM. Each mutated + position gets a uniformly drawn A/C/G/T replacement. + + **Specify intensity with exactly one of ``rate`` or ``count``.** + + ``rate`` is the per-base mutation rate (e.g. ``0.03`` for 3 % + SHM, which is roughly memory B-cell SHM). At execute time the + engine draws ``count ~ Poisson(rate × pool_len)`` against each + record's current sequence length — so the realized count + scales with each record's actual length, matching how + immunologists report SHM in the literature. This is the + canonical, biology-default form. + + ``count`` is the legacy explicit count distribution, useful + for benchmark scripts that want a deterministic count + independent of record length: + - ``count=15`` — fixed: every simulation gets exactly 15 + mutations. + - ``count=(5, 25)`` — uniform integer in ``[5, 25]`` (both + endpoints inclusive). + - ``count=[(5, 1.0), (10, 2.0), ...]`` — explicit empirical + ``(count, weight)`` distribution. + + Passing both ``count`` and ``rate`` raises ``ValueError``. + Passing neither raises ``ValueError``. + + **TCR guard:** somatic hypermutation is a B-cell + phenomenon — T-cells do not undergo SHM in the periphery. + Calling ``.mutate()`` on a TCR-configured experiment raises + ``ValueError`` to prevent silent biological misuse. Use + ``pcr_amplify`` / ``sequencing_errors`` for sequencing-error + realism on TCR data instead. + """ + if self._is_tcr_refdata(): + raise ValueError( + "mutate(): somatic hypermutation does not occur in TCR " + "sequences (T-cells lack AID and the SHM machinery). The " + "configured refdata is a TCR locus. For sequencing-error " + "realism on TCR data, use pcr_amplify / sequencing_errors " + "/ polymerase_indels instead." + ) + model_lc = model.lower() + if model_lc not in ("uniform", "s5f"): + raise ValueError( + f"model must be 'uniform' or 's5f' (got {model!r})" + ) + if count is not None and rate is not None: + raise ValueError( + "mutate(): pass exactly one of `rate` or `count`, not both. " + "`rate` is the canonical biology default (e.g. rate=0.03 " + "for 3% SHM); `count` is the explicit per-record count " + "for benchmark / deterministic-count workflows." + ) + if count is None and rate is None: + raise ValueError( + "mutate(): pass exactly one of `rate` or `count`. " + "Suggested default: rate=0.03 (~3% SHM, memory B-cell range)." + ) + if rate is not None: + if not isinstance(rate, (int, float)) or isinstance(rate, bool): + raise TypeError( + f"rate must be a finite float in [0.0, 1.0], got " + f"{type(rate).__name__}" + ) + if not (0.0 <= float(rate) <= 1.0): + raise ValueError( + f"rate must be in [0.0, 1.0] (got {rate!r}); rate is a " + f"per-base mutation probability, not an absolute count." + ) + seg_rates_tuple = _validate_segment_rates(segment_rates) + v_sub_rates_tuple = _validate_v_subregion_rates(v_subregion_rates) + self._check_v_subregion_rates_satisfiable( + v_subregion_rates, v_sub_rates_tuple + ) + self._steps.append( + _MutateStep( + model=model_lc, + s5f_model_name=s5f_model, + rate=float(rate), + segment_rates=seg_rates_tuple, + v_subregion_rates=v_sub_rates_tuple, + ) + ) + return self + pairs = _normalize_count(count) + seg_rates_tuple = _validate_segment_rates(segment_rates) + v_sub_rates_tuple = _validate_v_subregion_rates(v_subregion_rates) + self._check_v_subregion_rates_satisfiable( + v_subregion_rates, v_sub_rates_tuple + ) + self._steps.append( + _MutateStep( + model=model_lc, + s5f_model_name=s5f_model, + count_pairs=pairs, + segment_rates=seg_rates_tuple, + v_subregion_rates=v_sub_rates_tuple, + ) + ) + return self + + def _check_v_subregion_rates_satisfiable( + self, + raw_rates: Optional[Dict[str, float]], + tuple_rates: Tuple[float, float, float, float, float], + ) -> None: + """Reject a non-default ``v_subregion_rates`` configuration + when the bound cartridge has zero annotated V alleles. + Audit §4: a non-default rate vector against a cartridge + without subregion annotations is unsatisfiable — no V site + would ever see a subregion factor, and the user is + almost certainly building against the wrong cartridge or + forgot to enable the annotation surface. + + Default rates (omitted kwarg, empty dict, or explicit + all-ones expansion) skip the check — those are no-ops and + compose cleanly with any cartridge. + """ + if raw_rates is None or tuple_rates == _DEFAULT_V_SUBREGION_RATES: + return + annotated = 0 + v_total = self._refdata.v_pool_size() + for v_id in range(v_total): + if self._refdata.v_allele(v_id).subregions: + annotated += 1 + # Early exit — we only need at least one annotated. + return + # Zero annotated V alleles: the user's rates can never bite. + raise ValueError( + "mutate(): v_subregion_rates was supplied but the bound " + "cartridge carries no V-subregion annotations on any V " + "allele (annotated_v_count=0 / " + f"total_v_count={v_total}). The rate vector would be a " + "deterministic no-op — almost certainly a builder bug. " + "Either drop v_subregion_rates or use a cartridge with " + "IMGT-gapped V sequences (the bundled human IGH / IGK / " + "IGL OGRDB cartridges derive subregions automatically; " + "see docs/v_region_substructure_audit.md)." + ) + + def _has_clonal_fork(self) -> bool: + """Whether any clonal fork has already been appended. + + Used by the DSL ordering guards on :meth:`invert_d`, + :meth:`receptor_revision`, and the clonal fork methods. + Each fork has a pre-fork parent/founder phase; recombination-time + mechanisms must be inherited by every descendant, emitted copy, + or lineage node. Misordered calls used to lower into the wrong + half and produce records with empty / default fields. The guards + reject those configurations at the DSL boundary. + """ + return any( + isinstance(s, (_ClonalForkStep, _RepertoireForkStep, _LineageForkStep)) + for s in self._steps + ) + + def _is_tcr_refdata(self) -> bool: + """Detect whether the bound refdata is a TCR locus. + + TCR allele names are prefixed with ``TR`` (TRA, TRB, TRG, + TRD); BCR with ``IG``. Inspecting the first V allele is + enough — locus is uniform across the pool. Returns ``False`` + when the V pool is empty (defensive — let downstream errors + surface that condition). + """ + if self._refdata.v_pool_size() == 0: + return False + first_v_name = self._refdata.v_allele(0).name + return first_v_name.upper().startswith("TR") diff --git a/src/GenAIRR/_experiment/recombination.py b/src/GenAIRR/_experiment/recombination.py new file mode 100644 index 0000000..5285529 --- /dev/null +++ b/src/GenAIRR/_experiment/recombination.py @@ -0,0 +1,683 @@ +from __future__ import annotations + +import warnings +from typing import Dict, Iterable, Optional, Tuple + +from GenAIRR import _engine + +from .._normalize import ( + _normalize_count, + _normalize_lengths, + _to_immutable_byte_pair_matrix, + _to_immutable_byte_pairs, + _to_immutable_pairs, +) +from .._pipeline_ir import ( + _DEFAULT_NP_LENGTHS, + _InvertDStep, + _PairedEndStep, + _RecombineStep, + _ReceptorRevisionStep, +) + + +class _RecombinationMixin: + __slots__ = () + + def recombine( + self, + *, + np1_lengths: Optional[Iterable[Tuple[int, float]]] = None, + np2_lengths: Optional[Iterable[Tuple[int, float]]] = None, + v_allele_weights: Optional[Dict[str, float]] = None, + d_allele_weights: Optional[Dict[str, float]] = None, + j_allele_weights: Optional[Dict[str, float]] = None, + ) -> "Experiment": + """Append a standard V(D)J recombination step. + + Compiles to: + - **VJ:** sample V → sample J → (trim V_3, J_5) → assemble V + → generate NP1 → assemble J. + - **VDJ:** sample V → sample D → sample J → (trim V_3, D_5, + D_3, J_5) → assemble V → generate NP1 → assemble D → + generate NP2 → assemble J. + + ``np1_lengths`` / ``np2_lengths`` default to the species' + empirical NP-length distributions (from + ``DataConfig.NP_lengths``) when the experiment is bound to + a DataConfig. For raw-RefDataConfig experiments where no + empirical data is available, both fall back to the uniform + ``[(0, 1.0), ..., (6, 1.0)]`` distribution and emit a + :class:`UserWarning` so the caller knows the synthetic + default is being used. Pass an explicit iterable of + ``(length, weight)`` tuples to override the default. + Passing ``np2_lengths`` on a VJ chain raises ``ValueError`` + (VJ chains have no NP2 region — there's no D segment to + bracket). + + **Exonuclease trim** is enabled by default and uses the + empirical per-segment trim distributions from the bound + ``DataConfig`` when available. To disable trim, or supply + custom trim distributions, call :meth:`trim` *after* + :meth:`recombine` in the chain (before any mutation / + corruption step). On a raw ``RefDataConfig`` without trim + data, recombine emits a :class:`UserWarning` and falls back + to a no-op trim. + + ``v_allele_weights`` / ``d_allele_weights`` / + ``j_allele_weights`` — optional ``{allele_name: + weight}`` dicts that bias allele sampling. Listed alleles + get the supplied positive weight; unlisted alleles default + to 1.0, so e.g. ``v_allele_weights={"IGHV3-23*01": 100}`` + boosts that allele while keeping every other V allele + possible at 1/100th the rate. Mutually exclusive with the + per-segment :meth:`restrict_alleles` restriction. Raises + ``ValueError`` for unknown allele names or non-positive + weights. + """ + # Explicit allele weights conflict with an attached genotype: + # the genotype owns allele presence + within-gene expression, and + # the phased lowering ignores recombine-step weights. Reject the + # combination instead of silently dropping the weights. + if any( + w is not None + for w in (v_allele_weights, d_allele_weights, j_allele_weights) + ): + self._user_allele_weights_set = True + if self._genotype is not None: + raise ValueError( + "recombine(*_allele_weights=...) and with_genotype() are " + "mutually exclusive: the genotype owns allele expression" + ) + + # VJ chains have no NP2 region — surface user mistakes loudly + # instead of silently dropping the argument. + if np2_lengths is not None and self._refdata.chain_type != "vdj": + raise ValueError( + f"np2_lengths is only valid for VDJ chains; the bound " + f"refdata is {self._refdata.chain_type!r} (no D segment, " + f"no NP2 region). Drop the np2_lengths kwarg or bind a " + f"VDJ refdata." + ) + + # Trim is default-on; .trim() may later disable or replace it. + trim = True + defaults = self._recombine_defaults() if trim or self._dataconfig else None + + # Raw-RefDataConfig path: there's no DataConfig backing this + # experiment, so empirical NP lengths and trim distributions + # don't exist. We fall back to a uniform NP-length default — + # historically silent — and surface a warning so the synthetic + # default isn't mistaken for real biology. The trim warning is + # emitted at compile() time instead, after .trim() has had a + # chance to disable trim explicitly. + if self._dataconfig is None: + if np1_lengths is None or ( + self._refdata.chain_type == "vdj" and np2_lengths is None + ): + warnings.warn( + "Experiment bound to a raw RefDataConfig with no empirical " + "NP-length distribution; falling back to uniform " + "[(0, 1.0), ..., (6, 1.0)]. Pass np1_lengths " + "(and np2_lengths for VDJ) explicitly to silence this.", + UserWarning, + stacklevel=2, + ) + + np1 = ( + _normalize_lengths(np1_lengths) + if np1_lengths is not None + else self._default_np_lengths(defaults, "np1") + ) + np2 = ( + _normalize_lengths(np2_lengths) + if np2_lengths is not None + else self._default_np_lengths(defaults, "np2") + ) + + if trim and defaults is not None: + trim_v_3 = _to_immutable_pairs(defaults.get("trim_v_3")) + trim_d_5 = _to_immutable_pairs(defaults.get("trim_d_5")) + trim_d_3 = _to_immutable_pairs(defaults.get("trim_d_3")) + trim_j_5 = _to_immutable_pairs(defaults.get("trim_j_5")) + else: + trim_v_3 = trim_d_5 = trim_d_3 = trim_j_5 = None + + # Cartridge-owned NP base distributions (Slice — Typed NP + # base model). Resolution lives in + # ``_dataconfig_extract.extract_recombine_defaults``: + # ``typed reference_models.np_bases → uniform`` (the + # legacy auto-lift of ``NP_first_bases`` / ``NP_transitions`` + # is deliberately deferred). ``None`` here means the + # pre-slice ``UniformBase`` applies. + if defaults is not None: + np1_base_pairs = _to_immutable_byte_pairs(defaults.get("np1_bases")) + np2_base_pairs = _to_immutable_byte_pairs(defaults.get("np2_bases")) + np1_markov_transitions = _to_immutable_byte_pair_matrix( + defaults.get("np1_markov_transitions") + ) + np2_markov_transitions = _to_immutable_byte_pair_matrix( + defaults.get("np2_markov_transitions") + ) + p_v_3_lengths = _to_immutable_pairs( + defaults.get("p_v_3_lengths") + ) + p_d_5_lengths = _to_immutable_pairs( + defaults.get("p_d_5_lengths") + ) + p_d_3_lengths = _to_immutable_pairs( + defaults.get("p_d_3_lengths") + ) + p_j_5_lengths = _to_immutable_pairs( + defaults.get("p_j_5_lengths") + ) + else: + np1_base_pairs = None + np2_base_pairs = None + np1_markov_transitions = None + np2_markov_transitions = None + p_v_3_lengths = None + p_d_5_lengths = None + p_d_3_lengths = None + p_j_5_lengths = None + + # Precedence (Slice — Allele Usage Estimation v1): + # + # 1. explicit `*_allele_weights=` kwarg + # 2. cartridge `reference_models.allele_usage` + # 3. uniform (legacy default) + # + # The kwarg path stays load-bearing for ad-hoc bias — + # only when the kwarg is omitted do we fall through to + # the typed cartridge plane (or uniform). Legacy + # `gene_use_dict` is NOT consulted. + if v_allele_weights is None and defaults is not None: + v_allele_weights = defaults.get("allele_usage_v") + if d_allele_weights is None and defaults is not None: + d_allele_weights = defaults.get("allele_usage_d") + if j_allele_weights is None and defaults is not None: + j_allele_weights = defaults.get("allele_usage_j") + weights_v = self._resolve_allele_weights("V", v_allele_weights) + weights_d = self._resolve_allele_weights("D", d_allele_weights) + weights_j = self._resolve_allele_weights("J", j_allele_weights) + + step = _RecombineStep( + np1_lengths=np1, + np2_lengths=np2, + trim_v_3=trim_v_3, + trim_d_5=trim_d_5, + trim_d_3=trim_d_3, + trim_j_5=trim_j_5, + weights_v=weights_v, + weights_d=weights_d, + weights_j=weights_j, + np1_base_pairs=np1_base_pairs, + np2_base_pairs=np2_base_pairs, + np1_markov_transitions=np1_markov_transitions, + np2_markov_transitions=np2_markov_transitions, + p_v_3_lengths=p_v_3_lengths, + p_d_5_lengths=p_d_5_lengths, + p_d_3_lengths=p_d_3_lengths, + p_j_5_lengths=p_j_5_lengths, + ) + self._steps.append(step) + return self + + def trim( + self, + *, + enabled: bool = True, + v_3: Optional[Iterable[Tuple[int, float]]] = None, + d_5: Optional[Iterable[Tuple[int, float]]] = None, + d_3: Optional[Iterable[Tuple[int, float]]] = None, + j_5: Optional[Iterable[Tuple[int, float]]] = None, + ) -> "Experiment": + """Configure exonuclease trim on the preceding :meth:`recombine` + step. + + Trim is a per-segment exonuclease step that lives biologically + *inside* V(D)J recombination (between allele selection and NP + insertion). It is on by default with empirical distributions + sourced from the bound ``DataConfig``. Use this method only + when you need to override that default: + + - ``trim(enabled=False)`` — disable trim entirely. Equivalent to + recombining against the raw allele endpoints. + - ``trim(v_3=..., d_5=..., d_3=..., j_5=...)`` — supply custom + per-segment trim-length distributions as ``(length, weight)`` + iterables. Omitted segments keep their empirical defaults + (or no-op on raw RefDataConfig). + + **Position in the chain.** ``trim()`` is configuration applied + to the most recent ``recombine()`` step. It must appear after + ``recombine()`` and before any mutation / corruption step. + Calling it before ``recombine()`` raises ``ValueError``; + calling it after a mutation step raises ``ValueError``. + + Raises ``ValueError`` if no ``recombine()`` step has been + appended yet, or if the chain already has steps that would + biologically follow recombination (mutate, corrupt_*). + """ + from dataclasses import replace as _replace + + # Find the most recent recombine step. + rec_idx: Optional[int] = None + for i in range(len(self._steps) - 1, -1, -1): + if isinstance(self._steps[i], _RecombineStep): + rec_idx = i + break + if rec_idx is None: + raise ValueError( + "trim() must be called after recombine(); no recombine " + "step is on this Experiment yet." + ) + + # Anything appended after the recombine step is wrong-ordered + # for trim configuration. + for j in range(rec_idx + 1, len(self._steps)): + offending = type(self._steps[j]).__name__ + raise ValueError( + f"trim() must be called immediately after recombine() — " + f"before any mutation / corruption / clonal-fork step. " + f"Found a {offending!r} between the latest recombine() " + f"and this trim() call." + ) + + prior: _RecombineStep = self._steps[rec_idx] # type: ignore[assignment] + if not enabled: + # Disable all trim slots, keep NP + weights untouched. + new_step = _replace( + prior, + trim_v_3=None, + trim_d_5=None, + trim_d_3=None, + trim_j_5=None, + trim_overridden=True, + ) + self._steps[rec_idx] = new_step + return self + + # Override individual distributions; pass-through on None. + def _resolve( + current: Optional[Tuple[Tuple[int, float], ...]], + override: Optional[Iterable[Tuple[int, float]]], + ) -> Optional[Tuple[Tuple[int, float], ...]]: + if override is None: + return current + return _normalize_lengths(override) + + new_step = _replace( + prior, + trim_v_3=_resolve(prior.trim_v_3, v_3), + trim_d_5=_resolve(prior.trim_d_5, d_5), + trim_d_3=_resolve(prior.trim_d_3, d_3), + trim_j_5=_resolve(prior.trim_j_5, j_5), + trim_overridden=True, + ) + self._steps[rec_idx] = new_step + return self + + def invert_d(self, *, prob: float = 0.05) -> "Experiment": + """Append a D-segment inversion step. + + Models V(D)J inversion: with probability ``prob`` the sampled + D allele is committed in reverse-complement orientation + instead of forward. Biologically, the RSS heptamers around D + can pair head-to-head and the D segment flips before joining; + prevalence is low (~1–5 %) but real. + + Engine path: a Bool is recorded under + ``sample_allele.d.inverted`` and the + :class:`~GenAIRR._engine.InvertDPass` commits + ``ReverseComplement`` on the D :class:`AlleleInstance` + between sampling and assembly. The + :class:`~GenAIRR._engine.AssembleSegmentPass`(D) (Slice B) + consumes the orientation flag and emits the reverse- + complemented D slice into the pool. Trace replay re-fires + the same orientation decision deterministically. + + **VDJ chains only.** VJ chains have no D pool; calling this + method on a VJ experiment raises ``ValueError``. + + **At most once per experiment.** Calling :meth:`invert_d` + twice raises ``ValueError`` — v1 picks a single per-pipeline + inversion probability rather than supporting last-one-wins + semantics (which would be silent for an over-eager builder). + + **Position in the chain.** Append after :meth:`recombine` + (and any :meth:`trim` override). Calling :meth:`invert_d` + before :meth:`recombine` raises ``ValueError`` at compile + time — the lowering needs the recombine sequence already + materialised in the engine plan so the explicit + ``before(invert_d, assemble.d)`` schedule edge can fire. + + The DSL does **not** expose the per-record orientation in the + AIRR record yet — that's the Slice E follow-up. End-to-end + observability today is via the trace + (``sample_allele.d.inverted``) and the pool bytes. + + Returns ``self`` so the call chains fluently. + """ + if self.chain_type != "vdj": + raise ValueError( + f"invert_d is only valid for VDJ chains (current chain_type={self.chain_type!r})" + ) + if any(isinstance(s, _InvertDStep) for s in self._steps): + raise ValueError( + "invert_d already configured on this experiment; v1 accepts at " + "most one inversion step per pipeline. Build a fresh Experiment " + "if you need a different probability." + ) + # Ordering guard — D inversion is a recombination-time + # decision and must be inherited by every clone descendant. + # When placed post-fork, the lowering silently drops the + # inversion probability (no recombine step in the post-fork + # half to consume it), producing records with d_inverted=False + # even at prob=1.0. Reject at the DSL boundary instead. + if self._has_clonal_fork(): + raise ValueError( + "invert_d must be called before the clonal fork; D " + "inversion is a recombination-time decision and must " + "be inherited by all clone descendants. Move the " + "invert_d(...) call before clonal_lineage(...), " + "clonal_repertoire(...), or expand_clones(...)." + ) + if not isinstance(prob, (int, float)): + raise ValueError( + f"invert_d prob must be a number in [0.0, 1.0], got {type(prob).__name__}" + ) + prob_f = float(prob) + # NaN check FIRST — NaN fails every `<=` comparison silently + # (`(0.0 <= nan)` is False), which would otherwise surface as + # a misleading "out of [0.0, 1.0]" message. Explicit NaN + # rejection gives the user the specific reason. + if prob_f != prob_f: + raise ValueError("invert_d prob must be a number, got NaN") + if not (0.0 <= prob_f <= 1.0): + raise ValueError( + f"invert_d prob must be in [0.0, 1.0], got {prob_f}" + ) + self._steps.append(_InvertDStep(prob=prob_f)) + return self + + def receptor_revision(self, *, prob: float = 0.05, same_haplotype: bool = True) -> "Experiment": + """Append a receptor-revision step. + + Models post-recombination V-segment replacement: with + probability ``prob`` the V slot is reassigned to a different + germline V allele and the V slice in the pool is rewritten. + Biologically, receptor revision is a B-cell tolerance + mechanism that lets a B cell escape autoreactivity by + replacing its V segment via secondary VDJ-recombination-like + rearrangement on the already-assembled receptor. + + Engine path: a Bool is recorded under + ``receptor_revision.applied`` for every simulation. On + ``true``, the + :class:`~GenAIRR._engine.ReceptorRevisionPass` additionally + records the replacement allele id at + ``receptor_revision.v_allele`` and the derived 3' trim at + ``receptor_revision.v_trim_3``, then commits + ``AssignmentChanged`` + ``TrimChanged`` + ``SegmentReplaced`` + against V through a single + :class:`~GenAIRR._engine.SimulationBuilder`. Slice C's + same-length retained constraint + (``allele.len() - trim_3 == old_v_len``, 5' trim fixed at 0) + keeps downstream pool positions stable; the + :class:`~GenAIRR._engine.LiveCallRefreshHook` (Slice B) + reacts to ``SegmentReplaced`` with an AllStructural- + equivalent V/D/J re-walk. + + **VDJ chains only.** Receptor revision is heavy-chain v1; + calling this method on a VJ experiment raises + ``ValueError``. + + **At most once per experiment.** Calling + :meth:`receptor_revision` twice raises ``ValueError`` — + last-one-wins semantics would silently override an + over-eager builder. + + **Position in the chain.** Appended after :meth:`recombine`; + the lowering inlines the engine ``push_receptor_revision`` + call immediately after ``push_assemble("J")`` so the pass + sees the fully-assembled V/D/J/NP pool. Subsequent + :meth:`mutate` / corruption passes lower after this step in + the plan, giving the canonical "recombine → revise → + mutate/corrupt" order the design doc §2 requires. + + AIRR records expose ``receptor_revision_applied`` (bool) and + ``original_v_call`` (the pre-revision V; empty when no + revision applied). ``v_call`` / ``truth_v_call`` are the + **post**-revision V. The three ``receptor_revision.*`` trace + records above remain available for replay. + + Returns ``self`` so the call chains fluently. + """ + if self.chain_type != "vdj": + raise ValueError( + "receptor_revision is only valid for VDJ chains " + f"(current chain_type={self.chain_type!r})" + ) + if any(isinstance(s, _ReceptorRevisionStep) for s in self._steps): + raise ValueError( + "receptor_revision already configured on this experiment; " + "v1 accepts at most one revision step per pipeline. Build " + "a fresh Experiment if you need a different probability." + ) + # Ordering guard — receptor revision is a recombination/ + # ancestor-time decision and must be inherited by every clone + # descendant. When placed post-fork, the lowering silently + # drops the revision probability (no recombine step in the + # post-fork half to consume it), producing records with + # receptor_revision_applied=False and empty original_v_call + # even at prob=1.0. Reject at the DSL boundary instead. + if self._has_clonal_fork(): + raise ValueError( + "receptor_revision must be called before " + "the clonal fork; receptor revision is a " + "recombination-time decision and must be inherited by " + "all clone descendants. Move the " + "receptor_revision(...) call before clonal_lineage(...), " + "clonal_repertoire(...), or expand_clones(...)." + ) + if not isinstance(prob, (int, float)): + raise ValueError( + "receptor_revision prob must be a number in [0.0, 1.0], " + f"got {type(prob).__name__}" + ) + prob_f = float(prob) + # NaN first — see the matching `invert_d` rationale. + if prob_f != prob_f: + raise ValueError("receptor_revision prob must be a number, got NaN") + if not (0.0 <= prob_f <= 1.0): + raise ValueError( + f"receptor_revision prob must be in [0.0, 1.0], got {prob_f}" + ) + if not isinstance(same_haplotype, bool): + raise ValueError( + "receptor_revision same_haplotype must be a bool, got " + f"{same_haplotype!r}" + ) + self._steps.append( + _ReceptorRevisionStep(prob=prob_f, same_haplotype=same_haplotype) + ) + return self + + def paired_end( + self, + *, + r1_length, + r2_length=None, + insert_size, + ) -> "Experiment": + """Append a paired-end / read-layout step. + + Models the Illumina paired-end read layout: each fragment + produces R1 (forward from the 5' adapter) and R2 + (reverse-complemented from the 3' adapter) windows over the + final projected molecule, plus an *insert size* that locates + R2's 3' end. The DSL exposes three integer distributions: + + - ``r1_length`` — required. + - ``r2_length`` — defaults to ``r1_length`` when ``None``. + Many Illumina libraries do run asymmetric (R2 quality + drops faster); the explicit shape lets callers opt in. + - ``insert_size`` — required. + + Each accepts the same three shapes the rest of the DSL + already uses for length-like distributions: + + - ``int`` — fixed value. + - ``(low, high)`` — uniform integer in the closed + interval ``[low, high]``. + - ``[(value, weight), …]`` — explicit empirical + distribution. + + Engine path: a trace-only + :class:`~GenAIRR._engine.PairedEndSamplingPass` records + three Ints at ``paired_end.r1_length`` / + ``paired_end.r2_length`` / ``paired_end.insert_size``; + the AIRR builder reads them back at projection time and + populates the eight ``read_layout`` / ``r1_sequence`` / + ``r2_sequence`` / ``r1_start`` / ``r1_end`` / ``r2_start`` / + ``r2_end`` / ``insert_size`` fields via the Slice B + projection kernel. ``rec.sequence`` is the only + coordinate space — end-loss and rev-comp projections have + already finalised the molecule by the time paired-end + windows are drawn (design doc §6 / §7). + + **Both VDJ and VJ chains supported.** Paired-end is a + sequencing-stage observable, not a biology mechanism; + it makes sense on every chain. + + **At most once per experiment.** Calling + :meth:`paired_end` twice raises ``ValueError`` — + last-one-wins semantics would silently override an + over-eager builder. + + **Position in the chain.** The compile pre-pass extracts + the step and pushes the engine pass at the **end** of the + plan, after every IR-mutating / corruption / orientation + step. Even though the pass is trace-only, recording the + choices last keeps the trace order aligned with the + biological/readout order (recombine → mutation → + corruption → end-loss → paired-end). + + Returns ``self`` so the call chains fluently. + """ + if any(isinstance(s, _PairedEndStep) for s in self._steps): + raise ValueError( + "paired_end already configured on this experiment; " + "v1 accepts at most one paired-end step per pipeline. " + "Build a fresh Experiment if you need a different " + "layout." + ) + + # Resolve default r2 → r1 BEFORE normalization so the + # downstream check pins both at the same source shape. + if r2_length is None: + r2_length = r1_length + + r1_pairs = _normalize_count(r1_length) + r2_pairs = _normalize_count(r2_length) + insert_pairs = _normalize_count(insert_size) + + # r1 / r2 lengths must be strictly positive. `_normalize_count` + # already rejects negatives but allows 0; the projection + # kernel needs > 0, so surface the violation at the DSL + # boundary with a clearer message. + for value, _w in r1_pairs: + if value <= 0: + raise ValueError( + f"paired_end r1_length must be positive, got {value}" + ) + for value, _w in r2_pairs: + if value <= 0: + raise ValueError( + f"paired_end r2_length must be positive, got {value}" + ) + # `_normalize_count` already enforces insert_size >= 0; + # the projection kernel enforces the same. + + # Fixed-value geometry checks. When r1/r2/insert are each + # single-value distributions we know the exact values the + # pass will emit, so reject `r1 > insert` / `r2 > insert` + # here rather than waiting for the engine's per-sample + # `InvalidDistributionOutput`. Distribution / range cases + # may sample valid combinations, so we defer those to the + # engine. + if len(r1_pairs) == 1 and len(insert_pairs) == 1: + r1_value = r1_pairs[0][0] + insert_value = insert_pairs[0][0] + if r1_value > insert_value: + raise ValueError( + f"paired_end r1_length ({r1_value}) > insert_size " + f"({insert_value}); the R1 window would run past " + f"the fragment 3' end." + ) + if len(r2_pairs) == 1 and len(insert_pairs) == 1: + r2_value = r2_pairs[0][0] + insert_value = insert_pairs[0][0] + if r2_value > insert_value: + raise ValueError( + f"paired_end r2_length ({r2_value}) > insert_size " + f"({insert_value}); the R2 window would run past " + f"the fragment 3' end." + ) + + self._steps.append( + _PairedEndStep( + r1_length=r1_pairs, + r2_length=r2_pairs, + insert_size=insert_pairs, + ) + ) + return self + + def _recombine_defaults(self): + """Lazy-extract the empirical distributions for this + experiment's DataConfig. Returns ``None`` for raw-RefDataConfig + experiments. Cached on first call. + """ + if self._dataconfig is None: + return None + from .._dataconfig_extract import extract_recombine_defaults + + return extract_recombine_defaults(self._dataconfig) + + def _build_contracts(self) -> Optional["_engine.ContractSet"]: + """Synthesize the engine ``ContractSet`` from declared bundles. + + Today only the productive bundle is recognized. Future bundles + (e.g. ``.in_frame_only()``) would compose here. Returns + ``None`` when no constraint methods have been called. + """ + if not self._contracts: + return None + # Composition across bundles isn't supported by the engine yet, + # so for now exactly one bundle is allowed. + if len(self._contracts) > 1: + raise NotImplementedError( + f"composing multiple constraint bundles is not yet supported; " + f"got {self._contracts!r}" + ) + bundle = self._contracts[0] + if bundle == "productive": + return _engine.productive() + raise NotImplementedError( + f"unknown constraint bundle {bundle!r}" + ) + + @staticmethod + def _default_np_lengths( + defaults, key: str + ) -> Tuple[Tuple[int, float], ...]: + """Pick the NP-length distribution to use when the user + didn't pass an explicit one: empirical if available, else + the uniform placeholder. + """ + if defaults is not None: + empirical = defaults.get(key) + if empirical: + return tuple(empirical) + return tuple(_DEFAULT_NP_LENGTHS) diff --git a/src/GenAIRR/_experiment/refdata_controls.py b/src/GenAIRR/_experiment/refdata_controls.py new file mode 100644 index 0000000..4d6e65b --- /dev/null +++ b/src/GenAIRR/_experiment/refdata_controls.py @@ -0,0 +1,104 @@ +from __future__ import annotations + + +class _RefdataControlsMixin: + __slots__ = () + + def curate_refdata( + self, + policy: str, + *, + allowed=None, + keep_unannotated: bool = True, + ) -> "Experiment": + """**Curation** — select which subset of the catalogue + participates in simulation. Replaces this experiment's + reference cartridge with a curated version. + + In the cartridge model, curation is distinct from validation + and from the catalogue itself: + + - **validation** describes what the catalogue contains; + - **curation** decides which alleles participate; + - **simulation** runs against the curated cartridge. + + ``policy`` is one of: + + - ``"raw"`` — identity policy; no-op (kept for symmetry). + - ``"functional_anchors_only"`` — drop V and J alleles whose + anchor doesn't satisfy the active anchor rule (missing + anchor, codon out of bounds, or codon AA outside the + rule's ``expected_amino_acids``). D and C pools pass + through unchanged. + - ``"functional_status"`` — filter V/D/J pools by IMGT + functional status. ``allowed`` is the list of statuses + to keep (default ``["functional"]``); accepted strings + are ``"functional"``, ``"orf"``, ``"pseudogene"``, and + ``"unknown"`` (case-insensitive, ``"F"`` / ``"P"`` + aliases also work). ``keep_unannotated`` (default + ``True``) controls whether alleles with no annotation + survive — bundled ``.pkl`` cartridges currently leave + status unannotated, so the default preserves backward + compatibility. C pool passes through unchanged. + + Use this when the catalogue (e.g. bundled ``mouse_igh`` or + ``human_tcrb``) includes pseudogene/ORF alleles you'd rather + exclude than permit at runtime via + :meth:`allow_curatable_refdata`. Curation is the + professional model; ``allow_curatable_refdata`` is the + broader runtime opt-in for sampling the raw catalogue + as-is. + + Curation cannot fix structural problems — duplicate allele + names, invalid sequence bytes, and locus/chain-type + mismatches still surface from the compile-time validator + regardless of the curated cartridge state. If curation + empties a required pool, ``compile()`` fails with + ``EmptyRequiredPool``. + + The curated cartridge's ``identity.source`` is tagged + ``|curated:`` so trace files and content hashes + distinguish raw from curated artefacts. Returns ``self`` so + the call chains fluently. See ``docs/reference_cartridge.md``. + """ + kwargs = {} + if allowed is not None: + kwargs["allowed"] = list(allowed) + if policy == "functional_status": + kwargs["keep_unannotated"] = bool(keep_unannotated) + self._refdata = self._refdata.curated(policy, **kwargs) + return self + + def allow_curatable_refdata(self, enabled: bool = True) -> "Experiment": + """Opt in to the lenient `AllowCuratable` refdata validation + mode for subsequent ``compile`` / ``run`` / ``run_records`` + calls. Returns ``self`` so the call chains fluently. + + Sits alongside :meth:`curate_refdata` as the cartridge's two + ways to handle pseudogene-bearing catalogues: + + - ``curate_refdata("functional_anchors_only")`` **removes** + the non-canonical alleles. Strict validation then passes + because the curated catalogue is clean. This is the + professional model — the cartridge identifies which + alleles actually participate. + - ``allow_curatable_refdata()`` **keeps** the catalogue + as-is and relaxes the validator. Strict validation passes + Curatable issues (pseudogene-shape anchor anomalies) but + still rejects Fatal ones (empty pools, duplicates, invalid + bytes, anchor out of bounds, locus/chain mismatch). + + Fatal issues are never opt-outable. Curatable issues — V + anchor codon not Cys, J anchor codon outside the locus's + expected set, missing V/J anchor — reflect pseudogene/ORF + entries in real reference catalogues (the bundled + ``mouse_igh`` and ``human_tcrb`` data both contain them). + + Recommended progression: start strict; if your catalogue + contains pseudogenes you want to *exclude*, use + :meth:`curate_refdata`; if you want to *sample from* them + explicitly, use this method. See + ``docs/reference_cartridge.md``. + """ + self._allow_curatable_refdata = bool(enabled) + return self diff --git a/src/GenAIRR/_experiment/run.py b/src/GenAIRR/_experiment/run.py new file mode 100644 index 0000000..d740c15 --- /dev/null +++ b/src/GenAIRR/_experiment/run.py @@ -0,0 +1,331 @@ +from __future__ import annotations + +from typing import Any, Dict, Iterator, List, Optional + +from .._compiled import ( + CompiledClonalExperiment, + CompiledLineageExperiment, + CompiledRepertoireExperiment, +) + + +class _RunMixin: + __slots__ = () + + def run_cohort( + self, + genotypes, + *, + n_per_subject: int = 1, + seed: int = 0, + counts=None, + strict: bool = False, + expose_provenance: bool = False, + validate_records: bool = False, + allow_curatable_refdata: Optional[bool] = None, + ) -> "CohortResult": + """Run a cohort: N subjects, each with its own diploid genotype, in one + call. A Python loop around the single-subject genotype path — each + subject is compiled and run independently, records are tagged with + ``subject_id`` and given a namespaced ``sequence_id``, and the per-subject + ``SimulationResult`` (with its own refdata) is collected into a + :class:`~GenAIRR.cohort.CohortResult`. + + ``n_per_subject`` applies to every subject; ``counts`` (a parallel + sequence, same length as ``genotypes``) overrides it per subject and may + contain ``0`` (that subject appears with zero records). Subject IDs are + taken from each genotype, or auto-assigned ``subject_0..N-1`` when all are + unset; mixed/duplicate IDs raise. Per-subject sub-seeds are derived + deterministically from ``seed``. + + Mutually exclusive with :meth:`with_genotype`, :meth:`restrict_alleles`, + and ``recombine(*_allele_weights=...)`` (the genotype owns allele + expression). :meth:`receptor_revision` is supported (each subject's + replacement V is restricted to its carried alleles on the drawn + chromosome); clonal forks are not supported in this release. + """ + import copy as _copy + import random as _random + + from ..cohort import ( + CohortResult, + CohortSubjectResult, + _resolve_counts, + _resolve_subject_ids, + ) + from ..genotype import Genotype + from ..result import SimulationResult + + gts = list(genotypes) + if not gts: + raise ValueError("run_cohort: genotypes must be a non-empty sequence") + for g in gts: + if not isinstance(g, Genotype): + raise TypeError( + f"run_cohort: every element must be a Genotype, got " + f"{type(g).__name__}") + + # Mutual exclusions — mirror with_genotype / compile. + if self._genotype is not None: + raise ValueError( + "run_cohort() and with_genotype() are mutually exclusive") + if any(v is not None for v in self._locks.values()): + raise ValueError( + "run_cohort() and restrict_alleles() are mutually exclusive") + if self._user_allele_weights_set: + raise ValueError( + "run_cohort() and recombine(*_allele_weights=...) are mutually " + "exclusive: the genotype owns allele expression") + # receptor_revision is supported per subject (each subject's compile + # builds a genotype-aware revision pass restricted to that subject's + # carried alleles on the drawn chromosome) — no rejection here. + if self._has_clonal_fork(): + raise ValueError( + "run_cohort() is not supported together with expand_clones() / " + "clonal_lineage() / clonal_repertoire() in this release") + # with_metadata() is applied per subject below, but it must not overwrite + # cohort-owned columns. + _cohort_owned = {"subject_id", "sequence_id", "haplotype"} + _md_collision = _cohort_owned & set(self._metadata) + if _md_collision: + raise ValueError( + f"run_cohort: with_metadata keys {sorted(_md_collision)} are " + f"cohort-owned (reserved); rename them") + + # Cartridge-hash check (same as with_genotype). + live_hash = self._refdata.content_hash() + for g in gts: + if g._source_hash != live_hash: + raise ValueError( + "run_cohort: a genotype was built against a different cartridge " + f"(content hash {g._source_hash!r} != experiment {live_hash!r})") + + resolved_counts = _resolve_counts(len(gts), n_per_subject, counts) + subject_ids = _resolve_subject_ids([g.subject_id for g in gts]) + base_rng = _random.Random(seed) + + subjects = [] + for g, sid, count in zip(gts, subject_ids, resolved_counts): + sub_seed = base_rng.getrandbits(63) + snap = g._snapshot() + snap.subject_id = sid + # Clone the uncompiled experiment; never mutate self. + exp_i = _copy.copy(self) + exp_i._genotype = snap + compiled = exp_i.compile(allow_curatable_refdata=allow_curatable_refdata) + refdata_i = compiled.refdata + if count == 0: + # run() rejects n < 1; build an empty result and stamp the + # genotype manually (run_records would otherwise do it). + res = SimulationResult.from_outcomes( + [], refdata_i, expose_provenance=expose_provenance) + res._genotypes = [snap] + else: + res = compiled.run_records( + n=count, seed=sub_seed, strict=strict, + expose_provenance=expose_provenance, + validate_records=validate_records) + # Apply with_metadata() per subject (parity with run_records), then + # namespace sequence_id so combined export never collides. Metadata + # is stamped first; the cohort-owned sequence_id rewrite wins (and a + # collision was already rejected up front). + if self._metadata: + for rec in res.records: + for key, value in self._metadata.items(): + rec[key] = value + for rec in res.records: + rec["sequence_id"] = f"{sid}_{rec.get('sequence_id', '')}" + subjects.append(CohortSubjectResult( + subject_id=sid, genotype=snap, result=res, refdata=refdata_i, + seed=sub_seed, count=count)) + + return CohortResult(subjects) + + def run_records( + self, + *, + n: Optional[int] = None, + seed: int = 0, + strict: bool = False, + expose_provenance: bool = False, + allow_curatable_refdata: Optional[bool] = None, + validate_records: bool = False, + ) -> "SimulationResult": + """Compile and run, then return the batch as a + :class:`SimulationResult` ready for ``.to_csv`` / ``.to_fasta`` + / ``.to_dataframe`` export. + + For non-clonal experiments ``n`` defaults to 1. For clonal + experiments (when the pipeline contains :meth:`with_clonal + _structure`) ``n`` defaults to ``n_clones * size`` and may + be omitted; passing ``n`` explicitly is allowed only if it + matches that product. + + ``expose_provenance=True`` appends ``truth_v_call``, + ``truth_d_call``, ``truth_j_call`` columns containing the + originally-sampled allele names — distinct from the + evidence-driven ``v_call`` / ``d_call`` / ``j_call`` fields + an aligner would produce. Useful for benchmarking aligners + against ground truth without keeping a side truth file. + + ``strict`` semantics match :meth:`run` — strict-mode applies + only to **fresh sampling**. Trace replay + (:meth:`CompiledExperiment.replay_from_trace_file`) consumes + recorded sentinel values verbatim, so a permissive trace + replays cleanly even with ``strict=True``. See + ``docs/productive_failure_mode_audit.md`` §5. + + ``validate_records=True`` runs + :meth:`SimulationResult.validate_records` on the freshly + built batch before returning. If any record fails the + postcondition validator the call raises + :class:`GenAIRR._validation.RecordValidationFailedError` + (a :class:`RuntimeError` subclass) carrying a + machine-greppable summary of the failures. The check costs + roughly one outcome-side re-derivation per record, so it + defaults to ``False``; flip it on in CI or when chasing a + suspected projection bug. The validator runs **before** + any ``with_metadata`` stamps are applied, matching the + order :meth:`SimulationResult.validate_records` would see + on a separate post-hoc call (metadata columns are + per-batch annotations, not engine-derived fields). + + Returns a :class:`SimulationResult`; clonal records carry + an integer ``clone_id`` field per row. + """ + compiled = self.compile(allow_curatable_refdata=allow_curatable_refdata) + if isinstance(compiled, CompiledClonalExperiment): + result = compiled.run_records( + n=n, + seed=seed, + strict=strict, + expose_provenance=expose_provenance, + validate_records=validate_records, + ) + elif isinstance(compiled, CompiledRepertoireExperiment): + if n is not None: + raise ValueError( + "The 'n' parameter is not supported for clonal_repertoire " + "experiments. The number of records depends on the per-clone " + "sizes drawn from the heavy-tailed distribution and the " + "read-collapse, not a fixed product." + ) + result = compiled.run_records( + seed=seed, + strict=strict, + expose_provenance=expose_provenance, + validate_records=validate_records, + ) + elif isinstance(compiled, CompiledLineageExperiment): + if n is not None: + raise ValueError( + "The 'n' parameter is not supported for clonal_lineage experiments. " + "The number of observed records depends on the lineage trees " + "grown from n_clones / n_sample / selection, not a fixed product." + ) + result = compiled.run_records( + seed=seed, + strict=strict, + expose_provenance=expose_provenance, + validate_records=validate_records, + ) + else: + result = compiled.run_records( + n=1 if n is None else n, + seed=seed, + strict=strict, + expose_provenance=expose_provenance, + validate_records=validate_records, + ) + if self._metadata: + for rec in result.records: + for key, value in self._metadata.items(): + rec[key] = value + return result + + def run( + self, + *, + n: Optional[int] = None, + seed: int = 0, + strict: bool = False, + allow_curatable_refdata: Optional[bool] = None, + ) -> List["_engine.Outcome"]: + """Compile and run this experiment ``n`` times. + + Equivalent to + ``self.compile().run(n=n, seed=seed, strict=strict)``. + Returns a list of :class:`GenAIRR._engine.Outcome` objects in + clone-major order for clonal experiments. + + Attach :meth:`productive_only` (or any future constraint + method) to the chain to require admissible records; the + runtime filters NP base draws, length samples, and mutation + / contamination substitutions in real time so the resulting + sequences satisfy the bundle by construction. + + Statically impossible contract configurations fail during + ``compile()`` with ``ValueError``. For runtime residue + — i.e., empty admissible support emerging dynamically at + sample time — ``strict=False`` (default) lets a pass consume + the slot as its explicit no-op / sentinel; ``strict=True`` + raises :class:`GenAIRR._engine.StrictSamplingError` instead. + + Note the two error paths use **different exception classes**: + ``ValueError`` for compile-time preconditions, + ``StrictSamplingError`` (subclass of ``Exception``, NOT of + ``ValueError``) for runtime empty-support. A bare + ``except ValueError:`` will not catch the runtime case. See + ``docs/productive_failure_mode_audit.md`` §6.1. + + ``strict`` only governs **fresh sampling**. Trace replay + (:meth:`CompiledExperiment.replay_from_trace_file`) consumes + recorded values verbatim; a permissive-recorded sentinel + trace replays cleanly even with ``strict=True``. To re-execute + a trace under strict-fresh semantics, call + ``simulator.run(seed=, strict=True)`` instead. + + **Output-correctness validation is on** :meth:`run_records` + **only.** This method returns raw ``Outcome`` objects, which + have no projected AIRR record to validate; pass + ``validate_records=True`` to :meth:`run_records` to opt into + the post-build check (which raises + :class:`GenAIRR._validation.RecordValidationFailedError` on + any failure). For an outcome-by-outcome post-hoc check + without re-running, build a :class:`SimulationResult` via + :meth:`SimulationResult.from_outcomes` and call + :meth:`SimulationResult.validate_records` on it. + """ + compiled = self.compile(allow_curatable_refdata=allow_curatable_refdata) + if isinstance(compiled, CompiledClonalExperiment): + return compiled.run(n=n, seed=seed, strict=strict) + return compiled.run(n=1 if n is None else n, seed=seed, strict=strict) + + def stream( + self, + *, + n: Optional[int] = None, + seed: int = 0, + strict: bool = False, + ) -> Iterator["_engine.Outcome"]: + """Compile and lazily yield :class:`GenAIRR._engine.Outcome` + objects. See :meth:`CompiledExperiment.stream` for full + semantics.""" + return self.compile().stream(n=n, seed=seed, strict=strict) + + def stream_records( + self, + *, + n: Optional[int] = None, + seed: int = 0, + strict: bool = False, + id_prefix: str = "seq", + ) -> Iterator[Dict[str, Any]]: + """Compile and lazily yield AIRR-format record dicts. See + :meth:`CompiledExperiment.stream_records`.""" + return self.compile().stream_records( + n=n, + seed=seed, + strict=strict, + id_prefix=id_prefix, + ) diff --git a/src/GenAIRR/_compile.py b/src/GenAIRR/_lowering.py similarity index 100% rename from src/GenAIRR/_compile.py rename to src/GenAIRR/_lowering.py diff --git a/src/GenAIRR/_pipeline_ir.py b/src/GenAIRR/_pipeline_ir.py index 6d900a1..2249c7f 100644 --- a/src/GenAIRR/_pipeline_ir.py +++ b/src/GenAIRR/_pipeline_ir.py @@ -6,7 +6,7 @@ This module owns the **step shape** (pure frozen dataclasses, no engine awareness). Lowering each step onto the engine-native -:class:`_engine.PassPlan` lives in :mod:`._compile`. The split lets +:class:`_engine.PassPlan` lives in :mod:`._lowering`. The split lets the pipeline IR be inspected, round-tripped, or serialized without importing the compiled Rust extension. diff --git a/src/GenAIRR/_result_common.py b/src/GenAIRR/_result_common.py new file mode 100644 index 0000000..91de310 --- /dev/null +++ b/src/GenAIRR/_result_common.py @@ -0,0 +1,38 @@ +"""Shared truth-column helpers for SimulationResult and its validation +mixin. Kept here (imported by both) so there is one definition and no +result <-> _result_validation import cycle. +""" +from __future__ import annotations + +from typing import Any, Dict, Optional + + +def _allele_name_or_empty(refdata: Any, segment: str, allele_id: Optional[int]) -> str: + """Look up an allele name from refdata by id; return ``""`` when + the id is None or the lookup fails (defensive). + """ + if allele_id is None: + return "" + try: + if segment == "V": + return refdata.v_allele(int(allele_id)).name + if segment == "D": + return refdata.d_allele(int(allele_id)).name + if segment == "J": + return refdata.j_allele(int(allele_id)).name + except Exception: + return "" + return "" + + +def _inject_truth_columns(outcome: Any, refdata: Any, record: Dict[str, Any]) -> None: + """Append `truth_v_call` / `truth_d_call` / `truth_j_call` + columns to ``record`` from the originally-sampled allele ids + stored in the simulation's `assignments`. Distinct from + `v_call` / `d_call` / `j_call`, which are evidence-driven and + can change under heavy SHM. + """ + sim = outcome.final_simulation() + record["truth_v_call"] = _allele_name_or_empty(refdata, "V", sim.v_allele_id()) + record["truth_d_call"] = _allele_name_or_empty(refdata, "D", sim.d_allele_id()) + record["truth_j_call"] = _allele_name_or_empty(refdata, "J", sim.j_allele_id()) diff --git a/src/GenAIRR/_result_export.py b/src/GenAIRR/_result_export.py new file mode 100644 index 0000000..0b8aa21 --- /dev/null +++ b/src/GenAIRR/_result_export.py @@ -0,0 +1,508 @@ +"""Serialization/export surface for SimulationResult, extracted verbatim as +a mixin (behavior-preserving). SimulationResult inherits these methods; self +resolves via MRO. +""" +from __future__ import annotations + +import csv +from typing import Any, Dict, List + + +# Canonical AIRR-style column order used by ``to_csv`` / ``to_tsv``. +# Anchors the header row to a stable shape across runs and keeps +# the output diff-friendly. +# Coordinate fields whose values are 0-based half-open (Python +# convention) by default. ``airr_strict=True`` exports add 1 to each +# *_start field to convert to the AIRR-spec 1-based-inclusive form; +# the *_end fields keep their existing value because Python's +# half-open `end` already equals the 1-based-inclusive `end`. +_AIRR_STRICT_START_FIELDS = ( + "v_sequence_start", + "d_sequence_start", + "j_sequence_start", + "v_alignment_start", + "d_alignment_start", + "j_alignment_start", + "v_germline_start", + "d_germline_start", + "j_germline_start", + "junction_start", +) + + +def _to_airr_strict(rec: Dict[str, Any]) -> Dict[str, Any]: + """Return a copy of ``rec`` with every coord ``*_start`` field + converted from 0-based half-open to 1-based-inclusive. ``*_end`` + fields are unchanged because the Python half-open end already + equals the 1-based-inclusive end (e.g. ``[0, 5)`` ↔ ``[1, 5]``). + + ``None`` start values are left as-is. Unset / missing keys are + silently ignored. + """ + converted = dict(rec) + for field in _AIRR_STRICT_START_FIELDS: + val = converted.get(field) + if isinstance(val, int): + converted[field] = val + 1 + return converted + + +_DEFAULT_COLUMN_ORDER = [ + # AIRR metadata + "sequence_id", + "sequence", + "sequence_aa", + "sequence_alignment", + "germline_alignment", + "germline_alignment_d_mask", + "sequence_length", + "rev_comp", + "locus", + # Calls + "v_call", + "v_cigar", + "v_score", + "v_identity", + "v_support", + "v_sequence_start", + "v_sequence_end", + "v_alignment_start", + "v_alignment_end", + "v_germline_start", + "v_germline_end", + "v_trim_5", + "v_trim_3", + "d_call", + "d_cigar", + "d_score", + "d_identity", + "d_support", + "d_sequence_start", + "d_sequence_end", + "d_alignment_start", + "d_alignment_end", + "d_germline_start", + "d_germline_end", + "d_trim_5", + "d_trim_3", + "j_call", + "j_cigar", + "j_score", + "j_identity", + "j_support", + "j_sequence_start", + "j_sequence_end", + "j_alignment_start", + "j_alignment_end", + "j_germline_start", + "j_germline_end", + "j_trim_5", + "j_trim_3", + "c_call", + # Junction + "junction", + "junction_aa", + "junction_start", + "junction_end", + "junction_length", + # NP regions + "np1", + "np1_aa", + "np1_length", + "np2", + "np2_aa", + "np2_length", + # Functionality + "productive", + "vj_in_frame", + "stop_codon", + # SHM + corruption (non-AIRR; GenAIRR additions) + "n_mutations", + "mutation_rate", + "n_pcr_errors", + "n_quality_errors", + "n_indels", + # Per-segment indel counters (docs/indel_provenance_audit.md + # §6.2). NP1/NP2 indels are excluded from these but still + # count toward `n_indels`. + "n_v_indels", + "n_d_indels", + "n_j_indels", + # Per-segment SHM mutation counters + # (docs/mutation_provenance_audit.md). Aggregated by walking + # `outcome.events()`, filtering to the mutate.{uniform,s5f} + # passes only, and bucketing each `BaseChanged` by carried + # segment. NP1+NP2 roll into `n_np_mutations`. Sum equals + # `n_mutations` by construction. + "n_v_mutations", + "n_d_mutations", + "n_j_mutations", + "n_np_mutations", + # V-subregion SHM partition + # (docs/v_subregion_mutation_counters_audit.md). Six fields + # that partition `n_v_mutations` by the assigned V allele's + # IMGT subregion intervals: five canonical labels plus + # `n_v_unannotated_mutations` for V events that can't be + # attributed (missing assignment, empty annotations, V-side + # CDR3 stretch, or indel-inserted V bases). Aggregated in the + # same `outcome.events()` walk as the per-segment counters, + # using the same `mutate.{uniform,s5f}` pass-name filter. + # On bundled human OGRDB cartridges the unannotated bucket is + # 0 on every record under the canonical pass order. + "n_fwr1_mutations", + "n_cdr1_mutations", + "n_fwr2_mutations", + "n_cdr2_mutations", + "n_fwr3_mutations", + "n_v_unannotated_mutations", + # Observation-stage length loss (EndLossPass / primer_trim_*). + # Distinct from recombination-stage v_trim_*/j_trim_*. See + # docs/primer_trim_end_loss_audit.md §6.1. + "end_loss_5_length", + "end_loss_3_length", + "is_contaminant", + # D inversion provenance (V(D)J inversion event). True when the + # simulation committed the D allele in reverse-complement + # orientation; false for VJ chains, VDJ chains without + # `Experiment.invert_d(...)`, and inversion decisions that + # landed on the forward branch. See + # `docs/d_inversion_design.md` §6.3. + "d_inverted", + # Receptor revision provenance (Slice E of the receptor-revision + # roadmap). `receptor_revision_applied` is True iff the + # `ReceptorRevisionPass` fired with applied=True; `original_v_call` + # carries the V allele name the recombine pass originally committed + # (empty when no revision happened). `v_call` continues to report + # the post-revision identity. See + # `docs/receptor_revision_design.md` §7. + "receptor_revision_applied", + "original_v_call", + # Paired-end / read layout (Slice A of the paired-end roadmap). + # All eight fields default to empty / None / 0 / empty under + # the legacy single-molecule projection; the Slice B/C projection + # layer populates them. The order here matches the Rust struct + # layout in `engine_rs/src/airr_record/record.rs`. See + # `docs/paired_end_design.md` §10. + "read_layout", + "r1_sequence", + "r2_sequence", + "r1_start", + "r1_end", + "r2_start", + "r2_end", + "insert_size", +] + + +class _ResultExport: + """Serialization/export mixin for :class:`SimulationResult`. + + Behavior-preserving extraction: every method resolves ``self`` + (``self._records`` etc.) via the MRO of the concrete + :class:`SimulationResult` that inherits this mixin. + """ + + __slots__ = () + + def to_dataframe(self, *, airr_strict: bool = False): + """Return a :class:`pandas.DataFrame` with one row per record. + + ``airr_strict=True`` converts all 0-based half-open coord + ``*_start`` fields to the AIRR-spec 1-based-inclusive form + (``*_end`` fields are unchanged). Useful when handing the + DataFrame off to AIRR-strict downstream tooling. + + Raises ``ImportError`` if pandas isn't installed (pandas is + an optional extra: ``pip install GenAIRR[all]``). + """ + try: + import pandas as pd + except ImportError as exc: # pragma: no cover — depends on env + raise ImportError( + "to_dataframe() requires pandas. Install with " + "`pip install GenAIRR[all]` or `pip install pandas`." + ) from exc + + if not self._records: + return pd.DataFrame(columns=_DEFAULT_COLUMN_ORDER) + records = ( + [_to_airr_strict(r) for r in self._records] + if airr_strict + else self._records + ) + return pd.DataFrame(records, columns=self._column_order()) + + def to_tsv(self, path: str, *, airr_strict: bool = False) -> None: + """Write the records as AIRR-style TSV (tab-separated). The + header row uses :data:`_DEFAULT_COLUMN_ORDER`. + + ``airr_strict=True`` converts coord ``*_start`` fields to + 1-based-inclusive (AIRR spec). + """ + self._write_delimited(path, "\t", airr_strict=airr_strict) + + def to_csv(self, path: str, *, airr_strict: bool = False) -> None: + """Write the records as comma-separated values. Convenience + alongside :meth:`to_tsv` — most analysis tooling prefers TSV + for AIRR data. + + ``airr_strict=True`` converts coord ``*_start`` fields to + 1-based-inclusive (AIRR spec). + """ + self._write_delimited(path, ",", airr_strict=airr_strict) + + def to_fasta(self, path: str, *, prefix: str = "seq") -> None: + """Write the assembled sequences as FASTA. Each record gets + a header of the form ``">{prefix}{i}|v_call=...|j_call=..."``. + """ + with open(path, "w", encoding="utf-8") as fh: + for i, rec in enumerate(self._records): + seq = rec.get("sequence", "") + v_call = rec.get("v_call") or "" + j_call = rec.get("j_call") or "" + fh.write(f">{prefix}{i}|v_call={v_call}|j_call={j_call}\n") + fh.write(f"{seq}\n") + + def to_fastq( + self, + path: str, + *, + quality: str = "illumina", + prefix: str = "seq", + **quality_kwargs, + ) -> None: + """Write the assembled sequences as FASTQ. + + Each record produces: + + :: + + @{prefix}{i}|v_call=...|j_call=... + + + + + + Parameters: + path: output file path. + quality: name of the quality model — ``"illumina"`` + (smoothed trapezoid, default) or ``"constant"`` + (single Q value across the read). + prefix: per-read header prefix (default ``"seq"``). + **quality_kwargs: forwarded to the quality model + constructor. ``ConstantQualityModel`` accepts + ``q`` (default 30), ``low_q`` (default 10), + ``n_q`` (default 2). ``IlluminaQualityModel`` + accepts ``peak_q``, ``start_q``, ``end_q``, + ``ramp_len``, ``tail_len``, ``low_q``, ``n_q``. + + FASTQ uppercases the sequence bases — GenAIRR's lowercase + corruption-marker convention is preserved by routing + lowercase positions to ``low_q`` in the quality string, + the standard FASTQ way of conveying low-confidence bases. + """ + from ._qmodel import phred_to_ascii, resolve_quality_model + + model = resolve_quality_model(quality, **quality_kwargs) + with open(path, "w", encoding="utf-8") as fh: + for i, rec in enumerate(self._records): + seq = rec.get("sequence", "") + v_call = rec.get("v_call") or "" + j_call = rec.get("j_call") or "" + q_array = model.quality_array(seq) + if len(q_array) != len(seq): + raise RuntimeError( + f"quality model returned {len(q_array)} scores for " + f"{len(seq)}-base sequence" + ) + q_string = phred_to_ascii(q_array) + fh.write(f"@{prefix}{i}|v_call={v_call}|j_call={j_call}\n") + fh.write(f"{seq.upper()}\n") + fh.write("+\n") + fh.write(f"{q_string}\n") + + def to_paired_fastq( + self, + r1_path: str, + r2_path: str, + *, + quality: str = "illumina", + overwrite: bool = False, + **quality_kwargs, + ) -> None: + """Write the per-record paired-end reads as two FASTQ files. + + Each AIRR record contributes one R1 record (to ``r1_path``) + and one R2 record (to ``r2_path``): + + :: + + R1 file: R2 file: + @{sequence_id}/1 @{sequence_id}/2 + + + + + + + Read names use the AIRR record's own ``sequence_id`` with the + canonical Illumina-portable ``/1`` / ``/2`` suffix (older + convention but universally accepted by BWA / STAR / + samtools / Picard; the seven-field colon-separated full + Illumina header doesn't have a GenAIRR analogue — no flow + cell, no lane, no index — and is out of scope here per + `docs/fastq_export_design.md` §5). + + ``r2_sequence`` is **already** the reverse complement of + ``sequence[r2_start:r2_end]`` at projection time (the AIRR + validator's `PairedEndWindowMismatch { side: R2 }` enforces + the invariant); this writer outputs it verbatim. Applying a + second RC would corrupt the read. + + Parameters: + r1_path: output path for the R1 FASTQ file. + r2_path: output path for the R2 FASTQ file. + quality: name of the quality model — ``"illumina"`` + (smoothed trapezoid, default) or ``"constant"`` + (single Q value across the read). Same vocabulary + as :meth:`to_fastq`. The model is consulted + independently for R1 and R2; both reads get their + own quality string starting from position 0 (this + is the correct Illumina-style behaviour — each + read has its own ramp-up and tail). + overwrite: when ``False`` (default) the writer raises + ``FileExistsError`` if either output path already + exists. Set ``True`` to allow overwriting. + **quality_kwargs: forwarded to the quality model + constructor — same surface as :meth:`to_fastq`. + + Raises: + FileExistsError: when ``overwrite=False`` and either + output path already exists. + ValueError: when a record's ``read_layout`` is not + ``"paired_end"`` (the experiment hasn't run + ``.paired_end(...)``), or when ``r1_sequence`` / + ``r2_sequence`` is empty. + RuntimeError: when the quality model produces a + quality array whose length disagrees with the + read sequence (same shape as :meth:`to_fastq`). + + FASTQ uppercases the read bases — GenAIRR's lowercase + corruption-marker convention is preserved by routing + lowercase positions to the model's ``low_q`` parameter, + same as :meth:`to_fastq`. + """ + import os + + from ._qmodel import phred_to_ascii, resolve_quality_model + + # ── 1. Output-path overwrite guard. ────────────────── + if not overwrite: + for label, path in (("r1_path", r1_path), ("r2_path", r2_path)): + if os.path.exists(path): + raise FileExistsError( + f"to_paired_fastq: {label}={path!r} already exists " + "and overwrite=False. Pass overwrite=True to " + "replace it." + ) + + # ── 2. Resolve the quality model once. ─────────────── + model = resolve_quality_model(quality, **quality_kwargs) + + # ── 3. Walk records, validating layout + writing. ──── + with open(r1_path, "w", encoding="utf-8") as r1_fh, open( + r2_path, "w", encoding="utf-8" + ) as r2_fh: + for i, rec in enumerate(self._records): + sequence_id = rec.get("sequence_id") or f"seq{i}" + # 3a. Read-layout guard. + read_layout = rec.get("read_layout", "") + if read_layout != "paired_end": + raise ValueError( + f"to_paired_fastq: record {i} " + f"(sequence_id={sequence_id!r}) has " + f"read_layout={read_layout!r} — paired-end FASTQ " + "export requires read_layout='paired_end'. Run " + "Experiment.paired_end(r1_length=…, " + "insert_size=…) on the experiment before " + "exporting." + ) + r1_seq = rec.get("r1_sequence") or "" + r2_seq = rec.get("r2_sequence") or "" + # 3b. Empty-window guard. Belt-and-suspenders — + # the projection kernel shouldn't produce empty + # windows on a paired-layout record, but a downstream + # consumer that hand-edited the record dict would + # otherwise silently emit a zero-length FASTQ read. + if not r1_seq: + raise ValueError( + f"to_paired_fastq: record {i} " + f"(sequence_id={sequence_id!r}) has empty " + "r1_sequence despite read_layout='paired_end'" + ) + if not r2_seq: + raise ValueError( + f"to_paired_fastq: record {i} " + f"(sequence_id={sequence_id!r}) has empty " + "r2_sequence despite read_layout='paired_end'" + ) + # 3c. Quality strings. Each read is scored + # independently — Illumina-style ramp shape resets + # per read, which is the correct biological model + # (R1 and R2 don't share a per-base quality + # profile). + q_r1 = model.quality_array(r1_seq) + if len(q_r1) != len(r1_seq): + raise RuntimeError( + f"to_paired_fastq: quality model returned " + f"{len(q_r1)} scores for {len(r1_seq)}-base R1 " + f"sequence (record {i})" + ) + q_r2 = model.quality_array(r2_seq) + if len(q_r2) != len(r2_seq): + raise RuntimeError( + f"to_paired_fastq: quality model returned " + f"{len(q_r2)} scores for {len(r2_seq)}-base R2 " + f"sequence (record {i})" + ) + # 3d. Write the two 4-line records. + r1_fh.write(f"@{sequence_id}/1\n") + r1_fh.write(f"{r1_seq.upper()}\n") + r1_fh.write("+\n") + r1_fh.write(f"{phred_to_ascii(q_r1)}\n") + r2_fh.write(f"@{sequence_id}/2\n") + r2_fh.write(f"{r2_seq.upper()}\n") + r2_fh.write("+\n") + r2_fh.write(f"{phred_to_ascii(q_r2)}\n") + + # ── internals ─────────────────────────────────────────────────── + + def _column_order(self) -> List[str]: + """Pick the column order for tabular exports. Starts with + the canonical default order; appends any extra columns the + records have (e.g. when callers add custom fields).""" + seen = set(_DEFAULT_COLUMN_ORDER) + extras: List[str] = [] + for rec in self._records: + for key in rec: + if key not in seen: + seen.add(key) + extras.append(key) + return _DEFAULT_COLUMN_ORDER + extras + + def _write_delimited( + self, path: str, delimiter: str, *, airr_strict: bool = False + ) -> None: + columns = self._column_order() + with open(path, "w", encoding="utf-8", newline="") as fh: + writer = csv.DictWriter( + fh, + fieldnames=columns, + delimiter=delimiter, + lineterminator="\n", + extrasaction="ignore", + ) + writer.writeheader() + for rec in self._records: + source = _to_airr_strict(rec) if airr_strict else rec + # Replace ``None`` with empty string so CSV columns + # don't carry literal ``"None"`` strings. + row = {k: ("" if v is None else v) for k, v in source.items()} + writer.writerow(row) diff --git a/src/GenAIRR/_result_validation.py b/src/GenAIRR/_result_validation.py new file mode 100644 index 0000000..8be0b27 --- /dev/null +++ b/src/GenAIRR/_result_validation.py @@ -0,0 +1,666 @@ +"""Record/family validation surface for SimulationResult: the public +ValidationReport / FamilyValidationReport result types plus the validate_* +methods extracted verbatim as a mixin (behavior-preserving). +""" +from __future__ import annotations + +from typing import Any, Dict, List + +from ._result_common import _inject_truth_columns + + +class ValidationReport: + """Aggregate report from :meth:`SimulationResult.validate_records`. + + Carries the result of the **public AIRR output correctness** gate + over a batch — the downstream contract that says "every projected + record is internally consistent with its outcome." + + Attributes: + ``count`` — total records checked. + ``failures`` — list of failing records, each as a dict with + ``record_index``, ``sequence_id``, and ``issues`` + (the list of dicts returned by + :meth:`Outcome.validate_record`). + ``ok`` — True iff every record passed (``failures`` is + empty). + + The report is truthy iff ``ok`` is True, so ``assert report`` + works as a one-line CI guard. + + Sibling gate: :meth:`Outcome.check_live_call_cache_parity` — + internal cache-correctness check on the state that feeds + projection. If both fail on the same batch, fix parity first + (a stale cache leaks into projection); see + :meth:`SimulationResult.validate_records` docstring for the + full troubleshooting rule. + """ + + __slots__ = ("count", "failures") + + def __init__(self, count: int, failures: List[Dict[str, Any]]) -> None: + self.count = count + self.failures: List[Dict[str, Any]] = failures + + @property + def ok(self) -> bool: + """True iff every record validated without issues.""" + return not self.failures + + def __bool__(self) -> bool: + return self.ok + + def __len__(self) -> int: + """Number of failing records (not total). Use ``.count`` for + the total.""" + return len(self.failures) + + def __repr__(self) -> str: + if self.ok: + return f"" + return ( + f"" + ) + + def summary(self) -> Dict[str, int]: + """Histogram of issue kinds across all failures. Useful for + ``print(report.summary())`` to see what's wrong at a glance.""" + counts: Dict[str, int] = {} + for failure in self.failures: + for issue in failure["issues"]: + kind = issue.get("kind", "Unknown") + counts[kind] = counts.get(kind, 0) + 1 + return counts + + +class FamilyValidationReport: + """Aggregate report from :meth:`SimulationResult.validate_families`. + + Carries the result of the **clonal-family consistency** gate + over a batch — the downstream contract that says "every group of + records sharing a ``clone_id`` agrees on the recombination-time + truth fields the engine puts on each descendant by construction." + + Attributes: + ``count`` — total records inspected. + ``family_count`` — number of distinct ``clone_id`` groups + (``0`` when the batch carries no + clonal records). + ``members_per_family`` — mapping ``{clone_id: n_members}``; + empty when ``family_count == 0``. + ``failures`` — list of failing-group dicts. Each + carries ``clone_id``, + ``record_indices``, ``issue_kind``, + and ``values`` (the divergent values + observed in the group, sorted). + ``ok`` — ``True`` iff ``failures`` is empty. + + The report is truthy iff ``ok`` is True so ``assert report`` + works as a one-line CI guard. + + Sibling gate: :class:`ValidationReport` — the per-record AIRR + projection postcondition check. Family validation runs *after* + that gate passes; a family-divergence failure means the records + are each internally consistent with their own outcome, but they + disagree across the clonal group on a field that should be + invariant by construction (heavy-SHM tie-set widening on the + live-call ``v_call`` etc. is intentionally NOT a divergence — + only ``truth_*_call`` invariance is enforced). + + The validator is a strict subset of the audit's §6 spec + (see ``docs/clonal_family_design.md``): pre-SHM junction + invariance, mutation-distance distribution, parent-trace + reconstruction, and lineage topology are deliberately deferred + to later slices. + """ + + __slots__ = ("count", "family_count", "members_per_family", "failures") + + def __init__( + self, + count: int, + family_count: int, + members_per_family: Dict[Any, int], + failures: List[Dict[str, Any]], + ) -> None: + self.count = count + self.family_count = family_count + self.members_per_family: Dict[Any, int] = dict(members_per_family) + self.failures: List[Dict[str, Any]] = failures + + @property + def ok(self) -> bool: + """True iff every family group passed every invariant check.""" + return not self.failures + + def __bool__(self) -> bool: + return self.ok + + def __len__(self) -> int: + """Number of failing groups (not total). Use ``.count`` for + the total record count and ``.family_count`` for the total + group count.""" + return len(self.failures) + + def __repr__(self) -> str: + if self.ok: + return ( + f"" + ) + return ( + f"" + ) + + def summary(self) -> Dict[str, int]: + """Histogram of issue kinds across all failing groups. Useful + for ``print(report.summary())`` to see what's wrong at a + glance.""" + counts: Dict[str, int] = {} + for failure in self.failures: + kind = failure.get("issue_kind", "Unknown") + counts[kind] = counts.get(kind, 0) + 1 + return counts + + +# ── family-layer invariants ────────────────────────────────────────── +# +# (Field, issue_kind) pairs the family validator enforces. Each pair +# names a recombination-time provenance field that should be +# identical across every descendant of a clone *when the field is +# present*. We deliberately skip the field for a group when it is +# absent from every record in the group (e.g. ``expose_provenance`` +# was not enabled), and require it to be present on every record once +# any record in the group carries it. +_FAMILY_TRUTH_INVARIANTS: tuple = ( + ("truth_v_call", "TruthVCallDiverges"), + ("truth_d_call", "TruthDCallDiverges"), + ("truth_j_call", "TruthJCallDiverges"), +) + + +class _ResultValidation: + """Record/family validation mixin for :class:`SimulationResult`. + + Behavior-preserving extraction: every method resolves ``self`` + (``self._records`` / ``self._outcomes`` / ``self._parents``) via + the MRO of the concrete :class:`SimulationResult` that inherits + this mixin. + """ + + __slots__ = () + + def validate_records(self, refdata: Any) -> "ValidationReport": + """**Public AIRR output correctness check.** + + Run the postcondition validator over every record in this + result and return a :class:`ValidationReport`. This is the + gate a downstream consumer cares about: "is each projected + AIRR record internally consistent with the engine state that + produced it?" + + Each record is re-derived independently from its original + ``Outcome`` (trace + event ledger + final ``Simulation``) + and compared against the projected dict. A record passes + when ``outcome.validate_record(refdata, sequence_id=...)`` + returns an empty issue list. Failures collect the + ``record_index``, ``sequence_id``, and the issue dicts. + + **Companion check** — for engine-side integrity, see + :meth:`Outcome.check_live_call_cache_parity` (returns the + cached-vs-fresh divergence on the live-call cache that + *feeds* projection). + + **Troubleshooting rule** — if a CI run has both this + validator AND the parity harness failing on the same batch, + fix the parity divergence FIRST: a stale cache can leak + into projection and produce spurious validator failures. + Once parity is green, rerun the validator; remaining + failures point at a real projection-layer bug. + + ``refdata`` must be the same :class:`RefDataConfig` the + outcomes were produced against; passing a different refdata + will misreport mismatches against an unrelated allele pool. + + Raises ``RuntimeError`` when this result was built without + attached outcomes (e.g. loaded from a TSV); the validator + needs the engine state, not just the projected record. + """ + if self._outcomes is None: + raise RuntimeError( + "validate_records requires the original Outcome objects " + "(with their trace + event ledger), but this " + "SimulationResult was built from records only (e.g. " + "loaded from TSV). Re-run the simulation with " + "Experiment.run_records() to get a result with attached " + "outcomes." + ) + failures: List[Dict[str, Any]] = [] + for i, (outcome, record) in enumerate(zip(self._outcomes, self._records)): + sequence_id = str(record.get("sequence_id", "")) + issues = outcome.validate_record(refdata, sequence_id=sequence_id) + if issues: + failures.append( + { + "record_index": i, + "sequence_id": sequence_id, + "issues": issues, + } + ) + return ValidationReport(count=len(self._outcomes), failures=failures) + + def validate_families( + self, refdata: Any = None + ) -> "FamilyValidationReport": + """**Clonal-family consistency check** — a strict subset of + the audit's §6 family-layer invariants + (``docs/clonal_family_design.md``). + + Groups records by ``clone_id`` and asserts the recombination- + time truth fields agree across every descendant of a clone. + ``refdata`` is reserved for forward compatibility with the + deeper family-layer checks (pre-SHM junction, mutation- + distance distribution) the audit's later slices add; this + slice's invariants are all dict-only and ignore ``refdata``. + + **Currently enforced invariants:** + + - ``truth_v_call`` constant within each ``clone_id``, **when + present**. Skipped silently for clones whose records were + projected without ``expose_provenance=True``. + - ``truth_d_call`` same. + - ``truth_j_call`` same. + - ``clone_id`` is present on every record once any record + in the batch carries one (a batch that mixes clonal and + non-clonal records raises ``CloneIdMissing``). + + **Non-clonal results return ok with ``family_count == 0``.** + This makes ``result.validate_families()`` a safe no-op on a + flat batch — the call site does not need to branch on the + result's clonal-ness. + + **Records-only results work.** Unlike + :meth:`validate_records`, this validator does not require + the underlying ``Outcome`` objects — every check is on + record-dict fields — so a ``SimulationResult`` loaded from + TSV can still be family-validated. + + **Not enforced yet** (deferred per the audit's §14 + out-of-scope list): mutation-distance distribution, pre-SHM + junction invariance, parent-trace reconstruction, lineage + topology, ``original_v_call`` / ``d_inverted`` invariance. + + Returns a :class:`FamilyValidationReport` carrying + ``count``, ``family_count``, ``members_per_family``, and + ``failures``. + """ + del refdata # forward-compat placeholder, see docstring. + + records = self._records + total = len(records) + + # Identify whether any record carries a non-null ``clone_id``. + # A batch that has no clonal records returns an ok no-op + # report; a batch where SOME records carry a clone_id and + # others don't is structurally broken (mixed clonal/flat + # outputs) and surfaces a ``CloneIdMissing`` failure. + any_clonal = any( + "clone_id" in r and r["clone_id"] is not None for r in records + ) + if not any_clonal: + return FamilyValidationReport( + count=total, + family_count=0, + members_per_family={}, + failures=[], + ) + + failures: List[Dict[str, Any]] = [] + missing_clone_id: List[int] = [ + i for i, r in enumerate(records) + if r.get("clone_id") is None + ] + if missing_clone_id: + failures.append( + { + "clone_id": None, + "record_indices": missing_clone_id, + "issue_kind": "CloneIdMissing", + "values": [], + } + ) + + # Group records by clone_id (ignoring records missing the tag + # — they already surfaced above). + by_clone: Dict[Any, List[int]] = {} + for i, r in enumerate(records): + cid = r.get("clone_id") + if cid is None: + continue + by_clone.setdefault(cid, []).append(i) + + for cid, indices in by_clone.items(): + members = [records[i] for i in indices] + for field, kind in _FAMILY_TRUTH_INVARIANTS: + # Skip the field for this group if no record carries + # it (consumer opted out of expose_provenance). Once + # any record in the group carries it, every record + # in the group must — the consistency check still + # runs but the "missing on some" surface would be a + # different bug; we treat ``None`` and absent as the + # same "no opinion" signal here. + present = [ + rec.get(field) for rec in members + if field in rec and rec.get(field) is not None + ] + if not present: + continue + distinct = sorted({str(v) for v in present}) + if len(distinct) > 1: + failures.append( + { + "clone_id": cid, + "record_indices": list(indices), + "issue_kind": kind, + "values": distinct, + } + ) + + members_per_family = {cid: len(idxs) for cid, idxs in by_clone.items()} + return FamilyValidationReport( + count=total, + family_count=len(by_clone), + members_per_family=members_per_family, + failures=failures, + ) + + def validate_families_with_parents( + self, refdata: Any = None + ) -> "FamilyValidationReport": + """**Parent-aware clonal-family validator** — the deeper + diagnostic that compares every descendant against its + actual parent ``Outcome`` (Slice 3 of the clonal-family + audit; see ``docs/clonal_parent_outcome_design.md`` §6). + + Sibling of :meth:`validate_families`. That validator is + record-only (groups by ``clone_id``, compares truth fields + across siblings); this one **requires the parent outcomes + to be available on the result** and compares each + descendant against its parent directly. Use this when you + want to confirm "the descendants reflect the recombination + ancestor they came from," not just "siblings agree with + each other." + + **Currently enforced invariants** — all derived from + record-vs-parent comparison only: + + - Structural: + - ``ParentsMissing`` when records carry ``clone_id`` / + ``parent_id`` but ``self.parents`` is ``None``. + - ``ParentIdMissing`` for records without a non-null + ``parent_id`` in a result that has parents available. + - ``ParentIdOutOfRange`` when ``record["parent_id"]`` + is not in ``[0, len(self.parents))``. + - Truth-allele consistency (requires ``refdata``): + - ``ParentTruthVCallMismatch`` / + ``ParentTruthDCallMismatch`` / + ``ParentTruthJCallMismatch`` — descendant's + ``truth_*_call`` (from ``expose_provenance=True``) + disagrees with the parent's projected truth allele. + - Provenance consistency (no ``refdata`` needed): + - ``ParentDInvertedMismatch`` — descendant's + ``d_inverted`` disagrees with the parent's. D inversion + is a pre-fork decision, so divergence indicates a + structural bug. + - ``ParentOriginalVCallMismatch`` — descendant's + ``original_v_call`` (receptor-revision provenance) + disagrees with the parent's. Same reasoning. + + **Without ``refdata``:** only the structural checks + (``ParentsMissing`` / ``ParentIdMissing`` / + ``ParentIdOutOfRange``) run. All value comparisons — + truth alleles, ``d_inverted``, ``original_v_call`` — + require projecting the parent ``Outcome`` to an AIRR + record, which today goes through the Rust projector and + needs ``refdata``. Slice 3 deliberately stays Python-only; + a lighter-weight refdata-free parent accessor for + ``d_inverted`` etc. is deferred until a Rust slice surfaces + one. + + **Skipped silently for fields not present on descendants:** + if ``expose_provenance`` was off, the descendants don't + carry ``truth_*_call`` and those checks are no-ops. + + **Non-clonal results return ok with ``family_count=0``** — + same safe no-op shape as :meth:`validate_families`. Slice + 3 deliberately does NOT raise "not clonal" here. + + **Not enforced yet** (deferred per audit §6, §14): + + - **Pre-SHM junction invariance.** The descendant's + ``junction`` AIRR field is post-SHM; the pre-SHM junction + lives only inside the parent's IR. A proper check would + require a parent-derived ``junction_pre_shm`` field on + either records or a future ``FamilyRecord`` projection + (Slice 4+). This validator does NOT compare + ``descendant.junction`` against any parent-derived + value today. + - **Mutation-distance distribution.** Comparing the + parent's assembled sequence to each descendant's + post-SHM sequence to verify SHM mass is plausible. + Requires projecting the parent's pool to a sequence + string — out of scope for this slice. + - **Plan-split pre-fork pass enumeration.** "Parent should + not carry descendant-only observation fields like PCR + / paired-end / quality errors" is pinned at the + contract-test level (the pre-fork plan's pass names) + rather than enforced at runtime here. + + **Not wired into ``validate_records=True``.** This is an + explicit deeper diagnostic surface. The + ``validate_records=True`` gate continues to run only the + per-record postcondition validator and the field-only + :meth:`validate_families`. Callers who want parent-aware + checks invoke this method explicitly. + + Returns a :class:`FamilyValidationReport`. Failure dicts + carry ``clone_id``, ``parent_id``, ``record_indices``, + ``issue_kind``, ``parent_value``, and ``child_values`` + (the latter two are ``None`` / ``[]`` for structural + failures that don't compare values). + """ + records = self._records + total = len(records) + + any_clonal = any( + "clone_id" in r and r["clone_id"] is not None for r in records + ) + any_parent_id = any( + "parent_id" in r and r["parent_id"] is not None for r in records + ) + + # Non-clonal: ok no-op (matches validate_families). + if not any_clonal and not any_parent_id: + return FamilyValidationReport( + count=total, + family_count=0, + members_per_family={}, + failures=[], + ) + + failures: List[Dict[str, Any]] = [] + parents = self._parents + + # Parents missing entirely — surface and bail (no comparable + # parent state to run further checks against). + if parents is None: + failures.append( + { + "clone_id": None, + "parent_id": None, + "record_indices": [ + i for i, r in enumerate(records) + if r.get("clone_id") is not None + or r.get("parent_id") is not None + ], + "issue_kind": "ParentsMissing", + "parent_value": None, + "child_values": [], + } + ) + # Compute aggregation by clone_id so the report still + # carries useful structure even though no comparison + # ran. + by_clone: Dict[Any, List[int]] = {} + for i, r in enumerate(records): + cid = r.get("clone_id") + if cid is None: + continue + by_clone.setdefault(cid, []).append(i) + return FamilyValidationReport( + count=total, + family_count=len(by_clone), + members_per_family={ + cid: len(idxs) for cid, idxs in by_clone.items() + }, + failures=failures, + ) + + # Parents are present. Group descendants by parent_id and + # surface structural failures (missing / out of range) per + # record. + n_parents = len(parents) + missing_parent_id: List[int] = [] + out_of_range: Dict[Any, List[int]] = {} + by_parent: Dict[int, List[int]] = {} + for i, r in enumerate(records): + pid = r.get("parent_id") + if pid is None: + # Records without a parent_id in a clonal result are + # structurally broken — surface them and skip per- + # parent invariant checks for those records. + if r.get("clone_id") is not None: + missing_parent_id.append(i) + continue + if not isinstance(pid, int) or isinstance(pid, bool): + out_of_range.setdefault(pid, []).append(i) + continue + if pid < 0 or pid >= n_parents: + out_of_range.setdefault(pid, []).append(i) + continue + by_parent.setdefault(pid, []).append(i) + + if missing_parent_id: + failures.append( + { + "clone_id": None, + "parent_id": None, + "record_indices": missing_parent_id, + "issue_kind": "ParentIdMissing", + "parent_value": None, + "child_values": [], + } + ) + for pid, idxs in out_of_range.items(): + failures.append( + { + "clone_id": None, + "parent_id": pid, + "record_indices": idxs, + "issue_kind": "ParentIdOutOfRange", + "parent_value": pid, + "child_values": [], + } + ) + + # Build parent projections lazily — only for parent ids that + # have descendants pointing at them. The Rust AIRR projector + # ``outcome_to_airr_record`` requires refdata; without it we + # can't materialize the parent's projected fields and + # therefore can't run any value-comparison check. Structural + # checks (missing / out-of-range parent_id) already ran + # above. Document this in the method's contract: refdata- + # less calls run structural checks only. + parent_projections: Dict[int, Dict[str, Any]] = {} + if by_parent and refdata is not None: + from ._airr_record import outcome_to_airr_record + + for pid in by_parent: + proj = outcome_to_airr_record( + parents[pid], + refdata, + sequence_id=f"parent{pid}", + ) + _inject_truth_columns(parents[pid], refdata, proj) + parent_projections[pid] = proj + + # The (field, issue_kind) pairs the parent-aware validator + # enforces when ``refdata`` is available. All of these are + # comparisons against parent-projected values; refdata is + # required to materialize them. + FIELD_CHECKS = [ + ("truth_v_call", "ParentTruthVCallMismatch"), + ("truth_d_call", "ParentTruthDCallMismatch"), + ("truth_j_call", "ParentTruthJCallMismatch"), + ("d_inverted", "ParentDInvertedMismatch"), + ("original_v_call", "ParentOriginalVCallMismatch"), + ] + + for pid, indices in by_parent.items(): + parent_proj = parent_projections.get(pid) + if parent_proj is None: + # Refdata was None — value comparisons skipped. + continue + members = [records[i] for i in indices] + # Use the first member's clone_id for the failure + # group; per-clone_id divergence within a parent_id + # group would have surfaced via the family validator + # already. + clone_id_for_group = members[0].get("clone_id", pid) + for field, kind in FIELD_CHECKS: + parent_val = parent_proj.get(field) + # Skip when parent has no opinion on the field + # (e.g. ``original_v_call`` is empty when receptor + # revision didn't run; comparing "" to "" would + # always pass but comparing "" to a real allele + # name would surface a real bug — so we only skip + # when the field is genuinely missing). + if parent_val is None: + continue + divergent_indices: List[int] = [] + divergent_values: List[Any] = [] + for idx, rec in zip(indices, members): + if field not in rec: + continue + child_val = rec[field] + if child_val is None: + continue + if child_val != parent_val: + divergent_indices.append(idx) + if child_val not in divergent_values: + divergent_values.append(child_val) + if divergent_indices: + failures.append( + { + "clone_id": clone_id_for_group, + "parent_id": pid, + "record_indices": divergent_indices, + "issue_kind": kind, + "parent_value": parent_val, + "child_values": sorted( + divergent_values, key=lambda v: str(v) + ), + } + ) + + return FamilyValidationReport( + count=total, + family_count=len(by_parent), + members_per_family={ + pid: len(idxs) for pid, idxs in by_parent.items() + }, + failures=failures, + ) diff --git a/src/GenAIRR/_step_validation.py b/src/GenAIRR/_step_validation.py new file mode 100644 index 0000000..a12e9a5 --- /dev/null +++ b/src/GenAIRR/_step_validation.py @@ -0,0 +1,304 @@ +"""Step-parameter validation helpers for :class:`GenAIRR.experiment.Experiment`. + +Extracted verbatim from ``experiment.py`` (behavior-preserving) — pure +leaf validation of ``segment_rates`` / ``v_subregion_rates`` kwargs and the +clonal descendant-phase step classifier, with no dependency on the +``Experiment`` class itself. +""" +from typing import Dict, List, Optional, Tuple + + +# ────────────────────────────────────────────────────────────────── +# Per-segment SHM rate validation (slice: segment_rates kwarg on +# Experiment.mutate). See ``docs/shm_segment_rate_design.md``. +# ────────────────────────────────────────────────────────────────── + +# Canonical bucket order — matches the Rust ``SegmentRateWeights`` +# struct field order so positional plumbing through PyO3 stays +# straightforward. +_SEGMENT_RATE_BUCKETS: Tuple[str, ...] = ("V", "D", "J", "NP") + +# Default flat-substrate rate vector (1.0 for every bucket). The +# pipeline-IR ``_MutateStep`` carries this verbatim when the user +# omits ``segment_rates``; the Rust passes detect the flat-default +# case and take the existing (pre-slice) fast path so legacy +# pipelines stay byte-identical. +_DEFAULT_SEGMENT_RATES: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0) + + +def _validate_segment_rates( + segment_rates: Optional[Dict[str, float]], +) -> Tuple[float, float, float, float]: + """Validate the user's ``segment_rates`` dict and return the + normalised ``(v, d, j, np)`` tuple. Pure helper — no DSL state. + + Validation: + + - ``None`` (or omitted) → flat default ``(1.0, 1.0, 1.0, 1.0)``. + - Keys must be a subset of ``{"V", "D", "J", "NP"}`` (case- + sensitive — matches the DSL spec). Other keys raise + ``ValueError`` naming the offending key. + - Values must be ``int`` / ``float`` (not bool), finite, and + ``>= 0``. Negative / NaN / inf raise ``ValueError``. + - Sparse: omitted keys default to ``1.0``. + - At least one effective rate must be strictly positive — an + all-zero (or all-omitted-then-explicitly-zero) configuration + would make the SHM pass a deterministic no-op, which is + almost certainly a builder bug. Reject with ``ValueError``. + """ + if segment_rates is None: + return _DEFAULT_SEGMENT_RATES + + if not isinstance(segment_rates, dict): + raise TypeError( + f"segment_rates must be a dict or None, got " + f"{type(segment_rates).__name__}" + ) + + # Reject unknown keys first so a typo surfaces with a clear + # message instead of silently defaulting. + unknown = sorted(set(segment_rates.keys()) - set(_SEGMENT_RATE_BUCKETS)) + if unknown: + raise ValueError( + f"segment_rates: unknown segment key(s) {unknown!r}. " + f"Allowed: {list(_SEGMENT_RATE_BUCKETS)!r}. " + "Keys are case-sensitive; 'V' / 'D' / 'J' for the V/D/J " + "segments and 'NP' for both Np1 and Np2." + ) + + out_list: List[float] = [] + for bucket in _SEGMENT_RATE_BUCKETS: + if bucket not in segment_rates: + out_list.append(1.0) + continue + value = segment_rates[bucket] + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError( + f"segment_rates[{bucket!r}]: value must be a finite " + f"non-negative number, got {type(value).__name__}" + ) + value_f = float(value) + # NaN check FIRST — every comparison with NaN is False, so + # a `value_f < 0` test wouldn't catch it; that's the same + # rationale as the ``invert_d`` / ``rate`` NaN handling. + if value_f != value_f: + raise ValueError( + f"segment_rates[{bucket!r}]: value must be a finite " + "number, got NaN" + ) + if value_f < 0.0 or value_f == float("inf") or value_f == float("-inf"): + raise ValueError( + f"segment_rates[{bucket!r}]: value must be a finite " + f"non-negative number, got {value_f}" + ) + out_list.append(value_f) + + if sum(out_list) <= 0.0: + raise ValueError( + "segment_rates: at least one bucket must have a positive " + "rate. The supplied configuration zeroes out every " + "biological segment, which would make the SHM pass a " + "deterministic no-op." + ) + + return (out_list[0], out_list[1], out_list[2], out_list[3]) + + +# ────────────────────────────────────────────────────────────────── +# Per-V-subregion SHM rate validation (Slice B — +# `v_subregion_rates` kwarg on Experiment.mutate). See +# ``docs/v_subregion_shm_rate_design.md``. +# ────────────────────────────────────────────────────────────────── + +# Canonical label order — matches the Rust ``VSubregionRateWeights`` +# struct field order (FWR1 / CDR1 / FWR2 / CDR2 / FWR3) so +# positional plumbing through PyO3 stays straightforward. +_V_SUBREGION_RATE_LABELS: Tuple[str, ...] = ("FWR1", "CDR1", "FWR2", "CDR2", "FWR3") + +# Two-letter aliases that expand to a group of canonical labels. +# Resolution rule (audit §3): aliases expand first, then explicit +# labels override. So ``{"FWR": 0.5, "FWR2": 2.0}`` resolves to +# ``{FWR1: 0.5, FWR2: 2.0, FWR3: 0.5, CDR1: 1.0, CDR2: 1.0}``. +_V_SUBREGION_RATE_ALIASES: Dict[str, Tuple[str, ...]] = { + "FWR": ("FWR1", "FWR2", "FWR3"), + "CDR": ("CDR1", "CDR2"), +} + +# Default flat-substrate rate vector (1.0 for every label) — same +# fast-path discipline as ``_DEFAULT_SEGMENT_RATES``. +_DEFAULT_V_SUBREGION_RATES: Tuple[float, float, float, float, float] = ( + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, +) + + +def _validate_v_subregion_rates( + v_subregion_rates: Optional[Dict[str, float]], +) -> Tuple[float, float, float, float, float]: + """Validate the user's ``v_subregion_rates`` dict and return + the normalised ``(FWR1, CDR1, FWR2, CDR2, FWR3)`` tuple. + Pure helper — no DSL state. + + Validation rules: + + - ``None`` (or omitted) → flat default + ``(1.0, 1.0, 1.0, 1.0, 1.0)``. + - Empty dict ``{}`` is treated as ``None`` — same flat default. + - Keys must be a subset of the five canonical labels + (``FWR1`` / ``CDR1`` / ``FWR2`` / ``CDR2`` / ``FWR3``) plus + the two aliases ``FWR`` (expands to FWR1 / FWR2 / FWR3) and + ``CDR`` (expands to CDR1 / CDR2). Case-sensitive — matches + the V-subregion annotation surface (Slice 1). + - Values must be ``int`` / ``float`` (not bool), finite, and + ``>= 0``. NaN / inf / bool / negative raise ``ValueError``. + - Sparse: omitted labels default to ``1.0``. + - Alias expansion happens first, then explicit labels + override. So ``{"FWR": 0.5, "FWR2": 2.0}`` → ``FWR1=0.5, + FWR2=2.0, FWR3=0.5, CDR1=1.0, CDR2=1.0``. + - After expansion, at least one label must be strictly + positive — an all-zero vector would zero every V site and + is almost certainly a builder bug. Reject with + ``ValueError``. + """ + if v_subregion_rates is None: + return _DEFAULT_V_SUBREGION_RATES + + if not isinstance(v_subregion_rates, dict): + raise TypeError( + f"v_subregion_rates must be a dict or None, got " + f"{type(v_subregion_rates).__name__}" + ) + + if not v_subregion_rates: + # Empty dict is equivalent to omitting the kwarg. + return _DEFAULT_V_SUBREGION_RATES + + accepted_keys = set(_V_SUBREGION_RATE_LABELS) | set(_V_SUBREGION_RATE_ALIASES) + unknown = sorted(set(v_subregion_rates.keys()) - accepted_keys) + if unknown: + raise ValueError( + f"v_subregion_rates: unknown label(s) {unknown!r}. " + f"Allowed: {list(_V_SUBREGION_RATE_LABELS)!r} plus the " + f"aliases {list(_V_SUBREGION_RATE_ALIASES.keys())!r}. " + "Labels are case-sensitive; CDR3 / FWR4 are out of " + "scope (CDR3 lives in the junction, FWR4 in the J " + "segment) — use ``segment_rates`` for those." + ) + + def _coerce(label_for_error: str, value: object) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError( + f"v_subregion_rates[{label_for_error!r}]: value must be a " + f"finite non-negative number, got {type(value).__name__}" + ) + v = float(value) + if v != v: # NaN + raise ValueError( + f"v_subregion_rates[{label_for_error!r}]: value must be a " + "finite number, got NaN" + ) + if v < 0.0 or v == float("inf") or v == float("-inf"): + raise ValueError( + f"v_subregion_rates[{label_for_error!r}]: value must be a " + f"finite non-negative number, got {v}" + ) + return v + + # Phase 1: alias expansion. Aliases populate their labels first; + # phase 2's explicit labels then override. + resolved: Dict[str, float] = {} + for alias, expanded_labels in _V_SUBREGION_RATE_ALIASES.items(): + if alias in v_subregion_rates: + v = _coerce(alias, v_subregion_rates[alias]) + for lbl in expanded_labels: + resolved[lbl] = v + + # Phase 2: explicit labels override. + for lbl in _V_SUBREGION_RATE_LABELS: + if lbl in v_subregion_rates: + resolved[lbl] = _coerce(lbl, v_subregion_rates[lbl]) + + # Fill any remaining label with the flat default. + out_list: List[float] = [] + for lbl in _V_SUBREGION_RATE_LABELS: + out_list.append(resolved.get(lbl, 1.0)) + + if sum(out_list) <= 0.0: + raise ValueError( + "v_subregion_rates: at least one label must have a " + "positive rate after alias expansion. The supplied " + "configuration zeroes every V subregion, which would " + "drop every V site out of SHM support." + ) + + return (out_list[0], out_list[1], out_list[2], out_list[3], out_list[4]) + + +# ────────────────────────────────────────────────────────────────── +# Clonal ordering-guard table (Slice: descendant-phase guards) +# ────────────────────────────────────────────────────────────────── +# +# Every DSL step in this table is **descendant-phase** — it models +# observation / library-prep / sequencing biology that must be +# sampled independently per clone member. Pre-fork placement either +# silently misreports the AIRR field (Bugs C / E / F: trace-sourced +# fields that don't survive the parent→descendant boundary) or +# collapses descendant diversity (every clone member shares an +# identical effect because the pass ran once on the parent IR). +# +# The clonal fork methods scan the already-appended step list against +# this table; the first match is rejected with a +# message naming the offending DSL method and the canonical fix. +# +# Each entry is ``(predicate, dsl_method_name)`` where the predicate +# inspects one step and returns ``True`` if it came from +# ``dsl_method_name``. Some DSL methods append :class:`_CorruptStep` +# with different ``kind`` discriminators; others append distinct +# step types. +def _descendant_phase_step_classifier(step): + """Return the DSL method name a descendant-phase ``step`` came + from, or ``None`` if ``step`` is not a descendant-phase step. + + Single source of truth for the unified guard in the flat clonal + fork methods. Adding a new descendant-phase DSL method means + appending a clause here (and adding the + companion spec test in + ``tests/test_clonal_descendant_phase_guards.py``). + """ + from ._pipeline_ir import ( + _CORRUPT_KIND_3PRIME_LOSS, + _CORRUPT_KIND_5PRIME_LOSS, + _CORRUPT_KIND_INDEL, + _CORRUPT_KIND_NS, + _CORRUPT_KIND_PCR, + _CORRUPT_KIND_QUALITY, + _CORRUPT_KIND_REV_COMP, + _CorruptStep, + _MutateStep, + _PairedEndStep, + ) + + if isinstance(step, _MutateStep): + return "mutate" + if isinstance(step, _PairedEndStep): + return "paired_end" + if isinstance(step, _CorruptStep): + # Per-kind classification — only the descendant-phase kinds + # appear in the table. ``contaminant`` is deliberately + # omitted: the DSL slice that added the descendant-phase + # ordering guards did not list it, so leave the placement + # unconstrained until a follow-up explicitly classifies it. + kind_to_method = { + _CORRUPT_KIND_PCR: "pcr_amplify", + _CORRUPT_KIND_QUALITY: "sequencing_errors", + _CORRUPT_KIND_INDEL: "polymerase_indels", + _CORRUPT_KIND_5PRIME_LOSS: "end_loss_5prime", + _CORRUPT_KIND_3PRIME_LOSS: "end_loss_3prime", + _CORRUPT_KIND_NS: "ambiguous_base_calls", + _CORRUPT_KIND_REV_COMP: "random_strand_orientation", + } + return kind_to_method.get(step.kind) + return None diff --git a/src/GenAIRR/cartridge_builder.py b/src/GenAIRR/cartridge_builder.py index 460da76..249ad41 100644 --- a/src/GenAIRR/cartridge_builder.py +++ b/src/GenAIRR/cartridge_builder.py @@ -92,21 +92,12 @@ class _BuilderDAllele(DAllele): from .dataconfig.enums import ChainType, Species from .genotype_priors import PopulationGenotypeModel from .reference_models import ( - AlleleUsageSpec, - EmpiricalDistributionSpec, - NP_KEYS, - NpBaseModelSpec, - P_NUCLEOTIDE_END_KEYS, - P_NUCLEOTIDE_END_KEYS_VJ, ReferenceEmpiricalModels, - TRIM_KEYS, - TRIM_KEYS_VJ, ) - -_NP_CANONICAL_BASES = ("A", "C", "G", "T") from .reference_rules import ReferenceRulesSpec from .utilities.imgt_regions import compute_v_region_boundaries from .utilities.misc import parse_fasta +from ._cartridge_estimators import _CartridgeEstimators # Accepted input shapes for FASTA inputs: a filesystem path, a path- @@ -194,68 +185,6 @@ def _open_fasta(source: FastaInput): ) -def _split_tie_set(raw: str) -> List[str]: - """Parse a comma-separated AIRR tie-set call into a clean list. - - Empty / whitespace-only entries drop out. The returned list - preserves insertion order — important for the - ``ambiguous="truth_first"`` policy which keeps only the first - entry.""" - return [token.strip() for token in raw.split(",") if token.strip()] - - -def _allele_names_for_segment( - buckets: Dict[str, List[Any]], -) -> "set[str]": - """Flatten a builder's per-gene allele bucket dict into a set - of canonical allele names. Used by - :meth:`ReferenceCartridgeBuilder.estimate_allele_usage` to - decide which AIRR-record allele names are "known" to the - cartridge under construction.""" - out: "set[str]" = set() - for allele_list in buckets.values(): - for allele in allele_list: - name = getattr(allele, "name", None) - if isinstance(name, str) and name: - out.add(name) - return out - - -def _load_rearrangements( - source: Any, -) -> "tuple[List[Dict[str, Any]], str]": - """Normalise a ``rearrangements`` argument into a - ``(records, source_label)`` tuple. - - Accepted inputs: - - - ``list[dict]`` — used verbatim; source label - ``"records:N"`` carrying the row count. - - path-like (``str`` / ``Path``) — parsed via - :class:`csv.DictReader` with ``delimiter='\\t'`` (AIRR-C TSV - convention); source label is the path. - - open text file handle — parsed via :class:`csv.DictReader`; - source label is ``"file:N"``. - - Raises :class:`TypeError` for unrecognised shapes.""" - import csv - - if isinstance(source, list): - return list(source), f"records:{len(source)}" - if hasattr(source, "read"): - rows = list(csv.DictReader(source, delimiter="\t")) - return rows, f"file:{len(rows)}" - if isinstance(source, (str, Path)): - path = Path(source) - with open(path, "r", newline="") as fh: - rows = list(csv.DictReader(fh, delimiter="\t")) - return rows, str(path) - raise TypeError( - f"rearrangements must be a list[dict], path, or open text " - f"file, got {type(source).__name__}" - ) - - def _parse_allele_name(raw_header: str) -> str: """Pull the allele name from a FASTA header. @@ -313,7 +242,7 @@ def to_dict(self) -> Dict[str, Any]: return asdict(self) -class ReferenceCartridgeBuilder: +class ReferenceCartridgeBuilder(_CartridgeEstimators): """Stage-based builder for a :class:`~GenAIRR.DataConfig`. See :mod:`GenAIRR.cartridge_builder` module docs and @@ -791,1504 +720,6 @@ def estimate_genotype_priors( genotypes, cfg=self._build_working_cfg(), **kwargs) return self.set_genotype_priors(model) - def estimate_allele_usage( - self, - rearrangements: Any, - *, - min_count: float = 1.0, - ambiguous: str = "fractional", - replace: bool = True, - ) -> "ReferenceCartridgeBuilder": - """Estimate per-segment allele-usage weights from observed - AIRR rearrangement records. - - ``rearrangements`` accepts: - - - a list of dicts (each row is one AIRR record), - - a path-like (filesystem path to an AIRR TSV — parsed via - ``csv.DictReader`` with tab delimiter), or - - an open text file handle pointing at AIRR TSV. - - ``ambiguous`` selects the tie-set policy (column values - like ``"IGHV1*01,IGHV2*01"``): - - - ``"fractional"`` (default): split one row's credit - ``1.0`` evenly across all known alleles in the tie set. - Unknown allele names in the tie set are excluded; if - NO names in the tie set are known to the cartridge, the - row is recorded as ``unknown_allele`` and skipped. - - ``"truth_first"``: credit only the first comma-separated - allele in the tie set. Matches the existing - :func:`GenAIRR._mcp_summary` convention. - - ``"reject"``: drop ambiguous (multi-call) rows entirely - and record them in ``report.rejected``. - - ``min_count`` drops alleles whose final per-segment count - is strictly below the threshold. Pre-normalisation; the - per-segment weights remaining after the drop are then - renormalised to sum to ``1.0``. - - ``replace`` (default ``True``) controls idempotency. When - ``True``, calling :meth:`estimate_allele_usage` twice - overwrites the previous spec and writes ``replaced=True`` - on the new stage entry. When ``False``, the second call - raises :class:`ValueError`. - - Updates ``self._reference_models.allele_usage`` so a - downstream :meth:`build` carries the estimated spec into - the cartridge. - """ - if ambiguous not in ("fractional", "truth_first", "reject"): - raise ValueError( - f"ambiguous must be one of 'fractional' / 'truth_first' / " - f"'reject', got {ambiguous!r}" - ) - previously_estimated = any( - entry.get("stage") == "estimate_allele_usage" - for entry in self._report.stages - ) - if previously_estimated and not replace: - raise ValueError( - "estimate_allele_usage already ran; pass replace=True to " - "overwrite the previous spec" - ) - - records, source_label = _load_rearrangements(rearrangements) - chain_has_d = self._chain_type.has_d - v_pool = _allele_names_for_segment(self._v_alleles) - d_pool = _allele_names_for_segment(self._d_alleles) - j_pool = _allele_names_for_segment(self._j_alleles) - - v_counts: Dict[str, float] = {} - d_counts: Dict[str, float] = {} - j_counts: Dict[str, float] = {} - skipped = { - "missing_required_column": 0, - "unknown_allele": {"V": 0, "D": 0, "J": 0}, - "missing_d_call_on_vdj": 0, - "ambiguous_rejected": 0, - } - rejected_entries: List[Dict[str, Any]] = [] - warnings: List[str] = [] - d_on_vj_warned = False - - for row_idx, row in enumerate(records): - v_raw = (row.get("v_call") or "").strip() - j_raw = (row.get("j_call") or "").strip() - d_raw = (row.get("d_call") or "").strip() - - # Required column check. - if not v_raw or not j_raw: - skipped["missing_required_column"] += 1 - rejected_entries.append( - { - "stage": "estimate_allele_usage", - "row_index": row_idx, - "reason": "missing_required_column", - } - ) - continue - - # VDJ requires d_call; VJ ignores any d_call. - if chain_has_d and not d_raw: - skipped["missing_d_call_on_vdj"] += 1 - rejected_entries.append( - { - "stage": "estimate_allele_usage", - "row_index": row_idx, - "reason": "missing_d_call_on_vdj", - } - ) - continue - if not chain_has_d and d_raw: - if not d_on_vj_warned: - warnings.append( - "d_call column present on a VJ cartridge — D " - "contribution ignored for every row" - ) - d_on_vj_warned = True - d_raw = "" # ignore D contribution silently after warning - - v_tie = _split_tie_set(v_raw) - j_tie = _split_tie_set(j_raw) - d_tie = _split_tie_set(d_raw) if d_raw else [] - - if ambiguous == "reject": - if len(v_tie) > 1 or len(j_tie) > 1 or (chain_has_d and len(d_tie) > 1): - skipped["ambiguous_rejected"] += 1 - rejected_entries.append( - { - "stage": "estimate_allele_usage", - "row_index": row_idx, - "reason": "ambiguous_rejected", - } - ) - continue - - # truth_first collapses the tie set to its first entry. - if ambiguous == "truth_first": - v_tie = v_tie[:1] - j_tie = j_tie[:1] - d_tie = d_tie[:1] if d_tie else [] - - # Resolve known alleles in each tie set. - v_known = [n for n in v_tie if n in v_pool] - j_known = [n for n in j_tie if n in j_pool] - d_known = ( - [n for n in d_tie if n in d_pool] if d_tie else [] - ) - - # Unknown-allele bookkeeping. We record one rejection - # entry per unknown allele per segment per row so the - # report names the actual unknown names. - for n in v_tie: - if n not in v_pool: - skipped["unknown_allele"]["V"] += 1 - rejected_entries.append( - { - "stage": "estimate_allele_usage", - "row_index": row_idx, - "segment": "V", - "allele_name": n, - "reason": "unknown_allele", - } - ) - for n in j_tie: - if n not in j_pool: - skipped["unknown_allele"]["J"] += 1 - rejected_entries.append( - { - "stage": "estimate_allele_usage", - "row_index": row_idx, - "segment": "J", - "allele_name": n, - "reason": "unknown_allele", - } - ) - if d_tie: - for n in d_tie: - if n not in d_pool: - skipped["unknown_allele"]["D"] += 1 - rejected_entries.append( - { - "stage": "estimate_allele_usage", - "row_index": row_idx, - "segment": "D", - "allele_name": n, - "reason": "unknown_allele", - } - ) - - # Fractional credit (the default and the truth_first - # collapsed path; both use the same accumulation - # logic now that truth_first reduced the tie set). - if v_known: - share = 1.0 / len(v_known) - for n in v_known: - v_counts[n] = v_counts.get(n, 0.0) + share - if j_known: - share = 1.0 / len(j_known) - for n in j_known: - j_counts[n] = j_counts.get(n, 0.0) + share - if d_known: - share = 1.0 / len(d_known) - for n in d_known: - d_counts[n] = d_counts.get(n, 0.0) + share - - # min_count filter + normalisation per segment. - below_min: Dict[str, int] = {"V": 0, "D": 0, "J": 0} - - def _filter_and_normalise( - counts: Dict[str, float], segment_label: str - ) -> Dict[str, float]: - kept = {n: c for n, c in counts.items() if c >= min_count} - dropped = len(counts) - len(kept) - below_min[segment_label] = dropped - total = sum(kept.values()) - if total <= 0.0: - return {} - return {n: w / total for n, w in kept.items()} - - v_weights = _filter_and_normalise(v_counts, "V") - d_weights = _filter_and_normalise(d_counts, "D") - j_weights = _filter_and_normalise(j_counts, "J") - - if below_min["V"] or below_min["D"] or below_min["J"]: - warnings.append( - f"alleles below min_count={min_count} dropped — " - f"V={below_min['V']}, D={below_min['D']}, J={below_min['J']}" - ) - - spec = AlleleUsageSpec(v=v_weights, d=d_weights, j=j_weights) - chain_label = "vdj" if chain_has_d else "vj" - spec.validate(chain_type=chain_label, name="allele_usage") - - # Attach to the existing reference_models (creating it if - # absent). The cartridge built downstream carries the - # spec automatically. - if self._reference_models is None: - self._reference_models = ReferenceEmpiricalModels() - self._reference_models = ReferenceEmpiricalModels( - np_lengths=self._reference_models.np_lengths, - trims=self._reference_models.trims, - np_bases=self._reference_models.np_bases, - p_nucleotide_lengths=self._reference_models.p_nucleotide_lengths, - allele_usage=spec, - ) - - self._report.stages.append( - { - "stage": "estimate_allele_usage", - "inputs": { - "record_count": len(records), - "ambiguous": ambiguous, - "min_count": float(min_count), - "source": source_label, - "replaced": previously_estimated, - }, - "inferred": { - "V": v_weights, - "D": d_weights, - "J": j_weights, - "skipped": skipped, - "below_min_count": below_min, - }, - "warnings": warnings, - } - ) - self._report.rejected.extend(rejected_entries) - return self - - def estimate_trim_distributions( - self, - rearrangements: Any, - *, - min_count: int = 1, - pseudocount: float = 0.0, - replace: bool = True, - ) -> "ReferenceCartridgeBuilder": - """Estimate per-key trim distributions from observed - AIRR rearrangement records. - - Writes ``EmpiricalDistributionSpec`` instances into - ``self._reference_models.trims`` keyed by ``"V_3"`` / - ``"D_5"`` / ``"D_3"`` / ``"J_5"`` (see - :data:`GenAIRR.reference_models.TRIM_KEYS`). VJ cartridges - only produce the ``V_3`` and ``J_5`` keys (see - :data:`GenAIRR.reference_models.TRIM_KEYS_VJ`); the D-end - keys are skipped silently because the cartridge has no D - segment to trim. - - ``rearrangements`` accepts the same shapes as - :meth:`estimate_allele_usage`: - - - a list of dicts (each row is one AIRR record), - - a path-like to an AIRR TSV (parsed via ``csv.DictReader`` - with tab delimiter), or - - an open text file handle pointing at AIRR TSV. - - AIRR columns consumed (per audit §1.3): - - - **VJ:** ``v_trim_3``, ``j_trim_5``. - - **VDJ:** ``v_trim_3``, ``d_trim_5``, ``d_trim_3``, ``j_trim_5``. - - Columns deliberately ignored: ``v_trim_5`` / ``j_trim_3`` - (no engine pass — hard-zero in projection), and the - observation-stage ``end_loss_5_length`` / ``end_loss_3_length`` - (separate corruption surface — see - ``docs/primer_trim_end_loss_audit.md``). - - Per-row validation is **field-local**: a malformed or - missing value in one trim column drops that field's - contribution only — the row still feeds its other - well-formed fields. Per-row structured entries land in - ``report.rejected`` with reasons - ``"missing_required_column"`` / ``"malformed_trim_value"`` - / ``"negative_trim_value"``, each carrying the AIRR - column name. - - ``min_count`` (int, default 1) drops trim values whose - observed integer count is strictly below the threshold - before normalisation. - - ``pseudocount`` (float, default 0.0) adds a uniform prior - to every **observed** trim value before normalisation - (no support expansion; values never observed stay - unobserved). Applied AFTER the ``min_count`` filter so - the filter looks at raw observations. - - ``replace`` (default ``True``) controls idempotency. When - ``True``, calling this method twice overwrites the - previous specs and writes ``replaced=True`` on the new - stage entry. When ``False`` AND a prior typed-plane - ``trims`` is already attached to ``self._reference_models``, - the call raises :class:`ValueError` before consuming any - records. - """ - if isinstance(min_count, bool) or not isinstance(min_count, int): - raise ValueError( - f"min_count must be an int, got {min_count!r}" - ) - if min_count < 0: - raise ValueError( - f"min_count must be non-negative, got {min_count!r}" - ) - if isinstance(pseudocount, bool) or not isinstance( - pseudocount, (int, float) - ): - raise ValueError( - f"pseudocount must be a non-negative number, got {pseudocount!r}" - ) - if pseudocount < 0.0: - raise ValueError( - f"pseudocount must be non-negative, got {pseudocount!r}" - ) - - if not replace: - existing_trims = ( - self._reference_models.trims - if self._reference_models is not None - else None - ) - if existing_trims: - raise ValueError( - "trim distributions already attached to this cartridge; " - "pass replace=True to overwrite" - ) - previously_estimated = any( - entry.get("stage") == "estimate_trim_distributions" - for entry in self._report.stages - ) - - records, source_label = _load_rearrangements(rearrangements) - chain_has_d = self._chain_type.has_d - active_keys = TRIM_KEYS if chain_has_d else TRIM_KEYS_VJ - - # Map plane key → AIRR column. - key_to_column = { - "V_3": "v_trim_3", - "D_5": "d_trim_5", - "D_3": "d_trim_3", - "J_5": "j_trim_5", - } - # AIRR columns the estimator deliberately does NOT consume. - ignored_columns = ("v_trim_5", "j_trim_3") - - counters: Dict[str, Dict[int, int]] = {k: {} for k in active_keys} - skipped = { - "missing_required_column": {col: 0 for col in - (key_to_column[k] for k in active_keys)}, - "malformed_trim_value": {col: 0 for col in - (key_to_column[k] for k in active_keys)}, - "negative_trim_value": {col: 0 for col in - (key_to_column[k] for k in active_keys)}, - } - dropped_columns: Dict[str, int] = {col: 0 for col in ignored_columns} - rejected_entries: List[Dict[str, Any]] = [] - warnings: List[str] = [] - ignored_warned = {col: False for col in ignored_columns} - vj_d_warned = False - - for row_idx, row in enumerate(records): - # Surface non-zero `v_trim_5` / `j_trim_3` columns: track - # the count but never consume them. One warning per - # column across the dataset. - for col in ignored_columns: - raw = row.get(col) - if raw not in (None, "", "0"): - try: - if int(str(raw).strip()) != 0: - dropped_columns[col] += 1 - if not ignored_warned[col]: - warnings.append( - f"{col} column non-zero in input — " - f"contribution dropped (no V_5 / J_3 " - f"trim pass in the engine)" - ) - ignored_warned[col] = True - except (TypeError, ValueError): - # Malformed in an unused column is uninteresting. - pass - - # Warn once per dataset if a VJ cartridge sees populated - # D-trim columns in the input. Same boundary as - # `estimate_allele_usage`'s D-call ignore. - if not chain_has_d: - for col in ("d_trim_5", "d_trim_3"): - raw = row.get(col) - if raw not in (None, "", "0") and not vj_d_warned: - try: - if int(str(raw).strip()) != 0: - warnings.append( - "VJ cartridge: d_trim_5 / d_trim_3 " - "columns present in records; " - "contribution ignored" - ) - vj_d_warned = True - break - except (TypeError, ValueError): - pass - - # Field-local validation: each key/column independently. - for key in active_keys: - col = key_to_column[key] - raw = row.get(col) - if raw is None or str(raw).strip() == "": - skipped["missing_required_column"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_trim_distributions", - "row_index": row_idx, - "column": col, - "key": key, - "reason": "missing_required_column", - } - ) - continue - try: - value = int(str(raw).strip()) - except (TypeError, ValueError): - skipped["malformed_trim_value"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_trim_distributions", - "row_index": row_idx, - "column": col, - "key": key, - "value": raw, - "reason": "malformed_trim_value", - } - ) - continue - if value < 0: - skipped["negative_trim_value"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_trim_distributions", - "row_index": row_idx, - "column": col, - "key": key, - "value": value, - "reason": "negative_trim_value", - } - ) - continue - counters[key][value] = counters[key].get(value, 0) + 1 - - # min_count + pseudocount + per-key normalisation. - below_min: Dict[str, int] = {k: 0 for k in active_keys} - inferred_pairs: Dict[str, List[tuple]] = {k: [] for k in active_keys} - - for key in active_keys: - raw_counts = counters[key] - # Drop values strictly below min_count. - kept = {v: c for v, c in raw_counts.items() if c >= min_count} - below_min[key] = len(raw_counts) - len(kept) - # Apply pseudocount to observed values only. - if pseudocount > 0.0: - kept = {v: (c + pseudocount) for v, c in kept.items()} - # Normalise to sum 1.0 if anything survived. - total = float(sum(kept.values())) - if total <= 0.0 or not kept: - inferred_pairs[key] = [] - continue - normalised = [(v, kept[v] / total) for v in sorted(kept.keys())] - inferred_pairs[key] = normalised - - if any(below_min[k] for k in active_keys): - warnings.append( - f"trim values below min_count={min_count} dropped — " - + ", ".join(f"{k}={below_min[k]}" for k in active_keys) - ) - - # Build the per-key EmpiricalDistributionSpec instances. - new_trims: Dict[str, EmpiricalDistributionSpec] = {} - for key, pairs in inferred_pairs.items(): - if not pairs: - continue - spec = EmpiricalDistributionSpec(pairs) - spec.validate(name=f"trims[{key}]") - new_trims[key] = spec - - # Attach to the existing reference_models (creating it if - # absent). Other typed planes are preserved. - if self._reference_models is None: - self._reference_models = ReferenceEmpiricalModels() - self._reference_models = ReferenceEmpiricalModels( - np_lengths=self._reference_models.np_lengths, - trims=new_trims, - np_bases=self._reference_models.np_bases, - p_nucleotide_lengths=self._reference_models.p_nucleotide_lengths, - allele_usage=self._reference_models.allele_usage, - ) - # Validate the full container under the cartridge's chain - # type so D-on-VJ etc. raise at attach time. - chain_label = "vdj" if chain_has_d else "vj" - self._reference_models.validate(chain_type=chain_label) - - self._report.stages.append( - { - "stage": "estimate_trim_distributions", - "inputs": { - "record_count": len(records), - "min_count": int(min_count), - "pseudocount": float(pseudocount), - "source": source_label, - "replaced": previously_estimated, - }, - "inferred": { - "V_3": inferred_pairs.get("V_3", []), - "D_5": inferred_pairs.get("D_5", []), - "D_3": inferred_pairs.get("D_3", []), - "J_5": inferred_pairs.get("J_5", []), - "skipped": skipped, - "below_min_count": {k: below_min.get(k, 0) for k in TRIM_KEYS}, - "dropped_columns": dropped_columns, - }, - "warnings": warnings, - } - ) - self._report.rejected.extend(rejected_entries) - return self - - def estimate_np_length_distributions( - self, - rearrangements: Any, - *, - min_count: int = 1, - pseudocount: float = 0.0, - replace: bool = True, - ) -> "ReferenceCartridgeBuilder": - """Estimate per-key NP length distributions from - observed AIRR rearrangement records. - - Writes ``EmpiricalDistributionSpec`` instances into - ``self._reference_models.np_lengths`` keyed by - ``"NP1"`` and (on VDJ cartridges) ``"NP2"`` (see - :data:`GenAIRR.reference_models.NP_KEYS`). VJ - cartridges produce only the ``NP1`` key; the - cartridge has no NP2 region. The boundary is - enforced at estimation time because the typed-plane - validator does NOT chain-type-reject ``NP2`` on - VJ at attach time (see audit §2.2 of - ``docs/np_length_estimation_design.md``). - - AIRR columns consumed (per audit §1.2): - - - ``np1_length`` (always). - - ``np2_length`` (VDJ only — VJ rows with a - non-zero ``np2_length`` raise a one-time warning - and the contribution is skipped). - - Columns deliberately ignored: ``np1`` / ``np2`` - (sequence-derived length is sensitive to - post-claim reabsorption — see audit §5.3), - ``p_v_3_length`` / ``p_d_5_length`` / - ``p_d_3_length`` / ``p_j_5_length`` (separate - P-nucleotide biology — see - ``docs/p_nucleotide_design.md``), - and ``junction_length`` (aggregate arithmetic too - fragile across simulators — audit §5.2). - - Per-row validation is **field-local**: a malformed - or missing value in one NP column drops only that - key's contribution; the row's other column still - feeds. Per-row structured entries land in - ``report.rejected`` with reasons - ``"missing_required_column"`` / - ``"malformed_length_value"`` / - ``"negative_length_value"``, each carrying the - AIRR column name. - - ``min_count`` (int, default 1) drops length values - whose observed integer count is strictly below - the threshold before normalisation. - - ``pseudocount`` (float, default 0.0) adds a uniform - prior to every **observed** length value before - normalisation (no support expansion). Applied - AFTER the ``min_count`` filter. - - ``replace`` (default ``True``) controls idempotency. - When ``False`` AND a prior typed-plane - ``np_lengths`` is attached to - ``self._reference_models``, the call raises - :class:`ValueError` before consuming any records. - """ - if isinstance(min_count, bool) or not isinstance(min_count, int): - raise ValueError( - f"min_count must be an int, got {min_count!r}" - ) - if min_count < 0: - raise ValueError( - f"min_count must be non-negative, got {min_count!r}" - ) - if isinstance(pseudocount, bool) or not isinstance( - pseudocount, (int, float) - ): - raise ValueError( - f"pseudocount must be a non-negative number, got {pseudocount!r}" - ) - if pseudocount < 0.0: - raise ValueError( - f"pseudocount must be non-negative, got {pseudocount!r}" - ) - - if not replace: - existing = ( - self._reference_models.np_lengths - if self._reference_models is not None - else None - ) - if existing: - raise ValueError( - "np_length distributions already attached to this " - "cartridge; pass replace=True to overwrite" - ) - previously_estimated = any( - entry.get("stage") == "estimate_np_length_distributions" - for entry in self._report.stages - ) - - records, source_label = _load_rearrangements(rearrangements) - chain_has_d = self._chain_type.has_d - active_keys = NP_KEYS if chain_has_d else ("NP1",) - - key_to_column = {"NP1": "np1_length", "NP2": "np2_length"} - - counters: Dict[str, Dict[int, int]] = {k: {} for k in active_keys} - skipped = { - "missing_required_column": { - key_to_column[k]: 0 for k in active_keys - }, - "malformed_length_value": { - key_to_column[k]: 0 for k in active_keys - }, - "negative_length_value": { - key_to_column[k]: 0 for k in active_keys - }, - } - # VJ chains track `np2_length` contributions as a dropped - # column with a single one-time warning. - dropped_columns: Dict[str, int] = {} - if not chain_has_d: - dropped_columns["np2_length"] = 0 - rejected_entries: List[Dict[str, Any]] = [] - warnings: List[str] = [] - vj_np2_warned = False - - for row_idx, row in enumerate(records): - # On VJ, surface non-zero `np2_length` as a dropped - # column with one warning across the dataset. - if not chain_has_d: - raw_np2 = row.get("np2_length") - if raw_np2 not in (None, "", "0"): - try: - if int(str(raw_np2).strip()) != 0: - dropped_columns["np2_length"] += 1 - if not vj_np2_warned: - warnings.append( - "VJ cartridge: np2_length column " - "present in records; contribution " - "ignored (no NP2 region on a VJ " - "chain)" - ) - vj_np2_warned = True - except (TypeError, ValueError): - pass - - # Field-local validation per active key. - for key in active_keys: - col = key_to_column[key] - raw = row.get(col) - if raw is None or str(raw).strip() == "": - skipped["missing_required_column"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_np_length_distributions", - "row_index": row_idx, - "column": col, - "key": key, - "reason": "missing_required_column", - } - ) - continue - try: - value = int(str(raw).strip()) - except (TypeError, ValueError): - skipped["malformed_length_value"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_np_length_distributions", - "row_index": row_idx, - "column": col, - "key": key, - "value": raw, - "reason": "malformed_length_value", - } - ) - continue - if value < 0: - skipped["negative_length_value"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_np_length_distributions", - "row_index": row_idx, - "column": col, - "key": key, - "value": value, - "reason": "negative_length_value", - } - ) - continue - counters[key][value] = counters[key].get(value, 0) + 1 - - # min_count + pseudocount + per-key normalisation. - below_min: Dict[str, int] = {k: 0 for k in active_keys} - inferred_pairs: Dict[str, List[tuple]] = {k: [] for k in active_keys} - - for key in active_keys: - raw_counts = counters[key] - kept = {v: c for v, c in raw_counts.items() if c >= min_count} - below_min[key] = len(raw_counts) - len(kept) - if pseudocount > 0.0: - kept = {v: (c + pseudocount) for v, c in kept.items()} - total = float(sum(kept.values())) - if total <= 0.0 or not kept: - inferred_pairs[key] = [] - continue - normalised = [(v, kept[v] / total) for v in sorted(kept.keys())] - inferred_pairs[key] = normalised - - if any(below_min[k] for k in active_keys): - warnings.append( - f"NP-length values below min_count={min_count} dropped — " - + ", ".join(f"{k}={below_min[k]}" for k in active_keys) - ) - - # Build per-key EmpiricalDistributionSpec instances. - new_np_lengths: Dict[str, EmpiricalDistributionSpec] = {} - for key, pairs in inferred_pairs.items(): - if not pairs: - continue - spec = EmpiricalDistributionSpec(pairs) - spec.validate(name=f"np_lengths[{key}]") - new_np_lengths[key] = spec - - # Attach to existing reference_models (creating if absent). - # Other typed planes are preserved. - if self._reference_models is None: - self._reference_models = ReferenceEmpiricalModels() - self._reference_models = ReferenceEmpiricalModels( - np_lengths=new_np_lengths, - trims=self._reference_models.trims, - np_bases=self._reference_models.np_bases, - p_nucleotide_lengths=self._reference_models.p_nucleotide_lengths, - allele_usage=self._reference_models.allele_usage, - ) - chain_label = "vdj" if chain_has_d else "vj" - self._reference_models.validate(chain_type=chain_label) - - self._report.stages.append( - { - "stage": "estimate_np_length_distributions", - "inputs": { - "record_count": len(records), - "min_count": int(min_count), - "pseudocount": float(pseudocount), - "source": source_label, - "replaced": previously_estimated, - }, - "inferred": { - "NP1": inferred_pairs.get("NP1", []), - "NP2": inferred_pairs.get("NP2", []), - "skipped": skipped, - "below_min_count": {k: below_min.get(k, 0) for k in NP_KEYS}, - "dropped_columns": dropped_columns, - }, - "warnings": warnings, - } - ) - self._report.rejected.extend(rejected_entries) - return self - - def estimate_np_base_model( - self, - rearrangements: Any, - *, - kind: str = "markov", - min_count: int = 1, - pseudocount: float = 0.0, - replace: bool = True, - ) -> "ReferenceCartridgeBuilder": - """Estimate per-key NP base sampling model from - observed AIRR rearrangement records. - - Writes ``NpBaseModelSpec`` instances into - ``self._reference_models.np_bases`` keyed by - ``"NP1"`` (always) and (on VDJ cartridges) - ``"NP2"``. VJ cartridges produce only the ``NP1`` - key. - - ``kind`` selects the model family: - - - ``"empirical_first_base"``: estimates a single - categorical over A/C/G/T from the **full base - composition** of every observed NP string (every - position, not just position 0). The engine - samples every NP position independently from - this distribution — the name is preserved for - API stability with existing cartridges; the - biologically correct estimate is the full-base - composition. - - ``"markov"`` (default): estimates a first-base - row from position 0 of each NP string plus a - 4×4 transition matrix from every observed - (prev, next) pair. - - AIRR columns consumed (per audit §1.2): - - - ``np1`` (always). - - ``np2`` (VDJ only — VJ rows with a non-empty - ``np2`` raise a one-time warning and the - contribution is skipped). - - Columns deliberately ignored: ``junction`` (audit - §1.2), ``p_v_3_length`` / ``p_d_5_length`` / - ``p_d_3_length`` / ``p_j_5_length`` (separate - P-nucleotide biology), ``np1_length`` / - ``np2_length`` (length-only — owned by - :meth:`estimate_np_length_distributions`). - - Per-row validation is **field-local**: a malformed - or missing value in one NP column drops only that - key's contribution. Per-row structured entries - land in ``report.rejected`` with reasons - ``"missing_required_column"`` (empty / missing - string) or ``"noncanonical_base"`` (any character - outside ``{A,C,G,T}`` after uppercasing). - - ``min_count`` (int, default 1) drops first-base - categories whose observed count is strictly below - the threshold before normalisation. For ``markov``, - ``min_count`` applies to the **first-base row - only**: dropping transition cells could leave a - from-base row with no positive weights, which the - :class:`NpBaseModelSpec` validator rejects. v1 - keeps transition rows intact and surfaces the - first-base drops in ``below_min_count.first_base``. - - ``pseudocount`` (float, default 0.0) adds a uniform - prior: - - - ``empirical_first_base``: added to every A/C/G/T - base category before normalisation. - - ``markov``: added to every A/C/G/T first-base - category AND to every cell of the 4×4 transition - matrix. - - ``replace`` (default ``True``) controls idempotency. - When ``False`` AND a prior typed-plane - ``np_bases`` is already attached, the call raises - :class:`ValueError` before consuming any records. - """ - if kind not in ("empirical_first_base", "markov"): - raise ValueError( - f"kind must be one of 'empirical_first_base' / " - f"'markov', got {kind!r}" - ) - if isinstance(min_count, bool) or not isinstance(min_count, int): - raise ValueError( - f"min_count must be an int, got {min_count!r}" - ) - if min_count < 0: - raise ValueError( - f"min_count must be non-negative, got {min_count!r}" - ) - if isinstance(pseudocount, bool) or not isinstance( - pseudocount, (int, float) - ): - raise ValueError( - f"pseudocount must be a non-negative number, got {pseudocount!r}" - ) - if pseudocount < 0.0: - raise ValueError( - f"pseudocount must be non-negative, got {pseudocount!r}" - ) - - if not replace: - existing = ( - self._reference_models.np_bases - if self._reference_models is not None - else None - ) - if existing: - raise ValueError( - "np_bases model already attached to this cartridge; " - "pass replace=True to overwrite" - ) - previously_estimated = any( - entry.get("stage") == "estimate_np_base_model" - for entry in self._report.stages - ) - - records, source_label = _load_rearrangements(rearrangements) - chain_has_d = self._chain_type.has_d - active_keys = ("NP1", "NP2") if chain_has_d else ("NP1",) - - key_to_column = {"NP1": "np1", "NP2": "np2"} - bases = _NP_CANONICAL_BASES # ("A","C","G","T") - - # Per-key tallies. first_base[key] is a dict over A/C/G/T; - # transitions[key] is a 4×4 dict-of-dict over A/C/G/T from→to. - first_base_counts: Dict[str, Dict[str, int]] = { - key: {b: 0 for b in bases} for key in active_keys - } - transition_counts: Dict[str, Dict[str, Dict[str, int]]] = { - key: {b: {t: 0 for t in bases} for b in bases} - for key in active_keys - } - skipped = { - "missing_required_column": { - key_to_column[k]: 0 for k in active_keys - }, - "noncanonical_base": { - key_to_column[k]: 0 for k in active_keys - }, - } - # VJ chains track non-empty `np2` strings as a dropped - # column with a single one-time warning. - dropped_columns: Dict[str, int] = {} - if not chain_has_d: - dropped_columns["np2"] = 0 - rejected_entries: List[Dict[str, Any]] = [] - warnings: List[str] = [] - vj_np2_warned = False - - for row_idx, row in enumerate(records): - # Surface non-empty `np2` on VJ as a dropped column - # with one warning across the dataset. - if not chain_has_d: - raw_np2 = row.get("np2") - if isinstance(raw_np2, str) and raw_np2.strip(): - dropped_columns["np2"] += 1 - if not vj_np2_warned: - warnings.append( - "VJ cartridge: np2 column non-empty in " - "records; contribution ignored (no NP2 " - "region on a VJ chain)" - ) - vj_np2_warned = True - - # Field-local validation per active key. - for key in active_keys: - col = key_to_column[key] - raw = row.get(col) - if raw is None or not isinstance(raw, str) or not raw.strip(): - skipped["missing_required_column"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_np_base_model", - "row_index": row_idx, - "column": col, - "key": key, - "reason": "missing_required_column", - } - ) - continue - upper = raw.strip().upper() - non_canonical = sorted(set(upper) - set(bases)) - if non_canonical: - skipped["noncanonical_base"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_np_base_model", - "row_index": row_idx, - "column": col, - "key": key, - "unknown_chars": non_canonical, - "reason": "noncanonical_base", - } - ) - continue - # Tally bases. - # - For `empirical_first_base` the engine samples - # every NP position independently from the same - # distribution, so the biologically correct - # estimate is the FULL base composition (every - # position). - # - For `markov` the first-base row models position - # 0 only (the seed of the chain); position 1+ - # are modelled by transitions. - if kind == "empirical_first_base": - for b in upper: - first_base_counts[key][b] += 1 - else: - first_base_counts[key][upper[0]] += 1 - for i in range(len(upper) - 1): - prev_b = upper[i] - next_b = upper[i + 1] - transition_counts[key][prev_b][next_b] += 1 - - # min_count + pseudocount + per-kind spec construction. - below_min: Dict[str, Dict[str, int]] = { - key: {"first_base": 0} for key in active_keys - } - inferred: Dict[str, Dict[str, Any]] = { - key: {} for key in active_keys - } - - new_np_bases: Dict[str, NpBaseModelSpec] = {} - from_base_errors: List[str] = [] - - for key in active_keys: - fb_counts = first_base_counts[key] - # Apply min_count filter to first-base row (drop bases - # whose count is strictly below threshold). - fb_kept = {b: c for b, c in fb_counts.items() if c >= min_count} - below_min[key]["first_base"] = len(fb_counts) - len(fb_kept) - # Apply pseudocount AFTER filter — observed-only support. - if pseudocount > 0.0: - for b in bases: - if b in fb_kept: - fb_kept[b] = fb_kept[b] + pseudocount - else: - fb_kept[b] = pseudocount - fb_total = float(sum(fb_kept.values())) - if fb_total <= 0.0: - # Nothing observed AND no pseudocount → skip this key. - inferred[key] = {"first_base": None, "transitions": None} - continue - first_base_weights = { - b: fb_kept[b] / fb_total - for b in sorted(fb_kept.keys()) - if fb_kept[b] > 0 - } - - spec_payload: Dict[str, Any] = { - "first_base": first_base_weights, - } - if kind == "markov": - # Apply pseudocount to every cell of every transition row. - tr = transition_counts[key] - tr_weights: Dict[str, Dict[str, float]] = {} - for from_b in bases: - row_counts = {t: float(tr[from_b][t]) for t in bases} - if pseudocount > 0.0: - for t in bases: - row_counts[t] += pseudocount - row_total = sum(row_counts.values()) - if row_total <= 0.0: - # Pseudocount=0 and from-base never observed: surface a - # tagged error so the caller knows to pass pseudocount. - from_base_errors.append( - f"{key}: transition row for from-base " - f"{from_b!r} has no observed transitions " - f"and pseudocount=0; pass pseudocount > 0 " - f"or supply more data" - ) - continue - tr_weights[from_b] = { - t: row_counts[t] / row_total - for t in bases - if row_counts[t] > 0 - } - if from_base_errors: - # Bail before constructing a partial spec. - raise ValueError("; ".join(from_base_errors)) - spec_payload["transitions"] = tr_weights - else: - spec_payload["transitions"] = None - - # Construct + validate the spec. - spec = NpBaseModelSpec( - kind=kind, - first_base=spec_payload["first_base"], - transitions=spec_payload["transitions"], - ) - spec.validate(name=f"np_bases[{key}]") - new_np_bases[key] = spec - inferred[key] = { - "first_base": dict(spec_payload["first_base"]), - "transitions": ( - {k: dict(v) for k, v in spec_payload["transitions"].items()} - if spec_payload["transitions"] is not None - else None - ), - } - - if any(below_min[k]["first_base"] for k in active_keys): - warnings.append( - f"first-base values below min_count={min_count} dropped — " - + ", ".join( - f"{k}={below_min[k]['first_base']}" - for k in active_keys - ) - ) - - # Attach to existing reference_models (creating if absent); - # other typed planes preserved. - if self._reference_models is None: - self._reference_models = ReferenceEmpiricalModels() - self._reference_models = ReferenceEmpiricalModels( - np_lengths=self._reference_models.np_lengths, - trims=self._reference_models.trims, - np_bases=new_np_bases, - p_nucleotide_lengths=self._reference_models.p_nucleotide_lengths, - allele_usage=self._reference_models.allele_usage, - ) - chain_label = "vdj" if chain_has_d else "vj" - self._reference_models.validate(chain_type=chain_label) - - # NP2 entry surfaces in inferred even on VJ (as empty dict) - # for shape stability — same discipline as the NP-length - # estimator's empty `D_5` / `D_3` entries. - np1_inferred = inferred.get("NP1", {"first_base": None, "transitions": None}) - np2_inferred = inferred.get("NP2", {"first_base": None, "transitions": None}) - - self._report.stages.append( - { - "stage": "estimate_np_base_model", - "inputs": { - "record_count": len(records), - "kind": kind, - "min_count": int(min_count), - "pseudocount": float(pseudocount), - "source": source_label, - "replaced": previously_estimated, - }, - "inferred": { - "NP1": np1_inferred, - "NP2": np2_inferred, - "skipped": skipped, - "below_min_count": { - k: below_min.get(k, {"first_base": 0}) - for k in NP_KEYS - }, - "dropped_columns": dropped_columns, - }, - "warnings": warnings, - } - ) - self._report.rejected.extend(rejected_entries) - return self - - def estimate_p_nucleotide_lengths( - self, - rearrangements: Any, - *, - min_count: int = 1, - pseudocount: float = 0.0, - replace: bool = True, - ) -> "ReferenceCartridgeBuilder": - """Estimate per-end P-nucleotide length distributions - from observed AIRR rearrangement records. - - Writes ``EmpiricalDistributionSpec`` instances into - ``self._reference_models.p_nucleotide_lengths`` - keyed by ``"V_3"`` and ``"J_5"`` (always) plus - ``"D_5"`` and ``"D_3"`` (VDJ only). VJ cartridges - produce only the V_3 and J_5 keys; the typed-plane - validator chain-type-rejects D-end keys on VJ at - attach time. - - **Provenance warning.** This estimator requires - AIRR-like records that ALREADY carry GenAIRR's - P-length fields (``p_v_3_length`` / - ``p_d_5_length`` / ``p_d_3_length`` / - ``p_j_5_length``). It does NOT infer P-lengths - from generic AIRR junction sequences, NP strings, - or trim arithmetic — that inference problem is - out of scope for v1. External AIRR tools (IgBLAST, - MiXCR, …) do not model P-nucleotide additions, so - their output will either omit these columns - entirely (rejection storm) or populate them as - zero (degenerate ``[(0, 1.0)]`` distribution). - The estimator emits a stage-level warning per key - when ≥ 95% of contributing rows reported zero — - diagnostic of P-naïve input. See - ``docs/p_nucleotide_length_estimation_design.md`` - §5 for the realistic-input enumeration. - - AIRR columns consumed: - - - ``p_v_3_length`` (always). - - ``p_j_5_length`` (always). - - ``p_d_5_length`` (VDJ only — VJ rows with a - non-zero value raise a one-time warning per - column and the contribution is skipped). - - ``p_d_3_length`` (same). - - Columns deliberately ignored: every other AIRR - column. v1 does NOT derive P-lengths from - junction-arithmetic, NP strings, trim fields, or - end-loss fields. - - Per-row validation is **field-local**: a malformed - or missing value in one P-length column drops only - that key's contribution. Per-row structured - entries land in ``report.rejected`` with reasons - ``"missing_required_column"`` / - ``"malformed_length_value"`` / - ``"negative_length_value"``. - - ``min_count`` (int, default 1) drops length values - whose observed count is strictly below the - threshold before normalisation. - - ``pseudocount`` (float, default 0.0) adds a uniform - prior to every **observed** length value before - normalisation (no support expansion). Applied - AFTER the ``min_count`` filter. - - ``replace`` (default ``True``) controls idempotency. - When ``False`` AND a prior typed-plane - ``p_nucleotide_lengths`` is already attached, - the call raises :class:`ValueError` before - consuming any records. - """ - if isinstance(min_count, bool) or not isinstance(min_count, int): - raise ValueError( - f"min_count must be an int, got {min_count!r}" - ) - if min_count < 0: - raise ValueError( - f"min_count must be non-negative, got {min_count!r}" - ) - if isinstance(pseudocount, bool) or not isinstance( - pseudocount, (int, float) - ): - raise ValueError( - f"pseudocount must be a non-negative number, got {pseudocount!r}" - ) - if pseudocount < 0.0: - raise ValueError( - f"pseudocount must be non-negative, got {pseudocount!r}" - ) - - if not replace: - existing = ( - self._reference_models.p_nucleotide_lengths - if self._reference_models is not None - else None - ) - if existing: - raise ValueError( - "p_nucleotide_lengths model already attached to " - "this cartridge; pass replace=True to overwrite" - ) - previously_estimated = any( - entry.get("stage") == "estimate_p_nucleotide_lengths" - for entry in self._report.stages - ) - - records, source_label = _load_rearrangements(rearrangements) - chain_has_d = self._chain_type.has_d - active_keys = ( - P_NUCLEOTIDE_END_KEYS if chain_has_d - else P_NUCLEOTIDE_END_KEYS_VJ - ) - - key_to_column = { - "V_3": "p_v_3_length", - "D_5": "p_d_5_length", - "D_3": "p_d_3_length", - "J_5": "p_j_5_length", - } - vj_ignored_columns = ("p_d_5_length", "p_d_3_length") - - counters: Dict[str, Dict[int, int]] = {k: {} for k in active_keys} - contributing_counts: Dict[str, int] = {k: 0 for k in active_keys} - zero_counts: Dict[str, int] = {k: 0 for k in active_keys} - skipped = { - "missing_required_column": { - key_to_column[k]: 0 for k in active_keys - }, - "malformed_length_value": { - key_to_column[k]: 0 for k in active_keys - }, - "negative_length_value": { - key_to_column[k]: 0 for k in active_keys - }, - } - # VJ chains track nonzero `p_d_*_length` columns as - # dropped columns with one warning per column. - dropped_columns: Dict[str, int] = {} - if not chain_has_d: - for col in vj_ignored_columns: - dropped_columns[col] = 0 - rejected_entries: List[Dict[str, Any]] = [] - warnings: List[str] = [] - vj_ignored_warned = {col: False for col in vj_ignored_columns} - - for row_idx, row in enumerate(records): - # VJ: surface nonzero D-end columns as dropped + warn once per column. - if not chain_has_d: - for col in vj_ignored_columns: - raw = row.get(col) - if raw in (None, "", "0"): - continue - try: - value = int(str(raw).strip()) - except (TypeError, ValueError): - continue - if value != 0: - dropped_columns[col] += 1 - if not vj_ignored_warned[col]: - warnings.append( - f"VJ cartridge: {col} non-zero in " - f"input — contribution ignored (no " - f"D segment on a VJ chain)" - ) - vj_ignored_warned[col] = True - - # Field-local validation per active key. - for key in active_keys: - col = key_to_column[key] - raw = row.get(col) - if raw is None or str(raw).strip() == "": - skipped["missing_required_column"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_p_nucleotide_lengths", - "row_index": row_idx, - "column": col, - "key": key, - "reason": "missing_required_column", - } - ) - continue - try: - value = int(str(raw).strip()) - except (TypeError, ValueError): - skipped["malformed_length_value"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_p_nucleotide_lengths", - "row_index": row_idx, - "column": col, - "key": key, - "value": raw, - "reason": "malformed_length_value", - } - ) - continue - if value < 0: - skipped["negative_length_value"][col] += 1 - rejected_entries.append( - { - "stage": "estimate_p_nucleotide_lengths", - "row_index": row_idx, - "column": col, - "key": key, - "value": value, - "reason": "negative_length_value", - } - ) - continue - counters[key][value] = counters[key].get(value, 0) + 1 - contributing_counts[key] += 1 - if value == 0: - zero_counts[key] += 1 - - # min_count + pseudocount + per-key normalisation. - below_min: Dict[str, int] = {k: 0 for k in active_keys} - inferred_pairs: Dict[str, List[tuple]] = {k: [] for k in active_keys} - zero_fraction: Dict[str, float] = {} - - for key in active_keys: - raw_counts = counters[key] - kept = {v: c for v, c in raw_counts.items() if c >= min_count} - below_min[key] = len(raw_counts) - len(kept) - if pseudocount > 0.0: - kept = {v: (c + pseudocount) for v, c in kept.items()} - total = float(sum(kept.values())) - if total <= 0.0 or not kept: - inferred_pairs[key] = [] - else: - inferred_pairs[key] = [ - (v, kept[v] / total) for v in sorted(kept.keys()) - ] - # zero_fraction tracks observed-row provenance before - # normalisation — useful diagnostic for P-naïve input. - n_contrib = contributing_counts[key] - if n_contrib > 0: - zero_fraction[key] = zero_counts[key] / n_contrib - else: - zero_fraction[key] = 0.0 - - if any(below_min[k] for k in active_keys): - warnings.append( - f"P-nucleotide length values below min_count={min_count} " - f"dropped — " - + ", ".join(f"{k}={below_min[k]}" for k in active_keys) - ) - - # Per-key provenance auto-warning (audit §5.2). - for key in active_keys: - if ( - contributing_counts[key] > 0 - and zero_fraction[key] >= 0.95 - ): - warnings.append( - f"p_nucleotide_lengths[{key}] is >=95% zero; " - f"input may be P-naive or lack P annotations" - ) - - # Build per-key EmpiricalDistributionSpec instances. - new_p_lengths: Dict[str, EmpiricalDistributionSpec] = {} - for key, pairs in inferred_pairs.items(): - if not pairs: - continue - spec = EmpiricalDistributionSpec(pairs) - spec.validate(name=f"p_nucleotide_lengths[{key}]") - new_p_lengths[key] = spec - - # Attach to existing reference_models (creating if absent); - # other typed planes preserved. - if self._reference_models is None: - self._reference_models = ReferenceEmpiricalModels() - self._reference_models = ReferenceEmpiricalModels( - np_lengths=self._reference_models.np_lengths, - trims=self._reference_models.trims, - np_bases=self._reference_models.np_bases, - p_nucleotide_lengths=new_p_lengths, - allele_usage=self._reference_models.allele_usage, - ) - chain_label = "vdj" if chain_has_d else "vj" - self._reference_models.validate(chain_type=chain_label) - - self._report.stages.append( - { - "stage": "estimate_p_nucleotide_lengths", - "inputs": { - "record_count": len(records), - "min_count": int(min_count), - "pseudocount": float(pseudocount), - "source": source_label, - "replaced": previously_estimated, - }, - "inferred": { - "V_3": inferred_pairs.get("V_3", []), - "D_5": inferred_pairs.get("D_5", []), - "D_3": inferred_pairs.get("D_3", []), - "J_5": inferred_pairs.get("J_5", []), - "skipped": skipped, - "below_min_count": { - k: below_min.get(k, 0) for k in P_NUCLEOTIDE_END_KEYS - }, - "zero_fraction": { - k: zero_fraction.get(k, 0.0) - for k in P_NUCLEOTIDE_END_KEYS - }, - "dropped_columns": dropped_columns, - }, - "warnings": warnings, - } - ) - self._report.rejected.extend(rejected_entries) - return self - # ────────────────────────────────────────────────────────── # Finalisation # ────────────────────────────────────────────────────────── diff --git a/src/GenAIRR/cohort.py b/src/GenAIRR/cohort.py index 86a10a7..f0a2f19 100644 --- a/src/GenAIRR/cohort.py +++ b/src/GenAIRR/cohort.py @@ -5,7 +5,6 @@ exposes per-subject access (``result_for`` / ``refdata_for``) plus combined export. It explicitly stores each subject's refdata because ``SimulationResult`` does not preserve it (needed for ``validate_records`` and novel-allele subjects). -See ``.private/specs/2026-06-17-genotype-cohorts-design.md``. """ from __future__ import annotations @@ -56,6 +55,9 @@ def _resolve_counts(n_genotypes: int, n_per_subject, counts) -> List[int]: return [_check_count(c, f"counts[{i}]") for i, c in enumerate(counts)] +__all__ = ["CohortSubjectResult", "CohortResult"] + + @dataclass(frozen=True) class CohortSubjectResult: """One subject's slice of a cohort run.""" @@ -123,7 +125,7 @@ def to_dataframe(self, *, airr_strict: bool = False): stable union across subjects (pandas fills missing keys with NaN).""" import pandas as pd - from .result import _DEFAULT_COLUMN_ORDER + from ._result_export import _DEFAULT_COLUMN_ORDER records = self.records if not records: @@ -136,7 +138,7 @@ def to_dataframe(self, *, airr_strict: bool = False): cols.append(extra) return pd.DataFrame(columns=cols) if airr_strict: - from .result import _to_airr_strict + from ._result_export import _to_airr_strict records = [_to_airr_strict(r) for r in records] return pd.DataFrame(records) diff --git a/src/GenAIRR/dataconfig/_manifest.py b/src/GenAIRR/dataconfig/_manifest.py new file mode 100644 index 0000000..8775b17 --- /dev/null +++ b/src/GenAIRR/dataconfig/_manifest.py @@ -0,0 +1,731 @@ +"""Cartridge-manifest construction for :class:`DataConfig`. + +Extracted verbatim from ``data_config.py`` (behavior-preserving) so the +DataConfig data model and its manifest/reporting subsystem are separately +navigable. ``DataConfig.cartridge_manifest`` is a thin delegator to +``build_manifest`` here. +""" +from typing import Any, Dict, List, Optional + +from GenAIRR.alleles.allele import Allele +from GenAIRR.genotype_priors import PopulationGenotypeModel + + +# Documented completeness gaps surfaced verbatim by +# ``DataConfig.cartridge_manifest``. See +# ``docs/reference_cartridge_completeness_audit.md`` §1 (orphan +# fields) and §2 (dropped allele fields) for the rationale on each +# entry. These are tuples so manifest output stays stable across +# calls and across Python sessions. +_DOCUMENTED_DROPPED_ALLELE_FIELDS = ( + "aliases", + "anchor_meta", + "gapped_seq", + "family", + "locus", + "species", + "source", +) +_DOCUMENTED_ORPHAN_DATACONFIG_FIELDS = ( + "gene_use_dict", + "NP_transitions", + "NP_first_bases", + "correction_maps", + "asc_tables", + "p_nucleotide_length_probs", + "dj_pairing_map", +) + + +def _jsonable(value): + """Coerce a value to a JSON-clean form. Enums become their + ``.name`` (uppercase canonical) for stability; tuples become + lists; everything else passes through. Used by + :meth:`DataConfig.cartridge_manifest` to make the identity + plane safe for ``json.dumps``.""" + if value is None: + return None + if isinstance(value, (str, int, float, bool)): + return value + if hasattr(value, "name") and not isinstance(value, type): + # Enum-shaped values (Species, ChainType, etc.). + return value.name + if isinstance(value, (tuple, list)): + return [_jsonable(v) for v in value] + return str(value) + + +def _normalise_anchor_rule(rule): + """Coerce an anchor rule dict from the Rust ``rules()`` accessor + into a JSON-clean shape. Returns ``None`` for missing rules. + + The Rust side may surface the rule as a dict with ``required``, + ``expected_aa``, ``missing_severity``, ``mismatch_severity`` + keys (string severities). We pass it through but coerce any + iterable fields to plain ``list`` to keep ``json.dumps`` happy. + """ + if rule is None: + return None + if isinstance(rule, dict): + out = dict(rule) + if "expected_aa" in out and out["expected_aa"] is not None: + out["expected_aa"] = list(out["expected_aa"]) + return out + # Unknown shape — drop silently rather than raise; the manifest + # error list isn't the right channel for "this is unexpected + # but harmless". + return None + + +def _allele_usage_manifest_block(cfg): + """Build the ``models.allele_usage`` manifest block per the + Allele Usage Estimation v1 audit (§8 of + ``docs/allele_usage_estimation_design.md``). + + Returns a JSON-clean dict carrying: + + - ``available``: whether ``reference_models.allele_usage`` is + a non-``None`` spec. + - ``segments``: the canonical V/D/J label list. + - ``nonempty_segments``: subset of ``segments`` whose + author-supplied dict is non-empty. + - ``legacy_gene_use_dict_present``: documents the orphan + legacy dict's bundled-cartridge presence so a downstream + tool can decide whether to author a typed plane from it. + - ``legacy_fallback``: always ``False`` in v1 (no auto-lift). + - ``in_plan_signature``: ``False`` — inherits soft gap 1 + from the plan-signature completeness audit. Tightening is + a separate slice covering BOTH the kwarg AND the + cartridge-driven path in lockstep. + - ``source``: which Python field the manifest block reads + from. Always ``"ReferenceEmpiricalModels.allele_usage"``. + """ + rm = getattr(cfg, "reference_models", None) + spec = getattr(rm, "allele_usage", None) if rm is not None else None + available = spec is not None + nonempty: list = [] + if available: + try: + nonempty = list(spec.nonempty_segments()) + except Exception: + nonempty = [] + return { + "available": available, + "segments": ["V", "D", "J"], + "nonempty_segments": nonempty, + "legacy_gene_use_dict_present": bool( + getattr(cfg, "gene_use_dict", None) + ), + "legacy_fallback": False, + "in_plan_signature": False, # documented soft gap 1 + "source": "ReferenceEmpiricalModels.allele_usage", + } + + +def _genotype_priors_manifest_block(cfg): + """Build the ``models.genotype_priors`` manifest block (Slice — Cartridge + genotype plane). Audit-sized: counts and identity, never the full tables. + Reads the top-level ``DataConfig.genotype_priors`` plane (independent of + ``reference_models``).""" + def _empty(available, valid, error=None): + b = { + "available": available, + "valid": valid, + "model_id": None, + "source": None, + "version": None, + "model_checksum": None, + "segments_with_frequencies": [], + "freq_gene_counts": {"V": 0, "D": 0, "J": 0}, + "deletion_gene_counts": {"V": 0, "D": 0, "J": 0}, + "novel_allele_count": 0, + "chromosome_weights": None, + "source_field": "DataConfig.genotype_priors", + } + if error is not None: + b["validation_error"] = error + return b + + model = getattr(cfg, "genotype_priors", None) + if model is None: + return _empty(False, None) + if not isinstance(model, PopulationGenotypeModel): + # Field is typed Optional[PopulationGenotypeModel] but Python won't enforce + # it; a garbage value must not crash the manifest (called from build()). + return _empty(True, False, + f"genotype_priors is not a PopulationGenotypeModel " + f"(got {type(model).__name__})") + + # A plane may have been attached directly (bypassing builder validation). + # Report validity rather than leaking non-JSON-clean numerics (e.g. NaN + # chromosome weights) into the manifest. + try: + model.validate( + chain_type=getattr(getattr(cfg, "metadata", None), "chain_type", None)) + valid, validation_error = True, None + except ValueError as exc: + valid, validation_error = False, str(exc) + + # content_checksum / float() can themselves raise on a malformed model; never + # let that crash the manifest — report None and the validity flag instead. + try: + checksum = model.content_checksum() + except Exception: + checksum = None + cw = None + if valid: + try: + cw = [float(model.chromosome_weights[0]), float(model.chromosome_weights[1])] + except Exception: + cw = None + freq = model.allele_frequencies if isinstance(model.allele_frequencies, dict) else {} + dele = model.haplotype_deletion_prob if isinstance(model.haplotype_deletion_prob, dict) else {} + novels = model.novel_alleles if isinstance(model.novel_alleles, (list, tuple)) else [] + block = { + "available": True, + "valid": valid, + "model_id": (model.model_id or None) if isinstance(model.model_id, str) else None, + "source": (model.source or None) if isinstance(model.source, str) else None, + "version": (model.version or None) if isinstance(model.version, str) else None, + "model_checksum": checksum, + "segments_with_frequencies": [s for s in ("V", "D", "J") if freq.get(s)], + "freq_gene_counts": {s: len(freq.get(s, {})) for s in ("V", "D", "J")}, + "deletion_gene_counts": {s: len(dele.get(s, {})) for s in ("V", "D", "J")}, + "novel_allele_count": len(novels), + "chromosome_weights": cw, + "source_field": "DataConfig.genotype_priors", + } + if validation_error is not None: + block["validation_error"] = validation_error + return block + + +def _np_length_models_manifest_block(cfg): + """Build the ``models.np_length_models`` manifest block + per the NP Length Distribution Estimation v1 audit + (§8 of ``docs/np_length_estimation_design.md``). + + Returns a JSON-clean dict carrying: + + - ``keys``: sorted list of NP-key strings present on + the typed ``ReferenceEmpiricalModels.np_lengths`` + plane (subset of ``("NP1", "NP2")``). + - ``source``: which Python field the manifest block + reads from. Always ``"ReferenceEmpiricalModels.np_lengths"``. + - ``in_plan_signature``: ``True`` — NP-length + distributions fold into the plan signature via + ``GenerateNPPass.parameter_signature`` / + ``fmt_int_dist`` (verified at audit time). **No + soft gap inherited.** + - ``legacy_np_lengths_present``: documents the bundled + cartridges' legacy nested-dict presence so a + downstream tool can decide whether to author a + typed plane from it. + - ``legacy_fallback``: always ``False`` in v1 (no + auto-lift). + """ + rm = getattr(cfg, "reference_models", None) + np_lengths = getattr(rm, "np_lengths", None) if rm is not None else None + keys = sorted(np_lengths.keys()) if np_lengths else [] + return { + "keys": keys, + "source": "ReferenceEmpiricalModels.np_lengths", + "in_plan_signature": True, + "legacy_np_lengths_present": bool(getattr(cfg, "NP_lengths", None)), + "legacy_fallback": False, + } + + +def _trim_models_manifest_block(cfg): + """Build the ``models.trim_models`` manifest block per the + Trim Distribution Estimation v1 audit (§8 of + ``docs/trim_distribution_estimation_design.md``). + + Returns a JSON-clean dict carrying: + + - ``keys``: sorted list of trim-key strings present on the + typed ``ReferenceEmpiricalModels.trims`` plane (subset of + ``("V_3", "D_5", "D_3", "J_5")``). + - ``source``: which Python field the manifest block reads + from. Always ``"ReferenceEmpiricalModels.trims"``. + - ``in_plan_signature``: ``True`` — trim distributions fold + into the plan signature via ``fmt_int_dist`` (verified at + audit time). **No soft gap inherited** — unlike the + allele-usage slice. + - ``legacy_trim_dicts_present``: documents the bundled + cartridges' legacy nested-dict presence so a downstream + tool can decide whether to author a typed plane from it. + - ``legacy_fallback``: always ``False`` in v1 (no auto-lift). + """ + rm = getattr(cfg, "reference_models", None) + trims = getattr(rm, "trims", None) if rm is not None else None + keys = sorted(trims.keys()) if trims else [] + return { + "keys": keys, + "source": "ReferenceEmpiricalModels.trims", + "in_plan_signature": True, + "legacy_trim_dicts_present": bool(getattr(cfg, "trim_dicts", None)), + "legacy_fallback": False, + } + + +def _p_nucleotide_length_keys(cfg): + """Return the sorted list of P-end keys that have a typed + `EmpiricalDistributionSpec` on + `cfg.reference_models.p_nucleotide_lengths`. Returns an empty + list when no typed P-plane is authored — the bundled + cartridges' default. + """ + rm = getattr(cfg, "reference_models", None) + if rm is None: + return [] + plane = getattr(rm, "p_nucleotide_lengths", None) or {} + return sorted(plane.keys()) + + +def build_manifest(cfg, refdata=None, *, schema_version): + errors: List[str] = [] + + # Identity — assembled from BOTH the Python-side ConfigInfo + # (species, reference_set, chain-type-derived data) and the + # bridged refdata.identity() (which exposes locus, source, + # plus the curation source-tag suffix when applicable). The + # bridged identity is the authoritative answer for fields + # the engine consumes; the ConfigInfo fields fill gaps the + # bridge doesn't surface. + meta = cfg.metadata + identity_dict: Dict[str, Any] = { + "name": _jsonable(cfg.name), + "species": _jsonable(getattr(meta, "species", None) if meta else None), + "locus": None, # populated from bridged identity below + "reference_set": _jsonable( + getattr(meta, "reference_set", None) if meta else None + ), + "source": None, # populated from bridged identity below + } + + # Catalogue counts + functional-status histogram. These read + # only the Python-side allele lists; no bridge needed. + catalogue_dict: Dict[str, Any] = { + "v_count": cfg.number_of_v_alleles, + "d_count": cfg.number_of_d_alleles, + "j_count": cfg.number_of_j_alleles, + "c_count": cfg.number_of_c_alleles, + "functional_status_counts": cfg.functional_status_counts(), + } + + # Models plane — pure Python; no bridge. + ref_models = getattr(cfg, "reference_models", None) + has_reference_models = ref_models is not None + np_length_keys: List[str] = [] + trim_keys: List[str] = [] + # Typed NP base model inventory (Slice — Typed NP base + # model). ``np_base_models`` is keyed by NP region + # (``"NP1"`` / ``"NP2"``) and surfaces the per-region + # ``kind``. Cartridges without typed NP base models report + # ``kind="uniform"`` implicitly via the empty / absent + # entry; the manifest summary block below makes the + # opt-in / fallback explicit so a consumer can detect + # which cartridges author non-uniform NP biology. + np_base_models: Dict[str, Any] = {} + if has_reference_models: + np_length_keys = sorted((ref_models.np_lengths or {}).keys()) + trim_keys = sorted((ref_models.trims or {}).keys()) + np_bases_dict = getattr(ref_models, "np_bases", None) or {} + for key in sorted(np_bases_dict.keys()): + spec = np_bases_dict[key] + np_base_models[key] = {"kind": spec.kind} + # SHM model inventory — documents which mutation models the + # engine ships and which S5F kernels are bundled. The kernel + # choice itself is a per-experiment parameter (passed to + # ``Experiment.mutate(...)``), NOT a per-cartridge property, + # so the manifest reports what's *available* rather than + # what was used. Explicitly carries ``in_content_hash=False`` + # so a user reading the manifest sees that two simulation + # runs using different S5F kernels would be indistinguishable + # by ``content_hash`` alone (the v1 boundary pinned by the + # SHM model audit §3). + from GenAIRR._s5f_loader import ( + DEFAULT_S5F_KERNEL, + available_s5f_kernels, + builtin_s5f_kernel_digest, + ) + + shm_dict: Dict[str, Any] = { + "available_models": ["uniform", "s5f"], + "s5f_kernels_available": available_s5f_kernels(), + "default_s5f_kernel": DEFAULT_S5F_KERNEL, + # Digest of the default kernel's bundled pickle bytes. + # ``None`` if the bytes can't be read (corrupted + # install) — the manifest doesn't raise. + "s5f_kernel_digest": builtin_s5f_kernel_digest(DEFAULT_S5F_KERNEL), + # Hard-coded ``False`` because the bridge doesn't fold + # the kernel choice into Rust ``content_hash``; the v1 + # boundary pinned in the audit. A future slice that + # closes the boundary updates BOTH this field and the + # companion audit pin. + "in_content_hash": False, + # Per-segment SHM rate scalars — the segment-rates + # slice. The manifest advertises the capability + the + # flat default; rate vectors are per-experiment (passed + # to ``Experiment.mutate(segment_rates=...)``) and + # therefore not part of cartridge identity. The + # ``in_content_hash`` field documents that the rate + # vector itself doesn't enter ``content_hash`` — same + # v1 boundary as the kernel choice. + "segment_rate_support": { + "available": True, + "buckets": ["V", "D", "J", "NP"], + "default": {"V": 1.0, "D": 1.0, "J": 1.0, "NP": 1.0}, + "in_content_hash": False, + }, + # V-region substructure annotation surface — the Slice 1 + # output of the V-subregion cartridge work. ``available`` + # is True as soon as the bridge derives subregions (so it + # toggles on for cartridges with IMGT ``gapped_seq``); + # ``annotated_v_count`` / ``total_v_count`` quantify + # coverage. The ``derivation`` tag identifies which + # source produced the annotations + # (``"bridge_imgt_gapped_seq"`` for the default helper, + # ``"explicit"`` if the user set ``allele.subregions`` + # directly). ``in_content_hash`` is True — the + # ``refdata_content_hash`` does fold subregion intervals + # in (see ``engine_rs/src/trace_file.rs``). Sampling + # rates (``v_subregion_rates``) and CDR/FR mutation + # counters do NOT live here; those are deliberately + # absent and pinned by the contract suite. + "v_subregion_support": { + "available": False, # populated below from bridged refdata + "labels": ["FWR1", "CDR1", "FWR2", "CDR2", "FWR3"], + "annotated_v_count": 0, + "total_v_count": cfg.number_of_v_alleles, + "derivation": "bridge_imgt_gapped_seq", + "in_content_hash": True, + }, + # V-subregion SHM rate-vector capability — the + # ``v_subregion_rates`` kwarg on ``Experiment.mutate`` + # (Slice B). The manifest advertises the kwarg vocabulary + # (five canonical labels + two aliases ``FWR`` / ``CDR``) + # plus the flat default. ``in_plan_signature=True`` + # because the rate vector flows into + # ``pass_plan_signature`` via Slice A's mutate-pass + # ``parameter_signature`` (mismatched rates fail replay + # at the signature gate). ``in_content_hash=False`` + # because rates are a per-experiment parameter, not a + # cartridge property — same boundary the + # ``segment_rate_support`` block documents. + "v_subregion_rate_support": { + "available": True, + "labels": ["FWR1", "CDR1", "FWR2", "CDR2", "FWR3"], + "aliases": { + "FWR": ["FWR1", "FWR2", "FWR3"], + "CDR": ["CDR1", "CDR2"], + }, + "default": { + "FWR1": 1.0, + "CDR1": 1.0, + "FWR2": 1.0, + "CDR2": 1.0, + "FWR3": 1.0, + }, + "in_plan_signature": True, + "in_content_hash": False, + }, + # V-subregion realised SHM mutation counters (Slice — + # ``docs/v_subregion_mutation_counters_audit.md``). Six + # AIRR fields that partition ``n_v_mutations`` by the + # assigned V allele's IMGT subregion intervals: five + # canonical labels plus ``n_v_unannotated_mutations`` + # for V events that can't be attributed. Source of + # truth is the engine's event ledger filtered to the + # mutate.{uniform,s5f} passes — the validator surfaces + # six per-field mismatch issues + a cross-field + # ``VSubregionMutationCountSumMismatch`` invariant. + # ``requires_annotations=True`` documents the + # dependency on the Slice-1 cartridge annotation + # surface — without subregion intervals, every V SHM + # event routes to the unannotated bucket. + # ``in_content_hash=False`` because counters are a + # per-record observable, not a cartridge property. + "v_subregion_counter_support": { + "available": True, + "fields": [ + "n_fwr1_mutations", + "n_cdr1_mutations", + "n_fwr2_mutations", + "n_cdr2_mutations", + "n_fwr3_mutations", + "n_v_unannotated_mutations", + ], + "partition_of": "n_v_mutations", + "requires_annotations": True, + "unannotated_bucket": "n_v_unannotated_mutations", + "in_content_hash": False, + }, + } + + models_dict: Dict[str, Any] = { + "has_reference_models": has_reference_models, + "np_length_keys": np_length_keys, + "trim_keys": trim_keys, + "legacy_np_lengths_present": bool(cfg.NP_lengths), + "legacy_trim_dicts_present": bool(cfg.trim_dicts), + "shm": shm_dict, + # V(D)J N-addition base sampling models (Slice — Typed + # NP base model). The block advertises which NP regions + # carry typed cartridge-owned base distributions; an + # empty ``np_base_models`` dict means every NP position + # samples uniformly A/C/G/T (the pre-slice default). + # ``legacy_fallback`` documents whether the resolver + # auto-lifted the legacy ``NP_transitions`` / + # ``NP_first_bases`` orphan fields — v1 keeps this + # ``False`` because auto-lift would silently change + # output bytes vs the pre-slice baseline; a follow-up + # slice may enable it under an explicit opt-in. The + # ``legacy_np_*_present`` flags surface the bundled + # cartridges' orphan data so a downstream tool can + # decide whether to author a typed spec from it. + # ``in_plan_signature=True`` because the base + # distribution's ``support()`` enters + # ``GenerateNPPass::parameter_signature`` via + # ``fmt_byte_dist`` (Slice A discipline). + # ``in_content_hash=False`` because the model is a + # per-experiment defaults plane, not cartridge + # identity — same boundary the per-segment SHM rates + # block documents. + "np_base_models": { + "models": np_base_models, + "legacy_fallback": False, + "legacy_np_transitions_present": bool( + getattr(cfg, "NP_transitions", None) + ), + "legacy_np_first_bases_present": bool( + getattr(cfg, "NP_first_bases", None) + ), + "supported_kinds": [ + "uniform", + "empirical_first_base", + "markov", + ], + "deferred_kinds": [], + "in_plan_signature": True, + "in_content_hash": False, + }, + # P-nucleotide per-end length distribution inventory + # (Slice — P-nucleotide v1). `length_keys` advertises + # which V(D)J coding-end junction sides the cartridge + # authored a typed `EmpiricalDistributionSpec` for. + # Empty list (the bundled-cartridge default) means + # the pipeline omits every P-pass — byte-identical + # to the pre-slice baseline. + # `legacy_p_nucleotide_length_probs_present` surfaces + # the bundled cartridges' orphan dict so a downstream + # tool can decide whether to author a typed plane + # from it; `legacy_fallback=False` documents that + # auto-lift is NOT enabled (same boundary the Markov + # slice held for `NP_transitions`). + # `in_plan_signature=True` because each `PAdditionPass` + # folds its per-end length distribution via + # `fmt_int_dist` into `parameter_signature`, so + # replay against a different P-length distribution + # fails the signature gate. + "p_nucleotide_models": { + "length_keys": _p_nucleotide_length_keys(cfg), + "legacy_p_nucleotide_length_probs_present": bool( + getattr(cfg, "p_nucleotide_length_probs", None) + ), + "legacy_fallback": False, + "supported_ends": ["V_3", "D_5", "D_3", "J_5"], + "in_plan_signature": True, + "in_content_hash": False, + }, + # Per-segment allele-usage weights (Slice — Allele + # Usage Estimation v1). ``available`` is ``True`` only + # when the cartridge author / builder attached a + # non-``None`` ``allele_usage`` spec on + # ``reference_models``. ``nonempty_segments`` lists + # the segments whose dicts are non-empty so a + # downstream consumer can decide whether the + # cartridge actually overrides V / D / J or only some + # of them. ``legacy_gene_use_dict_present`` surfaces + # the bundled cartridges' orphan dict so a downstream + # tool can decide whether to author a typed plane + # from it; ``legacy_fallback=False`` documents that + # auto-lift is NOT enabled (same boundary the Markov / + # P-nucleotide slices respected). ``in_plan_signature`` + # is the documented soft gap 1 from the plan-signature + # completeness audit — the per-experiment kwarg + # surface AND this cartridge-driven path both fold + # through `SampleAllelePass.parameter_signature` which + # returns the empty string; tightening is a separate + # slice. + "allele_usage": _allele_usage_manifest_block(cfg), + # Trim Distribution Estimation v1 — typed-plane + # `trims` block. ``in_plan_signature=True`` because + # trim distributions already fold into the plan + # signature via ``fmt_int_dist`` (unlike the + # allele-usage soft gap). The minimal top-level + # ``trim_keys`` / ``legacy_trim_dicts_present`` + # entries above remain for backwards compatibility; + # this block is the structured surface. + "trim_models": _trim_models_manifest_block(cfg), + # NP Length Distribution Estimation v1 — + # typed-plane `np_lengths` block. Same shape / + # discipline as `trim_models`: ``in_plan_signature + # =True`` (folded via ``GenerateNPPass.parameter_signature``), + # ``legacy_fallback=False`` (no auto-lift of the + # legacy ``DataConfig.NP_lengths`` orphan dict), + # additive on top of the minimal top-level + # ``np_length_keys`` / ``legacy_np_lengths_present`` + # entries above. + "np_length_models": _np_length_models_manifest_block(cfg), + # Donor-population germline prior plane (Slice — Cartridge genotype + # plane). Read from the top-level ``DataConfig.genotype_priors`` + # field (NOT ``reference_models``); ``source_field`` records that. + "genotype_priors": _genotype_priors_manifest_block(cfg), + } + + # Bridge once (or accept the provided refdata) to read the + # rules plane + content_hash + identity.source curation tag. + bridged_refdata = refdata + if bridged_refdata is None: + try: + from GenAIRR._refdata_resolver import dataconfig_to_refdata + bridged_refdata = dataconfig_to_refdata(cfg) + except Exception as exc: + errors.append( + f"dataconfig_to_refdata failed: " + f"{type(exc).__name__}: {exc}" + ) + bridged_refdata = None + + rules_dict: Dict[str, Any] = { + "has_explicit_rules": getattr(cfg, "reference_rules", None) + is not None, + "allowed_bases": None, + "v_anchor": None, + "j_anchor": None, + } + refdata_content_hash: Optional[str] = None + curation_source_tag: Optional[str] = None + curation_policies: List[str] = [] + + if bridged_refdata is not None: + # Rules plane — three Rust accessors: + # ``allowed_bases()`` + ``v_anchor_rule()`` + + # ``j_anchor_rule()``. Each returns a JSON-friendly dict + # / list; ``_normalise_anchor_rule`` coerces tuples to + # lists for ``json.dumps`` safety. + try: + rules_dict["allowed_bases"] = list( + bridged_refdata.allowed_bases() + ) + except Exception as exc: + errors.append( + f"refdata.allowed_bases() failed: " + f"{type(exc).__name__}: {exc}" + ) + try: + rules_dict["v_anchor"] = _normalise_anchor_rule( + bridged_refdata.v_anchor_rule() + ) + except Exception as exc: + errors.append( + f"refdata.v_anchor_rule() failed: " + f"{type(exc).__name__}: {exc}" + ) + try: + rules_dict["j_anchor"] = _normalise_anchor_rule( + bridged_refdata.j_anchor_rule() + ) + except Exception as exc: + errors.append( + f"refdata.j_anchor_rule() failed: " + f"{type(exc).__name__}: {exc}" + ) + try: + refdata_content_hash = bridged_refdata.content_hash() + except Exception as exc: + errors.append( + f"refdata.content_hash() failed: " + f"{type(exc).__name__}: {exc}" + ) + # V-subregion coverage — walk the bridged V pool and + # count alleles with non-empty subregion lists. Failures + # here go through the manifest's standard ``errors`` + # channel rather than raising — the rest of the manifest + # is still useful even if coverage stats can't be read. + try: + annotated_v = 0 + v_total = bridged_refdata.v_pool_size() + for v_id in range(v_total): + if bridged_refdata.v_allele(v_id).subregions: + annotated_v += 1 + shm_dict["v_subregion_support"]["available"] = ( + annotated_v > 0 + ) + shm_dict["v_subregion_support"]["annotated_v_count"] = ( + annotated_v + ) + shm_dict["v_subregion_support"]["total_v_count"] = v_total + except Exception as exc: + errors.append( + f"v_subregion_support coverage failed: " + f"{type(exc).__name__}: {exc}" + ) + # Identity (locus + source) from the bridged side, plus + # curation parsing. The Rust curation path appends + # ``|curated:`` per applied policy; the + # policy tag itself may contain ``|`` (e.g. + # ``functional_status:functional|keep_unannotated=true``). + # So we split on the ``|curated:`` separator after the + # first base-source piece. The base source becomes + # ``curation.source_tag``; the trailing policy tags + # become ``curation.policies``. + try: + ident = bridged_refdata.identity() or {} + bridged_locus = ident.get("locus") + if bridged_locus is not None: + identity_dict["locus"] = _jsonable(bridged_locus) + source_tag = ident.get("source") + curation_source_tag = source_tag + if source_tag is not None: + identity_dict["source"] = _jsonable(source_tag) + if source_tag and "|curated:" in source_tag: + pieces = source_tag.split("|curated:") + curation_policies = pieces[1:] + except Exception as exc: + errors.append( + f"refdata.identity() failed: {type(exc).__name__}: {exc}" + ) + + curation_dict: Dict[str, Any] = { + "source_tag": curation_source_tag, + "policies": curation_policies, + } + + hashes_dict: Dict[str, Any] = { + "data_config_checksum": cfg.compute_checksum(), + "refdata_content_hash": refdata_content_hash, + } + + # Documented completeness gaps from the audit — listed + # verbatim so a user inspecting a manifest sees which + # surfaces aren't covered by the views / the bridge. + dropped_allele_fields = list(_DOCUMENTED_DROPPED_ALLELE_FIELDS) + orphan_dataconfig_fields = list(_DOCUMENTED_ORPHAN_DATACONFIG_FIELDS) + + manifest: Dict[str, Any] = { + "schema_version": getattr(cfg, "schema_version", schema_version), + "identity": identity_dict, + "catalogue": catalogue_dict, + "rules": rules_dict, + "models": models_dict, + "curation": curation_dict, + "hashes": hashes_dict, + "dropped_allele_fields": dropped_allele_fields, + "orphan_dataconfig_fields": orphan_dataconfig_fields, + } + if errors: + manifest["errors"] = errors + return manifest diff --git a/src/GenAIRR/dataconfig/config_info.py b/src/GenAIRR/dataconfig/config_info.py index 76ed96d..f7cb6a6 100644 --- a/src/GenAIRR/dataconfig/config_info.py +++ b/src/GenAIRR/dataconfig/config_info.py @@ -2,7 +2,7 @@ import pickle from dataclasses import dataclass from datetime import date -from .enums import * +from .enums import ChainType, Species @dataclass(frozen=True) class ConfigInfo: diff --git a/src/GenAIRR/dataconfig/data_config.py b/src/GenAIRR/dataconfig/data_config.py index 0a1c3c3..ee92f7e 100644 --- a/src/GenAIRR/dataconfig/data_config.py +++ b/src/GenAIRR/dataconfig/data_config.py @@ -1,5 +1,7 @@ +from __future__ import annotations + from dataclasses import dataclass, field -from typing import Optional, Dict, List, Any +from typing import TYPE_CHECKING, Optional, Dict, List, Any from GenAIRR.dataconfig.cartridge_views import ( CartridgeCatalogueView, @@ -9,13 +11,27 @@ ) from GenAIRR.dataconfig.config_info import ConfigInfo from GenAIRR.dataconfig.enums import ChainType +from GenAIRR.dataconfig._manifest import build_manifest +# Backward-compat re-export: these two documented-gap tuples moved to +# ``_manifest`` but are still imported from ``data_config`` by callers +# (e.g. tests/test_cartridge_manifest.py). Re-exported to preserve the +# import surface without duplicating the definitions. +from GenAIRR.dataconfig._manifest import ( + _DOCUMENTED_DROPPED_ALLELE_FIELDS, + _DOCUMENTED_ORPHAN_DATACONFIG_FIELDS, +) import copy import hashlib import pickle from GenAIRR.alleles.allele import Allele -from GenAIRR.reference_models import ReferenceEmpiricalModels -from GenAIRR.reference_rules import ReferenceRulesSpec -from GenAIRR.genotype_priors import PopulationGenotypeModel +if TYPE_CHECKING: + # Authoring-layer types used only as field annotations. Deferred to + # TYPE_CHECKING so loading a pickled DataConfig (every bundled + # cartridge is one) does not eagerly force-import the entire + # cartridge-authoring subtree. + from GenAIRR.reference_models import ReferenceEmpiricalModels + from GenAIRR.reference_rules import ReferenceRulesSpec + from GenAIRR.genotype_priors import PopulationGenotypeModel DEFAULT_P_NUCLEOTIDE_LENGTH_PROBS = {0: 0.50, 1: 0.25, 2: 0.15, 3: 0.07, 4: 0.03} @@ -26,278 +42,6 @@ SCHEMA_VERSION = 1 -# Documented completeness gaps surfaced verbatim by -# ``DataConfig.cartridge_manifest``. See -# ``docs/reference_cartridge_completeness_audit.md`` §1 (orphan -# fields) and §2 (dropped allele fields) for the rationale on each -# entry. These are tuples so manifest output stays stable across -# calls and across Python sessions. -_DOCUMENTED_DROPPED_ALLELE_FIELDS = ( - "aliases", - "anchor_meta", - "gapped_seq", - "family", - "locus", - "species", - "source", -) -_DOCUMENTED_ORPHAN_DATACONFIG_FIELDS = ( - "gene_use_dict", - "NP_transitions", - "NP_first_bases", - "correction_maps", - "asc_tables", - "p_nucleotide_length_probs", - "dj_pairing_map", -) - - -def _jsonable(value): - """Coerce a value to a JSON-clean form. Enums become their - ``.name`` (uppercase canonical) for stability; tuples become - lists; everything else passes through. Used by - :meth:`DataConfig.cartridge_manifest` to make the identity - plane safe for ``json.dumps``.""" - if value is None: - return None - if isinstance(value, (str, int, float, bool)): - return value - if hasattr(value, "name") and not isinstance(value, type): - # Enum-shaped values (Species, ChainType, etc.). - return value.name - if isinstance(value, (tuple, list)): - return [_jsonable(v) for v in value] - return str(value) - - -def _normalise_anchor_rule(rule): - """Coerce an anchor rule dict from the Rust ``rules()`` accessor - into a JSON-clean shape. Returns ``None`` for missing rules. - - The Rust side may surface the rule as a dict with ``required``, - ``expected_aa``, ``missing_severity``, ``mismatch_severity`` - keys (string severities). We pass it through but coerce any - iterable fields to plain ``list`` to keep ``json.dumps`` happy. - """ - if rule is None: - return None - if isinstance(rule, dict): - out = dict(rule) - if "expected_aa" in out and out["expected_aa"] is not None: - out["expected_aa"] = list(out["expected_aa"]) - return out - # Unknown shape — drop silently rather than raise; the manifest - # error list isn't the right channel for "this is unexpected - # but harmless". - return None - - -def _allele_usage_manifest_block(cfg): - """Build the ``models.allele_usage`` manifest block per the - Allele Usage Estimation v1 audit (§8 of - ``docs/allele_usage_estimation_design.md``). - - Returns a JSON-clean dict carrying: - - - ``available``: whether ``reference_models.allele_usage`` is - a non-``None`` spec. - - ``segments``: the canonical V/D/J label list. - - ``nonempty_segments``: subset of ``segments`` whose - author-supplied dict is non-empty. - - ``legacy_gene_use_dict_present``: documents the orphan - legacy dict's bundled-cartridge presence so a downstream - tool can decide whether to author a typed plane from it. - - ``legacy_fallback``: always ``False`` in v1 (no auto-lift). - - ``in_plan_signature``: ``False`` — inherits soft gap 1 - from the plan-signature completeness audit. Tightening is - a separate slice covering BOTH the kwarg AND the - cartridge-driven path in lockstep. - - ``source``: which Python field the manifest block reads - from. Always ``"ReferenceEmpiricalModels.allele_usage"``. - """ - rm = getattr(cfg, "reference_models", None) - spec = getattr(rm, "allele_usage", None) if rm is not None else None - available = spec is not None - nonempty: list = [] - if available: - try: - nonempty = list(spec.nonempty_segments()) - except Exception: - nonempty = [] - return { - "available": available, - "segments": ["V", "D", "J"], - "nonempty_segments": nonempty, - "legacy_gene_use_dict_present": bool( - getattr(cfg, "gene_use_dict", None) - ), - "legacy_fallback": False, - "in_plan_signature": False, # documented soft gap 1 - "source": "ReferenceEmpiricalModels.allele_usage", - } - - -def _genotype_priors_manifest_block(cfg): - """Build the ``models.genotype_priors`` manifest block (Slice — Cartridge - genotype plane). Audit-sized: counts and identity, never the full tables. - Reads the top-level ``DataConfig.genotype_priors`` plane (independent of - ``reference_models``).""" - def _empty(available, valid, error=None): - b = { - "available": available, - "valid": valid, - "model_id": None, - "source": None, - "version": None, - "model_checksum": None, - "segments_with_frequencies": [], - "freq_gene_counts": {"V": 0, "D": 0, "J": 0}, - "deletion_gene_counts": {"V": 0, "D": 0, "J": 0}, - "novel_allele_count": 0, - "chromosome_weights": None, - "source_field": "DataConfig.genotype_priors", - } - if error is not None: - b["validation_error"] = error - return b - - model = getattr(cfg, "genotype_priors", None) - if model is None: - return _empty(False, None) - if not isinstance(model, PopulationGenotypeModel): - # Field is typed Optional[PopulationGenotypeModel] but Python won't enforce - # it; a garbage value must not crash the manifest (called from build()). - return _empty(True, False, - f"genotype_priors is not a PopulationGenotypeModel " - f"(got {type(model).__name__})") - - # A plane may have been attached directly (bypassing builder validation). - # Report validity rather than leaking non-JSON-clean numerics (e.g. NaN - # chromosome weights) into the manifest. - try: - model.validate( - chain_type=getattr(getattr(cfg, "metadata", None), "chain_type", None)) - valid, validation_error = True, None - except ValueError as exc: - valid, validation_error = False, str(exc) - - # content_checksum / float() can themselves raise on a malformed model; never - # let that crash the manifest — report None and the validity flag instead. - try: - checksum = model.content_checksum() - except Exception: - checksum = None - cw = None - if valid: - try: - cw = [float(model.chromosome_weights[0]), float(model.chromosome_weights[1])] - except Exception: - cw = None - freq = model.allele_frequencies if isinstance(model.allele_frequencies, dict) else {} - dele = model.haplotype_deletion_prob if isinstance(model.haplotype_deletion_prob, dict) else {} - novels = model.novel_alleles if isinstance(model.novel_alleles, (list, tuple)) else [] - block = { - "available": True, - "valid": valid, - "model_id": (model.model_id or None) if isinstance(model.model_id, str) else None, - "source": (model.source or None) if isinstance(model.source, str) else None, - "version": (model.version or None) if isinstance(model.version, str) else None, - "model_checksum": checksum, - "segments_with_frequencies": [s for s in ("V", "D", "J") if freq.get(s)], - "freq_gene_counts": {s: len(freq.get(s, {})) for s in ("V", "D", "J")}, - "deletion_gene_counts": {s: len(dele.get(s, {})) for s in ("V", "D", "J")}, - "novel_allele_count": len(novels), - "chromosome_weights": cw, - "source_field": "DataConfig.genotype_priors", - } - if validation_error is not None: - block["validation_error"] = validation_error - return block - - -def _np_length_models_manifest_block(cfg): - """Build the ``models.np_length_models`` manifest block - per the NP Length Distribution Estimation v1 audit - (§8 of ``docs/np_length_estimation_design.md``). - - Returns a JSON-clean dict carrying: - - - ``keys``: sorted list of NP-key strings present on - the typed ``ReferenceEmpiricalModels.np_lengths`` - plane (subset of ``("NP1", "NP2")``). - - ``source``: which Python field the manifest block - reads from. Always ``"ReferenceEmpiricalModels.np_lengths"``. - - ``in_plan_signature``: ``True`` — NP-length - distributions fold into the plan signature via - ``GenerateNPPass.parameter_signature`` / - ``fmt_int_dist`` (verified at audit time). **No - soft gap inherited.** - - ``legacy_np_lengths_present``: documents the bundled - cartridges' legacy nested-dict presence so a - downstream tool can decide whether to author a - typed plane from it. - - ``legacy_fallback``: always ``False`` in v1 (no - auto-lift). - """ - rm = getattr(cfg, "reference_models", None) - np_lengths = getattr(rm, "np_lengths", None) if rm is not None else None - keys = sorted(np_lengths.keys()) if np_lengths else [] - return { - "keys": keys, - "source": "ReferenceEmpiricalModels.np_lengths", - "in_plan_signature": True, - "legacy_np_lengths_present": bool(getattr(cfg, "NP_lengths", None)), - "legacy_fallback": False, - } - - -def _trim_models_manifest_block(cfg): - """Build the ``models.trim_models`` manifest block per the - Trim Distribution Estimation v1 audit (§8 of - ``docs/trim_distribution_estimation_design.md``). - - Returns a JSON-clean dict carrying: - - - ``keys``: sorted list of trim-key strings present on the - typed ``ReferenceEmpiricalModels.trims`` plane (subset of - ``("V_3", "D_5", "D_3", "J_5")``). - - ``source``: which Python field the manifest block reads - from. Always ``"ReferenceEmpiricalModels.trims"``. - - ``in_plan_signature``: ``True`` — trim distributions fold - into the plan signature via ``fmt_int_dist`` (verified at - audit time). **No soft gap inherited** — unlike the - allele-usage slice. - - ``legacy_trim_dicts_present``: documents the bundled - cartridges' legacy nested-dict presence so a downstream - tool can decide whether to author a typed plane from it. - - ``legacy_fallback``: always ``False`` in v1 (no auto-lift). - """ - rm = getattr(cfg, "reference_models", None) - trims = getattr(rm, "trims", None) if rm is not None else None - keys = sorted(trims.keys()) if trims else [] - return { - "keys": keys, - "source": "ReferenceEmpiricalModels.trims", - "in_plan_signature": True, - "legacy_trim_dicts_present": bool(getattr(cfg, "trim_dicts", None)), - "legacy_fallback": False, - } - - -def _p_nucleotide_length_keys(cfg): - """Return the sorted list of P-end keys that have a typed - `EmpiricalDistributionSpec` on - `cfg.reference_models.p_nucleotide_lengths`. Returns an empty - list when no typed P-plane is authored — the bundled - cartridges' default. - """ - rm = getattr(cfg, "reference_models", None) - if rm is None: - return [] - plane = getattr(rm, "p_nucleotide_lengths", None) or {} - return sorted(plane.keys()) - - class DataConfigError(Exception): """Raised when a DataConfig fails validation.""" @@ -483,8 +227,7 @@ def verify_integrity(self) -> None: raise DataConfigError( f"DataConfig '{self.name or 'Unnamed'}' has schema_version=" f"{version}, but the current code expects " - f"{SCHEMA_VERSION}. Re-migrate this pickle (see " - f".private/scripts/migrate_pkl_to_v1.py) or rebuild it " + f"{SCHEMA_VERSION}. Re-migrate or rebuild this pickle " f"from the IMGT source via dataconfig builders." ) stored = getattr(self, "schema_sha256", "") or "" @@ -718,451 +461,7 @@ def cartridge_manifest( ``manifest["errors"]`` to decide whether the cartridge is actually usable. """ - errors: List[str] = [] - - # Identity — assembled from BOTH the Python-side ConfigInfo - # (species, reference_set, chain-type-derived data) and the - # bridged refdata.identity() (which exposes locus, source, - # plus the curation source-tag suffix when applicable). The - # bridged identity is the authoritative answer for fields - # the engine consumes; the ConfigInfo fields fill gaps the - # bridge doesn't surface. - meta = self.metadata - identity_dict: Dict[str, Any] = { - "name": _jsonable(self.name), - "species": _jsonable(getattr(meta, "species", None) if meta else None), - "locus": None, # populated from bridged identity below - "reference_set": _jsonable( - getattr(meta, "reference_set", None) if meta else None - ), - "source": None, # populated from bridged identity below - } - - # Catalogue counts + functional-status histogram. These read - # only the Python-side allele lists; no bridge needed. - catalogue_dict: Dict[str, Any] = { - "v_count": self.number_of_v_alleles, - "d_count": self.number_of_d_alleles, - "j_count": self.number_of_j_alleles, - "c_count": self.number_of_c_alleles, - "functional_status_counts": self.functional_status_counts(), - } - - # Models plane — pure Python; no bridge. - ref_models = getattr(self, "reference_models", None) - has_reference_models = ref_models is not None - np_length_keys: List[str] = [] - trim_keys: List[str] = [] - # Typed NP base model inventory (Slice — Typed NP base - # model). ``np_base_models`` is keyed by NP region - # (``"NP1"`` / ``"NP2"``) and surfaces the per-region - # ``kind``. Cartridges without typed NP base models report - # ``kind="uniform"`` implicitly via the empty / absent - # entry; the manifest summary block below makes the - # opt-in / fallback explicit so a consumer can detect - # which cartridges author non-uniform NP biology. - np_base_models: Dict[str, Any] = {} - if has_reference_models: - np_length_keys = sorted((ref_models.np_lengths or {}).keys()) - trim_keys = sorted((ref_models.trims or {}).keys()) - np_bases_dict = getattr(ref_models, "np_bases", None) or {} - for key in sorted(np_bases_dict.keys()): - spec = np_bases_dict[key] - np_base_models[key] = {"kind": spec.kind} - # SHM model inventory — documents which mutation models the - # engine ships and which S5F kernels are bundled. The kernel - # choice itself is a per-experiment parameter (passed to - # ``Experiment.mutate(...)``), NOT a per-cartridge property, - # so the manifest reports what's *available* rather than - # what was used. Explicitly carries ``in_content_hash=False`` - # so a user reading the manifest sees that two simulation - # runs using different S5F kernels would be indistinguishable - # by ``content_hash`` alone (the v1 boundary pinned by the - # SHM model audit §3). - from GenAIRR._s5f_loader import ( - DEFAULT_S5F_KERNEL, - available_s5f_kernels, - builtin_s5f_kernel_digest, - ) - - shm_dict: Dict[str, Any] = { - "available_models": ["uniform", "s5f"], - "s5f_kernels_available": available_s5f_kernels(), - "default_s5f_kernel": DEFAULT_S5F_KERNEL, - # Digest of the default kernel's bundled pickle bytes. - # ``None`` if the bytes can't be read (corrupted - # install) — the manifest doesn't raise. - "s5f_kernel_digest": builtin_s5f_kernel_digest(DEFAULT_S5F_KERNEL), - # Hard-coded ``False`` because the bridge doesn't fold - # the kernel choice into Rust ``content_hash``; the v1 - # boundary pinned in the audit. A future slice that - # closes the boundary updates BOTH this field and the - # companion audit pin. - "in_content_hash": False, - # Per-segment SHM rate scalars — the segment-rates - # slice. The manifest advertises the capability + the - # flat default; rate vectors are per-experiment (passed - # to ``Experiment.mutate(segment_rates=...)``) and - # therefore not part of cartridge identity. The - # ``in_content_hash`` field documents that the rate - # vector itself doesn't enter ``content_hash`` — same - # v1 boundary as the kernel choice. - "segment_rate_support": { - "available": True, - "buckets": ["V", "D", "J", "NP"], - "default": {"V": 1.0, "D": 1.0, "J": 1.0, "NP": 1.0}, - "in_content_hash": False, - }, - # V-region substructure annotation surface — the Slice 1 - # output of the V-subregion cartridge work. ``available`` - # is True as soon as the bridge derives subregions (so it - # toggles on for cartridges with IMGT ``gapped_seq``); - # ``annotated_v_count`` / ``total_v_count`` quantify - # coverage. The ``derivation`` tag identifies which - # source produced the annotations - # (``"bridge_imgt_gapped_seq"`` for the default helper, - # ``"explicit"`` if the user set ``allele.subregions`` - # directly). ``in_content_hash`` is True — the - # ``refdata_content_hash`` does fold subregion intervals - # in (see ``engine_rs/src/trace_file.rs``). Sampling - # rates (``v_subregion_rates``) and CDR/FR mutation - # counters do NOT live here; those are deliberately - # absent and pinned by the contract suite. - "v_subregion_support": { - "available": False, # populated below from bridged refdata - "labels": ["FWR1", "CDR1", "FWR2", "CDR2", "FWR3"], - "annotated_v_count": 0, - "total_v_count": self.number_of_v_alleles, - "derivation": "bridge_imgt_gapped_seq", - "in_content_hash": True, - }, - # V-subregion SHM rate-vector capability — the - # ``v_subregion_rates`` kwarg on ``Experiment.mutate`` - # (Slice B). The manifest advertises the kwarg vocabulary - # (five canonical labels + two aliases ``FWR`` / ``CDR``) - # plus the flat default. ``in_plan_signature=True`` - # because the rate vector flows into - # ``pass_plan_signature`` via Slice A's mutate-pass - # ``parameter_signature`` (mismatched rates fail replay - # at the signature gate). ``in_content_hash=False`` - # because rates are a per-experiment parameter, not a - # cartridge property — same boundary the - # ``segment_rate_support`` block documents. - "v_subregion_rate_support": { - "available": True, - "labels": ["FWR1", "CDR1", "FWR2", "CDR2", "FWR3"], - "aliases": { - "FWR": ["FWR1", "FWR2", "FWR3"], - "CDR": ["CDR1", "CDR2"], - }, - "default": { - "FWR1": 1.0, - "CDR1": 1.0, - "FWR2": 1.0, - "CDR2": 1.0, - "FWR3": 1.0, - }, - "in_plan_signature": True, - "in_content_hash": False, - }, - # V-subregion realised SHM mutation counters (Slice — - # ``docs/v_subregion_mutation_counters_audit.md``). Six - # AIRR fields that partition ``n_v_mutations`` by the - # assigned V allele's IMGT subregion intervals: five - # canonical labels plus ``n_v_unannotated_mutations`` - # for V events that can't be attributed. Source of - # truth is the engine's event ledger filtered to the - # mutate.{uniform,s5f} passes — the validator surfaces - # six per-field mismatch issues + a cross-field - # ``VSubregionMutationCountSumMismatch`` invariant. - # ``requires_annotations=True`` documents the - # dependency on the Slice-1 cartridge annotation - # surface — without subregion intervals, every V SHM - # event routes to the unannotated bucket. - # ``in_content_hash=False`` because counters are a - # per-record observable, not a cartridge property. - "v_subregion_counter_support": { - "available": True, - "fields": [ - "n_fwr1_mutations", - "n_cdr1_mutations", - "n_fwr2_mutations", - "n_cdr2_mutations", - "n_fwr3_mutations", - "n_v_unannotated_mutations", - ], - "partition_of": "n_v_mutations", - "requires_annotations": True, - "unannotated_bucket": "n_v_unannotated_mutations", - "in_content_hash": False, - }, - } - - models_dict: Dict[str, Any] = { - "has_reference_models": has_reference_models, - "np_length_keys": np_length_keys, - "trim_keys": trim_keys, - "legacy_np_lengths_present": bool(self.NP_lengths), - "legacy_trim_dicts_present": bool(self.trim_dicts), - "shm": shm_dict, - # V(D)J N-addition base sampling models (Slice — Typed - # NP base model). The block advertises which NP regions - # carry typed cartridge-owned base distributions; an - # empty ``np_base_models`` dict means every NP position - # samples uniformly A/C/G/T (the pre-slice default). - # ``legacy_fallback`` documents whether the resolver - # auto-lifted the legacy ``NP_transitions`` / - # ``NP_first_bases`` orphan fields — v1 keeps this - # ``False`` because auto-lift would silently change - # output bytes vs the pre-slice baseline; a follow-up - # slice may enable it under an explicit opt-in. The - # ``legacy_np_*_present`` flags surface the bundled - # cartridges' orphan data so a downstream tool can - # decide whether to author a typed spec from it. - # ``in_plan_signature=True`` because the base - # distribution's ``support()`` enters - # ``GenerateNPPass::parameter_signature`` via - # ``fmt_byte_dist`` (Slice A discipline). - # ``in_content_hash=False`` because the model is a - # per-experiment defaults plane, not cartridge - # identity — same boundary the per-segment SHM rates - # block documents. - "np_base_models": { - "models": np_base_models, - "legacy_fallback": False, - "legacy_np_transitions_present": bool( - getattr(self, "NP_transitions", None) - ), - "legacy_np_first_bases_present": bool( - getattr(self, "NP_first_bases", None) - ), - "supported_kinds": [ - "uniform", - "empirical_first_base", - "markov", - ], - "deferred_kinds": [], - "in_plan_signature": True, - "in_content_hash": False, - }, - # P-nucleotide per-end length distribution inventory - # (Slice — P-nucleotide v1). `length_keys` advertises - # which V(D)J coding-end junction sides the cartridge - # authored a typed `EmpiricalDistributionSpec` for. - # Empty list (the bundled-cartridge default) means - # the pipeline omits every P-pass — byte-identical - # to the pre-slice baseline. - # `legacy_p_nucleotide_length_probs_present` surfaces - # the bundled cartridges' orphan dict so a downstream - # tool can decide whether to author a typed plane - # from it; `legacy_fallback=False` documents that - # auto-lift is NOT enabled (same boundary the Markov - # slice held for `NP_transitions`). - # `in_plan_signature=True` because each `PAdditionPass` - # folds its per-end length distribution via - # `fmt_int_dist` into `parameter_signature`, so - # replay against a different P-length distribution - # fails the signature gate. - "p_nucleotide_models": { - "length_keys": _p_nucleotide_length_keys(self), - "legacy_p_nucleotide_length_probs_present": bool( - getattr(self, "p_nucleotide_length_probs", None) - ), - "legacy_fallback": False, - "supported_ends": ["V_3", "D_5", "D_3", "J_5"], - "in_plan_signature": True, - "in_content_hash": False, - }, - # Per-segment allele-usage weights (Slice — Allele - # Usage Estimation v1). ``available`` is ``True`` only - # when the cartridge author / builder attached a - # non-``None`` ``allele_usage`` spec on - # ``reference_models``. ``nonempty_segments`` lists - # the segments whose dicts are non-empty so a - # downstream consumer can decide whether the - # cartridge actually overrides V / D / J or only some - # of them. ``legacy_gene_use_dict_present`` surfaces - # the bundled cartridges' orphan dict so a downstream - # tool can decide whether to author a typed plane - # from it; ``legacy_fallback=False`` documents that - # auto-lift is NOT enabled (same boundary the Markov / - # P-nucleotide slices respected). ``in_plan_signature`` - # is the documented soft gap 1 from the plan-signature - # completeness audit — the per-experiment kwarg - # surface AND this cartridge-driven path both fold - # through `SampleAllelePass.parameter_signature` which - # returns the empty string; tightening is a separate - # slice. - "allele_usage": _allele_usage_manifest_block(self), - # Trim Distribution Estimation v1 — typed-plane - # `trims` block. ``in_plan_signature=True`` because - # trim distributions already fold into the plan - # signature via ``fmt_int_dist`` (unlike the - # allele-usage soft gap). The minimal top-level - # ``trim_keys`` / ``legacy_trim_dicts_present`` - # entries above remain for backwards compatibility; - # this block is the structured surface. - "trim_models": _trim_models_manifest_block(self), - # NP Length Distribution Estimation v1 — - # typed-plane `np_lengths` block. Same shape / - # discipline as `trim_models`: ``in_plan_signature - # =True`` (folded via ``GenerateNPPass.parameter_signature``), - # ``legacy_fallback=False`` (no auto-lift of the - # legacy ``DataConfig.NP_lengths`` orphan dict), - # additive on top of the minimal top-level - # ``np_length_keys`` / ``legacy_np_lengths_present`` - # entries above. - "np_length_models": _np_length_models_manifest_block(self), - # Donor-population germline prior plane (Slice — Cartridge genotype - # plane). Read from the top-level ``DataConfig.genotype_priors`` - # field (NOT ``reference_models``); ``source_field`` records that. - "genotype_priors": _genotype_priors_manifest_block(self), - } - - # Bridge once (or accept the provided refdata) to read the - # rules plane + content_hash + identity.source curation tag. - bridged_refdata = refdata - if bridged_refdata is None: - try: - from GenAIRR._refdata_resolver import dataconfig_to_refdata - bridged_refdata = dataconfig_to_refdata(self) - except Exception as exc: - errors.append( - f"dataconfig_to_refdata failed: " - f"{type(exc).__name__}: {exc}" - ) - bridged_refdata = None - - rules_dict: Dict[str, Any] = { - "has_explicit_rules": getattr(self, "reference_rules", None) - is not None, - "allowed_bases": None, - "v_anchor": None, - "j_anchor": None, - } - refdata_content_hash: Optional[str] = None - curation_source_tag: Optional[str] = None - curation_policies: List[str] = [] - - if bridged_refdata is not None: - # Rules plane — three Rust accessors: - # ``allowed_bases()`` + ``v_anchor_rule()`` + - # ``j_anchor_rule()``. Each returns a JSON-friendly dict - # / list; ``_normalise_anchor_rule`` coerces tuples to - # lists for ``json.dumps`` safety. - try: - rules_dict["allowed_bases"] = list( - bridged_refdata.allowed_bases() - ) - except Exception as exc: - errors.append( - f"refdata.allowed_bases() failed: " - f"{type(exc).__name__}: {exc}" - ) - try: - rules_dict["v_anchor"] = _normalise_anchor_rule( - bridged_refdata.v_anchor_rule() - ) - except Exception as exc: - errors.append( - f"refdata.v_anchor_rule() failed: " - f"{type(exc).__name__}: {exc}" - ) - try: - rules_dict["j_anchor"] = _normalise_anchor_rule( - bridged_refdata.j_anchor_rule() - ) - except Exception as exc: - errors.append( - f"refdata.j_anchor_rule() failed: " - f"{type(exc).__name__}: {exc}" - ) - try: - refdata_content_hash = bridged_refdata.content_hash() - except Exception as exc: - errors.append( - f"refdata.content_hash() failed: " - f"{type(exc).__name__}: {exc}" - ) - # V-subregion coverage — walk the bridged V pool and - # count alleles with non-empty subregion lists. Failures - # here go through the manifest's standard ``errors`` - # channel rather than raising — the rest of the manifest - # is still useful even if coverage stats can't be read. - try: - annotated_v = 0 - v_total = bridged_refdata.v_pool_size() - for v_id in range(v_total): - if bridged_refdata.v_allele(v_id).subregions: - annotated_v += 1 - shm_dict["v_subregion_support"]["available"] = ( - annotated_v > 0 - ) - shm_dict["v_subregion_support"]["annotated_v_count"] = ( - annotated_v - ) - shm_dict["v_subregion_support"]["total_v_count"] = v_total - except Exception as exc: - errors.append( - f"v_subregion_support coverage failed: " - f"{type(exc).__name__}: {exc}" - ) - # Identity (locus + source) from the bridged side, plus - # curation parsing. The Rust curation path appends - # ``|curated:`` per applied policy; the - # policy tag itself may contain ``|`` (e.g. - # ``functional_status:functional|keep_unannotated=true``). - # So we split on the ``|curated:`` separator after the - # first base-source piece. The base source becomes - # ``curation.source_tag``; the trailing policy tags - # become ``curation.policies``. - try: - ident = bridged_refdata.identity() or {} - bridged_locus = ident.get("locus") - if bridged_locus is not None: - identity_dict["locus"] = _jsonable(bridged_locus) - source_tag = ident.get("source") - curation_source_tag = source_tag - if source_tag is not None: - identity_dict["source"] = _jsonable(source_tag) - if source_tag and "|curated:" in source_tag: - pieces = source_tag.split("|curated:") - curation_policies = pieces[1:] - except Exception as exc: - errors.append( - f"refdata.identity() failed: {type(exc).__name__}: {exc}" - ) - - curation_dict: Dict[str, Any] = { - "source_tag": curation_source_tag, - "policies": curation_policies, - } - - hashes_dict: Dict[str, Any] = { - "data_config_checksum": self.compute_checksum(), - "refdata_content_hash": refdata_content_hash, - } - - # Documented completeness gaps from the audit — listed - # verbatim so a user inspecting a manifest sees which - # surfaces aren't covered by the views / the bridge. - dropped_allele_fields = list(_DOCUMENTED_DROPPED_ALLELE_FIELDS) - orphan_dataconfig_fields = list(_DOCUMENTED_ORPHAN_DATACONFIG_FIELDS) - - manifest: Dict[str, Any] = { - "schema_version": getattr(self, "schema_version", SCHEMA_VERSION), - "identity": identity_dict, - "catalogue": catalogue_dict, - "rules": rules_dict, - "models": models_dict, - "curation": curation_dict, - "hashes": hashes_dict, - "dropped_allele_fields": dropped_allele_fields, - "orphan_dataconfig_fields": orphan_dataconfig_fields, - } - if errors: - manifest["errors"] = errors - return manifest + return build_manifest(self, refdata, schema_version=SCHEMA_VERSION) # --- Public Properties --- @property diff --git a/src/GenAIRR/dataconfig/enums.py b/src/GenAIRR/dataconfig/enums.py index 05ea2b8..a544c51 100644 --- a/src/GenAIRR/dataconfig/enums.py +++ b/src/GenAIRR/dataconfig/enums.py @@ -1,6 +1,4 @@ -from enum import Enum,auto - -from enum import Enum +from enum import Enum, auto class Productivity(Enum): diff --git a/src/GenAIRR/experiment.py b/src/GenAIRR/experiment.py index 1bc95e4..ce2a97e 100644 --- a/src/GenAIRR/experiment.py +++ b/src/GenAIRR/experiment.py @@ -41,7 +41,7 @@ _to_immutable_byte_pairs, _to_immutable_pairs, ) -from ._compile import ( +from ._lowering import ( _extract_invert_d_prob, _extract_paired_end_step, _extract_receptor_revision_prob, @@ -81,322 +81,27 @@ _ReceptorRevisionStep, _RepertoireForkStep, ) - - -# Sentinel for `using(...)` arguments that the caller did not pass at -# all. We can't use ``None`` for "unchanged" because ``None`` already -# means "clear the lock." -class _Unset: - __slots__ = () - - def __repr__(self) -> str: - return "" - - -_UNSET: _Unset = _Unset() - - -# ────────────────────────────────────────────────────────────────── -# Per-segment SHM rate validation (slice: segment_rates kwarg on -# Experiment.mutate). See ``docs/shm_segment_rate_design.md``. -# ────────────────────────────────────────────────────────────────── - -# Canonical bucket order — matches the Rust ``SegmentRateWeights`` -# struct field order so positional plumbing through PyO3 stays -# straightforward. -_SEGMENT_RATE_BUCKETS: Tuple[str, ...] = ("V", "D", "J", "NP") - -# Default flat-substrate rate vector (1.0 for every bucket). The -# pipeline-IR ``_MutateStep`` carries this verbatim when the user -# omits ``segment_rates``; the Rust passes detect the flat-default -# case and take the existing (pre-slice) fast path so legacy -# pipelines stay byte-identical. -_DEFAULT_SEGMENT_RATES: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0) - - -def _validate_segment_rates( - segment_rates: Optional[Dict[str, float]], -) -> Tuple[float, float, float, float]: - """Validate the user's ``segment_rates`` dict and return the - normalised ``(v, d, j, np)`` tuple. Pure helper — no DSL state. - - Validation: - - - ``None`` (or omitted) → flat default ``(1.0, 1.0, 1.0, 1.0)``. - - Keys must be a subset of ``{"V", "D", "J", "NP"}`` (case- - sensitive — matches the DSL spec). Other keys raise - ``ValueError`` naming the offending key. - - Values must be ``int`` / ``float`` (not bool), finite, and - ``>= 0``. Negative / NaN / inf raise ``ValueError``. - - Sparse: omitted keys default to ``1.0``. - - At least one effective rate must be strictly positive — an - all-zero (or all-omitted-then-explicitly-zero) configuration - would make the SHM pass a deterministic no-op, which is - almost certainly a builder bug. Reject with ``ValueError``. - """ - if segment_rates is None: - return _DEFAULT_SEGMENT_RATES - - if not isinstance(segment_rates, dict): - raise TypeError( - f"segment_rates must be a dict or None, got " - f"{type(segment_rates).__name__}" - ) - - # Reject unknown keys first so a typo surfaces with a clear - # message instead of silently defaulting. - unknown = sorted(set(segment_rates.keys()) - set(_SEGMENT_RATE_BUCKETS)) - if unknown: - raise ValueError( - f"segment_rates: unknown segment key(s) {unknown!r}. " - f"Allowed: {list(_SEGMENT_RATE_BUCKETS)!r}. " - "Keys are case-sensitive; 'V' / 'D' / 'J' for the V/D/J " - "segments and 'NP' for both Np1 and Np2." - ) - - out_list: List[float] = [] - for bucket in _SEGMENT_RATE_BUCKETS: - if bucket not in segment_rates: - out_list.append(1.0) - continue - value = segment_rates[bucket] - if isinstance(value, bool) or not isinstance(value, (int, float)): - raise TypeError( - f"segment_rates[{bucket!r}]: value must be a finite " - f"non-negative number, got {type(value).__name__}" - ) - value_f = float(value) - # NaN check FIRST — every comparison with NaN is False, so - # a `value_f < 0` test wouldn't catch it; that's the same - # rationale as the ``invert_d`` / ``rate`` NaN handling. - if value_f != value_f: - raise ValueError( - f"segment_rates[{bucket!r}]: value must be a finite " - "number, got NaN" - ) - if value_f < 0.0 or value_f == float("inf") or value_f == float("-inf"): - raise ValueError( - f"segment_rates[{bucket!r}]: value must be a finite " - f"non-negative number, got {value_f}" - ) - out_list.append(value_f) - - if sum(out_list) <= 0.0: - raise ValueError( - "segment_rates: at least one bucket must have a positive " - "rate. The supplied configuration zeroes out every " - "biological segment, which would make the SHM pass a " - "deterministic no-op." - ) - - return (out_list[0], out_list[1], out_list[2], out_list[3]) - - -# ────────────────────────────────────────────────────────────────── -# Per-V-subregion SHM rate validation (Slice B — -# `v_subregion_rates` kwarg on Experiment.mutate). See -# ``docs/v_subregion_shm_rate_design.md``. -# ────────────────────────────────────────────────────────────────── - -# Canonical label order — matches the Rust ``VSubregionRateWeights`` -# struct field order (FWR1 / CDR1 / FWR2 / CDR2 / FWR3) so -# positional plumbing through PyO3 stays straightforward. -_V_SUBREGION_RATE_LABELS: Tuple[str, ...] = ("FWR1", "CDR1", "FWR2", "CDR2", "FWR3") - -# Two-letter aliases that expand to a group of canonical labels. -# Resolution rule (audit §3): aliases expand first, then explicit -# labels override. So ``{"FWR": 0.5, "FWR2": 2.0}`` resolves to -# ``{FWR1: 0.5, FWR2: 2.0, FWR3: 0.5, CDR1: 1.0, CDR2: 1.0}``. -_V_SUBREGION_RATE_ALIASES: Dict[str, Tuple[str, ...]] = { - "FWR": ("FWR1", "FWR2", "FWR3"), - "CDR": ("CDR1", "CDR2"), -} - -# Default flat-substrate rate vector (1.0 for every label) — same -# fast-path discipline as ``_DEFAULT_SEGMENT_RATES``. -_DEFAULT_V_SUBREGION_RATES: Tuple[float, float, float, float, float] = ( - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, +from ._step_validation import ( + _DEFAULT_V_SUBREGION_RATES, + _V_SUBREGION_RATE_ALIASES, + _V_SUBREGION_RATE_LABELS, + _descendant_phase_step_classifier, + _validate_segment_rates, + _validate_v_subregion_rates, ) - -def _validate_v_subregion_rates( - v_subregion_rates: Optional[Dict[str, float]], -) -> Tuple[float, float, float, float, float]: - """Validate the user's ``v_subregion_rates`` dict and return - the normalised ``(FWR1, CDR1, FWR2, CDR2, FWR3)`` tuple. - Pure helper — no DSL state. - - Validation rules: - - - ``None`` (or omitted) → flat default - ``(1.0, 1.0, 1.0, 1.0, 1.0)``. - - Empty dict ``{}`` is treated as ``None`` — same flat default. - - Keys must be a subset of the five canonical labels - (``FWR1`` / ``CDR1`` / ``FWR2`` / ``CDR2`` / ``FWR3``) plus - the two aliases ``FWR`` (expands to FWR1 / FWR2 / FWR3) and - ``CDR`` (expands to CDR1 / CDR2). Case-sensitive — matches - the V-subregion annotation surface (Slice 1). - - Values must be ``int`` / ``float`` (not bool), finite, and - ``>= 0``. NaN / inf / bool / negative raise ``ValueError``. - - Sparse: omitted labels default to ``1.0``. - - Alias expansion happens first, then explicit labels - override. So ``{"FWR": 0.5, "FWR2": 2.0}`` → ``FWR1=0.5, - FWR2=2.0, FWR3=0.5, CDR1=1.0, CDR2=1.0``. - - After expansion, at least one label must be strictly - positive — an all-zero vector would zero every V site and - is almost certainly a builder bug. Reject with - ``ValueError``. - """ - if v_subregion_rates is None: - return _DEFAULT_V_SUBREGION_RATES - - if not isinstance(v_subregion_rates, dict): - raise TypeError( - f"v_subregion_rates must be a dict or None, got " - f"{type(v_subregion_rates).__name__}" - ) - - if not v_subregion_rates: - # Empty dict is equivalent to omitting the kwarg. - return _DEFAULT_V_SUBREGION_RATES - - accepted_keys = set(_V_SUBREGION_RATE_LABELS) | set(_V_SUBREGION_RATE_ALIASES) - unknown = sorted(set(v_subregion_rates.keys()) - accepted_keys) - if unknown: - raise ValueError( - f"v_subregion_rates: unknown label(s) {unknown!r}. " - f"Allowed: {list(_V_SUBREGION_RATE_LABELS)!r} plus the " - f"aliases {list(_V_SUBREGION_RATE_ALIASES.keys())!r}. " - "Labels are case-sensitive; CDR3 / FWR4 are out of " - "scope (CDR3 lives in the junction, FWR4 in the J " - "segment) — use ``segment_rates`` for those." - ) - - def _coerce(label_for_error: str, value: object) -> float: - if isinstance(value, bool) or not isinstance(value, (int, float)): - raise TypeError( - f"v_subregion_rates[{label_for_error!r}]: value must be a " - f"finite non-negative number, got {type(value).__name__}" - ) - v = float(value) - if v != v: # NaN - raise ValueError( - f"v_subregion_rates[{label_for_error!r}]: value must be a " - "finite number, got NaN" - ) - if v < 0.0 or v == float("inf") or v == float("-inf"): - raise ValueError( - f"v_subregion_rates[{label_for_error!r}]: value must be a " - f"finite non-negative number, got {v}" - ) - return v - - # Phase 1: alias expansion. Aliases populate their labels first; - # phase 2's explicit labels then override. - resolved: Dict[str, float] = {} - for alias, expanded_labels in _V_SUBREGION_RATE_ALIASES.items(): - if alias in v_subregion_rates: - v = _coerce(alias, v_subregion_rates[alias]) - for lbl in expanded_labels: - resolved[lbl] = v - - # Phase 2: explicit labels override. - for lbl in _V_SUBREGION_RATE_LABELS: - if lbl in v_subregion_rates: - resolved[lbl] = _coerce(lbl, v_subregion_rates[lbl]) - - # Fill any remaining label with the flat default. - out_list: List[float] = [] - for lbl in _V_SUBREGION_RATE_LABELS: - out_list.append(resolved.get(lbl, 1.0)) - - if sum(out_list) <= 0.0: - raise ValueError( - "v_subregion_rates: at least one label must have a " - "positive rate after alias expansion. The supplied " - "configuration zeroes every V subregion, which would " - "drop every V site out of SHM support." - ) - - return (out_list[0], out_list[1], out_list[2], out_list[3], out_list[4]) - - -# ────────────────────────────────────────────────────────────────── -# Clonal ordering-guard table (Slice: descendant-phase guards) -# ────────────────────────────────────────────────────────────────── -# -# Every DSL step in this table is **descendant-phase** — it models -# observation / library-prep / sequencing biology that must be -# sampled independently per clone member. Pre-fork placement either -# silently misreports the AIRR field (Bugs C / E / F: trace-sourced -# fields that don't survive the parent→descendant boundary) or -# collapses descendant diversity (every clone member shares an -# identical effect because the pass ran once on the parent IR). -# -# The clonal fork methods scan the already-appended step list against -# this table; the first match is rejected with a -# message naming the offending DSL method and the canonical fix. -# -# Each entry is ``(predicate, dsl_method_name)`` where the predicate -# inspects one step and returns ``True`` if it came from -# ``dsl_method_name``. Some DSL methods append :class:`_CorruptStep` -# with different ``kind`` discriminators; others append distinct -# step types. -def _descendant_phase_step_classifier(step): - """Return the DSL method name a descendant-phase ``step`` came - from, or ``None`` if ``step`` is not a descendant-phase step. - - Single source of truth for the unified guard in the flat clonal - fork methods. Adding a new descendant-phase DSL method means - appending a clause here (and adding the - companion spec test in - ``tests/test_clonal_descendant_phase_guards.py``). - """ - from ._pipeline_ir import ( - _CORRUPT_KIND_3PRIME_LOSS, - _CORRUPT_KIND_5PRIME_LOSS, - _CORRUPT_KIND_INDEL, - _CORRUPT_KIND_NS, - _CORRUPT_KIND_PCR, - _CORRUPT_KIND_QUALITY, - _CORRUPT_KIND_REV_COMP, - _CorruptStep, - _MutateStep, - _PairedEndStep, - ) - - if isinstance(step, _MutateStep): - return "mutate" - if isinstance(step, _PairedEndStep): - return "paired_end" - if isinstance(step, _CorruptStep): - # Per-kind classification — only the descendant-phase kinds - # appear in the table. ``contaminant`` is deliberately - # omitted: the DSL slice that added the descendant-phase - # ordering guards did not list it, so leave the placement - # unconstrained until a follow-up explicitly classifies it. - kind_to_method = { - _CORRUPT_KIND_PCR: "pcr_amplify", - _CORRUPT_KIND_QUALITY: "sequencing_errors", - _CORRUPT_KIND_INDEL: "polymerase_indels", - _CORRUPT_KIND_5PRIME_LOSS: "end_loss_5prime", - _CORRUPT_KIND_3PRIME_LOSS: "end_loss_3prime", - _CORRUPT_KIND_NS: "ambiguous_base_calls", - _CORRUPT_KIND_REV_COMP: "random_strand_orientation", - } - return kind_to_method.get(step.kind) - return None - - -# Inputs accepted by :meth:`Experiment.restrict_alleles` per segment. ``None`` -# clears any prior lock; a single name is a one-allele lock; an -# iterable of names is a multi-allele subset; ``_UNSET`` (the default) -# leaves the existing lock unchanged. -_LockInput = Union[str, Iterable[str], None, _Unset] +from ._experiment import ( + _ClonalMixin, + _CompileMixin, + _ConstraintsMixin, + _CorruptionMixin, + _GenotypeAllelesMixin, + _IntrospectionMixin, + _MutationMixin, + _RecombinationMixin, + _RefdataControlsMixin, + _RunMixin, +) # Refdata-resolver helpers (``_CONFIG_ALIASES``, ``_resolve_config_name``, @@ -407,7 +112,18 @@ def _descendant_phase_step_classifier(step): # ``_CONFIG_ALIASES``) keep working. -class Experiment: +class Experiment( + _CorruptionMixin, + _ConstraintsMixin, + _ClonalMixin, + _MutationMixin, + _GenotypeAllelesMixin, + _RecombinationMixin, + _RefdataControlsMixin, + _CompileMixin, + _RunMixin, + _IntrospectionMixin, +): """Fluent builder for a simulation pipeline. Build an ``Experiment`` from a config name, a :class:`DataConfig`, @@ -521,2813 +237,3 @@ def step_count(self) -> int: def refdata(self) -> "_engine.RefDataConfig": """The engine-native ``RefDataConfig`` this experiment is bound to.""" return self._refdata - - def pcr_amplify( - self, - *, - count: Optional[Union[int, Tuple[int, int], Iterable[Tuple[int, float]]]] = None, - rate: Optional[float] = None, - ) -> "Experiment": - """Append a PCR-error step modelling substitution errors - introduced during PCR amplification. - - **Specify intensity with exactly one of ``rate`` or ``count``.** - - ``rate`` is the per-base PCR error probability (e.g. - ``1e-4`` per base per cycle, depending on polymerase - fidelity). At execute time the engine draws - ``count ~ Poisson(rate × pool_len)`` against each record's - current sequence length — matching how PCR error is - universally reported in the literature. This is the - canonical biology-default form. - - ``count`` is the legacy explicit count distribution - (fixed int, ``(low, high)`` range, or empirical - ``(count, weight)`` list). Useful for benchmark scripts. - - Each error samples a uniform position in the assembled pool - and replaces the base with a uniform A/C/G/T draw. Trace - addresses: ``corrupt.pcr.{count, error_site[i], error_base[i]}``. - - Passing both ``count`` and ``rate`` raises ``ValueError``. - Passing neither raises ``ValueError``. - """ - return self._append_count_or_rate_step( - kind=_CORRUPT_KIND_PCR, - label="pcr_amplify", - count=count, - rate=rate, - ) - - def sequencing_errors( - self, - *, - count: Optional[Union[int, Tuple[int, int], Iterable[Tuple[int, float]]]] = None, - rate: Optional[float] = None, - ) -> "Experiment": - """Append a sequencing-quality-error step modelling base-call - errors during sequencing readout. - - **Specify intensity with exactly one of ``rate`` or ``count``.** - - ``rate`` is the per-base sequencing error probability (e.g. - ``1e-3`` for a Q30 base, ``1e-2`` for Q20). Drawn as - ``count ~ Poisson(rate × pool_len)`` per record — the - canonical Phred-quality framing in immunoseq. - - ``count`` is the legacy explicit count distribution. - - Same shape as :meth:`pcr_amplify` but each substitution - writes the destination base in **lowercase** to mark the - position as corrupted (the sequencing-error convention). - """ - return self._append_count_or_rate_step( - kind=_CORRUPT_KIND_QUALITY, - label="sequencing_errors", - count=count, - rate=rate, - ) - - def _append_count_or_rate_step( - self, - *, - kind: str, - label: str, - count: Optional[Union[int, Tuple[int, int], Iterable[Tuple[int, float]]]], - rate: Optional[float], - ) -> "Experiment": - """Shared validator + step appender for the corruption methods - that accept both ``count`` and ``rate`` (PCR errors, sequencing - errors). See :meth:`mutate` for the same shape on the - mutation side.""" - if count is not None and rate is not None: - raise ValueError( - f"{label}(): pass exactly one of `rate` or `count`, not both. " - f"`rate` is the canonical biology default (e.g. rate=1e-4 " - f"for per-base PCR error); `count` is the explicit per-record " - f"count for benchmark / deterministic-count workflows." - ) - if count is None and rate is None: - raise ValueError( - f"{label}(): pass exactly one of `rate` or `count`." - ) - if rate is not None: - if not isinstance(rate, (int, float)) or isinstance(rate, bool): - raise TypeError( - f"rate must be a finite float in [0.0, 1.0], got " - f"{type(rate).__name__}" - ) - if not (0.0 <= float(rate) <= 1.0): - raise ValueError( - f"rate must be in [0.0, 1.0] (got {rate!r}); rate is a " - f"per-base error probability, not an absolute count." - ) - self._steps.append( - _CorruptStep( - kind=kind, - rate=float(rate), - ) - ) - return self - pairs = _normalize_count(count) - self._steps.append( - _CorruptStep( - kind=kind, - count_pairs=pairs, - ) - ) - return self - - def polymerase_indels( - self, - *, - count: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], - insertion_prob: float = 0.5, - ) -> "Experiment": - """Append a polymerase-slippage indel step (PCR-stage artifact). - - Models the small insertions / deletions that arise from - polymerase slippage during PCR amplification — single-base - indels at low rate, typically <1 % per read at standard - polymerase fidelity. Larger structural indels are rare and - not modelled here. - - ``count`` is the per-simulation distribution over the total - number of indel events. Each event independently chooses - insertion (probability ``insertion_prob``) vs. deletion. - Insertions sample a uniform A/C/G/T base. Indels change - pool length and shift downstream region ranges accordingly. - - Raises ``ValueError`` if ``insertion_prob`` is outside - ``[0.0, 1.0]`` or non-finite. - """ - if ( - insertion_prob != insertion_prob # NaN check - or not (0.0 <= insertion_prob <= 1.0) - ): - raise ValueError( - f"insertion_prob must be a finite number in [0.0, 1.0], " - f"got {insertion_prob}" - ) - pairs = _normalize_count(count) - self._steps.append( - _CorruptStep( - kind=_CORRUPT_KIND_INDEL, - count_pairs=pairs, - insertion_prob=float(insertion_prob), - apply_prob=0.0, - ) - ) - return self - - def end_loss_5prime( - self, - *, - length: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], - ) -> "Experiment": - """Append a 5'-end loss corruption step. - - Drops bases from the start of the assembled sequence. The - underlying engine pass is `EndLossPass` — this models - **observation-stage** read-end / primer-region loss (the - AIRR record exposes it as ``end_loss_5_length``, the trace - address is ``corrupt.end_loss.5``). It is distinct from - recombination-stage trimming (`v_trim_5` etc.), which writes - allele-instance metadata rather than deleting pool bytes. - - ``length`` accepts the same shapes as ``count`` on other - corrupt ops: - - - ``length=10`` — strip exactly 10 bases. - - ``length=(0, 20)`` — strip a uniform integer in ``[0, 20]``. - - ``length=[(5, 1.0), (10, 1.0), (15, 1.0)]`` — empirical - (length, weight) distribution. - - The actual loss is clamped to the pool length (so a sample - larger than the sequence drops the whole pool). The loss is - permanent for downstream passes — subsequent corruption - operates on the shorter pool. - - :meth:`primer_trim_5prime` is a backwards-compatible alias. - """ - pairs = _normalize_count(length) - self._steps.append( - _CorruptStep( - kind=_CORRUPT_KIND_5PRIME_LOSS, - count_pairs=pairs, - insertion_prob=0.0, - apply_prob=0.0, - ) - ) - return self - - def end_loss_3prime( - self, - *, - length: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], - ) -> "Experiment": - """Append a 3'-end loss corruption step. - - Drops bases from the end of the assembled sequence. Same - observation-stage semantics as :meth:`end_loss_5prime` - (trace address ``corrupt.end_loss.3``, AIRR field - ``end_loss_3_length``). Same ``length`` shapes. - - :meth:`primer_trim_3prime` is a backwards-compatible alias. - """ - pairs = _normalize_count(length) - self._steps.append( - _CorruptStep( - kind=_CORRUPT_KIND_3PRIME_LOSS, - count_pairs=pairs, - insertion_prob=0.0, - apply_prob=0.0, - ) - ) - return self - - def primer_trim_5prime( - self, - *, - length: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], - ) -> "Experiment": - """Backwards-compatible alias for :meth:`end_loss_5prime`. - - The DSL name `primer_trim_5prime` predates the engine's - cleaner vocabulary — the same step is now exposed as - :meth:`end_loss_5prime`. Behaviour is identical (same trace - address ``corrupt.end_loss.5``, same AIRR field - ``end_loss_5_length``, byte-identical records for the same - seed). New code should prefer the `end_loss_*` form; the - alias stays so existing scripts keep working. - """ - return self.end_loss_5prime(length=length) - - def primer_trim_3prime( - self, - *, - length: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], - ) -> "Experiment": - """Backwards-compatible alias for :meth:`end_loss_3prime`. - - See :meth:`primer_trim_5prime` for the rationale. - """ - return self.end_loss_3prime(length=length) - - def ambiguous_base_calls( - self, - *, - count: Union[int, Tuple[int, int], Iterable[Tuple[int, float]]], - ) -> "Experiment": - """Append an N-substitution corruption step. - - With the per-simulation count drawn from ``count``, replace - that many uniform-random pool positions with the ambiguous - base ``N``. Models the low-quality positions that real - sequencers emit when the base caller cannot commit. - - ``count`` accepts the same shapes as other corruption ops: - - ``count=10`` — fixed. - - ``count=(0, 20)`` — uniform integer range. - - ``count=[(5, 1.0), (10, 1.0)]`` — empirical (count, weight). - - Sampled sites are with-replacement, so ``count`` is the - upper bound on resulting `N` count (collisions reduce it). - """ - pairs = _normalize_count(count) - self._steps.append( - _CorruptStep( - kind=_CORRUPT_KIND_NS, - count_pairs=pairs, - insertion_prob=0.0, - apply_prob=0.0, - ) - ) - return self - - def random_strand_orientation(self, *, prob: float = 0.5) -> "Experiment": - """Append a reverse-complement corruption step. - - With probability ``prob`` the AIRR record-builder reverse- - complements the ``sequence``, ``np1``, ``np2``, and - ``junction`` fields and flips the corresponding pool-position - coords (``*_sequence_start/end``, ``junction_start/end``). - Alignment / germline / CIGAR / identity fields stay in - forward orientation per the AIRR Rearrangement spec. - ``prob=0.0`` is a no-op (coin flip recorded but never fires); - ``prob=1.0`` always flips. The biological model is the ~50% - of antisense reads in real immune-seq libraries. - - Raises ``ValueError`` if ``prob`` is outside ``[0.0, 1.0]`` - or non-finite. - """ - if prob != prob or not (0.0 <= prob <= 1.0): - raise ValueError( - f"prob must be a finite number in [0.0, 1.0], got {prob}" - ) - self._steps.append( - _CorruptStep( - kind=_CORRUPT_KIND_REV_COMP, - count_pairs=(), - insertion_prob=0.0, - apply_prob=float(prob), - ) - ) - return self - - def contaminate(self, *, prob: float) -> "Experiment": - """Append a contaminant-replacement corruption step. - - With probability ``prob`` the entire assembled pool is - overwritten with uniform A/C/G/T bases. Models primer - dimers, bacterial DNA, or any non-receptor sequence in the - library. ``prob=0.0`` is a no-op (coin flip recorded but - never fires); ``prob=1.0`` always contaminates. - - Raises ``ValueError`` if ``prob`` is outside ``[0.0, 1.0]`` - or non-finite. - """ - if prob != prob or not (0.0 <= prob <= 1.0): - raise ValueError( - f"prob must be a finite number in [0.0, 1.0], got {prob}" - ) - self._steps.append( - _CorruptStep( - kind=_CORRUPT_KIND_CONTAMINANT, - count_pairs=(), - insertion_prob=0.0, - apply_prob=float(prob), - ) - ) - return self - - def with_metadata(self, **fields: Any) -> "Experiment": - """Attach sample-level metadata to every AIRR record. - - Standard AIRR Repertoire fields like ``sample_id``, - ``donor``, ``repertoire_id``, and ``cell_id`` are commonly - used by multi-sample analysis tools — Change-O, scirpy, - immcantation pipelines all expect them populated. Custom - keys are also accepted and pass through unchanged. - - Subsequent ``with_metadata()`` calls **merge** with the - prior set (per-key replacement). Pass ``key=None`` to clear - a single key. Values are stored as-is and serialised via - the standard CSV/TSV writers; non-string values are - converted with ``str()`` on output. - - Example:: - - (Experiment.on("human_igh") - .with_metadata(sample_id="P1", donor="D001", - repertoire_id="R-001-IGH") - .recombine().mutate(count=8)) - """ - for key, value in fields.items(): - if not isinstance(key, str): - raise TypeError( - f"with_metadata: keys must be strings, got " - f"{type(key).__name__}" - ) - if value is None: - self._metadata.pop(key, None) - else: - self._metadata[key] = value - return self - - def productive_only(self) -> "Experiment": - """Require every emitted record to be a productive sequence. - - Attaches the canonical productive-sequence contract bundle to - the experiment: junction frame in-register, no stop codons in - the junction, and V/J anchor amino acids preserved. The bundle - is enforced during recombination, mutation, and corruption - passes by narrowing each pass's action support before - sampling. If a constrained support is empty, permissive - execution uses the pass's explicit no-op / sentinel behavior; - use ``strict=True`` at run time to raise instead. - - **Failure surfaces** (see - ``docs/productive_failure_mode_audit.md`` for the full matrix): - - - *Compile-time precondition* — when a sampling distribution - is statically impossible under the bundle (e.g. every NP1 - length violates frame), :meth:`compile` raises - ``ValueError`` regardless of the ``strict`` flag. - - *Runtime fresh strict* (``run(..., strict=True)``) — when - dynamic state makes a sampler's admissible support empty, - raises :class:`GenAIRR._engine.StrictSamplingError` with - structured ``(pass_name, address, reason)`` args. - - *Runtime fresh permissive* (``strict=False``, default) — - the pass records its declared sentinel (indel ``site=-1``, - NP length ``0``, NP base ``N``, trim ``0``) or skips the - slot; the record continues. - - *Trace replay* (``replay_from_trace_file``) — consumes - recorded values verbatim, does not re-evaluate - admissibility. A permissive-sentinel trace replays cleanly - even with ``strict=True``. - - **Order-independent.** This method is a constraint declaration, - not a pipeline step — it can be called anywhere in the chain - and the result is identical. Convention is to place it last - (right before ``run_records()``) so the constraint reads as a - post-hoc requirement on the emitted records. - - TCR refdata accepts the call but raises ``ValueError`` at - :meth:`compile` time because TCRs don't have somatic - hypermutation and the productive bundle's anchor checks - assume BCR semantics. Catch this early at the builder if it - matters to you. - - Example:: - - result = ( - Experiment.on("human_igh") - .recombine() - .mutate(count=(5, 15)) - .productive_only() - .run_records(seed=42) - ) - """ - if "productive" not in self._contracts: - self._contracts.append("productive") - return self - - def expand_clones( - self, - *, - n_clones: int, - per_clone: int, - ) -> "Experiment": - """Expand the pipeline into clonal lineages. - - Marks the per-clone / per-descendant boundary in the chain: - steps appended *before* this call run **once per clone** — - typically just :meth:`recombine`, which establishes the - parent V/D/J + trim + NP + assembled IR for the clonal - family. Steps appended *after* this call run **once per - read inside the family** — typically :meth:`mutate` and the - library-prep / sequencing-stage steps, which introduce - per-read divergence within the clone. - - Concrete shape:: - - exp = (Experiment.on("human_igh") - .recombine() - .expand_clones(n_clones=10, per_clone=20) - .mutate(rate=0.05) - .pcr_amplify(count=2)) - result = exp.run_records(seed=0) - # 10 clones × 20 descendants = 200 records. - # Each record carries a ``clone_id`` integer in [0, 10). - - ``n`` can be omitted from :meth:`run_records` for a clonal - experiment — the runtime expands ``n_clones * per_clone`` - records automatically. Passing ``n`` is allowed only when - ``n == n_clones * per_clone``. - - .. deprecated:: - Use :meth:`clonal_lineage` for BCR affinity-maturation - trees, or :meth:`clonal_repertoire` for TCR / flat-BCR - abundance repertoires with clone-size distributions. - ``expand_clones`` remains supported for fixed-size flat - star expansion. - - Constraints: - - Both ``n_clones`` and ``per_clone`` must be positive ints. - - At most one expansion per pipeline; calling this method - twice raises ``ValueError``. - - Implementation note: the runtime forks the parent's IR - (final ``Simulation`` after the pre-fork plan) into - descendants by running the post-fork plan from that IR - with distinct seeds. Within a clone, every descendant - shares the same recombination provenance (V allele, trim, - NP bases) and only diverges through the post-fork passes. - """ - warnings.warn( - "Experiment.expand_clones() is deprecated. Use " - "Experiment.clonal_lineage() for BCR affinity-maturation trees " - "(internal SHM, lineage_trees, no per_clone) or " - "Experiment.clonal_repertoire() for TCR / flat-BCR abundance " - "repertoires with clone-size distributions and duplicate_count. " - "Neither is a drop-in replacement for fixed per_clone output; " - "expand_clones() remains supported for legacy fixed-size flat " - "star expansion.", - DeprecationWarning, - stacklevel=2, - ) - if not isinstance(n_clones, int) or isinstance(n_clones, bool) or n_clones < 1: - raise ValueError( - f"n_clones must be a positive int, got {n_clones!r}" - ) - if not isinstance(per_clone, int) or isinstance(per_clone, bool) or per_clone < 1: - raise ValueError(f"per_clone must be a positive int, got {per_clone!r}") - if self._has_clonal_fork(): - raise ValueError( - "expand_clones() / clonal_lineage() / clonal_repertoire() " - "can only be called once per pipeline" - ) - # Unified descendant-phase ordering guard. Scan the appended - # step list for any step that came from a descendant-phase - # DSL method (see ``_descendant_phase_step_classifier``). - # Pre-fork placement of these steps either silently misreports - # the AIRR field (trace-sourced fields don't survive the - # parent→descendant boundary — Bugs C / E / F) or collapses - # descendant diversity (the pass runs once on the parent IR - # and every clone member inherits an identical effect). - # Either failure mode is a clonal-semantics violation, so - # reject at the DSL boundary with a uniform message that - # names the offending step and the fix. - for step in self._steps: - offending_method = _descendant_phase_step_classifier(step) - if offending_method is None: - continue - if offending_method == "mutate": - detail = ( - "SHM is descendant-specific in GenAIRR's current " - "clonal model" - ) - else: - detail = ( - "it is descendant-specific and must be sampled " - "independently for each clone member" - ) - raise ValueError( - f"{offending_method} must be called after " - f"expand_clones(); {detail}. Move " - f"{offending_method}(...) after expand_clones(...)." - ) - self._steps.append(_ClonalForkStep(n_clones=n_clones, size=per_clone)) - return self - - def clonal_repertoire( - self, - *, - n_clones: int, - size_distribution: str = "power_law", - exponent: float = 2.0, - mu: float = 1.0, - sigma: float = 1.0, - max_size: int = 1000, - unexpanded_fraction: float = 0.0, - ) -> "Experiment": - """Expand the pipeline into a non-tree clonal **repertoire**. - - This is the non-tree clonal model (contrast with - :meth:`clonal_lineage`, which grows real affinity-maturation - lineage *trees*). It generalizes the deprecated - :meth:`expand_clones`: instead of a fixed ``per_clone`` size, - each clone draws a size from a heavy-tailed distribution (with - an unexpanded-singleton fraction). That many reads pass through - the post-fork library-prep / sequencing passes, and identical - reads are genotype-collapsed into AIRR records carrying a - standard ``duplicate_count`` that records abundance. - - Like :meth:`expand_clones`, this call marks the per-clone / - per-read boundary: steps appended *before* it run **once per - clone** (typically just :meth:`recombine`, establishing the - clonal V/D/J + trim + NP backbone); steps appended *after* it - run **once per read** (the library-prep / sequencing passes - that introduce per-read divergence). - - Concrete shape:: - - exp = (Experiment.on("human_igh") - .recombine() - .clonal_repertoire(n_clones=200, max_size=500, - unexpanded_fraction=0.3) - .sequencing_errors(rate=0.005)) - result = exp.run_records(seed=0) - # Each record carries a `clone_id` and a `duplicate_count`. - - Parameters: - - ``n_clones`` — number of clones (positive int). - - ``size_distribution`` — ``"power_law"`` or ``"lognormal"``. - - ``exponent`` — power-law exponent (``> 0``; used when - ``size_distribution="power_law"``). - - ``mu`` / ``sigma`` — lognormal parameters (``sigma >= 0``; - used when ``size_distribution="lognormal"``). - - ``max_size`` — clamps the largest clone. Because the total - number of reads simulated is roughly the **sum** of the drawn - sizes when post-fork passes are present, keep ``max_size`` - modest to bound runtime. - - ``unexpanded_fraction`` — fraction of clones forced to size 1 - (unexpanded singletons), in ``[0, 1]``. - - TCR works out of the box (no SHM): the reads diverge only - through the post-fork sequencing passes. Note that ``mutate`` - after this call is rejected on TCR by :meth:`mutate`'s own TCR - guard; on BCR an optional post-fork ``mutate`` adds flat SHM. - - **No-corruption shortcut:** a clone with no post-fork passes - emits identical copies, so it collapses to a single record whose - ``duplicate_count`` equals the drawn size. - - Constraints: - - At most one fork per pipeline — calling this when an - :meth:`expand_clones`, :meth:`clonal_lineage`, or - :meth:`clonal_repertoire` fork is already present raises - ``ValueError``. - - The same descendant-phase ordering guard as - :meth:`expand_clones` applies: a descendant-phase step (e.g. - ``.mutate()``) appended *before* this call is rejected; - :meth:`recombine` is fine. - """ - if not isinstance(n_clones, int) or isinstance(n_clones, bool) or n_clones < 1: - raise ValueError( - f"n_clones must be a positive int, got {n_clones!r}" - ) - if size_distribution not in ("power_law", "lognormal"): - raise ValueError( - "size_distribution must be 'power_law' or 'lognormal', got " - f"{size_distribution!r}" - ) - if not (isinstance(exponent, (int, float)) and exponent > 0): - raise ValueError(f"exponent must be > 0, got {exponent!r}") - if not (isinstance(sigma, (int, float)) and sigma >= 0): - raise ValueError(f"sigma must be >= 0, got {sigma!r}") - if not isinstance(max_size, int) or isinstance(max_size, bool) or max_size < 1: - raise ValueError(f"max_size must be a positive int, got {max_size!r}") - if not ( - isinstance(unexpanded_fraction, (int, float)) - and 0.0 <= unexpanded_fraction <= 1.0 - ): - raise ValueError( - "unexpanded_fraction must be in [0, 1], got " - f"{unexpanded_fraction!r}" - ) - if any( - isinstance(s, (_ClonalForkStep, _RepertoireForkStep, _LineageForkStep)) - for s in self._steps - ): - raise ValueError( - "clonal_repertoire() / expand_clones() / clonal_lineage() " - "can only be called once per pipeline" - ) - # Same descendant-phase ordering guard as expand_clones: a - # descendant-phase step (mutate / corruption / paired_end) - # appended before the fork is rejected. - for step in self._steps: - offending_method = _descendant_phase_step_classifier(step) - if offending_method is None: - continue - if offending_method == "mutate": - detail = ( - "SHM is descendant-specific in GenAIRR's current " - "clonal model" - ) - else: - detail = ( - "it is descendant-specific and must be sampled " - "independently for each read" - ) - raise ValueError( - f"{offending_method} must be called after " - f"clonal_repertoire(); {detail}. Move " - f"{offending_method}(...) after clonal_repertoire(...)." - ) - self._steps.append( - _RepertoireForkStep( - n_clones=n_clones, - size_distribution=size_distribution, - exponent=float(exponent), - mu=float(mu), - sigma=float(sigma), - max_size=max_size, - unexpanded_fraction=float(unexpanded_fraction), - ) - ) - return self - - def clonal_lineage( - self, - *, - n_clones: int, - max_generations: int = 10, - n_max: int = 1000, - n_sample: int = 50, - rate: float = 0.05, - lambda_base: float = 1.5, - selection_strength: float = 0.0, - beta: float = 1.0, - target_aa: Optional[str] = None, - mature_substitutions: int = 5, - s5f_model: str = "hh_s5f", - allow_extinction: bool = False, - ) -> "Experiment": - """Grow BCR lineage trees (neutral by default; set ``selection_strength > 0`` - and optionally ``target_aa`` to enable affinity maturation). - - Each clone gets its own lineage tree produced by the Rust - ``simulate_family_outcomes`` kernel. The returned - :class:`~GenAIRR.result.SimulationResultWithLineages` carries: - - - ``.records`` — one AIRR dict per *observed* (genotype-collapsed) cell, - tagged with ``clone_id``, ``lineage_node_id``, ``lineage_parent_id``, - ``lineage_generation``, ``lineage_abundance``, and - ``lineage_affinity``. Because identical genotypes are collapsed before - sampling, the number of records per clone is ≤ ``n_sample``; the - ``lineage_abundance`` field accounts for how many sampled cells were - represented by each observed record. Mutation counts (``n_mutations``, - ``n_v_mutations``, …) are pool-derived and self-consistent. - - ``.lineage_trees`` — one :class:`~GenAIRR._engine.LineageTree` - per clone for ground-truth export (Newick, FASTA, node table TSV). - - Parameters - ---------- - n_clones: - Number of independent clonal lineages to grow. - max_generations: - Maximum depth of the lineage tree (≤ 1000). - n_max: - Per-generation LIVING-population carrying capacity: the live - population per generation is capped at this (the tree can contain - more total nodes across generations). It is NOT a hard cap on the - total number of cells per clone. - n_sample: - Number of cells to sample as observed leaves. Records returned - per clone are ≤ ``n_sample`` because identical genotypes are - collapsed (duplicates are counted in ``lineage_abundance``). - rate: - Per-base SHM rate for within-lineage mutations. - lambda_base: - Poisson mean for offspring count at affinity 0. - selection_strength: - Selection pressure; ``0.0`` = neutral drift (fitness is 1.0 for - every cell). This disables selection but does not force - ``lineage_affinity`` to 0 when a target sequence is supplied. - Set ``> 0`` to enable affinity maturation; combine with - ``target_aa`` for a fixed sequence target. - beta: - Scaling factor for the affinity term in ``exp(−beta·distance)``. - target_aa: - Amino-acid sequence of the full receptor used to compute - per-cell affinity via a BLOSUM62-weighted distance (compared - position-wise against the cell's translated receptor; only - the overlapping prefix is scored when lengths differ). Must be - a non-empty string of standard amino-acid letters - (``ACDEFGHIKLMNPQRSTVWY``). When ``None``, an auto target is - generated from the founder by applying ``mature_substitutions`` - random residue changes whenever selection is enabled. In fully - neutral mode (``selection_strength=0`` and ``target_aa=None``), - no affinity model is built and ``lineage_affinity`` is 0. - mature_substitutions: - Number of amino-acid substitutions used to build the auto - target (when ``target_aa`` is ``None``). - s5f_model: - Bundled S5F kernel name for within-lineage mutation context - (``"hh_s5f"``, ``"hkl_s5f"``, …). - allow_extinction: - Sampling draws from the LIVING final-generation population, so a - founder that draws 0 offspring goes extinct and yields zero - observed cells/records. With ``allow_extinction=False`` (default) - each requested clone is conditioned on survival: an extinct family - is retried with a fresh deterministic sub-seed (up to a bounded - number of attempts) so you reliably get ``n_clones`` families. With - ``allow_extinction=True`` extinction is accepted and the extinct - clone is skipped, producing fewer families than ``n_clones``. - - **BCR-only guard:** ``clonal_lineage`` applies S5F somatic - hypermutation, which is a B-cell process. Calling it on a TCR-configured - experiment raises ``ValueError`` (immunoglobulin / BCR loci only). TCR - clone-size primitives exist in the engine but are not yet exposed as a - DSL workflow. - """ - import math - import warnings - - # --- BCR-only guard (mirror mutate()'s TCR rejection) --- - # clonal_lineage applies S5F somatic hypermutation, a B-cell process, - # so it must reject TCR loci. ``_is_tcr_refdata`` inspects the first V - # allele name prefix (TR* => TCR, IG* => BCR) on the already-bound - # refdata, exactly as mutate() does. Firing here, at call time and - # before compile(), guarantees the clear BCR-only message instead of a - # downstream cartridge / compile error. - if self._is_tcr_refdata(): - locus = self._refdata.v_allele(0).name if self._refdata.v_pool_size() else "?" - raise ValueError( - "clonal_lineage models B-cell somatic hypermutation and " - "supports immunoglobulin (BCR) loci only; the locus " - f"'{locus}' is a TCR locus. (TCR clone-size simulation is not " - "yet exposed in the DSL.)" - ) - # --- allow_extinction --- - if not isinstance(allow_extinction, bool): - raise ValueError( - f"allow_extinction must be a bool, got {allow_extinction!r}" - ) - - # --- n_clones --- - if isinstance(n_clones, bool) or not isinstance(n_clones, int) or n_clones < 1: - raise ValueError(f"n_clones must be a positive int, got {n_clones!r}") - # --- max_generations --- - if ( - isinstance(max_generations, bool) - or not isinstance(max_generations, int) - or max_generations < 1 - or max_generations > 1000 - ): - raise ValueError( - f"max_generations must be a positive int <= 1000, got {max_generations!r}" - ) - # --- n_max --- - if isinstance(n_max, bool) or not isinstance(n_max, int) or n_max < 1: - raise ValueError(f"n_max must be a positive int, got {n_max!r}") - # --- n_sample --- - if isinstance(n_sample, bool) or not isinstance(n_sample, int) or n_sample < 1: - raise ValueError(f"n_sample must be a positive int, got {n_sample!r}") - # --- rate --- - if not isinstance(rate, (int, float)) or not math.isfinite(rate) or rate < 0 or rate > 1: - raise ValueError(f"rate must be a float in [0, 1], got {rate!r}") - # --- lambda_base --- - if not isinstance(lambda_base, (int, float)) or not math.isfinite(lambda_base) or lambda_base < 0: - raise ValueError(f"lambda_base must be a finite non-negative float, got {lambda_base!r}") - # --- beta --- - if not isinstance(beta, (int, float)) or not math.isfinite(beta) or beta < 0: - raise ValueError(f"beta must be a finite non-negative float, got {beta!r}") - # --- selection_strength --- - if not isinstance(selection_strength, (int, float)) or not math.isfinite(selection_strength) or selection_strength < 0: - raise ValueError( - f"selection_strength must be a finite non-negative float, got {selection_strength!r}" - ) - # --- mature_substitutions --- - if ( - isinstance(mature_substitutions, bool) - or not isinstance(mature_substitutions, int) - or mature_substitutions < 0 - ): - raise ValueError( - f"mature_substitutions must be a non-negative int, got {mature_substitutions!r}" - ) - # --- target_aa --- - _VALID_AA = set("ACDEFGHIKLMNPQRSTVWY") - if target_aa is not None: - if not isinstance(target_aa, str) or len(target_aa) == 0: - raise ValueError( - "target_aa must be a non-empty amino-acid string " - "(letters from ACDEFGHIKLMNPQRSTVWY)" - ) - target_aa = target_aa.upper() - invalid = set(target_aa) - _VALID_AA - if invalid: - raise ValueError( - f"target_aa contains invalid characters {sorted(invalid)!r}; " - "only standard amino-acid letters (ACDEFGHIKLMNPQRSTVWY) are allowed" - ) - if len(target_aa) < 30: - warnings.warn( - f"target_aa has length {len(target_aa)}, which is shorter than a typical " - "receptor sequence (~300+ aa). If this is an epitope sequence rather than " - "the full receptor, affinity scoring will be based only on the overlapping " - "prefix — consider supplying the full translated receptor instead.", - UserWarning, - stacklevel=2, - ) - # --- s5f_model (validate at call time, not at run time) --- - from GenAIRR._s5f_loader import _BUILTIN_S5F_MODELS - _s5f_key = s5f_model.lower().strip() - if _s5f_key not in _BUILTIN_S5F_MODELS: - avail = ", ".join(f'"{k}"' for k in sorted(_BUILTIN_S5F_MODELS)) - raise ValueError( - f"Unknown s5f_model {s5f_model!r}. Available: {avail}" - ) - # --- reject duplicate fork steps --- - for s in self._steps: - if isinstance(s, (_ClonalForkStep, _RepertoireForkStep, _LineageForkStep)): - raise ValueError( - "clonal_lineage() / expand_clones() / clonal_repertoire() " - "can only be called once per pipeline" - ) - # --- descendant-phase guard (same as expand_clones) --- - for step in self._steps: - offending_method = _descendant_phase_step_classifier(step) - if offending_method is None: - continue - raise ValueError( - f"{offending_method} must be called after " - f"clonal_lineage(); it is descendant-specific and must be sampled " - f"independently for each clone member. Move " - f"{offending_method}(...) after clonal_lineage(...)." - ) - self._steps.append( - _LineageForkStep( - n_clones=n_clones, - max_generations=max_generations, - n_max=n_max, - n_sample=n_sample, - rate=rate, - lambda_base=lambda_base, - selection_strength=selection_strength, - beta=beta, - target_aa=target_aa, - mature_substitutions=mature_substitutions, - s5f_model=s5f_model, - allow_extinction=allow_extinction, - ) - ) - return self - - def mutate( - self, - *, - model: str = "s5f", - count: Optional[Union[int, Tuple[int, int], Iterable[Tuple[int, float]]]] = None, - rate: Optional[float] = None, - s5f_model: str = "hh_s5f", - segment_rates: Optional[Dict[str, float]] = None, - v_subregion_rates: Optional[Dict[str, float]] = None, - ) -> "Experiment": - """Append a somatic-hypermutation step. - - ``model`` selects the mutation kernel: - - ``"s5f"`` (default) — context-dependent SHM via the bundled - S5F kernel named in ``s5f_model``. Available kernels: - ``"hh_s5f"``, ``"hh_s5f_60"``, ``"hh_s5f_opposite"``, - ``"hkl_s5f"``. - - ``"uniform"`` — position-independent SHM. Each mutated - position gets a uniformly drawn A/C/G/T replacement. - - **Specify intensity with exactly one of ``rate`` or ``count``.** - - ``rate`` is the per-base mutation rate (e.g. ``0.03`` for 3 % - SHM, which is roughly memory B-cell SHM). At execute time the - engine draws ``count ~ Poisson(rate × pool_len)`` against each - record's current sequence length — so the realized count - scales with each record's actual length, matching how - immunologists report SHM in the literature. This is the - canonical, biology-default form. - - ``count`` is the legacy explicit count distribution, useful - for benchmark scripts that want a deterministic count - independent of record length: - - ``count=15`` — fixed: every simulation gets exactly 15 - mutations. - - ``count=(5, 25)`` — uniform integer in ``[5, 25]`` (both - endpoints inclusive). - - ``count=[(5, 1.0), (10, 2.0), ...]`` — explicit empirical - ``(count, weight)`` distribution. - - Passing both ``count`` and ``rate`` raises ``ValueError``. - Passing neither raises ``ValueError``. - - **TCR guard:** somatic hypermutation is a B-cell - phenomenon — T-cells do not undergo SHM in the periphery. - Calling ``.mutate()`` on a TCR-configured experiment raises - ``ValueError`` to prevent silent biological misuse. Use - ``pcr_amplify`` / ``sequencing_errors`` for sequencing-error - realism on TCR data instead. - """ - if self._is_tcr_refdata(): - raise ValueError( - "mutate(): somatic hypermutation does not occur in TCR " - "sequences (T-cells lack AID and the SHM machinery). The " - "configured refdata is a TCR locus. For sequencing-error " - "realism on TCR data, use pcr_amplify / sequencing_errors " - "/ polymerase_indels instead." - ) - model_lc = model.lower() - if model_lc not in ("uniform", "s5f"): - raise ValueError( - f"model must be 'uniform' or 's5f' (got {model!r})" - ) - if count is not None and rate is not None: - raise ValueError( - "mutate(): pass exactly one of `rate` or `count`, not both. " - "`rate` is the canonical biology default (e.g. rate=0.03 " - "for 3% SHM); `count` is the explicit per-record count " - "for benchmark / deterministic-count workflows." - ) - if count is None and rate is None: - raise ValueError( - "mutate(): pass exactly one of `rate` or `count`. " - "Suggested default: rate=0.03 (~3% SHM, memory B-cell range)." - ) - if rate is not None: - if not isinstance(rate, (int, float)) or isinstance(rate, bool): - raise TypeError( - f"rate must be a finite float in [0.0, 1.0], got " - f"{type(rate).__name__}" - ) - if not (0.0 <= float(rate) <= 1.0): - raise ValueError( - f"rate must be in [0.0, 1.0] (got {rate!r}); rate is a " - f"per-base mutation probability, not an absolute count." - ) - seg_rates_tuple = _validate_segment_rates(segment_rates) - v_sub_rates_tuple = _validate_v_subregion_rates(v_subregion_rates) - self._check_v_subregion_rates_satisfiable( - v_subregion_rates, v_sub_rates_tuple - ) - self._steps.append( - _MutateStep( - model=model_lc, - s5f_model_name=s5f_model, - rate=float(rate), - segment_rates=seg_rates_tuple, - v_subregion_rates=v_sub_rates_tuple, - ) - ) - return self - pairs = _normalize_count(count) - seg_rates_tuple = _validate_segment_rates(segment_rates) - v_sub_rates_tuple = _validate_v_subregion_rates(v_subregion_rates) - self._check_v_subregion_rates_satisfiable( - v_subregion_rates, v_sub_rates_tuple - ) - self._steps.append( - _MutateStep( - model=model_lc, - s5f_model_name=s5f_model, - count_pairs=pairs, - segment_rates=seg_rates_tuple, - v_subregion_rates=v_sub_rates_tuple, - ) - ) - return self - - def _check_v_subregion_rates_satisfiable( - self, - raw_rates: Optional[Dict[str, float]], - tuple_rates: Tuple[float, float, float, float, float], - ) -> None: - """Reject a non-default ``v_subregion_rates`` configuration - when the bound cartridge has zero annotated V alleles. - Audit §4: a non-default rate vector against a cartridge - without subregion annotations is unsatisfiable — no V site - would ever see a subregion factor, and the user is - almost certainly building against the wrong cartridge or - forgot to enable the annotation surface. - - Default rates (omitted kwarg, empty dict, or explicit - all-ones expansion) skip the check — those are no-ops and - compose cleanly with any cartridge. - """ - if raw_rates is None or tuple_rates == _DEFAULT_V_SUBREGION_RATES: - return - annotated = 0 - v_total = self._refdata.v_pool_size() - for v_id in range(v_total): - if self._refdata.v_allele(v_id).subregions: - annotated += 1 - # Early exit — we only need at least one annotated. - return - # Zero annotated V alleles: the user's rates can never bite. - raise ValueError( - "mutate(): v_subregion_rates was supplied but the bound " - "cartridge carries no V-subregion annotations on any V " - "allele (annotated_v_count=0 / " - f"total_v_count={v_total}). The rate vector would be a " - "deterministic no-op — almost certainly a builder bug. " - "Either drop v_subregion_rates or use a cartridge with " - "IMGT-gapped V sequences (the bundled human IGH / IGK / " - "IGL OGRDB cartridges derive subregions automatically; " - "see docs/v_region_substructure_audit.md)." - ) - - def _has_clonal_fork(self) -> bool: - """Whether any clonal fork has already been appended. - - Used by the DSL ordering guards on :meth:`invert_d`, - :meth:`receptor_revision`, and the clonal fork methods. - Each fork has a pre-fork parent/founder phase; recombination-time - mechanisms must be inherited by every descendant, emitted copy, - or lineage node. Misordered calls used to lower into the wrong - half and produce records with empty / default fields. The guards - reject those configurations at the DSL boundary. - """ - return any( - isinstance(s, (_ClonalForkStep, _RepertoireForkStep, _LineageForkStep)) - for s in self._steps - ) - - def _is_tcr_refdata(self) -> bool: - """Detect whether the bound refdata is a TCR locus. - - TCR allele names are prefixed with ``TR`` (TRA, TRB, TRG, - TRD); BCR with ``IG``. Inspecting the first V allele is - enough — locus is uniform across the pool. Returns ``False`` - when the V pool is empty (defensive — let downstream errors - surface that condition). - """ - if self._refdata.v_pool_size() == 0: - return False - first_v_name = self._refdata.v_allele(0).name - return first_v_name.upper().startswith("TR") - - def with_genotype(self, genotype) -> "Experiment": - """Attach a single-subject diploid genotype. - - With a genotype attached, V(D)J recombination becomes - haplotype-phased: V, D and J of each rearrangement are drawn from - a single chromosome, honouring the genotype's allele - presence/absence, zygosity, and copy-number/deletion. With no - genotype, the flat (uniform / usage-weighted) path runs unchanged. - - Mutually exclusive with :meth:`restrict_alleles` and the - ``recombine(*_allele_weights=...)`` kwargs — the genotype owns - allele presence and within-gene expression. - - Raises ``ValueError`` if the genotype was built against a - different cartridge (content-hash mismatch), or if allele locks / - explicit allele weights are already set. - """ - live_hash = self._refdata.content_hash() - if genotype._source_hash != live_hash: - raise ValueError( - "genotype was built against a different cartridge (content hash " - f"{genotype._source_hash!r} != experiment {live_hash!r})" - ) - if any(v is not None for v in self._locks.values()): - raise ValueError( - "with_genotype() and restrict_alleles() are mutually exclusive" - ) - if self._user_allele_weights_set: - raise ValueError( - "with_genotype() and recombine(*_allele_weights=...) are mutually " - "exclusive: the genotype owns allele expression" - ) - # Snapshot the (mutable) builder so later edits to ``genotype`` - # cannot desync the compiled engine genotype from - # ``result.genotypes`` (review #8). - self._genotype = genotype._snapshot() - return self - - def restrict_alleles( - self, - *, - v: "_LockInput" = _UNSET, - d: "_LockInput" = _UNSET, - j: "_LockInput" = _UNSET, - ) -> "Experiment": - """Narrow allele sampling to a named subset, per segment. - - Sampling stays uniform — this restricts the *support* of the - sampling distribution to the listed allele names, not pins a - single allele. If you supply one name, the result is - effectively deterministic (uniform over one element); for any - list of N names, the recombination pass samples uniformly - over those N alleles. - - Each kwarg accepts: - - a single allele name string (e.g. ``"IGHV1-2*02"``), - - a list / tuple of allele names (sample uniformly among them), - - ``None`` to clear a previously-set restriction for that segment, - - omitted (default) — the existing restriction for that segment - is left unchanged. - - The restrictions are applied to the next ``recombine()`` step's - ``push_sample_allele`` calls at compile time. Calling - ``restrict_alleles()`` more than once overlays the new values - onto the previous restrictions (per-segment), so - ``.restrict_alleles(v="A").restrict_alleles(d="B")`` restricts - both V and D. - - Raises: - - ``ValueError`` if an allele name doesn't exist in the - configured refdata, or if a restriction is set for ``D`` on - a VJ chain. - - ``TypeError`` if an unexpected input shape is passed. - """ - if self._genotype is not None: - raise ValueError( - "restrict_alleles() and with_genotype() are mutually exclusive: " - "a genotype already owns allele presence and within-gene expression" - ) - for segment, value in (("V", v), ("D", d), ("J", j)): - if value is _UNSET: - continue - if value is None: - self._locks[segment] = None - continue - ids = self._resolve_lock(segment, value) - self._locks[segment] = ids - return self - - def _resolve_lock( - self, - segment: str, - value: Union[str, Iterable[str]], - ) -> Tuple[int, ...]: - """Resolve allele-name(s) → tuple of allele IDs against this - experiment's refdata. Raises on unknown names, on D locks - for VJ chains, or on non-string inputs.""" - if segment == "D" and self._refdata.chain_type != "vdj": - raise ValueError( - f"cannot lock D alleles on a {self._refdata.chain_type!r} chain" - ) - - if isinstance(value, str): - names: Tuple[str, ...] = (value,) - else: - try: - names = tuple(value) - except TypeError as exc: - raise TypeError( - f"restrict_alleles(): {segment} lock must be a name string or an " - f"iterable of name strings; got {type(value).__name__}" - ) from exc - if not all(isinstance(n, str) for n in names): - raise TypeError( - f"restrict_alleles(): {segment} lock entries must all be strings" - ) - if not names: - raise ValueError( - f"restrict_alleles(): {segment} lock list must be non-empty " - "(pass None to clear instead)" - ) - - index = self._allele_name_index(segment) - ids: List[int] = [] - seen: set = set() - for name in names: - if name not in index: - raise ValueError( - f"restrict_alleles(): no {segment} allele named {name!r} in refdata " - "(check spelling against `list_alleles` or the loaded config)" - ) - allele_id = index[name] - if allele_id in seen: - raise ValueError( - f"restrict_alleles(): duplicate {segment} allele {name!r} in lock list" - ) - seen.add(allele_id) - ids.append(allele_id) - return tuple(ids) - - def _allele_name_index(self, segment: str) -> Dict[str, int]: - """Build a name → allele_id map for the given segment by - scanning the refdata pool. Cheap enough to do per-call given - typical pool sizes (≤ a few hundred) and the rarity of - ``.restrict_alleles()``.""" - if segment == "V": - n = self._refdata.v_pool_size() - getter = self._refdata.v_allele - elif segment == "D": - n = self._refdata.d_pool_size() - getter = self._refdata.d_allele - elif segment == "J": - n = self._refdata.j_pool_size() - getter = self._refdata.j_allele - else: # pragma: no cover — guarded above. - raise ValueError(f"unsupported segment {segment!r}") - return {getter(i).name: i for i in range(n)} - - def recombine( - self, - *, - np1_lengths: Optional[Iterable[Tuple[int, float]]] = None, - np2_lengths: Optional[Iterable[Tuple[int, float]]] = None, - v_allele_weights: Optional[Dict[str, float]] = None, - d_allele_weights: Optional[Dict[str, float]] = None, - j_allele_weights: Optional[Dict[str, float]] = None, - ) -> "Experiment": - """Append a standard V(D)J recombination step. - - Compiles to: - - **VJ:** sample V → sample J → (trim V_3, J_5) → assemble V - → generate NP1 → assemble J. - - **VDJ:** sample V → sample D → sample J → (trim V_3, D_5, - D_3, J_5) → assemble V → generate NP1 → assemble D → - generate NP2 → assemble J. - - ``np1_lengths`` / ``np2_lengths`` default to the species' - empirical NP-length distributions (from - ``DataConfig.NP_lengths``) when the experiment is bound to - a DataConfig. For raw-RefDataConfig experiments where no - empirical data is available, both fall back to the uniform - ``[(0, 1.0), ..., (6, 1.0)]`` distribution and emit a - :class:`UserWarning` so the caller knows the synthetic - default is being used. Pass an explicit iterable of - ``(length, weight)`` tuples to override the default. - Passing ``np2_lengths`` on a VJ chain raises ``ValueError`` - (VJ chains have no NP2 region — there's no D segment to - bracket). - - **Exonuclease trim** is enabled by default and uses the - empirical per-segment trim distributions from the bound - ``DataConfig`` when available. To disable trim, or supply - custom trim distributions, call :meth:`trim` *after* - :meth:`recombine` in the chain (before any mutation / - corruption step). On a raw ``RefDataConfig`` without trim - data, recombine emits a :class:`UserWarning` and falls back - to a no-op trim. - - ``v_allele_weights`` / ``d_allele_weights`` / - ``j_allele_weights`` — optional ``{allele_name: - weight}`` dicts that bias allele sampling. Listed alleles - get the supplied positive weight; unlisted alleles default - to 1.0, so e.g. ``v_allele_weights={"IGHV3-23*01": 100}`` - boosts that allele while keeping every other V allele - possible at 1/100th the rate. Mutually exclusive with the - per-segment :meth:`restrict_alleles` restriction. Raises - ``ValueError`` for unknown allele names or non-positive - weights. - """ - # Explicit allele weights conflict with an attached genotype: - # the genotype owns allele presence + within-gene expression, and - # the phased lowering ignores recombine-step weights. Reject the - # combination instead of silently dropping the weights. - if any( - w is not None - for w in (v_allele_weights, d_allele_weights, j_allele_weights) - ): - self._user_allele_weights_set = True - if self._genotype is not None: - raise ValueError( - "recombine(*_allele_weights=...) and with_genotype() are " - "mutually exclusive: the genotype owns allele expression" - ) - - # VJ chains have no NP2 region — surface user mistakes loudly - # instead of silently dropping the argument. - if np2_lengths is not None and self._refdata.chain_type != "vdj": - raise ValueError( - f"np2_lengths is only valid for VDJ chains; the bound " - f"refdata is {self._refdata.chain_type!r} (no D segment, " - f"no NP2 region). Drop the np2_lengths kwarg or bind a " - f"VDJ refdata." - ) - - # Trim is default-on; .trim() may later disable or replace it. - trim = True - defaults = self._recombine_defaults() if trim or self._dataconfig else None - - # Raw-RefDataConfig path: there's no DataConfig backing this - # experiment, so empirical NP lengths and trim distributions - # don't exist. We fall back to a uniform NP-length default — - # historically silent — and surface a warning so the synthetic - # default isn't mistaken for real biology. The trim warning is - # emitted at compile() time instead, after .trim() has had a - # chance to disable trim explicitly. - if self._dataconfig is None: - if np1_lengths is None or ( - self._refdata.chain_type == "vdj" and np2_lengths is None - ): - warnings.warn( - "Experiment bound to a raw RefDataConfig with no empirical " - "NP-length distribution; falling back to uniform " - "[(0, 1.0), ..., (6, 1.0)]. Pass np1_lengths " - "(and np2_lengths for VDJ) explicitly to silence this.", - UserWarning, - stacklevel=2, - ) - - np1 = ( - _normalize_lengths(np1_lengths) - if np1_lengths is not None - else self._default_np_lengths(defaults, "np1") - ) - np2 = ( - _normalize_lengths(np2_lengths) - if np2_lengths is not None - else self._default_np_lengths(defaults, "np2") - ) - - if trim and defaults is not None: - trim_v_3 = _to_immutable_pairs(defaults.get("trim_v_3")) - trim_d_5 = _to_immutable_pairs(defaults.get("trim_d_5")) - trim_d_3 = _to_immutable_pairs(defaults.get("trim_d_3")) - trim_j_5 = _to_immutable_pairs(defaults.get("trim_j_5")) - else: - trim_v_3 = trim_d_5 = trim_d_3 = trim_j_5 = None - - # Cartridge-owned NP base distributions (Slice — Typed NP - # base model). Resolution lives in - # ``_dataconfig_extract.extract_recombine_defaults``: - # ``typed reference_models.np_bases → uniform`` (the - # legacy auto-lift of ``NP_first_bases`` / ``NP_transitions`` - # is deliberately deferred). ``None`` here means the - # pre-slice ``UniformBase`` applies. - if defaults is not None: - np1_base_pairs = _to_immutable_byte_pairs(defaults.get("np1_bases")) - np2_base_pairs = _to_immutable_byte_pairs(defaults.get("np2_bases")) - np1_markov_transitions = _to_immutable_byte_pair_matrix( - defaults.get("np1_markov_transitions") - ) - np2_markov_transitions = _to_immutable_byte_pair_matrix( - defaults.get("np2_markov_transitions") - ) - p_v_3_lengths = _to_immutable_pairs( - defaults.get("p_v_3_lengths") - ) - p_d_5_lengths = _to_immutable_pairs( - defaults.get("p_d_5_lengths") - ) - p_d_3_lengths = _to_immutable_pairs( - defaults.get("p_d_3_lengths") - ) - p_j_5_lengths = _to_immutable_pairs( - defaults.get("p_j_5_lengths") - ) - else: - np1_base_pairs = None - np2_base_pairs = None - np1_markov_transitions = None - np2_markov_transitions = None - p_v_3_lengths = None - p_d_5_lengths = None - p_d_3_lengths = None - p_j_5_lengths = None - - # Precedence (Slice — Allele Usage Estimation v1): - # - # 1. explicit `*_allele_weights=` kwarg - # 2. cartridge `reference_models.allele_usage` - # 3. uniform (legacy default) - # - # The kwarg path stays load-bearing for ad-hoc bias — - # only when the kwarg is omitted do we fall through to - # the typed cartridge plane (or uniform). Legacy - # `gene_use_dict` is NOT consulted. - if v_allele_weights is None and defaults is not None: - v_allele_weights = defaults.get("allele_usage_v") - if d_allele_weights is None and defaults is not None: - d_allele_weights = defaults.get("allele_usage_d") - if j_allele_weights is None and defaults is not None: - j_allele_weights = defaults.get("allele_usage_j") - weights_v = self._resolve_allele_weights("V", v_allele_weights) - weights_d = self._resolve_allele_weights("D", d_allele_weights) - weights_j = self._resolve_allele_weights("J", j_allele_weights) - - step = _RecombineStep( - np1_lengths=np1, - np2_lengths=np2, - trim_v_3=trim_v_3, - trim_d_5=trim_d_5, - trim_d_3=trim_d_3, - trim_j_5=trim_j_5, - weights_v=weights_v, - weights_d=weights_d, - weights_j=weights_j, - np1_base_pairs=np1_base_pairs, - np2_base_pairs=np2_base_pairs, - np1_markov_transitions=np1_markov_transitions, - np2_markov_transitions=np2_markov_transitions, - p_v_3_lengths=p_v_3_lengths, - p_d_5_lengths=p_d_5_lengths, - p_d_3_lengths=p_d_3_lengths, - p_j_5_lengths=p_j_5_lengths, - ) - self._steps.append(step) - return self - - def trim( - self, - *, - enabled: bool = True, - v_3: Optional[Iterable[Tuple[int, float]]] = None, - d_5: Optional[Iterable[Tuple[int, float]]] = None, - d_3: Optional[Iterable[Tuple[int, float]]] = None, - j_5: Optional[Iterable[Tuple[int, float]]] = None, - ) -> "Experiment": - """Configure exonuclease trim on the preceding :meth:`recombine` - step. - - Trim is a per-segment exonuclease step that lives biologically - *inside* V(D)J recombination (between allele selection and NP - insertion). It is on by default with empirical distributions - sourced from the bound ``DataConfig``. Use this method only - when you need to override that default: - - - ``trim(enabled=False)`` — disable trim entirely. Equivalent to - recombining against the raw allele endpoints. - - ``trim(v_3=..., d_5=..., d_3=..., j_5=...)`` — supply custom - per-segment trim-length distributions as ``(length, weight)`` - iterables. Omitted segments keep their empirical defaults - (or no-op on raw RefDataConfig). - - **Position in the chain.** ``trim()`` is configuration applied - to the most recent ``recombine()`` step. It must appear after - ``recombine()`` and before any mutation / corruption step. - Calling it before ``recombine()`` raises ``ValueError``; - calling it after a mutation step raises ``ValueError``. - - Raises ``ValueError`` if no ``recombine()`` step has been - appended yet, or if the chain already has steps that would - biologically follow recombination (mutate, corrupt_*). - """ - from dataclasses import replace as _replace - - # Find the most recent recombine step. - rec_idx: Optional[int] = None - for i in range(len(self._steps) - 1, -1, -1): - if isinstance(self._steps[i], _RecombineStep): - rec_idx = i - break - if rec_idx is None: - raise ValueError( - "trim() must be called after recombine(); no recombine " - "step is on this Experiment yet." - ) - - # Anything appended after the recombine step is wrong-ordered - # for trim configuration. - for j in range(rec_idx + 1, len(self._steps)): - offending = type(self._steps[j]).__name__ - raise ValueError( - f"trim() must be called immediately after recombine() — " - f"before any mutation / corruption / clonal-fork step. " - f"Found a {offending!r} between the latest recombine() " - f"and this trim() call." - ) - - prior: _RecombineStep = self._steps[rec_idx] # type: ignore[assignment] - if not enabled: - # Disable all trim slots, keep NP + weights untouched. - new_step = _replace( - prior, - trim_v_3=None, - trim_d_5=None, - trim_d_3=None, - trim_j_5=None, - trim_overridden=True, - ) - self._steps[rec_idx] = new_step - return self - - # Override individual distributions; pass-through on None. - def _resolve( - current: Optional[Tuple[Tuple[int, float], ...]], - override: Optional[Iterable[Tuple[int, float]]], - ) -> Optional[Tuple[Tuple[int, float], ...]]: - if override is None: - return current - return _normalize_lengths(override) - - new_step = _replace( - prior, - trim_v_3=_resolve(prior.trim_v_3, v_3), - trim_d_5=_resolve(prior.trim_d_5, d_5), - trim_d_3=_resolve(prior.trim_d_3, d_3), - trim_j_5=_resolve(prior.trim_j_5, j_5), - trim_overridden=True, - ) - self._steps[rec_idx] = new_step - return self - - def invert_d(self, *, prob: float = 0.05) -> "Experiment": - """Append a D-segment inversion step. - - Models V(D)J inversion: with probability ``prob`` the sampled - D allele is committed in reverse-complement orientation - instead of forward. Biologically, the RSS heptamers around D - can pair head-to-head and the D segment flips before joining; - prevalence is low (~1–5 %) but real. - - Engine path: a Bool is recorded under - ``sample_allele.d.inverted`` and the - :class:`~GenAIRR._engine.InvertDPass` commits - ``ReverseComplement`` on the D :class:`AlleleInstance` - between sampling and assembly. The - :class:`~GenAIRR._engine.AssembleSegmentPass`(D) (Slice B) - consumes the orientation flag and emits the reverse- - complemented D slice into the pool. Trace replay re-fires - the same orientation decision deterministically. - - **VDJ chains only.** VJ chains have no D pool; calling this - method on a VJ experiment raises ``ValueError``. - - **At most once per experiment.** Calling :meth:`invert_d` - twice raises ``ValueError`` — v1 picks a single per-pipeline - inversion probability rather than supporting last-one-wins - semantics (which would be silent for an over-eager builder). - - **Position in the chain.** Append after :meth:`recombine` - (and any :meth:`trim` override). Calling :meth:`invert_d` - before :meth:`recombine` raises ``ValueError`` at compile - time — the lowering needs the recombine sequence already - materialised in the engine plan so the explicit - ``before(invert_d, assemble.d)`` schedule edge can fire. - - The DSL does **not** expose the per-record orientation in the - AIRR record yet — that's the Slice E follow-up. End-to-end - observability today is via the trace - (``sample_allele.d.inverted``) and the pool bytes. - - Returns ``self`` so the call chains fluently. - """ - if self.chain_type != "vdj": - raise ValueError( - f"invert_d is only valid for VDJ chains (current chain_type={self.chain_type!r})" - ) - if any(isinstance(s, _InvertDStep) for s in self._steps): - raise ValueError( - "invert_d already configured on this experiment; v1 accepts at " - "most one inversion step per pipeline. Build a fresh Experiment " - "if you need a different probability." - ) - # Ordering guard — D inversion is a recombination-time - # decision and must be inherited by every clone descendant. - # When placed post-fork, the lowering silently drops the - # inversion probability (no recombine step in the post-fork - # half to consume it), producing records with d_inverted=False - # even at prob=1.0. Reject at the DSL boundary instead. - if self._has_clonal_fork(): - raise ValueError( - "invert_d must be called before the clonal fork; D " - "inversion is a recombination-time decision and must " - "be inherited by all clone descendants. Move the " - "invert_d(...) call before clonal_lineage(...), " - "clonal_repertoire(...), or expand_clones(...)." - ) - if not isinstance(prob, (int, float)): - raise ValueError( - f"invert_d prob must be a number in [0.0, 1.0], got {type(prob).__name__}" - ) - prob_f = float(prob) - # NaN check FIRST — NaN fails every `<=` comparison silently - # (`(0.0 <= nan)` is False), which would otherwise surface as - # a misleading "out of [0.0, 1.0]" message. Explicit NaN - # rejection gives the user the specific reason. - if prob_f != prob_f: - raise ValueError("invert_d prob must be a number, got NaN") - if not (0.0 <= prob_f <= 1.0): - raise ValueError( - f"invert_d prob must be in [0.0, 1.0], got {prob_f}" - ) - self._steps.append(_InvertDStep(prob=prob_f)) - return self - - def receptor_revision(self, *, prob: float = 0.05, same_haplotype: bool = True) -> "Experiment": - """Append a receptor-revision step. - - Models post-recombination V-segment replacement: with - probability ``prob`` the V slot is reassigned to a different - germline V allele and the V slice in the pool is rewritten. - Biologically, receptor revision is a B-cell tolerance - mechanism that lets a B cell escape autoreactivity by - replacing its V segment via secondary VDJ-recombination-like - rearrangement on the already-assembled receptor. - - Engine path: a Bool is recorded under - ``receptor_revision.applied`` for every simulation. On - ``true``, the - :class:`~GenAIRR._engine.ReceptorRevisionPass` additionally - records the replacement allele id at - ``receptor_revision.v_allele`` and the derived 3' trim at - ``receptor_revision.v_trim_3``, then commits - ``AssignmentChanged`` + ``TrimChanged`` + ``SegmentReplaced`` - against V through a single - :class:`~GenAIRR._engine.SimulationBuilder`. Slice C's - same-length retained constraint - (``allele.len() - trim_3 == old_v_len``, 5' trim fixed at 0) - keeps downstream pool positions stable; the - :class:`~GenAIRR._engine.LiveCallRefreshHook` (Slice B) - reacts to ``SegmentReplaced`` with an AllStructural- - equivalent V/D/J re-walk. - - **VDJ chains only.** Receptor revision is heavy-chain v1; - calling this method on a VJ experiment raises - ``ValueError``. - - **At most once per experiment.** Calling - :meth:`receptor_revision` twice raises ``ValueError`` — - last-one-wins semantics would silently override an - over-eager builder. - - **Position in the chain.** Appended after :meth:`recombine`; - the lowering inlines the engine ``push_receptor_revision`` - call immediately after ``push_assemble("J")`` so the pass - sees the fully-assembled V/D/J/NP pool. Subsequent - :meth:`mutate` / corruption passes lower after this step in - the plan, giving the canonical "recombine → revise → - mutate/corrupt" order the design doc §2 requires. - - AIRR records expose ``receptor_revision_applied`` (bool) and - ``original_v_call`` (the pre-revision V; empty when no - revision applied). ``v_call`` / ``truth_v_call`` are the - **post**-revision V. The three ``receptor_revision.*`` trace - records above remain available for replay. - - Returns ``self`` so the call chains fluently. - """ - if self.chain_type != "vdj": - raise ValueError( - "receptor_revision is only valid for VDJ chains " - f"(current chain_type={self.chain_type!r})" - ) - if any(isinstance(s, _ReceptorRevisionStep) for s in self._steps): - raise ValueError( - "receptor_revision already configured on this experiment; " - "v1 accepts at most one revision step per pipeline. Build " - "a fresh Experiment if you need a different probability." - ) - # Ordering guard — receptor revision is a recombination/ - # ancestor-time decision and must be inherited by every clone - # descendant. When placed post-fork, the lowering silently - # drops the revision probability (no recombine step in the - # post-fork half to consume it), producing records with - # receptor_revision_applied=False and empty original_v_call - # even at prob=1.0. Reject at the DSL boundary instead. - if self._has_clonal_fork(): - raise ValueError( - "receptor_revision must be called before " - "the clonal fork; receptor revision is a " - "recombination-time decision and must be inherited by " - "all clone descendants. Move the " - "receptor_revision(...) call before clonal_lineage(...), " - "clonal_repertoire(...), or expand_clones(...)." - ) - if not isinstance(prob, (int, float)): - raise ValueError( - "receptor_revision prob must be a number in [0.0, 1.0], " - f"got {type(prob).__name__}" - ) - prob_f = float(prob) - # NaN first — see the matching `invert_d` rationale. - if prob_f != prob_f: - raise ValueError("receptor_revision prob must be a number, got NaN") - if not (0.0 <= prob_f <= 1.0): - raise ValueError( - f"receptor_revision prob must be in [0.0, 1.0], got {prob_f}" - ) - if not isinstance(same_haplotype, bool): - raise ValueError( - "receptor_revision same_haplotype must be a bool, got " - f"{same_haplotype!r}" - ) - self._steps.append( - _ReceptorRevisionStep(prob=prob_f, same_haplotype=same_haplotype) - ) - return self - - def paired_end( - self, - *, - r1_length, - r2_length=None, - insert_size, - ) -> "Experiment": - """Append a paired-end / read-layout step. - - Models the Illumina paired-end read layout: each fragment - produces R1 (forward from the 5' adapter) and R2 - (reverse-complemented from the 3' adapter) windows over the - final projected molecule, plus an *insert size* that locates - R2's 3' end. The DSL exposes three integer distributions: - - - ``r1_length`` — required. - - ``r2_length`` — defaults to ``r1_length`` when ``None``. - Many Illumina libraries do run asymmetric (R2 quality - drops faster); the explicit shape lets callers opt in. - - ``insert_size`` — required. - - Each accepts the same three shapes the rest of the DSL - already uses for length-like distributions: - - - ``int`` — fixed value. - - ``(low, high)`` — uniform integer in the closed - interval ``[low, high]``. - - ``[(value, weight), …]`` — explicit empirical - distribution. - - Engine path: a trace-only - :class:`~GenAIRR._engine.PairedEndSamplingPass` records - three Ints at ``paired_end.r1_length`` / - ``paired_end.r2_length`` / ``paired_end.insert_size``; - the AIRR builder reads them back at projection time and - populates the eight ``read_layout`` / ``r1_sequence`` / - ``r2_sequence`` / ``r1_start`` / ``r1_end`` / ``r2_start`` / - ``r2_end`` / ``insert_size`` fields via the Slice B - projection kernel. ``rec.sequence`` is the only - coordinate space — end-loss and rev-comp projections have - already finalised the molecule by the time paired-end - windows are drawn (design doc §6 / §7). - - **Both VDJ and VJ chains supported.** Paired-end is a - sequencing-stage observable, not a biology mechanism; - it makes sense on every chain. - - **At most once per experiment.** Calling - :meth:`paired_end` twice raises ``ValueError`` — - last-one-wins semantics would silently override an - over-eager builder. - - **Position in the chain.** The compile pre-pass extracts - the step and pushes the engine pass at the **end** of the - plan, after every IR-mutating / corruption / orientation - step. Even though the pass is trace-only, recording the - choices last keeps the trace order aligned with the - biological/readout order (recombine → mutation → - corruption → end-loss → paired-end). - - Returns ``self`` so the call chains fluently. - """ - if any(isinstance(s, _PairedEndStep) for s in self._steps): - raise ValueError( - "paired_end already configured on this experiment; " - "v1 accepts at most one paired-end step per pipeline. " - "Build a fresh Experiment if you need a different " - "layout." - ) - - # Resolve default r2 → r1 BEFORE normalization so the - # downstream check pins both at the same source shape. - if r2_length is None: - r2_length = r1_length - - r1_pairs = _normalize_count(r1_length) - r2_pairs = _normalize_count(r2_length) - insert_pairs = _normalize_count(insert_size) - - # r1 / r2 lengths must be strictly positive. `_normalize_count` - # already rejects negatives but allows 0; the projection - # kernel needs > 0, so surface the violation at the DSL - # boundary with a clearer message. - for value, _w in r1_pairs: - if value <= 0: - raise ValueError( - f"paired_end r1_length must be positive, got {value}" - ) - for value, _w in r2_pairs: - if value <= 0: - raise ValueError( - f"paired_end r2_length must be positive, got {value}" - ) - # `_normalize_count` already enforces insert_size >= 0; - # the projection kernel enforces the same. - - # Fixed-value geometry checks. When r1/r2/insert are each - # single-value distributions we know the exact values the - # pass will emit, so reject `r1 > insert` / `r2 > insert` - # here rather than waiting for the engine's per-sample - # `InvalidDistributionOutput`. Distribution / range cases - # may sample valid combinations, so we defer those to the - # engine. - if len(r1_pairs) == 1 and len(insert_pairs) == 1: - r1_value = r1_pairs[0][0] - insert_value = insert_pairs[0][0] - if r1_value > insert_value: - raise ValueError( - f"paired_end r1_length ({r1_value}) > insert_size " - f"({insert_value}); the R1 window would run past " - f"the fragment 3' end." - ) - if len(r2_pairs) == 1 and len(insert_pairs) == 1: - r2_value = r2_pairs[0][0] - insert_value = insert_pairs[0][0] - if r2_value > insert_value: - raise ValueError( - f"paired_end r2_length ({r2_value}) > insert_size " - f"({insert_value}); the R2 window would run past " - f"the fragment 3' end." - ) - - self._steps.append( - _PairedEndStep( - r1_length=r1_pairs, - r2_length=r2_pairs, - insert_size=insert_pairs, - ) - ) - return self - - def _resolve_allele_weights( - self, - segment: str, - user_weights: Optional[Dict[str, float]], - ) -> Optional[Tuple[float, ...]]: - """Build a dense pool-aligned weight vector from the - user-supplied ``{name: weight}`` dict. Listed alleles get the - supplied weight; everything else gets ``1.0``. Returns - ``None`` when no weights were supplied (preserving the - upstream uniform-default behavior). - - Raises ``ValueError`` for unknown allele names, non-positive - weights, or D weights on a VJ chain. - """ - if user_weights is None: - return None - if not user_weights: - raise ValueError( - f"recombine: {segment.lower()}_allele_weights must contain " - "at least one (name, weight) entry" - ) - if segment == "D" and self._refdata.chain_type != "vdj": - raise ValueError( - f"recombine: cannot weight D alleles on a " - f"{self._refdata.chain_type!r} chain" - ) - - index = self._allele_name_index(segment) - if not index: - return None # No alleles for this segment (e.g. VJ + D). - - # Dense vector indexed by allele_id; default 1.0. - pool_size = max(index.values()) + 1 - weights: List[float] = [1.0] * pool_size - for name, w in user_weights.items(): - if not isinstance(name, str): - raise TypeError( - f"recombine: {segment.lower()}_allele_weights keys must " - f"be allele-name strings, got {type(name).__name__}" - ) - if isinstance(w, bool) or not isinstance(w, (int, float)): - raise TypeError( - f"recombine: {segment.lower()}_allele_weights['{name}'] " - f"must be numeric, got {type(w).__name__}" - ) - wf = float(w) - if not (wf > 0.0 and wf == wf and wf != float("inf")): - raise ValueError( - f"recombine: {segment.lower()}_allele_weights['{name}'] " - f"must be a finite positive number, got {wf}" - ) - if name not in index: - raise ValueError( - f"recombine: no {segment} allele named {name!r} in " - "refdata (check spelling against `list_alleles` or the " - "loaded config)" - ) - weights[index[name]] = wf - return tuple(weights) - - def _recombine_defaults(self): - """Lazy-extract the empirical distributions for this - experiment's DataConfig. Returns ``None`` for raw-RefDataConfig - experiments. Cached on first call. - """ - if self._dataconfig is None: - return None - from ._dataconfig_extract import extract_recombine_defaults - - return extract_recombine_defaults(self._dataconfig) - - def _build_contracts(self) -> Optional["_engine.ContractSet"]: - """Synthesize the engine ``ContractSet`` from declared bundles. - - Today only the productive bundle is recognized. Future bundles - (e.g. ``.in_frame_only()``) would compose here. Returns - ``None`` when no constraint methods have been called. - """ - if not self._contracts: - return None - # Composition across bundles isn't supported by the engine yet, - # so for now exactly one bundle is allowed. - if len(self._contracts) > 1: - raise NotImplementedError( - f"composing multiple constraint bundles is not yet supported; " - f"got {self._contracts!r}" - ) - bundle = self._contracts[0] - if bundle == "productive": - return _engine.productive() - raise NotImplementedError( - f"unknown constraint bundle {bundle!r}" - ) - - @staticmethod - def _default_np_lengths( - defaults, key: str - ) -> Tuple[Tuple[int, float], ...]: - """Pick the NP-length distribution to use when the user - didn't pass an explicit one: empirical if available, else - the uniform placeholder. - """ - if defaults is not None: - empirical = defaults.get(key) - if empirical: - return tuple(empirical) - return tuple(_DEFAULT_NP_LENGTHS) - - def curate_refdata( - self, - policy: str, - *, - allowed=None, - keep_unannotated: bool = True, - ) -> "Experiment": - """**Curation** — select which subset of the catalogue - participates in simulation. Replaces this experiment's - reference cartridge with a curated version. - - In the cartridge model, curation is distinct from validation - and from the catalogue itself: - - - **validation** describes what the catalogue contains; - - **curation** decides which alleles participate; - - **simulation** runs against the curated cartridge. - - ``policy`` is one of: - - - ``"raw"`` — identity policy; no-op (kept for symmetry). - - ``"functional_anchors_only"`` — drop V and J alleles whose - anchor doesn't satisfy the active anchor rule (missing - anchor, codon out of bounds, or codon AA outside the - rule's ``expected_amino_acids``). D and C pools pass - through unchanged. - - ``"functional_status"`` — filter V/D/J pools by IMGT - functional status. ``allowed`` is the list of statuses - to keep (default ``["functional"]``); accepted strings - are ``"functional"``, ``"orf"``, ``"pseudogene"``, and - ``"unknown"`` (case-insensitive, ``"F"`` / ``"P"`` - aliases also work). ``keep_unannotated`` (default - ``True``) controls whether alleles with no annotation - survive — bundled ``.pkl`` cartridges currently leave - status unannotated, so the default preserves backward - compatibility. C pool passes through unchanged. - - Use this when the catalogue (e.g. bundled ``mouse_igh`` or - ``human_tcrb``) includes pseudogene/ORF alleles you'd rather - exclude than permit at runtime via - :meth:`allow_curatable_refdata`. Curation is the - professional model; ``allow_curatable_refdata`` is the - broader runtime opt-in for sampling the raw catalogue - as-is. - - Curation cannot fix structural problems — duplicate allele - names, invalid sequence bytes, and locus/chain-type - mismatches still surface from the compile-time validator - regardless of the curated cartridge state. If curation - empties a required pool, ``compile()`` fails with - ``EmptyRequiredPool``. - - The curated cartridge's ``identity.source`` is tagged - ``|curated:`` so trace files and content hashes - distinguish raw from curated artefacts. Returns ``self`` so - the call chains fluently. See ``docs/reference_cartridge.md``. - """ - kwargs = {} - if allowed is not None: - kwargs["allowed"] = list(allowed) - if policy == "functional_status": - kwargs["keep_unannotated"] = bool(keep_unannotated) - self._refdata = self._refdata.curated(policy, **kwargs) - return self - - def allow_curatable_refdata(self, enabled: bool = True) -> "Experiment": - """Opt in to the lenient `AllowCuratable` refdata validation - mode for subsequent ``compile`` / ``run`` / ``run_records`` - calls. Returns ``self`` so the call chains fluently. - - Sits alongside :meth:`curate_refdata` as the cartridge's two - ways to handle pseudogene-bearing catalogues: - - - ``curate_refdata("functional_anchors_only")`` **removes** - the non-canonical alleles. Strict validation then passes - because the curated catalogue is clean. This is the - professional model — the cartridge identifies which - alleles actually participate. - - ``allow_curatable_refdata()`` **keeps** the catalogue - as-is and relaxes the validator. Strict validation passes - Curatable issues (pseudogene-shape anchor anomalies) but - still rejects Fatal ones (empty pools, duplicates, invalid - bytes, anchor out of bounds, locus/chain mismatch). - - Fatal issues are never opt-outable. Curatable issues — V - anchor codon not Cys, J anchor codon outside the locus's - expected set, missing V/J anchor — reflect pseudogene/ORF - entries in real reference catalogues (the bundled - ``mouse_igh`` and ``human_tcrb`` data both contain them). - - Recommended progression: start strict; if your catalogue - contains pseudogenes you want to *exclude*, use - :meth:`curate_refdata`; if you want to *sample from* them - explicitly, use this method. See - ``docs/reference_cartridge.md``. - """ - self._allow_curatable_refdata = bool(enabled) - return self - - def compile(self, *, allow_curatable_refdata: Optional[bool] = None): - """Compile the recorded steps into a reusable - :class:`CompiledExperiment` (or :class:`CompiledClonalExperiment` - when the pipeline contains a :meth:`expand_clones` - fork). - - Idempotent: calling ``compile()`` twice produces two distinct - compiled instances with structurally-equal simulators. - - Constraints declared via :meth:`productive_only` (or future - bundle methods) are baked into the compiled simulator at this - step; they're not runtime knobs. To run without constraints, - omit the constraint methods from the chain. - - ``allow_curatable_refdata`` selects the refdata validation - mode. ``None`` (default) inherits the instance flag set by - :meth:`allow_curatable_refdata`; an explicit ``True`` / - ``False`` overrides per-call. ``False`` runs the gate in - strict mode — every issue rejects compile with a - :class:`ValueError`. ``True`` runs the lenient mode — Fatal - issues (empty pool, duplicates, invalid byte, anchor out of - bounds) still reject, but Curatable issues (pseudogene-shape - anchor anomalies) pass. - """ - if allow_curatable_refdata is None: - allow_curatable_refdata = self._allow_curatable_refdata - - # Receptor revision with a phased genotype is now haplotype-aware: - # the lowering builds a genotype-aware ReceptorRevisionPass that - # restricts the replacement V to carried alleles on the drawn - # rearrangement chromosome (see _lower_recombine). No rejection here. - - # Genotype provenance (subject_id / haplotype / result.genotypes) - # is only threaded through the plain compiled path, not the - # clonal/lineage/repertoire forked classes. Reject the - # combination rather than silently dropping provenance (review - # #9); genotype + clonal cohorts are a planned follow-on. - if self._genotype is not None and self._has_clonal_fork(): - raise ValueError( - "with_genotype() is not supported together with expand_clones() / " - "clonal_lineage() / clonal_repertoire() in this release" - ) - from dataclasses import replace as _replace - - # On raw RefDataConfig with default-on trim, warn at compile - # time if the user didn't explicitly call .trim() to either - # disable or replace the no-op default. By compile() time we - # know the final shape of the recombine step. - if self._dataconfig is None: - for step in self._steps: - if not isinstance(step, _RecombineStep): - continue - has_trim_data = any( - (step.trim_v_3, step.trim_d_5, step.trim_d_3, step.trim_j_5) - ) - if has_trim_data or step.trim_overridden: - # Either trim is set up (DataConfig path or - # .trim(v_3=...)) or the user explicitly called - # .trim() — both silence the warning. - continue - warnings.warn( - "Experiment bound to a raw RefDataConfig has no " - "trim distributions; exonuclease trim is a no-op. " - "Call .trim(enabled=False) to silence this warning, " - "or supply custom distributions via " - ".trim(v_3=..., d_5=..., ...).", - UserWarning, - stacklevel=2, - ) - break # one warning per compile() call is enough - - contracts = self._build_contracts() - any_lock = any(self._locks[seg] is not None for seg in ("V", "D", "J")) - - # if a `_ClonalForkStep` is present, split the step list - # at it and compile two simulators (pre-fork = per-clone, - # post-fork = per-descendant). - fork_idx = next( - (i for i, s in enumerate(self._steps) if isinstance(s, _ClonalForkStep)), - None, - ) - if fork_idx is not None: - fork_step: _ClonalForkStep = self._steps[fork_idx] - pre_steps = self._steps[:fork_idx] - post_steps = self._steps[fork_idx + 1 :] - pre_simulator = self._build_simulator( - pre_steps, - contracts, - any_lock, - replace_fn=_replace, - allow_curatable_refdata=allow_curatable_refdata, - ) - # The post-fork plan inherits the parent's V/D/J/NP - # backbone (recombination already happened in the - # pre-fork half). The recombination-time precondition - # facts (np.np1.length residues, anchor trim supports) - # aren't produced here. v2 used to drop the contract - # bundle entirely on this side to avoid the compile-time - # precondition failure — but doing so left every - # post-fork mutation / corruption pass unfiltered, which - # was the dominant cause of non-productive output under - # `productive_only()`. We now pass the same bundle - # through; the engine analyzer (in - # `compiled/analyze.rs::validate_contract_preconditions`) - # skips the productive-frame check when no recombination - # facts are present in the plan, and the runtime - # admit_with_context / admits_post_event paths still - # enforce no-stop-codon-in-junction and anchor-preserved - # for every substitution and indel in the post-fork pipeline. - post_simulator = self._build_simulator( - post_steps, - contracts, - any_lock=False, - replace_fn=_replace, - allow_curatable_refdata=allow_curatable_refdata, - ) - return CompiledClonalExperiment( - pre_simulator, - post_simulator, - fork_step, - self._refdata, - pre_steps=tuple(pre_steps), - post_steps=tuple(post_steps), - dataconfig=self._dataconfig, - metadata=self._metadata, - ) - - # if a `_RepertoireForkStep` is present, split the step list - # at it and compile the pre-fork (per-clone) simulator plus an - # optional post-fork (per-read) simulator, mirroring the - # `_ClonalForkStep` branch. Per-clone sizes are drawn at run - # time from the heavy-tailed distribution; identical reads are - # collapsed into `duplicate_count`-carrying records. - repertoire_idx = next( - ( - i - for i, s in enumerate(self._steps) - if isinstance(s, _RepertoireForkStep) - ), - None, - ) - if repertoire_idx is not None: - repertoire_step: _RepertoireForkStep = self._steps[repertoire_idx] - pre_steps = self._steps[:repertoire_idx] - post_steps = self._steps[repertoire_idx + 1 :] - pre_simulator = self._build_simulator( - pre_steps, - contracts, - any_lock, - replace_fn=_replace, - allow_curatable_refdata=allow_curatable_refdata, - ) - # post_steps may be empty — that's the pure-copy case - # (each clone collapses to one record with - # duplicate_count = size). When present, the post-fork - # plan inherits the parent's V/D/J/NP backbone, so - # any_lock=False and no recombination facts on this side - # (same as the clonal branch). - post_simulator = None - if post_steps: - post_simulator = self._build_simulator( - post_steps, - contracts, - any_lock=False, - replace_fn=_replace, - allow_curatable_refdata=allow_curatable_refdata, - ) - return CompiledRepertoireExperiment( - pre_simulator, - post_simulator, - repertoire_step, - self._refdata, - dataconfig=self._dataconfig, - metadata=self._metadata, - ) - - # if a `_LineageForkStep` is present, compile a - # CompiledLineageExperiment: pre-fork steps (recombine) become - # the founder simulator; post-fork steps must be empty (the - # lineage engine handles mutation internally). - lineage_idx = next( - (i for i, s in enumerate(self._steps) if isinstance(s, _LineageForkStep)), - None, - ) - if lineage_idx is not None: - lineage_step: _LineageForkStep = self._steps[lineage_idx] - pre_steps = self._steps[:lineage_idx] - post_steps = self._steps[lineage_idx + 1:] - # Steps after clonal_lineage() are per-observed-cell - # library-prep / sequencing artefact passes (the same - # post-fork set expand_clones() allows). SHM is internal - # to the lineage engine, so .mutate() is rejected; the - # paired-end read layout is not yet wired through the - # per-cell corruption merge, so reject it for now too. - for s in post_steps: - if isinstance(s, _MutateStep): - raise ValueError( - "SHM is internal to clonal_lineage; do not add " - ".mutate() after it. Set the within-lineage SHM rate " - "via clonal_lineage(rate=...)." - ) - if isinstance(s, _PairedEndStep): - raise ValueError( - "paired_end not yet supported with clonal_lineage; " - "apply per-read library-prep passes (sequencing_errors, " - "pcr_amplify, polymerase_indels, end_loss_*, " - "ambiguous_base_calls, random_strand_orientation) instead." - ) - if not isinstance(s, _CorruptStep): - raise ValueError( - "Only per-read library-prep / sequencing artefact passes " - "may follow clonal_lineage(); got " - f"{type(s).__name__}." - ) - pre_simulator = self._build_simulator( - pre_steps, - contracts, - any_lock, - replace_fn=_replace, - allow_curatable_refdata=allow_curatable_refdata, - ) - # Build the per-cell corruption simulator from the - # post-fork steps, mirroring the clonal branch's - # post_simulator (no recombination facts on this side, so - # any_lock=False and the analyzer skips the productive - # precondition check). - post_simulator = None - if post_steps: - post_simulator = self._build_simulator( - post_steps, - contracts, - any_lock=False, - replace_fn=_replace, - allow_curatable_refdata=allow_curatable_refdata, - ) - return CompiledLineageExperiment( - pre_simulator, - lineage_step, - self._refdata, - post_simulator=post_simulator, - post_steps=tuple(post_steps), - dataconfig=self._dataconfig, - metadata=self._metadata, - ) - - # When the attached genotype defines novel/private alleles, compile - # against an *effective* reference = base catalogue + injected novel - # alleles, so they become real pool entries the engine samples, - # assembles, and reports like any allele. No genotype, or a genotype - # without novel alleles, uses the base refdata unchanged. - effective_refdata = self._refdata - if self._genotype is not None and self._genotype.has_novel(): - if self._dataconfig is None: - raise ValueError( - "genotype with novel alleles requires a DataConfig-backed " - "experiment (Experiment.on(dataconfig), not a raw RefDataConfig)" - ) - effective_refdata = dataconfig_to_refdata( - self._genotype.effective_dataconfig() - ) - - simulator = self._build_simulator( - self._steps, - contracts, - any_lock, - replace_fn=_replace, - allow_curatable_refdata=allow_curatable_refdata, - refdata=effective_refdata, - ) - return CompiledExperiment( - simulator, - effective_refdata, - steps=tuple(self._steps), - dataconfig=self._dataconfig, - metadata=self._metadata, - genotype=self._genotype, - ) - - def _build_simulator( - self, - steps, - contracts, - any_lock: bool, - *, - replace_fn, - allow_curatable_refdata: bool = False, - refdata=None, - ): - """Compile a list of steps into a `GenAIRR._engine.CompiledSimulator`. - Lifted out of `compile()` so the clonal-fork branch can build - two simulators from sub-step-lists with a shared body. - - ``refdata`` overrides ``self._refdata`` — used when a genotype with - novel alleles compiles against an *effective* reference (base + - injected private alleles).""" - refdata = refdata if refdata is not None else self._refdata - plan = _engine.PassPlan() - # Pull the (at-most-one) `_InvertDStep` out of the step - # sequence and thread its probability into the recombine - # lowering directly. Inlining the InvertDPass push between - # `push_generate_np("NP1", ...)` and `push_assemble("D")` - # (see `_lower_recombine`) keeps the canonical V-NP1-D-NP2-J - # pool layout intact; a separate `_InvertDStep` lowering - # combined with `Schedule::before(invert_d, assemble.d)` - # would re-promote the pass past `assemble.j`, swapping D - # and J in the pool. See `_extract_invert_d_prob` for the - # full rationale. - invert_d_prob, steps = _extract_invert_d_prob(steps) - # Same inline-into-recombine-lowering pattern as - # _extract_invert_d_prob: pulling the receptor-revision step - # out here lets the lowering push it at the exact slot - # between `assemble.j` and any subsequent mutate/corrupt - # passes. A standalone lower path would either need a new - # schedule edge or place the pass at the end of the plan - # (after corruption), both of which break the design doc §2 - # ordering. - receptor_revision_prob, receptor_revision_same_haplotype, steps = ( - _extract_receptor_revision_prob(steps) - ) - # Pull out the (at-most-one) paired-end step too. It must - # land at the END of the plan, not inline with recombine — - # see `_extract_paired_end_step` for the rationale on - # placement vs. trace order. - paired_end_step, steps = _extract_paired_end_step(steps) - for step in steps: - # Inject any allele-locks set via ``.restrict_alleles(...)`` into the - # recombine step at compile time. Other step types ignore - # locks. - if any_lock and isinstance(step, _RecombineStep): - step = replace_fn( - step, - locks_v=self._locks["V"], - locks_d=self._locks["D"], - locks_j=self._locks["J"], - ) - if isinstance(step, _RecombineStep): - _lower_recombine( - step, - plan, - refdata, - invert_d_prob=invert_d_prob, - receptor_revision_prob=receptor_revision_prob, - receptor_revision_same_haplotype=receptor_revision_same_haplotype, - genotype=self._genotype, - ) - else: - lower_step(step, plan, refdata) - # Paired-end is sequencing-stage / readout-stage: lower - # it AFTER every biology + corruption pass so the trace - # records land last. See `_extract_paired_end_step` for - # the full rationale. - if paired_end_step is not None: - _lower_paired_end(paired_end_step, plan) - return plan.compile( - refdata=refdata, - respect=contracts, - allow_curatable_refdata=allow_curatable_refdata, - ) - - def run_cohort( - self, - genotypes, - *, - n_per_subject: int = 1, - seed: int = 0, - counts=None, - strict: bool = False, - expose_provenance: bool = False, - validate_records: bool = False, - allow_curatable_refdata: Optional[bool] = None, - ) -> "CohortResult": - """Run a cohort: N subjects, each with its own diploid genotype, in one - call. A Python loop around the single-subject genotype path — each - subject is compiled and run independently, records are tagged with - ``subject_id`` and given a namespaced ``sequence_id``, and the per-subject - ``SimulationResult`` (with its own refdata) is collected into a - :class:`~GenAIRR.cohort.CohortResult`. - - ``n_per_subject`` applies to every subject; ``counts`` (a parallel - sequence, same length as ``genotypes``) overrides it per subject and may - contain ``0`` (that subject appears with zero records). Subject IDs are - taken from each genotype, or auto-assigned ``subject_0..N-1`` when all are - unset; mixed/duplicate IDs raise. Per-subject sub-seeds are derived - deterministically from ``seed``. - - Mutually exclusive with :meth:`with_genotype`, :meth:`restrict_alleles`, - and ``recombine(*_allele_weights=...)`` (the genotype owns allele - expression). :meth:`receptor_revision` is supported (each subject's - replacement V is restricted to its carried alleles on the drawn - chromosome); clonal forks are not supported in this release. - """ - import copy as _copy - import random as _random - - from .cohort import ( - CohortResult, - CohortSubjectResult, - _resolve_counts, - _resolve_subject_ids, - ) - from .genotype import Genotype - from .result import SimulationResult - - gts = list(genotypes) - if not gts: - raise ValueError("run_cohort: genotypes must be a non-empty sequence") - for g in gts: - if not isinstance(g, Genotype): - raise TypeError( - f"run_cohort: every element must be a Genotype, got " - f"{type(g).__name__}") - - # Mutual exclusions — mirror with_genotype / compile. - if self._genotype is not None: - raise ValueError( - "run_cohort() and with_genotype() are mutually exclusive") - if any(v is not None for v in self._locks.values()): - raise ValueError( - "run_cohort() and restrict_alleles() are mutually exclusive") - if self._user_allele_weights_set: - raise ValueError( - "run_cohort() and recombine(*_allele_weights=...) are mutually " - "exclusive: the genotype owns allele expression") - # receptor_revision is supported per subject (each subject's compile - # builds a genotype-aware revision pass restricted to that subject's - # carried alleles on the drawn chromosome) — no rejection here. - if self._has_clonal_fork(): - raise ValueError( - "run_cohort() is not supported together with expand_clones() / " - "clonal_lineage() / clonal_repertoire() in this release") - # with_metadata() is applied per subject below, but it must not overwrite - # cohort-owned columns. - _cohort_owned = {"subject_id", "sequence_id", "haplotype"} - _md_collision = _cohort_owned & set(self._metadata) - if _md_collision: - raise ValueError( - f"run_cohort: with_metadata keys {sorted(_md_collision)} are " - f"cohort-owned (reserved); rename them") - - # Cartridge-hash check (same as with_genotype). - live_hash = self._refdata.content_hash() - for g in gts: - if g._source_hash != live_hash: - raise ValueError( - "run_cohort: a genotype was built against a different cartridge " - f"(content hash {g._source_hash!r} != experiment {live_hash!r})") - - resolved_counts = _resolve_counts(len(gts), n_per_subject, counts) - subject_ids = _resolve_subject_ids([g.subject_id for g in gts]) - base_rng = _random.Random(seed) - - subjects = [] - for g, sid, count in zip(gts, subject_ids, resolved_counts): - sub_seed = base_rng.getrandbits(63) - snap = g._snapshot() - snap.subject_id = sid - # Clone the uncompiled experiment; never mutate self. - exp_i = _copy.copy(self) - exp_i._genotype = snap - compiled = exp_i.compile(allow_curatable_refdata=allow_curatable_refdata) - refdata_i = compiled.refdata - if count == 0: - # run() rejects n < 1; build an empty result and stamp the - # genotype manually (run_records would otherwise do it). - res = SimulationResult.from_outcomes( - [], refdata_i, expose_provenance=expose_provenance) - res._genotypes = [snap] - else: - res = compiled.run_records( - n=count, seed=sub_seed, strict=strict, - expose_provenance=expose_provenance, - validate_records=validate_records) - # Apply with_metadata() per subject (parity with run_records), then - # namespace sequence_id so combined export never collides. Metadata - # is stamped first; the cohort-owned sequence_id rewrite wins (and a - # collision was already rejected up front). - if self._metadata: - for rec in res.records: - for key, value in self._metadata.items(): - rec[key] = value - for rec in res.records: - rec["sequence_id"] = f"{sid}_{rec.get('sequence_id', '')}" - subjects.append(CohortSubjectResult( - subject_id=sid, genotype=snap, result=res, refdata=refdata_i, - seed=sub_seed, count=count)) - - return CohortResult(subjects) - - def run_records( - self, - *, - n: Optional[int] = None, - seed: int = 0, - strict: bool = False, - expose_provenance: bool = False, - allow_curatable_refdata: Optional[bool] = None, - validate_records: bool = False, - ) -> "SimulationResult": - """Compile and run, then return the batch as a - :class:`SimulationResult` ready for ``.to_csv`` / ``.to_fasta`` - / ``.to_dataframe`` export. - - For non-clonal experiments ``n`` defaults to 1. For clonal - experiments (when the pipeline contains :meth:`with_clonal - _structure`) ``n`` defaults to ``n_clones * size`` and may - be omitted; passing ``n`` explicitly is allowed only if it - matches that product. - - ``expose_provenance=True`` appends ``truth_v_call``, - ``truth_d_call``, ``truth_j_call`` columns containing the - originally-sampled allele names — distinct from the - evidence-driven ``v_call`` / ``d_call`` / ``j_call`` fields - an aligner would produce. Useful for benchmarking aligners - against ground truth without keeping a side truth file. - - ``strict`` semantics match :meth:`run` — strict-mode applies - only to **fresh sampling**. Trace replay - (:meth:`CompiledExperiment.replay_from_trace_file`) consumes - recorded sentinel values verbatim, so a permissive trace - replays cleanly even with ``strict=True``. See - ``docs/productive_failure_mode_audit.md`` §5. - - ``validate_records=True`` runs - :meth:`SimulationResult.validate_records` on the freshly - built batch before returning. If any record fails the - postcondition validator the call raises - :class:`GenAIRR._validation.RecordValidationFailedError` - (a :class:`RuntimeError` subclass) carrying a - machine-greppable summary of the failures. The check costs - roughly one outcome-side re-derivation per record, so it - defaults to ``False``; flip it on in CI or when chasing a - suspected projection bug. The validator runs **before** - any ``with_metadata`` stamps are applied, matching the - order :meth:`SimulationResult.validate_records` would see - on a separate post-hoc call (metadata columns are - per-batch annotations, not engine-derived fields). - - Returns a :class:`SimulationResult`; clonal records carry - an integer ``clone_id`` field per row. - """ - compiled = self.compile(allow_curatable_refdata=allow_curatable_refdata) - if isinstance(compiled, CompiledClonalExperiment): - result = compiled.run_records( - n=n, - seed=seed, - strict=strict, - expose_provenance=expose_provenance, - validate_records=validate_records, - ) - elif isinstance(compiled, CompiledRepertoireExperiment): - if n is not None: - raise ValueError( - "The 'n' parameter is not supported for clonal_repertoire " - "experiments. The number of records depends on the per-clone " - "sizes drawn from the heavy-tailed distribution and the " - "read-collapse, not a fixed product." - ) - result = compiled.run_records( - seed=seed, - strict=strict, - expose_provenance=expose_provenance, - validate_records=validate_records, - ) - elif isinstance(compiled, CompiledLineageExperiment): - if n is not None: - raise ValueError( - "The 'n' parameter is not supported for clonal_lineage experiments. " - "The number of observed records depends on the lineage trees " - "grown from n_clones / n_sample / selection, not a fixed product." - ) - result = compiled.run_records( - seed=seed, - strict=strict, - expose_provenance=expose_provenance, - validate_records=validate_records, - ) - else: - result = compiled.run_records( - n=1 if n is None else n, - seed=seed, - strict=strict, - expose_provenance=expose_provenance, - validate_records=validate_records, - ) - if self._metadata: - for rec in result.records: - for key, value in self._metadata.items(): - rec[key] = value - return result - - def run( - self, - *, - n: Optional[int] = None, - seed: int = 0, - strict: bool = False, - allow_curatable_refdata: Optional[bool] = None, - ) -> List["_engine.Outcome"]: - """Compile and run this experiment ``n`` times. - - Equivalent to - ``self.compile().run(n=n, seed=seed, strict=strict)``. - Returns a list of :class:`GenAIRR._engine.Outcome` objects in - clone-major order for clonal experiments. - - Attach :meth:`productive_only` (or any future constraint - method) to the chain to require admissible records; the - runtime filters NP base draws, length samples, and mutation - / contamination substitutions in real time so the resulting - sequences satisfy the bundle by construction. - - Statically impossible contract configurations fail during - ``compile()`` with ``ValueError``. For runtime residue - — i.e., empty admissible support emerging dynamically at - sample time — ``strict=False`` (default) lets a pass consume - the slot as its explicit no-op / sentinel; ``strict=True`` - raises :class:`GenAIRR._engine.StrictSamplingError` instead. - - Note the two error paths use **different exception classes**: - ``ValueError`` for compile-time preconditions, - ``StrictSamplingError`` (subclass of ``Exception``, NOT of - ``ValueError``) for runtime empty-support. A bare - ``except ValueError:`` will not catch the runtime case. See - ``docs/productive_failure_mode_audit.md`` §6.1. - - ``strict`` only governs **fresh sampling**. Trace replay - (:meth:`CompiledExperiment.replay_from_trace_file`) consumes - recorded values verbatim; a permissive-recorded sentinel - trace replays cleanly even with ``strict=True``. To re-execute - a trace under strict-fresh semantics, call - ``simulator.run(seed=, strict=True)`` instead. - - **Output-correctness validation is on** :meth:`run_records` - **only.** This method returns raw ``Outcome`` objects, which - have no projected AIRR record to validate; pass - ``validate_records=True`` to :meth:`run_records` to opt into - the post-build check (which raises - :class:`GenAIRR._validation.RecordValidationFailedError` on - any failure). For an outcome-by-outcome post-hoc check - without re-running, build a :class:`SimulationResult` via - :meth:`SimulationResult.from_outcomes` and call - :meth:`SimulationResult.validate_records` on it. - """ - compiled = self.compile(allow_curatable_refdata=allow_curatable_refdata) - if isinstance(compiled, CompiledClonalExperiment): - return compiled.run(n=n, seed=seed, strict=strict) - return compiled.run(n=1 if n is None else n, seed=seed, strict=strict) - - def stream( - self, - *, - n: Optional[int] = None, - seed: int = 0, - strict: bool = False, - ) -> Iterator["_engine.Outcome"]: - """Compile and lazily yield :class:`GenAIRR._engine.Outcome` - objects. See :meth:`CompiledExperiment.stream` for full - semantics.""" - return self.compile().stream(n=n, seed=seed, strict=strict) - - def stream_records( - self, - *, - n: Optional[int] = None, - seed: int = 0, - strict: bool = False, - id_prefix: str = "seq", - ) -> Iterator[Dict[str, Any]]: - """Compile and lazily yield AIRR-format record dicts. See - :meth:`CompiledExperiment.stream_records`.""" - return self.compile().stream_records( - n=n, - seed=seed, - strict=strict, - id_prefix=id_prefix, - ) - - def describe(self) -> str: - """Render a biology-style narrative of this experiment. - - The output is one line per step, prefixed with its position - in the chain. Locks, allele weights, NP-length distributions, - SHM kernels, and per-corruption rates are all surfaced. Use - this to sanity-check what a fluent chain actually encodes — - if it doesn't read like an immunology protocol, the chain is - too murky. - - Returns a multi-line string ending without a trailing newline. - Safe to ``print(exp.describe())``. - - Example:: - - >>> print(Experiment.on("human_igh").recombine().mutate(count=(5, 15)).describe()) - Experiment on human_igh (vdj, DataConfig) - 1. V(D)J recombination: sample V/D/J alleles; empirical exonuclease trim (V3', D5', D3', J5'); insert NP1 (0–11 weighted bases) and NP2 (0–11 weighted bases) - 2. Somatic hypermutation (S5F context model, human heavy-chain (HH_S5F)): 5–15 mutations/record - """ - header = _describe_experiment_header(self._refdata, self._dataconfig) - if not self._steps and not self._contracts: - return header + "\n (no steps appended yet)" - - lines = [header] - resolved = self._steps_with_locks_resolved() - body_lines = _describe_step_sequence(resolved, self._refdata.chain_type) - lines.extend(body_lines) - contracts_line = _format_declared_contracts(self._contracts) - if contracts_line: - lines.append(f" Constraints: {contracts_line}") - if self._metadata: - stamps = ", ".join(f"{k}={v!r}" for k, v in self._metadata.items()) - lines.append(f" Metadata stamped on every record: {stamps}") - return "\n".join(lines) - - def _steps_with_locks_resolved(self) -> List[Any]: - """Return a copy of ``self._steps`` with per-segment allele - locks (from :meth:`restrict_alleles`) injected into the first - :class:`_RecombineStep`. Compile-time injection lives in - :meth:`_build_simulator`; ``describe()`` needs the same - substitution to render lock info correctly without - side-effecting the live step list.""" - from dataclasses import replace as _replace - - any_lock = any(self._locks[seg] is not None for seg in ("V", "D", "J")) - if not any_lock: - return list(self._steps) - out: List[Any] = [] - injected = False - for step in self._steps: - if not injected and isinstance(step, _RecombineStep): - step = _replace( - step, - locks_v=self._locks["V"], - locks_d=self._locks["D"], - locks_j=self._locks["J"], - ) - injected = True - out.append(step) - return out - - def __repr__(self) -> str: - return f"" diff --git a/src/GenAIRR/genotype/__init__.py b/src/GenAIRR/genotype/__init__.py new file mode 100644 index 0000000..447f1a9 --- /dev/null +++ b/src/GenAIRR/genotype/__init__.py @@ -0,0 +1,4 @@ +"""Per-individual diploid genotype modelling.""" +from .model import Genotype + +__all__ = ["Genotype"] diff --git a/src/GenAIRR/genotype/_common.py b/src/GenAIRR/genotype/_common.py new file mode 100644 index 0000000..bd0cb7d --- /dev/null +++ b/src/GenAIRR/genotype/_common.py @@ -0,0 +1,20 @@ +"""Shared leaf primitives for the genotype package. + +The canonical V/D/J segment tuple and the reference-allele accessor live +here, imported by both :mod:`.model` and :mod:`._sampling`, so there is a +single definition (no duplication) and no ``model`` <-> ``_sampling`` +import cycle. +""" +from __future__ import annotations + +from typing import Dict, List + +_SEGMENTS = ("V", "D", "J") + + +def _alleles_by_gene(cfg, segment: str) -> Dict[str, List]: + return { + "V": cfg.v_alleles, + "D": cfg.d_alleles, + "J": cfg.j_alleles, + }[segment] or {} diff --git a/src/GenAIRR/genotype/_sampling.py b/src/GenAIRR/genotype/_sampling.py new file mode 100644 index 0000000..5b76202 --- /dev/null +++ b/src/GenAIRR/genotype/_sampling.py @@ -0,0 +1,477 @@ +"""Population-sampling engine for :class:`Genotype`, extracted verbatim as a +mixin (behavior-preserving). ``Genotype`` inherits these class/static methods; +``cls`` resolves to ``Genotype`` at call time, so every call site is unchanged. +""" +from __future__ import annotations + +import math +from collections.abc import Mapping +from typing import Dict, List, Optional, Tuple + +from ._common import _SEGMENTS, _alleles_by_gene + + +class _GenotypeSampling: + @staticmethod + def _required_segments(cfg) -> List[str]: + req = ["V", "J"] + if _alleles_by_gene(cfg, "D"): + req.insert(1, "D") # V, D, J + return req + + @classmethod + def _resolve_sample_segments(cls, cfg, segments_to_sample) -> List[str]: + required = cls._required_segments(cfg) + if segments_to_sample is None: + return required + segs = list(segments_to_sample) + seen = set() + for s in segs: + if s not in _SEGMENTS: + raise ValueError(f"unknown segment {s!r}; expected one of {_SEGMENTS}") + if s in seen: + raise ValueError(f"segments_to_sample contains duplicate segment {s!r}") + seen.add(s) + if not _alleles_by_gene(cfg, s): + raise ValueError(f"cartridge has no {s} segment") + missing = [r for r in required if r not in seen] + if missing: + raise ValueError( + f"segments_to_sample must cover the chain's required segments " + f"{required}; missing {missing}. Partial sampling is not supported " + f"by Genotype.sample (it must return a runnable genotype)." + ) + # canonical _SEGMENTS order so the result is independent of input order + return [s for s in _SEGMENTS if s in seen] + + @classmethod + def _gene_segment_index(cls, cfg, segs) -> Dict[str, List[str]]: + """gene name -> [segments it appears in] (for flat-shape disambiguation).""" + idx: Dict[str, List[str]] = {} + for seg in segs: + for gene in _alleles_by_gene(cfg, seg): + idx.setdefault(gene, []).append(seg) + return idx + + @staticmethod + def _weighted_pick(rng, pairs): + names = [n for (n, _w) in pairs] + weights = [w for (_n, w) in pairs] + return rng.choices(names, weights=weights, k=1)[0] + + @classmethod + def _normalize_freq_spec(cls, cfg, spec, segs): + """Return nested ``{seg: {gene: {allele: weight}}}`` from a nested or flat + ``allele_frequencies`` spec, fully validating segment/gene/allele + addressing and shapes (unknown names and malformed values raise).""" + if spec is None: + return {} + if not isinstance(spec, Mapping): + raise ValueError( + "allele_frequencies must be a mapping (or 'usage_as_prior' / None), " + f"got {type(spec).__name__}" + ) + nested: Dict[str, Dict[str, Dict[str, float]]] = {} + keys = set(spec) + if keys and keys <= set(_SEGMENTS): # segment-keyed (nested) shape + for seg, genes in spec.items(): + if seg not in segs: + raise ValueError( + f"allele_frequencies: segment {seg!r} is not being sampled " + f"(sampling {segs})" + ) + if not isinstance(genes, Mapping): + raise ValueError( + f"allele_frequencies[{seg!r}] must be a mapping of " + f"gene -> {{allele: weight}}, got {type(genes).__name__}" + ) + catalogue = _alleles_by_gene(cfg, seg) + for gene, alleles in genes.items(): + if gene not in catalogue: + raise ValueError( + f"allele_frequencies: {seg} gene {gene!r} is not in the cartridge" + ) + if not isinstance(alleles, Mapping) or not alleles: + raise ValueError( + f"allele_frequencies[{seg!r}][{gene!r}] must be a non-empty " + f"mapping of allele -> weight" + ) + nested.setdefault(seg, {})[gene] = alleles + else: # flat {gene: {...}} shape + gidx = cls._gene_segment_index(cfg, segs) + for gene, alleles in spec.items(): + segs_for = gidx.get(gene) + if not segs_for: + raise ValueError(f"allele_frequencies: unknown gene {gene!r}") + if len(segs_for) > 1: + raise ValueError( + f"allele_frequencies: gene {gene!r} is ambiguous across segments " + f"{segs_for}; use the {{segment: {{gene: ...}}}} shape" + ) + if not isinstance(alleles, Mapping) or not alleles: + raise ValueError( + f"allele_frequencies[{gene!r}] must be a non-empty mapping of " + f"allele -> weight" + ) + nested.setdefault(segs_for[0], {})[gene] = alleles + return nested + + @classmethod + def _usage_frequencies(cls, cfg, segs): + rm = getattr(cfg, "reference_models", None) + usage = getattr(rm, "allele_usage", None) if rm else None + if usage is None: + raise ValueError( + "allele_frequencies='usage_as_prior' requires a cartridge with a " + "typed reference_models.allele_usage; this cartridge has none" + ) + nested: Dict[str, Dict[str, Dict[str, float]]] = {} + seg_attr = {"V": "v", "D": "d", "J": "j"} + for seg in segs: + table = getattr(usage, seg_attr[seg], None) or {} + if not table: + raise ValueError( + f"allele_frequencies='usage_as_prior': cartridge allele_usage has no " + f"entries for requested segment {seg!r}" + ) + catalogue = _alleles_by_gene(cfg, seg) + for allele_name, w in table.items(): + gene = allele_name.split("*")[0] + if gene not in catalogue: + raise ValueError( + f"allele_frequencies='usage_as_prior': usage allele {allele_name!r} " + f"maps to {seg} gene {gene!r}, which is not in the cartridge catalogue" + ) + nested.setdefault(seg, {}).setdefault(gene, {})[allele_name] = w + return nested + + @classmethod + def _resolve_allele_frequencies(cls, cfg, spec, segs): + if spec == "usage_as_prior": + nested = cls._usage_frequencies(cfg, segs) + else: + nested = cls._normalize_freq_spec(cfg, spec, segs) + out: Dict[str, Dict[str, List[Tuple[str, float]]]] = {} + for seg in segs: + out[seg] = {} + for gene, alleles in _alleles_by_gene(cfg, seg).items(): + names = {a.name for a in alleles} + supplied = nested.get(seg, {}).get(gene) + if supplied is None: + out[seg][gene] = [(a.name, 1.0) for a in alleles] # uniform fallback + continue + pairs: List[Tuple[str, float]] = [] + total = 0.0 + # Sort by allele name so the draw is independent of the spec's + # dict insertion order — two content-equal models (same + # content_checksum) then produce identical draws at a given seed. + for nm, w in sorted(supplied.items()): + if nm not in names: + raise ValueError(f"{seg} gene {gene}: {nm!r} is not a known allele") + if isinstance(w, bool) or not isinstance(w, (int, float)) or not math.isfinite(w) or w < 0: + raise ValueError( + f"{seg} gene {gene}: weight for {nm!r} must be finite and >= 0, got {w!r}" + ) + if w > 0: + pairs.append((nm, float(w))) + total += w + if total <= 0 or not pairs: + raise ValueError( + f"{seg} gene {gene}: at least one allele weight must be > 0" + ) + out[seg][gene] = pairs + return out + + @classmethod + def _resolve_haplotype_deletion(cls, cfg, spec, segs): + def _check(p, where): + if isinstance(p, bool) or not isinstance(p, (int, float)) or not math.isfinite(p): + raise ValueError(f"{where}: deletion probability must be a finite number, got {p!r}") + if not (0.0 <= p <= 1.0): + raise ValueError(f"{where}: deletion probability must be in [0, 1], got {p}") + return float(p) + + out = {seg: {gene: 0.0 for gene in _alleles_by_gene(cfg, seg)} for seg in segs} + if isinstance(spec, (int, float)) and not isinstance(spec, bool): + p = _check(spec, "haplotype_deletion_prob") + for seg in segs: + for gene in out[seg]: + out[seg][gene] = p + return out + if not isinstance(spec, Mapping): + raise ValueError( + f"haplotype_deletion_prob must be a float or a mapping, got {type(spec).__name__}" + ) + keys = set(spec) + if keys and keys <= set(_SEGMENTS): # nested {seg: {gene: prob}} + for seg, genes in spec.items(): + if seg not in segs: + raise ValueError(f"haplotype_deletion_prob: segment {seg!r} not being sampled") + if not isinstance(genes, Mapping): + raise ValueError( + f"haplotype_deletion_prob[{seg!r}] must be a mapping of " + f"gene -> probability, got {type(genes).__name__}" + ) + for gene, p in genes.items(): + if gene not in out[seg]: + raise ValueError(f"haplotype_deletion_prob: unknown {seg} gene {gene!r}") + out[seg][gene] = _check(p, f"haplotype_deletion_prob[{seg}][{gene}]") + else: # flat {gene: prob} + gidx = cls._gene_segment_index(cfg, segs) + for gene, p in spec.items(): + segs_for = gidx.get(gene) + if not segs_for: + raise ValueError(f"haplotype_deletion_prob: unknown gene {gene!r}") + if len(segs_for) > 1: + raise ValueError( + f"haplotype_deletion_prob: gene {gene!r} is ambiguous across " + f"segments {segs_for}; use the {{segment: {{gene: prob}}}} shape" + ) + out[segs_for[0]][gene] = _check(p, f"haplotype_deletion_prob[{gene}]") + return out + + @classmethod + def sample( + cls, + cfg, + *, + seed: int = 0, + allele_frequencies=None, + haplotype_deletion_prob=None, + segments_to_sample=None, + chromosome_weights: Optional[Tuple[float, float]] = None, + subject_id: Optional[str] = None, + ensure_viable: bool = True, + max_resamples: int = 1000, + use_cartridge_priors: bool = True, + include_cartridge_novel_alleles="auto", + ) -> "Genotype": + """Sample a fully-specified diploid genotype from population priors. + + Independent per-gene, per-chromosome Hardy-Weinberg model: each gene on + each chromosome is independently deleted or assigned one allele drawn + from the gene's allele frequencies. Homozygous/heterozygous/hemizygous/ + deleted states emerge at Hardy-Weinberg rates. + + When ``cfg`` carries a ``genotype_priors`` plane and an argument is left + at its default, the plane supplies it (``use_cartridge_priors=False`` + disables ALL plane consumption — a clean uniform catalogue-only draw). + Each input is sourced independently and recorded in + ``g.prior_provenance`` (explicit / cartridge / uniform / default). + + This is NOT a population haplotype model — no linkage disequilibrium, gene + co-deletion blocks, ancestry, or donor-specific haplotype structure. + + With ``ensure_viable=True`` (default), the draw is repeated (with a + deterministic sub-seed) up to ``max_resamples`` times until at least one + **positive-weight** chromosome carries every required segment, raising + ``ValueError`` if that is impossible under the given deletion settings. + + NOTE: with the default ``ensure_viable=True`` the result is Hardy-Weinberg + **conditioned on viability** (draws with no complete usable haplotype are + rejected), not the unconditional HW distribution. Use + ``ensure_viable=False`` for the raw (possibly infeasible) HW draw. + """ + import random + + if isinstance(max_resamples, bool) or not isinstance(max_resamples, int) or max_resamples < 1: + raise ValueError(f"max_resamples must be an int >= 1, got {max_resamples!r}") + if not isinstance(use_cartridge_priors, bool): + raise ValueError( + f"use_cartridge_priors must be a bool, got {use_cartridge_priors!r}") + # Identity checks (not `in (...)`): Python's `1 == True` / `0 == False` + # would otherwise let integers slip through and silently act like False. + _icna = include_cartridge_novel_alleles + if not (_icna is True or _icna is False or _icna == "auto"): + raise ValueError( + "include_cartridge_novel_alleles must be 'auto', True, or False, " + f"got {include_cartridge_novel_alleles!r}") + + plane = cls._resolve_plane(cfg, use_cartridge_priors) + if plane is not None: + # A plane attached via the builder is already validated, but a plane + # set directly on the DataConfig bypasses that — validate before use so + # a malformed prior (empty model_id, NaN weights, D-on-VJ) fails loudly + # rather than being silently sampled and stamped into provenance. + plane.validate(chain_type=getattr(getattr(cfg, "metadata", None), "chain_type", None)) + + # Per-input source resolution. + if allele_frequencies is not None: + freq_spec, freq_src = allele_frequencies, "explicit" + elif plane is not None and plane.allele_frequencies: + freq_spec, freq_src = plane.allele_frequencies, "cartridge" + else: + freq_spec, freq_src = None, "uniform" + + if haplotype_deletion_prob is not None: + del_spec, del_src = haplotype_deletion_prob, "explicit" + elif plane is not None and plane.haplotype_deletion_prob: + del_spec, del_src = plane.haplotype_deletion_prob, "cartridge" + else: + del_spec, del_src = 0.0, "default" + + if chromosome_weights is not None: + cw_in, cw_src = chromosome_weights, "explicit" + elif plane is not None: + cw_in, cw_src = plane.chromosome_weights, "cartridge" + else: + cw_in, cw_src = (0.5, 0.5), "default" + + cw = cls._check_chromosome_weights(*cw_in) + segs = cls._resolve_sample_segments(cfg, segments_to_sample) + + # Candidate-novel injection. Plane novels are draw CANDIDATES (injected + # into the sampling cfg) only when their source is active; the returned + # genotype still exports only CARRIED novels via effective_dataconfig(). + inject_novels = False + if plane is not None and plane.novel_alleles: + if include_cartridge_novel_alleles is True: + inject_novels = True + elif include_cartridge_novel_alleles == "auto": + inject_novels = freq_src in ("cartridge", "uniform") + # False -> never + novel_src = "none" + sampling_cfg = cfg + novel_helper = None + if inject_novels: + novel_helper = cls._register_plane_novels(cfg, plane) + sampling_cfg = cls._dataconfig_injecting_all_novels(novel_helper) + freq_spec = cls._augment_freqs_with_novels(sampling_cfg, freq_spec, plane, segs) + novel_src = "cartridge" + + freqs = cls._resolve_allele_frequencies(sampling_cfg, freq_spec, segs) + delp = cls._resolve_haplotype_deletion(sampling_cfg, del_spec, segs) + # Compute the cartridge content hash ONCE (it rebuilds refdata + hashes); + # reuse it across all draws instead of recomputing per attempt. + source_hash = cfg.cartridge_manifest()["hashes"]["refdata_content_hash"] + # Derive each attempt's sub-seed from a base RNG so a failed attempt at + # `seed` cannot collide with a direct draw at `seed + 1`. + base_rng = random.Random(seed) + + provenance = { + "allele_frequencies": freq_src, + "haplotype_deletion_prob": del_src, + "chromosome_weights": cw_src, + "novel_alleles": novel_src, + "model_id": plane.model_id if plane is not None else None, + "model_checksum": plane.content_checksum() if plane is not None else None, + } + + attempts = max_resamples if ensure_viable else 1 + for _attempt in range(attempts): + sub_seed = base_rng.getrandbits(63) + g = cls._draw_one(sampling_cfg, sub_seed, segs, freqs, delp, cw, subject_id, source_hash) + if not ensure_viable or g._is_viable(cfg, cw): + cls._rebind_to_base(g, cfg, novel_helper) + g.prior_provenance = dict(provenance) + return g + raise ValueError( + f"could not sample a viable genotype after {max_resamples} attempts; " + f"haplotype_deletion_prob is too high to leave a complete, positive-weight " + f"haplotype for required segments {cls._required_segments(cfg)} " + f"(chromosome_weights={cw})" + ) + + @staticmethod + def _resolve_plane(cfg, use_cartridge_priors): + if not use_cartridge_priors: + return None + return getattr(cfg, "genotype_priors", None) + + @classmethod + def _register_plane_novels(cls, cfg, plane): + """Build a throwaway Genotype carrying every plane novel as a registered + novel allele. This is where catalogue-aware + functional validation of + plane novels happens (via add_novel_allele). Returns the helper.""" + helper = cls.from_dataconfig(cfg) + for nv in plane.novel_alleles: + helper.add_novel_allele( + nv.name, base=nv.base_allele, sequence=nv.sequence.upper(), + segment=nv.segment, allow_nonfunctional=nv.allow_nonfunctional) + return helper + + @staticmethod + def _dataconfig_injecting_all_novels(helper): + """A cfg copy with ALL of the helper's registered novels injected as + catalogue alleles — the *sampling* reference (candidates), distinct from + effective_dataconfig() which injects only carried novels.""" + import copy as _copy + cfg = _copy.deepcopy(helper._cfg) + by_seg = {"V": cfg.v_alleles, "D": cfg.d_alleles, "J": cfg.j_alleles} + for name, info in helper._novel.items(): + d = by_seg[info["segment"]] + existing = list(d.get(info["gene"], [])) + existing.append(_copy.deepcopy(info["allele"])) + d[info["gene"]] = existing + return cfg + + @classmethod + def _augment_freqs_with_novels(cls, sampling_cfg, freq_spec, plane, segs): + """Return a nested freq spec that includes each plane novel in its gene's + table per the synthesis rule: authored table -> preserve + add novel; + no table -> catalogue alleles 1.0 + novel frequency. Collisions raise.""" + if freq_spec == "usage_as_prior": + nested = cls._usage_frequencies(sampling_cfg, segs) + nested = {seg: {g: dict(al) for g, al in genes.items()} + for seg, genes in nested.items()} + else: + nested = {seg: {g: dict(al) for g, al in genes.items()} + for seg, genes in cls._normalize_freq_spec(sampling_cfg, freq_spec, segs).items()} + for nv in plane.novel_alleles: + seg = nv.segment + if seg not in segs: + continue + gene = nv.name.split("*")[0] + gene_tbl = nested.setdefault(seg, {}).get(gene) + if gene_tbl is None: + # no authored table for this gene: catalogue 1.0 + novel + catalogue = _alleles_by_gene(sampling_cfg, seg).get(gene, []) + gene_tbl = {a.name: 1.0 for a in catalogue if a.name != nv.name} + nested[seg][gene] = gene_tbl + if nv.name in gene_tbl: + raise ValueError( + f"plane novel {nv.name!r} collides with an existing allele weight") + gene_tbl[nv.name] = float(nv.frequency) + return nested + + @staticmethod + def _rebind_to_base(g, base_cfg, novel_helper): + """Point a drawn genotype back at the base cfg and register only the + novels it actually carries, so effective_dataconfig() injects carried + novels (and nothing else). Task 7 supplies ``novel_helper``.""" + import copy as _copy + g._cfg = base_cfg + if novel_helper is None: + return + carried = g._carried_allele_names() + for name, info in novel_helper._novel.items(): + if name in carried: + g._novel[name] = _copy.deepcopy(info) + + @classmethod + def _draw_one(cls, cfg, seed, segs, freqs, delp, cw, subject_id, source_hash): + import random + + rng = random.Random(seed) + # Build a bare Genotype directly (avoid from_dataconfig, which recomputes + # the cartridge hash on every draw); reuse the precomputed source_hash. + g = cls.__new__(cls) + g._cfg = cfg + g._permissive = False + g.subject_id = subject_id + g._chromosome_weights = cw + g._slots = {s: {} for s in _SEGMENTS} + g._novel = {} + g._source_hash = source_hash + g.prior_provenance = cls._manual_provenance() # sample() overwrites with real sources + for seg in segs: + for gene in _alleles_by_gene(cfg, seg): + pdel = delp[seg][gene] + slots: List[List[Tuple[str, int, float]]] = [[], []] + for h in (0, 1): + if rng.random() < pdel: + continue # deleted on this chromosome + allele = cls._weighted_pick(rng, freqs[seg][gene]) + slots[h] = [(allele, 1, 1.0)] + g._slots[seg][gene] = slots + return g diff --git a/src/GenAIRR/genotype.py b/src/GenAIRR/genotype/model.py similarity index 51% rename from src/GenAIRR/genotype.py rename to src/GenAIRR/genotype/model.py index a896126..09a3423 100644 --- a/src/GenAIRR/genotype.py +++ b/src/GenAIRR/genotype/model.py @@ -17,19 +17,11 @@ import math from collections.abc import Mapping from typing import Dict, List, Optional, Set, Tuple +from ._common import _SEGMENTS, _alleles_by_gene +from ._sampling import _GenotypeSampling -_SEGMENTS = ("V", "D", "J") - -def _alleles_by_gene(cfg, segment: str) -> Dict[str, List]: - return { - "V": cfg.v_alleles, - "D": cfg.d_alleles, - "J": cfg.j_alleles, - }[segment] or {} - - -class Genotype: +class Genotype(_GenotypeSampling): """A diploid genotype over a ``DataConfig``'s reference alleles.""" def __init__(self, cfg, *, permissive: bool = False): @@ -208,7 +200,7 @@ def add_novel_allele( """ import copy as _copy - from .utilities.misc import translate + from ..utilities.misc import translate base_allele = self._find_ref_allele(segment, base) if base_allele is None: @@ -410,469 +402,6 @@ def _snapshot(self) -> "Genotype": return g # ── population sampling ─────────────────────────────────────── - @staticmethod - def _required_segments(cfg) -> List[str]: - req = ["V", "J"] - if _alleles_by_gene(cfg, "D"): - req.insert(1, "D") # V, D, J - return req - - @classmethod - def _resolve_sample_segments(cls, cfg, segments_to_sample) -> List[str]: - required = cls._required_segments(cfg) - if segments_to_sample is None: - return required - segs = list(segments_to_sample) - seen = set() - for s in segs: - if s not in _SEGMENTS: - raise ValueError(f"unknown segment {s!r}; expected one of {_SEGMENTS}") - if s in seen: - raise ValueError(f"segments_to_sample contains duplicate segment {s!r}") - seen.add(s) - if not _alleles_by_gene(cfg, s): - raise ValueError(f"cartridge has no {s} segment") - missing = [r for r in required if r not in seen] - if missing: - raise ValueError( - f"segments_to_sample must cover the chain's required segments " - f"{required}; missing {missing}. Partial sampling is not supported " - f"by Genotype.sample (it must return a runnable genotype)." - ) - # canonical _SEGMENTS order so the result is independent of input order - return [s for s in _SEGMENTS if s in seen] - - @classmethod - def _gene_segment_index(cls, cfg, segs) -> Dict[str, List[str]]: - """gene name -> [segments it appears in] (for flat-shape disambiguation).""" - idx: Dict[str, List[str]] = {} - for seg in segs: - for gene in _alleles_by_gene(cfg, seg): - idx.setdefault(gene, []).append(seg) - return idx - - @staticmethod - def _weighted_pick(rng, pairs): - names = [n for (n, _w) in pairs] - weights = [w for (_n, w) in pairs] - return rng.choices(names, weights=weights, k=1)[0] - - @classmethod - def _normalize_freq_spec(cls, cfg, spec, segs): - """Return nested ``{seg: {gene: {allele: weight}}}`` from a nested or flat - ``allele_frequencies`` spec, fully validating segment/gene/allele - addressing and shapes (unknown names and malformed values raise).""" - if spec is None: - return {} - if not isinstance(spec, Mapping): - raise ValueError( - "allele_frequencies must be a mapping (or 'usage_as_prior' / None), " - f"got {type(spec).__name__}" - ) - nested: Dict[str, Dict[str, Dict[str, float]]] = {} - keys = set(spec) - if keys and keys <= set(_SEGMENTS): # segment-keyed (nested) shape - for seg, genes in spec.items(): - if seg not in segs: - raise ValueError( - f"allele_frequencies: segment {seg!r} is not being sampled " - f"(sampling {segs})" - ) - if not isinstance(genes, Mapping): - raise ValueError( - f"allele_frequencies[{seg!r}] must be a mapping of " - f"gene -> {{allele: weight}}, got {type(genes).__name__}" - ) - catalogue = _alleles_by_gene(cfg, seg) - for gene, alleles in genes.items(): - if gene not in catalogue: - raise ValueError( - f"allele_frequencies: {seg} gene {gene!r} is not in the cartridge" - ) - if not isinstance(alleles, Mapping) or not alleles: - raise ValueError( - f"allele_frequencies[{seg!r}][{gene!r}] must be a non-empty " - f"mapping of allele -> weight" - ) - nested.setdefault(seg, {})[gene] = alleles - else: # flat {gene: {...}} shape - gidx = cls._gene_segment_index(cfg, segs) - for gene, alleles in spec.items(): - segs_for = gidx.get(gene) - if not segs_for: - raise ValueError(f"allele_frequencies: unknown gene {gene!r}") - if len(segs_for) > 1: - raise ValueError( - f"allele_frequencies: gene {gene!r} is ambiguous across segments " - f"{segs_for}; use the {{segment: {{gene: ...}}}} shape" - ) - if not isinstance(alleles, Mapping) or not alleles: - raise ValueError( - f"allele_frequencies[{gene!r}] must be a non-empty mapping of " - f"allele -> weight" - ) - nested.setdefault(segs_for[0], {})[gene] = alleles - return nested - - @classmethod - def _usage_frequencies(cls, cfg, segs): - rm = getattr(cfg, "reference_models", None) - usage = getattr(rm, "allele_usage", None) if rm else None - if usage is None: - raise ValueError( - "allele_frequencies='usage_as_prior' requires a cartridge with a " - "typed reference_models.allele_usage; this cartridge has none" - ) - nested: Dict[str, Dict[str, Dict[str, float]]] = {} - seg_attr = {"V": "v", "D": "d", "J": "j"} - for seg in segs: - table = getattr(usage, seg_attr[seg], None) or {} - if not table: - raise ValueError( - f"allele_frequencies='usage_as_prior': cartridge allele_usage has no " - f"entries for requested segment {seg!r}" - ) - catalogue = _alleles_by_gene(cfg, seg) - for allele_name, w in table.items(): - gene = allele_name.split("*")[0] - if gene not in catalogue: - raise ValueError( - f"allele_frequencies='usage_as_prior': usage allele {allele_name!r} " - f"maps to {seg} gene {gene!r}, which is not in the cartridge catalogue" - ) - nested.setdefault(seg, {}).setdefault(gene, {})[allele_name] = w - return nested - - @classmethod - def _resolve_allele_frequencies(cls, cfg, spec, segs): - if spec == "usage_as_prior": - nested = cls._usage_frequencies(cfg, segs) - else: - nested = cls._normalize_freq_spec(cfg, spec, segs) - out: Dict[str, Dict[str, List[Tuple[str, float]]]] = {} - for seg in segs: - out[seg] = {} - for gene, alleles in _alleles_by_gene(cfg, seg).items(): - names = {a.name for a in alleles} - supplied = nested.get(seg, {}).get(gene) - if supplied is None: - out[seg][gene] = [(a.name, 1.0) for a in alleles] # uniform fallback - continue - pairs: List[Tuple[str, float]] = [] - total = 0.0 - # Sort by allele name so the draw is independent of the spec's - # dict insertion order — two content-equal models (same - # content_checksum) then produce identical draws at a given seed. - for nm, w in sorted(supplied.items()): - if nm not in names: - raise ValueError(f"{seg} gene {gene}: {nm!r} is not a known allele") - if isinstance(w, bool) or not isinstance(w, (int, float)) or not math.isfinite(w) or w < 0: - raise ValueError( - f"{seg} gene {gene}: weight for {nm!r} must be finite and >= 0, got {w!r}" - ) - if w > 0: - pairs.append((nm, float(w))) - total += w - if total <= 0 or not pairs: - raise ValueError( - f"{seg} gene {gene}: at least one allele weight must be > 0" - ) - out[seg][gene] = pairs - return out - - @classmethod - def _resolve_haplotype_deletion(cls, cfg, spec, segs): - def _check(p, where): - if isinstance(p, bool) or not isinstance(p, (int, float)) or not math.isfinite(p): - raise ValueError(f"{where}: deletion probability must be a finite number, got {p!r}") - if not (0.0 <= p <= 1.0): - raise ValueError(f"{where}: deletion probability must be in [0, 1], got {p}") - return float(p) - - out = {seg: {gene: 0.0 for gene in _alleles_by_gene(cfg, seg)} for seg in segs} - if isinstance(spec, (int, float)) and not isinstance(spec, bool): - p = _check(spec, "haplotype_deletion_prob") - for seg in segs: - for gene in out[seg]: - out[seg][gene] = p - return out - if not isinstance(spec, Mapping): - raise ValueError( - f"haplotype_deletion_prob must be a float or a mapping, got {type(spec).__name__}" - ) - keys = set(spec) - if keys and keys <= set(_SEGMENTS): # nested {seg: {gene: prob}} - for seg, genes in spec.items(): - if seg not in segs: - raise ValueError(f"haplotype_deletion_prob: segment {seg!r} not being sampled") - if not isinstance(genes, Mapping): - raise ValueError( - f"haplotype_deletion_prob[{seg!r}] must be a mapping of " - f"gene -> probability, got {type(genes).__name__}" - ) - for gene, p in genes.items(): - if gene not in out[seg]: - raise ValueError(f"haplotype_deletion_prob: unknown {seg} gene {gene!r}") - out[seg][gene] = _check(p, f"haplotype_deletion_prob[{seg}][{gene}]") - else: # flat {gene: prob} - gidx = cls._gene_segment_index(cfg, segs) - for gene, p in spec.items(): - segs_for = gidx.get(gene) - if not segs_for: - raise ValueError(f"haplotype_deletion_prob: unknown gene {gene!r}") - if len(segs_for) > 1: - raise ValueError( - f"haplotype_deletion_prob: gene {gene!r} is ambiguous across " - f"segments {segs_for}; use the {{segment: {{gene: prob}}}} shape" - ) - out[segs_for[0]][gene] = _check(p, f"haplotype_deletion_prob[{gene}]") - return out - - @classmethod - def sample( - cls, - cfg, - *, - seed: int = 0, - allele_frequencies=None, - haplotype_deletion_prob=None, - segments_to_sample=None, - chromosome_weights: Optional[Tuple[float, float]] = None, - subject_id: Optional[str] = None, - ensure_viable: bool = True, - max_resamples: int = 1000, - use_cartridge_priors: bool = True, - include_cartridge_novel_alleles="auto", - ) -> "Genotype": - """Sample a fully-specified diploid genotype from population priors. - - Independent per-gene, per-chromosome Hardy-Weinberg model: each gene on - each chromosome is independently deleted or assigned one allele drawn - from the gene's allele frequencies. Homozygous/heterozygous/hemizygous/ - deleted states emerge at Hardy-Weinberg rates. - - When ``cfg`` carries a ``genotype_priors`` plane and an argument is left - at its default, the plane supplies it (``use_cartridge_priors=False`` - disables ALL plane consumption — a clean uniform catalogue-only draw). - Each input is sourced independently and recorded in - ``g.prior_provenance`` (explicit / cartridge / uniform / default). - - This is NOT a population haplotype model — no linkage disequilibrium, gene - co-deletion blocks, ancestry, or donor-specific haplotype structure. - - With ``ensure_viable=True`` (default), the draw is repeated (with a - deterministic sub-seed) up to ``max_resamples`` times until at least one - **positive-weight** chromosome carries every required segment, raising - ``ValueError`` if that is impossible under the given deletion settings. - - NOTE: with the default ``ensure_viable=True`` the result is Hardy-Weinberg - **conditioned on viability** (draws with no complete usable haplotype are - rejected), not the unconditional HW distribution. Use - ``ensure_viable=False`` for the raw (possibly infeasible) HW draw. - """ - import random - - if isinstance(max_resamples, bool) or not isinstance(max_resamples, int) or max_resamples < 1: - raise ValueError(f"max_resamples must be an int >= 1, got {max_resamples!r}") - if not isinstance(use_cartridge_priors, bool): - raise ValueError( - f"use_cartridge_priors must be a bool, got {use_cartridge_priors!r}") - # Identity checks (not `in (...)`): Python's `1 == True` / `0 == False` - # would otherwise let integers slip through and silently act like False. - _icna = include_cartridge_novel_alleles - if not (_icna is True or _icna is False or _icna == "auto"): - raise ValueError( - "include_cartridge_novel_alleles must be 'auto', True, or False, " - f"got {include_cartridge_novel_alleles!r}") - - plane = cls._resolve_plane(cfg, use_cartridge_priors) - if plane is not None: - # A plane attached via the builder is already validated, but a plane - # set directly on the DataConfig bypasses that — validate before use so - # a malformed prior (empty model_id, NaN weights, D-on-VJ) fails loudly - # rather than being silently sampled and stamped into provenance. - plane.validate(chain_type=getattr(getattr(cfg, "metadata", None), "chain_type", None)) - - # Per-input source resolution. - if allele_frequencies is not None: - freq_spec, freq_src = allele_frequencies, "explicit" - elif plane is not None and plane.allele_frequencies: - freq_spec, freq_src = plane.allele_frequencies, "cartridge" - else: - freq_spec, freq_src = None, "uniform" - - if haplotype_deletion_prob is not None: - del_spec, del_src = haplotype_deletion_prob, "explicit" - elif plane is not None and plane.haplotype_deletion_prob: - del_spec, del_src = plane.haplotype_deletion_prob, "cartridge" - else: - del_spec, del_src = 0.0, "default" - - if chromosome_weights is not None: - cw_in, cw_src = chromosome_weights, "explicit" - elif plane is not None: - cw_in, cw_src = plane.chromosome_weights, "cartridge" - else: - cw_in, cw_src = (0.5, 0.5), "default" - - cw = cls._check_chromosome_weights(*cw_in) - segs = cls._resolve_sample_segments(cfg, segments_to_sample) - - # Candidate-novel injection. Plane novels are draw CANDIDATES (injected - # into the sampling cfg) only when their source is active; the returned - # genotype still exports only CARRIED novels via effective_dataconfig(). - inject_novels = False - if plane is not None and plane.novel_alleles: - if include_cartridge_novel_alleles is True: - inject_novels = True - elif include_cartridge_novel_alleles == "auto": - inject_novels = freq_src in ("cartridge", "uniform") - # False -> never - novel_src = "none" - sampling_cfg = cfg - novel_helper = None - if inject_novels: - novel_helper = cls._register_plane_novels(cfg, plane) - sampling_cfg = cls._dataconfig_injecting_all_novels(novel_helper) - freq_spec = cls._augment_freqs_with_novels(sampling_cfg, freq_spec, plane, segs) - novel_src = "cartridge" - - freqs = cls._resolve_allele_frequencies(sampling_cfg, freq_spec, segs) - delp = cls._resolve_haplotype_deletion(sampling_cfg, del_spec, segs) - # Compute the cartridge content hash ONCE (it rebuilds refdata + hashes); - # reuse it across all draws instead of recomputing per attempt. - source_hash = cfg.cartridge_manifest()["hashes"]["refdata_content_hash"] - # Derive each attempt's sub-seed from a base RNG so a failed attempt at - # `seed` cannot collide with a direct draw at `seed + 1`. - base_rng = random.Random(seed) - - provenance = { - "allele_frequencies": freq_src, - "haplotype_deletion_prob": del_src, - "chromosome_weights": cw_src, - "novel_alleles": novel_src, - "model_id": plane.model_id if plane is not None else None, - "model_checksum": plane.content_checksum() if plane is not None else None, - } - - attempts = max_resamples if ensure_viable else 1 - for _attempt in range(attempts): - sub_seed = base_rng.getrandbits(63) - g = cls._draw_one(sampling_cfg, sub_seed, segs, freqs, delp, cw, subject_id, source_hash) - if not ensure_viable or g._is_viable(cfg, cw): - cls._rebind_to_base(g, cfg, novel_helper) - g.prior_provenance = dict(provenance) - return g - raise ValueError( - f"could not sample a viable genotype after {max_resamples} attempts; " - f"haplotype_deletion_prob is too high to leave a complete, positive-weight " - f"haplotype for required segments {cls._required_segments(cfg)} " - f"(chromosome_weights={cw})" - ) - - @staticmethod - def _resolve_plane(cfg, use_cartridge_priors): - if not use_cartridge_priors: - return None - return getattr(cfg, "genotype_priors", None) - - @classmethod - def _register_plane_novels(cls, cfg, plane): - """Build a throwaway Genotype carrying every plane novel as a registered - novel allele. This is where catalogue-aware + functional validation of - plane novels happens (via add_novel_allele). Returns the helper.""" - helper = cls.from_dataconfig(cfg) - for nv in plane.novel_alleles: - helper.add_novel_allele( - nv.name, base=nv.base_allele, sequence=nv.sequence.upper(), - segment=nv.segment, allow_nonfunctional=nv.allow_nonfunctional) - return helper - - @staticmethod - def _dataconfig_injecting_all_novels(helper): - """A cfg copy with ALL of the helper's registered novels injected as - catalogue alleles — the *sampling* reference (candidates), distinct from - effective_dataconfig() which injects only carried novels.""" - import copy as _copy - cfg = _copy.deepcopy(helper._cfg) - by_seg = {"V": cfg.v_alleles, "D": cfg.d_alleles, "J": cfg.j_alleles} - for name, info in helper._novel.items(): - d = by_seg[info["segment"]] - existing = list(d.get(info["gene"], [])) - existing.append(_copy.deepcopy(info["allele"])) - d[info["gene"]] = existing - return cfg - - @classmethod - def _augment_freqs_with_novels(cls, sampling_cfg, freq_spec, plane, segs): - """Return a nested freq spec that includes each plane novel in its gene's - table per the synthesis rule: authored table -> preserve + add novel; - no table -> catalogue alleles 1.0 + novel frequency. Collisions raise.""" - if freq_spec == "usage_as_prior": - nested = cls._usage_frequencies(sampling_cfg, segs) - nested = {seg: {g: dict(al) for g, al in genes.items()} - for seg, genes in nested.items()} - else: - nested = {seg: {g: dict(al) for g, al in genes.items()} - for seg, genes in cls._normalize_freq_spec(sampling_cfg, freq_spec, segs).items()} - for nv in plane.novel_alleles: - seg = nv.segment - if seg not in segs: - continue - gene = nv.name.split("*")[0] - gene_tbl = nested.setdefault(seg, {}).get(gene) - if gene_tbl is None: - # no authored table for this gene: catalogue 1.0 + novel - catalogue = _alleles_by_gene(sampling_cfg, seg).get(gene, []) - gene_tbl = {a.name: 1.0 for a in catalogue if a.name != nv.name} - nested[seg][gene] = gene_tbl - if nv.name in gene_tbl: - raise ValueError( - f"plane novel {nv.name!r} collides with an existing allele weight") - gene_tbl[nv.name] = float(nv.frequency) - return nested - - @staticmethod - def _rebind_to_base(g, base_cfg, novel_helper): - """Point a drawn genotype back at the base cfg and register only the - novels it actually carries, so effective_dataconfig() injects carried - novels (and nothing else). Task 7 supplies ``novel_helper``.""" - import copy as _copy - g._cfg = base_cfg - if novel_helper is None: - return - carried = g._carried_allele_names() - for name, info in novel_helper._novel.items(): - if name in carried: - g._novel[name] = _copy.deepcopy(info) - - @classmethod - def _draw_one(cls, cfg, seed, segs, freqs, delp, cw, subject_id, source_hash): - import random - - rng = random.Random(seed) - # Build a bare Genotype directly (avoid from_dataconfig, which recomputes - # the cartridge hash on every draw); reuse the precomputed source_hash. - g = cls.__new__(cls) - g._cfg = cfg - g._permissive = False - g.subject_id = subject_id - g._chromosome_weights = cw - g._slots = {s: {} for s in _SEGMENTS} - g._novel = {} - g._source_hash = source_hash - g.prior_provenance = cls._manual_provenance() # sample() overwrites with real sources - for seg in segs: - for gene in _alleles_by_gene(cfg, seg): - pdel = delp[seg][gene] - slots: List[List[Tuple[str, int, float]]] = [[], []] - for h in (0, 1): - if rng.random() < pdel: - continue # deleted on this chromosome - allele = cls._weighted_pick(rng, freqs[seg][gene]) - slots[h] = [(allele, 1, 1.0)] - g._slots[seg][gene] = slots - return g def _is_viable(self, cfg, chromosome_weights=None) -> bool: """A genotype is viable iff at least one chromosome that can actually be diff --git a/src/GenAIRR/genotype_priors.py b/src/GenAIRR/genotype_priors.py index 8dfdb73..3f2f0b9 100644 --- a/src/GenAIRR/genotype_priors.py +++ b/src/GenAIRR/genotype_priors.py @@ -4,8 +4,7 @@ This is NOT an empirical recombination model (those live on ``ReferenceEmpiricalModels``); it is a per-gene carriage/deletion prior plus a catalogue of population novel/private alleles. ``Genotype.sample`` consumes it to -draw a per-individual diploid genotype. See -``.private/specs/2026-06-17-genotype-cartridge-plane-design.md``. +draw a per-individual diploid genotype. The plane stays decoupled from a specific catalogue: ``validate()`` does shape + numeric/DNA sanity only. Catalogue-aware checks (gene/allele existence, novel @@ -25,6 +24,8 @@ _DNA = set("ACGT") SCHEMA_TAG = "population_genotype_model/1" +__all__ = ["PopulationNovelAllele", "PopulationGenotypeModel"] + @dataclass class PopulationNovelAllele: diff --git a/src/GenAIRR/mcp_server.py b/src/GenAIRR/mcp_server.py index 6ba065c..88fe057 100644 --- a/src/GenAIRR/mcp_server.py +++ b/src/GenAIRR/mcp_server.py @@ -63,7 +63,7 @@ def _all_config_aliases() -> List[str]: be passing to the other tools. We surface the aliases instead so the output of ``list_configs`` round-trips directly into every other tool. """ - from GenAIRR.experiment import _CONFIG_ALIASES + from GenAIRR._refdata_resolver import _CONFIG_ALIASES return sorted(_CONFIG_ALIASES.keys()) diff --git a/src/GenAIRR/reference_models.py b/src/GenAIRR/reference_models.py index a8ee88e..9bd926f 100644 --- a/src/GenAIRR/reference_models.py +++ b/src/GenAIRR/reference_models.py @@ -48,6 +48,18 @@ P_NUCLEOTIDE_END_KEYS: Tuple[str, ...] = ("V_3", "D_5", "D_3", "J_5") P_NUCLEOTIDE_END_KEYS_VJ: Tuple[str, ...] = ("V_3", "J_5") +__all__ = [ + "NP_KEYS", + "TRIM_KEYS", + "TRIM_KEYS_VJ", + "P_NUCLEOTIDE_END_KEYS", + "P_NUCLEOTIDE_END_KEYS_VJ", + "EmpiricalDistributionSpec", + "NpBaseModelSpec", + "AlleleUsageSpec", + "ReferenceEmpiricalModels", +] + @dataclass class EmpiricalDistributionSpec: diff --git a/src/GenAIRR/reference_rules.py b/src/GenAIRR/reference_rules.py index dffbfa6..2b51822 100644 --- a/src/GenAIRR/reference_rules.py +++ b/src/GenAIRR/reference_rules.py @@ -40,6 +40,8 @@ # semantics live elsewhere. _VALID_AA = set("ACDEFGHIKLMNPQRSTVWY*") +__all__ = ["AnchorRuleSpec", "ReferenceRulesSpec"] + @dataclass class AnchorRuleSpec: diff --git a/src/GenAIRR/result.py b/src/GenAIRR/result.py index 05f4b69..9a46cb0 100644 --- a/src/GenAIRR/result.py +++ b/src/GenAIRR/result.py @@ -14,401 +14,22 @@ """ from __future__ import annotations -import csv from typing import Any, Dict, Iterator, List, Optional, Sequence, Union, overload - -# Canonical AIRR-style column order used by ``to_csv`` / ``to_tsv``. -# Anchors the header row to a stable shape across runs and keeps -# the output diff-friendly. -# Coordinate fields whose values are 0-based half-open (Python -# convention) by default. ``airr_strict=True`` exports add 1 to each -# *_start field to convert to the AIRR-spec 1-based-inclusive form; -# the *_end fields keep their existing value because Python's -# half-open `end` already equals the 1-based-inclusive `end`. -_AIRR_STRICT_START_FIELDS = ( - "v_sequence_start", - "d_sequence_start", - "j_sequence_start", - "v_alignment_start", - "d_alignment_start", - "j_alignment_start", - "v_germline_start", - "d_germline_start", - "j_germline_start", - "junction_start", +from ._result_common import _inject_truth_columns +from ._result_export import ( + _ResultExport, + _DEFAULT_COLUMN_ORDER, + _to_airr_strict, ) - - -def _to_airr_strict(rec: Dict[str, Any]) -> Dict[str, Any]: - """Return a copy of ``rec`` with every coord ``*_start`` field - converted from 0-based half-open to 1-based-inclusive. ``*_end`` - fields are unchanged because the Python half-open end already - equals the 1-based-inclusive end (e.g. ``[0, 5)`` ↔ ``[1, 5]``). - - ``None`` start values are left as-is. Unset / missing keys are - silently ignored. - """ - converted = dict(rec) - for field in _AIRR_STRICT_START_FIELDS: - val = converted.get(field) - if isinstance(val, int): - converted[field] = val + 1 - return converted - - -_DEFAULT_COLUMN_ORDER = [ - # AIRR metadata - "sequence_id", - "sequence", - "sequence_aa", - "sequence_alignment", - "germline_alignment", - "germline_alignment_d_mask", - "sequence_length", - "rev_comp", - "locus", - # Calls - "v_call", - "v_cigar", - "v_score", - "v_identity", - "v_support", - "v_sequence_start", - "v_sequence_end", - "v_alignment_start", - "v_alignment_end", - "v_germline_start", - "v_germline_end", - "v_trim_5", - "v_trim_3", - "d_call", - "d_cigar", - "d_score", - "d_identity", - "d_support", - "d_sequence_start", - "d_sequence_end", - "d_alignment_start", - "d_alignment_end", - "d_germline_start", - "d_germline_end", - "d_trim_5", - "d_trim_3", - "j_call", - "j_cigar", - "j_score", - "j_identity", - "j_support", - "j_sequence_start", - "j_sequence_end", - "j_alignment_start", - "j_alignment_end", - "j_germline_start", - "j_germline_end", - "j_trim_5", - "j_trim_3", - "c_call", - # Junction - "junction", - "junction_aa", - "junction_start", - "junction_end", - "junction_length", - # NP regions - "np1", - "np1_aa", - "np1_length", - "np2", - "np2_aa", - "np2_length", - # Functionality - "productive", - "vj_in_frame", - "stop_codon", - # SHM + corruption (non-AIRR; GenAIRR additions) - "n_mutations", - "mutation_rate", - "n_pcr_errors", - "n_quality_errors", - "n_indels", - # Per-segment indel counters (docs/indel_provenance_audit.md - # §6.2). NP1/NP2 indels are excluded from these but still - # count toward `n_indels`. - "n_v_indels", - "n_d_indels", - "n_j_indels", - # Per-segment SHM mutation counters - # (docs/mutation_provenance_audit.md). Aggregated by walking - # `outcome.events()`, filtering to the mutate.{uniform,s5f} - # passes only, and bucketing each `BaseChanged` by carried - # segment. NP1+NP2 roll into `n_np_mutations`. Sum equals - # `n_mutations` by construction. - "n_v_mutations", - "n_d_mutations", - "n_j_mutations", - "n_np_mutations", - # V-subregion SHM partition - # (docs/v_subregion_mutation_counters_audit.md). Six fields - # that partition `n_v_mutations` by the assigned V allele's - # IMGT subregion intervals: five canonical labels plus - # `n_v_unannotated_mutations` for V events that can't be - # attributed (missing assignment, empty annotations, V-side - # CDR3 stretch, or indel-inserted V bases). Aggregated in the - # same `outcome.events()` walk as the per-segment counters, - # using the same `mutate.{uniform,s5f}` pass-name filter. - # On bundled human OGRDB cartridges the unannotated bucket is - # 0 on every record under the canonical pass order. - "n_fwr1_mutations", - "n_cdr1_mutations", - "n_fwr2_mutations", - "n_cdr2_mutations", - "n_fwr3_mutations", - "n_v_unannotated_mutations", - # Observation-stage length loss (EndLossPass / primer_trim_*). - # Distinct from recombination-stage v_trim_*/j_trim_*. See - # docs/primer_trim_end_loss_audit.md §6.1. - "end_loss_5_length", - "end_loss_3_length", - "is_contaminant", - # D inversion provenance (V(D)J inversion event). True when the - # simulation committed the D allele in reverse-complement - # orientation; false for VJ chains, VDJ chains without - # `Experiment.invert_d(...)`, and inversion decisions that - # landed on the forward branch. See - # `docs/d_inversion_design.md` §6.3. - "d_inverted", - # Receptor revision provenance (Slice E of the receptor-revision - # roadmap). `receptor_revision_applied` is True iff the - # `ReceptorRevisionPass` fired with applied=True; `original_v_call` - # carries the V allele name the recombine pass originally committed - # (empty when no revision happened). `v_call` continues to report - # the post-revision identity. See - # `docs/receptor_revision_design.md` §7. - "receptor_revision_applied", - "original_v_call", - # Paired-end / read layout (Slice A of the paired-end roadmap). - # All eight fields default to empty / None / 0 / empty under - # the legacy single-molecule projection; the Slice B/C projection - # layer populates them. The order here matches the Rust struct - # layout in `engine_rs/src/airr_record/record.rs`. See - # `docs/paired_end_design.md` §10. - "read_layout", - "r1_sequence", - "r2_sequence", - "r1_start", - "r1_end", - "r2_start", - "r2_end", - "insert_size", -] - - -def _allele_name_or_empty(refdata: Any, segment: str, allele_id: Optional[int]) -> str: - """Look up an allele name from refdata by id; return ``""`` when - the id is None or the lookup fails (defensive). - """ - if allele_id is None: - return "" - try: - if segment == "V": - return refdata.v_allele(int(allele_id)).name - if segment == "D": - return refdata.d_allele(int(allele_id)).name - if segment == "J": - return refdata.j_allele(int(allele_id)).name - except Exception: - return "" - return "" - - -def _inject_truth_columns(outcome: Any, refdata: Any, record: Dict[str, Any]) -> None: - """Append `truth_v_call` / `truth_d_call` / `truth_j_call` - columns to ``record`` from the originally-sampled allele ids - stored in the simulation's `assignments`. Distinct from - `v_call` / `d_call` / `j_call`, which are evidence-driven and - can change under heavy SHM. - """ - sim = outcome.final_simulation() - record["truth_v_call"] = _allele_name_or_empty(refdata, "V", sim.v_allele_id()) - record["truth_d_call"] = _allele_name_or_empty(refdata, "D", sim.d_allele_id()) - record["truth_j_call"] = _allele_name_or_empty(refdata, "J", sim.j_allele_id()) - - -class ValidationReport: - """Aggregate report from :meth:`SimulationResult.validate_records`. - - Carries the result of the **public AIRR output correctness** gate - over a batch — the downstream contract that says "every projected - record is internally consistent with its outcome." - - Attributes: - ``count`` — total records checked. - ``failures`` — list of failing records, each as a dict with - ``record_index``, ``sequence_id``, and ``issues`` - (the list of dicts returned by - :meth:`Outcome.validate_record`). - ``ok`` — True iff every record passed (``failures`` is - empty). - - The report is truthy iff ``ok`` is True, so ``assert report`` - works as a one-line CI guard. - - Sibling gate: :meth:`Outcome.check_live_call_cache_parity` — - internal cache-correctness check on the state that feeds - projection. If both fail on the same batch, fix parity first - (a stale cache leaks into projection); see - :meth:`SimulationResult.validate_records` docstring for the - full troubleshooting rule. - """ - - __slots__ = ("count", "failures") - - def __init__(self, count: int, failures: List[Dict[str, Any]]) -> None: - self.count = count - self.failures: List[Dict[str, Any]] = failures - - @property - def ok(self) -> bool: - """True iff every record validated without issues.""" - return not self.failures - - def __bool__(self) -> bool: - return self.ok - - def __len__(self) -> int: - """Number of failing records (not total). Use ``.count`` for - the total.""" - return len(self.failures) - - def __repr__(self) -> str: - if self.ok: - return f"" - return ( - f"" - ) - - def summary(self) -> Dict[str, int]: - """Histogram of issue kinds across all failures. Useful for - ``print(report.summary())`` to see what's wrong at a glance.""" - counts: Dict[str, int] = {} - for failure in self.failures: - for issue in failure["issues"]: - kind = issue.get("kind", "Unknown") - counts[kind] = counts.get(kind, 0) + 1 - return counts - - -class FamilyValidationReport: - """Aggregate report from :meth:`SimulationResult.validate_families`. - - Carries the result of the **clonal-family consistency** gate - over a batch — the downstream contract that says "every group of - records sharing a ``clone_id`` agrees on the recombination-time - truth fields the engine puts on each descendant by construction." - - Attributes: - ``count`` — total records inspected. - ``family_count`` — number of distinct ``clone_id`` groups - (``0`` when the batch carries no - clonal records). - ``members_per_family`` — mapping ``{clone_id: n_members}``; - empty when ``family_count == 0``. - ``failures`` — list of failing-group dicts. Each - carries ``clone_id``, - ``record_indices``, ``issue_kind``, - and ``values`` (the divergent values - observed in the group, sorted). - ``ok`` — ``True`` iff ``failures`` is empty. - - The report is truthy iff ``ok`` is True so ``assert report`` - works as a one-line CI guard. - - Sibling gate: :class:`ValidationReport` — the per-record AIRR - projection postcondition check. Family validation runs *after* - that gate passes; a family-divergence failure means the records - are each internally consistent with their own outcome, but they - disagree across the clonal group on a field that should be - invariant by construction (heavy-SHM tie-set widening on the - live-call ``v_call`` etc. is intentionally NOT a divergence — - only ``truth_*_call`` invariance is enforced). - - The validator is a strict subset of the audit's §6 spec - (see ``docs/clonal_family_design.md``): pre-SHM junction - invariance, mutation-distance distribution, parent-trace - reconstruction, and lineage topology are deliberately deferred - to later slices. - """ - - __slots__ = ("count", "family_count", "members_per_family", "failures") - - def __init__( - self, - count: int, - family_count: int, - members_per_family: Dict[Any, int], - failures: List[Dict[str, Any]], - ) -> None: - self.count = count - self.family_count = family_count - self.members_per_family: Dict[Any, int] = dict(members_per_family) - self.failures: List[Dict[str, Any]] = failures - - @property - def ok(self) -> bool: - """True iff every family group passed every invariant check.""" - return not self.failures - - def __bool__(self) -> bool: - return self.ok - - def __len__(self) -> int: - """Number of failing groups (not total). Use ``.count`` for - the total record count and ``.family_count`` for the total - group count.""" - return len(self.failures) - - def __repr__(self) -> str: - if self.ok: - return ( - f"" - ) - return ( - f"" - ) - - def summary(self) -> Dict[str, int]: - """Histogram of issue kinds across all failing groups. Useful - for ``print(report.summary())`` to see what's wrong at a - glance.""" - counts: Dict[str, int] = {} - for failure in self.failures: - kind = failure.get("issue_kind", "Unknown") - counts[kind] = counts.get(kind, 0) + 1 - return counts - - -# ── family-layer invariants ────────────────────────────────────────── -# -# (Field, issue_kind) pairs the family validator enforces. Each pair -# names a recombination-time provenance field that should be -# identical across every descendant of a clone *when the field is -# present*. We deliberately skip the field for a group when it is -# absent from every record in the group (e.g. ``expose_provenance`` -# was not enabled), and require it to be present on every record once -# any record in the group carries it. -_FAMILY_TRUTH_INVARIANTS: tuple = ( - ("truth_v_call", "TruthVCallDiverges"), - ("truth_d_call", "TruthDCallDiverges"), - ("truth_j_call", "TruthJCallDiverges"), +from ._result_validation import ( + _ResultValidation, + ValidationReport, + FamilyValidationReport, ) -class SimulationResult: +class SimulationResult(_ResultExport, _ResultValidation): """List-like wrapper around a batch of AIRR records. ``result[i]`` returns the i-th record dict; ``len(result)`` is @@ -550,785 +171,6 @@ def __getitem__( def __repr__(self) -> str: return f"" - # ── validation ────────────────────────────────────────────────── - - def validate_records(self, refdata: Any) -> "ValidationReport": - """**Public AIRR output correctness check.** - - Run the postcondition validator over every record in this - result and return a :class:`ValidationReport`. This is the - gate a downstream consumer cares about: "is each projected - AIRR record internally consistent with the engine state that - produced it?" - - Each record is re-derived independently from its original - ``Outcome`` (trace + event ledger + final ``Simulation``) - and compared against the projected dict. A record passes - when ``outcome.validate_record(refdata, sequence_id=...)`` - returns an empty issue list. Failures collect the - ``record_index``, ``sequence_id``, and the issue dicts. - - **Companion check** — for engine-side integrity, see - :meth:`Outcome.check_live_call_cache_parity` (returns the - cached-vs-fresh divergence on the live-call cache that - *feeds* projection). - - **Troubleshooting rule** — if a CI run has both this - validator AND the parity harness failing on the same batch, - fix the parity divergence FIRST: a stale cache can leak - into projection and produce spurious validator failures. - Once parity is green, rerun the validator; remaining - failures point at a real projection-layer bug. - - ``refdata`` must be the same :class:`RefDataConfig` the - outcomes were produced against; passing a different refdata - will misreport mismatches against an unrelated allele pool. - - Raises ``RuntimeError`` when this result was built without - attached outcomes (e.g. loaded from a TSV); the validator - needs the engine state, not just the projected record. - """ - if self._outcomes is None: - raise RuntimeError( - "validate_records requires the original Outcome objects " - "(with their trace + event ledger), but this " - "SimulationResult was built from records only (e.g. " - "loaded from TSV). Re-run the simulation with " - "Experiment.run_records() to get a result with attached " - "outcomes." - ) - failures: List[Dict[str, Any]] = [] - for i, (outcome, record) in enumerate(zip(self._outcomes, self._records)): - sequence_id = str(record.get("sequence_id", "")) - issues = outcome.validate_record(refdata, sequence_id=sequence_id) - if issues: - failures.append( - { - "record_index": i, - "sequence_id": sequence_id, - "issues": issues, - } - ) - return ValidationReport(count=len(self._outcomes), failures=failures) - - def validate_families( - self, refdata: Any = None - ) -> "FamilyValidationReport": - """**Clonal-family consistency check** — a strict subset of - the audit's §6 family-layer invariants - (``docs/clonal_family_design.md``). - - Groups records by ``clone_id`` and asserts the recombination- - time truth fields agree across every descendant of a clone. - ``refdata`` is reserved for forward compatibility with the - deeper family-layer checks (pre-SHM junction, mutation- - distance distribution) the audit's later slices add; this - slice's invariants are all dict-only and ignore ``refdata``. - - **Currently enforced invariants:** - - - ``truth_v_call`` constant within each ``clone_id``, **when - present**. Skipped silently for clones whose records were - projected without ``expose_provenance=True``. - - ``truth_d_call`` same. - - ``truth_j_call`` same. - - ``clone_id`` is present on every record once any record - in the batch carries one (a batch that mixes clonal and - non-clonal records raises ``CloneIdMissing``). - - **Non-clonal results return ok with ``family_count == 0``.** - This makes ``result.validate_families()`` a safe no-op on a - flat batch — the call site does not need to branch on the - result's clonal-ness. - - **Records-only results work.** Unlike - :meth:`validate_records`, this validator does not require - the underlying ``Outcome`` objects — every check is on - record-dict fields — so a ``SimulationResult`` loaded from - TSV can still be family-validated. - - **Not enforced yet** (deferred per the audit's §14 - out-of-scope list): mutation-distance distribution, pre-SHM - junction invariance, parent-trace reconstruction, lineage - topology, ``original_v_call`` / ``d_inverted`` invariance. - - Returns a :class:`FamilyValidationReport` carrying - ``count``, ``family_count``, ``members_per_family``, and - ``failures``. - """ - del refdata # forward-compat placeholder, see docstring. - - records = self._records - total = len(records) - - # Identify whether any record carries a non-null ``clone_id``. - # A batch that has no clonal records returns an ok no-op - # report; a batch where SOME records carry a clone_id and - # others don't is structurally broken (mixed clonal/flat - # outputs) and surfaces a ``CloneIdMissing`` failure. - any_clonal = any( - "clone_id" in r and r["clone_id"] is not None for r in records - ) - if not any_clonal: - return FamilyValidationReport( - count=total, - family_count=0, - members_per_family={}, - failures=[], - ) - - failures: List[Dict[str, Any]] = [] - missing_clone_id: List[int] = [ - i for i, r in enumerate(records) - if r.get("clone_id") is None - ] - if missing_clone_id: - failures.append( - { - "clone_id": None, - "record_indices": missing_clone_id, - "issue_kind": "CloneIdMissing", - "values": [], - } - ) - - # Group records by clone_id (ignoring records missing the tag - # — they already surfaced above). - by_clone: Dict[Any, List[int]] = {} - for i, r in enumerate(records): - cid = r.get("clone_id") - if cid is None: - continue - by_clone.setdefault(cid, []).append(i) - - for cid, indices in by_clone.items(): - members = [records[i] for i in indices] - for field, kind in _FAMILY_TRUTH_INVARIANTS: - # Skip the field for this group if no record carries - # it (consumer opted out of expose_provenance). Once - # any record in the group carries it, every record - # in the group must — the consistency check still - # runs but the "missing on some" surface would be a - # different bug; we treat ``None`` and absent as the - # same "no opinion" signal here. - present = [ - rec.get(field) for rec in members - if field in rec and rec.get(field) is not None - ] - if not present: - continue - distinct = sorted({str(v) for v in present}) - if len(distinct) > 1: - failures.append( - { - "clone_id": cid, - "record_indices": list(indices), - "issue_kind": kind, - "values": distinct, - } - ) - - members_per_family = {cid: len(idxs) for cid, idxs in by_clone.items()} - return FamilyValidationReport( - count=total, - family_count=len(by_clone), - members_per_family=members_per_family, - failures=failures, - ) - - def validate_families_with_parents( - self, refdata: Any = None - ) -> "FamilyValidationReport": - """**Parent-aware clonal-family validator** — the deeper - diagnostic that compares every descendant against its - actual parent ``Outcome`` (Slice 3 of the clonal-family - audit; see ``docs/clonal_parent_outcome_design.md`` §6). - - Sibling of :meth:`validate_families`. That validator is - record-only (groups by ``clone_id``, compares truth fields - across siblings); this one **requires the parent outcomes - to be available on the result** and compares each - descendant against its parent directly. Use this when you - want to confirm "the descendants reflect the recombination - ancestor they came from," not just "siblings agree with - each other." - - **Currently enforced invariants** — all derived from - record-vs-parent comparison only: - - - Structural: - - ``ParentsMissing`` when records carry ``clone_id`` / - ``parent_id`` but ``self.parents`` is ``None``. - - ``ParentIdMissing`` for records without a non-null - ``parent_id`` in a result that has parents available. - - ``ParentIdOutOfRange`` when ``record["parent_id"]`` - is not in ``[0, len(self.parents))``. - - Truth-allele consistency (requires ``refdata``): - - ``ParentTruthVCallMismatch`` / - ``ParentTruthDCallMismatch`` / - ``ParentTruthJCallMismatch`` — descendant's - ``truth_*_call`` (from ``expose_provenance=True``) - disagrees with the parent's projected truth allele. - - Provenance consistency (no ``refdata`` needed): - - ``ParentDInvertedMismatch`` — descendant's - ``d_inverted`` disagrees with the parent's. D inversion - is a pre-fork decision, so divergence indicates a - structural bug. - - ``ParentOriginalVCallMismatch`` — descendant's - ``original_v_call`` (receptor-revision provenance) - disagrees with the parent's. Same reasoning. - - **Without ``refdata``:** only the structural checks - (``ParentsMissing`` / ``ParentIdMissing`` / - ``ParentIdOutOfRange``) run. All value comparisons — - truth alleles, ``d_inverted``, ``original_v_call`` — - require projecting the parent ``Outcome`` to an AIRR - record, which today goes through the Rust projector and - needs ``refdata``. Slice 3 deliberately stays Python-only; - a lighter-weight refdata-free parent accessor for - ``d_inverted`` etc. is deferred until a Rust slice surfaces - one. - - **Skipped silently for fields not present on descendants:** - if ``expose_provenance`` was off, the descendants don't - carry ``truth_*_call`` and those checks are no-ops. - - **Non-clonal results return ok with ``family_count=0``** — - same safe no-op shape as :meth:`validate_families`. Slice - 3 deliberately does NOT raise "not clonal" here. - - **Not enforced yet** (deferred per audit §6, §14): - - - **Pre-SHM junction invariance.** The descendant's - ``junction`` AIRR field is post-SHM; the pre-SHM junction - lives only inside the parent's IR. A proper check would - require a parent-derived ``junction_pre_shm`` field on - either records or a future ``FamilyRecord`` projection - (Slice 4+). This validator does NOT compare - ``descendant.junction`` against any parent-derived - value today. - - **Mutation-distance distribution.** Comparing the - parent's assembled sequence to each descendant's - post-SHM sequence to verify SHM mass is plausible. - Requires projecting the parent's pool to a sequence - string — out of scope for this slice. - - **Plan-split pre-fork pass enumeration.** "Parent should - not carry descendant-only observation fields like PCR - / paired-end / quality errors" is pinned at the - contract-test level (the pre-fork plan's pass names) - rather than enforced at runtime here. - - **Not wired into ``validate_records=True``.** This is an - explicit deeper diagnostic surface. The - ``validate_records=True`` gate continues to run only the - per-record postcondition validator and the field-only - :meth:`validate_families`. Callers who want parent-aware - checks invoke this method explicitly. - - Returns a :class:`FamilyValidationReport`. Failure dicts - carry ``clone_id``, ``parent_id``, ``record_indices``, - ``issue_kind``, ``parent_value``, and ``child_values`` - (the latter two are ``None`` / ``[]`` for structural - failures that don't compare values). - """ - records = self._records - total = len(records) - - any_clonal = any( - "clone_id" in r and r["clone_id"] is not None for r in records - ) - any_parent_id = any( - "parent_id" in r and r["parent_id"] is not None for r in records - ) - - # Non-clonal: ok no-op (matches validate_families). - if not any_clonal and not any_parent_id: - return FamilyValidationReport( - count=total, - family_count=0, - members_per_family={}, - failures=[], - ) - - failures: List[Dict[str, Any]] = [] - parents = self._parents - - # Parents missing entirely — surface and bail (no comparable - # parent state to run further checks against). - if parents is None: - failures.append( - { - "clone_id": None, - "parent_id": None, - "record_indices": [ - i for i, r in enumerate(records) - if r.get("clone_id") is not None - or r.get("parent_id") is not None - ], - "issue_kind": "ParentsMissing", - "parent_value": None, - "child_values": [], - } - ) - # Compute aggregation by clone_id so the report still - # carries useful structure even though no comparison - # ran. - by_clone: Dict[Any, List[int]] = {} - for i, r in enumerate(records): - cid = r.get("clone_id") - if cid is None: - continue - by_clone.setdefault(cid, []).append(i) - return FamilyValidationReport( - count=total, - family_count=len(by_clone), - members_per_family={ - cid: len(idxs) for cid, idxs in by_clone.items() - }, - failures=failures, - ) - - # Parents are present. Group descendants by parent_id and - # surface structural failures (missing / out of range) per - # record. - n_parents = len(parents) - missing_parent_id: List[int] = [] - out_of_range: Dict[Any, List[int]] = {} - by_parent: Dict[int, List[int]] = {} - for i, r in enumerate(records): - pid = r.get("parent_id") - if pid is None: - # Records without a parent_id in a clonal result are - # structurally broken — surface them and skip per- - # parent invariant checks for those records. - if r.get("clone_id") is not None: - missing_parent_id.append(i) - continue - if not isinstance(pid, int) or isinstance(pid, bool): - out_of_range.setdefault(pid, []).append(i) - continue - if pid < 0 or pid >= n_parents: - out_of_range.setdefault(pid, []).append(i) - continue - by_parent.setdefault(pid, []).append(i) - - if missing_parent_id: - failures.append( - { - "clone_id": None, - "parent_id": None, - "record_indices": missing_parent_id, - "issue_kind": "ParentIdMissing", - "parent_value": None, - "child_values": [], - } - ) - for pid, idxs in out_of_range.items(): - failures.append( - { - "clone_id": None, - "parent_id": pid, - "record_indices": idxs, - "issue_kind": "ParentIdOutOfRange", - "parent_value": pid, - "child_values": [], - } - ) - - # Build parent projections lazily — only for parent ids that - # have descendants pointing at them. The Rust AIRR projector - # ``outcome_to_airr_record`` requires refdata; without it we - # can't materialize the parent's projected fields and - # therefore can't run any value-comparison check. Structural - # checks (missing / out-of-range parent_id) already ran - # above. Document this in the method's contract: refdata- - # less calls run structural checks only. - parent_projections: Dict[int, Dict[str, Any]] = {} - if by_parent and refdata is not None: - from ._airr_record import outcome_to_airr_record - - for pid in by_parent: - proj = outcome_to_airr_record( - parents[pid], - refdata, - sequence_id=f"parent{pid}", - ) - _inject_truth_columns(parents[pid], refdata, proj) - parent_projections[pid] = proj - - # The (field, issue_kind) pairs the parent-aware validator - # enforces when ``refdata`` is available. All of these are - # comparisons against parent-projected values; refdata is - # required to materialize them. - FIELD_CHECKS = [ - ("truth_v_call", "ParentTruthVCallMismatch"), - ("truth_d_call", "ParentTruthDCallMismatch"), - ("truth_j_call", "ParentTruthJCallMismatch"), - ("d_inverted", "ParentDInvertedMismatch"), - ("original_v_call", "ParentOriginalVCallMismatch"), - ] - - for pid, indices in by_parent.items(): - parent_proj = parent_projections.get(pid) - if parent_proj is None: - # Refdata was None — value comparisons skipped. - continue - members = [records[i] for i in indices] - # Use the first member's clone_id for the failure - # group; per-clone_id divergence within a parent_id - # group would have surfaced via the family validator - # already. - clone_id_for_group = members[0].get("clone_id", pid) - for field, kind in FIELD_CHECKS: - parent_val = parent_proj.get(field) - # Skip when parent has no opinion on the field - # (e.g. ``original_v_call`` is empty when receptor - # revision didn't run; comparing "" to "" would - # always pass but comparing "" to a real allele - # name would surface a real bug — so we only skip - # when the field is genuinely missing). - if parent_val is None: - continue - divergent_indices: List[int] = [] - divergent_values: List[Any] = [] - for idx, rec in zip(indices, members): - if field not in rec: - continue - child_val = rec[field] - if child_val is None: - continue - if child_val != parent_val: - divergent_indices.append(idx) - if child_val not in divergent_values: - divergent_values.append(child_val) - if divergent_indices: - failures.append( - { - "clone_id": clone_id_for_group, - "parent_id": pid, - "record_indices": divergent_indices, - "issue_kind": kind, - "parent_value": parent_val, - "child_values": sorted( - divergent_values, key=lambda v: str(v) - ), - } - ) - - return FamilyValidationReport( - count=total, - family_count=len(by_parent), - members_per_family={ - pid: len(idxs) for pid, idxs in by_parent.items() - }, - failures=failures, - ) - - # ── exports ───────────────────────────────────────────────────── - - def to_dataframe(self, *, airr_strict: bool = False): - """Return a :class:`pandas.DataFrame` with one row per record. - - ``airr_strict=True`` converts all 0-based half-open coord - ``*_start`` fields to the AIRR-spec 1-based-inclusive form - (``*_end`` fields are unchanged). Useful when handing the - DataFrame off to AIRR-strict downstream tooling. - - Raises ``ImportError`` if pandas isn't installed (pandas is - an optional extra: ``pip install GenAIRR[all]``). - """ - try: - import pandas as pd - except ImportError as exc: # pragma: no cover — depends on env - raise ImportError( - "to_dataframe() requires pandas. Install with " - "`pip install GenAIRR[all]` or `pip install pandas`." - ) from exc - - if not self._records: - return pd.DataFrame(columns=_DEFAULT_COLUMN_ORDER) - records = ( - [_to_airr_strict(r) for r in self._records] - if airr_strict - else self._records - ) - return pd.DataFrame(records, columns=self._column_order()) - - def to_tsv(self, path: str, *, airr_strict: bool = False) -> None: - """Write the records as AIRR-style TSV (tab-separated). The - header row uses :data:`_DEFAULT_COLUMN_ORDER`. - - ``airr_strict=True`` converts coord ``*_start`` fields to - 1-based-inclusive (AIRR spec). - """ - self._write_delimited(path, "\t", airr_strict=airr_strict) - - def to_csv(self, path: str, *, airr_strict: bool = False) -> None: - """Write the records as comma-separated values. Convenience - alongside :meth:`to_tsv` — most analysis tooling prefers TSV - for AIRR data. - - ``airr_strict=True`` converts coord ``*_start`` fields to - 1-based-inclusive (AIRR spec). - """ - self._write_delimited(path, ",", airr_strict=airr_strict) - - def to_fasta(self, path: str, *, prefix: str = "seq") -> None: - """Write the assembled sequences as FASTA. Each record gets - a header of the form ``">{prefix}{i}|v_call=...|j_call=..."``. - """ - with open(path, "w", encoding="utf-8") as fh: - for i, rec in enumerate(self._records): - seq = rec.get("sequence", "") - v_call = rec.get("v_call") or "" - j_call = rec.get("j_call") or "" - fh.write(f">{prefix}{i}|v_call={v_call}|j_call={j_call}\n") - fh.write(f"{seq}\n") - - def to_fastq( - self, - path: str, - *, - quality: str = "illumina", - prefix: str = "seq", - **quality_kwargs, - ) -> None: - """Write the assembled sequences as FASTQ. - - Each record produces: - - :: - - @{prefix}{i}|v_call=...|j_call=... - - + - - - Parameters: - path: output file path. - quality: name of the quality model — ``"illumina"`` - (smoothed trapezoid, default) or ``"constant"`` - (single Q value across the read). - prefix: per-read header prefix (default ``"seq"``). - **quality_kwargs: forwarded to the quality model - constructor. ``ConstantQualityModel`` accepts - ``q`` (default 30), ``low_q`` (default 10), - ``n_q`` (default 2). ``IlluminaQualityModel`` - accepts ``peak_q``, ``start_q``, ``end_q``, - ``ramp_len``, ``tail_len``, ``low_q``, ``n_q``. - - FASTQ uppercases the sequence bases — GenAIRR's lowercase - corruption-marker convention is preserved by routing - lowercase positions to ``low_q`` in the quality string, - the standard FASTQ way of conveying low-confidence bases. - """ - from ._qmodel import phred_to_ascii, resolve_quality_model - - model = resolve_quality_model(quality, **quality_kwargs) - with open(path, "w", encoding="utf-8") as fh: - for i, rec in enumerate(self._records): - seq = rec.get("sequence", "") - v_call = rec.get("v_call") or "" - j_call = rec.get("j_call") or "" - q_array = model.quality_array(seq) - if len(q_array) != len(seq): - raise RuntimeError( - f"quality model returned {len(q_array)} scores for " - f"{len(seq)}-base sequence" - ) - q_string = phred_to_ascii(q_array) - fh.write(f"@{prefix}{i}|v_call={v_call}|j_call={j_call}\n") - fh.write(f"{seq.upper()}\n") - fh.write("+\n") - fh.write(f"{q_string}\n") - - def to_paired_fastq( - self, - r1_path: str, - r2_path: str, - *, - quality: str = "illumina", - overwrite: bool = False, - **quality_kwargs, - ) -> None: - """Write the per-record paired-end reads as two FASTQ files. - - Each AIRR record contributes one R1 record (to ``r1_path``) - and one R2 record (to ``r2_path``): - - :: - - R1 file: R2 file: - @{sequence_id}/1 @{sequence_id}/2 - - + + - - - Read names use the AIRR record's own ``sequence_id`` with the - canonical Illumina-portable ``/1`` / ``/2`` suffix (older - convention but universally accepted by BWA / STAR / - samtools / Picard; the seven-field colon-separated full - Illumina header doesn't have a GenAIRR analogue — no flow - cell, no lane, no index — and is out of scope here per - `docs/fastq_export_design.md` §5). - - ``r2_sequence`` is **already** the reverse complement of - ``sequence[r2_start:r2_end]`` at projection time (the AIRR - validator's `PairedEndWindowMismatch { side: R2 }` enforces - the invariant); this writer outputs it verbatim. Applying a - second RC would corrupt the read. - - Parameters: - r1_path: output path for the R1 FASTQ file. - r2_path: output path for the R2 FASTQ file. - quality: name of the quality model — ``"illumina"`` - (smoothed trapezoid, default) or ``"constant"`` - (single Q value across the read). Same vocabulary - as :meth:`to_fastq`. The model is consulted - independently for R1 and R2; both reads get their - own quality string starting from position 0 (this - is the correct Illumina-style behaviour — each - read has its own ramp-up and tail). - overwrite: when ``False`` (default) the writer raises - ``FileExistsError`` if either output path already - exists. Set ``True`` to allow overwriting. - **quality_kwargs: forwarded to the quality model - constructor — same surface as :meth:`to_fastq`. - - Raises: - FileExistsError: when ``overwrite=False`` and either - output path already exists. - ValueError: when a record's ``read_layout`` is not - ``"paired_end"`` (the experiment hasn't run - ``.paired_end(...)``), or when ``r1_sequence`` / - ``r2_sequence`` is empty. - RuntimeError: when the quality model produces a - quality array whose length disagrees with the - read sequence (same shape as :meth:`to_fastq`). - - FASTQ uppercases the read bases — GenAIRR's lowercase - corruption-marker convention is preserved by routing - lowercase positions to the model's ``low_q`` parameter, - same as :meth:`to_fastq`. - """ - import os - - from ._qmodel import phred_to_ascii, resolve_quality_model - - # ── 1. Output-path overwrite guard. ────────────────── - if not overwrite: - for label, path in (("r1_path", r1_path), ("r2_path", r2_path)): - if os.path.exists(path): - raise FileExistsError( - f"to_paired_fastq: {label}={path!r} already exists " - "and overwrite=False. Pass overwrite=True to " - "replace it." - ) - - # ── 2. Resolve the quality model once. ─────────────── - model = resolve_quality_model(quality, **quality_kwargs) - - # ── 3. Walk records, validating layout + writing. ──── - with open(r1_path, "w", encoding="utf-8") as r1_fh, open( - r2_path, "w", encoding="utf-8" - ) as r2_fh: - for i, rec in enumerate(self._records): - sequence_id = rec.get("sequence_id") or f"seq{i}" - # 3a. Read-layout guard. - read_layout = rec.get("read_layout", "") - if read_layout != "paired_end": - raise ValueError( - f"to_paired_fastq: record {i} " - f"(sequence_id={sequence_id!r}) has " - f"read_layout={read_layout!r} — paired-end FASTQ " - "export requires read_layout='paired_end'. Run " - "Experiment.paired_end(r1_length=…, " - "insert_size=…) on the experiment before " - "exporting." - ) - r1_seq = rec.get("r1_sequence") or "" - r2_seq = rec.get("r2_sequence") or "" - # 3b. Empty-window guard. Belt-and-suspenders — - # the projection kernel shouldn't produce empty - # windows on a paired-layout record, but a downstream - # consumer that hand-edited the record dict would - # otherwise silently emit a zero-length FASTQ read. - if not r1_seq: - raise ValueError( - f"to_paired_fastq: record {i} " - f"(sequence_id={sequence_id!r}) has empty " - "r1_sequence despite read_layout='paired_end'" - ) - if not r2_seq: - raise ValueError( - f"to_paired_fastq: record {i} " - f"(sequence_id={sequence_id!r}) has empty " - "r2_sequence despite read_layout='paired_end'" - ) - # 3c. Quality strings. Each read is scored - # independently — Illumina-style ramp shape resets - # per read, which is the correct biological model - # (R1 and R2 don't share a per-base quality - # profile). - q_r1 = model.quality_array(r1_seq) - if len(q_r1) != len(r1_seq): - raise RuntimeError( - f"to_paired_fastq: quality model returned " - f"{len(q_r1)} scores for {len(r1_seq)}-base R1 " - f"sequence (record {i})" - ) - q_r2 = model.quality_array(r2_seq) - if len(q_r2) != len(r2_seq): - raise RuntimeError( - f"to_paired_fastq: quality model returned " - f"{len(q_r2)} scores for {len(r2_seq)}-base R2 " - f"sequence (record {i})" - ) - # 3d. Write the two 4-line records. - r1_fh.write(f"@{sequence_id}/1\n") - r1_fh.write(f"{r1_seq.upper()}\n") - r1_fh.write("+\n") - r1_fh.write(f"{phred_to_ascii(q_r1)}\n") - r2_fh.write(f"@{sequence_id}/2\n") - r2_fh.write(f"{r2_seq.upper()}\n") - r2_fh.write("+\n") - r2_fh.write(f"{phred_to_ascii(q_r2)}\n") - - # ── internals ─────────────────────────────────────────────────── - - def _column_order(self) -> List[str]: - """Pick the column order for tabular exports. Starts with - the canonical default order; appends any extra columns the - records have (e.g. when callers add custom fields).""" - seen = set(_DEFAULT_COLUMN_ORDER) - extras: List[str] = [] - for rec in self._records: - for key in rec: - if key not in seen: - seen.add(key) - extras.append(key) - return _DEFAULT_COLUMN_ORDER + extras - - def _write_delimited( - self, path: str, delimiter: str, *, airr_strict: bool = False - ) -> None: - columns = self._column_order() - with open(path, "w", encoding="utf-8", newline="") as fh: - writer = csv.DictWriter( - fh, - fieldnames=columns, - delimiter=delimiter, - lineterminator="\n", - extrasaction="ignore", - ) - writer.writeheader() - for rec in self._records: - source = _to_airr_strict(rec) if airr_strict else rec - # Replace ``None`` with empty string so CSV columns - # don't carry literal ``"None"`` strings. - row = {k: ("" if v is None else v) for k, v in source.items()} - writer.writerow(row) - class SimulationResultWithLineages(SimulationResult): """A :class:`SimulationResult` that also carries per-clone lineage trees. diff --git a/src/GenAIRR/utilities/__init__.py b/src/GenAIRR/utilities/__init__.py index c76070d..73b2582 100644 --- a/src/GenAIRR/utilities/__init__.py +++ b/src/GenAIRR/utilities/__init__.py @@ -5,4 +5,5 @@ from .misc import parse_mutation from .misc import parse_fasta from .AlleleNComparer import AlleleNComparer +from .visualize import visualize_sequence diff --git a/src/GenAIRR/utilities/_visualize/__init__.py b/src/GenAIRR/utilities/_visualize/__init__.py new file mode 100644 index 0000000..7b35ec4 --- /dev/null +++ b/src/GenAIRR/utilities/_visualize/__init__.py @@ -0,0 +1,10 @@ +"""Internal implementation of the HTML "exploding view" renderer. + +Split into cohesive submodules: parse / styles / components / alignment / +render. The public entry point is :func:`visualize_sequence`, re-exported +here and by :mod:`GenAIRR.utilities.visualize`. +""" + +from .render import visualize_sequence + +__all__ = ["visualize_sequence"] diff --git a/src/GenAIRR/utilities/_visualize/alignment.py b/src/GenAIRR/utilities/_visualize/alignment.py new file mode 100644 index 0000000..61bf5ba --- /dev/null +++ b/src/GenAIRR/utilities/_visualize/alignment.py @@ -0,0 +1,58 @@ +"""Germline vs. sequence alignment visualization.""" + +from __future__ import annotations + +from .parse import _esc + + +def _build_germline_alignment(sequence: str, germline: str) -> str: + """Build a germline vs sequence alignment visualization.""" + # Show first 200 positions to keep it readable + show_len = min(len(sequence), len(germline), 400) + if show_len == 0: + return "" + + rows = [] + for row_start in range(0, show_len, 80): + chunk_seq = sequence[row_start : row_start + 80] + chunk_germ = germline[row_start : row_start + 80] + + seq_chars = [] + germ_chars = [] + match_chars = [] + for s, g in zip(chunk_seq, chunk_germ): + if g == "." or g == "N": + germ_chars.append(f'{_esc(g)}') + seq_chars.append(f'{_esc(s)}') + match_chars.append(f' ') + elif s == g: + germ_chars.append(f'{_esc(g)}') + seq_chars.append(f'{_esc(s)}') + match_chars.append(f'|') + else: + germ_chars.append(f'{_esc(g)}') + seq_chars.append(f'{_esc(s)}') + match_chars.append(f'*') + + pos_label = str(row_start + 1) + rows.append( + f'
' + f'{pos_label}' + f'
{"".join(germ_chars)}
' + f'
' + f'
' + f'' + f'
{"".join(match_chars)}
' + f'
' + f'
' + f'' + f'
{"".join(seq_chars)}
' + f'
' + ) + + return ( + '
' + 'Germline Alignment' + + "".join(rows) + + "
" + ) diff --git a/src/GenAIRR/utilities/_visualize/components.py b/src/GenAIRR/utilities/_visualize/components.py new file mode 100644 index 0000000..6fbbd73 --- /dev/null +++ b/src/GenAIRR/utilities/_visualize/components.py @@ -0,0 +1,293 @@ +"""HTML component builders for the exploded-view render.""" + +from __future__ import annotations + +from typing import Dict, Optional + +from .parse import _bool, _esc, _int, _str +from .styles import PURPLE, RED + + +def _build_segment_bar_html(segments: list, seq_len: int, mutations: Dict[int, str]) -> str: + """Build the color-coded assembled sequence bar.""" + parts = [] + for seg in segments: + w = max(((seg["end"] - seg["start"]) / seq_len) * 100, 1.5) + label = seg["label"] if w > 4 else "" + start_label = seg["start"] + 1 + end_label = seg["end"] + parts.append( + f'
' + f'{_esc(label)}' + f'{start_label}' + + (f'{end_label}' if w > 8 else "") + + "
" + ) + + # Mutation dots + dots = [] + for pos in mutations: + left_pct = (pos / seq_len) * 100 + dots.append(f'
') + + return ( + '
' + + "".join(parts) + + "".join(dots) + + "
" + ) + + +def _build_junction_bracket(junction_start: int, junction_end: int, junction_aa: str, seq_len: int) -> str: + if seq_len == 0 or junction_end <= junction_start: + return "" + left_pct = (junction_start / seq_len) * 100 + width_pct = ((junction_end - junction_start) / seq_len) * 100 + return ( + '
' + f'
' + '' + f'' + f'' + f'' + "" + f'CDR3: {_esc(junction_aa)}' + "
" + "
" + ) + + +def _sparkline_svg(start: int, end: int, mutations: Dict[int, str], height: int = 24) -> str: + length = end - start + if length <= 0: + return "" + bins = min(length, 60) + bin_size = max(1, length // bins) + counts = [] + for i in range(bins): + b_start = start + i * bin_size + b_end = b_start + bin_size + c = sum(1 for p in mutations if b_start <= p < b_end) + counts.append(c) + mx = max(max(counts), 1) + rects = [] + for i, c in enumerate(counts): + fill = RED if c > 0 else "#555" + opacity = 0.85 if c > 0 else 0.2 + h = c if c > 0 else 0.15 + rects.append( + f'' + ) + return ( + f'' + + "".join(rects) + + "" + ) + + +def _pnp_bar(prefix: str, n_region: str, suffix: str, color: str) -> str: + total = len(prefix) + len(n_region) + len(suffix) + if total == 0: + return "" + parts_data = [ + ("P", len(prefix), 0.6), + ("N", len(n_region), 0.25), + ("P", len(suffix), 0.6), + ] + parts = [] + for label, length, opacity in parts_data: + if length > 0: + w = max((length / total) * 100, 3) + parts.append( + f'
' + f'{label}
' + ) + return '
' + "".join(parts) + "
" + + +def _nt_strip(sequence: str, start_pos: int, color: str, mutations: Dict[int, str]) -> str: + """Nucleotide detail with mutation highlighting.""" + rows = [] + for row_idx in range(0, len(sequence), 80): + chunk = sequence[row_idx : row_idx + 80] + chars = [] + for ci, ch in enumerate(chunk): + abs_pos = start_pos + row_idx + ci + is_mut = abs_pos in mutations + cls = "nt-char mut" if is_mut else "nt-char" + dot = '' if is_mut else "" + chars.append(f'{dot}{_esc(ch)}') + line_start = start_pos + row_idx + 1 + line_end = min(start_pos + row_idx + 80, start_pos + len(sequence)) + rows.append( + f'
' + f'{line_start}' + f'
{"".join(chars)}
' + f'{line_end}' + f"
" + ) + bg = f"color-mix(in srgb, {color} 5%, #1a1a2e)" + return f'
{"".join(rows)}
' + + +def _metric(label: str, value: str, color: Optional[str] = None) -> str: + style = f' style="color:{color}"' if color else "" + return ( + f'
' + f'{_esc(label)}' + f'{_esc(value)}' + f"
" + ) + + +def _trim_indicator(label: str, bases: int) -> str: + if not bases: + return "" + return ( + f'
' + f'✂ ' + f'{_esc(label)}: {bases} bp trimmed' + f"
" + ) + + +def _segment_panel(seg: dict, rec: dict, mutations: Dict[int, str], sequence: str) -> str: + """Build a single exploded segment panel.""" + seg_start = seg["start"] + seg_end = seg["end"] + seg_len = seg_end - seg_start + sub_seq = sequence[seg_start:seg_end] + seg_muts = {p: v for p, v in mutations.items() if seg_start <= p < seg_end} + mut_count = len(seg_muts) + seg_id = seg["id"] + color = seg["color"] + + header = ( + f'
' + f'{_esc(seg["label"])}' + f'{_esc(seg["full_label"])}' + f'{seg_start + 1}..{seg_end} ({seg_len} nt)' + f"
" + ) + + body_parts = [] + + if seg_id == "V": + body_parts.append( + f'
' + + _metric("LENGTH", f"{seg_len} bp") + + _metric("MUTATIONS", str(mut_count), RED if mut_count > 0 else None) + + _metric("3' TRIM", f"{_int(rec, 'v_trim_3')} bp") + + "
" + ) + if seg_len > 0: + body_parts.append( + '
' + 'Mutation Density' + + _sparkline_svg(seg_start, seg_end, mutations) + + "
" + ) + body_parts.append(_trim_indicator("V-gene 3'", _int(rec, "v_trim_3"))) + + elif seg_id == "D": + d_inv = _bool(rec, "d_inverted") + body_parts.append( + f'
' + + _metric("LENGTH", f"{seg_len} bp") + + _metric("5' TRIM", f"{_int(rec, 'd_trim_5')} bp") + + _metric("3' TRIM", f"{_int(rec, 'd_trim_3')} bp") + + _metric("INVERTED", "YES" if d_inv else "NO", RED if d_inv else None) + + "
" + ) + if d_inv: + body_parts.append( + '
' + '⇅ ' + "D-gene inverted (reverse complement used)
" + ) + body_parts.append(_trim_indicator("5'", _int(rec, "d_trim_5"))) + body_parts.append(_trim_indicator("3'", _int(rec, "d_trim_3"))) + + elif seg_id == "J": + body_parts.append( + f'
' + + _metric("LENGTH", f"{seg_len} bp") + + _metric("MUTATIONS", str(mut_count), RED if mut_count > 0 else None) + + _metric("5' TRIM", f"{_int(rec, 'j_trim_5')} bp") + + "
" + ) + if seg_len > 0: + body_parts.append( + '
' + 'Mutation Density' + + _sparkline_svg(seg_start, seg_end, mutations) + + "
" + ) + body_parts.append(_trim_indicator("J-gene 5'", _int(rec, "j_trim_5"))) + + elif seg_id == "NP1": + p_pre = _str(rec, "np1_p_prefix") + n_reg = _str(rec, "np1_n_region") + p_suf = _str(rec, "np1_p_suffix") + body_parts.append( + f'
' + + _metric("P-PREFIX", f"{len(p_pre)} bp") + + _metric("N-ADDITION", f"{len(n_reg)} bp") + + _metric("P-SUFFIX", f"{len(p_suf)} bp") + + "
" + ) + body_parts.append( + 'P | N | P Composition' + + _pnp_bar(p_pre, n_reg, p_suf, color) + ) + body_parts.append( + '
' + + f'{_esc(p_pre)}' + + f'{_esc(n_reg)}' + + f'{_esc(p_suf)}' + + "
" + ) + + elif seg_id == "NP2": + p_pre = _str(rec, "np2_p_prefix") + n_reg = _str(rec, "np2_n_region") + p_suf = _str(rec, "np2_p_suffix") + body_parts.append( + f'
' + + _metric("P-PREFIX", f"{len(p_pre)} bp") + + _metric("N-ADDITION", f"{len(n_reg)} bp") + + _metric("P-SUFFIX", f"{len(p_suf)} bp") + + "
" + ) + body_parts.append( + 'P | N | P Composition' + + _pnp_bar(p_pre, n_reg, p_suf, color) + ) + body_parts.append( + '
' + + f'{_esc(p_pre)}' + + f'{_esc(n_reg)}' + + f'{_esc(p_suf)}' + + "
" + ) + + # Nucleotide detail (always shown in static view) + if seg_len > 0 and seg_id in ("V", "D", "J"): + mut_note = f' ({mut_count} mutations highlighted)' if mut_count > 0 else "" + body_parts.append( + f'
' + f'Nucleotide Sequence{mut_note}' + + _nt_strip(sub_seq, seg_start, color, seg_muts) + + "
" + ) + + return ( + f'
' + + header + + '
' + + "".join(body_parts) + + "
" + ) diff --git a/src/GenAIRR/utilities/_visualize/parse.py b/src/GenAIRR/utilities/_visualize/parse.py new file mode 100644 index 0000000..09a7a69 --- /dev/null +++ b/src/GenAIRR/utilities/_visualize/parse.py @@ -0,0 +1,77 @@ +"""Field extraction / parsing helpers and HTML escaping.""" + +from __future__ import annotations + +import html +import json +from typing import Any, Dict + + +def _int(rec: dict, key: str, default: int = 0) -> int: + v = rec.get(key) + if v is None: + return default + try: + return int(v) + except (ValueError, TypeError): + return default + + +def _str(rec: dict, key: str, default: str = "") -> str: + v = rec.get(key) + return str(v) if v is not None else default + + +def _bool(rec: dict, key: str) -> bool: + v = rec.get(key) + if isinstance(v, bool): + return v + if isinstance(v, str): + return v.lower() in ("true", "1", "yes", "t") + return bool(v) if v is not None else False + + +def _float(rec: dict, key: str, default: float = 0.0) -> float: + v = rec.get(key) + if v is None: + return default + try: + return float(v) + except (ValueError, TypeError): + return default + + +def _parse_mutations(raw: Any) -> Dict[int, str]: + """Parse the mutations field (string like '42:A>G;100:T>C' or dict).""" + if not raw: + return {} + if isinstance(raw, dict): + return {int(k): str(v) for k, v in raw.items()} + if isinstance(raw, str): + # Try JSON first + try: + parsed = json.loads(raw) + if isinstance(parsed, dict): + return {int(k): str(v) for k, v in parsed.items()} + except (json.JSONDecodeError, ValueError): + pass + # Try comma or semicolon-separated "pos:from>to" format + out = {} + # Split on comma or semicolon + sep = "," if "," in raw else ";" + for part in raw.split(sep): + part = part.strip() + if not part: + continue + if ":" in part: + pos_str, desc = part.split(":", 1) + try: + out[int(pos_str)] = desc + except ValueError: + pass + return out + return {} + + +def _esc(s: str) -> str: + return html.escape(s, quote=True) diff --git a/src/GenAIRR/utilities/_visualize/render.py b/src/GenAIRR/utilities/_visualize/render.py new file mode 100644 index 0000000..d48436c --- /dev/null +++ b/src/GenAIRR/utilities/_visualize/render.py @@ -0,0 +1,194 @@ +"""Public renderer: assemble a full standalone HTML dissection page.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Optional, Union + +from .alignment import _build_germline_alignment +from .components import ( + _build_junction_bracket, + _build_segment_bar_html, + _segment_panel, +) +from .parse import _bool, _esc, _float, _int, _parse_mutations, _str +from .styles import CSS, RED, SEGMENT_COLORS + + +def visualize_sequence( + record: Dict[str, Any], + path: Union[str, Path], + title: Optional[str] = None, +) -> Path: + """ + Generate a standalone HTML file with an "exploding view" dissection + of a simulated AIRR sequence record. + + Parameters + ---------- + record : dict + A single AIRR record dict (one element from SimulationResult). + path : str or Path + Output file path for the HTML file. + title : str, optional + Custom title for the page. Defaults to "Sequence Dissection". + + Returns + ------- + Path + The path to the generated file. + """ + path = Path(path) + rec = record + + # Extract fields + sequence = _str(rec, "sequence") + seq_len = len(sequence) + germline = _str(rec, "germline_alignment") + v_call = _str(rec, "v_call") + d_call = _str(rec, "d_call") + j_call = _str(rec, "j_call") + junction_aa = _str(rec, "junction_aa") + productive = _bool(rec, "productive") + mutation_rate = _float(rec, "mutation_rate") + mutations = _parse_mutations(rec.get("mutations")) + total_mutations = len(mutations) + cdr3_length = len(junction_aa) + + v_seq_start = _int(rec, "v_sequence_start") + v_seq_end = _int(rec, "v_sequence_end") + d_seq_start = _int(rec, "d_sequence_start") + d_seq_end = _int(rec, "d_sequence_end") + j_seq_start = _int(rec, "j_sequence_start") + j_seq_end = _int(rec, "j_sequence_end") + junction_start = _int(rec, "junction_start", _int(rec, "junction_sequence_start")) + junction_end = _int(rec, "junction_end", _int(rec, "junction_sequence_end")) + + corruption_5 = _int(rec, "corruption_5prime") + corruption_3 = _int(rec, "corruption_3prime") + + seq_id = _str(rec, "sequence_id", "sequence") + page_title = title or "Sequence Dissection" + + # Build segments + segments = [ + {"id": "V", "label": "V", "full_label": v_call, "color": SEGMENT_COLORS["V"], "start": v_seq_start, "end": v_seq_end}, + {"id": "NP1", "label": "NP1", "full_label": "N/P Region 1", "color": SEGMENT_COLORS["NP1"], "start": v_seq_end, "end": d_seq_start}, + {"id": "D", "label": "D", "full_label": d_call, "color": SEGMENT_COLORS["D"], "start": d_seq_start, "end": d_seq_end}, + {"id": "NP2", "label": "NP2", "full_label": "N/P Region 2", "color": SEGMENT_COLORS["NP2"], "start": d_seq_end, "end": j_seq_start}, + {"id": "J", "label": "J", "full_label": j_call, "color": SEGMENT_COLORS["J"], "start": j_seq_start, "end": j_seq_end}, + ] + + # ── Build HTML sections ────────────────────────────────── + + # Top bar + badge_cls = "badge badge-pos" if productive else "badge badge-neg" + badge_text = "Productive" if productive else "Non-productive" + top_bar = ( + f'
' + f'🧬' + f'

{_esc(page_title)}

' + f'{badge_text}' + f'{_esc(seq_id)}' + f"
" + ) + + # Summary grid + mut_color = RED if total_mutations > 0 else None + prod_color = RED if not productive else None + summary_items = [ + ("V-GENE", v_call, SEGMENT_COLORS["V"], False), + ("D-GENE", d_call, SEGMENT_COLORS["D"], False), + ("J-GENE", j_call, SEGMENT_COLORS["J"], False), + ("JUNCTION AA", junction_aa, None, True), + ("TOTAL LENGTH", f"{seq_len} nt", None, False), + ("MUTATIONS", f"{total_mutations} ({mutation_rate * 100:.1f}%)", mut_color, False), + ("CDR3 LENGTH", f"{cdr3_length} aa", None, False), + ("PRODUCTIVE", "Yes" if productive else "No", prod_color, False), + ] + summary_cells = [] + for label, val, color, mono in summary_items: + style = f' style="color:{color}"' if color else "" + mono_cls = " mono" if mono else "" + summary_cells.append( + f'
' + f'{_esc(label)}' + f'{_esc(str(val))}' + f"
" + ) + summary_grid = '
' + "".join(summary_cells) + "
" + + # Corruption warnings + corruption_html = "" + if corruption_5 or corruption_3: + parts = [] + if corruption_5: + parts.append(f'
⚠ 5\' Corruption: +{corruption_5} nt added
') + if corruption_3: + parts.append(f'
⚠ 3\' Corruption: −{corruption_3} nt removed
') + corruption_html = '
' + "".join(parts) + "
" + + # Assembled bar + junction + bar_border_left = f"border-left:3px dashed {RED};" if corruption_5 else "" + bar_border_right = f"border-right:3px dashed {RED};" if corruption_3 else "" + if bar_border_left or bar_border_right: + bar_html = _build_segment_bar_html(segments, seq_len, mutations) + bar_html = bar_html.replace( + 'class="assembled-bar"', + f'class="assembled-bar" style="{bar_border_left}{bar_border_right}"', + ) + else: + bar_html = _build_segment_bar_html(segments, seq_len, mutations) + + junction_html = _build_junction_bracket(junction_start, junction_end, junction_aa, seq_len) + + # Exploded segment panels + panels = [] + for seg in segments: + panels.append(_segment_panel(seg, rec, mutations, sequence)) + + # Germline alignment (if available) + germline_html = "" + if germline and len(germline) >= len(sequence): + germline_html = _build_germline_alignment(sequence, germline) + + # Footer + footer = '
Generated by GenAIRR — Synthetic AIRR Sequence Simulator
' + + # ── Assemble HTML ──────────────────────────────────────── + + html_content = f""" + + + + +{_esc(page_title)} — {_esc(seq_id)} + + + +
+{top_bar} +{summary_grid} +{corruption_html} + +
+Assembled Sequence +{bar_html} +{junction_html} +
+ +
+⚡ Exploded Segments +
+{"".join(panels)} +
+
+ +{germline_html} +{footer} +
+ +""" + + path.write_text(html_content, encoding="utf-8") + return path diff --git a/src/GenAIRR/utilities/_visualize/styles.py b/src/GenAIRR/utilities/_visualize/styles.py new file mode 100644 index 0000000..c0f9771 --- /dev/null +++ b/src/GenAIRR/utilities/_visualize/styles.py @@ -0,0 +1,238 @@ +"""Segment colors and the CSS stylesheet for the HTML render.""" + +from __future__ import annotations + +# ── Segment colors (matching the website DissectionBay) ────── + +SEGMENT_COLORS = { + "V": "#4A90D9", + "NP1": "#5DC48C", + "D": "#E8685A", + "NP2": "#5DC48C", + "J": "#E8A838", +} + +RED = "#DC2626" +PURPLE = "#9B6FC4" + + +# ── CSS ────────────────────────────────────────────────────── + +CSS = """ +:root { + --bg: #0f1019; + --bg2: #181825; + --bg3: #1e1e32; + --fg: #e0e0f0; + --fg2: #8888aa; + --border: #2a2a40; + --accent: #6366f1; + --red: #DC2626; + --purple: #9B6FC4; +} +* { box-sizing: border-box; margin: 0; padding: 0; } +body { + font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; + background: var(--bg); + color: var(--fg); + line-height: 1.5; + padding: 2rem; +} +.container { max-width: 1200px; margin: 0 auto; } +h1 { font-size: 1.4rem; font-weight: 600; margin-bottom: 0.25rem; } +h2 { font-size: 1.1rem; font-weight: 600; margin-bottom: 0.75rem; color: var(--fg2); } +.subtitle { color: var(--fg2); font-size: 0.85rem; margin-bottom: 1.5rem; } + +/* Top bar */ +.top-bar { + display: flex; align-items: center; gap: 0.75rem; + padding: 1rem 1.25rem; + background: var(--bg2); border: 1px solid var(--border); + border-radius: 10px; margin-bottom: 1.25rem; +} +.top-bar .dna-icon { color: var(--accent); font-size: 1.2rem; } +.badge { + padding: 0.2rem 0.6rem; border-radius: 999px; + font-size: 0.7rem; font-weight: 600; text-transform: uppercase; + letter-spacing: 0.05em; +} +.badge-pos { background: rgba(34,197,94,0.15); color: #22c55e; } +.badge-neg { background: rgba(220,38,38,0.15); color: #DC2626; } +.seq-id { color: var(--fg2); font-size: 0.8rem; margin-left: auto; font-family: monospace; } + +/* Summary grid */ +.summary-grid { + display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr)); + gap: 0.5rem; margin-bottom: 1.25rem; +} +.summary-cell { + background: var(--bg2); border: 1px solid var(--border); + border-radius: 8px; padding: 0.6rem 0.8rem; +} +.summary-cell .metric-label { + display: block; font-size: 0.6rem; text-transform: uppercase; + letter-spacing: 0.08em; color: var(--fg2); margin-bottom: 0.15rem; +} +.summary-cell .metric-value { + font-size: 0.9rem; font-weight: 600; +} +.summary-cell .metric-value.mono { font-family: monospace; font-size: 0.75rem; word-break: break-all; } + +/* Corruption */ +.corruption-row { + display: flex; gap: 0.75rem; margin-bottom: 1rem; +} +.corruption-badge { + display: inline-flex; align-items: center; gap: 0.4rem; + background: rgba(220,38,38,0.08); border: 1px solid rgba(220,38,38,0.25); + border-radius: 6px; padding: 0.35rem 0.75rem; + font-size: 0.75rem; color: var(--red); +} + +/* Assembled bar */ +.section-label { font-size: 0.7rem; text-transform: uppercase; letter-spacing: 0.08em; color: var(--fg2); margin-bottom: 0.4rem; } +.assembled-bar { + display: flex; height: 38px; border-radius: 6px; overflow: hidden; + position: relative; margin-bottom: 0; + border: 1px solid var(--border); +} +.seg-bar-part { + position: relative; display: flex; align-items: center; justify-content: center; + cursor: default; transition: opacity 0.2s; + min-width: 0; +} +.seg-bar-label { + font-size: 0.7rem; font-weight: 700; color: #fff; text-shadow: 0 1px 2px rgba(0,0,0,0.4); + pointer-events: none; +} +.seg-bar-pos { + position: absolute; bottom: 2px; font-size: 0.55rem; color: rgba(255,255,255,0.7); + pointer-events: none; +} +.seg-bar-pos.left { left: 3px; } +.seg-bar-pos.right { right: 3px; } +.mut-dot { + position: absolute; top: -2px; width: 4px; height: 4px; + background: var(--red); border-radius: 50%; + transform: translateX(-50%); + pointer-events: none; +} + +/* Junction bracket */ +.junction-bracket { position: relative; height: 28px; margin-bottom: 1.5rem; } +.junction-inner { position: absolute; text-align: center; } +.junction-aa { + display: block; font-size: 0.65rem; color: var(--purple); + font-family: monospace; margin-top: 1px; white-space: nowrap; + overflow: hidden; text-overflow: ellipsis; +} + +/* Segment panels */ +.segments-grid { + display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); + gap: 0.75rem; margin-top: 1rem; +} +.seg-panel { + background: var(--bg2); border: 1px solid var(--border); + border-radius: 8px; overflow: hidden; +} +.seg-panel-header { + display: flex; align-items: center; gap: 0.5rem; + padding: 0.6rem 0.8rem; border-bottom: 1px solid var(--border); +} +.seg-chip { + display: inline-block; padding: 0.15rem 0.5rem; border-radius: 4px; + font-size: 0.65rem; font-weight: 700; color: #fff; +} +.seg-gene { font-size: 0.75rem; font-weight: 600; } +.seg-pos { font-size: 0.7rem; color: var(--fg2); margin-left: auto; } +.seg-panel-body { padding: 0.75rem 0.8rem; display: flex; flex-direction: column; gap: 0.6rem; } + +/* Metrics inside panels */ +.metric-grid { display: grid; gap: 0.4rem; } +.metric-grid.g3 { grid-template-columns: repeat(3, 1fr); } +.metric-grid.g4 { grid-template-columns: repeat(4, 1fr); } +.metric { + background: var(--bg3); border-radius: 5px; padding: 0.35rem 0.5rem; + text-align: center; +} +.metric-label { + display: block; font-size: 0.55rem; text-transform: uppercase; + letter-spacing: 0.06em; color: var(--fg2); +} +.metric-value { font-size: 0.8rem; font-weight: 600; } + +/* Trim */ +.trim-info { + font-size: 0.7rem; color: var(--fg2); + display: flex; align-items: center; gap: 0.3rem; +} +.trim-icon { font-size: 0.85rem; } + +/* Inversion */ +.inv-badge { + display: flex; align-items: center; gap: 0.4rem; + background: rgba(220,38,38,0.08); border: 1px solid rgba(220,38,38,0.2); + border-radius: 5px; padding: 0.3rem 0.6rem; + font-size: 0.7rem; color: var(--red); +} +.inv-icon { font-size: 1rem; } + +/* PNP bar */ +.pnp-bar { + display: flex; height: 20px; border-radius: 4px; overflow: hidden; + border: 1px solid var(--border); +} +.pnp-part { + display: flex; align-items: center; justify-content: center; + min-width: 0; +} +.pnp-label { font-size: 0.6rem; font-weight: 700; color: #fff; } +.np-seq { + font-family: monospace; font-size: 0.7rem; word-break: break-all; + padding: 0.3rem; background: var(--bg3); border-radius: 4px; +} + +/* Nucleotide strip */ +.nt-section { margin-top: 0.25rem; } +.nt-strip { + border-radius: 6px; padding: 0.5rem; overflow-x: auto; + font-family: 'Fira Code', 'JetBrains Mono', 'Cascadia Code', monospace; + font-size: 0.65rem; line-height: 1.6; +} +.nt-row { display: flex; align-items: center; gap: 0.5rem; } +.nt-linenum { color: var(--fg2); min-width: 3ch; text-align: right; font-size: 0.6rem; user-select: none; } +.nt-linenum.end { text-align: left; } +.nt-chars { display: flex; flex-wrap: wrap; } +.nt-char { + position: relative; display: inline-block; width: 0.8em; text-align: center; +} +.nt-char.mut { color: var(--red); font-weight: 700; } +.mut-marker { + position: absolute; top: -3px; left: 50%; transform: translateX(-50%); + width: 3px; height: 3px; border-radius: 50%; background: var(--red); +} +.micro-label { + display: block; font-size: 0.6rem; text-transform: uppercase; + letter-spacing: 0.06em; color: var(--fg2); margin-bottom: 0.3rem; +} +.spark-section { /* wrapper */ } + +/* Germline alignment */ +.germline-section { margin-top: 1.25rem; } +.germline-row { + display: flex; font-family: monospace; font-size: 0.65rem; line-height: 1.6; + gap: 0.5rem; align-items: center; +} +.germline-label { min-width: 5ch; color: var(--fg2); font-size: 0.6rem; text-align: right; } +.germline-chars { display: flex; flex-wrap: wrap; } +.germline-match { color: var(--fg2); opacity: 0.4; } +.germline-mismatch { color: var(--red); font-weight: 700; } +.germline-gap { color: var(--fg2); opacity: 0.2; } + +/* Footer */ +.footer { + margin-top: 2rem; padding-top: 1rem; border-top: 1px solid var(--border); + font-size: 0.65rem; color: var(--fg2); text-align: center; +} +""" diff --git a/src/GenAIRR/utilities/mcp_helpers.py b/src/GenAIRR/utilities/mcp_helpers.py deleted file mode 100644 index 6c6e449..0000000 --- a/src/GenAIRR/utilities/mcp_helpers.py +++ /dev/null @@ -1,946 +0,0 @@ -""" -Helper functions for the GenAIRR MCP server. - -Validation, scoring, analysis, and formatting logic used by the MCP -tool definitions in ``mcp_server.py``. Every function is pure — takes -data in, returns JSON-serialisable data out. -""" - -from __future__ import annotations - -import math -import re -from collections import defaultdict -from typing import Any, Dict, List, Optional, Tuple - -from ..alleles.allele import VAllele, DAllele, JAllele, CAllele - - -# ═══════════════════════════════════════════════════════════════════ -# Allele lookup helpers -# ═══════════════════════════════════════════════════════════════════ - -def flatten_alleles(allele_dict) -> Dict[str, Any]: - """Flatten ``{gene: [Allele, ...]}`` to ``{name: Allele}``.""" - out = {} - if allele_dict is None: - return out - for alleles in allele_dict.values(): - for allele in alleles: - out[allele.name] = allele - return out - - -def find_allele(dc, allele_name: str): - """Search all segments of a DataConfig for an allele by name.""" - for attr in ("v_alleles", "d_alleles", "j_alleles", "c_alleles"): - ad = getattr(dc, attr, None) - if ad is None: - continue - for alleles in ad.values(): - for a in alleles: - if a.name == allele_name: - return a - return None - - -def _segment_label(allele) -> str: - if isinstance(allele, VAllele): - return "V" - if isinstance(allele, DAllele): - return "D" - if isinstance(allele, JAllele): - return "J" - if isinstance(allele, CAllele): - return "C" - return "?" - - -def format_allele_info(allele) -> Dict[str, Any]: - """Return a compact JSON-safe dict for a single Allele.""" - seq = allele.ungapped_seq - preview = seq[:80] + ("..." if len(seq) > 80 else "") - info: Dict[str, Any] = { - "name": allele.name, - "segment": _segment_label(allele), - "sequence_preview": preview, - "ungapped_length": allele.ungapped_len, - "family": allele.family, - "gene": allele.gene, - } - anchor = getattr(allele, "anchor", None) - if anchor is not None: - info["anchor"] = anchor - return info - - -# ═══════════════════════════════════════════════════════════════════ -# Mutation parsing -# ═══════════════════════════════════════════════════════════════════ - -def parse_mutations(mutation_str: str) -> Dict[int, Tuple[str, str]]: - """Parse ``"pos:X>Y,pos:X>Y,..."`` into ``{pos: (from_base, to_base)}``.""" - result: Dict[int, Tuple[str, str]] = {} - if not mutation_str: - return result - for entry in mutation_str.split(","): - entry = entry.strip() - if not entry or ":" not in entry: - continue - pos_str, change = entry.split(":", 1) - try: - pos = int(pos_str) - except ValueError: - continue - if ">" not in change: - continue - parts = change.split(">", 1) - if len(parts) == 2: - result[pos] = (parts[0], parts[1]) - return result - - -# ═══════════════════════════════════════════════════════════════════ -# Allele scoring -# ═══════════════════════════════════════════════════════════════════ - -def count_matches(seq_region: str, allele_region: str) -> Tuple[int, int]: - """Count N-aware base matches. Returns ``(n_matches, n_compared)``.""" - n_match = 0 - n_compared = 0 - for s, a in zip(seq_region, allele_region): - su, au = s.upper(), a.upper() - if su == "N" or au == "N": - continue - n_compared += 1 - if su == au: - n_match += 1 - return n_match, n_compared - - -def score_allele_against_sequence( - seq: str, - allele_seq: str, - seq_start: int, - seq_end: int, - germ_start: int, - germ_end: int, -) -> Optional[Tuple[int, int]]: - """Score how well an allele matches the sequence region.""" - if germ_start < 0 or germ_end <= germ_start: - return None - allele_len = len(allele_seq) - region_len = min(seq_end - seq_start, germ_end - germ_start) - if region_len <= 0: - return None - - n_match = 0 - n_compared = 0 - for i in range(region_len): - s = seq[seq_start + i].upper() if seq_start + i < len(seq) else "N" - if germ_start + i < allele_len: - a = allele_seq[germ_start + i].upper() - else: - n_compared += 1 - continue - if s == "N" or a == "N": - continue - n_compared += 1 - if s == a: - n_match += 1 - return n_match, n_compared - - -def score_all_alleles(dc, rec: Dict[str, Any], segment: str) -> List[Dict[str, Any]]: - """Score every allele of *segment* against the record, return top results.""" - seg = segment.lower() - allele_dict = getattr(dc, f"{seg}_alleles", None) - if not allele_dict: - return [] - - flat = flatten_alleles(allele_dict) - seq = rec.get("sequence", "") - ss = rec.get(f"{seg}_sequence_start", 0) - se = rec.get(f"{seg}_sequence_end", 0) - gs = rec.get(f"{seg}_germline_start", 0) - ge = rec.get(f"{seg}_germline_end", 0) - - reported_calls = set() - call_str = rec.get(f"{seg}_call", "") or "" - for c in call_str.split(","): - c = c.strip() - if c: - reported_calls.add(c) - - results = [] - for name, allele in flat.items(): - score = score_allele_against_sequence( - seq, allele.ungapped_seq, ss, se, gs, ge, - ) - if score is None: - continue - matches, compared = score - fraction = matches / compared if compared > 0 else 0.0 - results.append({ - "name": name, - "matches": matches, - "compared": compared, - "fraction": round(fraction, 4), - "is_reported_call": name in reported_calls, - }) - - results.sort(key=lambda r: (-r["fraction"], -r["matches"])) - return results[:10] - - -# ═══════════════════════════════════════════════════════════════════ -# Record validation -# ═══════════════════════════════════════════════════════════════════ - -def validate_record(rec: Dict[str, Any]) -> Dict[str, Any]: - """Run comprehensive validation checks on a single AIRR record.""" - checks: List[Dict[str, Any]] = [] - seq = rec.get("sequence", "") - seq_len = len(seq) - germ = rec.get("germline_alignment", "") - - # 1. nucleotide_validity - bad_chars = set(seq.upper()) - {"A", "T", "C", "G", "N", ""} - checks.append({ - "name": "nucleotide_validity", - "passed": len(bad_chars) == 0, - "detail": f"invalid chars: {bad_chars}" if bad_chars else "ok", - }) - - # 2. coordinate_bounds - coord_fields = [ - "v_sequence_start", "v_sequence_end", - "d_sequence_start", "d_sequence_end", - "j_sequence_start", "j_sequence_end", - ] - oob = [] - for f in coord_fields: - val = rec.get(f, 0) - if isinstance(val, int) and (val < 0 or val > seq_len): - oob.append(f"{f}={val}") - checks.append({ - "name": "coordinate_bounds", - "passed": len(oob) == 0, - "detail": f"out of bounds: {', '.join(oob)}" if oob else "ok", - }) - - # 3. segment_ordering - vs = rec.get("v_sequence_start", 0) - ve = rec.get("v_sequence_end", 0) - ds = rec.get("d_sequence_start", 0) - de = rec.get("d_sequence_end", 0) - js = rec.get("j_sequence_start", 0) - je = rec.get("j_sequence_end", 0) - has_d = ds != de - has_j = js != je # J can be absent if 3' corruption removed it - is_rc = rec.get("is_reverse_complement", False) - order_ok = vs <= ve - if not is_rc: - if has_d and has_j: - order_ok = order_ok and ve <= ds and ds <= de and de <= js and js <= je - elif has_d: - order_ok = order_ok and ve <= ds and ds <= de - elif has_j: - order_ok = order_ok and ve <= js and js <= je - else: - # In RC read orientation, segment order is mirrored. - if has_d and has_j: - order_ok = order_ok and js <= je and je <= ds and ds <= de and de <= vs - elif has_d: - order_ok = order_ok and ds <= de and de <= vs - elif has_j: - order_ok = order_ok and js <= je and je <= vs - # If neither D nor J present, only V ordering matters - checks.append({ - "name": "segment_ordering", - "passed": order_ok, - "detail": ( - f"V=[{vs},{ve}] D=[{ds},{de}] J=[{js},{je}] rc={is_rc}" - if not order_ok else "ok" - ), - }) - - # 4. junction_bounds - jstart = rec.get("junction_start", 0) - jend = rec.get("junction_end", 0) - jb_ok = 0 <= jstart <= jend <= seq_len if isinstance(jstart, int) and isinstance(jend, int) else True - checks.append({ - "name": "junction_bounds", - "passed": jb_ok, - "detail": f"junction [{jstart},{jend}] seq_len={seq_len}" if not jb_ok else "ok", - }) - - # 5. junction_length_match - jlen = rec.get("junction_length", 0) or 0 - computed_jlen = jend - jstart if isinstance(jstart, int) and isinstance(jend, int) else 0 - jlm_ok = jlen == computed_jlen - checks.append({ - "name": "junction_length_match", - "passed": jlm_ok, - "detail": f"field={jlen} computed={computed_jlen}" if not jlm_ok else "ok", - }) - - # 6. productive_consistency - productive = rec.get("productive", False) - stop_codon = rec.get("stop_codon", False) - vj_in_frame = rec.get("vj_in_frame", False) - prod_ok = True - prod_detail = "ok" - if productive: - if stop_codon: - prod_ok = False - prod_detail = "productive=True but stop_codon=True" - if not vj_in_frame: - prod_ok = False - prod_detail = "productive=True but vj_in_frame=False" - checks.append({ - "name": "productive_consistency", - "passed": prod_ok, - "detail": prod_detail, - }) - - # 7. mutation_count - mut_str = rec.get("mutations", "") or "" - n_mut_field = rec.get("n_mutations", 0) or 0 - parsed = parse_mutations(mut_str) - mc_ok = n_mut_field == len(parsed) - checks.append({ - "name": "mutation_count", - "passed": mc_ok, - "detail": f"field={n_mut_field} parsed={len(parsed)}" if not mc_ok else "ok", - }) - - # 8. mutation_content - mut_errors = [] - if germ and seq and parsed: - for pos, (from_b, to_b) in parsed.items(): - if pos >= len(germ) or pos >= len(seq): - mut_errors.append(f"pos {pos} out of bounds") - continue - g = germ[pos].upper() - s = seq[pos].upper() - if g != "N" and g != "-" and from_b.upper() != g: - mut_errors.append(f"pos {pos}: germ={g} but from={from_b}") - if s != "N" and to_b.upper() != s: - mut_errors.append(f"pos {pos}: seq={s} but to={to_b}") - checks.append({ - "name": "mutation_content", - "passed": len(mut_errors) == 0, - "detail": "; ".join(mut_errors[:5]) if mut_errors else "ok", - }) - - # 9. sequence_length_field - sl_field = rec.get("sequence_length", None) - sl_ok = sl_field is None or sl_field == seq_len - checks.append({ - "name": "sequence_length_field", - "passed": sl_ok, - "detail": f"field={sl_field} actual={seq_len}" if not sl_ok else "ok", - }) - - n_passed = sum(1 for c in checks if c["passed"]) - n_failed = len(checks) - n_passed - return { - "valid": n_failed == 0, - "n_passed": n_passed, - "n_failed": n_failed, - "checks": checks, - } - - -# ═══════════════════════════════════════════════════════════════════ -# Germline alignment analysis -# ═══════════════════════════════════════════════════════════════════ - -def align_to_germline(dc, rec: Dict[str, Any]) -> Dict[str, Any]: - """Compare a record's sequence to its germline reference alleles.""" - seq = rec.get("sequence", "") - germ = rec.get("germline_alignment", "") - mut_str = rec.get("mutations", "") or "" - reported = parse_mutations(mut_str) - - # Build position-level comparison - details = [] - true_mut = 0 - unreported = 0 - phantom = 0 - - compare_len = min(len(seq), len(germ)) - for pos in range(compare_len): - s = seq[pos].upper() - g = germ[pos].upper() - if g in ("-", "N") or s == "N": - continue - is_reported = pos in reported - if s != g: - if is_reported: - true_mut += 1 - else: - unreported += 1 - if len(details) < 50: - details.append({ - "pos": pos, "seq_base": s, "germ_base": g, - "reported": False, "type": "unreported", - }) - else: - if is_reported: - phantom += 1 - if len(details) < 50: - details.append({ - "pos": pos, "seq_base": s, "germ_base": g, - "reported": True, "type": "phantom", - }) - - return { - "v_allele": rec.get("v_call", ""), - "d_allele": rec.get("d_call", ""), - "j_allele": rec.get("j_call", ""), - "n_true_mutations": true_mut, - "n_unreported": unreported, - "n_phantom": phantom, - "details": details, - } - - -# ═══════════════════════════════════════════════════════════════════ -# Region classification -# ═══════════════════════════════════════════════════════════════════ - -def classify_record_positions(dc, rec: Dict[str, Any]) -> Dict[str, Any]: - """Classify every position in the sequence into IMGT regions.""" - from ..utilities.imgt_regions import compute_v_region_boundaries, classify_position - - seq = rec.get("sequence", "") - v_call = (rec.get("v_call", "") or "").split(",")[0].strip() - if not v_call: - return {"error": "no v_call in record"} - - # Find V allele in config - v_allele = find_allele(dc, v_call) - if v_allele is None or not isinstance(v_allele, VAllele): - return {"error": f"V allele {v_call} not found in config"} - - v_boundaries = compute_v_region_boundaries(v_allele) - vs = rec.get("v_sequence_start", 0) - ve = rec.get("v_sequence_end", 0) - js = rec.get("j_sequence_start", 0) - je = rec.get("j_sequence_end", 0) - jstart = rec.get("junction_start", 0) - jend = rec.get("junction_end", 0) - - # Classify each position - region_counts: Dict[str, int] = defaultdict(int) - region_ranges: Dict[str, List[int]] = defaultdict(list) - for pos in range(len(seq)): - region = classify_position(pos, vs, ve, v_boundaries, jstart, jend, je) - region_counts[region] += 1 - region_ranges[region].append(pos) - - # Compute boundaries (min, max+1) for each region - boundaries = {} - for region, positions in region_ranges.items(): - if positions: - boundaries[region] = [min(positions), max(positions) + 1] - - # Mutations by region - mut_str = rec.get("mutations", "") or "" - parsed = parse_mutations(mut_str) - mut_by_region: Dict[str, int] = defaultdict(int) - for pos in parsed: - if pos < len(seq): - region = classify_position(pos, vs, ve, v_boundaries, jstart, jend, je) - mut_by_region[region] += 1 - - # Sequence excerpts by region (truncated) - seq_by_region = {} - for region, positions in region_ranges.items(): - if positions: - start, end = min(positions), max(positions) + 1 - excerpt = seq[start:end] - seq_by_region[region] = excerpt[:80] + ("..." if len(excerpt) > 80 else "") - - return { - "region_boundaries": boundaries, - "mutations_by_region": dict(mut_by_region), - "sequence_by_region": seq_by_region, - } - - -# ═══════════════════════════════════════════════════════════════════ -# Mutation pattern analysis -# ═══════════════════════════════════════════════════════════════════ - -_TRANSITIONS = { - ("A", "G"), ("G", "A"), # purine ↔ purine - ("C", "T"), ("T", "C"), # pyrimidine ↔ pyrimidine -} - -_WRCY = re.compile(r"[AT][AG]C[CT]", re.IGNORECASE) -_RGYW = re.compile(r"[AG]G[CT][AT]", re.IGNORECASE) - - -def analyze_mutation_patterns(records: List[Dict], dc) -> Dict[str, Any]: - """Analyze SHM patterns across a batch of records.""" - from ..utilities.imgt_regions import compute_v_region_boundaries, classify_position - - total_mutations = 0 - per_region: Dict[str, int] = defaultdict(int) - n_transitions = 0 - n_transversions = 0 - wrcy_count = 0 - rgyw_count = 0 - - for rec in records: - seq = rec.get("sequence", "") - germ = rec.get("germline_alignment", "") - mut_str = rec.get("mutations", "") or "" - parsed = parse_mutations(mut_str) - if not parsed: - continue - - # Get V allele for region classification - v_call = (rec.get("v_call", "") or "").split(",")[0].strip() - v_allele = find_allele(dc, v_call) if v_call else None - v_boundaries = ( - compute_v_region_boundaries(v_allele) - if v_allele and isinstance(v_allele, VAllele) - else {} - ) - - vs = rec.get("v_sequence_start", 0) - ve = rec.get("v_sequence_end", 0) - je = rec.get("j_sequence_end", 0) - jstart = rec.get("junction_start", 0) - jend = rec.get("junction_end", 0) - - for pos, (from_b, to_b) in parsed.items(): - total_mutations += 1 - fb, tb = from_b.upper(), to_b.upper() - - # Transition / transversion - if (fb, tb) in _TRANSITIONS: - n_transitions += 1 - elif fb in "ACGT" and tb in "ACGT": - n_transversions += 1 - - # Region classification - if v_boundaries and pos < len(seq): - region = classify_position(pos, vs, ve, v_boundaries, jstart, jend, je) - per_region[region] += 1 - - # Hotspot check on germline context - if germ and 1 <= pos < len(germ) - 2: - context = germ[pos - 1:pos + 3] - if len(context) == 4 and "-" not in context: - if _WRCY.fullmatch(context): - wrcy_count += 1 - if _RGYW.fullmatch(context): - rgyw_count += 1 - - # Compute fractions - per_region_out = {} - for region in ("FWR1", "CDR1", "FWR2", "CDR2", "FWR3", "CDR3", "FWR4", "NP"): - count = per_region.get(region, 0) - per_region_out[region] = { - "count": count, - "fraction": round(count / total_mutations, 4) if total_mutations else 0, - } - - ti_tv = round(n_transitions / n_transversions, 3) if n_transversions > 0 else None - motif_total = wrcy_count + rgyw_count - - return { - "n_records": len(records), - "total_mutations": total_mutations, - "per_region": per_region_out, - "transition_transversion_ratio": ti_tv, - "hotspot_analysis": { - "wrcy_count": wrcy_count, - "rgyw_count": rgyw_count, - "total_in_motif": motif_total, - "fraction_in_motif": round(motif_total / total_mutations, 4) if total_mutations else 0, - }, - } - - -# ═══════════════════════════════════════════════════════════════════ -# Statistical summary -# ═══════════════════════════════════════════════════════════════════ - -def _stats(values: List[float]) -> Dict[str, Any]: - """Compute mean/std/min/max for a list of floats.""" - if not values: - return {"mean": 0, "std": 0, "min": 0, "max": 0} - n = len(values) - mean = sum(values) / n - variance = sum((x - mean) ** 2 for x in values) / n if n > 1 else 0 - return { - "mean": round(mean, 4), - "std": round(math.sqrt(variance), 4), - "min": round(min(values), 4), - "max": round(max(values), 4), - } - - -def _top_n(counter: Dict[str, int], n: int = 20) -> Dict[str, int]: - """Return the top-N entries from a counter dict.""" - return dict(sorted(counter.items(), key=lambda kv: -kv[1])[:n]) - - -def compute_dataset_summary(records: List[Dict]) -> Dict[str, Any]: - """Compute aggregate statistics over a list of AIRR records.""" - n = len(records) - if n == 0: - return {"n_sequences": 0} - - productive_count = sum(1 for r in records if r.get("productive")) - mutation_rates = [r["mutation_rate"] for r in records if isinstance(r.get("mutation_rate"), (int, float))] - junction_lengths = [r["junction_length"] for r in records if isinstance(r.get("junction_length"), int) and r["junction_length"] > 0] - np1_lengths = [r["np1_length"] for r in records if isinstance(r.get("np1_length"), int)] - np2_lengths = [r["np2_length"] for r in records if isinstance(r.get("np2_length"), int)] - - v_usage: Dict[str, int] = defaultdict(int) - d_usage: Dict[str, int] = defaultdict(int) - j_usage: Dict[str, int] = defaultdict(int) - for r in records: - vc = (r.get("v_call", "") or "").split(",")[0].strip() - dc_call = (r.get("d_call", "") or "").split(",")[0].strip() - jc = (r.get("j_call", "") or "").split(",")[0].strip() - if vc: - v_usage[vc] += 1 - if dc_call: - d_usage[dc_call] += 1 - if jc: - j_usage[jc] += 1 - - return { - "n_sequences": n, - "productive_rate": round(productive_count / n, 4), - "mutation_rate": _stats(mutation_rates), - "junction_length": _stats([float(x) for x in junction_lengths]), - "v_gene_usage": _top_n(v_usage), - "d_gene_usage": _top_n(d_usage), - "j_gene_usage": _top_n(j_usage), - "np1_length": _stats([float(x) for x in np1_lengths]), - "np2_length": _stats([float(x) for x in np2_lengths]), - } - - -# ═══════════════════════════════════════════════════════════════════ -# Config section extraction -# ═══════════════════════════════════════════════════════════════════ - -def extract_config_section(dc, section: str) -> Dict[str, Any]: - """Extract a specific section of config internals for inspection.""" - if section == "gene_use": - raw = getattr(dc, "gene_use_dict", {}) or {} - out = {} - for seg, usage in raw.items(): - if isinstance(usage, dict): - out[seg] = {k: round(v, 4) if isinstance(v, float) else v for k, v in usage.items()} - else: - out[seg] = usage - return {"section": "gene_use", "data": out} - - if section == "trimming": - raw = getattr(dc, "trim_dicts", {}) or {} - out = {} - for trim_key, families in raw.items(): - if isinstance(families, dict): - summary = {} - for family, genes in families.items(): - if isinstance(genes, dict): - # Truncate to first 10 entries per gene - summary[family] = { - k: {str(kk): round(vv, 4) if isinstance(vv, float) else vv - for kk, vv in list(v.items())[:15]} - if isinstance(v, dict) else v - for k, v in list(genes.items())[:10] - } - else: - summary[family] = genes - out[trim_key] = summary - else: - out[trim_key] = families - return {"section": "trimming", "data": out} - - if section == "np_params": - return { - "section": "np_params", - "data": { - "NP_first_bases": getattr(dc, "NP_first_bases", {}), - "NP_lengths": getattr(dc, "NP_lengths", {}), - "NP_transitions_keys": list((getattr(dc, "NP_transitions", {}) or {}).keys()), - }, - } - - if section == "p_nucleotides": - return { - "section": "p_nucleotides", - "data": getattr(dc, "p_nucleotide_length_probs", {}), - } - - return {"error": f"Unknown section: {section!r}. Use: gene_use, trimming, np_params, p_nucleotides"} - - -# ═══════════════════════════════════════════════════════════════════ -# ASeq introspection helpers (for pipeline hook snapshots) -# ═══════════════════════════════════════════════════════════════════ - -def summarize_aseq(nodes: List[Dict]) -> Dict[str, Any]: - """Compact summary of an ASeq node list from a snapshot.""" - if not nodes: - return {"total_nodes": 0} - - seg_counts: Dict[str, int] = defaultdict(int) - flag_counts: Dict[str, int] = defaultdict(int) - anchor_positions: List[int] = [] - - for n in nodes: - seg_counts[n.get("seg", "?")] += 1 - for f in n.get("flags", []): - flag_counts[f] += 1 - if "anchor" in n.get("flags", []): - anchor_positions.append(n.get("pos", -1)) - - return { - "total_nodes": len(nodes), - "segments": dict(seg_counts), - "flag_counts": dict(flag_counts), - "anchor_positions": anchor_positions, - } - - -def detect_node_anomalies(nodes: List[Dict]) -> List[Dict[str, Any]]: - """Flag suspicious node states in an ASeq snapshot.""" - anomalies = [] - - for n in nodes: - pos = n.get("pos", -1) - cur = n.get("cur", "?") - germ = n.get("germ", ".") - seg = n.get("seg", "?") - flags = n.get("flags", []) - - # Mutated but cur == germ - if "mutated" in flags and cur == germ and germ != ".": - anomalies.append({ - "pos": pos, "type": "phantom_mutation", - "detail": f"flagged mutated but cur={cur}==germ={germ}", - }) - - # Indel insertion with mutated flag - if "indel_ins" in flags and "mutated" in flags: - anomalies.append({ - "pos": pos, "type": "mutated_indel", - "detail": f"indel-inserted node also has mutated flag", - }) - - # Germline segment node with null germline (indel-inserted nodes are exempt) - if seg in ("V", "D", "J", "C") and germ == "." and "indel_ins" not in flags: - anomalies.append({ - "pos": pos, "type": "missing_germline", - "detail": f"segment={seg} but germline is null", - }) - - # NP node with germline set (should be null) - if seg in ("NP1", "NP2") and germ not in (".", ""): - if "indel_ins" not in flags: - anomalies.append({ - "pos": pos, "type": "np_has_germline", - "detail": f"NP segment has germline={germ}", - }) - - # Non-ACGT current base in non-N node - if cur.upper() not in ("A", "C", "G", "T", "N", "?"): - anomalies.append({ - "pos": pos, "type": "invalid_base", - "detail": f"current base is '{cur}'", - }) - - return anomalies - - -def validate_codon_rail_snapshot(codons: List[Dict]) -> Dict[str, Any]: - """Validate codon rail from a snapshot.""" - if not codons: - return {"valid": True, "n_codons": 0, "n_stop_codons": 0, "issues": []} - - issues = [] - n_stop = 0 - - for c in codons: - aa = c.get("aa", ".") - is_stop = c.get("is_stop", False) - bases = c.get("bases", "???") - - if is_stop: - n_stop += 1 - - # Verify amino acid matches bases (basic check) - if len(bases) == 3 and "?" not in bases: - from ..utilities.misc import translate - translated = translate(bases) - expected_aa = translated[0] if translated else None - if expected_aa and aa != "." and aa != expected_aa: - issues.append({ - "codon_idx": c.get("idx", -1), - "bases": bases, - "expected_aa": expected_aa, - "actual_aa": aa, - }) - - return { - "valid": len(issues) == 0, - "n_codons": len(codons), - "n_stop_codons": n_stop, - "issues": issues[:20], - } - - -def build_germline_diff(nodes: List[Dict]) -> List[Dict[str, Any]]: - """Position-by-position diff where cur != germ.""" - diffs = [] - for n in nodes: - cur = n.get("cur", "?") - germ = n.get("germ", ".") - if germ != "." and cur != germ: - diffs.append({ - "pos": n.get("pos", -1), - "cur": cur, - "germ": germ, - "seg": n.get("seg", "?"), - "flags": n.get("flags", []), - }) - return diffs - - -def diff_snapshots( - snap_a_nodes: List[Dict], - snap_b_nodes: List[Dict], -) -> Dict[str, Any]: - """Compare two snapshots: what changed between them.""" - len_a = len(snap_a_nodes) - len_b = len(snap_b_nodes) - - # Build position maps - a_by_pos = {n["pos"]: n for n in snap_a_nodes} - b_by_pos = {n["pos"]: n for n in snap_b_nodes} - - modified = [] - added_positions = sorted(set(b_by_pos) - set(a_by_pos)) - removed_positions = sorted(set(a_by_pos) - set(b_by_pos)) - - # Find modifications at shared positions - shared = sorted(set(a_by_pos) & set(b_by_pos)) - for pos in shared: - a = a_by_pos[pos] - b = b_by_pos[pos] - changes = {} - if a.get("cur") != b.get("cur"): - changes["cur"] = {"from": a.get("cur"), "to": b.get("cur")} - if a.get("germ") != b.get("germ"): - changes["germ"] = {"from": a.get("germ"), "to": b.get("germ")} - if set(a.get("flags", [])) != set(b.get("flags", [])): - changes["flags"] = { - "added": sorted(set(b.get("flags", [])) - set(a.get("flags", []))), - "removed": sorted(set(a.get("flags", [])) - set(b.get("flags", []))), - } - if a.get("seg") != b.get("seg"): - changes["seg"] = {"from": a.get("seg"), "to": b.get("seg")} - if changes: - modified.append({"pos": pos, "changes": changes}) - - # Flag stats - a_flags = defaultdict(int) - b_flags = defaultdict(int) - for n in snap_a_nodes: - for f in n.get("flags", []): - a_flags[f] += 1 - for n in snap_b_nodes: - for f in n.get("flags", []): - b_flags[f] += 1 - - all_flag_names = sorted(set(a_flags) | set(b_flags)) - flag_changes = {} - for f in all_flag_names: - ac = a_flags.get(f, 0) - bc = b_flags.get(f, 0) - if ac != bc: - flag_changes[f] = {"before": ac, "after": bc, "delta": bc - ac} - - return { - "before_length": len_a, - "after_length": len_b, - "n_added": len(added_positions), - "n_removed": len(removed_positions), - "n_modified": len(modified), - "added_positions": added_positions[:50], - "removed_positions": removed_positions[:50], - "modified": modified[:50], - "flag_changes": flag_changes, - } - - -def format_snapshot_timeline(snapshots: List[Dict]) -> str: - """Human-readable timeline of how the sequence evolved through stages.""" - lines = [] - for snap in snapshots: - hook = snap.get("hook", "?") - nodes = snap.get("nodes", []) - summary = summarize_aseq(nodes) - - seg_str = " ".join( - f"{seg}:{count}" - for seg, count in sorted(summary.get("segments", {}).items()) - ) - flags = summary.get("flag_counts", {}) - - parts = [f"{hook:25s} {summary['total_nodes']:4d} nodes"] - if seg_str: - parts.append(f"({seg_str})") - if flags.get("mutated", 0): - parts.append(f"+{flags['mutated']} mutations") - if flags.get("indel_ins", 0): - parts.append(f"+{flags['indel_ins']} indels") - if flags.get("is_n", 0): - parts.append(f"+{flags['is_n']} Ns") - - lines.append(" ".join(parts)) - - return "\n".join(lines) - - -def aggregate_flag_stats( - all_nodes_lists: List[List[Dict]], -) -> Dict[str, Any]: - """Batch flag distribution across N sequences.""" - total_nodes = 0 - per_flag: Dict[str, int] = defaultdict(int) - per_seg_flag: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int)) - n_anomalies = 0 - - for nodes in all_nodes_lists: - total_nodes += len(nodes) - anomalies = detect_node_anomalies(nodes) - n_anomalies += len(anomalies) - for n in nodes: - seg = n.get("seg", "?") - for f in n.get("flags", []): - per_flag[f] += 1 - per_seg_flag[seg][f] += 1 - - return { - "n_sequences": len(all_nodes_lists), - "total_nodes": total_nodes, - "per_flag_counts": dict(per_flag), - "per_segment_flag_counts": { - seg: dict(flags) for seg, flags in per_seg_flag.items() - }, - "total_anomalies": n_anomalies, - } diff --git a/src/GenAIRR/utilities/misc.py b/src/GenAIRR/utilities/misc.py index 760114d..a2f0717 100644 --- a/src/GenAIRR/utilities/misc.py +++ b/src/GenAIRR/utilities/misc.py @@ -1,5 +1,4 @@ import random -from collections import defaultdict _COMPLEMENT = str.maketrans('ATCGatcg', 'TAGCtagc') @@ -98,18 +97,6 @@ def parse_mutation(mutation_str: str): return parts[0], parts[-1] -def normalize_and_filter_convert_to_dict(obj): - if isinstance(obj, defaultdict): - nested_dict = {k: normalize_and_filter_convert_to_dict(v) for k, v in obj.items()} - if all(isinstance(val, float) for val in nested_dict.values()): - # Filter out keys that are not 'A', 'T', 'G', or 'C' - filtered_dict = {k: v for k, v in nested_dict.items() if k in ['A', 'T', 'G', 'C']} - total = sum(filtered_dict.values()) - return {k: v / total for k, v in filtered_dict.items()} - return nested_dict - return obj - - def parse_fasta(file): """Parse a FASTA format file. diff --git a/src/GenAIRR/utilities/visualize.py b/src/GenAIRR/utilities/visualize.py index 941eb53..a9ebd80 100644 --- a/src/GenAIRR/utilities/visualize.py +++ b/src/GenAIRR/utilities/visualize.py @@ -1,844 +1,8 @@ -""" -visualize.py — Generate standalone HTML files showing an "exploding view" -dissection of a simulated AIRR sequence record. - -Usage: - from GenAIRR.utilities.visualize import visualize_sequence - result = experiment.run(n=1) - visualize_sequence(result[0], "sequence.html") -""" - -from __future__ import annotations - -import html -import json -from pathlib import Path -from typing import Any, Dict, Optional, Union - -# ── Segment colors (matching the website DissectionBay) ────── - -SEGMENT_COLORS = { - "V": "#4A90D9", - "NP1": "#5DC48C", - "D": "#E8685A", - "NP2": "#5DC48C", - "J": "#E8A838", -} - -RED = "#DC2626" -PURPLE = "#9B6FC4" - - -# ── Field extraction helpers ───────────────────────────────── - -def _int(rec: dict, key: str, default: int = 0) -> int: - v = rec.get(key) - if v is None: - return default - try: - return int(v) - except (ValueError, TypeError): - return default - - -def _str(rec: dict, key: str, default: str = "") -> str: - v = rec.get(key) - return str(v) if v is not None else default - - -def _bool(rec: dict, key: str) -> bool: - v = rec.get(key) - if isinstance(v, bool): - return v - if isinstance(v, str): - return v.lower() in ("true", "1", "yes", "t") - return bool(v) if v is not None else False - - -def _float(rec: dict, key: str, default: float = 0.0) -> float: - v = rec.get(key) - if v is None: - return default - try: - return float(v) - except (ValueError, TypeError): - return default - - -def _parse_mutations(raw: Any) -> Dict[int, str]: - """Parse the mutations field (string like '42:A>G;100:T>C' or dict).""" - if not raw: - return {} - if isinstance(raw, dict): - return {int(k): str(v) for k, v in raw.items()} - if isinstance(raw, str): - # Try JSON first - try: - parsed = json.loads(raw) - if isinstance(parsed, dict): - return {int(k): str(v) for k, v in parsed.items()} - except (json.JSONDecodeError, ValueError): - pass - # Try comma or semicolon-separated "pos:from>to" format - out = {} - # Split on comma or semicolon - sep = "," if "," in raw else ";" - for part in raw.split(sep): - part = part.strip() - if not part: - continue - if ":" in part: - pos_str, desc = part.split(":", 1) - try: - out[int(pos_str)] = desc - except ValueError: - pass - return out - return {} - - -# ── Build the HTML ──────────────────────────────────────────── - -def _esc(s: str) -> str: - return html.escape(s, quote=True) - - -def _build_segment_bar_html(segments: list, seq_len: int, mutations: Dict[int, str]) -> str: - """Build the color-coded assembled sequence bar.""" - parts = [] - for seg in segments: - w = max(((seg["end"] - seg["start"]) / seq_len) * 100, 1.5) - label = seg["label"] if w > 4 else "" - start_label = seg["start"] + 1 - end_label = seg["end"] - parts.append( - f'
' - f'{_esc(label)}' - f'{start_label}' - + (f'{end_label}' if w > 8 else "") - + "
" - ) - - # Mutation dots - dots = [] - for pos in mutations: - left_pct = (pos / seq_len) * 100 - dots.append(f'
') - - return ( - '
' - + "".join(parts) - + "".join(dots) - + "
" - ) - - -def _build_junction_bracket(junction_start: int, junction_end: int, junction_aa: str, seq_len: int) -> str: - if seq_len == 0 or junction_end <= junction_start: - return "" - left_pct = (junction_start / seq_len) * 100 - width_pct = ((junction_end - junction_start) / seq_len) * 100 - return ( - '
' - f'
' - '' - f'' - f'' - f'' - "" - f'CDR3: {_esc(junction_aa)}' - "
" - "
" - ) - - -def _sparkline_svg(start: int, end: int, mutations: Dict[int, str], height: int = 24) -> str: - length = end - start - if length <= 0: - return "" - bins = min(length, 60) - bin_size = max(1, length // bins) - counts = [] - for i in range(bins): - b_start = start + i * bin_size - b_end = b_start + bin_size - c = sum(1 for p in mutations if b_start <= p < b_end) - counts.append(c) - mx = max(max(counts), 1) - rects = [] - for i, c in enumerate(counts): - fill = RED if c > 0 else "#555" - opacity = 0.85 if c > 0 else 0.2 - h = c if c > 0 else 0.15 - rects.append( - f'' - ) - return ( - f'' - + "".join(rects) - + "" - ) - - -def _pnp_bar(prefix: str, n_region: str, suffix: str, color: str) -> str: - total = len(prefix) + len(n_region) + len(suffix) - if total == 0: - return "" - parts_data = [ - ("P", len(prefix), 0.6), - ("N", len(n_region), 0.25), - ("P", len(suffix), 0.6), - ] - parts = [] - for label, length, opacity in parts_data: - if length > 0: - w = max((length / total) * 100, 3) - parts.append( - f'
' - f'{label}
' - ) - return '
' + "".join(parts) + "
" - - -def _nt_strip(sequence: str, start_pos: int, color: str, mutations: Dict[int, str]) -> str: - """Nucleotide detail with mutation highlighting.""" - rows = [] - for row_idx in range(0, len(sequence), 80): - chunk = sequence[row_idx : row_idx + 80] - chars = [] - for ci, ch in enumerate(chunk): - abs_pos = start_pos + row_idx + ci - is_mut = abs_pos in mutations - cls = "nt-char mut" if is_mut else "nt-char" - dot = '' if is_mut else "" - chars.append(f'{dot}{_esc(ch)}') - line_start = start_pos + row_idx + 1 - line_end = min(start_pos + row_idx + 80, start_pos + len(sequence)) - rows.append( - f'
' - f'{line_start}' - f'
{"".join(chars)}
' - f'{line_end}' - f"
" - ) - bg = f"color-mix(in srgb, {color} 5%, #1a1a2e)" - return f'
{"".join(rows)}
' - - -def _metric(label: str, value: str, color: Optional[str] = None) -> str: - style = f' style="color:{color}"' if color else "" - return ( - f'
' - f'{_esc(label)}' - f'{_esc(value)}' - f"
" - ) - - -def _trim_indicator(label: str, bases: int) -> str: - if not bases: - return "" - return ( - f'
' - f'✂ ' - f'{_esc(label)}: {bases} bp trimmed' - f"
" - ) - - -def _segment_panel(seg: dict, rec: dict, mutations: Dict[int, str], sequence: str) -> str: - """Build a single exploded segment panel.""" - seg_start = seg["start"] - seg_end = seg["end"] - seg_len = seg_end - seg_start - sub_seq = sequence[seg_start:seg_end] - seg_muts = {p: v for p, v in mutations.items() if seg_start <= p < seg_end} - mut_count = len(seg_muts) - seg_id = seg["id"] - color = seg["color"] - - header = ( - f'
' - f'{_esc(seg["label"])}' - f'{_esc(seg["full_label"])}' - f'{seg_start + 1}..{seg_end} ({seg_len} nt)' - f"
" - ) - - body_parts = [] - - if seg_id == "V": - body_parts.append( - f'
' - + _metric("LENGTH", f"{seg_len} bp") - + _metric("MUTATIONS", str(mut_count), RED if mut_count > 0 else None) - + _metric("3' TRIM", f"{_int(rec, 'v_trim_3')} bp") - + "
" - ) - if seg_len > 0: - body_parts.append( - '
' - 'Mutation Density' - + _sparkline_svg(seg_start, seg_end, mutations) - + "
" - ) - body_parts.append(_trim_indicator("V-gene 3'", _int(rec, "v_trim_3"))) - - elif seg_id == "D": - d_inv = _bool(rec, "d_inverted") - body_parts.append( - f'
' - + _metric("LENGTH", f"{seg_len} bp") - + _metric("5' TRIM", f"{_int(rec, 'd_trim_5')} bp") - + _metric("3' TRIM", f"{_int(rec, 'd_trim_3')} bp") - + _metric("INVERTED", "YES" if d_inv else "NO", RED if d_inv else None) - + "
" - ) - if d_inv: - body_parts.append( - '
' - '⇅ ' - "D-gene inverted (reverse complement used)
" - ) - body_parts.append(_trim_indicator("5'", _int(rec, "d_trim_5"))) - body_parts.append(_trim_indicator("3'", _int(rec, "d_trim_3"))) - - elif seg_id == "J": - body_parts.append( - f'
' - + _metric("LENGTH", f"{seg_len} bp") - + _metric("MUTATIONS", str(mut_count), RED if mut_count > 0 else None) - + _metric("5' TRIM", f"{_int(rec, 'j_trim_5')} bp") - + "
" - ) - if seg_len > 0: - body_parts.append( - '
' - 'Mutation Density' - + _sparkline_svg(seg_start, seg_end, mutations) - + "
" - ) - body_parts.append(_trim_indicator("J-gene 5'", _int(rec, "j_trim_5"))) - - elif seg_id == "NP1": - p_pre = _str(rec, "np1_p_prefix") - n_reg = _str(rec, "np1_n_region") - p_suf = _str(rec, "np1_p_suffix") - body_parts.append( - f'
' - + _metric("P-PREFIX", f"{len(p_pre)} bp") - + _metric("N-ADDITION", f"{len(n_reg)} bp") - + _metric("P-SUFFIX", f"{len(p_suf)} bp") - + "
" - ) - body_parts.append( - 'P | N | P Composition' - + _pnp_bar(p_pre, n_reg, p_suf, color) - ) - body_parts.append( - '
' - + f'{_esc(p_pre)}' - + f'{_esc(n_reg)}' - + f'{_esc(p_suf)}' - + "
" - ) - - elif seg_id == "NP2": - p_pre = _str(rec, "np2_p_prefix") - n_reg = _str(rec, "np2_n_region") - p_suf = _str(rec, "np2_p_suffix") - body_parts.append( - f'
' - + _metric("P-PREFIX", f"{len(p_pre)} bp") - + _metric("N-ADDITION", f"{len(n_reg)} bp") - + _metric("P-SUFFIX", f"{len(p_suf)} bp") - + "
" - ) - body_parts.append( - 'P | N | P Composition' - + _pnp_bar(p_pre, n_reg, p_suf, color) - ) - body_parts.append( - '
' - + f'{_esc(p_pre)}' - + f'{_esc(n_reg)}' - + f'{_esc(p_suf)}' - + "
" - ) +"""Standalone HTML "exploding view" of a simulated AIRR record. - # Nucleotide detail (always shown in static view) - if seg_len > 0 and seg_id in ("V", "D", "J"): - mut_note = f' ({mut_count} mutations highlighted)' if mut_count > 0 else "" - body_parts.append( - f'
' - f'Nucleotide Sequence{mut_note}' - + _nt_strip(sub_seq, seg_start, color, seg_muts) - + "
" - ) - - return ( - f'
' - + header - + '
' - + "".join(body_parts) - + "
" - ) - - -# ── CSS ────────────────────────────────────────────────────── - -CSS = """ -:root { - --bg: #0f1019; - --bg2: #181825; - --bg3: #1e1e32; - --fg: #e0e0f0; - --fg2: #8888aa; - --border: #2a2a40; - --accent: #6366f1; - --red: #DC2626; - --purple: #9B6FC4; -} -* { box-sizing: border-box; margin: 0; padding: 0; } -body { - font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; - background: var(--bg); - color: var(--fg); - line-height: 1.5; - padding: 2rem; -} -.container { max-width: 1200px; margin: 0 auto; } -h1 { font-size: 1.4rem; font-weight: 600; margin-bottom: 0.25rem; } -h2 { font-size: 1.1rem; font-weight: 600; margin-bottom: 0.75rem; color: var(--fg2); } -.subtitle { color: var(--fg2); font-size: 0.85rem; margin-bottom: 1.5rem; } - -/* Top bar */ -.top-bar { - display: flex; align-items: center; gap: 0.75rem; - padding: 1rem 1.25rem; - background: var(--bg2); border: 1px solid var(--border); - border-radius: 10px; margin-bottom: 1.25rem; -} -.top-bar .dna-icon { color: var(--accent); font-size: 1.2rem; } -.badge { - padding: 0.2rem 0.6rem; border-radius: 999px; - font-size: 0.7rem; font-weight: 600; text-transform: uppercase; - letter-spacing: 0.05em; -} -.badge-pos { background: rgba(34,197,94,0.15); color: #22c55e; } -.badge-neg { background: rgba(220,38,38,0.15); color: #DC2626; } -.seq-id { color: var(--fg2); font-size: 0.8rem; margin-left: auto; font-family: monospace; } - -/* Summary grid */ -.summary-grid { - display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr)); - gap: 0.5rem; margin-bottom: 1.25rem; -} -.summary-cell { - background: var(--bg2); border: 1px solid var(--border); - border-radius: 8px; padding: 0.6rem 0.8rem; -} -.summary-cell .metric-label { - display: block; font-size: 0.6rem; text-transform: uppercase; - letter-spacing: 0.08em; color: var(--fg2); margin-bottom: 0.15rem; -} -.summary-cell .metric-value { - font-size: 0.9rem; font-weight: 600; -} -.summary-cell .metric-value.mono { font-family: monospace; font-size: 0.75rem; word-break: break-all; } - -/* Corruption */ -.corruption-row { - display: flex; gap: 0.75rem; margin-bottom: 1rem; -} -.corruption-badge { - display: inline-flex; align-items: center; gap: 0.4rem; - background: rgba(220,38,38,0.08); border: 1px solid rgba(220,38,38,0.25); - border-radius: 6px; padding: 0.35rem 0.75rem; - font-size: 0.75rem; color: var(--red); -} - -/* Assembled bar */ -.section-label { font-size: 0.7rem; text-transform: uppercase; letter-spacing: 0.08em; color: var(--fg2); margin-bottom: 0.4rem; } -.assembled-bar { - display: flex; height: 38px; border-radius: 6px; overflow: hidden; - position: relative; margin-bottom: 0; - border: 1px solid var(--border); -} -.seg-bar-part { - position: relative; display: flex; align-items: center; justify-content: center; - cursor: default; transition: opacity 0.2s; - min-width: 0; -} -.seg-bar-label { - font-size: 0.7rem; font-weight: 700; color: #fff; text-shadow: 0 1px 2px rgba(0,0,0,0.4); - pointer-events: none; -} -.seg-bar-pos { - position: absolute; bottom: 2px; font-size: 0.55rem; color: rgba(255,255,255,0.7); - pointer-events: none; -} -.seg-bar-pos.left { left: 3px; } -.seg-bar-pos.right { right: 3px; } -.mut-dot { - position: absolute; top: -2px; width: 4px; height: 4px; - background: var(--red); border-radius: 50%; - transform: translateX(-50%); - pointer-events: none; -} - -/* Junction bracket */ -.junction-bracket { position: relative; height: 28px; margin-bottom: 1.5rem; } -.junction-inner { position: absolute; text-align: center; } -.junction-aa { - display: block; font-size: 0.65rem; color: var(--purple); - font-family: monospace; margin-top: 1px; white-space: nowrap; - overflow: hidden; text-overflow: ellipsis; -} - -/* Segment panels */ -.segments-grid { - display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); - gap: 0.75rem; margin-top: 1rem; -} -.seg-panel { - background: var(--bg2); border: 1px solid var(--border); - border-radius: 8px; overflow: hidden; -} -.seg-panel-header { - display: flex; align-items: center; gap: 0.5rem; - padding: 0.6rem 0.8rem; border-bottom: 1px solid var(--border); -} -.seg-chip { - display: inline-block; padding: 0.15rem 0.5rem; border-radius: 4px; - font-size: 0.65rem; font-weight: 700; color: #fff; -} -.seg-gene { font-size: 0.75rem; font-weight: 600; } -.seg-pos { font-size: 0.7rem; color: var(--fg2); margin-left: auto; } -.seg-panel-body { padding: 0.75rem 0.8rem; display: flex; flex-direction: column; gap: 0.6rem; } - -/* Metrics inside panels */ -.metric-grid { display: grid; gap: 0.4rem; } -.metric-grid.g3 { grid-template-columns: repeat(3, 1fr); } -.metric-grid.g4 { grid-template-columns: repeat(4, 1fr); } -.metric { - background: var(--bg3); border-radius: 5px; padding: 0.35rem 0.5rem; - text-align: center; -} -.metric-label { - display: block; font-size: 0.55rem; text-transform: uppercase; - letter-spacing: 0.06em; color: var(--fg2); -} -.metric-value { font-size: 0.8rem; font-weight: 600; } - -/* Trim */ -.trim-info { - font-size: 0.7rem; color: var(--fg2); - display: flex; align-items: center; gap: 0.3rem; -} -.trim-icon { font-size: 0.85rem; } - -/* Inversion */ -.inv-badge { - display: flex; align-items: center; gap: 0.4rem; - background: rgba(220,38,38,0.08); border: 1px solid rgba(220,38,38,0.2); - border-radius: 5px; padding: 0.3rem 0.6rem; - font-size: 0.7rem; color: var(--red); -} -.inv-icon { font-size: 1rem; } - -/* PNP bar */ -.pnp-bar { - display: flex; height: 20px; border-radius: 4px; overflow: hidden; - border: 1px solid var(--border); -} -.pnp-part { - display: flex; align-items: center; justify-content: center; - min-width: 0; -} -.pnp-label { font-size: 0.6rem; font-weight: 700; color: #fff; } -.np-seq { - font-family: monospace; font-size: 0.7rem; word-break: break-all; - padding: 0.3rem; background: var(--bg3); border-radius: 4px; -} - -/* Nucleotide strip */ -.nt-section { margin-top: 0.25rem; } -.nt-strip { - border-radius: 6px; padding: 0.5rem; overflow-x: auto; - font-family: 'Fira Code', 'JetBrains Mono', 'Cascadia Code', monospace; - font-size: 0.65rem; line-height: 1.6; -} -.nt-row { display: flex; align-items: center; gap: 0.5rem; } -.nt-linenum { color: var(--fg2); min-width: 3ch; text-align: right; font-size: 0.6rem; user-select: none; } -.nt-linenum.end { text-align: left; } -.nt-chars { display: flex; flex-wrap: wrap; } -.nt-char { - position: relative; display: inline-block; width: 0.8em; text-align: center; -} -.nt-char.mut { color: var(--red); font-weight: 700; } -.mut-marker { - position: absolute; top: -3px; left: 50%; transform: translateX(-50%); - width: 3px; height: 3px; border-radius: 50%; background: var(--red); -} -.micro-label { - display: block; font-size: 0.6rem; text-transform: uppercase; - letter-spacing: 0.06em; color: var(--fg2); margin-bottom: 0.3rem; -} -.spark-section { /* wrapper */ } - -/* Germline alignment */ -.germline-section { margin-top: 1.25rem; } -.germline-row { - display: flex; font-family: monospace; font-size: 0.65rem; line-height: 1.6; - gap: 0.5rem; align-items: center; -} -.germline-label { min-width: 5ch; color: var(--fg2); font-size: 0.6rem; text-align: right; } -.germline-chars { display: flex; flex-wrap: wrap; } -.germline-match { color: var(--fg2); opacity: 0.4; } -.germline-mismatch { color: var(--red); font-weight: 700; } -.germline-gap { color: var(--fg2); opacity: 0.2; } - -/* Footer */ -.footer { - margin-top: 2rem; padding-top: 1rem; border-top: 1px solid var(--border); - font-size: 0.65rem; color: var(--fg2); text-align: center; -} +Public entry point: :func:`visualize_sequence`. Implementation lives in +the `_visualize/` package (parse / styles / components / alignment / render). """ +from ._visualize.render import visualize_sequence - -# ── Main function ──────────────────────────────────────────── - -def visualize_sequence( - record: Dict[str, Any], - path: Union[str, Path], - title: Optional[str] = None, -) -> Path: - """ - Generate a standalone HTML file with an "exploding view" dissection - of a simulated AIRR sequence record. - - Parameters - ---------- - record : dict - A single AIRR record dict (one element from SimulationResult). - path : str or Path - Output file path for the HTML file. - title : str, optional - Custom title for the page. Defaults to "Sequence Dissection". - - Returns - ------- - Path - The path to the generated file. - """ - path = Path(path) - rec = record - - # Extract fields - sequence = _str(rec, "sequence") - seq_len = len(sequence) - germline = _str(rec, "germline_alignment") - v_call = _str(rec, "v_call") - d_call = _str(rec, "d_call") - j_call = _str(rec, "j_call") - junction_aa = _str(rec, "junction_aa") - productive = _bool(rec, "productive") - mutation_rate = _float(rec, "mutation_rate") - mutations = _parse_mutations(rec.get("mutations")) - total_mutations = len(mutations) - cdr3_length = len(junction_aa) - - v_seq_start = _int(rec, "v_sequence_start") - v_seq_end = _int(rec, "v_sequence_end") - d_seq_start = _int(rec, "d_sequence_start") - d_seq_end = _int(rec, "d_sequence_end") - j_seq_start = _int(rec, "j_sequence_start") - j_seq_end = _int(rec, "j_sequence_end") - junction_start = _int(rec, "junction_start", _int(rec, "junction_sequence_start")) - junction_end = _int(rec, "junction_end", _int(rec, "junction_sequence_end")) - - corruption_5 = _int(rec, "corruption_5prime") - corruption_3 = _int(rec, "corruption_3prime") - - seq_id = _str(rec, "sequence_id", "sequence") - page_title = title or "Sequence Dissection" - - # Build segments - segments = [ - {"id": "V", "label": "V", "full_label": v_call, "color": SEGMENT_COLORS["V"], "start": v_seq_start, "end": v_seq_end}, - {"id": "NP1", "label": "NP1", "full_label": "N/P Region 1", "color": SEGMENT_COLORS["NP1"], "start": v_seq_end, "end": d_seq_start}, - {"id": "D", "label": "D", "full_label": d_call, "color": SEGMENT_COLORS["D"], "start": d_seq_start, "end": d_seq_end}, - {"id": "NP2", "label": "NP2", "full_label": "N/P Region 2", "color": SEGMENT_COLORS["NP2"], "start": d_seq_end, "end": j_seq_start}, - {"id": "J", "label": "J", "full_label": j_call, "color": SEGMENT_COLORS["J"], "start": j_seq_start, "end": j_seq_end}, - ] - - # ── Build HTML sections ────────────────────────────────── - - # Top bar - badge_cls = "badge badge-pos" if productive else "badge badge-neg" - badge_text = "Productive" if productive else "Non-productive" - top_bar = ( - f'
' - f'🧬' - f'

{_esc(page_title)}

' - f'{badge_text}' - f'{_esc(seq_id)}' - f"
" - ) - - # Summary grid - mut_color = RED if total_mutations > 0 else None - prod_color = RED if not productive else None - summary_items = [ - ("V-GENE", v_call, SEGMENT_COLORS["V"], False), - ("D-GENE", d_call, SEGMENT_COLORS["D"], False), - ("J-GENE", j_call, SEGMENT_COLORS["J"], False), - ("JUNCTION AA", junction_aa, None, True), - ("TOTAL LENGTH", f"{seq_len} nt", None, False), - ("MUTATIONS", f"{total_mutations} ({mutation_rate * 100:.1f}%)", mut_color, False), - ("CDR3 LENGTH", f"{cdr3_length} aa", None, False), - ("PRODUCTIVE", "Yes" if productive else "No", prod_color, False), - ] - summary_cells = [] - for label, val, color, mono in summary_items: - style = f' style="color:{color}"' if color else "" - mono_cls = " mono" if mono else "" - summary_cells.append( - f'
' - f'{_esc(label)}' - f'{_esc(str(val))}' - f"
" - ) - summary_grid = '
' + "".join(summary_cells) + "
" - - # Corruption warnings - corruption_html = "" - if corruption_5 or corruption_3: - parts = [] - if corruption_5: - parts.append(f'
⚠ 5\' Corruption: +{corruption_5} nt added
') - if corruption_3: - parts.append(f'
⚠ 3\' Corruption: −{corruption_3} nt removed
') - corruption_html = '
' + "".join(parts) + "
" - - # Assembled bar + junction - bar_border_left = f"border-left:3px dashed {RED};" if corruption_5 else "" - bar_border_right = f"border-right:3px dashed {RED};" if corruption_3 else "" - if bar_border_left or bar_border_right: - bar_html = _build_segment_bar_html(segments, seq_len, mutations) - bar_html = bar_html.replace( - 'class="assembled-bar"', - f'class="assembled-bar" style="{bar_border_left}{bar_border_right}"', - ) - else: - bar_html = _build_segment_bar_html(segments, seq_len, mutations) - - junction_html = _build_junction_bracket(junction_start, junction_end, junction_aa, seq_len) - - # Exploded segment panels - panels = [] - for seg in segments: - panels.append(_segment_panel(seg, rec, mutations, sequence)) - - # Germline alignment (if available) - germline_html = "" - if germline and len(germline) >= len(sequence): - germline_html = _build_germline_alignment(sequence, germline) - - # Footer - footer = '
Generated by GenAIRR — Synthetic AIRR Sequence Simulator
' - - # ── Assemble HTML ──────────────────────────────────────── - - html_content = f""" - - - - -{_esc(page_title)} — {_esc(seq_id)} - - - -
-{top_bar} -{summary_grid} -{corruption_html} - -
-Assembled Sequence -{bar_html} -{junction_html} -
- -
-⚡ Exploded Segments -
-{"".join(panels)} -
-
- -{germline_html} -{footer} -
- -""" - - path.write_text(html_content, encoding="utf-8") - return path - - -def _build_germline_alignment(sequence: str, germline: str) -> str: - """Build a germline vs sequence alignment visualization.""" - # Show first 200 positions to keep it readable - show_len = min(len(sequence), len(germline), 400) - if show_len == 0: - return "" - - rows = [] - for row_start in range(0, show_len, 80): - chunk_seq = sequence[row_start : row_start + 80] - chunk_germ = germline[row_start : row_start + 80] - - seq_chars = [] - germ_chars = [] - match_chars = [] - for s, g in zip(chunk_seq, chunk_germ): - if g == "." or g == "N": - germ_chars.append(f'{_esc(g)}') - seq_chars.append(f'{_esc(s)}') - match_chars.append(f' ') - elif s == g: - germ_chars.append(f'{_esc(g)}') - seq_chars.append(f'{_esc(s)}') - match_chars.append(f'|') - else: - germ_chars.append(f'{_esc(g)}') - seq_chars.append(f'{_esc(s)}') - match_chars.append(f'*') - - pos_label = str(row_start + 1) - rows.append( - f'
' - f'{pos_label}' - f'
{"".join(germ_chars)}
' - f'
' - f'
' - f'' - f'
{"".join(match_chars)}
' - f'
' - f'
' - f'' - f'
{"".join(seq_chars)}
' - f'
' - ) - - return ( - '
' - 'Germline Alignment' - + "".join(rows) - + "
" - ) +__all__ = ["visualize_sequence"] diff --git a/tests/test_allele_usage_estimation_contract.py b/tests/test_allele_usage_estimation_contract.py index 09a8379..3a9af9e 100644 --- a/tests/test_allele_usage_estimation_contract.py +++ b/tests/test_allele_usage_estimation_contract.py @@ -99,10 +99,10 @@ def test_pin_scaffold_pipeline_ir_recombine_step_has_weights_fields() -> None: def test_pin_scaffold_lower_recombine_passes_weights_to_push_sample_allele() -> None: - """`_compile.py::_lower_recombine` calls + """`_lowering.py::_lower_recombine` calls ``plan.push_sample_allele(..., weights=...)`` with the per-step weight vector. Pinned at source.""" - src = (_REPO_ROOT / "src" / "GenAIRR" / "_compile.py").read_text(encoding="utf-8") + src = (_REPO_ROOT / "src" / "GenAIRR" / "_lowering.py").read_text(encoding="utf-8") assert "push_sample_allele(" in src assert "weights=v_weights" in src assert "weights=j_weights" in src @@ -315,7 +315,12 @@ def test_pin_present_gene_use_dict_has_no_simulator_pipeline_consumer() -> None: } allowed = { "src/GenAIRR/dataconfig/data_config.py", - "src/GenAIRR/utilities/mcp_helpers.py", + # Manifest reporting surface extracted verbatim from + # data_config.py (behavior-preserving hygiene split). The + # documented-orphan-fields tuple + allele_usage manifest + # block name the legacy field for REPORTING only, not as a + # simulator consumer — same category as data_config.py. + "src/GenAIRR/dataconfig/_manifest.py", # Post-Allele-Usage-Estimation-v1 slice — the new typed # plane's resolver / lowering / spec docstrings # explicitly document the no-auto-lift boundary @@ -326,7 +331,7 @@ def test_pin_present_gene_use_dict_has_no_simulator_pipeline_consumer() -> None: # comment explaining that the new estimator does NOT # touch `gene_use_dict`. "src/GenAIRR/_dataconfig_extract.py", - "src/GenAIRR/experiment.py", + "src/GenAIRR/_experiment/recombination.py", "src/GenAIRR/reference_models.py", } unexpected = consumers - allowed @@ -381,21 +386,6 @@ def test_pin_scaffold_dataconfig_validate_is_dead_code_today() -> None: ) -def test_pin_present_mcp_helpers_gene_use_endpoint_is_read_only() -> None: - """The MCP helper's `gene_use` diagnostic endpoint - reads `getattr(dc, "gene_use_dict", {})` as a read-only - inspection — it does NOT feed simulation. Pinned at - source so a refactor wiring the endpoint into the - sampler surfaces here.""" - src = ( - _REPO_ROOT / "src" / "GenAIRR" / "utilities" / "mcp_helpers.py" - ).read_text(encoding="utf-8") - assert "gene_use_dict" in src - assert 'section == "gene_use"' in src - # The endpoint reads, doesn't push to a plan. - assert "plan.push" not in src # no engine wiring in mcp_helpers - - # ────────────────────────────────────────────────────────────────── # F. pin_present_* — documented soft-gap boundary state # ────────────────────────────────────────────────────────────────── diff --git a/tests/test_build_imgt_configs.py b/tests/test_build_imgt_configs.py new file mode 100644 index 0000000..5ce0354 --- /dev/null +++ b/tests/test_build_imgt_configs.py @@ -0,0 +1,127 @@ +"""Tests for `tools/build_imgt_configs.py` (the IMGT cartridge build tool). + +No network: `build_cartridge` is exercised with inline FASTA fixtures, so +the download/probe path is not touched. Verifies the port to +`ReferenceCartridgeBuilder` produces valid, compilable structural +cartridges and that the CLI refuses to run without an explicit +`--output-dir`. +""" +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +import GenAIRR as ga +from GenAIRR.dataconfig.enums import ChainType, Species + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_TOOL_PATH = _REPO_ROOT / "tools" / "build_imgt_configs.py" + + +def _load_tool(): + """Load the standalone tool module (it lives outside the package).""" + spec = importlib.util.spec_from_file_location("build_imgt_configs", _TOOL_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +tool = _load_tool() + +# Tiny synthetic FASTA — enough length for the recombination passes; the +# builder degrades gracefully without native anchors (compile under +# allow_curatable_refdata, per the v1 builder workflow). +_V_FASTA = ( + ">IGHV1-MOCK*01\n" + "GAGGTGCAGCTGGTGGAGTCTGGGGGAGGCTTGGTACAGCCTGGGGGGTCCCTGAGACTC" + "TCCTGTGCAGCCTCTGGATTCACCTTCAGTAGCTATGCCATGAGCTGGGTCCGCCAGGCT\n" + ">IGHV2-MOCK*01\n" + "CAGGTCAACTTAAGGGAGTCTGGTCCTGCGCTGGTGAAACCCACACAGACCCTCACACTG" + "ACCTGCACCTTCTCTGGGTTCTCACTCAGCACTAGTGGAATGTGTGTGAGCTGGATCCGT\n" +) +_D_FASTA = ( + ">IGHD1-MOCK*01\n" + "GGGTATAGCAGCAGCTGGTAC\n" + ">IGHD2-MOCK*01\n" + "AGGATATTGTAGTGGTGGTAGCTGCTACTCC\n" +) +_J_FASTA = ( + ">IGHJ1-MOCK*01\n" + "TACTACTACGGTATGGACGTCTGGGGCCAAGGGACCACGGTCACCGTCTCCTCAG\n" + ">IGHJ2-MOCK*01\n" + "ACTACTGGTACTTCGATCTCTGGGGCCGTGGCACCCTGGTCACTGTCTCCTCAG\n" +) + + +def test_tool_module_imports_and_carries_imgt_maps() -> None: + """The module loads and exposes the IMGT layout tables.""" + assert tool.IMGT_BASE.startswith("https://") + assert "Homo_sapiens" in tool.SPECIES_MAP + assert tool.SPECIES_MAP["Homo_sapiens"][0] is Species.HUMAN + assert set(tool.LOCUS_DEFS) == {"IGH", "IGK", "IGL", "TRA", "TRB", "TRG", "TRD"} + + +def test_build_cartridge_vdj_produces_valid_dataconfig() -> None: + """A VDJ (heavy) build yields a structural DataConfig with the + expected catalogue and identity.""" + cfg = tool.build_cartridge( + Species.HUMAN, ChainType.BCR_HEAVY, "IGH", "MOCK", + v_path=_V_FASTA, j_path=_J_FASTA, d_path=_D_FASTA, + ) + assert isinstance(cfg, ga.DataConfig) + assert cfg.name == "MOCK_IGH_IMGT" + assert len(cfg.v_alleles) == 2 + assert len(cfg.d_alleles) == 2 + assert len(cfg.j_alleles) == 2 + assert cfg.metadata is not None + assert cfg.metadata.has_d is True + # Structural only: no data-derived typed empirical models. + assert cfg.reference_models is None + + +def test_build_cartridge_vj_light_chain_has_no_d() -> None: + """A VJ (light) build passes no D FASTA and yields a D-less config.""" + cfg = tool.build_cartridge( + Species.HUMAN, ChainType.BCR_LIGHT_KAPPA, "IGK", "MOCK", + v_path=_V_FASTA, j_path=_J_FASTA, d_path=None, + ) + assert isinstance(cfg, ga.DataConfig) + assert cfg.name == "MOCK_IGK_IMGT" + assert cfg.metadata.has_d is False + assert not cfg.d_alleles + + +def test_build_cartridge_missing_d_for_d_chain_returns_none() -> None: + """A D-bearing chain with no D FASTA is skipped (returns None), not + an exception — matches the tool's per-locus fault tolerance.""" + cfg = tool.build_cartridge( + Species.HUMAN, ChainType.BCR_HEAVY, "IGH", "MOCK", + v_path=_V_FASTA, j_path=_J_FASTA, d_path=None, + ) + assert cfg is None + + +def test_built_cartridge_compiles_and_runs_through_experiment() -> None: + """The built structural cartridge is drop-in for Experiment.on(cfg).""" + cfg = tool.build_cartridge( + Species.HUMAN, ChainType.BCR_HEAVY, "IGH", "MOCK", + v_path=_V_FASTA, j_path=_J_FASTA, d_path=_D_FASTA, + ) + result = ( + ga.Experiment.on(cfg) + .recombine() + .allow_curatable_refdata() + .run_records(n=3, seed=1) + ) + assert len(result) == 3 + + +def test_cli_requires_output_dir() -> None: + """--output-dir is mandatory; the parser rejects its absence.""" + with pytest.raises(SystemExit): + tool._parser().parse_args(["--dry-run"]) + ns = tool._parser().parse_args(["--output-dir", "/tmp/x", "--dry-run"]) + assert ns.output_dir == "/tmp/x" + assert ns.dry_run is True diff --git a/tests/test_docs_mkdocs_migration_contract.py b/tests/test_docs_mkdocs_migration_contract.py index bcb6ac7..4c0423c 100644 --- a/tests/test_docs_mkdocs_migration_contract.py +++ b/tests/test_docs_mkdocs_migration_contract.py @@ -48,7 +48,7 @@ _DEPLOY_WORKFLOW = _REPO_ROOT / ".github" / "workflows" / "deploy-docs.yml" # Tests in this file pin the migration plan's structure. The plan is a -# private AI-session artifact (see docs/.gitignore); CI checkouts skip +# private planning artifact (kept out of git); CI checkouts skip # the whole file rather than report dozens of irrelevant failures. pytestmark = pytest.mark.skipif( not _MIGRATION_PLAN.is_file(), diff --git a/tests/test_docs_website_contract.py b/tests/test_docs_website_contract.py index f6c9712..464e65f 100644 --- a/tests/test_docs_website_contract.py +++ b/tests/test_docs_website_contract.py @@ -105,7 +105,7 @@ def test_pin_post_cutover_deploy_docs_workflow_builds_mkdocs_and_uploads_site() def test_pin_scaffold_docs_dir_carries_audit_design_md_files() -> None: """`audit-docs/` carries ≥ 35 markdown files (audits + designs + hubs). Originally housed in `docs/` (which is now - private — Claude session artefacts + impl-time notes); moved + private — design + impl-time notes); moved to `audit-docs/` so contract tests can reference them in CI. The pin is intentionally loose so adding new audits doesn't fail this test.""" @@ -162,33 +162,6 @@ def test_pin_scaffold_old_docs_dir_exists_as_abandoned_earlier_attempt() -> None ) -def test_pin_scaffold_docs_superpowers_subdir_holds_session_artifacts() -> None: - """`docs/superpowers/plans/` carries Claude-session - planning artefacts, NOT user-facing documentation. - Pinned defensively so a future "is this a doc?" sweep - treats it correctly.""" - superpowers = _DOCS / "superpowers" - if not superpowers.is_dir(): - # If somebody moves it, that's fine — pin only the docs - # subdirs that exist today. - return - plans = superpowers / "plans" - if plans.is_dir(): - plans_files = sorted(plans.glob("*.md")) - # The presence of a date-prefixed planning markdown is - # the canonical evidence (e.g. 2026-05-18-mcp-redesign-v2.md). - date_prefixed = [ - f for f in plans_files - if re.match(r"^\d{4}-\d{2}-\d{2}-", f.name) - ] - assert date_prefixed, ( - "docs/superpowers/plans/ has no date-prefixed planning " - "files — the session-artefact discipline broke; this " - "directory's contents may have been recategorised as " - "docs without an audit-doc update" - ) - - def test_pin_scaffold_docs_build_subdir_holds_wheel_artefacts_not_docs() -> None: """`docs/build/` holds Python wheel build artefacts. Skip when docs/ is absent (CI checkouts where docs/ is gitignored).""" diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 86eab28..fab52fe 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,6 +1,6 @@ """Tests for the redesigned GenAIRR MCP server. -Layered per the spec at docs/superpowers/specs/2026-05-18-mcp-redesign-v2-design.md: +Layered by tier: - Tier 1: happy path per tool (14 tests) - Tier 2: error envelope per error code (~6 tests) - Tier 3: one end-to-end agent-style chain (1 test) diff --git a/tests/test_np_base_model_estimation_contract.py b/tests/test_np_base_model_estimation_contract.py index 85ba036..6914d33 100644 --- a/tests/test_np_base_model_estimation_contract.py +++ b/tests/test_np_base_model_estimation_contract.py @@ -95,11 +95,11 @@ def test_pin_scaffold_unclaimed_np_string_walks_structural_region_only() -> None def test_pin_scaffold_result_column_order_includes_np1_and_np2_strings() -> None: - """`result.py`'s canonical column order declares both + """`_result_export.py`'s canonical column order declares both `np1` and `np2` string fields. The estimator inherits the column names directly.""" src = ( - _REPO_ROOT / "src" / "GenAIRR" / "result.py" + _REPO_ROOT / "src" / "GenAIRR" / "_result_export.py" ).read_text(encoding="utf-8") assert '"np1"' in src assert '"np2"' in src diff --git a/tests/test_np_length_estimation_contract.py b/tests/test_np_length_estimation_contract.py index 43d3e8c..cd4f917 100644 --- a/tests/test_np_length_estimation_contract.py +++ b/tests/test_np_length_estimation_contract.py @@ -105,12 +105,12 @@ def test_pin_scaffold_python_airr_projection_emits_np1_length_np2_length() -> No def test_pin_scaffold_result_column_order_includes_np_length_fields() -> None: - """`result.py`'s canonical column order declares both + """`_result_export.py`'s canonical column order declares both length fields. The estimator's output is consumed by `Experiment.on(cfg).recombine().run_records()` which will emit AIRR rows carrying these columns.""" src = ( - _REPO_ROOT / "src" / "GenAIRR" / "result.py" + _REPO_ROOT / "src" / "GenAIRR" / "_result_export.py" ).read_text(encoding="utf-8") assert '"np1_length"' in src assert '"np2_length"' in src diff --git a/tests/test_p_nucleotide_contract.py b/tests/test_p_nucleotide_contract.py index e2855b5..0ba029e 100644 --- a/tests/test_p_nucleotide_contract.py +++ b/tests/test_p_nucleotide_contract.py @@ -47,9 +47,8 @@ _GENERATE_NP_SAMPLING = ( _REPO_ROOT / "engine_rs" / "src" / "passes" / "generate_np" / "sampling.rs" ) -_COMPILE_PY = _REPO_ROOT / "src" / "GenAIRR" / "_compile.py" +_LOWERING_PY = _REPO_ROOT / "src" / "GenAIRR" / "_lowering.py" _DATACONFIG_PY = _REPO_ROOT / "src" / "GenAIRR" / "dataconfig" / "data_config.py" -_MCP_HELPERS_PY = _REPO_ROOT / "src" / "GenAIRR" / "utilities" / "mcp_helpers.py" _AIRR_RECORD_RS = _REPO_ROOT / "engine_rs" / "src" / "airr_record" / "record.rs" _VALIDATE_RS = _REPO_ROOT / "engine_rs" / "src" / "airr_record" / "validate.rs" @@ -161,7 +160,7 @@ def test_pin_present_pipeline_order_has_p_addition_at_audited_positions() -> Non sequence must be read under the post-inversion orientation. See `docs/p_nucleotide_design.md` §9.3 for the corrected ordering.""" - src = _COMPILE_PY.read_text(encoding="utf-8") + src = _LOWERING_PY.read_text(encoding="utf-8") assert 'push_p_addition("V_3"' in src assert 'push_p_addition("D_5"' in src assert 'push_p_addition("D_3"' in src @@ -196,7 +195,7 @@ def test_pin_scaffold_invert_d_commits_before_assemble_d() -> None: `PAdditionPass(end=D_5)` inserted between them reads the post-inversion orientation of D's effective_seq. Pin the current ordering.""" - src = _COMPILE_PY.read_text(encoding="utf-8") + src = _LOWERING_PY.read_text(encoding="utf-8") # In the VDJ branch, push_invert_d appears before # push_assemble("D"). Use rough textual ordering as the pin. invert_idx = src.find("push_invert_d(") @@ -212,31 +211,6 @@ def test_pin_scaffold_invert_d_commits_before_assemble_d() -> None: # ────────────────────────────────────────────────────────────────── -def test_pin_scaffold_mcp_p_nucleotides_endpoint_is_read_only() -> None: - """Audit §Pre-flight / §7.3 — the MCP helper's - `p_nucleotides` inspection section reads - `getattr(dc, "p_nucleotide_length_probs", {})` as a - diagnostic surface; it does NOT feed simulation. Pinned - so a refactor that wires the helper into the sampler - surfaces immediately.""" - src = _MCP_HELPERS_PY.read_text(encoding="utf-8") - assert 'section == "p_nucleotides"' in src - assert 'getattr(dc, "p_nucleotide_length_probs"' in src - # The endpoint is a read-only data surface — no - # mutation / plan push / engine call. - for forbidden in ( - "push_p_addition", - "plan.push", - "Experiment.on", - ): - # The text is allowed to contain `plan.push` only if it's - # part of a comment/docstring; we keep this check loose to - # avoid false positives on unrelated tooling code. - # The load-bearing assertion is that getattr is the - # source — which is. - pass - - # ────────────────────────────────────────────────────────────────── # 8. Present — p_nucleotide_length_probs exists with default # ────────────────────────────────────────────────────────────────── @@ -316,8 +290,6 @@ def test_pin_present_p_nucleotide_length_probs_has_no_simulator_consumer() -> No # Allowed consumers: # - the dataclass declaration itself # (`src/GenAIRR/dataconfig/data_config.py`) - # - the MCP read-only diagnostic endpoint - # (`src/GenAIRR/utilities/mcp_helpers.py`) result = subprocess.run( [ "grep", @@ -339,7 +311,11 @@ def test_pin_present_p_nucleotide_length_probs_has_no_simulator_consumer() -> No } allowed = { "src/GenAIRR/dataconfig/data_config.py", - "src/GenAIRR/utilities/mcp_helpers.py", + # Manifest reporting surface extracted verbatim from + # data_config.py (behavior-preserving hygiene split); names + # the legacy field for REPORTING only, not as a simulator + # consumer — same category as data_config.py. + "src/GenAIRR/dataconfig/_manifest.py", # Post-slice — the typed-plane resolver explicitly # documents the no-auto-lift boundary in docstring + # comments referencing the legacy field name. Those diff --git a/tests/test_paired_end_schema.py b/tests/test_paired_end_schema.py index 32cc7a8..1cc46e0 100644 --- a/tests/test_paired_end_schema.py +++ b/tests/test_paired_end_schema.py @@ -282,7 +282,8 @@ def test_slice_b_geometry_check_helpers_landed_in_engine_source() -> None: / "engine_rs" / "src" / "airr_record" - / "validate.rs" + / "validate" + / "paired_end.rs" ).read_text(encoding="utf-8") # The geometry check function name + the four issue variants # it can produce. Each must appear in the dispatch path. @@ -294,8 +295,8 @@ def test_slice_b_geometry_check_helpers_landed_in_engine_source() -> None: "RecordValidationIssue::ReadLayoutMismatch", ): assert required in validate_src, ( - f"validate.rs no longer references {required!r}; Slice " - f"B's geometry dispatch has drifted." + f"validate/paired_end.rs no longer references {required!r}; " + f"Slice B's geometry dispatch has drifted." ) diff --git a/tests/test_plan_signature_completeness_contract.py b/tests/test_plan_signature_completeness_contract.py index 1531565..a88181d 100644 --- a/tests/test_plan_signature_completeness_contract.py +++ b/tests/test_plan_signature_completeness_contract.py @@ -624,11 +624,11 @@ def test_pin_legacy_address_schema_version_is_one() -> None: def test_pin_legacy_frozen_address_spellings_test_exists() -> None: """The `frozen_address_spellings_for_choice_address_schema_v1` - Rust unit test in `address.rs` pins one representative + Rust unit test in `address/tests.rs` pins one representative of every typed variant. The Python contract pin here asserts the Rust test fixture is present (the Rust test itself is the actual enforcement).""" - src = (_REPO_ROOT / "engine_rs" / "src" / "address.rs").read_text(encoding="utf-8") + src = (_REPO_ROOT / "engine_rs" / "src" / "address" / "tests.rs").read_text(encoding="utf-8") assert "fn frozen_address_spellings_for_choice_address_schema_v1" in src diff --git a/tests/test_reference_cartridge_authoring_contract.py b/tests/test_reference_cartridge_authoring_contract.py index 447ef88..da881cf 100644 --- a/tests/test_reference_cartridge_authoring_contract.py +++ b/tests/test_reference_cartridge_authoring_contract.py @@ -244,34 +244,23 @@ def test_pin_present_build_report_docstring_now_references_new_builder() -> None importlib.import_module("GenAIRR.dataconfig.make.random") -def test_pin_present_private_build_script_now_raises_explicit_legacy_error() -> None: - """Post-slice — the `.private/scripts/build_imgt_configs.py` - script raises an explicit `NotImplementedError` at module- - load time pointing at the new builder, rather than failing - with a deep `ModuleNotFoundError` deep inside an obsolete - import. Flipped from the prior dead-import present-pin.""" - script = _REPO_ROOT / ".private" / "scripts" / "build_imgt_configs.py" - if not script.exists(): - pytest.skip(".private/scripts/build_imgt_configs.py absent") - src = script.read_text(encoding="utf-8") - # The dead import is no longer the load-time failure point. - # The explicit raise comes first. - assert "raise NotImplementedError(" in src, ( - "private build script no longer raises explicit " - "NotImplementedError at module load — verify the dead-" - "reference cleanup landed and the script's failure mode " - "is still self-documenting" - ) +def test_pin_present_imgt_build_tool_uses_reference_cartridge_builder() -> None: + """The IMGT cartridge build tool is a first-class, tracked + maintainer tool (`tools/build_imgt_configs.py`) built on the + current `ReferenceCartridgeBuilder` — not the removed + `RandomDataConfigBuilder`. (Flipped from the prior pin on the + legacy private stub, now that the real tool exists.)""" + tool = _REPO_ROOT / "tools" / "build_imgt_configs.py" + assert tool.exists(), "tools/build_imgt_configs.py is missing" + src = tool.read_text(encoding="utf-8") assert "ReferenceCartridgeBuilder" in src, ( - "private build script does not reference the new builder — " - "the porting hint regressed" + "IMGT build tool no longer uses ReferenceCartridgeBuilder" + ) + assert "RandomDataConfigBuilder" not in src, ( + "IMGT build tool still references the removed RandomDataConfigBuilder" ) - raise_pos = src.find("raise NotImplementedError(") - legacy_import_pos = src.find("from GenAIRR.dataconfig.make.random import") - assert legacy_import_pos == -1 or raise_pos < legacy_import_pos, ( - "the dead import would fire BEFORE the explicit raise — " - "the legacy guard needs to come first or the dead import " - "needs to be moved into the unreachable body" + assert "from GenAIRR.dataconfig.make" not in src, ( + "IMGT build tool still imports the removed dataconfig.make namespace" ) diff --git a/tests/test_reference_cartridge_authoring_implementation.py b/tests/test_reference_cartridge_authoring_implementation.py index 86606cf..2a9f8f5 100644 --- a/tests/test_reference_cartridge_authoring_implementation.py +++ b/tests/test_reference_cartridge_authoring_implementation.py @@ -398,14 +398,15 @@ def test_manual_dataconfig_construction_remains_supported() -> None: # ────────────────────────────────────────────────────────────────── -def test_dead_reference_cleanup_data_config_docstring_and_private_script() -> None: - """Integration check that the lockstep cleanup landed: +def test_dead_reference_cleanup_data_config_docstring_and_build_tool() -> None: + """Integration check that the RandomDataConfigBuilder retirement + landed in lockstep: - ``DataConfig.build_report`` docstring names the new builder. - - ``.private/scripts/build_imgt_configs.py`` raises - ``NotImplementedError`` at module-load time rather - than ``ModuleNotFoundError`` on the dead import. + - the IMGT cartridge build tool (`tools/build_imgt_configs.py`) + uses ``ReferenceCartridgeBuilder`` and no longer references + the removed ``RandomDataConfigBuilder``. The contract file's `pin_present_*` pins are the authoritative source; this test is the integration @@ -420,11 +421,11 @@ def test_dead_reference_cleanup_data_config_docstring_and_private_script() -> No assert "RandomDataConfigBuilder" not in dc_src assert "ReferenceCartridgeBuilder" in dc_src - script = repo / ".private" / "scripts" / "build_imgt_configs.py" - if script.exists(): - src = script.read_text(encoding="utf-8") - assert "raise NotImplementedError(" in src - assert "ReferenceCartridgeBuilder" in src + tool = repo / "tools" / "build_imgt_configs.py" + assert tool.exists(), "tools/build_imgt_configs.py is missing" + tool_src = tool.read_text(encoding="utf-8") + assert "ReferenceCartridgeBuilder" in tool_src + assert "RandomDataConfigBuilder" not in tool_src # ────────────────────────────────────────────────────────────────── diff --git a/tests/test_trim_distribution_estimation_contract.py b/tests/test_trim_distribution_estimation_contract.py index 69b4b2a..b3c7cbe 100644 --- a/tests/test_trim_distribution_estimation_contract.py +++ b/tests/test_trim_distribution_estimation_contract.py @@ -336,7 +336,7 @@ def test_pin_scaffold_trim_pass_count_matches_segment_end_pairs() -> None: so a fifth `TrimPass` (e.g. `(V,Five)`) without the matching plane key surfaces here.""" src = ( - _REPO_ROOT / "src" / "GenAIRR" / "_compile.py" + _REPO_ROOT / "src" / "GenAIRR" / "_lowering.py" ).read_text(encoding="utf-8") # The lowering site dispatches via .push_trim("V","3"), etc. assert 'push_trim("V", "3"' in src @@ -408,7 +408,7 @@ def test_pin_scaffold_end_loss_dsl_wires_to_end_loss_pass_not_trim() -> None: NOT `TrimPass`. The estimator MUST NOT consume DSL state from the end-loss surface.""" src = ( - _REPO_ROOT / "src" / "GenAIRR" / "experiment.py" + _REPO_ROOT / "src" / "GenAIRR" / "_experiment" / "corruption.py" ).read_text(encoding="utf-8") # Both DSL methods exist. assert "def end_loss_5prime(" in src diff --git a/tests/test_v_region_substructure_contract.py b/tests/test_v_region_substructure_contract.py index 1e99b64..42cead9 100644 --- a/tests/test_v_region_substructure_contract.py +++ b/tests/test_v_region_substructure_contract.py @@ -556,7 +556,6 @@ def test_pin_present_imgt_regions_consumed_by_bridge_and_mcp() -> None: consumers = sorted(set(result.stdout.strip().splitlines())) expected = { "src/GenAIRR/_refdata_resolver.py", - "src/GenAIRR/utilities/mcp_helpers.py", # `ReferenceCartridgeBuilder.infer_v_subregions` uses the # same derivation helper at authoring time — same # boundary the bridge resolver uses at load time. See diff --git a/tests/test_v_subregion_mutation_counters_contract.py b/tests/test_v_subregion_mutation_counters_contract.py index 6e27bde..425ffbe 100644 --- a/tests/test_v_subregion_mutation_counters_contract.py +++ b/tests/test_v_subregion_mutation_counters_contract.py @@ -203,7 +203,7 @@ def test_pin_scaffold_validator_recomputes_independently() -> None: The new V-subregion validator extension inherits the same discipline.""" src = ( - _REPO_ROOT / "engine_rs" / "src" / "airr_record" / "validate.rs" + _REPO_ROOT / "engine_rs" / "src" / "airr_record" / "validate" / "counters.rs" ).read_text(encoding="utf-8") # The four existing per-segment mismatch issue kinds. for kind in ( diff --git a/tools/README.md b/tools/README.md new file mode 100644 index 0000000..746118e --- /dev/null +++ b/tools/README.md @@ -0,0 +1,34 @@ +# tools/ + +Maintainer tooling for GenAIRR. Not part of the shipped `GenAIRR` wheel — +these scripts are run from a source checkout. + +## `build_imgt_configs.py` + +Fetches IMGT V-QUEST germline FASTA and builds structural GenAIRR +cartridges (`__IMGT.pkl`) via +`GenAIRR.ReferenceCartridgeBuilder`. + +```bash +# Build one species into a scratch dir (never touches the bundled set): +python tools/build_imgt_configs.py --output-dir ./built_cartridges --species Mus_musculus + +# Probe locus availability only: +python tools/build_imgt_configs.py --output-dir ./built_cartridges --dry-run +``` + +**Structural output.** Built cartridges carry the germline V/D/J +catalogue plus IMGT V-subregion (FWR/CDR) annotations, but **no** +data-derived empirical distributions (trim / NP length / NP base model / +gene usage) — those cannot be inferred from germline FASTA. At simulation +time the engine uses uniform defaults for those parameters. Fit real +distributions into the returned `DataConfig`, or estimate them with the +builder's `estimate_*` methods from your own AIRR data, if you need +empirically-grounded parameters. Only the bundled human IGH/IGK/IGL and +TCRB cartridges ship with real data-derived distributions. + +`--output-dir` is **required**: a freshly-built structural cartridge +differs from the shipped `src/GenAIRR/data/builtin_dataconfigs/` set, so +the tool never overwrites it by default. Point `--output-dir` there +explicitly only when you deliberately intend to regenerate the bundled +cartridges (and expect the golden baselines to move). diff --git a/tools/build_imgt_configs.py b/tools/build_imgt_configs.py new file mode 100644 index 0000000..29f2e02 --- /dev/null +++ b/tools/build_imgt_configs.py @@ -0,0 +1,371 @@ +#!/usr/bin/env python3 +"""Fetch IMGT V-QUEST germline FASTA and build GenAIRR cartridges. + +Maintainer tool (not part of the shipped ``GenAIRR`` wheel). It downloads +the IMGT V-QUEST reference directory FASTA for each species/locus and +builds a structural :class:`GenAIRR.DataConfig` cartridge per +species+chain via :class:`GenAIRR.ReferenceCartridgeBuilder`. + +**Structural output.** Cartridges carry the germline V/D/J allele +catalogue plus IMGT-derived V-subregion (FWR/CDR) annotations. They do +**not** carry data-derived empirical distributions (trim lengths, NP +lengths, NP base model, gene usage): those are not inferable from +germline FASTA alone. At simulation time the engine falls back to its +uniform defaults for those parameters — i.e. a uniform representation of +the scenarios a repertoire can express. Fit real distributions into the +returned ``DataConfig`` (or estimate them with the builder's +``estimate_*`` methods from your own AIRR data) if you need +empirically-grounded parameters. + +Only the bundled ``HUMAN_IGH`` / ``HUMAN_IGK`` / ``HUMAN_IGL`` / +``HUMAN_TCRB`` cartridges ship with real data-derived distributions; +everything this tool builds is uniform-by-construction. + +Usage:: + + # Build one species into a scratch dir (never touches the bundled set): + python tools/build_imgt_configs.py --output-dir ./built_cartridges \\ + --species Mus_musculus + + # Probe availability only, no build: + python tools/build_imgt_configs.py --output-dir ./built_cartridges --dry-run + +``--output-dir`` is required: writing a freshly-built (structural) +cartridge over the shipped ``src/GenAIRR/data/builtin_dataconfigs/`` set +would change simulation behaviour and break the golden tests, so it is +never the default. Point ``--output-dir`` there explicitly only when you +intend to regenerate the bundled set. +""" +from __future__ import annotations + +import argparse +import logging +import os +import pickle +import sys +import time +import urllib.error +import urllib.request + +from GenAIRR import DataConfig, ReferenceCartridgeBuilder +from GenAIRR.dataconfig.enums import ChainType, Species + +logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(message)s") +logger = logging.getLogger("build_imgt_configs") + +# ───────────────────────────────────────────────────────────── +# IMGT V-QUEST reference directory layout +# ───────────────────────────────────────────────────────────── + +IMGT_BASE = ( + "https://www.imgt.org/download/V-QUEST/IMGT_V-QUEST_reference_directory" +) + +# IMGT V-QUEST species directories carried by GenAIRR's bundled set. +IMGT_SPECIES = [ + "Aotus_nancymaae", + "Bos_taurus", + "Camelus_dromedarius", + "Canis_lupus_familiaris", + "Capra_hircus", + "Danio_rerio", + "Equus_caballus", + "Felis_catus", + "Gallus_gallus", + "Gorilla_gorilla_gorilla", + "Homo_sapiens", + "Macaca_fascicularis", + "Macaca_mulatta", + "Mus_musculus", + "Mus_musculus_C57BL6J", + "Mustela_putorius_furo", + "Oncorhynchus_mykiss", + "Ornithorhynchus_anatinus", + "Oryctolagus_cuniculus", + "Ovis_aries", + "Rattus_norvegicus", + "Salmo_salar", + "Sus_scrofa", + "Vicugna_pacos", +] + +# IMGT directory name → (Species enum, short label used in the cartridge name). +SPECIES_MAP = { + "Homo_sapiens": (Species.HUMAN, "HUMAN"), + "Mus_musculus": (Species.MOUSE, "MOUSE"), + "Mus_musculus_C57BL6J": (Species.MOUSE_C57BL6J, "MOUSE_C57BL6J"), + "Rattus_norvegicus": (Species.RAT, "RAT"), + "Oryctolagus_cuniculus": (Species.RABBIT, "RABBIT"), + "Macaca_mulatta": (Species.RHESUS_MACAQUE, "RHESUS"), + "Macaca_fascicularis": (Species.CYNOMOLGUS_MACAQUE, "CYNOMOLGUS"), + "Gorilla_gorilla_gorilla": (Species.GORILLA, "GORILLA"), + "Aotus_nancymaae": (Species.OWL_MONKEY, "OWL_MONKEY"), + "Bos_taurus": (Species.COW, "COW"), + "Ovis_aries": (Species.SHEEP, "SHEEP"), + "Capra_hircus": (Species.GOAT, "GOAT"), + "Sus_scrofa": (Species.PIG, "PIG"), + "Equus_caballus": (Species.HORSE, "HORSE"), + "Canis_lupus_familiaris": (Species.DOG, "DOG"), + "Felis_catus": (Species.CAT, "CAT"), + "Camelus_dromedarius": (Species.DROMEDARY_CAMEL, "DROMEDARY"), + "Vicugna_pacos": (Species.ALPACA, "ALPACA"), + "Mustela_putorius_furo": (Species.FERRET, "FERRET"), + "Ornithorhynchus_anatinus": (Species.PLATYPUS, "PLATYPUS"), + "Gallus_gallus": (Species.CHICKEN, "CHICKEN"), + "Danio_rerio": (Species.ZEBRAFISH, "ZEBRAFISH"), + "Oncorhynchus_mykiss": (Species.TROUT, "TROUT"), + "Salmo_salar": (Species.SALMON, "SALMON"), +} + +# IMGT locus → (ChainType, chain label, V file, D file or None, J file). +LOCUS_DEFS = { + "IGH": (ChainType.BCR_HEAVY, "IGH", "IGHV.fasta", "IGHD.fasta", "IGHJ.fasta"), + "IGK": (ChainType.BCR_LIGHT_KAPPA, "IGK", "IGKV.fasta", None, "IGKJ.fasta"), + "IGL": (ChainType.BCR_LIGHT_LAMBDA, "IGL", "IGLV.fasta", None, "IGLJ.fasta"), + "TRA": (ChainType.TCR_ALPHA, "TCRA", "TRAV.fasta", None, "TRAJ.fasta"), + "TRB": (ChainType.TCR_BETA, "TCRB", "TRBV.fasta", "TRBD.fasta", "TRBJ.fasta"), + "TRG": (ChainType.TCR_GAMMA, "TCRG", "TRGV.fasta", None, "TRGJ.fasta"), + "TRD": (ChainType.TCR_DELTA, "TCRD", "TRDV.fasta", "TRDD.fasta", "TRDJ.fasta"), +} + +# ───────────────────────────────────────────────────────────── +# Download / probe helpers +# ───────────────────────────────────────────────────────────── + + +def download_file(url: str, dest: str, retries: int = 3, delay: float = 1.0) -> bool: + """Download ``url`` to ``dest`` with retries. Returns ``True`` on success, + ``False`` on a 404 (locus absent for the species) or exhausted retries.""" + for attempt in range(retries): + try: + req = urllib.request.Request( + url, headers={"User-Agent": "GenAIRR-DataConfig-Builder/1.0"} + ) + with urllib.request.urlopen(req, timeout=30) as resp: + data = resp.read() + os.makedirs(os.path.dirname(dest), exist_ok=True) + with open(dest, "wb") as fh: + fh.write(data) + return True + except urllib.error.HTTPError as exc: + if exc.code == 404: + return False + logger.warning("HTTP %d for %s (attempt %d/%d)", exc.code, url, attempt + 1, retries) + except Exception as exc: # network hiccup — retry + logger.warning("Error downloading %s: %s (attempt %d/%d)", url, exc, attempt + 1, retries) + time.sleep(delay * (attempt + 1)) + return False + + +def probe_loci(imgt_species: str, cache_dir: str) -> dict: + """Probe which loci exist for a species and cache their FASTA locally. + + Returns ``{locus_name: {"V": path, "J": path, "D": path?}}`` for every + locus with (at least) a V and J file. A locus with no V or no J is + skipped entirely; a D-bearing locus without a D file keeps V/J only. + """ + results: dict = {} + for locus_name, (_chain, _label, v_file, d_file, j_file) in LOCUS_DEFS.items(): + receptor = "IG" if locus_name.startswith("IG") else "TR" + base_url = f"{IMGT_BASE}/{imgt_species}/{receptor}" + species_dir = os.path.join(cache_dir, imgt_species, receptor) + + segments: dict = {} + + v_dest = os.path.join(species_dir, v_file) + if os.path.exists(v_dest) or download_file(f"{base_url}/{v_file}", v_dest): + segments["V"] = v_dest + else: + continue # no V → skip the locus + + j_dest = os.path.join(species_dir, j_file) + if os.path.exists(j_dest) or download_file(f"{base_url}/{j_file}", j_dest): + segments["J"] = j_dest + else: + continue # no J → skip + + if d_file is not None: + d_dest = os.path.join(species_dir, d_file) + if os.path.exists(d_dest) or download_file(f"{base_url}/{d_file}", d_dest): + segments["D"] = d_dest + + results[locus_name] = segments + return results + + +def fasta_has_sequences(path: str) -> bool: + """True if ``path`` exists and contains at least one FASTA record.""" + if not os.path.exists(path): + return False + with open(path) as fh: + return any(line.startswith(">") for line in fh) + + +# ───────────────────────────────────────────────────────────── +# Build (ported to ReferenceCartridgeBuilder) +# ───────────────────────────────────────────────────────────── + + +def build_cartridge( + species_enum: Species, + chain_type: ChainType, + chain_label: str, + species_label: str, + v_path: str, + j_path: str, + d_path: str | None = None, + reference_set: str = "IMGT", +) -> DataConfig | None: + """Build one structural cartridge from IMGT FASTA. + + Runs the ``ReferenceCartridgeBuilder`` chain + (``from_fasta → infer_identity → infer_v_subregions → build``) and + returns the resulting :class:`GenAIRR.DataConfig`, or ``None`` if the + build fails (e.g. a D-bearing locus whose D FASTA was unavailable). + """ + if chain_type.has_d and d_path is None: + logger.warning( + "%s_%s: chain has a D segment but no D FASTA — skipping", + species_label, chain_label, + ) + return None + try: + return ( + ReferenceCartridgeBuilder.from_fasta( + v_fasta=v_path, + j_fasta=j_path, + d_fasta=d_path, + chain_type=chain_type, + ) + .infer_identity( + species=species_enum, + locus=chain_label, + reference_set=reference_set, + name=f"{species_label}_{chain_label}_{reference_set}", + source="IMGT", + ) + .infer_v_subregions() + .build() + ) + except Exception as exc: # malformed FASTA, anchor issues, etc. + logger.error("Failed to build %s_%s: %s", species_label, chain_label, exc) + return None + + +# ───────────────────────────────────────────────────────────── +# CLI +# ───────────────────────────────────────────────────────────── + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Build structural GenAIRR cartridges from IMGT V-QUEST FASTA.", + ) + parser.add_argument( + "--output-dir", required=True, + help="Directory to write built __IMGT.pkl cartridges. " + "REQUIRED (never defaults to the shipped builtin_dataconfigs set).", + ) + parser.add_argument( + "--cache-dir", default="/tmp/imgt_cache", + help="Directory to cache downloaded IMGT FASTA (default: /tmp/imgt_cache).", + ) + parser.add_argument( + "--species", nargs="*", default=None, + help="Only process these IMGT species directory names (default: all).", + ) + parser.add_argument( + "--dry-run", action="store_true", + help="Probe and report available loci; do not build or write anything.", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + + os.makedirs(args.output_dir, exist_ok=True) + os.makedirs(args.cache_dir, exist_ok=True) + + species_list = args.species if args.species else IMGT_SPECIES + built: list[str] = [] + skipped: list[tuple[str, str]] = [] + failed: list[str] = [] + + for imgt_species in species_list: + if imgt_species not in SPECIES_MAP: + logger.warning("No Species enum mapping for %s — skipping", imgt_species) + skipped.append((imgt_species, "no enum mapping")) + continue + + species_enum, species_label = SPECIES_MAP[imgt_species] + logger.info("=== %s (%s) ===", imgt_species, species_label) + + loci = probe_loci(imgt_species, args.cache_dir) + if not loci: + logger.info(" no loci found") + skipped.append((imgt_species, "no loci")) + continue + + for locus_name, segments in sorted(loci.items()): + chain_type, chain_label, _, _, _ = LOCUS_DEFS[locus_name] + config_name = f"{species_label}_{chain_label}_IMGT" + + if not fasta_has_sequences(segments["V"]): + skipped.append((config_name, "empty V")) + continue + if not fasta_has_sequences(segments["J"]): + skipped.append((config_name, "empty J")) + continue + + d_path = segments.get("D") + logger.info( + " %s: V=%s J=%s D=%s", locus_name, + os.path.basename(segments["V"]), + os.path.basename(segments["J"]), + os.path.basename(d_path) if d_path else "none", + ) + if args.dry_run: + built.append(config_name) + continue + + config = build_cartridge( + species_enum, chain_type, chain_label, species_label, + v_path=segments["V"], j_path=segments["J"], d_path=d_path, + ) + if config is None: + failed.append(config_name) + continue + + n_v = sum(len(a) for a in config.v_alleles.values()) if config.v_alleles else 0 + n_d = sum(len(a) for a in config.d_alleles.values()) if config.d_alleles else 0 + n_j = sum(len(a) for a in config.j_alleles.values()) if config.j_alleles else 0 + + pkl_path = os.path.join(args.output_dir, f"{config_name}.pkl") + with open(pkl_path, "wb") as fh: + pickle.dump(config, fh, protocol=pickle.HIGHEST_PROTOCOL) + logger.info( + " built %s: %d V, %d D, %d J alleles → %s (%.1f KB)", + config_name, n_v, n_d, n_j, pkl_path, + os.path.getsize(pkl_path) / 1024, + ) + built.append(config_name) + + print("\n" + "=" * 60) + print(f"BUILT: {len(built)}") + for name in sorted(built): + print(f" {name}") + if skipped: + print(f"\nSKIPPED: {len(skipped)}") + for name, reason in skipped: + print(f" {name}: {reason}") + if failed: + print(f"\nFAILED: {len(failed)}") + for name in failed: + print(f" {name}") + print("=" * 60) + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main())