File size: 1,067 Bytes
1f0c7b9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import glob
from sample_for_pocket import call, OUT_DIR
from rdkit import Chem
import json
def load(value: str):
full_pdb_path = '../MolCRAFT/data/test_set/' + value
sdf_path = glob.glob(full_pdb_path.rstrip('.pdb') + '*.sdf')
assert len(sdf_path) == 1
sdf_path = sdf_path[0]
return [full_pdb_path, sdf_path]
pdb_file = 'ABL2_HUMAN_274_551_0/4xli_B_rec.pdb'
protein_fn, sdf_fn = load(pdb_file)
print(protein_fn, sdf_fn)
ref_mol = Chem.SDMolSupplier(sdf_fn, removeHs=False)[0]
ref_lig_pos = ref_mol.GetConformer().GetPositions()
print('[DEBUG] read raw file')
print(ref_lig_pos.shape, ref_lig_pos.mean(axis=0))
def generate(value: str):
protein_path, ligand_path = load(value)
call(protein_path, ligand_path)
out_fns = sorted(glob.glob(f'{OUT_DIR}/*.sdf'))
print(out_fns)
out_fn = out_fns[0]
sdf_mol = Chem.SDMolSupplier(out_fn, removeHs=False)[0]
# get all properties from sdf_mol
props = sdf_mol.GetPropsAsDict()
return [protein_path, out_fn], json.dumps(props, indent=4)
generate(pdb_file)
|