Backbone noise-level axis for score and every driver family (replaces the score refusal, #2509/#2517)

.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.

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),

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,38 +630,57 @@ 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):
     if int(row.shape[0]) != length:
       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,15 +929,19 @@ 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),
         default_batch_size=1,
         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