Constrained diffusion on one chain in a protein complex

[caseysm]

Hi,

This is a really impressive tool and I am thinking of many cool things this can be used for in the future. I was wondering if the following would be possible with Chroma.

Given a pdb file that is a 2 protein complex (chain A and chain B), how would I do constrained diffusion on chain A? In other words, I want to hold chain B constant, and I want to constrain the output of Chroma to obey the topology of chain A.

Thanks for the help!

[zak-generate]

Hi! We have the Substructure conditioner for this purpose. Here is a minimal working example that is slightly adapted from our tutorial notebook:

from generate.models import Chroma
from generate.data import Protein
from generate.layers.structure import conditioners
device = 'cuda' if torch.cuda.is_available() else 'cpu'

chroma = Chroma()
MY_2CHAIN_PROTEIN = '3FGR'
protein = Protein(MY_2CHAIN_PROTEIN, canonicalize=True, device=device)

# Configure Substructure Conditioner
conditioner = conditioners.SubstructureConditioner(
        protein,
        backbone_model=chroma.backbone_network, 
        selection = 'chain A').to(device)

# Draw a Sample
torch.manual_seed(0)
sample = chroma.sample(
             protein_init=protein,
             conditioner=conditioner,
             langevin_factor=4.0,
             langevin_isothermal=True,
             inverse_temperature=8.0,
             steps=500)

display(sample)

hope this helps!

[wujiewang]

Also, it is recommended to use sde_func='langevin' for infilling task. See https://github.com/generatebio/chroma/issues/24 for more info.

[wujiewang]

closing this issue for now. feel free to reopen or submit a new issue if you have follow-up questions.