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

src/proteinfoundation/flow_matching/product_space_flow_matcher.py

@@ -132,6 +132,46 @@ def interpolate(
         }
         return x_t
 
+    def seed_state(
+        self,
+        clean: dict[str, torch.Tensor],
+        mask: Bool[Tensor, "* n"],
+        ts: dict[str, Float[Tensor, "nsteps + 1"]],
+        start_step: int,
+        device: torch.device | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Partially-noised product-space state at ``start_step`` from a clean sample — the
+        SDEdit / partial-diffusion analog of :func:`sample_noise` (de-novo init).
+
+        Builds the flow-matching forward marginal at the start step, per channel::
+
+            x_t = (1 - t) * noise + t * clean        t = ts[data_mode][start_step]
+
+        which is exactly the ``t`` :func:`partial_simulation` consumes on its first step
+        (``for step in range(start_step, end_step): t = ts[dm][step]``). Integrating
+        ``partial_simulation`` from ``start_step`` -> ``nsteps`` then yields a *variant* of
+        ``clean`` rather than a de-novo sample. ``start_step`` near ``nsteps`` (t -> 1) stays
+        close to the seed; a smaller ``start_step`` (t -> 0) explores more. There is no
+        physical-Angstrom noise knob -- ``start_step`` (or a ``renoise_frac`` mapped to it)
+        is a unitless schedule index and must be calibrated empirically (Ca-RMSD sweep).
+
+        ``clean`` holds the flow's t=1 endpoints: ``bb_ca`` in **nanometers** (Angstrom * 0.1)
+        and ``local_latents`` in the autoencoder ``mean`` / ``z_latent`` space. Only the
+        *noise* is zero-COM'd (per channel config, inside ``sample_noise``); ``clean`` is used
+        as-is, so the caller's frame is preserved -- place a docked binder in the target frame
+        (do NOT independently zero-COM it) or zero-COM a lone chain to match the noise prior.
+        Self-conditioning state starts empty exactly as in de-novo generation.
+        """
+        n = mask.shape[-1]
+        batch_shape = tuple(mask.shape[:-1])
+        noise = self.sample_noise(n=n, shape=batch_shape, device=device, mask=mask)
+        ones = torch.ones(batch_shape, device=device) if batch_shape else torch.ones((), device=device)
+        t = {
+            data_mode: ts[data_mode][start_step].to(device=device) * ones
+            for data_mode in self.data_modes
+        }
+        return self.interpolate(x_0=noise, x_1=clean, t=t, mask=mask)
+
     def process_batch(self, batch: dict) -> tuple[Tensor, Tensor, tuple, int, torch.dtype]:
         """
         Extracts clean sample, mask, batch size, protein length n, dtype and device

src/proteinfoundation/search/fk_steering.py

@@ -92,6 +92,42 @@ def search(self, batch: dict) -> dict:
                 f"nsteps ({search_ctx.nsteps}); otherwise final samples are partially "
                 f"denoised and rewards are computed on incomplete structures."
             )
+        # ── partial diffusion (SDEdit): optional seed from an existing structure ──
+        # `search.seed` present  -> start the reverse process from a partially-noised copy
+        #   of a real binder (variant generation / refinement) instead of pure noise.
+        # `start_step` (or `renoise_frac`) picks how far back to noise: larger start_step
+        #   (closer to nsteps) stays nearer the seed. See fm.seed_state for the math.
+        seed_cfg = self.inf_cfg.search.get("seed", None)
+        seed_clean = None
+        if seed_cfg is not None:
+            from proteinfoundation.utils.pdb_utils import encode_seed
+
+            if "start_step" in seed_cfg:
+                start_step = int(seed_cfg["start_step"])
+            elif "renoise_frac" in seed_cfg:
+                # unitless mixing fraction -> schedule index; nonlinear per-channel schedules
+                # make this approximate, so prefer explicit start_step and calibrate by Ca-RMSD.
+                start_step = int(round((1.0 - float(seed_cfg["renoise_frac"])) * search_ctx.nsteps))
+            else:
+                raise ValueError("search.seed requires 'start_step' or 'renoise_frac'")
+            if not (0 < start_step < search_ctx.nsteps):
+                raise ValueError(
+                    f"partial-diffusion start_step must be in (0, nsteps={search_ctx.nsteps}); "
+                    f"got {start_step} (0 == pure noise / de novo, nsteps == unchanged seed)"
+                )
+            # re-space checkpoints onto [start_step, nsteps] so FK still resamples,
+            # but the trajectory begins mid-way (end stays nsteps, validated above).
+            step_checkpoints = [start_step] + [c for c in step_checkpoints if c > start_step]
+            if step_checkpoints[-1] != search_ctx.nsteps:
+                step_checkpoints.append(search_ctx.nsteps)
+            seed_clean = encode_seed(
+                self.proteina.autoencoder,
+                pdb_path=seed_cfg["pdb_path"],
+                chain=seed_cfg.get("chain", None),
+                target_chain=seed_cfg.get("target_chain", None),  # align seed to target-centered frame
+                device=search_ctx.device,
+            )
+
         n_steps_total = len(step_checkpoints) - 1
 
         # ── initialise noise + tags ─────────────────────────────────────
@@ -101,14 +137,36 @@ def search(self, batch: dict) -> dict:
             f"[FKSteering] Starting | nsamples={nsamples}, "
             f"beam_width={beam_width}, n_branch={n_branch}, "
             f"temperature={temperature}, checkpoints={step_checkpoints}"
+            + (f", SEED start_step={step_checkpoints[0]}" if seed_clean is not None else "")
         )
         mask = init_mask.repeat_interleave(beam_width, dim=0)
-        xt = self.proteina.fm.sample_noise(
-            n,
-            shape=(nsamples * beam_width,),
-            device=search_ctx.device,
-            mask=mask,
-        )
+        if seed_clean is None:
+            xt = self.proteina.fm.sample_noise(
+                n,
+                shape=(nsamples * beam_width,),
+                device=search_ctx.device,
+                mask=mask,
+            )
+        else:
+            # broadcast the single seed to every replica, then noise to start_step
+            total = nsamples * beam_width
+            for k, v in seed_clean.items():
+                if v.shape[-2] != n:
+                    raise ValueError(
+                        f"seed channel '{k}' has {v.shape[-2]} residues but the binder mask "
+                        f"has n={n}; the seed structure must match the requested binder length"
+                    )
+            clean = {
+                k: v.to(search_ctx.device)[None].expand(total, *v.shape).contiguous()
+                for k, v in seed_clean.items()
+            }
+            xt = self.proteina.fm.seed_state(
+                clean=clean,
+                mask=mask,
+                ts=search_ctx.ts,
+                start_step=step_checkpoints[0],
+                device=search_ctx.device,
+            )
         x_1_pred = None
         metadata_tags = make_initial_search_tags("fk", nsamples, beam_width)
 

src/proteinfoundation/utils/pdb_utils.py

@@ -581,3 +581,101 @@ def load_target_from_pdb(target_spec, pdb_path, target_hotspots=None, convert_an
         target_hotspots_mask,
         target_chain,
     )
+
+
+def encode_seed(
+    autoencoder,
+    pdb_path: str,
+    chain: str | None = None,
+    target_chain: str | None = None,
+    device=None,
+) -> dict:
+    """Encode a protein chain into a clean product-space seed for partial diffusion (SDEdit).
+
+    Returns a ``dict`` ready for ``ProductSpaceFlowMatcher.seed_state``::
+
+        {"bb_ca": [n, 3] Ca coords in NANOMETERS, "local_latents": [n, latent_dim]}
+
+    ``bb_ca`` and ``local_latents`` are the flow's t=1 endpoints: the decoder maps the
+    generated ``local_latents`` back to structure, and the encoder ``mean`` here is the
+    deterministic point in that same latent space, so seeding is space-consistent. The
+    deterministic ``mean`` (not the reparameterized ``z_latent`` sample) is used so the seed
+    is reproducible.
+
+    Frame (docked pose): during generation the *target* is centered on its CA centre-of-mass
+    (:class:`CoordsTensorCenteringTransform`, ``data_mode="bb_ca"``), so a seed binder left in
+    its raw PDB frame would be mis-placed relative to the centered target. Pass ``target_chain``
+    (the target chain in *this* pdb -- e.g. the complex the binder was docked in): the binder
+    is then shifted into that same target-CA-COM frame, so the docked pose is preserved by
+    default. The latent is translation-invariant, so only ``bb_ca`` is shifted. If
+    ``target_chain`` is None the binder is returned in its own frame and a warning is emitted --
+    fine for *unconditioned* SDEdit, but for target-conditioned generation the docked pose may
+    drift unless an epitope/contact constraint re-docks it.
+
+    The encoder's ``FeatureFactory`` runs with ``strict_feats: False``; the load-bearing batch
+    keys are the clean atom37 coords (A and nm), the atom mask and the residue types.
+    """
+    import torch
+
+    from proteinfoundation.utils.coors_utils import ang_to_nm
+
+    target_spec = chain if chain is not None else "A"
+    # Angstrom coords -> the feature factory reads both `coords` (A, for bb/sidechain
+    # angles) and `coords_nm` (nm, for a37coors + pair dists).
+    atom_mask, coords_ang, residue_type, _, _ = load_target_from_pdb(
+        target_spec=target_spec, pdb_path=pdb_path, convert_ang_to_nm=False
+    )
+    n = residue_type.shape[0]
+    dev = device if device is not None else next(autoencoder.parameters()).device
+
+    coords_ang = coords_ang.to(dev).float()          # [n, 37, 3], Angstrom
+    coords_nm = ang_to_nm(coords_ang)                 # [n, 37, 3], nm (repo's exact conversion)
+    atom_mask = atom_mask.to(dev)                     # [n, 37] bool
+    residue_type = residue_type.to(dev).long()        # [n]
+    resmask = torch.ones(1, n, dtype=torch.bool, device=dev)  # [1, n] no padding (single chain)
+
+    # AF2 atom37 layout: index 1 is CA.
+    ca_nm = coords_nm[:, 1, :]                         # [n, 3], nm
+
+    # Full batch matching what the encoder's FeatureFactory reads (nn_130m feats).
+    batch = {
+        "coords": coords_ang[None],                   # [1, n, 37, 3] Angstrom  (bb/sidechain angles)
+        "coords_nm": coords_nm[None],                 # [1, n, 37, 3] nm        (a37coors_nm, pair dists)
+        "coord_mask": atom_mask[None],                # [1, n, 37]
+        "ca_coors_nm": ca_nm[None],                   # [1, n, 3]
+        "residue_type": residue_type[None],           # [1, n]
+        "residue_pdb_idx": torch.arange(n, device=dev)[None],  # [1, n] contiguous
+        "mask": resmask,                              # [1, n]
+        "mask_dict": {                                # feats index these sub-keys directly
+            "residue_type": resmask,                  # [1, n] padding mask (ResidueTypeSeqFeat)
+            "coords": atom_mask[None, ..., None].expand(1, n, 37, 3),  # [1, n, 37, 3]
+        },
+    }
+
+    autoencoder.eval()
+    with torch.no_grad():
+        out = autoencoder.encode(batch)               # {"mean", "log_scale", "z_latent"}
+    latents = out["mean"][0]                           # [n, latent_dim], deterministic (translation-invariant)
+
+    # Place the binder in the target's centered frame so the docked pose survives, matching
+    # CoordsTensorCenteringTransform(data_mode="bb_ca") applied to the target at generation time.
+    if target_chain is not None:
+        from proteinfoundation.utils.align_utils import mean_w_mask
+
+        t_mask, t_ang, _, _, _ = load_target_from_pdb(
+            target_spec=target_chain, pdb_path=pdb_path, convert_ang_to_nm=False
+        )
+        t_ca_nm = ang_to_nm(t_ang.to(dev).float())[:, 1, :]          # [nt, 3] target CA, nm
+        t_ca_mask = t_mask.to(dev)[:, 1]                             # [nt] CA present
+        com = mean_w_mask(t_ca_nm, t_ca_mask, keepdim=False).reshape(3)  # target CA centre-of-mass
+        ca_nm = ca_nm - com                                         # shift binder into target-CA-COM frame
+    else:
+        from loguru import logger
+
+        logger.warning(
+            "encode_seed: no target_chain given -> seed binder kept in its own PDB frame. "
+            "For target-conditioned generation the docked pose may be inconsistent with the "
+            "centered target; pass target_chain (the target chain in this pdb) to auto-align."
+        )
+
+    return {"bb_ca": ca_nm, "local_latents": latents}

tests/test_seed.py

@@ -0,0 +1,91 @@
+"""Tests for SDEdit / partial-diffusion seeding.
+
+Covers ``ProductSpaceFlowMatcher.seed_state`` -- the SDEdit forward marginal that replaces
+pure-noise init in ``search/fk_steering.py``. Pure flow-matching math, no model weights and
+no config files required (a minimal two-channel flow matcher is built inline).
+
+The encoder path (``utils.pdb_utils.encode_seed``) and the ``partial_simulation`` start_step
+integration require model weights and are exercised by a live run, not here.
+"""
+
+import torch
+from omegaconf import OmegaConf
+
+from proteinfoundation.flow_matching.product_space_flow_matcher import ProductSpaceFlowMatcher
+
+
+def _fm():
+    """Minimal two-channel (bb_ca + local_latents) flow matcher, no weights."""
+    cfg = OmegaConf.create(
+        {
+            "product_flowmatcher": {
+                "bb_ca": {"zero_com_noise": True, "guidance_enabled": False, "dim": 3},
+                "local_latents": {"zero_com_noise": False, "guidance_enabled": False, "dim": 8},
+            }
+        }
+    )
+    fm = ProductSpaceFlowMatcher(cfg)
+    fm.eval()
+    return fm
+
+
+def _ts(nsteps):
+    # ts[0] = 0, ts[-1] = 1; seed_state only indexes ts[dm][start_step], so a linear
+    # schedule suffices to test the interpolant math independent of get_schedule internals.
+    sched = torch.linspace(0.0, 1.0, nsteps + 1)
+    return {"bb_ca": sched, "local_latents": sched}
+
+
+def _clean(nsamples, n):
+    return {
+        "bb_ca": torch.randn(nsamples, n, 3),
+        "local_latents": torch.randn(nsamples, n, 8),
+    }
+
+
+def test_seed_state_t1_returns_clean():
+    """start_step == nsteps (t = 1): the seeded state is exactly the clean input (no noise)."""
+    nsamples, n, nsteps = 2, 10, 100
+    fm, ts = _fm(), _ts(100)
+    mask = torch.ones(nsamples, n, dtype=torch.bool)
+    clean = _clean(nsamples, n)
+    s = fm.seed_state(clean=clean, mask=mask, ts=ts, start_step=nsteps)
+    assert torch.allclose(s["bb_ca"], clean["bb_ca"], atol=1e-5)
+    assert torch.allclose(s["local_latents"], clean["local_latents"], atol=1e-5)
+
+
+def test_seed_state_t0_is_pure_noise():
+    """start_step == 0 (t = 0): the seeded state is pure noise, independent of the clean input."""
+    nsamples, n, nsteps = 2, 10, 100
+    fm, ts = _fm(), _ts(nsteps)
+    mask = torch.ones(nsamples, n, dtype=torch.bool)
+    clean = _clean(nsamples, n)
+    clean = {k: v + 100.0 for k, v in clean.items()}  # push clean far away
+    s = fm.seed_state(clean=clean, mask=mask, ts=ts, start_step=0)
+    assert not torch.allclose(s["bb_ca"], clean["bb_ca"], atol=1.0)
+    assert not torch.allclose(s["local_latents"], clean["local_latents"], atol=1.0)
+
+
+def test_seed_state_larger_start_stays_closer():
+    """Larger start_step (t -> 1) stays closer to the clean input than a smaller one."""
+    nsamples, n, nsteps = 2, 10, 200
+    fm, ts = _fm(), _ts(nsteps)
+    mask = torch.ones(nsamples, n, dtype=torch.bool)
+    clean = _clean(nsamples, n)
+    torch.manual_seed(0)
+    near = (fm.seed_state(clean=clean, mask=mask, ts=ts, start_step=180)["local_latents"] - clean["local_latents"]).norm()
+    torch.manual_seed(0)
+    far = (fm.seed_state(clean=clean, mask=mask, ts=ts, start_step=40)["local_latents"] - clean["local_latents"]).norm()
+    assert near < far
+
+
+def test_seed_state_preserves_shape_and_modes():
+    """Output has both channels with the clean shapes (contract for fk_steering init)."""
+    nsamples, n, nsteps = 3, 12, 50
+    fm, ts = _fm(), _ts(nsteps)
+    mask = torch.ones(nsamples, n, dtype=torch.bool)
+    clean = _clean(nsamples, n)
+    s = fm.seed_state(clean=clean, mask=mask, ts=ts, start_step=25)
+    assert set(s) == {"bb_ca", "local_latents"}
+    assert s["bb_ca"].shape == (nsamples, n, 3)
+    assert s["local_latents"].shape == (nsamples, n, 8)