diff --git a/.praxia/docs/specs/261011_noise-level-axis.md b/.praxia/docs/specs/261011_noise-level-axis.md new file mode 100644 index 00000000..ee626358 --- /dev/null +++ b/.praxia/docs/specs/261011_noise-level-axis.md @@ -0,0 +1,70 @@ +--- +title: Backbone noise-level axis for design and score +description: Every family honours a multi-level backbone_noise sweep in both sample and score, with one shared key scheme; replaces the #214 refusal +task_id: 261011_noise-axis +status: active +--- + +# Backbone noise-level axis for design and score (261011) + +User direction (261011): "we need to have a noise-level axis available for both design and score". +This supersedes PR #214's decision to *refuse* nonzero `backbone_noise` on plain score (debt #2509). + +## Current state (recon 261011, spot-checked) + +| Path | Noise applied? | Axis in output? | +|---|---|---| +| stock `runner.sample` Path A | yes, per level | yes: `(S, N, D, T, L)` (`host/kernel_dispatch.py:661`) | +| stock `runner.sample` Path B (encoding fusion) | yes, fused over D | no (by design) | +| stock `runner.score` plain, encode-once | no (`scoring/score.py:279` hard-codes 0.0) | no, `(S, C)` | +| stock `runner.score` plain, per-candidate | no (`score_fn` not passed noise) | no | +| stock `_score_fused_multistate` | no | no, `(1, C)` | +| stock averaged (`average_node_features`) | yes, fused over D | no (by design); D>=2 crashes (#2517) | +| pottsmpnn (sample + score) | level[0] only, extra levels silently dropped | no | +| lasermpnn (sample + score) | level[0] only, extra levels silently dropped | no | +| protonpottsmpnn (all purposes) | never (silently ignored) | no | + +Design derives the per-structure noise key as `fold_in(_base_sampling_key(spec), structure_idx)` +(`kernel_dispatch.py:236,562`); the one draw is shared by every level, temperature and sample. + +## Decisions + +- **N-D1 shape.** Score returns `(S, C)` when the spec has exactly one noise level (any value, including + nonzero), and `(S, C, D)` when it has two or more. `C` before `D` mirrors design's samples-before-noise + order. The shape depends only on the spec, never on data. Fused multi-state: `(1, C)` / `(1, C, D)`. + Logits, when returned, gain the same axis in the same position. The levels are always reported + (`noise_levels`) in the result metadata. +- **N-D2 key.** Score draws noise with exactly design's key: `fold_in(_base_sampling_key(spec), structure_idx)`, + shared across candidates and levels. So scoring a design's output at the same spec/seed sees the identical + noised backbone. The per-(structure, candidate) keys that drive decoding order are unchanged, so noise-0 + scores stay bit-identical to today. +- **N-D3 encode-once.** Encode once per (structure, level) and score every candidate against it; the per-candidate + branch passes `backbone_noise=level` with the N-D2 noise key (not the candidate key). +- **N-D4 averaged path.** Unchanged semantics (it fuses D). Fix #2517 (move the R3 host check out of the vmap) + if it is local; otherwise leave it and keep D>=2 refused there with a clear error. +- **N-D5 driver families.** pottsmpnn, lasermpnn and protonpottsmpnn take every level. Each family keeps its + existing noise rule and key (`0x9077` Potts, `0x6262` LASEr); the level index is NOT folded into the key, so + levels share one draw scaled by the level (as stock design does) and a single-level run is bit-identical to + today. Output gains a noise axis only when D >= 2, after the sample/candidate axis, declared in `sample_axes` + / `result_schema`. ProtonPotts reuses `apply_potts_backbone_noise` (upstream PottsMPNN-pH featurizes through + the same `structure_noise` path; its pH engine passes 0.0, `potts_mpnn_ph.py:1119`) with its own key constant, + applied to featurized coordinates before any table/energy/pH-design computation. +- **N-D6 sinks.** `agent/shaping.py::_score_structures` accepts the 3-D form and emits `noise_index`; + `SCORING_SCHEMA_VERSION` bumps to `scoring_v2` (adds `noise_levels`). Docstrings fixed (runner.sample still says + `(B*N, L)`). +- **N-D7 out of scope here.** Browser pages (no backbone noise exists in any JS page today; it needs a + coordinate-noise step before the encoder graph and can only be distributionally, not bitwise, graded). + Multistate PoE campaign path (`sampling/multistate_poe.py`) keeps its own key scheme. + +## Acceptance + +1. Noise 0 (default) is bit-identical to main for every path and family (pin tests + the 9-row re-wave). +2. Score with D >= 2 returns `(S, C, D)`; slice `[..., d]` equals a single-level run at level `d`. +3. For the stock family, the noised backbone (or encoder output) at (structure s, level d) is identical between + `runner.sample` and `runner.score` for the same spec/seed. +4. Nonzero noise changes the score (control: the old silent-drop bug would fail this). +5. Each driver family: D >= 2 yields the noise axis; level 0 slice equals the noise-free run; a level > 0 changes it. +6. #214's refusal test file is replaced by these tests. + +Ledger: `host/runner.py`, `scoring/score.py`, the family drivers are in row closures, so all 9 rows go stale; +one re-wave after the branch is complete, ledger last. diff --git a/src/aminx/agent/shaping.py b/src/aminx/agent/shaping.py index 93b5666a..66e588f8 100644 --- a/src/aminx/agent/shaping.py +++ b/src/aminx/agent/shaping.py @@ -31,7 +31,7 @@ DEFAULT_INLINE_CAP = 10_000 -_SKIP_ROUTE = frozenset({"metadata", "schema_version"}) +_SKIP_ROUTE = frozenset({"metadata", "schema_version", "noise_levels"}) _KINDS = frozenset({"sample", "score", "inspect", "jacobian"}) _SEQUENCE_NDIM = 5 # Residue-length axes of per-structure list results (inspect features, jacobian outputs). The @@ -227,6 +227,9 @@ def shape_result( fused = metadata.get("fused_structure_ids") if fused is not None: summary["fused_structure_ids"] = to_jsonable(fused) + noise_levels = metadata.get("noise_levels") + if noise_levels is not None: + summary["noise_levels"] = to_jsonable(noise_levels) if "lineage" in metadata: summary["lineage"] = to_jsonable(metadata["lineage"]) summary["side_file"] = _write_side_file(output_dir, run_id, store) @@ -320,13 +323,18 @@ def _score_structures( structure_ids: Sequence[str], warnings: list[str], ) -> list[dict[str, object]]: - """Pair negative-log-likelihood rows with ``sequences_to_score``.""" + """Pair negative-log-likelihood rows with ``sequences_to_score``. + + ``scores`` is ``(S, C)`` for a single noise level, or ``(S, C, D)`` with two or more (scoring_v2): + the 3-D form emits one entry per (candidate, noise level), tagged with ``noise_index``. + """ array = np.asarray(scores) if array.ndim == 1: array = array.reshape(1, -1) - if array.ndim != 2: - warnings.append(f"score array has ndim {array.ndim}; expected 2") + if array.ndim not in (2, 3): + warnings.append(f"score array has ndim {array.ndim}; expected 2 or 3") return [] + has_noise_axis = array.ndim == 3 sequences = [str(item) for item in getattr(spec, "sequences_to_score", ()) or ()] if len(sequences) < int(array.shape[1]): warnings.append( @@ -337,7 +345,13 @@ def _score_structures( scored: list[dict[str, object]] = [] for column in range(int(array.shape[1])): sequence = sequences[column] if column < len(sequences) else "" - scored.append({"sequence": sequence, "nll": float(array[row_index, column])}) + if has_noise_axis: + scored.extend( + {"sequence": sequence, "noise_index": noise_index, "nll": float(value)} + for noise_index, value in enumerate(array[row_index, column]) + ) + else: + scored.append({"sequence": sequence, "nll": float(array[row_index, column])}) rows.append( { "structure_id": _structure_id_at(structure_ids, row_index), diff --git a/src/aminx/families/laser_mpnn/driver.py b/src/aminx/families/laser_mpnn/driver.py index 55f74410..b7e001eb 100644 --- a/src/aminx/families/laser_mpnn/driver.py +++ b/src/aminx/families/laser_mpnn/driver.py @@ -214,6 +214,9 @@ class _Prepared(NamedTuple): laser_indices: np.ndarray canonical_indices: np.ndarray order: np.ndarray + # One plain ``_Prepared`` per noise level when the spec sweeps two or more + # (spec 261011 N-D5); empty for a single level. + variants: tuple[_Prepared, ...] = () def _graph(model: LASErMPNN) -> GraphStructure: @@ -336,6 +339,13 @@ def _score_prepared(model: LASErMPNN, item: _Prepared, purpose: str) -> jax.Arra return jnp.stack(rows) +def _score_noised(model: LASErMPNN, item: _Prepared, purpose: str) -> jax.Array: + """One structure's scores; with a noise sweep the level axis sits at 1 (after N_cand).""" + if not item.variants: + return _score_prepared(model, item, purpose) + return jnp.stack([_score_prepared(model, v, purpose) for v in item.variants], axis=1) + + def _proofread_batch( model: LASErMPNN, prepared: tuple[_Prepared, ...], @@ -512,7 +522,7 @@ def __call__( seed=self._seed, dropout_key=dropout_key, ) - stacked = jnp.stack([_score_prepared(self._model, item, self._purpose) for item in prepared]) + stacked = jnp.stack([_score_noised(self._model, item, self._purpose) for item in prepared]) key = "logits" if self._purpose == "score:logits" else "nll" payload: dict[str, jax.Array] = {key: stacked} if self._options.output_fasta or self._options.output_fasta_only: @@ -524,6 +534,37 @@ def __call__( return payload +def backbone_noise_levels(spec: Any) -> tuple[float, ...]: # noqa: ANN401 + """Every backbone noise level this spec asks for, or ``(0.0,)``. + + Levels share ONE draw (the key ``0x6262`` fold is not indexed by level) scaled + by the level, so a single-level run is bit-identical to the first-level-only + behaviour. A noise axis exists in the output iff there are two or more. + """ + for bundle in getattr(spec, "noise", ()) or (): + if bundle.feature_type == "backbone" and bundle.enabled and bundle.noise_levels: + return tuple(float(level) for level in bundle.noise_levels) + return (0.0,) + + +def noise_axis_dims(dims: tuple[str, ...], levels: Sequence[float]) -> tuple[str, ...]: + """Insert ``D_noise`` right after the leading candidate/sample axis when D >= 2.""" + if len(levels) < 2: + return dims + return (dims[0], "D_noise", *dims[1:]) + + +def refuse_proofread_noise(spec: Any, purpose: str) -> None: # noqa: ANN401 + """Proofread reduces over an ensemble of one noised structure: no noise axis exists.""" + levels = backbone_noise_levels(spec) + if purpose.startswith("score:proofread_") and len(levels) >= 2: + msg = ( + f"lasermpnn {purpose}: a multi-level backbone_noise sweep ({len(levels)} levels) " + "has no noise axis on proofread results; pass a single level." + ) + raise ValueError(msg) + + def backbone_noise_level(spec: Any) -> float: # noqa: ANN401 """The backbone noise level this spec asks for, or 0.0. @@ -589,22 +630,6 @@ def _prepare(path: Path, spec: Any, options: LaserOptions) -> _Prepared: # noqa ignore_ligand=options.ignore_ligand, dtype=_working_dtype(), ) - level = backbone_noise_level(spec) - if level > 0.0: - # Keyed off random_seed so a run is reproducible, and folded with a constant - # so the noise draw cannot coincide with the decoding-order draw below, - # which reads the same seed. - noise_key = jax.random.fold_in(jax.random.PRNGKey(int(spec.random_seed)), 0x6262) - # LaserFeatures is a frozen dataclass, not a NamedTuple, so dataclasses - # .replace rather than ._replace. - features = dataclasses.replace( - features, - backbone_coords=apply_backbone_noise( - features.backbone_coords, - level, - key=noise_key, - ), - ) length = int(features.sequence_indices.shape[0]) rows = [_letters(sequence) for sequence in _spec_sequences(spec)] for index, row in enumerate(rows): @@ -612,15 +637,50 @@ def _prepare(path: Path, spec: Any, options: LaserOptions) -> _Prepared: # noqa msg = f"sequences_to_score[{index}] length {row.shape[0]} != structure length {length}" raise ValueError(msg) laser = np.stack(rows) - return _Prepared( + canonical = _canonical_index(laser) + variants = tuple( + _Prepared( + noised, + laser, + canonical, + decoding_order( + np.asarray(noised.chain_mask), + np.asarray(noised.extra_atom_contact_mask), + seed=int(spec.random_seed), + ), + ) + for noised in ( + noised_features(features, spec, level, noise_fold=0) for level in backbone_noise_levels(spec) + ) + ) + if len(variants) == 1: + return variants[0] + return variants[0]._replace(variants=variants) + + +def noised_features( + features: LaserFeatures, + spec: Any, # noqa: ANN401 + level: float, + *, + noise_fold: int = 0, +) -> LaserFeatures: + """``features`` with LASEr backbone noise at ``level``; ``level <= 0`` is a no-op. + + Keyed off ``random_seed`` so a run is reproducible, and folded with a constant + so the noise draw cannot coincide with the decoding-order draw, which reads the + same seed. The level is NOT folded in: every level of a sweep shares one draw. + ``noise_fold`` separates the tied partner (1) from the primary structure (0). + """ + if level <= 0.0: + return features + base = jax.random.fold_in(jax.random.PRNGKey(int(spec.random_seed)), 0x6262) + key = base if noise_fold == 0 else jax.random.fold_in(base, noise_fold) + # LaserFeatures is a frozen dataclass, not a NamedTuple, so dataclasses + # .replace rather than ._replace. + return dataclasses.replace( features, - laser, - _canonical_index(laser), - decoding_order( - np.asarray(features.chain_mask), - np.asarray(features.extra_atom_contact_mask), - seed=int(spec.random_seed), - ), + backbone_coords=apply_backbone_noise(features.backbone_coords, level, key=key), ) @@ -869,8 +929,8 @@ def axes(self, spec: Any, purpose: str, batch: FamilyBatch) -> list[AxisSpec]: from aminx.families.laser_mpnn.sample_host import sample_axes # noqa: PLC0415 return sample_axes(spec, batch) - del spec, purpose - return [ + del purpose + specs = [ AxisSpec( name="structures", cardinality=max(len(batch.input_indices), 1), @@ -878,6 +938,10 @@ def axes(self, spec: Any, purpose: str, batch: FamilyBatch) -> list[AxisSpec]: heterogeneous=True, ), ] + n_noise = len(backbone_noise_levels(spec)) + if n_noise >= 2: + specs.append(AxisSpec(name="noise", cardinality=n_noise, default_batch_size=n_noise)) + return specs def stages(self, spec: Any, purpose: str, model: eqx.Module) -> FamilyStages: # noqa: ANN401 """Per-batch sample, teacher-forced score, or a proofread reduction.""" @@ -894,16 +958,27 @@ def result_schema(self, spec: Any, purpose: str) -> Mapping[str, SinkArraySpec]: return sample_schema(spec) options = _options(spec) + refuse_proofread_noise(spec, purpose) + levels = backbone_noise_levels(spec) + # ``fasta`` is the candidate's own sequence: it does not depend on the noise. fasta = { "fasta": SinkArraySpec(dims=("N_cand", "L_total"), dtype="int32", attrs={}), } if purpose == "score:logits": body = { - "logits": SinkArraySpec(dims=("N_cand", "L_total", "alphabet"), dtype="float32", attrs={}), + "logits": SinkArraySpec( + dims=noise_axis_dims(("N_cand", "L_total", "alphabet"), levels), + dtype="float32", + attrs={}, + ), } elif purpose == "score:nll": body = { - "nll": SinkArraySpec(dims=("N_cand", "L_total"), dtype="float32", attrs={}), + "nll": SinkArraySpec( + dims=noise_axis_dims(("N_cand", "L_total"), levels), + dtype="float32", + attrs={}, + ), } elif purpose == "score:proofread_unconditional": # No candidate axis. Proofreading reads the structure's own sequence, so diff --git a/src/aminx/families/laser_mpnn/sample_host.py b/src/aminx/families/laser_mpnn/sample_host.py index c2c317b0..0c25f7ae 100644 --- a/src/aminx/families/laser_mpnn/sample_host.py +++ b/src/aminx/families/laser_mpnn/sample_host.py @@ -6,7 +6,6 @@ from __future__ import annotations -import dataclasses from collections.abc import Iterator, Mapping from pathlib import Path from typing import Any, NamedTuple, cast @@ -18,11 +17,12 @@ from aminx.families.laser_mpnn.driver import ( CANONICAL_OF_LASER, - apply_backbone_noise, - backbone_noise_level, + backbone_noise_levels, budget_mask_for, decoding_order, load_joint_decode, + noise_axis_dims, + noised_features, ) from aminx.families.laser_mpnn.featurize import ( LaserFeatures, @@ -46,6 +46,9 @@ class _Sampled(NamedTuple): features: LaserFeatures chain_mask: np.ndarray tied: LaserFeatures | None + # One plain ``_Sampled`` per noise level when the spec sweeps two or more + # (spec 261011 N-D5); empty for a single level. + variants: tuple[_Sampled, ...] = () def _x64_enabled() -> bool: @@ -95,9 +98,20 @@ def _featurize_one( options: LaserOptions, *, noise_fold: int, + level: float | None = None, ) -> LaserFeatures: """B0 features plus LASEr backbone noise. Fold 0 matches the score-path key.""" - features = featurize( + features = _featurize_raw(path, options) + return noised_features( + features, + spec, + backbone_noise_levels(spec)[0] if level is None else level, + noise_fold=noise_fold, + ) + + +def _featurize_raw(path: Path, options: LaserOptions) -> LaserFeatures: + return featurize( path, fix_from_bfactor=options.fix_from_bfactor, use_water=options.use_water, @@ -105,15 +119,6 @@ def _featurize_one( ignore_ligand=options.ignore_ligand, dtype=_working_dtype(), ) - level = backbone_noise_level(spec) - if level > 0.0: - base = jax.random.fold_in(jax.random.PRNGKey(int(spec.random_seed)), 0x6262) - key = base if noise_fold == 0 else jax.random.fold_in(base, noise_fold) - features = dataclasses.replace( - features, - backbone_coords=apply_backbone_noise(features.backbone_coords, level, key=key), - ) - return features def _chain_mask(features: LaserFeatures, spec: Any, options: LaserOptions) -> np.ndarray: # noqa: ANN401 @@ -134,11 +139,15 @@ def _tied_partner( options: LaserOptions, spec: Any, # noqa: ANN401 length: int, + *, + level: float | None = None, ) -> LaserFeatures | None: if not options.tied_second_input: return None try: - partner = _featurize_one(Path(options.tied_second_input), spec, options, noise_fold=1) + partner = _featurize_one( + Path(options.tied_second_input), spec, options, noise_fold=1, level=level, + ) except LaserInputError as exc: msg = f"tied_second_input: {exc}" raise ValueError(msg) from exc @@ -152,22 +161,30 @@ def _tied_partner( def sample_batches(spec: Any) -> Iterator[FamilyBatch]: # noqa: ANN401 """One structure per batch. Tied partners are featurized here, decoded later.""" options = _options(spec) + levels = backbone_noise_levels(spec) pending: list[tuple[int, str]] = [] for index, item in enumerate(_inputs(spec)): if not isinstance(item, (str, Path)): pending.append((index, "lasermpnn_requires_pdb")) continue try: - features = _featurize_one(Path(item), spec, options, noise_fold=0) + per_level = [] + for level in levels: + features = _featurize_one(Path(item), spec, options, noise_fold=0, level=level) + length = int(features.sequence_indices.shape[0]) + per_level.append( + _Sampled( + features, + _chain_mask(features, spec, options), + _tied_partner(options, spec, length, level=level), + ), + ) except LaserInputError as exc: pending.append((index, str(exc))) continue - length = int(features.sequence_indices.shape[0]) - sampled = _Sampled( - features, - _chain_mask(features, spec, options), - _tied_partner(options, spec, length), - ) + sampled = per_level[0] + if len(per_level) > 1: + sampled = sampled._replace(variants=tuple(per_level)) yield FamilyBatch( input_indices=(index,), arrays={"sampled": (sampled,)}, @@ -198,6 +215,9 @@ def sample_axes(spec: Any, batch: FamilyBatch) -> list[AxisSpec]: # noqa: ANN40 AxisSpec(name="samples", cardinality=n_samples, default_batch_size=n_samples), AxisSpec(name="temperatures", cardinality=n_temp, default_batch_size=n_temp), ] + n_noise = len(backbone_noise_levels(spec)) + if n_noise >= 2: + specs.append(AxisSpec(name="noise", cardinality=n_noise, default_batch_size=n_noise)) BatchPlanner().plan(specs) return specs @@ -212,6 +232,15 @@ def _include_tied_logits(spec: Any, options: LaserOptions) -> bool: # noqa: ANN def sample_schema(spec: Any) -> dict[str, SinkArraySpec]: # noqa: ANN401 + """§5.6 sample channels, with ``D_noise`` after ``N`` when the spec sweeps >= 2 levels.""" + levels = backbone_noise_levels(spec) + return { + name: array._replace(dims=noise_axis_dims(array.dims, levels)) + for name, array in _base_sample_schema(spec).items() + } + + +def _base_sample_schema(spec: Any) -> dict[str, SinkArraySpec]: # noqa: ANN401 """§5.6 sample channels. ``sidechain_coords`` is omitted for FASTA-only runs.""" options = _options(spec) fasta = { @@ -398,6 +427,21 @@ def _one( chunk_start: int, chunk_count: int, dropout_key: jax.Array, + ) -> dict[str, np.ndarray]: + """One level, or every level with the SAME dropout key, noise axis at 1.""" + if not item.variants: + return self._one_level(item, chunk_start, chunk_count, dropout_key) + per_level = [ + self._one_level(variant, chunk_start, chunk_count, dropout_key) for variant in item.variants + ] + return {key: np.stack([row[key] for row in per_level], axis=1) for key in per_level[0]} + + def _one_level( + self, + item: _Sampled, + chunk_start: int, + chunk_count: int, + dropout_key: jax.Array, ) -> dict[str, np.ndarray]: temperatures = _temperatures(self._spec) n_temp = len(temperatures) diff --git a/src/aminx/families/potts_mpnn/driver.py b/src/aminx/families/potts_mpnn/driver.py index c4e1e2bc..b8febde0 100644 --- a/src/aminx/families/potts_mpnn/driver.py +++ b/src/aminx/families/potts_mpnn/driver.py @@ -40,6 +40,7 @@ tied_featurize_port, ) from aminx.families.potts_mpnn.model import PottsMPNN, _encoder_states, cast_floating +from aminx.families.potts_mpnn.noise_levels import potts_backbone_noise_levels from aminx.families.potts_mpnn.sample_host import ( SampleStages, force_optimize_num_samples, @@ -95,6 +96,9 @@ class _Prepared(NamedTuple): expt: np.ndarray partitions: tuple[_Graph, ...] chain_order: tuple[str, ...] + # One entry per noise level when the spec has two or more (spec 261011 N-D5); + # each is a plain ``_Prepared`` at that level. Empty for a single level. + variants: tuple[_Prepared, ...] = () class _SubtractPartitions: @@ -181,7 +185,13 @@ def __call__( ids: list[jax.Array] = [] expt: list[jax.Array] = [] for item in prepared: - values = _score_one(self._model, item, mean_norm=self._mean_norm) + if item.variants: + values = jnp.stack( + [_score_one(self._model, v, mean_norm=self._mean_norm) for v in item.variants], + axis=-1, + ) + else: + values = _score_one(self._model, item, mean_norm=self._mean_norm) rows.append(values) ids.append(jnp.arange(values.shape[0], dtype=jnp.int32)) if item.include_reference_in_output: @@ -310,7 +320,7 @@ def axes(self, spec: Any, purpose: str, batch: FamilyBatch) -> list[AxisSpec]: """Structure, mutant, partition, or sample/temperature axes.""" if purpose == "sample": return sample_axes(spec, batch) - del spec + n_noise = len(potts_backbone_noise_levels(spec)) prepared = cast("tuple[_Prepared, ...]", batch.arrays.get("prepared", ())) n_mut = int(prepared[0].graph.sequences.shape[0]) if prepared else 1 n_part = len(prepared[0].partitions) if prepared else 0 @@ -328,6 +338,8 @@ def axes(self, spec: Any, purpose: str, batch: FamilyBatch) -> list[AxisSpec]: default_batch_size=max(n_mut, 1), ), ] + if n_noise >= 2: + specs.append(AxisSpec(name="noise", cardinality=n_noise, default_batch_size=n_noise)) if purpose == "score:ddg" and n_part > 0: specs.append( AxisSpec( @@ -356,15 +368,23 @@ def result_schema(self, spec: Any, purpose: str) -> Mapping[str, SinkArraySpec]: """§4.5 arrays. Per-position dims are ``L_total`` on the sample path.""" if purpose == "sample": return sample_schema(spec) - del spec + noised = len(potts_backbone_noise_levels(spec)) >= 2 if purpose == "score:energy": return { - "energy": SinkArraySpec(dims=("N_cand",), dtype="float32", attrs={}), + "energy": SinkArraySpec( + dims=("N_cand", "D_noise") if noised else ("N_cand",), + dtype="float32", + attrs={}, + ), "candidate_ids": SinkArraySpec(dims=("N_cand",), dtype="int32", attrs={}), } if purpose == "score:ddg": return { - "ddg": SinkArraySpec(dims=("N_mut",), dtype="float32", attrs={}), + "ddg": SinkArraySpec( + dims=("N_mut", "D_noise") if noised else ("N_mut",), + dtype="float32", + attrs={}, + ), "mutant_ids": SinkArraySpec(dims=("N_mut",), dtype="int32", attrs={}), "ddg_expt": SinkArraySpec(dims=("N_mut",), dtype="float32", attrs={}), } @@ -396,6 +416,7 @@ def _pack( name: str, *, spec: Any = None, # noqa: ANN401 + level: float | None = None, ) -> _Prepared: reference = "".join(chain.sequence for chain in chains) if purpose == "score:energy": @@ -413,7 +434,9 @@ def _pack( raise ValueError(msg) rows = (reference, *[item.sequence for item in candidates]) touched = set().union(*(item.chains for item in candidates)) - partitions = _partition_graphs(parsed, chains, options, name, rows, touched, spec=spec) + partitions = _partition_graphs( + parsed, chains, options, name, rows, touched, spec=spec, level=level, + ) expt = np.asarray([item.expt for item in candidates], dtype=np.float32) return _Prepared( graph=_graph(features, rows), @@ -579,6 +602,7 @@ def _partition_graphs( touched: set[str], *, spec: Any = None, # noqa: ANN401 + level: float | None = None, ) -> tuple[_Graph, ...]: """Re-featurize each flagged partition and score those chains only. @@ -600,7 +624,7 @@ def _partition_graphs( designed = sorted(partition) part_features = tied_featurize_port([parsed], {name: (designed, [])})[0] if spec is not None: - part_features = _noised(part_features, spec, stream=stream) + part_features = _noised(part_features, spec, stream=stream, level=level) part_rows = tuple(_project(row, chains, designed) for row in rows) graphs.append(_graph(part_features, part_rows)) return tuple(graphs) @@ -762,15 +786,23 @@ def apply_potts_backbone_noise( return base + float(level) * shift -def _noised(features: PottsFeatures, spec: Any, *, stream: int = 0) -> PottsFeatures: # noqa: ANN401 +def _noised( + features: PottsFeatures, + spec: Any, # noqa: ANN401 + *, + stream: int = 0, + level: float | None = None, +) -> PottsFeatures: """Apply the spec's backbone noise to featurized coordinates, if any. Returns ``features`` unchanged when the knob is off, which is the default, so every run that does not ask for noise is bit-identical to before this existed. ``stream`` separates independent draws (binding partitions); stream 0 is the - complex and keeps exactly the key it always had. + complex and keeps exactly the key it always had. ``level`` overrides the + spec's first level; every level of a multi-level sweep uses the SAME key. """ - level = potts_backbone_noise_level(spec) + if level is None: + level = potts_backbone_noise_level(spec) if level <= 0.0: return features key = jax.random.fold_in(jax.random.PRNGKey(int(spec.random_seed)), 0x9077) @@ -898,6 +930,13 @@ def _is_sample(spec: Any) -> bool: # noqa: ANN401 def _sample_batches(spec: Any) -> Iterator[FamilyBatch]: # noqa: ANN401 """One PDB per batch for purpose ``sample``.""" options = _options(spec) + levels = potts_backbone_noise_levels(spec) + if len(levels) >= 2 and options.emit_etab: + msg = ( + "pottsmpnn: emit_etab has no noise axis; a multi-level backbone_noise sweep " + f"({len(levels)} levels) cannot carry potts_etab. Use one level or disable emit_etab." + ) + raise ValueError(msg) pending: list[tuple[int, str]] = [] for index, item in enumerate(_inputs(spec)): if not isinstance(item, (str, Path)): @@ -908,16 +947,24 @@ def _sample_batches(spec: Any) -> Iterator[FamilyBatch]: # noqa: ANN401 parsed_list = parse_pdb_upstream(path, skip_gaps=options.skip_gaps) parsed = parsed_list[0] name = str(parsed["name"]) - features = _noised(_featurize_one(parsed, options, name), spec) - if knn_boundary_tie(features.present, features.L_total): - log.warning("knn_boundary_tie for %s (L_total=%s)", name, features.L_total) - chains = _chain_sequences(parsed, features) - ready = prepare_sample( - features, - tuple((chain.letter, chain.sequence) for chain in chains), - options, - spec, - ) + base = _featurize_one(parsed, options, name) + readies = [] + for level in levels: + features = _noised(base, spec, level=level) + if knn_boundary_tie(features.present, features.L_total): + log.warning("knn_boundary_tie for %s (L_total=%s)", name, features.L_total) + chains = _chain_sequences(parsed, features) + readies.append( + prepare_sample( + features, + tuple((chain.letter, chain.sequence) for chain in chains), + options, + spec, + ), + ) + ready = readies[0] + if len(readies) > 1: + ready = ready._replace(noise_variants=tuple(readies)) except PottsInputError as exc: pending.append((index, str(exc))) continue @@ -946,7 +993,9 @@ def _prepare_from_spec(item: Any, spec: Any) -> _Prepared: # noqa: ANN401 parsed_list = parse_pdb_upstream(path, skip_gaps=options.skip_gaps) parsed = parsed_list[0] name = str(parsed["name"]) - features = _noised(_featurize_one(parsed, options, name), spec) + base = _featurize_one(parsed, options, name) + levels = potts_backbone_noise_levels(spec) + features = _noised(base, spec, level=levels[0]) if knn_boundary_tie(features.present, features.L_total): log.warning("knn_boundary_tie for %s (L_total=%s)", name, features.L_total) chains = _chain_sequences(parsed, features) @@ -957,7 +1006,23 @@ def _prepare_from_spec(item: Any, spec: Any) -> _Prepared: # noqa: ANN401 ) else: candidates = _candidates(name, chains, options, purpose, ()) - return _pack(parsed, features, chains, candidates, options, purpose, name, spec=spec) + variants = tuple( + _pack( + parsed, + _noised(base, spec, level=level), + chains, + candidates, + options, + purpose, + name, + spec=spec, + level=level, + ) + for level in levels + ) + if len(variants) == 1: + return variants[0] + return variants[0]._replace(variants=variants) FAMILY_DRIVERS.register("pottsmpnn")(PottsMPNNDriver()) diff --git a/src/aminx/families/potts_mpnn/noise_levels.py b/src/aminx/families/potts_mpnn/noise_levels.py new file mode 100644 index 00000000..12ad52e3 --- /dev/null +++ b/src/aminx/families/potts_mpnn/noise_levels.py @@ -0,0 +1,19 @@ +"""Every backbone noise level a spec asks for (spec 261011 N-D5).""" + +from __future__ import annotations + +from typing import Any + + +def potts_backbone_noise_levels(spec: Any) -> tuple[float, ...]: # noqa: ANN401 + """All levels of the first enabled backbone noise bundle, or ``(0.0,)``. + + Levels share ONE noise draw scaled by the level (the level index is never + folded into the key), so a single-level run is bit-identical to a run that + only ever read ``noise_levels[0]``. The shape of the output depends only on + ``len(levels)``: a noise axis exists iff there are two or more. + """ + for bundle in getattr(spec, "noise", ()) or (): + if bundle.feature_type == "backbone" and bundle.enabled and bundle.noise_levels: + return tuple(float(level) for level in bundle.noise_levels) + return (0.0,) diff --git a/src/aminx/families/potts_mpnn/sample_host.py b/src/aminx/families/potts_mpnn/sample_host.py index 9241873c..b53a553d 100644 --- a/src/aminx/families/potts_mpnn/sample_host.py +++ b/src/aminx/families/potts_mpnn/sample_host.py @@ -33,6 +33,7 @@ ) from aminx.families.potts_mpnn.featurize import PottsFeatures, pad from aminx.families.potts_mpnn.model import PottsMPNN, _encoder_states, cast_floating +from aminx.families.potts_mpnn.noise_levels import potts_backbone_noise_levels from aminx.families.potts_mpnn.refine import BindingTables, PottsRefine, check_tied_only_groups from aminx.host.family_driver import FamilyBatch, SinkArraySpec from aminx.host.omit_aa_bias import omit_letter_indices @@ -210,6 +211,8 @@ class _Ready(NamedTuple): name: str l_total: int tied: bool + # Per-level readies when the spec sweeps two or more noise levels (261011 N-D5). + noise_variants: tuple[_Ready, ...] = () def build_tie_groups_np( @@ -384,6 +387,9 @@ def sample_axes(spec: Any, batch: FamilyBatch) -> list[AxisSpec]: # noqa: ANN40 specs.append( AxisSpec(name="temperatures", cardinality=n_temp, default_batch_size=n_temp), ) + n_noise = len(potts_backbone_noise_levels(spec)) + if n_noise >= 2: + specs.append(AxisSpec(name="noise", cardinality=n_noise, default_batch_size=n_noise)) BatchPlanner().plan(specs) return specs @@ -394,19 +400,22 @@ def sample_schema(spec: Any) -> dict[str, SinkArraySpec]: # noqa: ANN401 optimizing = bool(options.optimize_pdb or options.optimize_fasta) and ( options.optimization_mode != "none" ) - sequence = SinkArraySpec(dims=("N", "L_total"), dtype="int32", attrs={}) + noised = len(potts_backbone_noise_levels(spec)) >= 2 + lead = ("N", "D_noise") if noised else ("N",) + sequence = SinkArraySpec(dims=(*lead, "L_total"), dtype="int32", attrs={}) if optimizing: schema: dict[str, SinkArraySpec] = {"refined_sequence": sequence} else: schema = { "sequence": sequence, - "sample_energy": SinkArraySpec(dims=("N",), dtype="float32", attrs={}), - "sample_rank": SinkArraySpec( - dims=("N",), - dtype="int32", - attrs={"level": "structure"}, - ), + "sample_energy": SinkArraySpec(dims=lead, dtype="float32", attrs={}), } + # The runner ranks along axis N, so with a noise axis the ranking is per level. + schema["sample_rank"] = SinkArraySpec( + dims=lead, + dtype="int32", + attrs={"level": "structure"}, + ) if options.optimization_mode != "none": schema["refined_sequence"] = sequence if options.emit_etab: @@ -449,6 +458,22 @@ def _one( chunk_start: int, chunk_count: int, dropout_key: jax.Array, + ) -> dict[str, jax.Array]: + """Run one level, or every level with the SAME dropout key, noise axis at 1.""" + if not ready.noise_variants: + return self._one_level(ready, chunk_start, chunk_count, dropout_key) + per_level = [ + self._one_level(variant, chunk_start, chunk_count, dropout_key) + for variant in ready.noise_variants + ] + return {key: jnp.stack([row[key] for row in per_level], axis=1) for key in per_level[0]} + + def _one_level( + self, + ready: _Ready, + chunk_start: int, + chunk_count: int, + dropout_key: jax.Array, ) -> dict[str, jax.Array]: options = self._options encoded = _encode( diff --git a/src/aminx/families/protonpotts_mpnn/driver.py b/src/aminx/families/protonpotts_mpnn/driver.py index dc769b2b..b6bba1dd 100644 --- a/src/aminx/families/protonpotts_mpnn/driver.py +++ b/src/aminx/families/protonpotts_mpnn/driver.py @@ -41,7 +41,9 @@ import numpy as np from xtrax.tiling import AxisSpec, BatchPlanner +from aminx.families.potts_mpnn.driver import apply_potts_backbone_noise from aminx.families.potts_mpnn.model import PottsMPNN +from aminx.families.potts_mpnn.noise_levels import potts_backbone_noise_levels from aminx.families.protonpotts_mpnn.energy import protonpotts_energies, protonpotts_table from aminx.families.protonpotts_mpnn.features import ( PdbFeatureError, @@ -76,6 +78,9 @@ _REFERENCE_PURPOSES = frozenset({"score:energy", "score:selectivity"}) _KEY = re.compile(r"^(?P[^:\s]+):(?P-?\d+)(?P[A-Za-z]?)$") _VOCABULARY = ",".join(PROTONPOTTS_V6.symbols) +# Fold constant of the backbone-noise key. Distinct from Potts' 0x9077 and LASEr's 0x6262 so the three +# families never share a draw, and shared by every level of a sweep (the level is not folded in). +_NOISE_FOLD = 0x7070 class _Graph(NamedTuple): @@ -93,6 +98,9 @@ class _Prepared(NamedTuple): graph: _Graph include_reference_in_output: bool kept: tuple[tuple[str, int, str, str], ...] = () + # One plain ``_Prepared`` per backbone-noise level when the spec sweeps two or more (spec 261011 + # N-D5); empty for a single level. + variants: tuple[_Prepared, ...] = () class ProtonPottsInputError(ValueError): @@ -119,6 +127,15 @@ class _ScoreStages: def __init__(self, model: PottsMPNN) -> None: self._model = model + def _block(self, item: _Prepared, *, reference_mode: bool) -> tuple[jax.Array, jax.Array]: + """Energies (or ddG) of one structure's rows at one noise level, and the rows shown.""" + graph = item.graph + block = _JIT_ENERGIES(self._model, *_graph_args(graph), jnp.asarray(graph.sequences)) + rows = jnp.asarray(graph.sequences, dtype=jnp.int32) + if reference_mode: + return block.astype(jnp.float32), rows + return (block[1:] - block[0]).astype(jnp.float32), rows[1:] + def __call__( self, batch: FamilyBatch, @@ -136,17 +153,16 @@ def __call__( expt: list[jax.Array] = [] reference_mode = prepared[0].include_reference_in_output for item in prepared: - graph = item.graph - block = _JIT_ENERGIES(self._model, *_graph_args(graph), jnp.asarray(graph.sequences)) - rows = jnp.asarray(graph.sequences, dtype=jnp.int32) - if reference_mode: - values, shown = block.astype(jnp.float32), rows - else: - values, shown = (block[1:] - block[0]).astype(jnp.float32), rows[1:] + per_level = [ + self._block(variant, reference_mode=reference_mode) for variant in (item.variants or (item,)) + ] + shown = per_level[0][1] + # Noise axis (D >= 2) sits after the candidate axis; one level keeps the plain (C,) vector. + values = per_level[0][0] if len(per_level) == 1 else jnp.stack([v for v, _ in per_level], axis=1) energies.append(values) ids.append(jnp.arange(values.shape[0], dtype=jnp.int32)) tokens.append(shown) - expt.append(jnp.full(values.shape, jnp.nan, dtype=jnp.float32)) + expt.append(jnp.full((values.shape[0],), jnp.nan, dtype=jnp.float32)) if reference_mode: return { "energy": jnp.stack(energies), @@ -187,13 +203,23 @@ def __call__( ids: list[np.ndarray] = [] tokens: list[np.ndarray] = [] for item, pins in zip(prepared, pins_by_structure, strict=True): - table, e_idx = _JIT_TABLE(self._model, *_graph_args(item.graph)) forced = np.stack([force_pins(row, pins) for row in item.graph.sequences]) - gaps = [selectivity_gaps(table, e_idx, row, pins) for row in forced] - totals.append(np.asarray([total for total, _ in gaps], dtype=np.float32)) - per_centre.append( - np.asarray([per for _, per in gaps], dtype=np.float32).reshape(len(gaps), -1), - ) + level_totals: list[np.ndarray] = [] + level_centres: list[np.ndarray] = [] + for variant in item.variants or (item,): + table, e_idx = _JIT_TABLE(self._model, *_graph_args(variant.graph)) + gaps = [selectivity_gaps(table, e_idx, row, pins) for row in forced] + level_totals.append(np.asarray([total for total, _ in gaps], dtype=np.float32)) + level_centres.append( + np.asarray([per for _, per in gaps], dtype=np.float32).reshape(len(gaps), -1), + ) + if item.variants: + # Noise axis right after the candidate axis: (C, D) and (C, D, N_centre). + totals.append(np.stack(level_totals, axis=1)) + per_centre.append(np.stack(level_centres, axis=1)) + else: + totals.append(level_totals[0]) + per_centre.append(level_centres[0]) ids.append(np.arange(forced.shape[0], dtype=np.int32)) tokens.append(forced) return { @@ -319,6 +345,7 @@ def batches(self, spec: Any) -> Iterator[FamilyBatch]: # noqa: ANN401 """One structure per batch, featurized by the v6 host port.""" purpose = _purpose(spec) options = _options(spec) + _refuse_design_noise_axis(spec, purpose) if purpose == "sample": _check_sample_count(spec) if purpose == "score:selectivity" and not options.explicit_centers: @@ -334,7 +361,7 @@ def batches(self, spec: Any) -> Iterator[FamilyBatch]: # noqa: ANN401 continue path = Path(item) try: - prepared = _prepare(path, options, purpose, _spec_sequences(spec)) + prepared = _prepare(path, options, purpose, _spec_sequences(spec), spec=spec) except PdbFeatureError as exc: pending.append((index, str(exc))) continue @@ -411,8 +438,7 @@ def _design_one( return designs def axes(self, spec: Any, purpose: str, batch: FamilyBatch) -> list[AxisSpec]: # noqa: ANN401 - """Structure and candidate axes.""" - del spec + """Structure and candidate axes, plus ``noise`` when the spec sweeps two or more levels.""" if purpose == "sample": designs = cast("tuple[list[PHDesign], ...]", batch.arrays.get("designs", ())) n_rows = max((len(d) for d in designs), default=1) @@ -429,6 +455,9 @@ def axes(self, spec: Any, purpose: str, batch: FamilyBatch) -> list[AxisSpec]: ), AxisSpec(name="mutants", cardinality=max(n_rows, 1), default_batch_size=max(n_rows, 1)), ] + n_noise = len(potts_backbone_noise_levels(spec)) + if n_noise >= 2 and purpose != "sample": + specs.append(AxisSpec(name="noise", cardinality=n_noise, default_batch_size=n_noise)) BatchPlanner().plan(specs) return specs @@ -449,24 +478,31 @@ def result_schema(self, spec: Any, purpose: str) -> Mapping[str, SinkArraySpec]: """ vocab = {"vocabulary": _VOCABULARY} names = _id_attrs(spec, purpose) + _refuse_design_noise_axis(spec, purpose) + noised = len(potts_backbone_noise_levels(spec)) >= 2 + cand = ("N_cand", "D_noise") if noised else ("N_cand",) if purpose == "score:energy": return { - "energy": SinkArraySpec(dims=("N_cand",), dtype="float32", attrs={}), + "energy": SinkArraySpec(dims=cand, dtype="float32", attrs={}), "candidate_ids": SinkArraySpec(dims=("N_cand",), dtype="int32", attrs=names), "candidate_tokens": SinkArraySpec(dims=("N_cand", "L_total"), dtype="int32", attrs=vocab), } if purpose == "score:ddg": return { - "ddg": SinkArraySpec(dims=("N_mut",), dtype="float32", attrs={}), + "ddg": SinkArraySpec( + dims=("N_mut", "D_noise") if noised else ("N_mut",), + dtype="float32", + attrs={}, + ), "mutant_ids": SinkArraySpec(dims=("N_mut",), dtype="int32", attrs=names), "ddg_expt": SinkArraySpec(dims=("N_mut",), dtype="float32", attrs={}), "mutant_tokens": SinkArraySpec(dims=("N_mut", "L_total"), dtype="int32", attrs=vocab), } if purpose == "score:selectivity": return { - "selectivity": SinkArraySpec(dims=("N_cand",), dtype="float32", attrs={}), + "selectivity": SinkArraySpec(dims=cand, dtype="float32", attrs={}), "selectivity_per_centre": SinkArraySpec( - dims=("N_cand", "N_centre"), + dims=(*cand, "N_centre"), dtype="float32", attrs={}, ), @@ -497,6 +533,36 @@ def _purpose(spec: Any) -> str: # noqa: ANN401 return f"score:{kind}" +def _refuse_design_noise_axis(spec: Any, purpose: str) -> None: # noqa: ANN401 + """pH design cannot carry a noise axis, so a multi-level sweep is refused rather than truncated. + + Designs are sorted by final energy per run, and the placement plan is chosen from the energy table of the + (noised) structure, so design ``i`` at level 0 and design ``i`` at level 1 need not share a plan, a pin + layout or even a count. There is no honest ``(N_design, D)`` array; run one level per call instead. + """ + levels = potts_backbone_noise_levels(spec) + if purpose == "sample" and len(levels) >= 2: + msg = ( + f"protonpottsmpnn sample (pH design): a multi-level backbone_noise sweep ({len(levels)} levels) has no " + "noise axis, because designs are energy-sorted and placement plans are re-chosen per noised " + "structure; pass a single level (one call per level)." + ) + raise ValueError(msg) + + +def _noised_coords(coords: np.ndarray, level: float, spec: Any) -> np.ndarray: # noqa: ANN401 + """Potts' per-atom Gaussian (``augment_eps * randn``) on featurized coordinates; ``level <= 0`` is a no-op. + + Upstream PottsMPNN-pH featurizes through the same ``structure_noise`` path as PottsMPNN (its pH engine + passes 0.0). Keyed off ``random_seed`` with its own fold constant; the level is not folded in, so levels + share one draw scaled by the level. + """ + if level <= 0.0: + return coords + key = jax.random.fold_in(jax.random.PRNGKey(int(spec.random_seed)), _NOISE_FOLD) + return np.asarray(apply_potts_backbone_noise(coords, level, key=key), dtype=np.float32) + + def _check_sample_count(spec: Any) -> None: # noqa: ANN401 """``sample`` designs are set by ``samples_per_site``; ``num_samples`` > 1 would repeat each design.""" count = int(getattr(spec, "num_samples", 1)) @@ -693,6 +759,8 @@ def _prepare( options: ProtonPottsOptions, purpose: str, sequences: Sequence[str], + *, + spec: Any = None, # noqa: ANN401 ) -> _Prepared: loaded = loaded_residues(path) labels = _reference_labels(loaded, options) @@ -723,11 +791,20 @@ def _prepare( pad_valid=np.ones(length, dtype=np.bool_), sequences=rows, ) - return _Prepared( - graph=graph, - include_reference_in_output=purpose in _REFERENCE_PURPOSES, - kept=tuple((chain, int(number), icode, name) for chain, number, icode, name in kept), + kept_rows = tuple((chain, int(number), icode, name) for chain, number, icode, name in kept) + # No spec (a direct call) means no noise, exactly as before this axis existed. + levels = potts_backbone_noise_levels(spec) if spec is not None else (0.0,) + variants = tuple( + _Prepared( + graph=graph._replace(coords=_noised_coords(graph.coords, level, spec)), + include_reference_in_output=purpose in _REFERENCE_PURPOSES, + kept=kept_rows, + ) + for level in levels ) + if len(variants) == 1: + return variants[0] + return variants[0]._replace(variants=variants) FAMILY_DRIVERS.register("protonpottsmpnn")(ProtonPottsDriver()) diff --git a/src/aminx/host/family_runner.py b/src/aminx/host/family_runner.py index eb907b18..bf07bb5d 100644 --- a/src/aminx/host/family_runner.py +++ b/src/aminx/host/family_runner.py @@ -353,7 +353,8 @@ def _stage_sample_rank( if not parts: continue energy = np.concatenate(parts, axis=0) if len(parts) > 1 else parts[0] - rank = np.argsort(energy, kind="stable").astype(np.int32) + # Rank along the sample axis: (N,) as always, or (N, D_noise) per noise level (spec 261011). + rank = np.argsort(energy, axis=0, kind="stable").astype(np.int32) if sink is None: bucket = memory.setdefault(index, {}) bucket["sample_rank"] = [rank] diff --git a/src/aminx/host/runner.py b/src/aminx/host/runner.py index 7400dd9a..f7381b5f 100644 --- a/src/aminx/host/runner.py +++ b/src/aminx/host/runner.py @@ -16,6 +16,7 @@ import numpy as np from aminx.host._sampling_grid_lineage import ( + _base_sampling_key, _grid_iteration_arrays, _grid_manifest_row_hash, _grid_sample_indices, @@ -303,18 +304,18 @@ def sample( dict[str, Any] Dictionary with keys: sequences : jax.Array - Sampled sequences, shape (B*N, L) where B is batch count, - N is num_samples per structure, L is sequence length. - Padded to max length across all sequences, masked with 'mask' key. + Sampled sequences, shape (S, N, D, T, L): S structures, N samples per + structure, D backbone-noise levels, T temperatures, L (padded) sequence length. + Masked with 'mask' key. mask : jax.Array - Sequence validity mask (1 for valid, 0 for padding). Shape: (B*N, L). + Sequence validity mask (1 for valid, 0 for padding). Same shape as ``sequences``. schema_version : str Schema version for results ('grid_v1' or 'sampling_v1'). metadata : dict Metadata including specification, skipped_inputs, structure_ids, and optional lineage info (grid mode). logits : jax.Array, optional - Per-position logits if return_logits=True. Shape: (B*N, L, 21). + Per-position logits if return_logits=True. Shape: (S, N, D, T, L, 21). pseudo_perplexity : jax.Array, optional Per-sequence pseudo-perplexity scores if computed. @@ -643,9 +644,14 @@ def score_sequence_averaged( ) bundles_per_noise.append(bundle) - # R3 assertion: runs on concrete arrays (outside JIT) so eqx.tree_equal returns - # a Python bool, not an abstract tracer. This is the safe location for this check. - if len(bundles_per_noise) > 1: + # R3 assertion: needs concrete arrays so eqx.tree_equal returns a Python bool. Under the + # runner's jax.vmap the inputs are tracers and bool() raises TracerBoolConversionError for + # D >= 2 (debt #2517), so the check is hoisted: ``runner.score`` calls + # ``score_sequence_averaged.check_r3`` once, eagerly, on the first concrete structure, and + # this in-call check only runs when the inputs are concrete (a direct, un-vmapped call). + if len(bundles_per_noise) > 1 and not any( + isinstance(leaf, jax.core.Tracer) for leaf in jax.tree.leaves(bundles_per_noise) + ): _check_r3_invariance(bundles_per_noise) # Delegate to JIT-compiled core @@ -660,6 +666,38 @@ def score_sequence_averaged( use_rolling_state=use_rolling_state, ) + def check_r3( + sequence: jax.Array, + structure_coordinates: jax.Array, + mask: jax.Array, + residue_index: jax.Array, + chain_index: jax.Array, + ) -> None: + """Run the R3 invariance check eagerly on concrete arrays (hoisted out of the vmap, #2517). + + The invariant -- conditioning / geometry / ligand / wave are identical across the D noise + levels -- is a property of the bundle builder, not of the data, so one concrete structure + suffices. Raises ``AssertionError`` on a violation; a no-op for fewer than two levels. + """ + levels = spec.run_spec.sampling.backbone_noise or (0.0,) + if len(levels) < 2: + return + bundles = [ + build_inference_bundle( + coords=structure_coordinates, + mask=mask, + residue_index=residue_index, + chain_index=chain_index, + sequence=sequence, + backbone_noise=float(level), + mode="score_conditional", + inference=True, + )[0] + for level in levels + ] + _check_r3_invariance(bundles) + + score_sequence_averaged.check_r3 = check_r3 # type: ignore[attr-defined] return score_sequence_averaged @@ -752,6 +790,46 @@ def _one_candidate( return int(per_structure) * batch_size +def _score_noise_levels(spec: Any) -> tuple[float, ...]: # noqa: ANN401 + """The spec's backbone-noise levels as floats; ``(0.0,)`` when none are set.""" + return tuple(float(level) for level in (spec.backbone_noise or (0.0,))) + + +def _score_noise_key(spec: Any, structure_idx: Any) -> jax.Array: # noqa: ANN401 + """The per-structure backbone-noise key, identical to ``runner.sample``'s (N-D2). + + Design derives ``fold_in(_base_sampling_key(spec), structure_idx)`` + (``host/kernel_dispatch.py``, ``encode_key``) and hands it to the encoder; score uses the + same derivation so scoring at a design's spec/seed sees the identical noised backbone. + ``structure_idx`` is the index WITHIN the structure batch, as in design (design folds + ``jnp.arange(batch_size)``, not a global index). Scoring has no grid lineage, so + ``grid_lineage=None`` (the non-grid design key). + """ + return jax.random.fold_in(_base_sampling_key(spec, grid_lineage=None), structure_idx) + + +def _over_noise_levels( + run: Callable[..., Any], + noise_levels: tuple[float, ...], + noise_key: jax.Array, +) -> Any: # noqa: ANN401 + """Run ``run(noise, key)`` at each noise level; D >= 2 gains a noise axis after candidates. + + ``run`` returns a pytree whose leaves have a leading candidate axis. One level calls it + directly (output unchanged, ``(C, ...)``); two or more map it over the levels with + ``jax.vmap`` and move the level axis to position 1, giving ``(C, D, ...)``. A single level + of exactly ``0.0`` passes ``(None, None)`` so the call is the historical noise-free one, + bit for bit. + """ + any_noise = any(level != 0.0 for level in noise_levels) + if len(noise_levels) == 1: + return run(*((noise_levels[0], noise_key) if any_noise else (None, None))) + key = noise_key if any_noise else None + levels = jnp.asarray(noise_levels, dtype=jnp.float32) + out = jax.vmap(lambda level: run(level, key))(levels) + return jax.tree.map(lambda leaf: jnp.moveaxis(leaf, 0, 1), out) + + def score( # noqa: PLR0915 spec: ScoringSpecification | None = None, **kwargs: Any, # noqa: ANN401 @@ -775,15 +853,27 @@ def score( # noqa: PLR0915 dict[str, Any] Dictionary with keys: scores : jax.Array - Negative log-likelihood scores. Shape: (num_structures, num_sequences). + Negative log-likelihood scores. Shape: (num_structures, num_sequences) when + ``spec.backbone_noise`` has exactly one level, else + (num_structures, num_sequences, D) with one trailing entry per noise level + (candidates before noise, mirroring ``sample``'s samples-before-noise order). + The shape depends only on the spec. The fused multi-state path + (``state_position_map``) has one fused "structure": (1, C) / (1, C, D). With + ``average_node_features=True`` the levels are fused (D -> 1), so there is never a + noise axis. mask : jax.Array Structure validity mask (1 for valid, 0 for padding). Shape: (num_structures,). schema_version : str - Schema version for results ('scoring_v1'). + Schema version for results ('scoring_v2'). + noise_levels : tuple[float, ...] + The backbone-noise levels scored, in axis order (also in ``metadata``). Noise is + drawn from ``runner.sample``'s per-structure key, so a score at the same + spec/seed sees the same noised backbone as the design did. metadata : dict - Metadata including specification, structure_ids, and skipped_inputs. + Metadata including specification, structure_ids, noise_levels, and skipped_inputs. logits : jax.Array, optional - Per-position logits if return_logits=True. Shape: (num_structures, num_sequences, L, 21). + Per-position logits if return_logits=True. Shape: (num_structures, num_sequences, L, 21), + with the same trailing-D axis position as scores: (S, C, D, L, 21) for D >= 2. decoding_orders : None, optional Present (and always ``None``) if return_decoding_orders=True. Every score() path is full-context teacher forcing (``1 - I`` mask, no waves), so there is no @@ -864,6 +954,13 @@ def score( # noqa: PLR0915 ) raise ValueError(msg) + # Backbone noise axis (N-D1..N-D3, spec 261011_noise-level-axis). Replaces the #214 refusal. + # The plain per-structure and fused multi-state paths apply every level in + # spec.backbone_noise, drawing the noise from design's per-structure key + # (``_score_noise_key``); the averaged-feature path fuses the levels (D -> 1) and has no + # noise axis. One level -> ``(S, C)``; two or more -> ``(S, C, D)``. + noise_levels = _score_noise_levels(spec) + from aminx.scoring.score import make_score_fn, make_score_split_fns # noqa: PLC0415 from aminx.utils.aa_convert import string_to_protein_sequence # noqa: PLC0415 @@ -888,10 +985,10 @@ def score( # noqa: PLR0915 score_fn = make_score_fn(model) # type: ignore[arg-type] # Encode-once path (#147). Only the plain per-structure branch below uses it: the encoding is - # candidate-independent exactly when no backbone noise reaches the encoder, which is how this - # branch has always called the core (it forwards no noise level). The averaged-feature path - # encodes at spec.backbone_noise levels with keys derived from each candidate's key, and the - # fused multi-state path (state_position_map) stacks states into one encode; both keep their + # candidate-independent because the noise is drawn from the per-STRUCTURE key (N-D2), so it + # is encoded once per (structure, noise level) (N-D3). The averaged-feature path encodes at + # spec.backbone_noise levels with keys derived from each candidate's key, and the fused + # multi-state path (state_position_map) stacks states into one encode; both keep their # per-candidate core and are NOT hoisted. split_fns = ( make_score_split_fns(model) # type: ignore[arg-type] @@ -922,6 +1019,7 @@ def score( # noqa: PLR0915 protein_iterator, score_fn, sequence_indices_list, + noise_levels=noise_levels, ) from aminx.sampling.conditional_logits import _plan_axis_strategy # noqa: PLC0415 @@ -964,6 +1062,7 @@ def _score_rung_of(batch: Any) -> int | None: # noqa: ANN401 ) padding_check = _PaddingCheck("score", rung_fn=_score_rung_of) + r3_checked = False for _batch_idx, batched_ensemble in enumerate(protein_iterator): padding_check(batched_ensemble) batch_size = batched_ensemble.coordinates.shape[0] @@ -1017,6 +1116,9 @@ def _score_rung_of(batch: Any) -> int | None: # noqa: ANN401 prng_key, subkey = jax.random.split(prng_key) flat_keys.append(subkey) batch_keys = jnp.stack(flat_keys, axis=0).reshape(batch_size, n_candidates, -1) + # One backbone-noise key per structure in the batch (N-D2): design's key, shared by every + # candidate and every noise level. Never consumed when no level is nonzero. + batch_noise_keys = jax.vmap(lambda idx: _score_noise_key(spec, idx))(jnp.arange(batch_size)) # Per-structure ligand tensors (B, L, M, ...) or None. Only prepared when the spec asks for # a ligand, so the ligand-free path is byte-for-byte what it was. @@ -1056,7 +1158,8 @@ def _score_rung_of(batch: Any) -> int | None: # noqa: ANN401 score_ligand = tuple(trim_residue_axis(part, rung, 1) for part in score_ligand) if split_fns is not None: - activation_bytes = _candidate_activation_bytes( + # The noise-level axis is vmapped, so its activations coexist: scale by D. + activation_bytes = len(noise_levels) * _candidate_activation_bytes( split_fns, spec.multi_state_strategy, SimpleNamespace( @@ -1092,6 +1195,7 @@ def _score_structure( struct_chain_index: jax.Array, struct_keys: jax.Array, struct_ligand: tuple[jax.Array | None, jax.Array | None, jax.Array | None], + struct_noise_key: jax.Array, _candidate_iterator: Any = candidate_iterator, # noqa: ANN401 _stacked_sequences: jax.Array = score_sequences, ) -> tuple[jax.Array, jax.Array, jax.Array | None]: @@ -1106,52 +1210,77 @@ def _score_structure( ) if split_fns is not None: - # Encode this structure ONCE (the ligand tensors belong to the encode side); the - # candidate map below closes over the encoding and runs only decode + NLL. - encoding = split_fns.encode_structure( - struct_coords, - struct_mask, - struct_residue_index, - struct_chain_index, - **ligand_kwargs, - ) - def _score_one_candidate( - item: dict[str, jax.Array], + def _run_level( + noise: float | jax.Array | None, + noise_key: jax.Array | None, ) -> tuple[jax.Array, jax.Array, jax.Array | None]: - return split_fns.score_candidate( - item["key"], - jax.nn.one_hot(item["seq"], 21), - encoding, + noise_kwargs = {} if noise is None else {"backbone_noise": noise, "noise_key": noise_key} + # Encode this structure ONCE per noise level (the ligand tensors belong to the encode + # side); the candidate map below closes over the encoding and runs only decode + NLL. + encoding = split_fns.encode_structure( struct_coords, struct_mask, struct_residue_index, struct_chain_index, - multi_state_strategy=spec.multi_state_strategy, + **ligand_kwargs, + **noise_kwargs, ) - else: + def _score_one_candidate( + item: dict[str, jax.Array], + ) -> tuple[jax.Array, jax.Array, jax.Array | None]: + return split_fns.score_candidate( + item["key"], + jax.nn.one_hot(item["seq"], 21), + encoding, + struct_coords, + struct_mask, + struct_residue_index, + struct_chain_index, + multi_state_strategy=spec.multi_state_strategy, + ) - def _score_one_candidate( - item: dict[str, jax.Array], - ) -> tuple[jax.Array, jax.Array, jax.Array | None]: - seq_one_hot = jax.nn.one_hot(item["seq"], 21) - return score_fn( # type: ignore[misc] - item["key"], - seq_one_hot, - struct_coords, - struct_mask, - struct_residue_index, - struct_chain_index, - multi_state_strategy=spec.multi_state_strategy, - **ligand_kwargs, + return _candidate_iterator( + _score_one_candidate, + {"key": struct_keys, "seq": _stacked_sequences}, ) + return _over_noise_levels(_run_level, noise_levels, struct_noise_key) + + # Averaged-feature path: it fuses the noise levels itself (D -> 1) from spec.backbone_noise, + # so it takes no noise key and has no noise axis. + def _score_one_candidate( + item: dict[str, jax.Array], + ) -> tuple[jax.Array, jax.Array, jax.Array | None]: + seq_one_hot = jax.nn.one_hot(item["seq"], 21) + return score_fn( # type: ignore[misc] + item["key"], + seq_one_hot, + struct_coords, + struct_mask, + struct_residue_index, + struct_chain_index, + multi_state_strategy=spec.multi_state_strategy, + **ligand_kwargs, + ) + return _candidate_iterator( _score_one_candidate, {"key": struct_keys, "seq": _stacked_sequences}, ) + if spec.average_node_features and len(noise_levels) > 1 and not r3_checked: + # Hoisted R3 check (#2517): concrete arrays, once, outside the vmap below. + score_fn.check_r3( # type: ignore[union-attr] + score_sequences[0], + score_coords[0], + score_mask[0], + score_residue_index[0], + score_chain_index[0], + ) + r3_checked = True + # The third output (a decoding order) is discarded: scoring is full-context, see the # `decoding_orders` note in this function's docstring (debt #1717). batch_scores, batch_logits, _ = jax.vmap(_score_structure)( @@ -1161,6 +1290,7 @@ def _score_one_candidate( score_chain_index, batch_keys, score_ligand, + batch_noise_keys, ) if rung is not None: batch_logits = repad_residue_axis(batch_logits, struct_len, -2) @@ -1172,22 +1302,25 @@ def _score_one_candidate( resolved_structure_ids.extend(batch_structure_ids) structure_offset += batch_size - # Concatenate results: shape (num_structures, num_sequences) + # Concatenate results: (num_structures, num_sequences), plus a trailing noise-level axis + # (num_structures, num_sequences, D) when the spec has two or more levels (N-D1). scores_array = jnp.concatenate(all_scores, axis=0) results = { "scores": scores_array, "mask": jnp.ones(scores_array.shape[0], dtype=jnp.int32), "schema_version": SCORING_SCHEMA_VERSION, + "noise_levels": noise_levels, "metadata": { "specification": spec, "skipped_inputs": getattr(protein_iterator, "skipped_frames", []), "structure_ids": resolved_structure_ids, + "noise_levels": noise_levels, }, } if spec.run_spec.sampling.return_logits and all_logits is not None: - # Shape: (num_structures, num_sequences, L, 21) + # Shape: (num_structures, num_sequences, L, 21), or (S, C, D, L, 21) with D >= 2 levels. results["logits"] = jnp.concatenate(all_logits, axis=0) if spec.run_spec.sampling.return_decoding_orders: @@ -1196,11 +1329,13 @@ def _score_one_candidate( return results -def _score_fused_multistate( # not length-bucketed in S8 v1 +def _score_fused_multistate( # noqa: PLR0915 # not length-bucketed in S8 v1 spec: ScoringSpecification, protein_iterator: Any, # noqa: ANN401 score_fn: Any, # noqa: ANN401 sequence_indices_list: list[jax.Array], + *, + noise_levels: tuple[float, ...] = (0.0,), ) -> dict[str, Any]: """Score sequences against ONE genuinely fused multi-state ensemble. @@ -1217,6 +1352,9 @@ def _score_fused_multistate( # not length-bucketed in S8 v1 Returns a `scores`/`logits` shaped with a leading dimension of 1 (one fused "structure"), not `len(spec.inputs)` independent ones -- the semantics are genuinely different from the default branch's per-structure output, by design. + `scores` is `(1, C)`, or `(1, C, D)` for two or more `noise_levels`; the noise draw uses the + design key for structure 0 (`_score_noise_key`). The averaged-feature scorer fuses levels + itself (D -> 1), so with `average_node_features` the levels are recorded but not mapped. """ # Draining the Grain protein_iterator into host-side lists is not the "hand-rolled # loop over JAX computation" this project's composability convention forbids -- no @@ -1302,41 +1440,66 @@ def _score_fused_multistate( # not length-bucketed in S8 v1 ) candidate_iterator = make_axis_dispatch_via_xtrax(strategy, axis=N_CANDIDATES.name) - def _score_one_candidate( - item: dict[str, jax.Array], - ) -> tuple[jax.Array, jax.Array, jax.Array | None]: - seq_one_hot = jax.nn.one_hot(item["seq"], 21) - return score_fn( - item["key"], - seq_one_hot, + if spec.average_node_features and len(noise_levels) > 1: + # Hoisted R3 check (#2517): eager, on concrete arrays, outside the candidate vmap. + score_fn.check_r3( + stacked_sequences[0], stacked_coords, stacked_mask, stacked_residue_index, stacked_chain_index, - state_position_map=state_position_map, - multi_state_strategy=spec.multi_state_strategy, - multi_state_temperature=spec.multi_state_temperature, ) - all_scores, all_logits, _ = candidate_iterator( - _score_one_candidate, - {"key": candidate_keys, "seq": stacked_sequences}, + def _run_level( + noise: float | jax.Array | None, + noise_key: jax.Array | None, + ) -> tuple[jax.Array, jax.Array, jax.Array | None]: + noise_kwargs = {} if noise is None else {"backbone_noise": noise, "noise_key": noise_key} + + def _score_one_candidate( + item: dict[str, jax.Array], + ) -> tuple[jax.Array, jax.Array, jax.Array | None]: + seq_one_hot = jax.nn.one_hot(item["seq"], 21) + return score_fn( + item["key"], + seq_one_hot, + stacked_coords, + stacked_mask, + stacked_residue_index, + stacked_chain_index, + state_position_map=state_position_map, + multi_state_strategy=spec.multi_state_strategy, + multi_state_temperature=spec.multi_state_temperature, + **noise_kwargs, + ) + + return candidate_iterator( + _score_one_candidate, + {"key": candidate_keys, "seq": stacked_sequences}, + ) + + all_scores, all_logits, _ = _over_noise_levels( + _run_level, + (0.0,) if spec.average_node_features else noise_levels, + _score_noise_key(spec, 0), ) # Leading dim of 1: ONE fused "structure", not len(spec.inputs) independent ones. results: dict[str, Any] = { - "scores": all_scores[None, :], + "scores": all_scores[None], "mask": jnp.ones(1, dtype=jnp.int32), "schema_version": SCORING_SCHEMA_VERSION, + "noise_levels": noise_levels, "metadata": { "specification": spec, "skipped_inputs": getattr(protein_iterator, "skipped_frames", []), + "noise_levels": noise_levels, "structure_ids": ["fused_multistate"], "fused_structure_ids": structure_ids, }, } if spec.return_logits: - results["logits"] = all_logits[None, :] + results["logits"] = all_logits[None] if spec.return_decoding_orders: # Full-context scoring has no decoding order; see score()'s docstring (debt #1717). results["decoding_orders"] = None diff --git a/src/aminx/host/schema_versions.py b/src/aminx/host/schema_versions.py index 594ec603..1ba6e9ed 100644 --- a/src/aminx/host/schema_versions.py +++ b/src/aminx/host/schema_versions.py @@ -30,7 +30,9 @@ SAMPLING_SCHEMA_VERSION: Final[str] = "sampling_v1" # Scoring (full-context inference) schema version. -SCORING_SCHEMA_VERSION: Final[str] = "scoring_v1" +# scoring_v2 (261011): scores gain a trailing noise-level axis when the spec has >= 2 +# backbone_noise levels, (S, C, D), and results record ``noise_levels``. +SCORING_SCHEMA_VERSION: Final[str] = "scoring_v2" # Inspection (detailed model internals export) schema version. INSPECTION_SCHEMA_VERSION: Final[str] = "inspection_v1" diff --git a/src/aminx/scoring/score.py b/src/aminx/scoring/score.py index 952c2acf..2e5dde16 100644 --- a/src/aminx/scoring/score.py +++ b/src/aminx/scoring/score.py @@ -138,6 +138,7 @@ def score_sequence( ligand_coords: jax.Array | None = None, ligand_atom_types: jax.Array | None = None, ligand_mask: jax.Array | None = None, + noise_key: jax.Array | None = None, **kwargs, # Accept but ignore extra kwargs (e.g., _k_neighbors for backward compat) ) -> tuple[jax.Array, jax.Array, jax.Array]: @@ -201,7 +202,15 @@ def score_sequence( state_weights=state_weights, ) - logits = score_conditional.kernel(model, prng_key, bundle, config, stage_set) + if noise_key is None: + logits = score_conditional.kernel(model, prng_key, bundle, config, stage_set) + else: + # Backbone noise is drawn from ``noise_key`` (the design-shared per-structure key, + # N-D2), NOT from the per-candidate ``prng_key``: that one only drives decoding. The + # decode half is split exactly as ``kernel`` does it. + _, k_dec = score_conditional.split_keys(prng_key) + enc = score_conditional.encode(model, noise_key, bundle, config) + logits = score_conditional.score_from_encoding(model, k_dec, enc, bundle, config, stage_set) # Compute score using single-source NLL formula nll = _nll_from_logits(logits, sequence, mask) @@ -240,12 +249,11 @@ def make_score_split_fns( decoding_order)`` -- the same three outputs as ``make_score_fn``'s function, with the same key derivation (decoding-order draw, then ``score_conditional.split_keys``). - ``encode_structure`` has no PRNG key and no ``backbone_noise`` on purpose. Sharing one - encoding across candidates is only valid when the encoding does not depend on the - candidate's key, i.e. when ``backbone_noise`` is 0 -- which is how ``runner.score`` calls - the un-split core (it never forwards a noise level on this path). With nonzero noise the - encoder consumes the per-candidate key and every candidate would need its own encoding; - use ``make_score_fn`` (or the averaged path) for that. + ``encode_structure`` takes an optional ``backbone_noise`` level and ``noise_key`` + (defaults ``0.0`` / ``None``, which reproduce the old noise-free encoding exactly). Sharing + one encoding across candidates is valid because the noise is drawn from a per-STRUCTURE key + (the same ``fold_in(base_key, structure_idx)`` design uses, N-D2), never from a + per-candidate key; the runner encodes once per (structure, noise level). Args: model: Protein or Ligand Equinox checkpoint. @@ -268,23 +276,29 @@ def encode_structure( ligand_coords: jax.Array | None = None, ligand_atom_types: jax.Array | None = None, ligand_mask: jax.Array | None = None, + backbone_noise: float | jax.Array = 0.0, + noise_key: jax.Array | None = None, ) -> EncoderOutput: # Sequence / AR mask are left at their defaults: the encoder reads only geometry, ligand - # and backbone_noise (0 here), never the conditioning bundle. + # and backbone_noise, never the conditioning bundle. bundle, config = build_inference_bundle( coords=structure_coordinates, mask=mask, residue_index=residue_index, chain_index=chain_index, - backbone_noise=0.0, + backbone_noise=backbone_noise, ligand_coords=ligand_coords, ligand_atom_types=ligand_atom_types, ligand_mask=ligand_mask, mode="score_conditional", inference=True, ) - # The key only feeds backbone-noise injection, which is off (see docstring above). - k_enc, _ = score_conditional.split_keys(jax.random.PRNGKey(0)) + # The key only feeds backbone-noise injection. With no noise_key (noise off) keep the + # historical constant key so the noise-free encoding is bit-identical to before. + if noise_key is None: + k_enc, _ = score_conditional.split_keys(jax.random.PRNGKey(0)) + else: + k_enc = noise_key return score_conditional.encode(model, k_enc, bundle, config) @partial(jax.jit, static_argnames=("multi_state_strategy",)) diff --git a/tests/agent/test_shaping.py b/tests/agent/test_shaping.py index 36628eeb..2c73aff8 100644 --- a/tests/agent/test_shaping.py +++ b/tests/agent/test_shaping.py @@ -204,7 +204,7 @@ def test_score_pairs_sequences_and_parks_logits(tmp_path: Path) -> None: results = { "scores": np.array([[1.0, 2.0]], dtype=np.float32), "logits": logits, - "schema_version": "scoring_v1", + "schema_version": "scoring_v2", "metadata": { "structure_ids": ["fused_multistate"], "fused_structure_ids": ["1ubq", "1mbn"], @@ -226,6 +226,33 @@ def test_score_pairs_sequences_and_parks_logits(tmp_path: Path) -> None: np.testing.assert_array_equal(loaded["logits"], logits) +def test_score_noise_axis_emits_noise_index(tmp_path: Path) -> None: + """A 3-D (S, C, D) score array emits one entry per (candidate, noise level) with noise_index.""" + spec = ScoringSpecification(inputs=[str(_PDB)], sequences_to_score=["AAAA", "CCCC"]) + scores = np.arange(1 * 2 * 3, dtype=np.float32).reshape(1, 2, 3) + results = { + "scores": scores, + "schema_version": "scoring_v2", + "noise_levels": (0.0, 0.1, 0.3), + "metadata": { + "structure_ids": ["1ubq"], + "skipped_inputs": [], + "noise_levels": (0.0, 0.1, 0.3), + }, + } + shaped = _shape(tmp_path, results, kind="score", spec=spec, inline_cap=10, run_id="score-noise") + (structure,) = shaped["structures"] + assert [(e["sequence"], e["noise_index"], e["nll"]) for e in structure["scores"]] == [ + ("AAAA", 0, 0.0), + ("AAAA", 1, 1.0), + ("AAAA", 2, 2.0), + ("CCCC", 0, 3.0), + ("CCCC", 1, 4.0), + ("CCCC", 2, 5.0), + ] + assert shaped["noise_levels"] == [0.0, 0.1, 0.3] + + def test_inspect_features_use_structure_keys(tmp_path: Path) -> None: """Inspect arrays are keyed by feature and structure id.""" logits = np.arange(4 * 21, dtype=np.float32).reshape(4, 21) diff --git a/tests/families/laser_mpnn/test_noise_axis.py b/tests/families/laser_mpnn/test_noise_axis.py new file mode 100644 index 00000000..bcafe28f --- /dev/null +++ b/tests/families/laser_mpnn/test_noise_axis.py @@ -0,0 +1,177 @@ +# ruff: noqa: S101 +"""LASEr noise-level axis (spec 261011 N-D5, acceptance 1 and 5). + +``backbone_noise`` takes every level. Levels share ONE draw (key ``0x6262``, the level is not folded in) +scaled by the level, so a single level keeps today's shape and values, and two or more levels add a +``D`` axis right after the candidate/sample axis whose slice ``d`` equals a single-level run at level ``d``. +Proofread results have no noise axis, so a multi-level sweep is refused there. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import equinox as eqx +import jax +import numpy as np +import pytest + +pytest.importorskip("prody", reason="needs the laser extra") + +from aminx.families.laser_mpnn.driver import LASErMPNN, LaserDriver # noqa: E402 +from aminx.host.family_driver import FAMILY_DRIVERS # noqa: E402 +from aminx.host.runner import sample, score # noqa: E402 +from aminx.run.options import LaserOptions # noqa: E402 +from aminx.run.specs import SamplingSpecification, ScoringSpecification # noqa: E402 +from tests.families.laser_mpnn.test_driver import _write_pdb # noqa: E402 + +_LEVEL = 0.5 + + +@pytest.fixture +def registered() -> Iterator[LaserDriver]: + driver = LaserDriver() + FAMILY_DRIVERS.register("lasermpnn")(driver) + yield driver + FAMILY_DRIVERS.discard("lasermpnn") + + +@pytest.fixture +def model_path(tmp_path: Path) -> Path: + path = tmp_path / "tiny.eqx" + eqx.tree_serialise_leaves(path, LASErMPNN(key=jax.random.PRNGKey(0))) + return path + + +@pytest.fixture +def pdb(tmp_path: Path) -> Path: + path = tmp_path / "complex.pdb" + _write_pdb(path) + return path + + +def _score_spec(pdb: Path, model_path: Path, kind: str, **kwargs: Any) -> ScoringSpecification: # noqa: ANN401 + return ScoringSpecification( + inputs=str(pdb), + model_family="lasermpnn", + checkpoint_id="lasermpnn_test", + model_local_path=model_path, + output_kind=kind, # type: ignore[arg-type] + sequences_to_score=["AA", "AG"], + random_seed=3, + **kwargs, + ) + + +def _nll(pdb: Path, model_path: Path, **kwargs: Any) -> np.ndarray: # noqa: ANN401 + result = score(_score_spec(pdb, model_path, "nll", **kwargs)) + return np.asarray(result["structures"]["0"]["arrays"]["nll"], dtype=np.float64) + + +def _sample(pdb: Path, model_path: Path, **kwargs: Any) -> dict[str, Any]: # noqa: ANN401 + spec = SamplingSpecification( + inputs=str(pdb), + model_family="lasermpnn", + checkpoint_id="lasermpnn_test", + model_local_path=model_path, + num_samples=2, + temperature=0.8, + random_seed=5, + laser=LaserOptions(chi_temp=1.0, disabled_residues=("X",)), + **kwargs, + ) + return sample(spec)["structures"]["0"]["arrays"] + + +def test_score_single_level_keeps_shape_and_values( + registered: LaserDriver, + pdb: Path, + model_path: Path, +) -> None: + del registered + clean = _nll(pdb, model_path) + scalar = _nll(pdb, model_path, backbone_noise=_LEVEL) + one = _nll(pdb, model_path, backbone_noise=(_LEVEL,)) + assert clean.shape == scalar.shape == one.shape # (C, L) + np.testing.assert_array_equal(scalar, one) + assert not np.allclose(scalar, clean), "nonzero single-level noise must still reach the nll" + + +def test_score_two_levels_add_axis_and_slices_match_single_runs( + registered: LaserDriver, + pdb: Path, + model_path: Path, +) -> None: + del registered + clean = _nll(pdb, model_path) + noised = _nll(pdb, model_path, backbone_noise=_LEVEL) + both = _nll(pdb, model_path, backbone_noise=(0.0, _LEVEL)) + n_cand, length = clean.shape + assert both.shape == (n_cand, 2, length), both.shape # (C, D, L): D right after the candidate axis + np.testing.assert_array_equal(both[:, 0], clean) + np.testing.assert_array_equal(both[:, 1], noised) + assert not np.allclose(both[:, 0], both[:, 1]) + + +def test_score_schema_declares_noise_axis_only_for_two_levels( + registered: LaserDriver, + pdb: Path, + model_path: Path, +) -> None: + one = _score_spec(pdb, model_path, "nll", backbone_noise=_LEVEL) + two = _score_spec(pdb, model_path, "nll", backbone_noise=(0.0, _LEVEL)) + assert registered.result_schema(one, "score:nll")["nll"].dims == ("N_cand", "L_total") + assert registered.result_schema(two, "score:nll")["nll"].dims == ("N_cand", "D_noise", "L_total") + logits = registered.result_schema(two, "score:logits")["logits"].dims + assert logits == ("N_cand", "D_noise", "L_total", "alphabet") + + +def test_proofread_refuses_a_noise_sweep( + registered: LaserDriver, + pdb: Path, + model_path: Path, +) -> None: + """Proofread reduces over an ensemble; there is no honest noise axis, so D >= 2 is refused.""" + spec = _score_spec(pdb, model_path, "proofread_unconditional", backbone_noise=(0.0, _LEVEL)) + with pytest.raises(ValueError, match="proofread"): + registered.result_schema(spec, "score:proofread_unconditional") + + +def test_sample_two_levels_add_axis_and_level_zero_is_noise_free( + registered: LaserDriver, + pdb: Path, + model_path: Path, +) -> None: + del registered + clean = _sample(pdb, model_path) + noised = _sample(pdb, model_path, backbone_noise=_LEVEL) + both = _sample(pdb, model_path, backbone_noise=(0.0, _LEVEL)) + + n = clean["sequence"].shape[0] + length = clean["sequence"].shape[1] + assert both["sequence"].shape == (n, 2, length), both["sequence"].shape + assert both["seq_log_prob"].shape == (n, 2, length) + assert both["chi_deg"].shape == (n, 2, length, 4) + np.testing.assert_array_equal(both["sequence"][:, 0], clean["sequence"]) + for name in ("seq_log_prob", "chi_deg"): + np.testing.assert_allclose(both[name][:, 0], clean[name], rtol=0, atol=0, equal_nan=True) + np.testing.assert_allclose(both[name][:, 1], noised[name], rtol=0, atol=0, equal_nan=True) + assert not np.allclose(both["seq_log_prob"][:, 0], both["seq_log_prob"][:, 1]) + + +def test_sample_schema_declares_noise_axis_only_for_two_levels(registered: LaserDriver) -> None: + def _spec(level: object) -> SamplingSpecification: + return SamplingSpecification( + inputs="a.pdb", + model_family="lasermpnn", + backbone_noise=level, # type: ignore[arg-type] + ) + + one = registered.result_schema(_spec(_LEVEL), "sample") + two = registered.result_schema(_spec((0.0, _LEVEL)), "sample") + assert one["sequence"].dims == ("N", "L_total") + assert two["sequence"].dims == ("N", "D_noise", "L_total") + assert two["chi_deg"].dims == ("N", "D_noise", "L_total", "4") + assert two["sidechain_coords"].dims == ("N", "D_noise", "L_total", "atom", "xyz") diff --git a/tests/host/test_length_bucketing_runner.py b/tests/host/test_length_bucketing_runner.py index 61c7e6a8..2cde11aa 100644 --- a/tests/host/test_length_bucketing_runner.py +++ b/tests/host/test_length_bucketing_runner.py @@ -197,8 +197,8 @@ def test_averaged_score_path_is_not_bucketed() -> None: This path draws backbone noise at the coordinate shape, so a trim would change the noise draw and the score. Before the fix it was trimmed despite being documented as - unbucketed. One noise level: with two or more, this path fails its R3 check under - vmap (TracerBoolConversionError) with bucketing on or off alike (aminx debt #2517). + unbucketed. One noise level here; the two-level case (formerly a TracerBoolConversionError, + debt #2517, fixed by hoisting the R3 check) is in test_score_noise_axis.py. """ pytest.importorskip("aminx") from aminx.host.runner import score # noqa: PLC0415 @@ -228,10 +228,10 @@ def once(length_bucketing: bool) -> np.ndarray: def test_g_invariance_comparison_can_fail() -> None: """Negative control: the rtol=1e-6 comparison detects a real change. - Backbone noise is not a usable control: the plain score path forwards no noise - level (runner.score, encode-once path #147), so 1ubq scores 4.4554 at noise 0, - 0.1 and 1.0 alike (aminx debt #2509). Scoring a different sequence through the - bucketed path does change the score, and the comparison must catch it. + Backbone noise would be a poor control here: the plain score path now applies it, but + only through the same bucketed/un-bucketed code on both sides, so it would not isolate the + comparison. Scoring a different sequence through the bucketed path does change the + score, and the comparison must catch it. """ pytest.importorskip("aminx") bucketed, opted_out = _score_pair( diff --git a/tests/host/test_score_noise_axis.py b/tests/host/test_score_noise_axis.py new file mode 100644 index 00000000..15cbf67a --- /dev/null +++ b/tests/host/test_score_noise_axis.py @@ -0,0 +1,209 @@ +"""``runner.score`` honours a backbone-noise level axis (spec 261011_noise-level-axis, N-D1..N-D4). + +Replaces the #214 refusal test (debt #2509): the plain and fused multi-state score paths now +apply every ``spec.backbone_noise`` level, drawing the noise from the key ``runner.sample`` uses +for the same spec/seed. + +Acceptance covered: (1) noise 0 is unchanged, (2) D >= 2 gives ``(S, C, D)`` whose ``[..., d]`` +slice equals a single-level run, (3) score and sample noise the backbone identically, (4) nonzero +noise changes the score (a control the old silent-drop behaviour would fail). +""" + +from __future__ import annotations + +import jax +import numpy as np +import pytest + +from aminx.host._sampling_grid_lineage import _base_sampling_key +from aminx.host.runner import _score_noise_key, score +from aminx.run.specs import SamplingSpecification, ScoringSpecification + +_STRUCTURE = "tests/data/1ubq.pdb" +_PROTEIN_CHECKPOINT = "proteinmpnn_v_48_020" +_UBIQUITIN = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG" +_CANDIDATES = [_UBIQUITIN, "A" * len(_UBIQUITIN)] +_LEVELS = (0.0, 0.1, 0.3) +_ATOL = 2e-5 # float32 reduction-order noise between a vmapped and an un-vmapped level (cf. encode-once) + + +def _spec(**kwargs: object) -> ScoringSpecification: + kwargs.setdefault("random_seed", 7) + kwargs.setdefault("max_length", 128) + return ScoringSpecification( + inputs=[_STRUCTURE], + checkpoint_id=_PROTEIN_CHECKPOINT, + sequences_to_score=list(_CANDIDATES), + return_logits=True, + **kwargs, # type: ignore[arg-type] + ) + + +def _sample_spec() -> SamplingSpecification: + return SamplingSpecification( + inputs=[_STRUCTURE], + checkpoint_id=_PROTEIN_CHECKPOINT, + max_length=128, + random_seed=7, + backbone_noise=_LEVELS, + ) + + +def test_noise_key_is_designs_key_without_weights() -> None: + """N-D2 (weight-free): score's per-structure key equals design's ``encode_key`` for the same seed.""" + score_spec = _spec(backbone_noise=_LEVELS) + sample_spec = _sample_spec() + for structure_idx in (0, 1, 5): + design_key = jax.random.fold_in( + _base_sampling_key(sample_spec, grid_lineage=None), + structure_idx, + ) + np.testing.assert_array_equal( + jax.random.key_data(_score_noise_key(score_spec, structure_idx)), + jax.random.key_data(design_key), + ) + # Negative control: a different seed or a different structure index gives a different key. + other_seed = _spec(backbone_noise=_LEVELS, random_seed=8) + assert not np.array_equal( + jax.random.key_data(_score_noise_key(other_seed, 0)), + jax.random.key_data(_score_noise_key(score_spec, 0)), + ) + assert not np.array_equal( + jax.random.key_data(_score_noise_key(score_spec, 1)), + jax.random.key_data(_score_noise_key(score_spec, 0)), + ) + + +@pytest.mark.slow +@pytest.mark.requires_weights +def test_score_and_sample_noise_the_backbone_identically() -> None: + """Acceptance 3: the encoder output at (structure 0, level d) is the same in score and sample. + + Seam: design encodes with ``plan.encode(bundle, fold_in(base_key, structure_idx), config)``; + score encodes with ``make_score_split_fns(...).encode_structure(..., noise, noise_key)``. + """ + from aminx.host.plan import make_inference_plan # noqa: PLC0415 + from aminx.host.prep import prep_protein_stream_and_model # noqa: PLC0415 + from aminx.inference.bundle_builder import build_inference_bundle # noqa: PLC0415 + from aminx.scoring.score import make_score_split_fns # noqa: PLC0415 + + score_spec = _spec(backbone_noise=_LEVELS) + sample_spec = _sample_spec() + iterator, model = prep_protein_stream_and_model(score_spec) + (ensemble,) = list(iterator) + coords, mask = ensemble.coordinates[0], ensemble.mask[0] + residue_index, chain_index = ensemble.residue_index[0], ensemble.chain_index[0] + + plan = make_inference_plan(model, sample_spec, purpose="sample") + split_fns = make_score_split_fns(model) + design_key = jax.random.fold_in(_base_sampling_key(sample_spec, grid_lineage=None), 0) + score_key = _score_noise_key(score_spec, 0) + + encodings = [] + for level in _LEVELS: + bundle, config = build_inference_bundle( + coords=coords, + mask=mask, + residue_index=residue_index, + chain_index=chain_index, + backbone_noise=level, + mode="score_conditional", + inference=True, + ) + designed = plan.encode(bundle, design_key, config) + scored = split_fns.encode_structure( + coords, + mask, + residue_index, + chain_index, + backbone_noise=level, + noise_key=score_key, + ) + np.testing.assert_allclose( + np.asarray(scored.node_features), + np.asarray(designed.node_features), + atol=_ATOL, + ) + np.testing.assert_allclose( + np.asarray(scored.edge_features), + np.asarray(designed.edge_features), + atol=_ATOL, + ) + encodings.append(np.asarray(scored.node_features)) + # Negative control: noise really moved the encoding, so the equalities above are not vacuous. + assert not np.allclose(encodings[0], encodings[2], atol=1e-4) + + +@pytest.mark.slow +@pytest.mark.requires_weights +def test_multi_level_score_has_noise_axis_and_slices_match_single_level() -> None: + """Acceptance 2 (+ N-D1): ``(S, C, D)`` scores, ``(S, C, D, L, 21)`` logits; slices match.""" + multi = score(_spec(backbone_noise=_LEVELS)) + n_candidates = len(_CANDIDATES) + assert np.asarray(multi["scores"]).shape == (1, n_candidates, len(_LEVELS)) + logits = np.asarray(multi["logits"]) + assert logits.shape[:3] == (1, n_candidates, len(_LEVELS)) + assert logits.shape[-1] == 21 + assert tuple(multi["noise_levels"]) == _LEVELS + assert tuple(multi["metadata"]["noise_levels"]) == _LEVELS + assert multi["schema_version"] == "scoring_v2" + + for d, level in enumerate(_LEVELS): + single = score(_spec(backbone_noise=(level,))) + assert np.asarray(single["scores"]).shape == (1, n_candidates) # one level keeps 2-D + np.testing.assert_allclose( + np.asarray(multi["scores"])[..., d], + np.asarray(single["scores"]), + atol=_ATOL, + ) + np.testing.assert_allclose(logits[:, :, d], np.asarray(single["logits"]), atol=_ATOL) + + +@pytest.mark.slow +@pytest.mark.requires_weights +def test_nonzero_noise_changes_score_and_zero_is_unchanged() -> None: + """Acceptances 1 and 4: level 0 equals the default run; a nonzero level moves every candidate.""" + default = np.asarray(score(_spec())["scores"]) + explicit_zero = np.asarray(score(_spec(backbone_noise=0.0))["scores"]) + np.testing.assert_array_equal(explicit_zero, default) # bit-identical, same code path + noised = np.asarray(score(_spec(backbone_noise=0.3))["scores"]) + assert noised.shape == default.shape + assert np.all(np.abs(noised - default) > 1e-4) # the old silent-drop bug leaves these equal + + +@pytest.mark.slow +@pytest.mark.requires_weights +def test_noised_score_is_deterministic_in_the_seed() -> None: + """Same spec/seed gives the same noised score; a different seed gives a different draw.""" + first = np.asarray(score(_spec(backbone_noise=0.3))["scores"]) + again = np.asarray(score(_spec(backbone_noise=0.3))["scores"]) + other = np.asarray(score(_spec(backbone_noise=0.3, random_seed=8))["scores"]) + np.testing.assert_array_equal(first, again) + assert not np.allclose(first, other, atol=1e-5) + + +@pytest.mark.slow +@pytest.mark.requires_weights +def test_fused_multistate_score_gets_the_noise_axis() -> None: + """Fused multi-state keeps its leading 1: ``(1, C)`` for one level, ``(1, C, D)`` for several.""" + # The fused path needs the map width to equal the padded length (as test_score_multistate_fusion does). + length = len(_UBIQUITIN) + positions = [list(range(length))] + multi = score(_spec(backbone_noise=_LEVELS, state_position_map=positions, max_length=length)) + assert np.asarray(multi["scores"]).shape == (1, len(_CANDIDATES), len(_LEVELS)) + single = score(_spec(backbone_noise=(0.3,), state_position_map=positions, max_length=length)) + assert np.asarray(single["scores"]).shape == (1, len(_CANDIDATES)) + np.testing.assert_allclose( + np.asarray(multi["scores"])[..., 2], + np.asarray(single["scores"]), + atol=_ATOL, + ) + + +@pytest.mark.slow +@pytest.mark.requires_weights +def test_averaged_score_with_two_levels_no_longer_crashes() -> None: + """N-D4 / #2517: the R3 check is hoisted out of the vmap, so D >= 2 fuses to one ``(S, C)`` score.""" + result = score(_spec(average_node_features=True, backbone_noise=(0.1, 0.3))) + assert np.asarray(result["scores"]).shape == (1, len(_CANDIDATES)) + assert tuple(result["noise_levels"]) == (0.1, 0.3) diff --git a/tests/knob_semantics/test_potts_noise_axis.py b/tests/knob_semantics/test_potts_noise_axis.py new file mode 100644 index 00000000..cd36fc3c --- /dev/null +++ b/tests/knob_semantics/test_potts_noise_axis.py @@ -0,0 +1,185 @@ +# ruff: noqa: S101 +"""PottsMPNN noise-level axis (spec 261011 N-D5, acceptance 1 and 5). + +``backbone_noise`` takes every level. Levels share ONE draw (key ``0x9077``, the level is not folded in) +scaled by the level, so: + +* one level, however large, is bit-identical to the old first-level-only behaviour and keeps its shape; +* two or more levels add a ``D`` axis right after the candidate/sample axis, and slice ``d`` equals a + single-level run at level ``d`` (level 0.0 is the noise-free run). + +The model is a random-init ``PottsMPNN`` serialised to ``tmp_path``; no checkpoint is needed. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import equinox as eqx +import jax +import numpy as np +import pytest + +from aminx.families.potts_mpnn.driver import PottsMPNNDriver +from aminx.families.potts_mpnn.model import PottsMPNN +from aminx.host.family_driver import FAMILY_DRIVERS +from aminx.host.runner import sample, score +from aminx.run.options import PottsMPNNOptions +from aminx.run.specs import SamplingSpecification, ScoringSpecification + +_THREE = {"A": "ALA", "C": "CYS", "D": "ASP", "E": "GLU", "F": "PHE", "G": "GLY"} +_LEVEL = 0.5 + + +def _write_pdb(path: Path, sequence: str) -> None: + lines: list[str] = [] + serial = 1 + for index, amino in enumerate(sequence, start=1): + for offset, atom in enumerate(("N", "CA", "C", "O")): + x, y, z = float(index * 3 + offset), float(ord("A")), float(offset) + lines.append( + f"{'ATOM':<6.6}{serial:5d} {atom:>4.4}{' ':1.1}{_THREE[amino]:>3.3} {'A':1.1}" + f"{index:4d}{' ':1.1} {x:8.3f}{y:8.3f}{z:8.3f}", + ) + serial += 1 + lines.append("END") + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +@pytest.fixture +def registered() -> Iterator[PottsMPNNDriver]: + driver = PottsMPNNDriver() + FAMILY_DRIVERS.register("pottsmpnn")(driver) + yield driver + FAMILY_DRIVERS.discard("pottsmpnn") + + +@pytest.fixture +def model_path(tmp_path: Path) -> Path: + path = tmp_path / "tiny.eqx" + eqx.tree_serialise_leaves(path, PottsMPNN(key=jax.random.PRNGKey(0))) + return path + + +@pytest.fixture +def pdb(tmp_path: Path) -> Path: + path = tmp_path / "toy.pdb" + _write_pdb(path, "ACDEFG") + return path + + +def _score_spec(pdb: Path, model_path: Path, **kwargs: Any) -> ScoringSpecification: # noqa: ANN401 + return ScoringSpecification( + inputs=str(pdb), + model_family="pottsmpnn", + checkpoint_id="pottsmpnn_vanilla_20", + model_local_path=model_path, + output_kind="energy", + sequences_to_score=["ADDEFG"], + potts_mpnn=PottsMPNNOptions(), + **kwargs, + ) + + +def _energy(pdb: Path, model_path: Path, **kwargs: Any) -> np.ndarray: # noqa: ANN401 + spec = _score_spec(pdb, model_path, **kwargs) + return np.asarray(score(spec)["structures"]["0"]["arrays"]["energy"], dtype=np.float64) + + +def _sample( + pdb: Path, + model_path: Path, + options: PottsMPNNOptions, + **kwargs: Any, # noqa: ANN401 +) -> dict[str, Any]: + spec = SamplingSpecification( + inputs=str(pdb), + model_family="pottsmpnn", + checkpoint_id="pottsmpnn_vanilla_20", + model_local_path=model_path, + num_samples=3, + samples_chunk_size=2, + return_logits=False, + potts_mpnn=options, + **kwargs, + ) + return sample(spec)["structures"]["0"]["arrays"] + + +def test_score_single_level_keeps_shape_and_values( + registered: PottsMPNNDriver, + pdb: Path, + model_path: Path, +) -> None: + """A one-level tuple is the old scalar: same shape, same numbers, and it still changes the energy.""" + del registered + clean = _energy(pdb, model_path) + scalar = _energy(pdb, model_path, backbone_noise=_LEVEL) + one = _energy(pdb, model_path, backbone_noise=(_LEVEL,)) + assert clean.shape == scalar.shape == one.shape == (2,) + np.testing.assert_array_equal(scalar, one) + assert not np.allclose(scalar, clean), "nonzero single-level noise must still reach the energy" + + +def test_score_two_levels_add_axis_and_slices_match_single_runs( + registered: PottsMPNNDriver, + pdb: Path, + model_path: Path, +) -> None: + del registered + clean = _energy(pdb, model_path) + noised = _energy(pdb, model_path, backbone_noise=_LEVEL) + both = _energy(pdb, model_path, backbone_noise=(0.0, _LEVEL)) + assert both.shape == (2, 2), both.shape # (C, D) + np.testing.assert_array_equal(both[:, 0], clean) + np.testing.assert_allclose(both[:, 1], noised, rtol=1e-6, atol=1e-6) + assert not np.allclose(both[:, 0], both[:, 1]) + + +def test_score_schema_declares_noise_axis_only_for_two_levels( + registered: PottsMPNNDriver, + pdb: Path, + model_path: Path, +) -> None: + one_spec = _score_spec(pdb, model_path, backbone_noise=_LEVEL) + two_spec = _score_spec(pdb, model_path, backbone_noise=(0.0, _LEVEL)) + assert registered.result_schema(one_spec, "score:energy")["energy"].dims == ("N_cand",) + assert registered.result_schema(two_spec, "score:energy")["energy"].dims == ("N_cand", "D_noise") + + +def test_sample_two_levels_add_axis_and_level_zero_is_noise_free( + registered: PottsMPNNDriver, + pdb: Path, + model_path: Path, +) -> None: + del registered + options = PottsMPNNOptions(optimization_mode="none") + clean = _sample(pdb, model_path, options) + noised = _sample(pdb, model_path, options, backbone_noise=_LEVEL) + both = _sample(pdb, model_path, options, backbone_noise=(0.0, _LEVEL)) + + n = clean["sequence"].shape[0] + assert clean["sample_energy"].shape == (n,) + assert both["sequence"].shape == (n, 2, 6), both["sequence"].shape + assert both["sample_energy"].shape == (n, 2), both["sample_energy"].shape + # The runner ranks along axis N, so D >= 2 yields one ranking per level. + assert both["sample_rank"].shape == (n, 2), both["sample_rank"].shape + np.testing.assert_array_equal(both["sample_rank"][:, 0], clean["sample_rank"]) + np.testing.assert_array_equal(both["sequence"][:, 0], clean["sequence"]) + np.testing.assert_allclose(both["sample_energy"][:, 0], clean["sample_energy"], rtol=1e-6) + np.testing.assert_allclose(both["sample_energy"][:, 1], noised["sample_energy"], rtol=1e-5) + assert not np.allclose(both["sample_energy"][:, 0], both["sample_energy"][:, 1]) + + +def test_sample_refuses_emit_etab_with_a_noise_sweep( + registered: PottsMPNNDriver, + pdb: Path, + model_path: Path, +) -> None: + """``potts_etab`` has no sample axis to put a noise axis after, so D >= 2 is refused, not truncated.""" + del registered + options = PottsMPNNOptions(optimization_mode="none", emit_etab=True) + with pytest.raises(ValueError, match="emit_etab"): + _sample(pdb, model_path, options, backbone_noise=(0.0, _LEVEL)) diff --git a/tests/protonpotts/test_noise_axis.py b/tests/protonpotts/test_noise_axis.py new file mode 100644 index 00000000..b80bbd9a --- /dev/null +++ b/tests/protonpotts/test_noise_axis.py @@ -0,0 +1,247 @@ +"""ProtonPottsMPNN backbone-noise axis (spec 261011 N-D5, acceptance 1 and 5). + +Before this, the ProtonPotts driver ignored ``backbone_noise`` for every purpose: a nonzero level changed +nothing and nothing said so. Upstream PottsMPNN-pH featurizes through the same ``structure_noise`` path as +PottsMPNN (its pH engine passes 0.0), so the driver now reuses ``apply_potts_backbone_noise`` on the featurized +coordinates with its own key (fold ``0x7070``). Levels share one draw scaled by the level. + +* the control the old silent drop would fail: a nonzero level changes ``score:energy``; +* one level keeps today's shape; two or more add a ``D`` axis right after the candidate axis, slice 0 (level + 0.0) is the noise-free run and slice ``d`` equals a single-level run at level ``d``; +* pH design (``sample``) cannot carry the axis and refuses D >= 2. + +The model is random, so these check the driver contract, not agreement with upstream. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import equinox as eqx +import jax +import numpy as np +import pytest + +from aminx.families.potts_mpnn.model import PottsMPNN +from aminx.families.protonpotts_mpnn import ProtonPottsDriver +from aminx.families.protonpotts_mpnn.vocab import PROTONPOTTS_V6 +from aminx.host.family_driver import FAMILY_DRIVERS +from aminx.host.runner import sample, score +from aminx.run.options import ProtonPottsOptions +from aminx.run.specs import SamplingSpecification, ScoringSpecification + +RESIDUES = ["ALA", "HIS", "ASP", "GLU", "HIS", "ALA", "GLY", "LYS", "SER", "THR"] +_LEVEL = 0.5 + + +def _atom(serial: int, name: str, resname: str, number: int, xyz: tuple[float, float, float]) -> str: + padded = f" {name:<3}" + return ( + f"ATOM {serial:>5} {padded} {resname:>3} A{number:>4} " + f"{xyz[0]:>8.3f}{xyz[1]:>8.3f}{xyz[2]:>8.3f}{1.0:>6.2f}{20.0:>6.2f} {name[0]:>2}" + ) + + +@pytest.fixture +def pdb(tmp_path: Path) -> Path: + rng = np.random.default_rng(0) + lines: list[str] = [] + for i, resname in enumerate(RESIDUES, start=1): + centre = np.array([3.8 * i, 2.0 * np.sin(i), 2.0 * np.cos(i)]) + for k, atom in enumerate(("N", "CA", "C", "O")): + xyz = centre + rng.normal(scale=0.4, size=3) + np.array([0.6 * k, 0.0, 0.0]) + lines.append(_atom(len(lines) + 1, atom, resname, i, (xyz[0], xyz[1], xyz[2]))) + path = tmp_path / "toy.pdb" + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path + + +@pytest.fixture +def registered() -> Iterator[ProtonPottsDriver]: + driver = ProtonPottsDriver() + FAMILY_DRIVERS.register("protonpottsmpnn")(driver) + yield driver + FAMILY_DRIVERS.discard("protonpottsmpnn") + + +@pytest.fixture +def model_path(tmp_path: Path) -> Path: + path = tmp_path / "tiny.eqx" + eqx.tree_serialise_leaves(path, PottsMPNN(key=jax.random.PRNGKey(0), alphabet=PROTONPOTTS_V6)) + return path + + +@pytest.fixture +def variants(tmp_path: Path) -> str: + path = tmp_path / "variants.json" + path.write_text(json.dumps({"his_protonated": {"A:2": "HIS-P"}}), encoding="utf-8") + return str(path) + + +def _score(pdb: Path, model_path: Path, kind: str, **kwargs: Any) -> dict[str, np.ndarray]: # noqa: ANN401 + options = kwargs.pop("options") + spec = ScoringSpecification( + inputs=str(pdb), + model_family="protonpottsmpnn", + model_local_path=model_path, + output_kind=kind, # type: ignore[arg-type] + protonpotts=options, + **kwargs, + ) + return score(spec)["structures"]["0"]["arrays"] + + +def _energy(pdb: Path, model_path: Path, variants: str, **kwargs: Any) -> np.ndarray: # noqa: ANN401 + options = ProtonPottsOptions(variants_json=variants) + arrays = _score(pdb, model_path, "energy", options=options, **kwargs) + return np.asarray(arrays["energy"], dtype=np.float64) + + +def test_nonzero_noise_changes_energy_which_the_old_silent_ignore_would_fail( + registered: ProtonPottsDriver, + pdb: Path, + model_path: Path, + variants: str, +) -> None: + del registered + clean = _energy(pdb, model_path, variants) + noised = _energy(pdb, model_path, variants, backbone_noise=_LEVEL) + assert clean.shape == noised.shape == (2,) # reference + one variant, no noise axis for one level + assert not np.allclose(clean, noised), "backbone_noise reached nothing: ProtonPotts ignored it" + np.testing.assert_array_equal(noised, _energy(pdb, model_path, variants, backbone_noise=(_LEVEL,))) + np.testing.assert_array_equal(clean, _energy(pdb, model_path, variants, backbone_noise=0.0)) + + +def test_noise_is_seeded( + registered: ProtonPottsDriver, + pdb: Path, + model_path: Path, + variants: str, +) -> None: + del registered + first = _energy(pdb, model_path, variants, backbone_noise=_LEVEL, random_seed=3) + again = _energy(pdb, model_path, variants, backbone_noise=_LEVEL, random_seed=3) + other = _energy(pdb, model_path, variants, backbone_noise=_LEVEL, random_seed=4) + np.testing.assert_array_equal(first, again) + assert not np.allclose(first, other) + + +def test_energy_two_levels_add_axis_and_slices_match_single_runs( + registered: ProtonPottsDriver, + pdb: Path, + model_path: Path, + variants: str, +) -> None: + del registered + clean = _energy(pdb, model_path, variants) + noised = _energy(pdb, model_path, variants, backbone_noise=_LEVEL) + both = _energy(pdb, model_path, variants, backbone_noise=(0.0, _LEVEL)) + assert both.shape == (2, 2), both.shape # (C, D) + np.testing.assert_array_equal(both[:, 0], clean) + np.testing.assert_allclose(both[:, 1], noised, rtol=1e-6, atol=1e-6) + assert not np.allclose(both[:, 0], both[:, 1]) + + +def test_ddg_two_levels_add_axis( + registered: ProtonPottsDriver, + pdb: Path, + model_path: Path, + variants: str, +) -> None: + del registered + options = ProtonPottsOptions(variants_json=variants) + clean = _score(pdb, model_path, "ddg", options=options)["ddg"] + both = _score(pdb, model_path, "ddg", options=options, backbone_noise=(0.0, _LEVEL)) + assert np.asarray(clean).shape == (1,) + assert both["ddg"].shape == (1, 2) + assert both["ddg_expt"].shape == (1,) # reference values do not depend on the noise + np.testing.assert_array_equal(both["ddg"][:, 0], clean) + + +def test_selectivity_two_levels_add_axis( + registered: ProtonPottsDriver, + pdb: Path, + model_path: Path, + variants: str, +) -> None: + del registered + options = ProtonPottsOptions( + binder_chain="A", + explicit_centers=((2, "HIS-P"),), + center_types=(), + variants_json=variants, + ) + clean = _score(pdb, model_path, "selectivity", options=options) + noised = _score(pdb, model_path, "selectivity", options=options, backbone_noise=_LEVEL) + both = _score(pdb, model_path, "selectivity", options=options, backbone_noise=(0.0, _LEVEL)) + assert clean["selectivity"].shape == (2,) + assert clean["selectivity_per_centre"].shape == (2, 1) + assert both["selectivity"].shape == (2, 2) + assert both["selectivity_per_centre"].shape == (2, 2, 1) + assert both["candidate_tokens"].shape == clean["candidate_tokens"].shape + np.testing.assert_array_equal(both["selectivity"][:, 0], clean["selectivity"]) + np.testing.assert_allclose(both["selectivity"][:, 1], noised["selectivity"], rtol=1e-5, atol=1e-6) + assert not np.allclose(clean["selectivity"], noised["selectivity"]) + + +def test_schema_declares_noise_axis_only_for_two_levels( + registered: ProtonPottsDriver, + pdb: Path, + model_path: Path, + variants: str, +) -> None: + options = ProtonPottsOptions(variants_json=variants) + + def _spec(level: object) -> ScoringSpecification: + return ScoringSpecification( + inputs=str(pdb), + model_family="protonpottsmpnn", + model_local_path=model_path, + output_kind="energy", + protonpotts=options, + backbone_noise=level, # type: ignore[arg-type] + ) + + assert registered.result_schema(_spec(_LEVEL), "score:energy")["energy"].dims == ("N_cand",) + two = registered.result_schema(_spec((0.0, _LEVEL)), "score:energy")["energy"].dims + assert two == ("N_cand", "D_noise") + + +def test_ph_design_refuses_a_noise_sweep_but_takes_one_level( + registered: ProtonPottsDriver, + pdb: Path, + model_path: Path, +) -> None: + """Designs are energy-sorted and the placement plan is chosen from the noised table: no honest D axis.""" + options = ProtonPottsOptions( + binder_chain="A", + center_types=("HIS-P",), + block_size=2, + samples_per_site=2, + neighbour_k=0, + max_mutations=0, + block_max_rounds=3, + temperature=0.05, + ) + + def _spec(**kwargs: Any) -> SamplingSpecification: # noqa: ANN401 + return SamplingSpecification( + inputs=str(pdb), + model_family="protonpottsmpnn", + model_local_path=model_path, + random_seed=42, + protonpotts=options, + **kwargs, + ) + + with pytest.raises(ValueError, match="noise axis"): + registered.result_schema(_spec(backbone_noise=(0.0, _LEVEL)), "sample") + with pytest.raises(ValueError, match="noise axis"): + sample(_spec(backbone_noise=(0.0, _LEVEL))) + clean = sample(_spec())["structures"]["0"]["arrays"] + one = sample(_spec(backbone_noise=_LEVEL))["structures"]["0"]["arrays"] + assert clean["sequence"].shape == one["sequence"].shape # one level: unchanged shape + assert not np.allclose(clean["final_potts_energy"], one["final_potts_energy"])