update paths in training tests

make_conda.sh

@@ -1,17 +1,24 @@
 #!/bin/bash 
-
+set -e
 # change to what you want 
 ENV_NAME="ca_rfd_release"
 
 
 # Create new conda env 
-conda create -n $ENV_NAME python=3.9 -y 
+#conda create -n $ENV_NAME python=3.9 -y
+conda create -n $ENV_NAME -y python=3.9 -c defaults -c conda-forge
 
 # hook and activate 
-eval "$('conda' 'shell.bash' 'hook')"
+# eval "$('conda' 'shell.bash' 'hook')"
+eval "$(conda shell.bash hook)"
 conda activate $ENV_NAME 
 conda info --envs
-
+# Verify activation
+#CURRENT_ENV=$(conda info --active-env-name)
+if [ "$CONDA_DEFAULT_ENV" != "$ENV_NAME" ]; then
+    echo "Failed to activate $ENV_NAME. Currently in: $CURRENT_ENV"
+    exit 1
+fi
 #### Install packages #### 
 # They are installed in stages
 
@@ -23,11 +30,13 @@ conda install -y \
 
 # pytorch + dependancies
 conda install -y \
+   -c nvidia \
    -c nvidia/label/cuda-12.1.0 \
    -c pytorch \
    -c pyg \
    -c dglteam/label/cu121 \
    -c anaconda \
+   -c conda-forge \
    pip \
    ipython=8.8.0 \
    "ipykernel>=6.22.0" \
@@ -39,12 +48,12 @@ conda install -y \
    pytorch==2.2 \
    pytorch-cuda==12.1 \
    dgl==2.0.0.cu121 \
-   einops=0.7.0 \
    pyg
 
 conda install -y \
    -c conda-forge \
    openbabel=3.1.1 \
+   einops=0.7.0
 
 # pip extras
 pip install \

rf_diffusion/conditioning.py

@@ -313,26 +313,25 @@ def __call__(self,
         ref_dict = conditions_dict['ref_dict']
         motif_indep  = ref_dict['motif_indep']
         ij_visible   = ref_dict['ij_visible']
-        con_hal_idx0 = torch.from_numpy(ref_dict['con_hal_idx0'])
-        con_ref_idx0 = torch.from_numpy(ref_dict['con_ref_idx0'])
+        con_hal_idx0 = torch.from_numpy(ref_dict['con_hal_idx0']).long()
+        con_ref_idx0 = torch.from_numpy(ref_dict['con_ref_idx0']).long()
 
         ### Compute 1D/2D mask information ### 
         ######################################
 
         # hal structure (design from diffusion)
         L_sm = indep.is_sm.sum()
         L_hal_w_ligand = indep.length()
-        add_con_hal0 = torch.arange(L_hal_w_ligand-L_sm, L_hal_w_ligand)
+        add_con_hal0 = torch.arange(L_hal_w_ligand-L_sm, L_hal_w_ligand, dtype=torch.long)
         con_hal_idx0 = torch.cat((con_hal_idx0, add_con_hal0))
         # reference structure (original motif pdb)
         L_ref_w_ligand = motif_indep.length()
-        add_con_ref0 = torch.arange(L_ref_w_ligand-L_sm, L_ref_w_ligand)
+        add_con_ref0 = torch.arange(L_ref_w_ligand-L_sm, L_ref_w_ligand, dtype=torch.long)
         con_ref_idx0 = torch.cat((con_ref_idx0, add_con_ref0))
 
 
         # Get constraints between motif chunks - coarse
         coarse_ij_visibility = get_coarse_ij_matrix(ij_visible)
-
         is_motif = torch.zeros(L_hal_w_ligand, dtype=torch.bool)
         is_motif[con_hal_idx0] = True # contains protein motif + ligand motif
 

rf_diffusion/config/training/test_ca_rfd_sm_train.yaml

@@ -2,12 +2,15 @@ defaults:
   - RFdiffusion_CA
 
 zero_weights: True 
+rundir: /tmp
 
 diffuser: 
   type: 'legacy'
 
 dataloader:
   DATAPKL_AA: goldens/tiny_dataset_diffusion_20230214.pkl
+  PDB_DIR: test_data/ca_rfd/PDB-2021AUG02
+  MOL_DIR: test_data/ca_rfd/RF2_allatom/rcsb/pkl
 
 spoof_item: "{   'chosen_dataset': 'sm_complex',
                  'mask_gen_seed': 46265381,

rf_diffusion/config/training/test_ca_rfd_unconditional_train.yaml

@@ -2,12 +2,15 @@ defaults:
   - RFdiffusion_CA
 
 zero_weights: True 
+rundir: /tmp
 
 diffuser: 
   type: 'legacy'
 
 dataloader:
   DATAPKL_AA: goldens/tiny_dataset_diffusion_20230214.pkl
+  PDB_DIR: test_data/ca_rfd/PDB-2021AUG02
+  MOL_DIR: test_data/ca_rfd/RF2_allatom/rcsb/pkl
 
 spoof_item: "{   'chosen_dataset': 'pdb_aa',
      'mask_gen_seed': 46265381,

rf_diffusion/run_inference.py

@@ -149,31 +149,33 @@ def add_carfd_sidechains(px0: torch.Tensor, ref_dict: dict) -> torch.Tensor:
     motif_idx_hal = ref_dict['con_hal_idx0']
 
     # grab the sequence & structure of non-SM motif residues  
-    motif_seq_ref = motif_indep.seq[motif_idx_ref]
-    motif_crds_ref = motif_indep.xyz[motif_idx_ref]
-
-    # compute sidechain torsions in the original (ref) motif
-    converter = rf2aa.util_module.XYZConverter()
-    tors, tors_alt, tors_mask, tors_planar = converter.get_torsions(motif_crds_ref[None], 
-                                                                    motif_seq_ref[None])
-
-
-    # grab predicted ("hal") motif backbone coordinates, and extend sidechains onto them 
-    # according to the torsions from perfect motif
-    motif_bb_hal = px0[motif_idx_hal,:3,:]
-    motif_allatom_hal = converter.compute_all_atom(motif_seq_ref[None], motif_bb_hal[None], tors)
-    (RTframes, xyz_full) = motif_allatom_hal 
-    # slice it in
-    px0[motif_idx_hal,:14] = xyz_full[0,:,:14]
-
-    # now align the output to the original motif on backbone
-    motif_bb_ref = motif_crds_ref[:,:3]
-    com_ref = motif_bb_ref.reshape(-1,3).mean(dim=0)
-    com_hal = motif_bb_hal.reshape(-1,3).mean(dim=0)
-    bb_rms, _, R = th_kabsch(motif_bb_ref.reshape(-1,3), motif_bb_hal.reshape(-1,3))
-    # bring to origin, rotate, then translate to reference/native position
-    px0 = torch.einsum('lai,ij->laj', px0 - com_hal, R) + com_ref
+    # need a big IF here
+    if len(motif_idx_ref) > 0:
+        motif_seq_ref = motif_indep.seq[motif_idx_ref]
+        motif_crds_ref = motif_indep.xyz[motif_idx_ref]
+
+        # compute sidechain torsions in the original (ref) motif
+        converter = rf2aa.util_module.XYZConverter()
+        tors, tors_alt, tors_mask, tors_planar = converter.get_torsions(motif_crds_ref[None], 
+                                                                        motif_seq_ref[None])
+
 
+        # grab predicted ("hal") motif backbone coordinates, and extend sidechains onto them 
+        # according to the torsions from perfect motif
+        motif_bb_hal = px0[motif_idx_hal,:3,:]
+        motif_allatom_hal = converter.compute_all_atom(motif_seq_ref[None], motif_bb_hal[None], tors)
+        (RTframes, xyz_full) = motif_allatom_hal 
+        # slice it in
+        px0[motif_idx_hal,:14] = xyz_full[0,:,:14]
+
+        # now align the output to the original motif on backbone
+        motif_bb_ref = motif_crds_ref[:,:3]
+        com_ref = motif_bb_ref.reshape(-1,3).mean(dim=0)
+        com_hal = motif_bb_hal.reshape(-1,3).mean(dim=0)
+        bb_rms, _, R = th_kabsch(motif_bb_ref.reshape(-1,3), motif_bb_hal.reshape(-1,3))
+        # bring to origin, rotate, then translate to reference/native position
+        px0 = torch.einsum('lai,ij->laj', px0 - com_hal, R) + com_ref
+    #else: # no motif
     return px0
 
 

rf_diffusion/test_ca_rfd_sm_train.py

@@ -35,6 +35,8 @@ def get_ca_config(overrides=[]):
     initialize(config_path="config/training")
     conf = compose(config_name='test_ca_rfd_sm_train.yaml', overrides=overrides, return_hydra_config=True)
     conf.dataloader.DATAPKL_AA = relative_to_absolute(conf.dataloader.DATAPKL_AA)
+    conf.dataloader.PDB_DIR = relative_to_absolute(conf.dataloader.PDB_DIR)
+    conf.dataloader.MOL_DIR = relative_to_absolute(conf.dataloader.MOL_DIR)
     return conf
 
 def rfold_side_effect(*args, **kwargs): 
@@ -404,4 +406,3 @@ def test_prev(self):
         torch.testing.assert_close(got_msa_prev, want_msa_prev)
         torch.testing.assert_close(got_pair_prev, want_pair_prev)
         torch.testing.assert_close(got_state_prev, want_state_prev)
-

rf_diffusion/test_ca_rfd_unconditional_train.py

@@ -33,6 +33,8 @@ def get_ca_config(overrides=[]):
     initialize(config_path="config/training")
     conf = compose(config_name='test_ca_rfd_unconditional_train.yaml', overrides=overrides, return_hydra_config=True)
     conf.dataloader.DATAPKL_AA = relative_to_absolute(conf.dataloader.DATAPKL_AA)
+    conf.dataloader.PDB_DIR = relative_to_absolute(conf.dataloader.PDB_DIR)
+    conf.dataloader.MOL_DIR = relative_to_absolute(conf.dataloader.MOL_DIR)
     return conf
 
 def rfold_side_effect(*args, **kwargs): 
@@ -379,4 +381,3 @@ def test_prev(self):
         torch.testing.assert_close(got_msa_prev, want_msa_prev)
         torch.testing.assert_close(got_pair_prev, want_pair_prev)
         torch.testing.assert_close(got_state_prev, want_state_prev)
-