import json
import os
import shutil
import random
import sys
import time
from typing import List, Tuple, Optional
import Bio.PDB
import Bio.SeqUtils
import pandas as pd
import numpy as np
import requests
from rdkit import Chem
from rdkit.Chem import AllChem
BASE_FOLDER = "/tmp/"
OUTPUT_FOLDER = f"{BASE_FOLDER}/processed"
# https://storage.googleapis.com/plinder/2024-06/v2/index/annotation_table.parquet
PLINDER_ANNOTATIONS = f'{BASE_FOLDER}/raw_data/2024-06_v2_index_annotation_table.parquet'
# https://storage.googleapis.com/plinder/2024-06/v2/splits/split.parquet
PLINDER_SPLITS = f'{BASE_FOLDER}/raw_data/2024-06_v2_splits_split.parquet'
# https://console.cloud.google.com/storage/browser/_details/plinder/2024-06/v2/links/kind%3Dapo/links.parquet
PLINDER_LINKED_APO_MAP = f"{BASE_FOLDER}/raw_data/2024-06_v2_links_kind=apo_links.parquet"
# https://console.cloud.google.com/storage/browser/_details/plinder/2024-06/v2/links/kind%3Dpred/links.parquet
PLINDER_LINKED_PRED_MAP = f"{BASE_FOLDER}/raw_data/2024-06_v2_links_kind=pred_links.parquet"
# https://storage.googleapis.com/plinder/2024-06/v2/linked_structures/apo.zip
PLINDER_LINKED_APO_STRUCTURES = f"{BASE_FOLDER}/raw_data/2024-06_v2_linked_structures_apo"
# https://storage.googleapis.com/plinder/2024-06/v2/linked_structures/pred.zip
PLINDER_LINKED_PRED_STRUCTURES = f"{BASE_FOLDER}/raw_data/2024-06_v2_linked_structures_pred"
GSUTIL_PATH = f"{BASE_FOLDER}/google-cloud-sdk/bin/gsutil"
def get_cached_systems_to_train(recompute=False):
output_path = os.path.join(OUTPUT_FOLDER, "to_train.pickle")
if os.path.exists(output_path) and not recompute:
return pd.read_pickle(output_path)
"""
full:
loaded 1357906 409726 163816 433865
loaded 990260 409726 125818 106411
joined splits 409726
Has splits 311008
unique systems 311008
split
train 309140
test 1036
val 832
Name: count, dtype: int64
Has affinity 36856
Has affinity by splits split
train 36598
test 142
val 116
Name: count, dtype: int64
Total systems before pred 311008
Total systems after pred 311008
Has pred 83487
Has apo 75127
Has both 51506
Has either 107108
columns Index(['system_id', 'entry_pdb_id', 'ligand_binding_affinity',
'entry_release_date', 'system_pocket_UniProt',
'system_num_protein_chains', 'system_num_ligand_chains', 'uniqueness',
'split', 'cluster', 'cluster_for_val_split',
'system_pass_validation_criteria', 'system_pass_statistics_criteria',
'system_proper_num_ligand_chains', 'system_proper_pocket_num_residues',
'system_proper_num_interactions',
'system_proper_ligand_max_molecular_weight',
'system_has_binding_affinity', 'system_has_apo_or_pred', '_bucket_id',
'linked_pred_id', 'linked_apo_id'],
dtype='object')
total systems 311008
"""
systems = pd.read_parquet(PLINDER_ANNOTATIONS,
columns=['system_id', 'entry_pdb_id', 'ligand_binding_affinity',
'entry_release_date', 'system_pocket_UniProt', 'entry_resolution',
'system_num_protein_chains', 'system_num_ligand_chains'])
splits = pd.read_parquet(PLINDER_SPLITS)
linked_pred = pd.read_parquet(PLINDER_LINKED_PRED_MAP)
linked_apo = pd.read_parquet(PLINDER_LINKED_APO_MAP)
print("loaded", len(systems), len(splits), len(linked_pred), len(linked_apo))
# remove duplicated
systems = systems.drop_duplicates(subset=['system_id'])
splits = splits.drop_duplicates(subset=['system_id'])
linked_pred = linked_pred.drop_duplicates(subset=['reference_system_id'])
linked_apo = linked_apo.drop_duplicates(subset=['reference_system_id'])
print("loaded", len(systems), len(splits), len(linked_pred), len(linked_apo))
# join splits
systems = pd.merge(systems, splits, on='system_id', how='inner')
print("joined splits", len(systems))
systems['_bucket_id'] = systems['entry_pdb_id'].str[1:3]
# leave only with train/val/test splits
systems = systems[systems['split'].isin(['train', 'val', 'test'])]
print("Has splits", len(systems))
print("unique systems", systems['system_id'].nunique())
print(systems["split"].value_counts())
print("Has affinity", len(systems[systems['ligand_binding_affinity'].notna()]))
# print has affinity by splits
print("Has affinity by splits", systems[systems['ligand_binding_affinity'].notna()]['split'].value_counts())
print("Total systems before pred", len(systems))
# join linked structures - allow to not have linked structures
systems = pd.merge(systems, linked_pred[['reference_system_id', 'id']],
left_on='system_id', right_on='reference_system_id',
how='left')
print("Total systems after pred", len(systems))
# Rename the 'id' column from linked_pred to 'linked_pred_id'
systems.rename(columns={'id': 'linked_pred_id'}, inplace=True)
# Merge the result with linked_apo on the same condition
systems = pd.merge(systems, linked_apo[['reference_system_id', 'id']],
left_on='system_id', right_on='reference_system_id',
how='left')
# Rename the 'id' column from linked_apo to 'linked_apo_id'
systems.rename(columns={'id': 'linked_apo_id'}, inplace=True)
# Drop the reference_system_id columns that were added during the merge
systems.drop(columns=['reference_system_id_x', 'reference_system_id_y'], inplace=True)
cluster_sizes = systems["cluster"].value_counts()
systems["cluster_size"] = systems["cluster"].map(cluster_sizes)
# print(systems[['system_id', 'cluster', 'cluster_size']])
print("Has pred", systems['linked_pred_id'].notna().sum())
print("Has apo", systems['linked_apo_id'].notna().sum())
print("Has both", (systems['linked_pred_id'].notna() & systems['linked_apo_id'].notna()).sum())
print("Has either", (systems['linked_pred_id'].notna() | systems['linked_apo_id'].notna()).sum())
print("columns", systems.columns)
systems.to_pickle(output_path)
return systems
def create_conformers(smiles, output_path, num_conformers=100, multiplier_samples=1):
target_mol = Chem.MolFromSmiles(smiles)
target_mol = Chem.AddHs(target_mol)
params = AllChem.ETKDGv3()
params.numThreads = 0 # Use all available threads
params.pruneRmsThresh = 0.1 # Pruning threshold for RMSD
conformer_ids = AllChem.EmbedMultipleConfs(target_mol, numConfs=num_conformers * multiplier_samples, params=params)
# Optional: Optimize each conformer using MMFF94 force field
# for conf_id in conformer_ids:
# AllChem.UFFOptimizeMolecule(target_mol, confId=conf_id)
# remove hydrogen atoms
target_mol = Chem.RemoveHs(target_mol)
# Save aligned conformers to a file (optional)
w = Chem.SDWriter(output_path)
for i, conf_id in enumerate(conformer_ids):
if i >= num_conformers:
break
w.write(target_mol, confId=conf_id)
w.close()
def do_robust_chain_object_renumber(chain: Bio.PDB.Chain.Chain, new_chain_id: str) -> Optional[Bio.PDB.Chain.Chain]:
all_residues = [res for res in chain.get_residues()
if "CA" in res and Bio.SeqUtils.seq1(res.get_resname()) not in ("X", "", " ")]
if not all_residues:
return None
res_and_res_id = [(res, res.get_id()[1]) for res in all_residues]
min_res_id = min([i[1] for i in res_and_res_id])
if min_res_id < 1:
print("Negative res id", chain, min_res_id)
factor = -1 * min_res_id + 1
res_and_res_id = [(res, res_id + factor) for res, res_id in res_and_res_id]
res_and_res_id_no_collisions = []
for res, res_id in res_and_res_id[::-1]:
if res_and_res_id_no_collisions and res_and_res_id_no_collisions[-1][1] == res_id:
# there is a collision, usually an insertion residue
res_and_res_id_no_collisions = [(i, j + 1) for i, j in res_and_res_id_no_collisions]
res_and_res_id_no_collisions.append((res, res_id))
first_res_id = min([i[1] for i in res_and_res_id_no_collisions])
factor = 1 - first_res_id # start from 1
new_chain = Bio.PDB.Chain.Chain(new_chain_id)
res_and_res_id_no_collisions.sort(key=lambda x: x[1])
for res, res_id in res_and_res_id_no_collisions:
chain.detach_child(res.id)
res.id = (" ", res_id + factor, " ")
new_chain.add(res)
return new_chain
def robust_renumber_protein(pdb_path: str, output_path: str):
if pdb_path.endswith(".pdb"):
pdb_parser = Bio.PDB.PDBParser(QUIET=True)
pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
elif pdb_path.endswith(".cif"):
pdb_struct = Bio.PDB.MMCIFParser().get_structure("original_pdb", pdb_path)
else:
raise ValueError("Unknown file type", pdb_path)
assert len(list(pdb_struct)) == 1, "can't extract if more than one model"
model = next(iter(pdb_struct))
chains = list(model.get_chains())
new_model = Bio.PDB.Model.Model(0)
chain_ids = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
for chain, chain_id in zip(chains, chain_ids):
new_chain = do_robust_chain_object_renumber(chain, chain_id)
if new_chain is None:
continue
new_model.add(new_chain)
new_struct = Bio.PDB.Structure.Structure("renumbered_pdb")
new_struct.add(new_model)
io = Bio.PDB.PDBIO()
io.set_structure(new_struct)
io.save(output_path)
def _get_extra(extra_to_save: int, res_before: List[int], res_after: List[int]) -> set:
take_from_before = random.randint(0, extra_to_save)
take_from_after = extra_to_save - take_from_before
if take_from_before > len(res_before):
take_from_after = extra_to_save - len(res_before)
take_from_before = len(res_before)
if take_from_after > len(res_after):
take_from_before = extra_to_save - len(res_after)
take_from_after = len(res_after)
extra_to_add = set()
if take_from_before > 0:
extra_to_add.update(res_before[-take_from_before:])
extra_to_add.update(res_after[:take_from_after])
return extra_to_add
def crop_protein_cont(gt_pdb_path: str, ligand_pos: np.ndarray, output_path: str, max_length: int,
distance_threshold: float):
protein = Chem.MolFromPDBFile(gt_pdb_path, sanitize=False)
ligand_size = ligand_pos.shape[0]
pdb_parser = Bio.PDB.PDBParser(QUIET=True)
gt_model = next(iter(pdb_parser.get_structure("gt_pdb", gt_pdb_path)))
all_res_ids_by_chain = {chain.id: sorted([res.id[1] for res in chain.get_residues() if "CA" in res])
for chain in gt_model.get_chains()}
protein_conf = protein.GetConformer()
protein_pos = protein_conf.GetPositions()
protein_atoms = list(protein.GetAtoms())
assert len(protein_pos) == len(protein_atoms), f"Positions and atoms mismatch in {gt_pdb_path}"
inter_dists = ligand_pos[:, np.newaxis, :] - protein_pos[np.newaxis, :, :]
inter_dists = np.sqrt((inter_dists ** 2).sum(-1))
min_inter_dist_per_protein_atom = inter_dists.min(axis=0)
res_to_save_count = max_length - ligand_size
used_protein_idx = np.where(min_inter_dist_per_protein_atom < distance_threshold)[0]
pocket_residues_by_chain = {}
for idx in used_protein_idx:
res = protein_atoms[idx].GetPDBResidueInfo()
if res.GetIsHeteroAtom():
continue
if res.GetChainId() not in pocket_residues_by_chain:
pocket_residues_by_chain[res.GetChainId()] = set()
# get residue chain
pocket_residues_by_chain[res.GetChainId()].add(res.GetResidueNumber())
if not pocket_residues_by_chain:
print("No pocket residues found")
return -1
# print("pocket_residues_by_chain", pocket_residues_by_chain)
complete_pocket = []
extended_pocket_per_chain = {}
for chain_id, pocket_residues in pocket_residues_by_chain.items():
max_pocket_res = max(pocket_residues)
min_pocket_res = min(pocket_residues)
extended_pocket_per_chain[chain_id] = {res_id for res_id in all_res_ids_by_chain[chain_id]
if min_pocket_res <= res_id <= max_pocket_res}
for res_id in extended_pocket_per_chain[chain_id]:
complete_pocket.append((chain_id, res_id))
# print("extended_pocket_per_chain", pocket_residues_by_chain)
if len(complete_pocket) > res_to_save_count:
total_res_ids = sum([len(res_ids) for res_ids in all_res_ids_by_chain.values()])
total_pocket_res = sum([len(res_ids) for res_ids in pocket_residues_by_chain.values()])
print(f"Too many residues all: {total_res_ids} pocket:{total_pocket_res} {len(complete_pocket)} "
f"(ligand size: {ligand_size})")
return -1
extra_to_save = res_to_save_count - len(complete_pocket)
# divide extra_to_save between chains
for chain_id, pocket_residues in extended_pocket_per_chain.items():
extra_to_save_per_chain = extra_to_save // len(extended_pocket_per_chain)
res_before = [res_id for res_id in all_res_ids_by_chain[chain_id] if res_id < min(pocket_residues)]
res_after = [res_id for res_id in all_res_ids_by_chain[chain_id] if res_id > max(pocket_residues)]
extra_to_add = _get_extra(extra_to_save_per_chain, res_before, res_after)
for res_id in extra_to_add:
complete_pocket.append((chain_id, res_id))
total_res_ids = sum([len(res_ids) for res_ids in all_res_ids_by_chain.values()])
total_pocket_res = sum([len(res_ids) for res_ids in pocket_residues_by_chain.values()])
total_extended_res = sum([len(res_ids) for res_ids in extended_pocket_per_chain.values()])
print(f"Found valid pocket all: {total_res_ids} pocket:{total_pocket_res} {total_extended_res} "
f"{len(complete_pocket)} (ligand size: {ligand_size}) extra: {extra_to_save}")
# print("all_res_ids_by_chain", all_res_ids_by_chain)
# print("complete_pocket", sorted(complete_pocket))
res_to_remove = []
for res in gt_model.get_residues():
if (res.parent.id, res.id[1]) not in complete_pocket or res.id[0].strip() != "" or res.id[2].strip() != "":
res_to_remove.append(res)
for res in res_to_remove:
gt_model[res.parent.id].detach_child(res.id)
io = Bio.PDB.PDBIO()
io.set_structure(gt_model)
io.save(output_path)
return len(complete_pocket)
def crop_protein_simple(gt_pdb_path: str, ligand_pos: np.ndarray, output_path: str, max_length: int):
protein = Chem.MolFromPDBFile(gt_pdb_path, sanitize=False)
ligand_size = ligand_pos.shape[0]
res_to_save_count = max_length - ligand_size
pdb_parser = Bio.PDB.PDBParser(QUIET=True)
gt_model = next(iter(pdb_parser.get_structure("gt_pdb", gt_pdb_path)))
protein_conf = protein.GetConformer()
protein_pos = protein_conf.GetPositions()
protein_atoms = list(protein.GetAtoms())
assert len(protein_pos) == len(protein_atoms), f"Positions and atoms mismatch in {gt_pdb_path}"
inter_dists = ligand_pos[:, np.newaxis, :] - protein_pos[np.newaxis, :, :]
inter_dists = np.sqrt((inter_dists ** 2).sum(-1))
min_inter_dist_per_protein_atom = inter_dists.min(axis=0)
protein_idx_by_dist = np.argsort(min_inter_dist_per_protein_atom)
pocket_residues_by_chain = {}
total_found = 0
for idx in protein_idx_by_dist:
res = protein_atoms[idx].GetPDBResidueInfo()
if res.GetIsHeteroAtom():
continue
if res.GetChainId() not in pocket_residues_by_chain:
pocket_residues_by_chain[res.GetChainId()] = set()
# get residue chain
pocket_residues_by_chain[res.GetChainId()].add(res.GetResidueNumber())
total_found = sum([len(res_ids) for res_ids in pocket_residues_by_chain.values()])
if total_found >= res_to_save_count:
break
print("saved with simple", total_found)
if not pocket_residues_by_chain:
print("No pocket residues found")
return -1
res_to_remove = []
for res in gt_model.get_residues():
if res.id[1] not in pocket_residues_by_chain.get(res.parent.id, set()) \
or res.id[0].strip() != "" or res.id[2].strip() != "":
res_to_remove.append(res)
for res in res_to_remove:
gt_model[res.parent.id].detach_child(res.id)
io = Bio.PDB.PDBIO()
io.set_structure(gt_model)
io.save(output_path)
return total_found
def cif_to_pdb(cif_path: str, pdb_path: str):
protein = Bio.PDB.MMCIFParser().get_structure("s_cif", cif_path)
io = Bio.PDB.PDBIO()
io.set_structure(protein)
io.save(pdb_path)
def get_chain_object_to_seq(chain: Bio.PDB.Chain.Chain) -> str:
res_id_to_res = {res.get_id()[1]: res for res in chain.get_residues() if "CA" in res}
if len(res_id_to_res) == 0:
print("skipping empty chain", chain.get_id())
return ""
seq = ""
for i in range(1, max(res_id_to_res) + 1):
if i in res_id_to_res:
seq += Bio.SeqUtils.seq1(res_id_to_res[i].get_resname())
else:
seq += "X"
return seq
def get_sequence_from_pdb(pdb_path: str) -> Tuple[str, List[int]]:
pdb_parser = Bio.PDB.PDBParser(QUIET=True)
pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
# chain_to_seq = {chain.id: get_chain_object_to_seq(chain) for chain in pdb_struct.get_chains()}
all_chain_seqs = [ get_chain_object_to_seq(chain) for chain in pdb_struct.get_chains()]
chain_lengths = [len(seq) for seq in all_chain_seqs]
return ("X" * 20).join(all_chain_seqs), chain_lengths
from Bio import PDB
from Bio import pairwise2
def extract_sequence(chain):
seq = ''
residues = []
for res in chain.get_residues():
seq_res = Bio.SeqUtils.seq1(res.get_resname())
if seq_res in ('X', "", " "):
continue
seq += seq_res
residues.append(res)
return seq, residues
def map_residues(alignment, residues_gt, residues_pred):
idx_gt = 0
idx_pred = 0
mapping = []
for i in range(len(alignment.seqA)):
aa_gt = alignment.seqA[i]
aa_pred = alignment.seqB[i]
res_gt = None
res_pred = None
if aa_gt != '-':
res_gt = residues_gt[idx_gt]
idx_gt += 1
if aa_pred != '-':
res_pred = residues_pred[idx_pred]
idx_pred +=1
if res_gt and res_pred:
mapping.append((res_gt, res_pred))
return mapping
class ResidueSelect(PDB.Select):
def __init__(self, residues_to_select):
self.residues_to_select = set(residues_to_select)
def accept_residue(self, residue):
return residue in self.residues_to_select
def align_gt_and_input(gt_pdb_path, input_pdb_path, output_gt_path, output_input_path):
parser = PDB.PDBParser(QUIET=True)
gt_structure = parser.get_structure('gt', gt_pdb_path)
pred_structure = parser.get_structure('pred', input_pdb_path)
matched_residues_gt = []
matched_residues_pred = []
used_chain_pred = []
total_mapping_size = 0
for chain_gt in gt_structure.get_chains():
seq_gt, residues_gt = extract_sequence(chain_gt)
best_alignment = None
best_chain_pred = None
best_score = -1
best_residues_pred = None
# Find the best matching chain in pred
for chain_pred in pred_structure.get_chains():
print("checking", chain_pred.get_id(), chain_gt.get_id())
if chain_pred in used_chain_pred:
continue
seq_pred, residues_pred = extract_sequence(chain_pred)
print(seq_gt)
print(seq_pred)
alignments = pairwise2.align.globalxx(seq_gt, seq_pred, one_alignment_only=True)
if not alignments:
continue
print("checking2", chain_pred.get_id(), chain_gt.get_id())
alignment = alignments[0]
score = alignment.score
if score > best_score:
best_score = score
best_alignment = alignment
best_chain_pred = chain_pred
best_residues_pred = residues_pred
if best_alignment:
mapping = map_residues(best_alignment, residues_gt, best_residues_pred)
total_mapping_size += len(mapping)
used_chain_pred.append(best_chain_pred)
for res_gt, res_pred in mapping:
matched_residues_gt.append(res_gt)
matched_residues_pred.append(res_pred)
else:
print(f"No matching chain found for chain {chain_gt.get_id()}")
print(f"Total mapping size: {total_mapping_size}")
# Write new PDB files with only matched residues
io = PDB.PDBIO()
io.set_structure(gt_structure)
io.save(output_gt_path, ResidueSelect(matched_residues_gt))
io.set_structure(pred_structure)
io.save(output_input_path, ResidueSelect(matched_residues_pred))
def validate_matching_input_gt(gt_pdb_path, input_pdb_path):
gt_residues = [res for res in PDB.PDBParser().get_structure('gt', gt_pdb_path).get_residues()]
input_residues = [res for res in PDB.PDBParser().get_structure('input', input_pdb_path).get_residues()]
if len(gt_residues) != len(input_residues):
print(f"Residue count mismatch: {len(gt_residues)} vs {len(input_residues)}")
return -1
for res_gt, res_input in zip(gt_residues, input_residues):
if res_gt.get_resname() != res_input.get_resname():
print(f"Residue name mismatch: {res_gt.get_resname()} vs {res_input.get_resname()}")
return -1
return len(input_residues)
def prepare_system(row, system_folder, output_models_folder, output_jsons_folder, should_overwrite=False):
output_json_path = os.path.join(output_jsons_folder, f"{row['system_id']}.json")
if os.path.exists(output_json_path) and not should_overwrite:
return "Already exists"
plinder_gt_pdb_path = os.path.join(system_folder, f"receptor.pdb")
plinder_gt_ligand_paths = []
plinder_gt_ligands_folder = os.path.join(system_folder, "ligand_files")
gt_output_path = os.path.join(output_models_folder, f"{row['system_id']}_gt.pdb")
gt_output_relative_path = "plinder_models/" + f"{row['system_id']}_gt.pdb"
tmp_input_path = os.path.join(output_models_folder, f"tmp_{row['system_id']}_input.pdb")
protein_input_path = os.path.join(output_models_folder, f"{row['system_id']}_input.pdb")
protein_input_relative_path = "plinder_models/" + f"{row['system_id']}_input.pdb"
print("Copying ground truth files")
if not os.path.exists(plinder_gt_pdb_path):
print("no receptor", plinder_gt_pdb_path)
return "No receptor"
tmp_gt_pdb_path = os.path.join(output_models_folder, f"tmp_{row['system_id']}_gt.pdb")
robust_renumber_protein(plinder_gt_pdb_path, tmp_gt_pdb_path)
ligand_pos_list = []
for ligand_file in os.listdir(plinder_gt_ligands_folder):
if not ligand_file.endswith(".sdf"):
continue
plinder_gt_ligand_paths.append(os.path.join(plinder_gt_ligands_folder, ligand_file))
loaded_ligand = Chem.MolFromMolFile(os.path.join(plinder_gt_ligands_folder, ligand_file))
ligand_pos_list.append(loaded_ligand.GetConformer().GetPositions())
if loaded_ligand is None:
print("failed to load", plinder_gt_ligand_paths[-1])
return "Failed to load ligand"
# Crop ground truth protein, also removes insertion codes
ligand_pos = np.concatenate(ligand_pos_list, axis=0)
res_count_in_protein = crop_protein_cont(tmp_gt_pdb_path, ligand_pos, gt_output_path, max_length=350,
distance_threshold=5)
if res_count_in_protein == -1:
print("Failed to crop protein continously, using simple crop")
crop_protein_simple(tmp_gt_pdb_path, ligand_pos, gt_output_path, max_length=350)
os.remove(tmp_gt_pdb_path)
# Generate input protein structure
input_protein_source = None
if pd.notna(row["linked_apo_id"]):
apo_pdb_path = os.path.join(PLINDER_LINKED_APO_STRUCTURES, f"{row['linked_apo_id']}.cif")
try:
robust_renumber_protein(apo_pdb_path, tmp_input_path)
input_protein_source = "apo"
print("Using input apo", row['linked_apo_id'])
except Exception as e:
print("Problem with apo", e, row["linked_apo_id"], apo_pdb_path)
if not os.path.exists(tmp_input_path) and pd.notna(row["linked_pred_id"]):
pred_pdb_path = os.path.join(PLINDER_LINKED_PRED_STRUCTURES, f"{row['linked_pred_id']}.cif")
try:
# cif_to_pdb(pred_pdb_path, tmp_input_path)
robust_renumber_protein(pred_pdb_path, tmp_input_path)
input_protein_source = "pred"
print("Using input pred", row['linked_pred_id'])
except:
print("Problem with pred")
if not os.path.exists(tmp_input_path):
print("No linked structure found, running ESM")
url = "https://api.esmatlas.com/foldSequence/v1/pdb/"
sequence, chain_lengths = get_sequence_from_pdb(gt_output_path)
if len(sequence) <= 400:
try:
response = requests.post(url, data=sequence)
response.raise_for_status()
pdb_text = response.text
with open(tmp_input_path, "w") as f:
f.write(pdb_text)
# divide to chains
if len(chain_lengths) > 1:
pdb_parser = Bio.PDB.PDBParser(QUIET=True)
pdb_struct = pdb_parser.get_structure("original_pdb", tmp_input_path)
pdb_model = next(iter(pdb_struct))
chain_ids = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"[:len(chain_lengths)]
start_ind = 1
esm_chain = next(pdb_model.get_chains())
new_model = Bio.PDB.Model.Model(0)
for chain_length, chain_id in zip(chain_lengths, chain_ids):
end_ind = start_ind + chain_length
new_chain = Bio.PDB.Chain.Chain(chain_id)
for res in esm_chain.get_residues():
if start_ind <= res.id[1] <= end_ind:
new_chain.add(res)
new_model.add(new_chain)
start_ind = end_ind + 20 # 20 is the gap in esm
io = Bio.PDB.PDBIO()
io.set_structure(new_model)
io.save(tmp_input_path)
input_protein_source = "esm"
print("Using input ESM")
except requests.exceptions.RequestException as e:
print(f"An error occurred in ESM: {e}")
# return "No linked structure found"
else:
print("Sequence too long for ESM")
if not os.path.exists(tmp_input_path):
print("Using input GT")
shutil.copyfile(gt_output_path, tmp_input_path)
input_protein_source = "gt"
align_gt_and_input(gt_output_path, tmp_input_path, gt_output_path, protein_input_path)
protein_size = validate_matching_input_gt(gt_output_path, protein_input_path)
assert protein_size > -1, "Failed to validate matching input and gt"
os.remove(tmp_input_path)
rel_gt_lig_paths = []
rel_ref_lig_paths = []
input_smiles = []
for i, ligand_path in enumerate(sorted(plinder_gt_ligand_paths)):
gt_ligand_output_path = os.path.join(output_models_folder, f"{row['system_id']}_ligand_gt_{i}.sdf")
# rel_gt_lig_paths.append(f"plinder_models/{row['system_id']}_ref_ligand_{i}.sdf")
rel_gt_lig_paths.append(f"plinder_models/{row['system_id']}_ligand_gt_{i}.sdf")
shutil.copyfile(ligand_path, gt_ligand_output_path)
loaded_ligand = Chem.MolFromMolFile(gt_ligand_output_path)
input_smiles.append(Chem.MolToSmiles(loaded_ligand))
ref_ligand_output_path = os.path.join(output_models_folder, f"{row['system_id']}_ligand_ref_{i}.sdf")
rel_ref_lig_paths.append(f"plinder_models/{row['system_id']}_ligand_ref_{i}.sdf")
create_conformers(input_smiles[-1], ref_ligand_output_path, num_conformers=1)
# check if file is empty
if os.path.getsize(ref_ligand_output_path) == 0:
print("Empty ref ligand, copying from gt", ref_ligand_output_path)
shutil.copyfile(gt_ligand_output_path, ref_ligand_output_path)
affinity = row["ligand_binding_affinity"]
if not pd.notna(affinity):
affinity = None
json_data = {
"input_structure": protein_input_relative_path,
"gt_structure": gt_output_relative_path,
"gt_sdf_list": rel_gt_lig_paths,
"input_smiles_list": input_smiles,
"resolution": row.fillna(99)["entry_resolution"],
"release_year": row["entry_release_date"],
"affinity": affinity,
"protein_seq_len": protein_size,
"uniprot": row["system_pocket_UniProt"],
"ligand_num_atoms": ligand_pos.shape[0],
"cluster": row["cluster"],
"cluster_size": row["cluster_size"],
"input_protein_source": input_protein_source,
"ref_sdf_list": rel_ref_lig_paths,
"pdb_id": row["system_id"],
}
open(output_json_path, "w").write(json.dumps(json_data, indent=4))
return "success"
# use linked structures
# input_structure_to_use = None
# apo_linked_structure = os.path.join(linked_structures_folder, "apo", system_id)
# pred_linked_structure = os.path.join(linked_structures_folder, "pred", system_id)
# if os.path.exists(apo_linked_structure):
# for folder in os.listdir(apo_linked_structure):
# if not os.path.isdir(os.path.join(pred_linked_structure, folder)):
# continue
# for filename in os.listdir(os.path.join(apo_linked_structure, folder)):
# if filename.endswith(".cif"):
# input_structure_to_use = os.path.join(apo_linked_structure, folder, filename)
# break
# if input_structure_to_use:
# break
# print(system_id, "found apo", input_structure_to_use)
# elif os.path.exists(pred_linked_structure):
# for folder in os.listdir(pred_linked_structure):
# if not os.path.isdir(os.path.join(pred_linked_structure, folder)):
# continue
# for filename in os.listdir(os.path.join(pred_linked_structure, folder)):
# if filename.endswith(".cif"):
# input_structure_to_use = os.path.join(pred_linked_structure, folder, filename)
# break
# if input_structure_to_use:
# break
# print(system_id, "found pred", input_structure_to_use)
# else:
# print(system_id, "no linked structure found")
# return "No linked structure found"
def main(prefix_bucket_id: str = "*"):
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
systems = get_cached_systems_to_train()
print("total systems", len(systems))
print("clusters", systems["cluster"].value_counts())
# systems = systems[systems["system_num_protein_chains"] > 1]
# return
print("splits", systems["split"].value_counts())
val_or_test = systems[(systems["split"] == "val") | (systems["split"] == "test")]
print("validation or test", len(val_or_test))
output_models_folder = os.path.join(OUTPUT_FOLDER, "plinder_models")
output_train_jsons_folder = os.path.join(OUTPUT_FOLDER, "plinder_jsons_train")
output_val_jsons_folder = os.path.join(OUTPUT_FOLDER, "plinder_jsons_val")
output_test_jsons_folder = os.path.join(OUTPUT_FOLDER, "plinder_jsons_test")
output_info = os.path.join(OUTPUT_FOLDER, "plinder_generation_info.csv")
if prefix_bucket_id != "*":
output_info = os.path.join(OUTPUT_FOLDER, f"plinder_generation_info_{prefix_bucket_id}.csv")
os.makedirs(output_models_folder, exist_ok=True)
os.makedirs(output_train_jsons_folder, exist_ok=True)
os.makedirs(output_val_jsons_folder, exist_ok=True)
os.makedirs(output_test_jsons_folder, exist_ok=True)
split_to_folder = {
"train": output_train_jsons_folder,
"val": output_val_jsons_folder,
"test": output_test_jsons_folder
}
output_info_file = open(output_info, "a+")
for bucket_id, bucket_systems in systems.groupby('_bucket_id', sort=True):
if prefix_bucket_id != "*" and not str(bucket_id).startswith(prefix_bucket_id):
continue
# if bucket_id != "z2":
# continue
# systems_folder = "{BASE_FOLDER}/processed/tmp_z2/systems"
print("Starting bucket", bucket_id, len(bucket_systems))
print(len(bucket_systems), bucket_systems["system_num_ligand_chains"].value_counts())
tmp_output_models_folder = os.path.join(OUTPUT_FOLDER, f"tmp_{bucket_id}")
os.makedirs(tmp_output_models_folder, exist_ok=True)
os.system(f'{GSUTIL_PATH} -m cp -r "gs://plinder/2024-06/v2/systems/{bucket_id}.zip" {tmp_output_models_folder}')
systems_folder = os.path.join(tmp_output_models_folder, "systems")
os.system(f'unzip -o {os.path.join(tmp_output_models_folder, f"{bucket_id}.zip")} -d {systems_folder}')
for i, row in bucket_systems.iterrows():
# if not str(row['system_id']).startswith("4z22__1__1.A__1.C"):
# continue
print("doing", row['system_id'], row["system_num_protein_chains"], row["system_num_ligand_chains"])
system_folder = os.path.join(systems_folder, row['system_id'])
try:
success = prepare_system(row, system_folder, output_models_folder, split_to_folder[row["split"]])
print("done", row['system_id'], success)
output_info_file.write(f"{bucket_id},{row['system_id']},{success}\n")
except Exception as e:
print("Failed", row['system_id'], e)
output_info_file.write(f"{bucket_id},{row['system_id']},Failed\n")
output_info_file.flush()
shutil.rmtree(tmp_output_models_folder)
if __name__ == '__main__':
prefix_bucket_id = "*"
if len(sys.argv) > 1:
prefix_bucket_id = sys.argv[1]
main(prefix_bucket_id)