feat: steering to avoid chain breaks and clashes #171

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

nw13slx merged 71 commits into main from luwinkler/cli_steering

Jan 21, 2026

.gitignore

-Original file line number
+Diff line change
@@ Expand Up / @@ -4,6 +4,8 @@ @@
     #vs code
     .vscode
     .vscode/
+    .github/copilot-instructions.md
+    outputs/
     # Jetbrains
     .idea
@@ Expand Down Expand Up / @@ -132,3 +134,14 @@ cython_debug/ @@
     *.pth
     *.pt
+    # wandb & various
+    wandb
+    wandb/
+    .amltconfig
+    */wandb/*
+    *.npz
+    *.pkl
+    *amlt*
+    *outputs*
+    *cache*

README.md

-Original file line number
+Diff line change
@@ Expand Up / @@ -20,6 +20,7 @@ This repository contains inference code and model weights. @@
     ## Table of Contents
     - [Installation](#installation)
     - [Sampling structures](#sampling-structures)
+    - [Steering to avoid chain breaks and clashes](#steering-to-avoid-chain-breaks-and-clashes)
     - [Azure AI Foundry](#azure-ai-foundry)
     - [Training data](#training-data)
     - [Get in touch](#get-in-touch)
@@ Expand Down Expand Up @@
     This code only supports sampling structures of monomers. You can try to sample multimers using the [linker trick](https://x.com/ag_smith/status/1417063635000598528), but in our limited experiments, this has not worked well.
+    ## Steering to avoid chain breaks and clashes
+    BioEmu includes a [steering system](https://arxiv.org/abs/2501.06848) that uses [Sequential Monte Carlo (SMC)](https://www.stats.ox.ac.uk/~doucet/doucet_defreitas_gordon_smcbookintro.pdf) to guide the diffusion process toward more physically plausible protein structures.
+    Empirically, using three (or up to 10) steering particles per output sample greatly reduces the number of unphysical samples (steric clashes or chain breaks) produced by the model.
+    Steering applies potential energy functions during denoising to favor conformations that satisfy physical constraints.
+    Algorithmically, steering simulates multiple *candidate samples* per desired output sample and resamples between these particles according to the favorability of the provided potentials.
+    ### Quick start with steering
+    Enable steering with physical constraints using the CLI:
+    ```bash
+    python -m bioemu.sample \
+        --sequence GYDPETGTWG \
+        --num_samples 100 \
+        --output_dir ~/steered-samples \
+        --steering_config src/bioemu/config/steering/physical_steering.yaml \
+        --denoiser_config src/bioemu/config/denoiser/stochastic_dpm.yaml
+    ```
+    Or using the Python API:
+    ```python
+    from bioemu.sample import main as sample
+    sample(
+        sequence='GYDPETGTWG',
+        num_samples=100,
+        output_dir='~/steered-samples',
+        denoiser_config="../src/bioemu/config/denoiser/stochastic_dpm.yaml",  # Use stochastic DPM
+        steering_config="../src/bioemu/config/steering/physicality_steering.yaml",  # Use physicality steering
+    )
+    ```
+    ### Key steering parameters
+    - `num_steering_particles`: Number of particles per sample (1 = no steering, >1 enables steering)
+    - `steering_start_time`: When to start steering (0.0-1.0, default: 0.1) with reverse sampling 1 -> 0
+    - `steering_end_time`: When to stop steering (0.0-1.0, default: 0.) with reverse sampling 1 -> 0
+    - `resampling_interval`: How often to resample particles (default: 1)
+    - `steering_config`: Path to potentials configuration file (required for steering)
+    ### Available potentials
+    The [`physical_steering.yaml`](./src/bioemu/config/steering/physical_steering.yaml) configuration provides potentials for physical realism:
+    - **ChainBreak**: Prevents backbone discontinuities
+    - **ChainClash**: Avoids steric clashes between non-neighboring residues
+    - **DisulfideBridge**: Encourages disulfide bond formation between specified cysteine pairs
+    You can create a custom `steering_config.yaml` YAML file instantiating your own potential to steer the system with your own potentials.
     ## Azure AI Foundry
     BioEmu is also available on [Azure AI Foundry](https://ai.azure.com/). See [How to run BioEmu on Azure AI Foundry](AZURE_AI_FOUNDRY.md) for more details.
@@ Expand Down @@

notebooks/disulfide_bridge_steering_example.py

-Original file line number
+Diff line change
@@ -0,0 +1,111 @@
+    """Script to compare sampling with and without physicality steering."""
+    import logging
+    from pathlib import Path
+    import matplotlib.pyplot as plt
+    import numpy as np
+    import torch
+    from bioemu.sample import main as sample_main
+    logging.basicConfig(level=logging.INFO)
+    logger = logging.getLogger(__name__)
+    def bridge_distances(pos: torch.Tensor, bridge_indices: list[tuple[int, int]]) -> torch.Tensor:
+        """Compute Ca-Ca distances for specified disulfide bridge indices.
+        Args:
+                pos (torch.Tensor): Tensor of shape (N, L, 3)
+        """
+        import torch
+        distances = []
+        for i, j in bridge_indices:
+            dist_ij = torch.norm(pos[:, i, :] - pos[:, j, :], dim=-1)  # (N,)
+            distances.append(dist_ij)
+        return torch.stack(distances, dim=-1)  # (N, num_bridges)
+    if __name__ == "__main__":
+        # https://www.uniprot.org/uniprotkb/P01542/entry#sequences
+        # TTCCPSIVARSNFNVCRLPGTPEALCATYTGCIIIPGATCPGDYAN
+        # PTM = [(3,40), (4,32), (16, 26)]
+        bridge_indices = [(2, 39), (3, 31), (15, 25)]  # adjusted by -1 to be 0-indexed
+        """Sample 128 structures with and without physicality steering."""
+        # Configuration
+        sequence = "TTCCPSIVARSNFNVCRLPGTPEALCATYTGCIIIPGATCPGDYAN"  # Example sequence
+        num_samples = 128
+        base_output_dir = Path("comparison_outputs_disulfide")
+        # Sample WITHOUT steering
+        logger.info("=" * 80)
+        logger.info("Sampling WITHOUT steering...")
+        logger.info("=" * 80)
+        output_dir_no_steering = base_output_dir / "no_steering"
+        sample_main(
+            sequence=sequence,
+            num_samples=num_samples,
+            output_dir=output_dir_no_steering,
+            batch_size_100=500,
+            denoiser_config="../src/bioemu/config/denoiser/stochastic_dpm.yaml",  # Use stochastic DPM
+            steering_config=None,  # No steering
+        )
+        pos_unsteered = torch.from_numpy(
+            np.load(list(output_dir_no_steering.glob("batch_*.npz"))[0])["pos"]
+        )
+        unsteered_bridge_distances = bridge_distances(pos_unsteered, bridge_indices)
+        # Sample WITH steering
+        logger.info("=" * 80)
+        logger.info("Sampling WITH physicality steering...")
+        logger.info("=" * 80)
+        output_dir_with_steering = base_output_dir / "with_steering"
+        sample_main(
+            sequence=sequence,
+            num_samples=num_samples,
+            output_dir=output_dir_with_steering,
+            denoiser_config="../src/bioemu/config/denoiser/stochastic_dpm.yaml",  # Use stochastic DPM
+            steering_config="../src/bioemu/config/steering/disulfide_bridge_steering.yaml",  # Use disulfide bridge steering
+        )
+        pos_steered = torch.from_numpy(
+            np.load(list(output_dir_with_steering.glob("batch_*.npz"))[0])["pos"]
+        )
+        steered_bridge_distances = bridge_distances(
+            pos_steered, bridge_indices
+        )  # pos_rot_steered in Angstrom
+        logger.info("=" * 80)
+        logger.info("Comparison complete!")
+        logger.info(f"Results without steering: {output_dir_no_steering}")
+        logger.info(f"Results with steering: {output_dir_with_steering}")
+        logger.info("=" * 80)
+        # Distances are in Angstrom
+        fig, ax = plt.subplots(1, 2, figsize=(16, 8))
+        ax[0].hist(
+            unsteered_bridge_distances.numpy().flatten(), bins=50, alpha=0.5, label="No Steering"
+        )
+        ax[0].hist(
+            steered_bridge_distances.numpy().flatten(), bins=50, alpha=0.5, label="With Steering"
+        )
+        ax[0].legend()
+        ax[0].set_xlim(0, 5)
+        ax[0].set_xlabel("Cα-Cα Distance (nM)")
+        ax[0].grid()
+        ax[1].hist(
+            unsteered_bridge_distances.numpy().flatten(), bins=100, alpha=0.5, label="No Steering"
+        )
+        ax[1].hist(
+            steered_bridge_distances.numpy().flatten(), bins=100, alpha=0.5, label="With Steering"
+        )
+        ax[1].legend()
+        ax[1].set_xlim(0.25, 1)
+        ax[1].set_xlabel("Cα-Cα Distance (nM)")
+        ax[1].grid()

notebooks/physical_steering_example.py

-Original file line number
+Diff line change
@@ -0,0 +1,54 @@
+    """Script to compare sampling with and without physicality steering."""
+    import logging
+    from pathlib import Path
+    from bioemu.sample import main as sample_main
+    logging.basicConfig(level=logging.INFO)
+    logger = logging.getLogger(__name__)
+    def main():
+        """Sample 128 structures with and without physicality steering."""
+        # Configuration
+        sequence = "MTEIAQKLKESNEPILYLAERYGFESQQTLTRTFKNYFDVPPHKYRMTNMQGESRFLHPL"  # Example sequence
+        num_samples = 128
+        base_output_dir = Path("comparison_outputs")
+        # Sample WITHOUT steering
+        logger.info("=" * 80)
+        logger.info("Sampling WITHOUT steering...")
+        logger.info("=" * 80)
+        output_dir_no_steering = base_output_dir / "no_steering"
+        sample_main(
+            sequence=sequence,
+            num_samples=num_samples,
+            output_dir=output_dir_no_steering,
+            denoiser_config="../src/bioemu/config/denoiser/stochastic_dpm.yaml",  # Use stochastic DPM
+            steering_config=None,  # No steering
+        )
+        # Sample WITH steering
+        logger.info("=" * 80)
+        logger.info("Sampling WITH physicality steering...")
+        logger.info("=" * 80)
+        output_dir_with_steering = base_output_dir / "with_steering"
+        sample_main(
+            sequence=sequence,
+            num_samples=num_samples,
+            output_dir=output_dir_with_steering,
+            denoiser_config="../src/bioemu/config/denoiser/stochastic_dpm.yaml",  # Use stochastic DPM
+            steering_config="../src/bioemu/config/steering/physical_steering.yaml",  # Use physicality steering
+        )
+        logger.info("=" * 80)
+        logger.info("Comparison complete!")
+        logger.info(f"Results without steering: {output_dir_no_steering}")
+        logger.info(f"Results with steering: {output_dir_with_steering}")
+        logger.info("=" * 80)
+    if __name__ == "__main__":
+        main()

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: steering to avoid chain breaks and clashes #171

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!