# Copyright 2021 AlQuraishi Laboratory # Copyright 2021 DeepMind Technologies Limited # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import sys from env_consts import TEST_INPUT_DIR, TEST_OUTPUT_DIR, CKPT_PATH import json import logging import numpy as np import os import pickle from dockformer.data.data_modules import OpenFoldSingleDataset logging.basicConfig() logger = logging.getLogger(__file__) logger.setLevel(level=logging.INFO) import torch torch_versions = torch.__version__.split(".") torch_major_version = int(torch_versions[0]) torch_minor_version = int(torch_versions[1]) if ( torch_major_version > 1 or (torch_major_version == 1 and torch_minor_version >= 12) ): # Gives a large speedup on Ampere-class GPUs torch.set_float32_matmul_precision("high") torch.set_grad_enabled(False) from dockformer.config import model_config from dockformer.utils.script_utils import (load_models_from_command_line, run_model, save_output_structure, get_latest_checkpoint) from dockformer.utils.tensor_utils import tensor_tree_map def list_files_with_extensions(dir, extensions): return [f for f in os.listdir(dir) if f.endswith(extensions)] def override_config(base_config, overriding_config): for k, v in overriding_config.items(): if isinstance(v, dict): base_config[k] = override_config(base_config[k], v) else: base_config[k] = v return base_config def run_on_folder(input_dir: str, output_dir: str, run_config_path: str, skip_relaxation=True, long_sequence_inference=False, skip_exists=False): config_preset = "initial_training" save_outputs = False device_name = "cuda" if torch.cuda.is_available() else "cpu" run_config = json.load(open(run_config_path)) ckpt_path = CKPT_PATH if ckpt_path is None: ckpt_path = get_latest_checkpoint(os.path.join(run_config["train_output_dir"], "checkpoint")) print("Using checkpoint: ", ckpt_path) config = model_config(config_preset, long_sequence_inference=long_sequence_inference) config = override_config(config, run_config.get("override_conf", {})) model_generator = load_models_from_command_line( config, model_device=device_name, model_checkpoint_path=ckpt_path, output_dir=output_dir) print("Model loaded") model, output_directory = next(model_generator) dataset = OpenFoldSingleDataset(data_dir=input_dir, config=config.data, mode="predict") for i, processed_feature_dict in enumerate(dataset): tag = dataset.get_metadata_for_idx(i)["input_name"] print("Processing", tag) output_name = f"{tag}_predicted" protein_output_path = os.path.join(output_directory, f'{output_name}_protein.pdb') if os.path.exists(protein_output_path) and skip_exists: print("skipping exists", output_name) continue # turn into a batch of size 1 processed_feature_dict = {key: value.unsqueeze(0).to(device_name) for key, value in processed_feature_dict.items()} out = run_model(model, processed_feature_dict, tag, output_dir) # Toss out the recycling dimensions --- we don't need them anymore processed_feature_dict = tensor_tree_map( lambda x: np.array(x[..., -1].cpu()), processed_feature_dict ) out = tensor_tree_map(lambda x: np.array(x.cpu()), out) affinity_output_path = os.path.join(output_directory, f'{output_name}_affinity.json') # affinity = torch.sum(torch.softmax(torch.tensor(out["affinity_2d_logits"]), -1) * torch.linspace(0, 15, 32), # dim=-1).item() affinity_2d = torch.sum(torch.softmax(torch.tensor(out["affinity_2d_logits"]), -1) * torch.linspace(0, 15, 32), dim=-1).item() affinity_1d = torch.sum(torch.softmax(torch.tensor(out["affinity_1d_logits"]), -1) * torch.linspace(0, 15, 32), dim=-1).item() affinity_cls = torch.sum(torch.softmax(torch.tensor(out["affinity_cls_logits"]), -1) * torch.linspace(0, 15, 32), dim=-1).item() affinity_2d_max = torch.linspace(0, 15, 32)[torch.argmax(torch.tensor(out["affinity_2d_logits"]))].item() affinity_1d_max = torch.linspace(0, 15, 32)[torch.argmax(torch.tensor(out["affinity_1d_logits"]))].item() affinity_cls_max = torch.linspace(0, 15, 32)[torch.argmax(torch.tensor(out["affinity_cls_logits"]))].item() print("Affinity: ", affinity_2d, affinity_cls, affinity_1d) with open(affinity_output_path, "w") as f: json.dump({"affinity_2d": affinity_2d, "affinity_1d": affinity_1d, "affinity_cls": affinity_cls, "affinity_2d_max": affinity_2d_max, "affinity_1d_max": affinity_1d_max, "affinity_cls_max": affinity_cls_max}, f) # binding_site = torch.sigmoid(torch.tensor(out["binding_site_logits"])) * 100 # binding_site = binding_site[:processed_feature_dict["aatype"].shape[1]].flatten() # predicted_contacts = torch.sigmoid(torch.tensor(out["inter_contact_logits"])) * 100 # binding_site = torch.max(predicted_contacts, dim=2).values.flatten() ligand_output_path = os.path.join(output_directory, f"{output_name}_ligand_{{i}}.sdf") protein_mask = processed_feature_dict["protein_mask"][0].astype(bool) ligand_mask = processed_feature_dict["ligand_mask"][0].astype(bool) save_output_structure( aatype=processed_feature_dict["aatype"][0][protein_mask], residue_index=processed_feature_dict["in_chain_residue_index"][0], chain_index=processed_feature_dict["chain_index"][0], plddt=out["plddt"][0][protein_mask], final_atom_protein_positions=out["final_atom_positions"][0][protein_mask], final_atom_mask=out["final_atom_mask"][0][protein_mask], ligand_atype=processed_feature_dict["ligand_atype"][0].astype(int), ligand_chiralities=processed_feature_dict["ligand_chirality"][0].astype(int), ligand_charges= processed_feature_dict["ligand_charge"][0].astype(int), ligand_bonds=processed_feature_dict["ligand_bonds"][0].astype(int), ligand_idx=processed_feature_dict["ligand_idx"][0].astype(int), ligand_bonds_idx=processed_feature_dict["ligand_bonds_idx"][0].astype(int), final_ligand_atom_positions=out["final_atom_positions"][0][ligand_mask][:, 1, :], # only ca index protein_output_path=protein_output_path, ligand_output_path=ligand_output_path, ) logger.info(f"Output written to {protein_output_path}...") if not skip_relaxation: # Relax the prediction. logger.info(f"Running relaxation on {protein_output_path}...") from dockformer.utils.relax import relax_complex try: relax_complex(protein_output_path, ligand_output_path, os.path.join(output_directory, f'{output_name}_protein_relaxed.pdb'), os.path.join(output_directory, f'{output_name}_ligand_relaxed.sdf')) except Exception as e: logger.error(f"Failed to relax {protein_output_path} due to {e}...") if save_outputs: output_dict_path = os.path.join( output_directory, f'{output_name}_output_dict.pkl' ) with open(output_dict_path, "wb") as fp: pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL) logger.info(f"Model output written to {output_dict_path}...") if __name__ == "__main__": config_path = sys.argv[1] if len(sys.argv) > 1 else os.path.join(os.path.dirname(__file__), "run_config.json") input_dir, output_dir = TEST_INPUT_DIR, TEST_OUTPUT_DIR options = {"skip_relaxation": True, "long_sequence_inference": False} if len(sys.argv) > 3: input_dir = sys.argv[2] output_dir = sys.argv[3] if "--relax" in sys.argv: options["skip_relaxation"] = False if "--long" in sys.argv: options["long_sequence_inference"] = True if "--allow-skip" in sys.argv: options["skip_exists"] = True run_on_folder(input_dir, output_dir, config_path, **options)