From bb739a8ab3d9bb707ffc6d95be92ac3bb2a6511c Mon Sep 17 00:00:00 2001 From: siddhant Date: Mon, 13 Jul 2026 20:17:37 +0530 Subject: [PATCH] Add partial-diffusion (SDEdit) seeding to fk_steering binder generation Seed the reverse process from an existing binder at an intermediate flow time instead of pure noise (RFdiffusion partial_t analog), enabling variant generation / refinement under FK steering. Backward-compatible: no search.seed => unchanged. - ProductSpaceFlowMatcher.seed_state: SDEdit forward marginal at start_step - fk_steering: reads search.seed, re-spaces step_checkpoints, seeds via seed_state - utils.pdb_utils.encode_seed: chain -> (bb_ca nm, local_latents); optional target_chain aligns the seed to the target CA-COM frame so docked pose survives - tests/test_seed.py: weights-free seed_state tests --- .../product_space_flow_matcher.py | 40 ++++++++ src/proteinfoundation/search/fk_steering.py | 70 +++++++++++-- src/proteinfoundation/utils/pdb_utils.py | 98 +++++++++++++++++++ tests/test_seed.py | 91 +++++++++++++++++ 4 files changed, 293 insertions(+), 6 deletions(-) create mode 100644 tests/test_seed.py diff --git a/src/proteinfoundation/flow_matching/product_space_flow_matcher.py b/src/proteinfoundation/flow_matching/product_space_flow_matcher.py index 4477aed..55a2481 100644 --- a/src/proteinfoundation/flow_matching/product_space_flow_matcher.py +++ b/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 diff --git a/src/proteinfoundation/search/fk_steering.py b/src/proteinfoundation/search/fk_steering.py index 1a91c49..9e1d70f 100644 --- a/src/proteinfoundation/search/fk_steering.py +++ b/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) diff --git a/src/proteinfoundation/utils/pdb_utils.py b/src/proteinfoundation/utils/pdb_utils.py index 657e72b..15e0d7e 100644 --- a/src/proteinfoundation/utils/pdb_utils.py +++ b/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} diff --git a/tests/test_seed.py b/tests/test_seed.py new file mode 100644 index 0000000..38318b4 --- /dev/null +++ b/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)