ColabDesign

AfDesign supports several optimization methods for sequence design. The choice of optimizer affects whether the output is a valid discrete sequence or a continuous probability profile (PSSM).

Optimization Types

• pssm_semigreedy (Recommended): Uses a designed PSSM to bias a semi-greedy optimization.
• 3stage: A gradient-based optimization (GD) following the sequence: logits $\rightarrow$ soft $\rightarrow$ hard.
• pssm: GD optimization (logits $\rightarrow$ soft) to generate a sequence profile (PSSM).
• semigreedy: Attempts random mutations and accepts those that decrease loss.
• logits: GD optimization of continuous logits.
• soft: GD optimization of softmax(logits) (probabilities).
• hard: GD optimization of one_hot(logits) (discrete).

WARNING: Optimizers like pssm, logits, and soft do not produce a valid one-hot sequence directly. To obtain a valid sequence, use the other optimizers or redesign the output backbone using a tool like ProteinMPNN.

RFdiffusion requires a 'blueprint' to define the desired secondary structure elements (SSEs) and their interactions. You can generate these in two ways:

1. Manual Mode: You define the number of SSE elements. You then specify the type for each element (Diagonal: H for helix, E for sheet, C for coil, ? for undefined) and the interactions between them (Off-diagonal: 0 for no contact, 1 for contact, ? for undefined). You also specify the length of each SSE and a buff_length (buffer) between them.
2. Automated Mode: You provide an existing PDB file. The system extracts the secondary structure and contact map from the PDB to create the blueprint automatically.

SSE Types:

• H: Helix
• E: Sheet
• C: Coil
• ?: Undefined

# Automated mode example
# blueprint_mode = "automated"
# pdb = "6MRR"
# chain = "A"
# trim_loops = True

# Manual mode example
# blueprint_mode = "manual"
# elements = 5

RFdiffusion uses contigs strings to define the structural constraints of the design. The syntax uses : to separate multiple contigs and / to define segments within a single contig.

Unconditional Design

• Monomer: contigs='100' (diffuse a single chain of length 100).
• Hetero-oligomer: contigs='50:100' (diffuse two chains of lengths 50 and 100).
• Homo-oligomer: contigs='50' with copies=2 (diffuse two copies of a 50-residue chain with symmetry).

Binder Design

• Targeted Binder: contigs='A:50' with pdb='4N5T' (diffuse a 50-residue binder to chain A of the provided PDB).
• Hotspot-targeted Binder: contigs='E6-155:70-100' with pdb='5KQV' and hotspot='E64,E88,E96' (diffuse a binder of length 70-100 to chain E, targeting specific residues).

Motif Scaffolding

• Loop Scaffolding: contigs='40/A163-181/40' with pdb='5TPN' (diffuse a 40-residue loop between two segments of the PDB).
• Segment Scaffolding: contigs='A3-30/36/A33-68' with pdb='6MRR' (diffuse a 36-residue loop between two specific PDB ranges).

Partial Diffusion

• Full Noise: contigs='' with pdb='6MRR' (noise all coordinates).
• Fixed Segment: contigs='A1-10' with pdb='6MRR' (keep first 10 positions fixed, noise the rest).
• Fixed Chain: contigs='A' with pdb='1SSC' (fix chain A, noise the rest).

Tips

• Use a dash for ranges: contigs='50-100' samples lengths within that range.
• Leave pdb='' blank to trigger an upload prompt.

# Examples of contig syntax
contigs='100'           # Monomer
contigs='50:100'       # Hetero-oligomer
contigs='A:50'         # Binder to chain A
contigs='40/A163-181/40' # Motif scaffolding
contigs='A1-10'        # Partial diffusion (fixed residues)

Partial hallucination mixes supervised (fixbb) and unsupervised (hallucination) losses. This allows you to constrain parts of the sequence/structure while allowing others to be hallucinated.

Key features:

• loss_callback: Pass a function to define custom losses (e.g., Radius of Gyration).
• pos: Define specific residue positions to constrain in prep_inputs.
• rewire: Define loop lengths between constrained segments.
• restart(mode=["soft", "gumbel", "wildtype"]): A hybrid initialization mode.

import jax
import jax.numpy as jnp
from colabdesign import mk_afdesign_model, clear_mem
from colabdesign.af.alphafold.common import residue_constants

# Define a custom Radius of Gyration (rg) loss
def rg_loss(inputs, outputs):
    positions = outputs["structure_module"]["final_atom_positions"]
    ca = positions[:, residue_constants.atom_order["CA"]]
    center = ca.mean(0)
    rg = jnp.sqrt(jnp.square(ca - center).sum(-1).mean() + 1e-8)
    rg_th = 2.38 * ca.shape[0] ** 0.365
    rg = jax.nn.elu(rg - rg_th)
    return {"rg": rg}

clear_mem()
# Initialize partial model
af_model = mk_afdesign_model(
    protocol="partial",
    loss_callback=rg_loss,
    use_templates=False
)

# Set weight for the custom loss
af_model.opt["weights"]["rg"] = 0.1

# Prepare inputs with constraints
af_model.prep_inputs(
    pdb_filename="6MRR",
    chain="A",
    pos="3-30,33-68",  # Constrain these positions
    length=100
)

# Set loop length between segments
af_model.rewire(loops=[36])

# Hybrid initialization and design
af_model.restart(mode=["soft", "gumbel", "wildtype"])
af_model.design_3stage(100, 100, 10)

To design heterodimers (two different chains), use a custom loss_callback to penalize the radius of gyration (rg) of each chain independently. This ensures both protomers are compact. After design, use a 'Homooligomer Filter' to verify that each protomer can fold independently and that they correctly form the intended heterodimer without unintended homooligomerization.

# 1. Define heterodimer loss
def hd_loss(inputs, outputs):
    positions = outputs["structure_module"]["final_atom_positions"]
    ca1 = positions[:LENGTH1, residue_constants.atom_order["CA"]]
    ca2 = positions[LENGTH2:, residue_constants.atom_order["CA"]]
    # ... calculate rg for ca1 and ca2 ...
    return {"hd": rg1 + rg2}

# 2. Initialize and design
clear_mem()
af_model = mk_afdesign_model(protocol="hallucination", loss_callback=hd_loss)
af_model.prep_inputs(length=LENGTH1 + LENGTH2)
af_model.restart(mode=["gumbel", "soft"])
af_model.opt["weights"]["hd"] = 0.1
af_model.design_logits(100)