import argparse from pathlib import Path import numpy as np import torch import torch.nn.functional as F from Bio.PDB import PDBParser from rdkit import Chem from torch_scatter import scatter_mean from openbabel import openbabel openbabel.obErrorLog.StopLogging() # suppress OpenBabel messages import utils from lightning_modules import LigandPocketDDPM from constants import FLOAT_TYPE, INT_TYPE from analysis.molecule_builder import build_molecule, process_molecule def prepare_from_sdf_files(sdf_files, atom_encoder): ligand_coords = [] atom_one_hot = [] for file in sdf_files: rdmol = Chem.SDMolSupplier(str(file), sanitize=False)[0] ligand_coords.append( torch.from_numpy(rdmol.GetConformer().GetPositions()).float() ) types = torch.tensor([atom_encoder[a.GetSymbol()] for a in rdmol.GetAtoms()]) atom_one_hot.append( F.one_hot(types, num_classes=len(atom_encoder)) ) return torch.cat(ligand_coords, dim=0), torch.cat(atom_one_hot, dim=0) def prepare_ligand_from_pdb(biopython_atoms, atom_encoder): coord = torch.tensor(np.array([a.get_coord() for a in biopython_atoms]), dtype=FLOAT_TYPE) types = torch.tensor([atom_encoder[a.element.capitalize()] for a in biopython_atoms]) one_hot = F.one_hot(types, num_classes=len(atom_encoder)) return coord, one_hot def prepare_substructure(ref_ligand, fix_atoms, pdb_model): if fix_atoms[0].endswith(".sdf"): # ligand as sdf file coord, one_hot = prepare_from_sdf_files(fix_atoms, model.lig_type_encoder) else: # ligand contained in PDB; given in : format chain, resi = ref_ligand.split(':') ligand = utils.get_residue_with_resi(pdb_model[chain], int(resi)) fixed_atoms = [a for a in ligand.get_atoms() if a.get_name() in set(fix_atoms)] coord, one_hot = prepare_ligand_from_pdb(fixed_atoms, model.lig_type_encoder) return coord, one_hot def inpaint_ligand(model, pdb_file, n_samples, ligand, fix_atoms, add_n_nodes=None, center='ligand', sanitize=False, largest_frag=False, relax_iter=0, timesteps=None, resamplings=1, save_traj=False): """ Generate ligands given a pocket Args: model: Lightning model pdb_file: PDB filename n_samples: number of samples ligand: reference ligand given in : format if the ligand is contained in the PDB file, or path to an SDF file that contains the ligand; used to define the pocket fix_atoms: ligand atoms that should be fixed, e.g. "C1 N6 C5 C12" center: 'ligand' or 'pocket' add_n_nodes: number of ligand nodes to add, sampled randomly if 'None' sanitize: whether to sanitize molecules or not largest_frag: only return the largest fragment relax_iter: number of force field optimization steps timesteps: number of denoising steps, use training value if None resamplings: number of resampling iterations save_traj: save intermediate states to visualize a denoising trajectory Returns: list of molecules """ if save_traj and n_samples > 1: raise NotImplementedError("Can only visualize trajectory with " "n_samples=1.") frames = timesteps if save_traj else 1 sanitize = False if save_traj else sanitize relax_iter = 0 if save_traj else relax_iter largest_frag = False if save_traj else largest_frag # Load PDB pdb_model = PDBParser(QUIET=True).get_structure('', pdb_file)[0] # Define pocket based on reference ligand residues = utils.get_pocket_from_ligand(pdb_model, ligand) pocket = model.prepare_pocket(residues, repeats=n_samples) # Get fixed ligand substructure x_fixed, one_hot_fixed = prepare_substructure(ligand, fix_atoms, pdb_model) n_fixed = len(x_fixed) if add_n_nodes is None: num_nodes_lig = model.ddpm.size_distribution.sample_conditional( n1=None, n2=pocket['size']) num_nodes_lig = torch.clamp(num_nodes_lig, min=n_fixed) else: num_nodes_lig = torch.ones(n_samples, dtype=int) * n_fixed + add_n_nodes ligand_mask = utils.num_nodes_to_batch_mask( len(num_nodes_lig), num_nodes_lig, model.device) ligand = { 'x': torch.zeros((len(ligand_mask), model.x_dims), device=model.device, dtype=FLOAT_TYPE), 'one_hot': torch.zeros((len(ligand_mask), model.atom_nf), device=model.device, dtype=FLOAT_TYPE), 'size': num_nodes_lig, 'mask': ligand_mask } # fill in fixed atoms lig_fixed = torch.zeros_like(ligand_mask) for i in range(n_samples): sele = (ligand_mask == i) x_new = ligand['x'][sele] x_new[:n_fixed] = x_fixed ligand['x'][sele] = x_new h_new = ligand['one_hot'][sele] h_new[:n_fixed] = one_hot_fixed ligand['one_hot'][sele] = h_new fixed_new = lig_fixed[sele] fixed_new[:n_fixed] = 1 lig_fixed[sele] = fixed_new # Pocket's center of mass pocket_com_before = scatter_mean(pocket['x'], pocket['mask'], dim=0) # Run sampling xh_lig, xh_pocket, lig_mask, pocket_mask = model.ddpm.inpaint( ligand, pocket, lig_fixed, center=center, resamplings=resamplings, timesteps=timesteps, return_frames=frames) # Treat intermediate states as molecules for downstream processing if save_traj: xh_lig = utils.reverse_tensor(xh_lig) xh_pocket = utils.reverse_tensor(xh_pocket) lig_mask = torch.arange(xh_lig.size(0), device=model.device ).repeat_interleave(len(lig_mask)) pocket_mask = torch.arange(xh_pocket.size(0), device=model.device ).repeat_interleave(len(pocket_mask)) xh_lig = xh_lig.view(-1, xh_lig.size(2)) xh_pocket = xh_pocket.view(-1, xh_pocket.size(2)) # Move generated molecule back to the original pocket position pocket_com_after = scatter_mean(xh_pocket[:, :model.x_dims], pocket_mask, dim=0) xh_pocket[:, :model.x_dims] += \ (pocket_com_before - pocket_com_after)[pocket_mask] xh_lig[:, :model.x_dims] += \ (pocket_com_before - pocket_com_after)[lig_mask] # Build mol objects x = xh_lig[:, :model.x_dims].detach().cpu() atom_type = xh_lig[:, model.x_dims:].argmax(1).detach().cpu() molecules = [] for mol_pc in zip(utils.batch_to_list(x, lig_mask), utils.batch_to_list(atom_type, lig_mask)): mol = build_molecule(*mol_pc, model.dataset_info, add_coords=True) mol = process_molecule(mol, add_hydrogens=False, sanitize=sanitize, relax_iter=relax_iter, largest_frag=largest_frag) if mol is not None: molecules.append(mol) return molecules if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('checkpoint', type=Path) parser.add_argument('--pdbfile', type=str) parser.add_argument('--ref_ligand', type=str, default=None) parser.add_argument('--fix_atoms', type=str, nargs='+', default=None) parser.add_argument('--center', type=str, default='ligand', choices={'ligand', 'pocket'}) parser.add_argument('--outfile', type=Path) parser.add_argument('--n_samples', type=int, default=20) parser.add_argument('--add_n_nodes', type=int, default=None) parser.add_argument('--relax', action='store_true') parser.add_argument('--sanitize', action='store_true') parser.add_argument('--resamplings', type=int, default=20) parser.add_argument('--timesteps', type=int, default=50) parser.add_argument('--save_traj', action='store_true') args = parser.parse_args() pdb_id = Path(args.pdbfile).stem device = 'cuda' if torch.cuda.is_available() else 'cpu' # Load model model = LigandPocketDDPM.load_from_checkpoint( args.checkpoint, map_location=device) model = model.to(device) molecules = inpaint_ligand(model, args.pdbfile, args.n_samples, args.ref_ligand, args.fix_atoms, args.add_n_nodes, center=args.center, sanitize=args.sanitize, largest_frag=False, relax_iter=(200 if args.relax else 0), timesteps=args.timesteps, resamplings=args.resamplings, save_traj=args.save_traj) # Make SDF files utils.write_sdf_file(args.outfile, molecules)