restructuring library as a python module

README.md

@@ -92,6 +92,8 @@ conda activate SE3nv
 cd env/SE3Transformer
 pip install --no-cache-dir -r requirements.txt
 python setup.py install
+cd ../.. # change into the root directory of the repository
+pip install -e . # install the rfdiffusion module from the root of the repository
 ```
 Anytime you run diffusion you should be sure to activate this conda environment by running the following command:
 ```

config/inference/base.yaml

@@ -19,6 +19,8 @@ inference:
   final_step: 1
   deterministic: False
   trb_save_ckpt_path: null
+  schedule_directory_path: null
+  model_directory_path: null
 
 contigmap:
   contigs: null

examples/design_cyclic_oligos.sh

@@ -7,7 +7,7 @@
 # We decay this potential with quadratic form, so that it is applied more strongly initially
 # We specify a total length of 480aa, so each chain is 80 residues long
 
-python ../run_inference.py \
+python ../scripts/run_inference.py \
 --config-name=symmetry \
 inference.symmetry="C6" \
 inference.num_designs=10 \

examples/design_dihedral_oligos.sh

@@ -7,7 +7,7 @@
 # We decay this potential with quadratic form, so that it is applied more strongly initially
 # We specify a total length of 320aa, so each chain is 80 residues long
 
-python ../run_inference.py \
+python ../scripts/run_inference.py \
 --config-name=symmetry \
 inference.symmetry="D2" \
 inference.num_designs=10 \

examples/design_enzyme.sh

@@ -13,4 +13,4 @@
 # We specify the identity of the substrate molecule (from which to apply the potential)
 # We then specify the we want to use the model weights fine tuned to scaffold small motifs (i.e. three single residues)
 
-../run_inference.py inference.output_prefix=example_outputs/design_enzyme inference.input_pdb=input_pdbs/5an7.pdb 'contigmap.contigs=[10-100/A1083-1083/10-100/A1051-1051/10-100/A1180-1180/10-100]' potentials.guide_scale=1 'potentials.guiding_potentials=["type:substrate_contacts,s:1,r_0:8,rep_r_0:5.0,rep_s:2,rep_r_min:1"]' potentials.substrate=LLK inference.ckpt_override_path=../models/ActiveSite_ckpt.pt
+../scripts/run_inference.py inference.output_prefix=example_outputs/design_enzyme inference.input_pdb=input_pdbs/5an7.pdb 'contigmap.contigs=[10-100/A1083-1083/10-100/A1051-1051/10-100/A1180-1180/10-100]' potentials.guide_scale=1 'potentials.guiding_potentials=["type:substrate_contacts,s:1,r_0:8,rep_r_0:5.0,rep_s:2,rep_r_min:1"]' potentials.substrate=LLK inference.ckpt_override_path=../models/ActiveSite_ckpt.pt

examples/design_motifscaffolding.sh

@@ -8,4 +8,4 @@
 #   - 10-40 residues (randomly sampled)
 # We generate 10 designs
 
-../run_inference.py inference.output_prefix=example_outputs/design_motifscaffolding inference.input_pdb=input_pdbs/5TPN.pdb 'contigmap.contigs=[10-40/A163-181/10-40]' inference.num_designs=10
+../scripts/run_inference.py inference.output_prefix=example_outputs/design_motifscaffolding inference.input_pdb=input_pdbs/5TPN.pdb 'contigmap.contigs=[10-40/A163-181/10-40]' inference.num_designs=10

examples/design_motifscaffolding_inpaintseq.sh

@@ -12,4 +12,4 @@
 # We generate 10 designs
 # We then specify that residues 163-168 (inclusive), 170-171 (inclusive) and 179 (inclusive) on the A chain of the input, should be masked in the input
 
-../run_inference.py inference.output_prefix=example_outputs/design_motifscaffolding_inpaintseq inference.input_pdb=input_pdbs/5TPN.pdb 'contigmap.contigs=[10-40/A163-181/10-40]' inference.num_designs=10 'contigmap.inpaint_seq=[A163-168/A170-171/A179]'
+../scripts/run_inference.py inference.output_prefix=example_outputs/design_motifscaffolding_inpaintseq inference.input_pdb=input_pdbs/5TPN.pdb 'contigmap.contigs=[10-40/A163-181/10-40]' inference.num_designs=10 'contigmap.inpaint_seq=[A163-168/A170-171/A179]'

examples/design_nickel.sh

@@ -17,4 +17,4 @@
 
 ckpt='../models/Base_epoch8_ckpt.pt'
 
-python ../run_inference.py inference.symmetry="C4" inference.num_designs=15 inference.output_prefix=example_outputs/design_nickel 'potentials.guiding_potentials=["type:olig_contacts,weight_intra:1,weight_inter:0.06"]' potentials.olig_intra_all=True potentials.olig_inter_all=True potentials.guide_scale=2 potentials.guide_decay="quadratic" inference.input_pdb=input_pdbs/nickel_symmetric_motif.pdb 'contigmap.contigs=[50/A2-4/50/0 50/A7-9/50/0 50/A12-14/50/0 50/A17-19/50/0]' inference.ckpt_override_path=$ckpt 
+python ../scripts/run_inference.py inference.symmetry="C4" inference.num_designs=15 inference.output_prefix=example_outputs/design_nickel 'potentials.guiding_potentials=["type:olig_contacts,weight_intra:1,weight_inter:0.06"]' potentials.olig_intra_all=True potentials.olig_inter_all=True potentials.guide_scale=2 potentials.guide_decay="quadratic" inference.input_pdb=input_pdbs/nickel_symmetric_motif.pdb 'contigmap.contigs=[50/A2-4/50/0 50/A7-9/50/0 50/A12-14/50/0 50/A17-19/50/0]' inference.ckpt_override_path=$ckpt

examples/design_partialdiffusion.sh

@@ -7,4 +7,4 @@
 # But, in either case, the contig length must sum to the length of the input pdb file
 # We generate 10 designs, and noise and denoise 10 steps (20% of the full trajectory)
 
-../run_inference.py inference.output_prefix=example_outputs/design_partialdiffusion inference.input_pdb=input_pdbs/2KL8.pdb 'contigmap.contigs=[79-79]' inference.num_designs=10 diffuser.partial_T=10
+../scripts/run_inference.py inference.output_prefix=example_outputs/design_partialdiffusion inference.input_pdb=input_pdbs/2KL8.pdb 'contigmap.contigs=[79-79]' inference.num_designs=10 diffuser.partial_T=10

examples/design_partialdiffusion_withseq.sh

@@ -11,4 +11,4 @@
 # This contig will lead to the whole input pdb being noise by 10 steps (partial_T=10)
 # However, we provide the sequence of the peptide (the last 20 residues in the contig), with provide_seq=[172-205]. This is 0-indexed
 
-../run_inference.py inference.output_prefix=example_outputs/design_partialdiffusion_peptidewithsequence inference.input_pdb=input_pdbs/peptide_complex_ideal_helix.pdb 'contigmap.contigs=["172-172/0 34-34"]' diffuser.partial_T=10 inference.num_designs=10 'contigmap.provide_seq=[172-205]'
+../scripts/run_inference.py inference.output_prefix=example_outputs/design_partialdiffusion_peptidewithsequence inference.input_pdb=input_pdbs/peptide_complex_ideal_helix.pdb 'contigmap.contigs=["172-172/0 34-34"]' diffuser.partial_T=10 inference.num_designs=10 'contigmap.provide_seq=[172-205]'

examples/design_ppi.sh

@@ -8,4 +8,4 @@
 # We tell diffusion to target three specific residues on the target, specifically residues 59, 83 and 91 of the A chain
 # We make 10 designs, and reduce the noise added during inference to 0, to improve the quality of the designs
 
-../run_inference.py inference.output_prefix=example_outputs/design_ppi inference.input_pdb=input_pdbs/insulin_target.pdb 'contigmap.contigs=[A1-150/0 70-100]' 'ppi.hotspot_res=[A59,A83,A91]' inference.num_designs=10 denoiser.noise_scale_ca=0 denoiser.noise_scale_frame=0
+../scripts/run_inference.py inference.output_prefix=example_outputs/design_ppi inference.input_pdb=input_pdbs/insulin_target.pdb 'contigmap.contigs=[A1-150/0 70-100]' 'ppi.hotspot_res=[A59,A83,A91]' inference.num_designs=10 denoiser.noise_scale_ca=0 denoiser.noise_scale_frame=0

examples/design_ppi_scaffolded.sh

@@ -7,4 +7,4 @@
 # We then provide a path to a directory of different scaffolds (we've provided some for you to use, from Cao et al., 2022)
 # We generate 10 designs, and reduce the noise added during inference to 0 (which improves the quality of designs)
 
-../run_inference.py scaffoldguided.target_path=input_pdbs/insulin_target.pdb inference.output_prefix=example_outputs/design_ppi_scaffolded scaffoldguided.scaffoldguided=True 'ppi.hotspot_res=[A59,A83,A91]' scaffoldguided.target_pdb=True scaffoldguided.target_ss=target_folds/insulin_target_ss.pt scaffoldguided.target_adj=target_folds/insulin_target_adj.pt scaffoldguided.scaffold_dir=./ppi_scaffolds/ inference.num_designs=10 denoiser.noise_scale_ca=0 denoiser.noise_scale_frame=0
+../scripts/run_inference.py scaffoldguided.target_path=input_pdbs/insulin_target.pdb inference.output_prefix=example_outputs/design_ppi_scaffolded scaffoldguided.scaffoldguided=True 'ppi.hotspot_res=[A59,A83,A91]' scaffoldguided.target_pdb=True scaffoldguided.target_ss=target_folds/insulin_target_ss.pt scaffoldguided.target_adj=target_folds/insulin_target_adj.pt scaffoldguided.scaffold_dir=./ppi_scaffolds/ inference.num_designs=10 denoiser.noise_scale_ca=0 denoiser.noise_scale_frame=0

examples/design_tetrahedral_oligos.sh

@@ -7,7 +7,7 @@
 # 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
 
-python ../run_inference.py \
+python ../scripts/run_inference.py \
 --config-name=symmetry \
 inference.symmetry="tetrahedral" \
 inference.num_designs=10 \

examples/design_timbarrel.sh

@@ -8,4 +8,4 @@
 # We add 0-5 residues (randomly sampled) to the N and the C-terminus
 # This will allow the generation of diverse TIM barrels with slightly different length helices and strands
 
-../run_inference.py inference.output_prefix=example_outputs/design_tim_barrel scaffoldguided.scaffoldguided=True scaffoldguided.target_pdb=False scaffoldguided.scaffold_dir=tim_barrel_scaffold/ inference.num_designs=200 denoiser.noise_scale_ca=0.5 denoiser.noise_scale_frame=0.5 scaffoldguided.sampled_insertion=0-5 scaffoldguided.sampled_N=0-5 scaffoldguided.sampled_C=0-5
+../scripts/run_inference.py inference.output_prefix=example_outputs/design_tim_barrel scaffoldguided.scaffoldguided=True scaffoldguided.target_pdb=False scaffoldguided.scaffold_dir=tim_barrel_scaffold/ inference.num_designs=200 denoiser.noise_scale_ca=0.5 denoiser.noise_scale_frame=0.5 scaffoldguided.sampled_insertion=0-5 scaffoldguided.sampled_N=0-5 scaffoldguided.sampled_C=0-5

examples/design_unconditional.sh

@@ -4,4 +4,4 @@
 # We tell RFdiffusion that designs should be 100-200 residues in length (randomly sampled each design)
 # We generate 10 such designs
 
-../run_inference.py inference.output_prefix=example_outputs/design_unconditional 'contigmap.contigs=[100-200]' inference.num_designs=10
+../scripts/run_inference.py inference.output_prefix=example_outputs/design_unconditional 'contigmap.contigs=[100-200]' inference.num_designs=10

Attention_module.py → rfdiffusion/Attention_module.py

@@ -3,7 +3,7 @@
 import torch.nn.functional as F
 import math
 from opt_einsum import contract as einsum
-from util_module import init_lecun_normal
+from rfdiffusion.util_module import init_lecun_normal
 
 class FeedForwardLayer(nn.Module):
     def __init__(self, d_model, r_ff, p_drop=0.1):

Embeddings.py → rfdiffusion/Embeddings.py

@@ -3,10 +3,10 @@
 import torch.nn.functional as F
 from opt_einsum import contract as einsum
 import torch.utils.checkpoint as checkpoint
-from util import get_tips
-from util_module import Dropout, create_custom_forward, rbf, init_lecun_normal
-from Attention_module import Attention, FeedForwardLayer, AttentionWithBias
-from Track_module import PairStr2Pair
+from rfdiffusion.util import get_tips
+from rfdiffusion.util_module import Dropout, create_custom_forward, rbf, init_lecun_normal
+from rfdiffusion.Attention_module import Attention, FeedForwardLayer, AttentionWithBias
+from rfdiffusion.Track_module import PairStr2Pair
 from icecream import ic
 import math
 

RoseTTAFoldModel.py → rfdiffusion/RoseTTAFoldModel.py

@@ -1,11 +1,10 @@
 import torch
 import torch.nn as nn
-from Embeddings import MSA_emb, Extra_emb, Templ_emb, Recycling, Timestep_emb
-from Track_module import IterativeSimulator
-from AuxiliaryPredictor import DistanceNetwork, MaskedTokenNetwork, ExpResolvedNetwork, LDDTNetwork
-from util import INIT_CRDS
+from rfdiffusion.Embeddings import MSA_emb, Extra_emb, Templ_emb, Recycling, Timestep_emb
+from rfdiffusion.Track_module import IterativeSimulator
+from rfdiffusion.AuxiliaryPredictor import DistanceNetwork, MaskedTokenNetwork, ExpResolvedNetwork, LDDTNetwork
 from opt_einsum import contract as einsum
-from icecream import ic 
+
 
 class RoseTTAFoldModule(nn.Module):
     def __init__(self, 

SE3_network.py → rfdiffusion/SE3_network.py

@@ -5,7 +5,7 @@
 #from equivariant_attention.modules import GConvSE3, GNormSE3
 #from equivariant_attention.fibers import Fiber
 
-from util_module import init_lecun_normal_param
+from rfdiffusion.util_module import init_lecun_normal_param
 from se3_transformer.model import SE3Transformer
 from se3_transformer.model.fiber import Fiber
 

Track_module.py → rfdiffusion/Track_module.py

@@ -1,12 +1,7 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from opt_einsum import contract as einsum
 import torch.utils.checkpoint as checkpoint
-from util import cross_product_matrix
-from util_module import *
-from Attention_module import *
-from SE3_network import SE3TransformerWrapper
+from rfdiffusion.util_module import *
+from rfdiffusion.Attention_module import *
+from rfdiffusion.SE3_network import SE3TransformerWrapper
 
 # Components for three-track blocks
 # 1. MSA -> MSA update (biased attention. bias from pair & structure)

diff_util.py → rfdiffusion/diff_util.py

@@ -2,7 +2,7 @@
 import numpy as np 
 import random 
 
-from chemical import INIT_CRDS
+from rfdiffusion.chemical import INIT_CRDS
 from icecream import ic 
 
 

diffusion.py → rfdiffusion/diffusion.py

@@ -4,22 +4,14 @@
 import numpy as np
 import os
 import logging
-from typing import List
 
 from scipy.spatial.transform import Rotation as scipy_R
 
-from util import rigid_from_3_points
+from rfdiffusion.util import rigid_from_3_points
 
-from util import torsion_indices as TOR_INDICES
-from util import torsion_can_flip as TOR_CAN_FLIP
-from util import reference_angles as REF_ANGLES
+from rfdiffusion.util_module import ComputeAllAtomCoords
 
-from util_module import ComputeAllAtomCoords
-
-from diff_util import th_min_angle, th_interpolate_angles
-
-from chemical import INIT_CRDS
-import igso3
+from rfdiffusion import igso3
 import time
 
 torch.set_printoptions(sci_mode=False)

inference/model_runners.py → rfdiffusion/inference/model_runners.py

@@ -1,35 +1,29 @@
 import torch
 import numpy as np
 from omegaconf import DictConfig, OmegaConf
-from RoseTTAFoldModel import RoseTTAFoldModule
-from kinematics import get_init_xyz, xyz_to_t2d
-from diffusion import Diffuser
-from chemical import seq2chars
-from util_module import ComputeAllAtomCoords
-from contigs import ContigMap
-from inference import utils as iu
-from potentials.manager import PotentialManager
-from inference import symmetry
+from rfdiffusion.RoseTTAFoldModel import RoseTTAFoldModule
+from rfdiffusion.kinematics import get_init_xyz, xyz_to_t2d
+from rfdiffusion.diffusion import Diffuser
+from rfdiffusion.chemical import seq2chars
+from rfdiffusion.util_module import ComputeAllAtomCoords
+from rfdiffusion.contigs import ContigMap
+from rfdiffusion.inference import utils as iu, symmetry
+from rfdiffusion.potentials.manager import PotentialManager
 import logging
 import torch.nn.functional as nn
-import util
-import hydra
+from rfdiffusion import util
 from hydra.core.hydra_config import HydraConfig
 import os
 
+from rfdiffusion.model_input_logger import pickle_function_call
 import sys
-SCRIPT_DIR=os.path.dirname(os.path.realpath(__file__))
-sys.path.append(SCRIPT_DIR + '../') # to access RF structure prediction stuff 
 
-from model_input_logger import pickle_function_call
+SCRIPT_DIR=os.path.dirname(os.path.realpath(__file__))
 
 TOR_INDICES  = util.torsion_indices
 TOR_CAN_FLIP = util.torsion_can_flip
 REF_ANGLES   = util.reference_angles
 
-# Check for cache schedule
-if not os.path.exists(f'{SCRIPT_DIR}/../schedules'):
-    os.mkdir(f'{SCRIPT_DIR}/../schedules')
 
 class Sampler:
 
@@ -66,6 +60,13 @@ def initialize(self, conf: DictConfig) -> None:
         ### Select Appropriate Model ###
         ################################
 
+        if conf.inference.model_directory_path is not None:
+            model_directory = conf.inference.model_directory_path
+        else:
+            model_directory = f"{SCRIPT_DIR}/../../models"
+
+        print(f"Reading models from {model_directory}")
+
         # Initialize inference only helper objects to Sampler
         if conf.inference.ckpt_override_path is not None:
             self.ckpt_path = conf.inference.ckpt_override_path
@@ -77,18 +78,18 @@ def initialize(self, conf: DictConfig) -> None:
                     # this is only used for partial diffusion
                     assert conf.diffuser.partial_T is not None, "The provide_seq input is specifically for partial diffusion"
                 if conf.scaffoldguided.scaffoldguided:
-                    self.ckpt_path=f'{SCRIPT_DIR}/../models/InpaintSeq_Fold_ckpt.pt'
+                    self.ckpt_path = f'{model_directory}/InpaintSeq_Fold_ckpt.pt'
                 else:
-                    self.ckpt_path = f'{SCRIPT_DIR}/../models/InpaintSeq_ckpt.pt'
+                    self.ckpt_path = f'{model_directory}/InpaintSeq_ckpt.pt'
             elif conf.ppi.hotspot_res is not None and conf.scaffoldguided.scaffoldguided is False:
                 # use complex trained model
-                self.ckpt_path = f'{SCRIPT_DIR}/../models/Complex_base_ckpt.pt'
+                self.ckpt_path = f'{model_directory}/Complex_base_ckpt.pt'
             elif conf.scaffoldguided.scaffoldguided is True:
                 # use complex and secondary structure-guided model
-                self.ckpt_path = f'{SCRIPT_DIR}/../models/Complex_Fold_base_ckpt.pt'
+                self.ckpt_path = f'{model_directory}/Complex_Fold_base_ckpt.pt'
             else:
                 # use default model
-                self.ckpt_path = f'{SCRIPT_DIR}/../models/Base_ckpt.pt'
+                self.ckpt_path = f'{model_directory}/Base_ckpt.pt'
         # for saving in trb file:
         assert self._conf.inference.trb_save_ckpt_path is None, "trb_save_ckpt_path is not the place to specify an input model. Specify in inference.ckpt_override_path"
         self._conf['inference']['trb_save_ckpt_path']=self.ckpt_path
@@ -117,7 +118,16 @@ def initialize(self, conf: DictConfig) -> None:
         self.potential_conf = self._conf.potentials
         self.diffuser_conf = self._conf.diffuser
         self.preprocess_conf = self._conf.preprocess
-        self.diffuser = Diffuser(**self._conf.diffuser, cache_dir=f'{SCRIPT_DIR}/../schedules')
+
+        if conf.inference.schedule_directory_path is not None:
+            schedule_directory = conf.inference.schedule_directory_path
+        else:
+            schedule_directory = f"{SCRIPT_DIR}/../../schedules"
+
+        # Check for cache schedule
+        if not os.path.exists(schedule_directory):
+            os.mkdir(schedule_directory)
+        self.diffuser = Diffuser(**self._conf.diffuser, cache_dir=schedule_directory)
 
         ###########################
         ### Initialise Symmetry ###
@@ -138,7 +148,7 @@ def initialize(self, conf: DictConfig) -> None:
         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.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
 
@@ -494,7 +504,7 @@ def _preprocess(self, seq, xyz_t, t, repack=False):
         ### alpha_t ###
         ###############
         seq_tmp = t1d[...,:-1].argmax(dim=-1).reshape(-1,L)
-        alpha, _, alpha_mask, _ = util.get_torsions(xyz_t.reshape(-1,L,27,3), seq_tmp, TOR_INDICES, TOR_CAN_FLIP, REF_ANGLES)
+        alpha, _, alpha_mask, _ = util.get_torsions(xyz_t.reshape(-1, L, 27, 3), seq_tmp, TOR_INDICES, TOR_CAN_FLIP, REF_ANGLES)
         alpha_mask = torch.logical_and(alpha_mask, ~torch.isnan(alpha[...,0]))
         alpha[torch.isnan(alpha)] = 0.0
         alpha = alpha.reshape(1,-1,L,10,2)

inference/utils.py → rfdiffusion/inference/utils.py

@@ -1,21 +1,16 @@
 import numpy as np
 import os
-import sys
 from omegaconf import DictConfig
-from kinematics import xyz_to_t2d
 import torch
 import torch.nn.functional as nn
-from diffusion import get_beta_schedule
+from rfdiffusion.diffusion import get_beta_schedule
 from scipy.spatial.transform import Rotation as scipy_R
-from util import rigid_from_3_points
-from util_module import ComputeAllAtomCoords
-from potentials.manager import PotentialManager
-import util
+from rfdiffusion.util import rigid_from_3_points
+from rfdiffusion.util_module import ComputeAllAtomCoords
+from rfdiffusion import util
 import random
 import logging
-import string
-from inference import model_runners
-import hydra
+from rfdiffusion.inference import model_runners
 import glob
 
 ###########################################################

kinematics.py → rfdiffusion/kinematics.py

@@ -1,6 +1,6 @@
 import numpy as np
 import torch
-from chemical import INIT_CRDS
+from rfdiffusion.chemical import INIT_CRDS
 
 PARAMS = {
     "DMIN"    : 2.0,