diff --git a/make_conda.sh b/make_conda.sh index 04a0d7b..e1e1c71 100755 --- a/make_conda.sh +++ b/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 \ diff --git a/rf_diffusion/conditioning.py b/rf_diffusion/conditioning.py index 79de6ec..65f821c 100644 --- a/rf_diffusion/conditioning.py +++ b/rf_diffusion/conditioning.py @@ -313,8 +313,8 @@ 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 ### ###################################### @@ -322,17 +322,16 @@ def __call__(self, # 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 diff --git a/rf_diffusion/config/training/test_ca_rfd_sm_train.yaml b/rf_diffusion/config/training/test_ca_rfd_sm_train.yaml index e8e1c28..8ff9aab 100644 --- a/rf_diffusion/config/training/test_ca_rfd_sm_train.yaml +++ b/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, diff --git a/rf_diffusion/config/training/test_ca_rfd_unconditional_train.yaml b/rf_diffusion/config/training/test_ca_rfd_unconditional_train.yaml index a41f231..b02e19b 100644 --- a/rf_diffusion/config/training/test_ca_rfd_unconditional_train.yaml +++ b/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, diff --git a/rf_diffusion/run_inference.py b/rf_diffusion/run_inference.py index 2c1d80b..f79bab6 100755 --- a/rf_diffusion/run_inference.py +++ b/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 diff --git a/rf_diffusion/test_ca_rfd_sm_train.py b/rf_diffusion/test_ca_rfd_sm_train.py index 8f36bbd..ddc8d19 100644 --- a/rf_diffusion/test_ca_rfd_sm_train.py +++ b/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) - diff --git a/rf_diffusion/test_ca_rfd_unconditional_train.py b/rf_diffusion/test_ca_rfd_unconditional_train.py index a8a23bb..0f4b041 100644 --- a/rf_diffusion/test_ca_rfd_unconditional_train.py +++ b/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) -