Merge remote-tracking branch 'origin/main' into fix_scaffoldguided_design

Author: woodsh17

.github/workflows/main.yml

@@ -69,6 +69,13 @@ jobs:
         uv pip install --no-cache-dir -e . --no-deps
         rm -rf ~/.cache # /app/RFdiffusion/tests
 
+    - name: Preseed DGL backend
+      shell: bash
+      run: |
+        mkdir -p "$HOME/.dgl"
+        printf '{"backend": "pytorch"}' > "$HOME/.dgl/config.conf"
+        echo "DGLBACKEND=pytorch" >> "$GITHUB_ENV"
+
     - name: Download weights
       run: |
         mkdir models
@@ -87,8 +94,29 @@ jobs:
     - name: Setup and Run ppi_scaffolds tests
       run: |
         tar -xvf examples/ppi_scaffolds_subset.tar.gz -C examples
-        cd tests && uv run python test_diffusion.py
+        total_chunks=$(nproc)
+        cd tests 
+
+        #launch all chunks in background and record PIDs + labels
+        pids=""
+        for chunk_index in $(seq 1 $total_chunks); do
+          echo "Running chunk $chunk_index of $total_chunks"
+          uv run python test_diffusion.py --total_chunks $total_chunks --chunk_index $chunk_index &
+          pids="$pids $!"
+        done
+
+        # wait for each and track failures
+        fail=0
+        for pid in $pids; do
+          if ! wait "$pid"; then
+            echo "A chunk (PID $pid) failed"
+            fail=1
+          else
+            echo "A chunk (PID $pid) passed"
+          fi
+        done
 
+        exit "$fail"
 
     # - name: Test with pytest
     #   run: |

README.md

@@ -21,6 +21,9 @@ RFdiffusion is an open source method for structure generation, with or without c
 - Binder design
 - Design diversification ("partial diffusion", sampling around a design)
 
+## [Documentation](https://sites.google.com/omsf.io/rfdiffusion/overview)
+View the RFdiffusion documentation resource maintained by Rosetta Commons [here](https://sites.google.com/omsf.io/rfdiffusion/overview). 
+
 ----
 
 # Table of contents
@@ -54,7 +57,9 @@ RFdiffusion is an open source method for structure generation, with or without c
 
 # Getting started / installation
 
-Thanks to Sergey Ovchinnikov, RFdiffusion is available as a [Google Colab Notebook](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/rf/examples/diffusion.ipynb) if you would like to run it there!
+Thanks to Sergey Ovchinnikov,many of the features of RFdiffusion are available as a [Google Colab Notebook](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/rf/examples/diffusion.ipynb) if you would like to run it there!
+
+There is also an official, Rosetta Commons-maintainted Docker image that you can find [here](https://hub.docker.com/r/rosettacommons/rfdiffusion). Thank you to Sergey Lyskov, Ajasja Ljubetic, and Hope Woods for creating this resource. 
 
 We strongly recommend reading this README carefully before getting started with RFdiffusion, and working through some of the examples in the Colab Notebook.
 

examples/design_macrocyclic_binder.sh

@@ -1,18 +1,15 @@
 #!/bin/bash
 
-prefix=./outputs/diffused_binder_cyclic2
+# Note that in the example below the indices in the
+# input_pdbs/7zkr_GABARAP.pdb file have been shifted
+# by +2 in chain A relative to pdbID 7zkr.
 
-# Note that the indices in this pdb file have been 
-# shifted by +2 in chain A relative to pdbID 7zkr.
-pdb='./input_pdbs/7zkr_GABARAP.pdb'
-
-num_designs=10
-script="../scripts/run_inference.py"
-$script --config-name base \
-inference.output_prefix=$prefix \
-inference.num_designs=$num_designs \
+../scripts/run_inference.py \
+--config-name base \
+inference.output_prefix=example_outputs/diffused_binder_cyclic2 \
+inference.num_designs=10 \
 'contigmap.contigs=[12-18 A3-117/0]' \
-inference.input_pdb=$pdb \
+inference.input_pdb=./input_pdbs/7zkr_GABARAP.pdb \
 inference.cyclic=True \
 diffuser.T=50 \
 inference.cyc_chains='a' \

examples/design_macrocyclic_monomer.sh

@@ -1,17 +1,15 @@
 #!/bin/bash 
 
-prefix=./outputs/uncond_cycpep
-# Note that the indices in this pdb file have been 
-# shifted by +2 in chain A relative to pdbID 7zkr.
-pdb='./input_pdbs/7zkr_GABARAP.pdb'
+# Note that in the example below the indices in the 
+# input_pdbs/7zkr_GABARAP.pdb file have been shifted
+# by +2 in chain A relative to pdbID 7zkr.
 
-num_designs=10
-script="../scripts/run_inference.py"
-$script --config-name base \
-inference.output_prefix=$prefix \
-inference.num_designs=$num_designs \
+../scripts/run_inference.py \
+--config-name base \
+inference.output_prefix=example_outputs/uncond_cycpep \
+inference.num_designs=10 \
 'contigmap.contigs=[12-18]' \
-inference.input_pdb=$pdb \
+inference.input_pdb=input_pdbs/7zkr_GABARAP.pdb \
 inference.cyclic=True \
 diffuser.T=50 \
 inference.cyc_chains='a'

examples/design_tetrahedral_oligos.sh

@@ -5,6 +5,6 @@
 # This external potential promotes contacts both within (with a relative weight of 1) and between chains (relative weight 0.1)
 # We specify that we want to apply these potentials to all chains, with a guide scale of 2.0 (a sensible starting point)
 # We decay this potential with quadratic form, so that it is applied more strongly initially
-# We specify a total length of 1200aa, so each chain is 100 residues long
+# We specify a total length of 1200aa, so each chain is 100 residues long - length updated to 600aa, so each chain is 50 residues long for testing to run faster
 
-python ../scripts/run_inference.py --config-name=symmetry inference.symmetry="tetrahedral" inference.num_designs=10 inference.output_prefix="example_outputs/tetrahedral_oligo" 'potentials.guiding_potentials=["type:olig_contacts,weight_intra:1,weight_inter:0.1"]' potentials.olig_intra_all=True potentials.olig_inter_all=True potentials.guide_scale=2.0 potentials.guide_decay="quadratic" 'contigmap.contigs=[1200-1200]'
+python ../scripts/run_inference.py --config-name=symmetry inference.symmetry="tetrahedral" inference.num_designs=10 inference.output_prefix="example_outputs/tetrahedral_oligo" 'potentials.guiding_potentials=["type:olig_contacts,weight_intra:1,weight_inter:0.1"]' potentials.olig_intra_all=True potentials.olig_inter_all=True potentials.guide_scale=2.0 potentials.guide_decay="quadratic" 'contigmap.contigs=[600-600]'

rfdiffusion/contigs.py

@@ -80,6 +80,7 @@ def __init__(
                 self.inpaint,
                 self.inpaint_hal,
                 self.inpaint_rf,
+                self.sampled_mask_length_bound,
             ) = self.expand_sampled_mask()
             self.ref = self.inpaint + self.receptor
             self.hal = self.inpaint_hal + self.receptor_hal
@@ -241,6 +242,8 @@ def expand_sampled_mask(self):
         inpaint_chain_idx = -1
         receptor_chain_break = []
         inpaint_chain_break = []
+        _receptor_mask_length_bound = []
+        _inpaint_mask_length_bound = []
         for con in self.sampled_mask:
             if (
                 all([i[0].isalpha() for i in con.split("/")[:-1]])
@@ -286,6 +289,7 @@ def expand_sampled_mask(self):
                         receptor_chain_break.append(
                             (receptor_idx - 1, 200)
                         )  # 200 aa chain break
+                _receptor_mask_length_bound.append(len(receptor))
             else:
                 inpaint_chain_idx += 1
                 for subcon in con.split("/"):
@@ -320,6 +324,7 @@ def expand_sampled_mask(self):
                         )
                         inpaint_idx += int(subcon.split("-")[0])
                 inpaint_chain_break.append((inpaint_idx - 1, 200))
+                _inpaint_mask_length_bound.append(len(inpaint))
 
         if self.topo is True or inpaint_hal == []:
             receptor_hal = [(i[0], i[1]) for i in receptor_hal]
@@ -335,14 +340,21 @@ def expand_sampled_mask(self):
             inpaint_rf[ch_break[0] :] += ch_break[1]
         for ch_break in receptor_chain_break[:-1]:
             receptor_rf[ch_break[0] :] += ch_break[1]
-
+        sampled_mask_length_bound = []
+        sampled_mask_length_bound.extend(_inpaint_mask_length_bound)
+        if _inpaint_mask_length_bound:
+            inpaint_last_bound = _inpaint_mask_length_bound[-1]
+        else:
+            inpaint_last_bound = 0
+        sampled_mask_length_bound.extend(map(lambda x: x + inpaint_last_bound, _receptor_mask_length_bound))
         return (
             receptor,
             receptor_hal,
             receptor_rf.tolist(),
             inpaint,
             inpaint_hal,
             inpaint_rf.tolist(),
+            sampled_mask_length_bound
         )
 
     def get_inpaint_seq_str(self, inpaint_s, ss=False):

rfdiffusion/inference/model_runners.py

@@ -14,6 +14,7 @@
 from rfdiffusion import util
 from hydra.core.hydra_config import HydraConfig
 import os
+import string
 
 from rfdiffusion.model_input_logger import pickle_function_call
 import sys
@@ -144,13 +145,14 @@ def initialize(self, conf: DictConfig) -> None:
             self.symmetry = None
 
         self.allatom = ComputeAllAtomCoords().to(self.device)
-        
+
         if self.inf_conf.input_pdb is None:
             # set default pdb
             script_dir=os.path.dirname(os.path.realpath(__file__))
             self.inf_conf.input_pdb=os.path.join(script_dir, '../../examples/input_pdbs/1qys.pdb')
         self.target_feats = iu.process_target(self.inf_conf.input_pdb, parse_hetatom=True, center=False)
         self.chain_idx = None
+        self.idx_pdb = None
 
         ##############################
         ### Handle Partial Noising ###
@@ -330,8 +332,42 @@ def sample_init(self, return_forward_trajectory=False):
         contig_map=self.contig_map
 
         self.diffusion_mask = self.mask_str
-        self.chain_idx=['A' if i < self.binderlen else 'B' for i in range(L_mapped)]
-        
+        length_bound = self.contig_map.sampled_mask_length_bound.copy()
+
+        first_res = 0
+        self.chain_idx = []
+        self.idx_pdb = []
+        all_chains = {contig_ref[0] for contig_ref in self.contig_map.ref}
+        available_chains = sorted(list(set(string.ascii_letters) - all_chains))
+
+        # Iterate over each chain
+        for last_res in length_bound:
+            chain_ids = {contig_ref[0] for contig_ref in self.contig_map.ref[first_res: last_res]}
+            # If we are designing this chain, it will have a '-' in the contig map
+            # Renumber this chain from 1
+            if "_" in chain_ids:
+                self.idx_pdb += [idx + 1 for idx in range(last_res - first_res)]
+                chain_ids = chain_ids - {"_"}
+                # If there are no fixed residues that have a chain id, pick the first available letter
+                if not chain_ids:
+                    if not available_chains:
+                        raise ValueError(f"No available chains!  You are trying to design a new chain, and you have "
+                                         f"already used all upper- and lower-case chain ids (up to 52 chains): "
+                                         f"{','.join(all_chains)}.")
+                    chain_id = available_chains[0]
+                    available_chains.remove(chain_id)
+                # Otherwise, use the chain of the fixed (motif) residues
+                else:
+                    assert len(chain_ids) == 1, f"Error: Multiple chain IDs in chain: {chain_ids}"
+                    chain_id = list(chain_ids)[0]
+                self.chain_idx += [chain_id] * (last_res - first_res)
+            # If this is a fixed chain, maintain the chain and residue numbering
+            else:
+                self.idx_pdb += [contig_ref[1] for contig_ref in self.contig_map.ref[first_res: last_res]]
+                assert len(chain_ids) == 1, f"Error: Multiple chain IDs in chain: {chain_ids}"
+                self.chain_idx += [list(chain_ids)[0]] * (last_res - first_res)
+            first_res = last_res
+
         ####################################
         ### Generate initial coordinates ###
         ####################################

scripts/run_inference.py

@@ -141,6 +141,7 @@ def main(conf: HydraConfig) -> None:
             sampler.binderlen,
             chain_idx=sampler.chain_idx,
             bfacts=bfacts,
+            idx_pdb=sampler.idx_pdb
         )
 
         # run metadata