Add partial-diffusion (SDEdit) seeding for binder generation and refinement

[ssiddhantsharma]

Summary

Adds partial diffusion (SDEdit) to the flow-matching binder pipeline: instead of starting
the reverse process from pure noise, seed it from an existing binder at an intermediate
flow time and denoise the rest of the way. This enables variant generation and
refinement of a known design under FK steering — the flow-matching analog of
RFdiffusion's partial_t.

Fully backward-compatible: with no search.seed in the config, generation is unchanged
(de novo from noise).

Motivation

The search stack already supports partial reverse integration (partial_simulation runs
between step_checkpoints), but generation always begins from pure noise, so there was no way
to start from — and stay near — an existing binder. Partial diffusion is the natural way to
(a) sample sequence/backbone variants of a good design and (b) refine a design toward a reward
(affinity/selectivity) under FK steering, which de-novo generation can't do.

What's added

File	Change
flow_matching/product_space_flow_matcher.py	seed_state(clean, mask, ts, start_step) — the SDEdit forward marginal x_t = (1−t)·noise + t·clean at t = ts[dm][start_step], built from the existing interpolate + sample_noise. The partial-diffusion analog of de-novo sample_noise.
search/fk_steering.py	Reads `search.seed {pdb_path, chain, target_chain, start_step
utils/pdb_utils.py	encode_seed(autoencoder, pdb_path, chain, target_chain=None) — encodes a chain to (bb_ca [n,3] nm, local_latents [n,8]) via autoencoder.encode()["mean"]. When target_chain is given, shifts the seed into the target's CA-COM frame (matching CoordsTensorCenteringTransform) so the docked pose is preserved; otherwise warns.
tests/test_seed.py	Weights-free tests for seed_state.

Design notes

• partial_simulation(start_step, end_step, ts) already supports partial reverse — seed_state
  builds x_t at exactly the t its first step consumes (for step in range(start_step, end_step): t = ts[dm][step]), so no changes to the integrator.
• The OT interpolant is (1−t)·x_0 + t·x_1 (x_0 = noise, x_1 = clean), so larger start_step
  (t → 1) stays nearer the seed; smaller explores more. start_step is a unitless schedule
  index — calibrate renoise_frac ↔ Cα-RMSD empirically.
• The latent is translation-invariant, so target_chain alignment shifts only bb_ca.

Usage

generation:
  search:
    algorithm: fk-steering
    step_checkpoints: [0, 25, 50]
    seed:
      pdb_path: /path/to/complex.pdb
      chain: A1-72           # binder chain contig
      target_chain: B73-222  # target chain in the same pdb -> aligns seed to the
                             # target-centered frame (omit for unconditioned SDEdit)
      start_step: 40         # or renoise_frac: 0.2  (larger start_step => nearer the seed)

Validation

Unit tests (tests/test_seed.py, weights-free, python -m pytest tests/test_seed.py -v):

tests/test_seed.py::test_seed_state_t1_returns_clean PASSED              [ 25%]
tests/test_seed.py::test_seed_state_t0_is_pure_noise PASSED             [ 50%]
tests/test_seed.py::test_seed_state_larger_start_stays_closer PASSED    [ 75%]
tests/test_seed.py::test_seed_state_preserves_shape_and_modes PASSED    [100%]
============================== 4 passed in 17.26s ==============================

Component / integration (against the released complexa + complexa_ae checkpoints):

• encode_seed on the real AE: latent_dim=8, encode→decode sequence recovery = 1.000,
  latents unit-scale (global std ≈ 1.2) — consistent with the OT noise prior and the decoder's
  latent space.
• End-to-end: real complexa generate (fk-steering + AF2 reward, target-conditioned) with
  search.seed set — the seed is encoded, step_checkpoints re-spaced [0,25,50]→[40,50],
  seed_state initializes the trajectory, generation completes.
• Target-frame alignment: verified the seed shift equals the target CA-COM and the latent
  is unchanged by it, at the origin and with the complex translated +50 Å off-origin
  (aligned seed identical in both frames → docked pose invariant to absolute placement).