data/prepare_crossdocked.py

import sys
import torch
import shutil
from pathlib import Path

from rdkit import Chem
from tqdm import tqdm


basedir = sys.argv[1]
structure_dir = Path(basedir, 'crossdocked_pocket10')

test_set = torch.load(Path(basedir, 'split_by_name.pt'))['test']

receptor_dir = Path(basedir, 'receptor_pdbs')
receptor_dir.mkdir(exist_ok=True)

ref_ligand_dir = Path(basedir, 'reference_ligands')
ref_ligand_dir.mkdir(exist_ok=True)

methods = ['cvae', 'sbdd', 'p2m']
for method in methods:
    method_lig_dir = Path(basedir, f'{method}_processed')
    method_lig_dir.mkdir(exist_ok=True)

for pocket_idx, (receptor_name, ligand_name) in enumerate(tqdm(test_set)):

    # copy receptor file and remove underscores
    new_rec_name = Path(receptor_name).stem.replace('_', '-')
    shutil.copy(Path(structure_dir, receptor_name), Path(receptor_dir, new_rec_name + '.pdb'))

    # copy and rename reference ligands
    new_lig_name = new_rec_name + '_' + Path(ligand_name).stem.replace('_', '-')
    shutil.copy(Path(structure_dir, ligand_name), Path(ref_ligand_dir, new_lig_name + '.sdf'))

    for method in methods:

        method_pocket_dir = Path(basedir, method, f'pocket_{pocket_idx}')

        generated_mols = [Chem.SDMolSupplier(str(file), sanitize=False)[0]
                          for file in method_pocket_dir.glob(f'mol_*.sdf')]

        # only select first 100 molecules
        generated_mols = generated_mols[:100]
        if len(generated_mols) < 1:
            print('No molecule found for this pocket')
            continue
        if len(generated_mols) < 100:
            print('Less than 100 molecules found for this pocket')

        # write a combined sdf file
        sdf_path = Path(basedir, f'{method}_processed', f'{new_rec_name}_mols-pocket-{pocket_idx}.sdf')
        with Chem.SDWriter(str(sdf_path)) as w:
            for mol in generated_mols:
                w.write(mol)