|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import json
|
| import logging
|
| import os
|
| from typing import Union
|
|
|
| import torch
|
| import torch.distributed as dist
|
| from protenix.config import parse_configs, parse_sys_args
|
| from protenix.config.extend_types import ListValue, RequiredValue
|
| from protenix.utils.distributed import DIST_WRAPPER
|
| from protenix.utils.logger import get_logger
|
|
|
| from pxdbench.pxd_configs.eval import eval_configs
|
| from pxdbench.tasks import get_task_class
|
| from pxdbench.utils import convert_cif_to_pdb
|
|
|
| logger = get_logger(__name__)
|
|
|
|
|
| def run_task(input_data: dict, configs, device_id: int = 0, seed: int = None):
|
| """
|
|
|
| Args:
|
| input_data (dict): A dictionary containing the following keys:
|
| task (str): The task to run. Example: "binder".
|
| pdb_dir (str): The directory containing the PDB files.
|
| pdb_names (list): The names of the PDB files.
|
| cond_chains (list): The chains to condition on.
|
| binder_chains (list): The chains to bind.
|
| out_dir (str): The root directory. Outputs will be saved here.
|
| configs (dict): A dictionary containing the configuration for the task.
|
| device_id (int): Device ID to run the task on.
|
| seed (int): Random seed to use.
|
| Returns:
|
| dict: A dictionary containing the results of the task.
|
| """
|
| if device_id >= 0 and not torch.cuda.is_available():
|
| raise ValueError("device_id must be -1 (CPU) or a valid GPU ID")
|
| task = input_data["task"]
|
| task_cls = get_task_class(task)
|
| task_cfg = configs.get(task)
|
| task = task_cls(input_data, task_cfg, device_id, seed)
|
| return task.run()
|
|
|
|
|
| def split_string(s):
|
| return s.split(",")
|
|
|
|
|
| def get_file_name_list(file_name_list: Union[str, list]):
|
| """
|
| Get file names from a string or a file.
|
| Args:
|
| file_name_list (str or list): A string or a list containing file names.
|
| If a string, it should be a comma-separated list of file names.
|
| If a file, it should contain one file name per line.
|
| Returns:
|
| list: A list of file names.
|
| """
|
| if isinstance(file_name_list, list):
|
| return file_name_list
|
| if os.path.exists(file_name_list):
|
| file_names = []
|
| with open(file_name_list, "r") as f:
|
| for line in f.readlines():
|
| line = line.strip()
|
| if line:
|
| file_names.append(line)
|
| return file_names
|
| else:
|
| file_name_list = file_name_list.split(",")
|
| return file_name_list
|
|
|
|
|
| def find_files_with_ext(folder_path, ext="cif"):
|
| pdb_files = []
|
| for filename in os.listdir(folder_path):
|
| if filename.endswith(f".{ext}"):
|
| pdb_files.append(filename[: -(len(ext) + 1)])
|
| return sorted(pdb_files)
|
|
|
|
|
| def get_chains_from_pdb(pdb_path):
|
| chains = set()
|
| with open(pdb_path, "r") as f:
|
| for line in f:
|
| if line.startswith("ATOM") or line.startswith("HETATM"):
|
| chain_id = line[21].strip()
|
| if chain_id:
|
| chains.add(chain_id)
|
| return chains
|
|
|
|
|
| def prepare_tasks_from_mmcif(
|
| data_dir: str,
|
| file_name_list: list,
|
| binder_chains: list[str],
|
| cond_chains: list[str],
|
| dump_dir: str,
|
| ):
|
| """
|
| Prepare tasks from mmCIF files.
|
| Args:
|
| data_dir (str): The directory containing the mmCIF files.
|
| file_name_list (str or list): A string or a list containing file names.
|
| If a string, it should be a comma-separated list of file names.
|
| If a file, it should contain one file name per line.
|
| binder_chains (list): The binder chain IDs.
|
| cond_chains (list): The chains to condition on. Could be auto infered by binder_chains.
|
| root_dir (str): The root directory. Outputs will be saved here.
|
| Returns:
|
| dict: A dictionary containing the results of the task.
|
| """
|
| n_cond_chains, n_binder_chains = None, None
|
|
|
| valid_file_name_list = []
|
| for file_name in file_name_list:
|
| mmcif_path = os.path.join(data_dir, f"{file_name}.cif")
|
| if not os.path.exists(mmcif_path):
|
| logger.warning(f"Could not find {mmcif_path}, skip!")
|
| continue
|
| else:
|
| valid_file_name_list.append(file_name)
|
| pdb_path = os.path.join(data_dir, "converted_pdbs", f"{file_name}.pdb")
|
| os.makedirs(os.path.dirname(pdb_path), exist_ok=True)
|
| new_cond_chains, new_binder_chains = convert_cif_to_pdb(
|
| cif_path=mmcif_path,
|
| out_pdb_path=pdb_path,
|
| binder_chains=binder_chains,
|
| )
|
| if n_cond_chains is None:
|
| n_cond_chains = list(new_cond_chains)
|
| else:
|
| assert set(n_cond_chains) == set(new_cond_chains)
|
| if n_binder_chains is None:
|
| n_binder_chains = list(new_binder_chains)
|
| else:
|
| assert set(n_binder_chains) == set(new_binder_chains)
|
|
|
| logger.info(
|
| f"Found {len(valid_file_name_list)} valid files, cond chains: {n_cond_chains}, binder chains: {n_binder_chains}"
|
| )
|
| input_data = {
|
| "task": "binder",
|
| "name": os.path.basename(data_dir),
|
| "pdb_dir": os.path.join(data_dir, "converted_pdbs"),
|
| "pdb_names": valid_file_name_list,
|
| "cond_chains": n_cond_chains,
|
| "binder_chains": n_binder_chains,
|
| "out_dir": dump_dir,
|
| }
|
| return input_data
|
|
|
|
|
| def prepare_tasks_from_pdb(
|
| data_dir: str,
|
| file_name_list: list,
|
| binder_chains: list[str],
|
| cond_chains: list[str],
|
| dump_dir: str,
|
| ):
|
| n_cond_chains = None
|
| valid_file_name_list = []
|
| for file_name in file_name_list:
|
| pdb_path = os.path.join(data_dir, f"{file_name}.pdb")
|
| if not os.path.exists(pdb_path):
|
| logger.warning(f"Could not find {pdb_path}, skip!")
|
| continue
|
| else:
|
| valid_file_name_list.append(file_name)
|
|
|
| if cond_chains == [""]:
|
| chains = get_chains_from_pdb(pdb_path)
|
| new_cond_chains = chains - set(binder_chains)
|
| else:
|
| new_cond_chains = cond_chains
|
| if n_cond_chains is None:
|
| n_cond_chains = list(new_cond_chains)
|
| else:
|
| assert set(n_cond_chains) == set(new_cond_chains)
|
|
|
| return {
|
| "task": "binder",
|
| "pdb_dir": data_dir,
|
| "name": os.path.basename(data_dir),
|
| "pdb_names": valid_file_name_list,
|
| "cond_chains": n_cond_chains,
|
| "binder_chains": binder_chains,
|
| "out_dir": dump_dir,
|
| }
|
|
|
|
|
| def prepare_tasks_from_json(json_path):
|
| """
|
| Prepare tasks from a JSON file.
|
|
|
| Args:
|
| json_path (str): Path to the JSON file.
|
|
|
| Returns:
|
| list: List of tasks.
|
| """
|
| with open(json_path, "r") as f:
|
| tasks = json.load(f)
|
| if isinstance(tasks, dict):
|
| tasks = [tasks]
|
| return tasks
|
|
|
|
|
| class EvalRunner(object):
|
| def __init__(self, configs):
|
| self.configs = configs
|
| self.init_env()
|
|
|
| def init_env(self) -> None:
|
| self.print(
|
| f"Distributed environment: world size: {DIST_WRAPPER.world_size}, "
|
| + f"global rank: {DIST_WRAPPER.rank}, local rank: {DIST_WRAPPER.local_rank}"
|
| )
|
| self.use_cuda = torch.cuda.device_count() > 0
|
| if self.use_cuda:
|
| self.device = torch.device("cuda:{}".format(DIST_WRAPPER.local_rank))
|
| os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
| all_gpu_ids = ",".join(str(x) for x in range(torch.cuda.device_count()))
|
| devices = os.getenv("CUDA_VISIBLE_DEVICES", all_gpu_ids)
|
| logging.info(
|
| f"LOCAL_RANK: {DIST_WRAPPER.local_rank} - CUDA_VISIBLE_DEVICES: [{devices}]"
|
| )
|
| torch.cuda.set_device(self.device)
|
| else:
|
| self.device = torch.device("cpu")
|
| if DIST_WRAPPER.world_size > 1:
|
| dist.init_process_group(backend="nccl")
|
| logging.info("Finished init ENV.")
|
|
|
| def print(self, msg: str):
|
| if DIST_WRAPPER.rank == 0:
|
| logger.info(msg)
|
|
|
| def run(self, input_data_list):
|
| for input_data in input_data_list:
|
| run_task(
|
| input_data,
|
| self.configs,
|
| device_id=DIST_WRAPPER.local_rank,
|
| seed=self.configs.seed,
|
| )
|
| logging.info("Eval done!")
|
|
|
|
|
| def main():
|
|
|
| config_dict = {
|
| "file_name_list": "",
|
| "json_path": "",
|
| "data_dir": "",
|
| "dump_dir": RequiredValue(str),
|
| "cond_chains": ListValue([""]),
|
| "binder_chains": ListValue([""]),
|
| "is_mmcif": False,
|
| "orig_seqs_json": "",
|
| "seed": 2025,
|
| **eval_configs,
|
| }
|
| configs = parse_configs(config_dict, arg_str=parse_sys_args())
|
|
|
|
|
| if configs.json_path:
|
| logger.info("Prepare tasks from json file.")
|
| input_data_list = prepare_tasks_from_json(configs.json_path)
|
| if DIST_WRAPPER.world_size > 1:
|
| input_data_list = input_data_list[
|
| DIST_WRAPPER.rank :: DIST_WRAPPER.world_size
|
| ]
|
| for input_data in input_data_list:
|
| input_data["out_dir"] += f"_rank{DIST_WRAPPER.rank}"
|
|
|
| else:
|
| if not configs.file_name_list:
|
| if configs.is_mmcif:
|
| fn_list = find_files_with_ext(configs.data_dir, "cif")
|
| else:
|
| fn_list = find_files_with_ext(configs.data_dir, "pdb")
|
| else:
|
| fn_list = get_file_name_list(configs.file_name_list)
|
|
|
| if DIST_WRAPPER.world_size > 1:
|
| fn_list = fn_list[DIST_WRAPPER.rank :: DIST_WRAPPER.world_size]
|
|
|
| if configs.is_mmcif:
|
| logger.info("Prepare tasks from mmCIF files.")
|
| input_data = prepare_tasks_from_mmcif(
|
| data_dir=configs.data_dir,
|
| file_name_list=fn_list,
|
| binder_chains=configs.binder_chains,
|
| cond_chains=configs.cond_chains,
|
| dump_dir=configs.dump_dir,
|
| )
|
| else:
|
| logger.info("Prepare tasks from PDB files.")
|
| input_data = prepare_tasks_from_pdb(
|
| data_dir=configs.data_dir,
|
| file_name_list=fn_list,
|
| binder_chains=configs.binder_chains,
|
| cond_chains=configs.cond_chains,
|
| dump_dir=configs.dump_dir,
|
| )
|
| if len(configs.orig_seqs_json) > 0:
|
| input_data["orig_seqs_json"] = configs.orig_seqs_json
|
|
|
| if DIST_WRAPPER.world_size > 1:
|
| input_data["out_dir"] += f"_rank{DIST_WRAPPER.rank}"
|
| input_data_list = [input_data]
|
|
|
|
|
| runner = EvalRunner(configs)
|
| runner.run(input_data_list)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
