|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import argparse
|
| import json
|
| import os
|
| import random
|
| import subprocess
|
| import sys
|
| from collections import defaultdict
|
| from glob import glob
|
| from typing import Any, Dict, List
|
|
|
| import numpy as np
|
| import pandas as pd
|
| import torch
|
| from Bio.PDB import MMCIFParser as BioMMCIFParser
|
| from Bio.PDB import *
|
| from Bio.PDB.Polypeptide import is_aa
|
| from biotite.structure import get_residue_starts
|
| from natsort import natsorted
|
|
|
| from protenix.data import ccd
|
| from protenix.data.parser import MMCIFParser
|
|
|
| three_to_one = {
|
| "ALA": "A",
|
| "ARG": "R",
|
| "ASN": "N",
|
| "ASP": "D",
|
| "CYS": "C",
|
| "GLN": "Q",
|
| "GLU": "E",
|
| "GLY": "G",
|
| "HIS": "H",
|
| "ILE": "I",
|
| "LEU": "L",
|
| "LYS": "K",
|
| "MET": "M",
|
| "PHE": "F",
|
| "PRO": "P",
|
| "SER": "S",
|
| "THR": "T",
|
| "TRP": "W",
|
| "TYR": "Y",
|
| "VAL": "V",
|
| }
|
|
|
|
|
| def seed_everything(seed, deterministic):
|
| random.seed(seed)
|
| np.random.seed(seed)
|
| torch.random.manual_seed(seed)
|
| torch.cuda.manual_seed_all(seed)
|
| if deterministic:
|
| torch.backends.cudnn.benchmark = False
|
|
|
| torch.backends.cudnn.deterministic = True
|
|
|
| torch.use_deterministic_algorithms(True)
|
|
|
| os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
|
|
|
|
| def str2bool(v):
|
| if isinstance(v, bool):
|
| return v
|
| if v.lower() in ("yes", "true", "t", "y", "1"):
|
| return True
|
| elif v.lower() in ("no", "false", "f", "n", "0"):
|
| return False
|
| else:
|
| raise argparse.ArgumentTypeError("Error")
|
|
|
|
|
| def convert_cif_to_pdb(
|
| cif_path: str,
|
| out_pdb_path: str = None,
|
| binder_chains: list[str] = None,
|
| trim_chain_ids=True,
|
| resname_mapping: dict = {"xpb": "GLY"},
|
| ) -> None:
|
| """
|
| Convert a CIF file to a PDB file.
|
| Args:
|
| cif_path: Path to the CIF file.
|
| out_pdb_path: Path to save the PDB file. If None, will save to the same directory as the CIF file.
|
| binder_chains: List of chain IDs to consider as binders. If None, will use all chains.
|
| trim_chain_ids: If True, will trim the chain IDs to the first character.
|
| resname_mapping: Dictionary of residue names to replace. If None, will not replace residue names.
|
| Returns:
|
| List of condition chain IDs and list of binder chain IDs.
|
| """
|
| if out_pdb_path is None:
|
| assert cif_path.endswith(".cif")
|
| out_pdb_path = cif_path[: -len(".cif")] + ".pdb"
|
|
|
| parser = BioMMCIFParser(QUIET=True)
|
| structure = parser.get_structure("protein", cif_path)
|
|
|
| if trim_chain_ids:
|
|
|
| original_ids = []
|
| trimmed_ids = []
|
| for model in structure:
|
| for chain in model:
|
| original_ids.append(chain.id)
|
| trimmed_ids.append(chain.id[0])
|
| if len(set(trimmed_ids)) < len(set(original_ids)):
|
| raise ValueError(
|
| "Chain ID collision detected after trimming to 1 character:\n"
|
| f"Original IDs: {sorted(set(original_ids))}\n"
|
| f"Trimmed IDs: {sorted(set(trimmed_ids))}"
|
| )
|
|
|
| new_cond_chains, new_binder_chains = [], []
|
| for model in structure:
|
| for chain in model:
|
| new_chain_id = chain.id[0] if trim_chain_ids else chain.id
|
| if binder_chains is not None and chain.id in binder_chains:
|
| new_binder_chains.append(new_chain_id)
|
| else:
|
| new_cond_chains.append(new_chain_id)
|
| if trim_chain_ids:
|
| chain.id = chain.id[0]
|
| if resname_mapping is not None:
|
| for res in chain:
|
| if res.resname in resname_mapping:
|
| res.resname = resname_mapping[res.resname]
|
| if binder_chains is not None and len(new_binder_chains) == 0:
|
| raise ValueError(f"binder chains {binder_chains} not found in the cif file.")
|
|
|
| if os.path.exists(out_pdb_path):
|
| print(
|
| f"[WARNING] PDB file {out_pdb_path} already exists when trying to convert a CIF file to it"
|
| )
|
| io = PDBIO()
|
| io.set_structure(structure)
|
| io.save(out_pdb_path)
|
| return new_cond_chains, new_binder_chains
|
|
|
|
|
| def find_cif_files(folder_path):
|
| pdb_files = []
|
| for filename in os.listdir(folder_path):
|
| if filename.endswith(".cif"):
|
| full_path = os.path.join(folder_path, filename)
|
| pdb_files.append(full_path)
|
| return pdb_files
|
|
|
|
|
| def find_cond_chains(cif_path):
|
|
|
| mmcif_parser = BioMMCIFParser()
|
| structure = mmcif_parser.get_structure("protein", cif_path)
|
| cond_chains = [chain.id for chain in structure[0]]
|
| return cond_chains[:-1]
|
|
|
|
|
| def find_binder_chains(cif_path, condition_chains):
|
| try:
|
| mmcif_parser = BioMMCIFParser()
|
| structure = mmcif_parser.get_structure("protein", cif_path)
|
| except:
|
| print("find binder chains fail: ", cif_path, condition_chains)
|
| raise ValueError()
|
|
|
|
|
|
|
| all_chain_ids = []
|
| for chain in structure[0]:
|
| residues = [
|
| res for res in chain if res.resname == "xpb" or is_aa(res, standard=False)
|
| ]
|
| if len(residues) > 0:
|
| all_chain_ids.append(chain.id)
|
| assert all(c in all_chain_ids for c in condition_chains)
|
| binder_chains = list(set(all_chain_ids) - set(condition_chains))
|
| return binder_chains
|
|
|
|
|
| def convert_cifs_to_pdbs(
|
| input_dir: str,
|
| out_pdb_dir: str = None,
|
| condition_chains: list[str] = None,
|
| resname_mapping: dict = {"xpb": "GLY"},
|
| ):
|
| """
|
| Converts all mmCIF (.cif) files in a directory to PDB format.
|
|
|
| This function scans the input directory for `.cif` files, infers binder chains
|
| from the first file, and converts all files to `.pdb` format using the same
|
| binder chains. The output PDB files are saved in the specified output directory.
|
|
|
| Args:
|
| input_dir (str): Path to the directory containing .cif files.
|
| out_pdb_dir (Optional[str]): Directory to save the converted .pdb files.
|
| If None, a `converted_pdbs` subdirectory will be created in `input_dir`.
|
| condition_chains (list[str]): List of condition chain IDs used to infer binder chains.
|
| resname_mapping (dict): Dictionary of residue names to replace.
|
|
|
| Returns:
|
| tuple[list[str], str, list[str], list[str]]:
|
| - Output PDB dir
|
| - List of output PDB names.
|
| - List of condition chain IDs.
|
| - List of binder chain IDs.
|
| """
|
| if not os.path.exists(input_dir):
|
| raise FileNotFoundError(input_dir)
|
| assert os.path.isdir(input_dir), "The input should be a directory"
|
| all_cif_files = find_cif_files(input_dir)
|
| if len(all_cif_files) == 0:
|
| print(f"[WARNING] Can not find cif files in {input_dir}")
|
| return [], None, None
|
| if condition_chains is None:
|
| condition_chains = find_cond_chains(all_cif_files[0])
|
| binder_chains = find_binder_chains(all_cif_files[0], condition_chains)
|
| if out_pdb_dir is None:
|
| out_pdb_dir = os.path.join(input_dir, "converted_pdbs")
|
| os.makedirs(out_pdb_dir, exist_ok=True)
|
|
|
| pdb_names = []
|
| for cif_file in all_cif_files:
|
| assert cif_file.endswith(".cif")
|
| cur_binder_chains = find_binder_chains(cif_file, condition_chains)
|
| assert set(binder_chains) == set(cur_binder_chains), (
|
| f"Binder chains in {cif_file} differ from those in the first file: "
|
| f"{set(cur_binder_chains)} != {set(binder_chains)}"
|
| )
|
| prefix = os.path.basename(cif_file)[: -len(".cif")]
|
| pdb_names.append(prefix)
|
| pdb_file = os.path.join(out_pdb_dir, f"{prefix}.pdb")
|
| new_cond_chains, new_binder_chains = convert_cif_to_pdb(
|
| cif_path=cif_file,
|
| out_pdb_path=pdb_file,
|
| binder_chains=binder_chains,
|
| resname_mapping=resname_mapping,
|
| )
|
| pdb_names = sorted(pdb_names)
|
| return out_pdb_dir, pdb_names, new_cond_chains, new_binder_chains
|
|
|
|
|
| def merge_list_of_dicts_on_key(
|
| list1: List[Dict[str, Any]], list2: List[Dict[str, Any]], key: str
|
| ) -> List[Dict[str, Any]]:
|
| """
|
| Merge two lists of dictionaries based on a common key.
|
|
|
| Args:
|
| list1: First list of dictionaries.
|
| list2: Second list of dictionaries.
|
| key: The key to merge on.
|
|
|
| Returns:
|
| A list of merged dictionaries.
|
| """
|
| index2 = {d[key]: d for d in list2}
|
| merged = []
|
| for d1 in list1:
|
| k = d1[key]
|
| if k in index2:
|
| merged.append({**d1, **index2[k]})
|
| else:
|
| merged.append(d1)
|
| return merged
|
|
|
|
|
| def concat_dict_values(dict_list: List[Dict[str, Any]]) -> Dict[str, List[Any]]:
|
| result = defaultdict(list)
|
| for d in dict_list:
|
| for k, v in d.items():
|
| result[k].append(v)
|
| return dict(result)
|
|
|
|
|
| def save_eval_results(
|
| sample_df,
|
| summary_dict,
|
| root_dir,
|
| sample_fn: str = "sample_level_output.csv",
|
| summary_fn: str = "summary_output.json",
|
| ):
|
| sample_save_path = os.path.join(root_dir, sample_fn)
|
| summary_save_path = os.path.join(root_dir, summary_fn)
|
| sample_df.to_csv(sample_save_path, index=False)
|
| with open(summary_save_path, "w") as f:
|
| json.dump(summary_dict, f, indent=4)
|
| return sample_save_path, summary_save_path
|
|
|
|
|
| def extract_chain_sequence(pdb_file, chain_id="R"):
|
| parser = PDBParser(QUIET=True)
|
| structure = parser.get_structure("X", pdb_file)
|
| for model in structure:
|
| if chain_id in model:
|
| chain = model[chain_id]
|
| seq = []
|
| current_resid = None
|
| for residue in chain:
|
|
|
| if current_resid is not None:
|
| gap = residue.id[1] - current_resid - 1
|
| seq.extend("X" * gap)
|
| current_resid = residue.id[1]
|
|
|
| try:
|
| seq.append(three_to_one[residue.get_resname()])
|
| except KeyError:
|
| seq.append("X")
|
|
|
| return "".join(seq)
|
|
|
|
|
| def extract_chain_sequence_from_mmcif(cif_file, chain_id="B"):
|
| parser = MMCIFParser(cif_file)
|
| atom_array = parser.get_structure(
|
| altloc="first", model=1, bond_lenth_threshold=None
|
| )
|
| chain_atom_array = atom_array[atom_array.chain_id == chain_id]
|
| starts = get_residue_starts(chain_atom_array, add_exclusive_stop=True)
|
| res_names = chain_atom_array.res_name[starts[:-1]].tolist()
|
| seq = ccd.res_names_to_sequence(res_names)
|
| return seq
|
|
|
|
|
| def prepare_tasks(task_json_path, input_dir, save_dir=None, task_indices=None):
|
| with open(task_json_path, "r") as f:
|
| data = json.load(f)
|
| if task_indices is not None:
|
| data = [data[i] for i in task_indices]
|
|
|
| inputs = []
|
| for x in data:
|
| if "condition" in x or "sequences" in x:
|
| task = "binder"
|
| else:
|
| task = "monomer"
|
|
|
|
|
| input_data_dir = glob(
|
| os.path.join(input_dir, x["name"], "seed_*", "predictions")
|
| )
|
| if len(input_data_dir) == 0:
|
| print(f"Could not find data to eval for name: {x['name']}")
|
| continue
|
| input_data_dir = input_data_dir[0]
|
| rel_path = os.path.relpath(input_data_dir, input_dir)
|
| out_dir = (
|
| input_data_dir if save_dir is None else os.path.join(save_dir, rel_path)
|
| )
|
| summary_path = os.path.join(out_dir, "summary_output.json")
|
| if os.path.exists(summary_path):
|
| print(f"Found existing {summary_path}, skip!")
|
| continue
|
|
|
|
|
| pdb_dir, pdb_names, cond_chains, binder_chains = convert_cifs_to_pdbs(
|
| input_data_dir,
|
| out_pdb_dir=os.path.join(out_dir, "converted_pdbs"),
|
| )
|
|
|
| if len(binder_chains) != 1:
|
| raise ValueError(
|
| f"Multiple binder chains are not supported for now! cond chains: {cond_chains}, binder chains: {binder_chains}"
|
| )
|
|
|
| inputs.append(
|
| {
|
| "task": task,
|
| "name": x["name"],
|
| "pdb_dir": pdb_dir,
|
| "pdb_names": pdb_names,
|
| "cond_chains": cond_chains,
|
| "binder_chains": binder_chains,
|
| "out_dir": out_dir,
|
| }
|
| )
|
|
|
| inputs = natsorted(inputs, key=lambda x: x["name"])
|
| for x in inputs:
|
| print(x)
|
| return inputs
|
|
|
