import math from argparse import Namespace from typing import Optional from time import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F from torch.utils.data import DataLoader import pytorch_lightning as pl import wandb from torch_scatter import scatter_add, scatter_mean from Bio.PDB import PDBParser from Bio.PDB.Polypeptide import three_to_one from constants import dataset_params, FLOAT_TYPE, INT_TYPE from equivariant_diffusion.dynamics import EGNNDynamics from equivariant_diffusion.en_diffusion import EnVariationalDiffusion from equivariant_diffusion.conditional_model import ConditionalDDPM, \ SimpleConditionalDDPM from dataset import ProcessedLigandPocketDataset import utils from analysis.visualization import save_xyz_file, visualize, visualize_chain from analysis.metrics import BasicMolecularMetrics, CategoricalDistribution, \ MoleculeProperties from analysis.molecule_builder import build_molecule, process_molecule from analysis.docking import smina_score class LigandPocketDDPM(pl.LightningModule): def __init__( self, outdir, dataset, datadir, batch_size, lr, egnn_params: Namespace, diffusion_params, num_workers, augment_noise, augment_rotation, clip_grad, eval_epochs, eval_params, visualize_sample_epoch, visualize_chain_epoch, auxiliary_loss, loss_params, mode, node_histogram, pocket_representation='CA', virtual_nodes=False ): super(LigandPocketDDPM, self).__init__() self.save_hyperparameters() ddpm_models = {'joint': EnVariationalDiffusion, 'pocket_conditioning': ConditionalDDPM, 'pocket_conditioning_simple': SimpleConditionalDDPM} assert mode in ddpm_models self.mode = mode assert pocket_representation in {'CA', 'full-atom'} self.pocket_representation = pocket_representation self.dataset_name = dataset self.datadir = datadir self.outdir = outdir self.batch_size = batch_size self.eval_batch_size = eval_params.eval_batch_size \ if 'eval_batch_size' in eval_params else batch_size self.lr = lr self.loss_type = diffusion_params.diffusion_loss_type self.eval_epochs = eval_epochs self.visualize_sample_epoch = visualize_sample_epoch self.visualize_chain_epoch = visualize_chain_epoch self.eval_params = eval_params self.num_workers = num_workers self.augment_noise = augment_noise self.augment_rotation = augment_rotation self.dataset_info = dataset_params[dataset] self.T = diffusion_params.diffusion_steps self.clip_grad = clip_grad if clip_grad: self.gradnorm_queue = utils.Queue() # Add large value that will be flushed. self.gradnorm_queue.add(3000) self.lig_type_encoder = self.dataset_info['atom_encoder'] self.lig_type_decoder = self.dataset_info['atom_decoder'] self.pocket_type_encoder = self.dataset_info['aa_encoder'] \ if self.pocket_representation == 'CA' \ else self.dataset_info['atom_encoder'] self.pocket_type_decoder = self.dataset_info['aa_decoder'] \ if self.pocket_representation == 'CA' \ else self.dataset_info['atom_decoder'] smiles_list = None if eval_params.smiles_file is None \ else np.load(eval_params.smiles_file) self.ligand_metrics = BasicMolecularMetrics(self.dataset_info, smiles_list) self.molecule_properties = MoleculeProperties() self.ligand_type_distribution = CategoricalDistribution( self.dataset_info['atom_hist'], self.lig_type_encoder) if self.pocket_representation == 'CA': self.pocket_type_distribution = CategoricalDistribution( self.dataset_info['aa_hist'], self.pocket_type_encoder) else: self.pocket_type_distribution = None self.train_dataset = None self.val_dataset = None self.test_dataset = None self.virtual_nodes = virtual_nodes self.data_transform = None self.max_num_nodes = len(node_histogram) - 1 if virtual_nodes: # symbol = 'virtual' symbol = 'Ne' # visualize as Neon atoms self.lig_type_encoder[symbol] = len(self.lig_type_encoder) self.virtual_atom = self.lig_type_encoder[symbol] self.lig_type_decoder.append(symbol) self.data_transform = utils.AppendVirtualNodes( self.max_num_nodes, self.lig_type_encoder, symbol) # Update dataset_info dictionary. This is necessary for using the # visualization functions. self.dataset_info['atom_encoder'] = self.lig_type_encoder self.dataset_info['atom_decoder'] = self.lig_type_decoder self.atom_nf = len(self.lig_type_decoder) self.aa_nf = len(self.pocket_type_decoder) self.x_dims = 3 net_dynamics = EGNNDynamics( atom_nf=self.atom_nf, residue_nf=self.aa_nf, n_dims=self.x_dims, joint_nf=egnn_params.joint_nf, device=egnn_params.device if torch.cuda.is_available() else 'cpu', hidden_nf=egnn_params.hidden_nf, act_fn=torch.nn.SiLU(), n_layers=egnn_params.n_layers, attention=egnn_params.attention, tanh=egnn_params.tanh, norm_constant=egnn_params.norm_constant, inv_sublayers=egnn_params.inv_sublayers, sin_embedding=egnn_params.sin_embedding, normalization_factor=egnn_params.normalization_factor, aggregation_method=egnn_params.aggregation_method, edge_cutoff_ligand=egnn_params.__dict__.get('edge_cutoff_ligand'), edge_cutoff_pocket=egnn_params.__dict__.get('edge_cutoff_pocket'), edge_cutoff_interaction=egnn_params.__dict__.get('edge_cutoff_interaction'), update_pocket_coords=(self.mode == 'joint'), reflection_equivariant=egnn_params.reflection_equivariant, edge_embedding_dim=egnn_params.__dict__.get('edge_embedding_dim'), ) self.ddpm = ddpm_models[self.mode]( dynamics=net_dynamics, atom_nf=self.atom_nf, residue_nf=self.aa_nf, n_dims=self.x_dims, timesteps=diffusion_params.diffusion_steps, noise_schedule=diffusion_params.diffusion_noise_schedule, noise_precision=diffusion_params.diffusion_noise_precision, loss_type=diffusion_params.diffusion_loss_type, norm_values=diffusion_params.normalize_factors, size_histogram=node_histogram, virtual_node_idx=self.lig_type_encoder[symbol] if virtual_nodes else None ) self.auxiliary_loss = auxiliary_loss self.lj_rm = self.dataset_info['lennard_jones_rm'] if self.auxiliary_loss: self.clamp_lj = loss_params.clamp_lj self.auxiliary_weight_schedule = WeightSchedule( T=diffusion_params.diffusion_steps, max_weight=loss_params.max_weight, mode=loss_params.schedule) def configure_optimizers(self): return torch.optim.AdamW(self.ddpm.parameters(), lr=self.lr, amsgrad=True, weight_decay=1e-12) def setup(self, stage: Optional[str] = None): if stage == 'fit': self.train_dataset = ProcessedLigandPocketDataset( Path(self.datadir, 'train.npz'), transform=self.data_transform) self.val_dataset = ProcessedLigandPocketDataset( Path(self.datadir, 'val.npz'), transform=self.data_transform) elif stage == 'test': self.test_dataset = ProcessedLigandPocketDataset( Path(self.datadir, 'test.npz'), transform=self.data_transform) else: raise NotImplementedError def train_dataloader(self): return DataLoader(self.train_dataset, self.batch_size, shuffle=True, num_workers=self.num_workers, collate_fn=self.train_dataset.collate_fn, pin_memory=True) def val_dataloader(self): return DataLoader(self.val_dataset, self.batch_size, shuffle=False, num_workers=self.num_workers, collate_fn=self.val_dataset.collate_fn, pin_memory=True) def test_dataloader(self): return DataLoader(self.test_dataset, self.batch_size, shuffle=False, num_workers=self.num_workers, collate_fn=self.test_dataset.collate_fn, pin_memory=True) def get_ligand_and_pocket(self, data): ligand = { 'x': data['lig_coords'].to(self.device, FLOAT_TYPE), 'one_hot': data['lig_one_hot'].to(self.device, FLOAT_TYPE), 'size': data['num_lig_atoms'].to(self.device, INT_TYPE), 'mask': data['lig_mask'].to(self.device, INT_TYPE), } if self.virtual_nodes: ligand['num_virtual_atoms'] = data['num_virtual_atoms'].to( self.device, INT_TYPE) pocket = { 'x': data['pocket_coords'].to(self.device, FLOAT_TYPE), 'one_hot': data['pocket_one_hot'].to(self.device, FLOAT_TYPE), 'size': data['num_pocket_nodes'].to(self.device, INT_TYPE), 'mask': data['pocket_mask'].to(self.device, INT_TYPE) } return ligand, pocket def forward(self, data): ligand, pocket = self.get_ligand_and_pocket(data) # Note: \mathcal{L} terms in the paper represent log-likelihoods while # our loss terms are a negative(!) log-likelihoods delta_log_px, error_t_lig, error_t_pocket, SNR_weight, \ loss_0_x_ligand, loss_0_x_pocket, loss_0_h, neg_log_const_0, \ kl_prior, log_pN, t_int, xh_lig_hat, info = \ self.ddpm(ligand, pocket, return_info=True) if self.loss_type == 'l2' and self.training: actual_ligand_size = ligand['size'] - ligand['num_virtual_atoms'] if self.virtual_nodes else ligand['size'] # normalize loss_t denom_lig = self.x_dims * actual_ligand_size + \ self.ddpm.atom_nf * ligand['size'] error_t_lig = error_t_lig / denom_lig denom_pocket = (self.x_dims + self.ddpm.residue_nf) * pocket['size'] error_t_pocket = error_t_pocket / denom_pocket loss_t = 0.5 * (error_t_lig + error_t_pocket) # normalize loss_0 loss_0_x_ligand = loss_0_x_ligand / (self.x_dims * actual_ligand_size) loss_0_x_pocket = loss_0_x_pocket / (self.x_dims * pocket['size']) loss_0 = loss_0_x_ligand + loss_0_x_pocket + loss_0_h # VLB objective or evaluation step else: # Note: SNR_weight should be negative loss_t = -self.T * 0.5 * SNR_weight * (error_t_lig + error_t_pocket) loss_0 = loss_0_x_ligand + loss_0_x_pocket + loss_0_h loss_0 = loss_0 + neg_log_const_0 nll = loss_t + loss_0 + kl_prior # Correct for normalization on x. if not (self.loss_type == 'l2' and self.training): nll = nll - delta_log_px # always the same number of nodes if virtual nodes are added if not self.virtual_nodes: # Transform conditional nll into joint nll # Note: # loss = -log p(x,h|N) and log p(x,h,N) = log p(x,h|N) + log p(N) # Therefore, log p(x,h|N) = -loss + log p(N) # => loss_new = -log p(x,h,N) = loss - log p(N) nll = nll - log_pN # Add auxiliary loss term if self.auxiliary_loss and self.loss_type == 'l2' and self.training: x_lig_hat = xh_lig_hat[:, :self.x_dims] h_lig_hat = xh_lig_hat[:, self.x_dims:] weighted_lj_potential = \ self.auxiliary_weight_schedule(t_int.long()) * \ self.lj_potential(x_lig_hat, h_lig_hat, ligand['mask']) nll = nll + weighted_lj_potential info['weighted_lj'] = weighted_lj_potential.mean(0) info['error_t_lig'] = error_t_lig.mean(0) info['error_t_pocket'] = error_t_pocket.mean(0) info['SNR_weight'] = SNR_weight.mean(0) info['loss_0'] = loss_0.mean(0) info['kl_prior'] = kl_prior.mean(0) info['delta_log_px'] = delta_log_px.mean(0) info['neg_log_const_0'] = neg_log_const_0.mean(0) info['log_pN'] = log_pN.mean(0) return nll, info def lj_potential(self, atom_x, atom_one_hot, batch_mask): adj = batch_mask[:, None] == batch_mask[None, :] adj = adj ^ torch.diag(torch.diag(adj)) # remove self-edges edges = torch.where(adj) # Compute pair-wise potentials r = torch.sum((atom_x[edges[0]] - atom_x[edges[1]])**2, dim=1).sqrt() # Get optimal radii lennard_jones_radii = torch.tensor(self.lj_rm, device=r.device) # unit conversion pm -> A lennard_jones_radii = lennard_jones_radii / 100.0 # normalization lennard_jones_radii = lennard_jones_radii / self.ddpm.norm_values[0] atom_type_idx = atom_one_hot.argmax(1) rm = lennard_jones_radii[atom_type_idx[edges[0]], atom_type_idx[edges[1]]] sigma = 2 ** (-1 / 6) * rm out = 4 * ((sigma / r) ** 12 - (sigma / r) ** 6) if self.clamp_lj is not None: out = torch.clamp(out, min=None, max=self.clamp_lj) # Compute potential per atom out = scatter_add(out, edges[0], dim=0, dim_size=len(atom_x)) # Sum potentials of all atoms return scatter_add(out, batch_mask, dim=0) def log_metrics(self, metrics_dict, split, batch_size=None, **kwargs): for m, value in metrics_dict.items(): self.log(f'{m}/{split}', value, batch_size=batch_size, **kwargs) def training_step(self, data, *args): if self.augment_noise > 0: raise NotImplementedError # Add noise eps ~ N(0, augment_noise) around points. eps = sample_center_gravity_zero_gaussian(x.size(), x.device) x = x + eps * args.augment_noise if self.augment_rotation: raise NotImplementedError x = utils.random_rotation(x).detach() try: nll, info = self.forward(data) except RuntimeError as e: # this is not supported for multi-GPU if self.trainer.num_devices < 2 and 'out of memory' in str(e): print('WARNING: ran out of memory, skipping to the next batch') return None else: raise e loss = nll.mean(0) info['loss'] = loss self.log_metrics(info, 'train', batch_size=len(data['num_lig_atoms'])) return info def _shared_eval(self, data, prefix, *args): nll, info = self.forward(data) loss = nll.mean(0) info['loss'] = loss self.log_metrics(info, prefix, batch_size=len(data['num_lig_atoms']), sync_dist=True) return info def validation_step(self, data, *args): self._shared_eval(data, 'val', *args) def test_step(self, data, *args): self._shared_eval(data, 'test', *args) def validation_epoch_end(self, validation_step_outputs): # Perform validation on single GPU if not self.trainer.is_global_zero: return suffix = '' if self.mode == 'joint' else '_given_pocket' if (self.current_epoch + 1) % self.eval_epochs == 0: tic = time() sampling_results = getattr(self, 'sample_and_analyze' + suffix)( self.eval_params.n_eval_samples, self.val_dataset, batch_size=self.eval_batch_size) self.log_metrics(sampling_results, 'val') print(f'Evaluation took {time() - tic:.2f} seconds') if (self.current_epoch + 1) % self.visualize_sample_epoch == 0: tic = time() getattr(self, 'sample_and_save' + suffix)( self.eval_params.n_visualize_samples) print(f'Sample visualization took {time() - tic:.2f} seconds') if (self.current_epoch + 1) % self.visualize_chain_epoch == 0: tic = time() getattr(self, 'sample_chain_and_save' + suffix)( self.eval_params.keep_frames) print(f'Chain visualization took {time() - tic:.2f} seconds') @torch.no_grad() def sample_and_analyze(self, n_samples, dataset=None, batch_size=None): print(f'Analyzing sampled molecules at epoch {self.current_epoch}...') batch_size = self.batch_size if batch_size is None else batch_size batch_size = min(batch_size, n_samples) # each item in molecules is a tuple (position, atom_type_encoded) molecules = [] atom_types = [] aa_types = [] for i in range(math.ceil(n_samples / batch_size)): n_samples_batch = min(batch_size, n_samples - len(molecules)) num_nodes_lig, num_nodes_pocket = \ self.ddpm.size_distribution.sample(n_samples_batch) xh_lig, xh_pocket, lig_mask, _ = self.ddpm.sample( n_samples_batch, num_nodes_lig, num_nodes_pocket, device=self.device) x = xh_lig[:, :self.x_dims].detach().cpu() atom_type = xh_lig[:, self.x_dims:].argmax(1).detach().cpu() lig_mask = lig_mask.cpu() molecules.extend(list( zip(utils.batch_to_list(x, lig_mask), utils.batch_to_list(atom_type, lig_mask)) )) atom_types.extend(atom_type.tolist()) aa_types.extend( xh_pocket[:, self.x_dims:].argmax(1).detach().cpu().tolist()) return self.analyze_sample(molecules, atom_types, aa_types) def analyze_sample(self, molecules, atom_types, aa_types, receptors=None): # Distribution of node types kl_div_atom = self.ligand_type_distribution.kl_divergence(atom_types) \ if self.ligand_type_distribution is not None else -1 kl_div_aa = self.pocket_type_distribution.kl_divergence(aa_types) \ if self.pocket_type_distribution is not None else -1 # Convert into rdmols rdmols = [build_molecule(*graph, self.dataset_info) for graph in molecules] # Other basic metrics (validity, connectivity, uniqueness, novelty), (_, connected_mols) = \ self.ligand_metrics.evaluate_rdmols(rdmols) qed, sa, logp, lipinski, diversity = \ self.molecule_properties.evaluate_mean(connected_mols) out = { 'kl_div_atom_types': kl_div_atom, 'kl_div_residue_types': kl_div_aa, 'Validity': validity, 'Connectivity': connectivity, 'Uniqueness': uniqueness, 'Novelty': novelty, 'QED': qed, 'SA': sa, 'LogP': logp, 'Lipinski': lipinski, 'Diversity': diversity } # Simple docking score if receptors is not None: # out['smina_score'] = np.mean(smina_score(rdmols, receptors)) out['smina_score'] = np.mean(smina_score(connected_mols, receptors)) return out def get_full_path(self, receptor_name): pdb, suffix = receptor_name.split('.') receptor_name = f'{pdb.upper()}-{suffix}.pdb' return Path(self.datadir, 'val', receptor_name) @torch.no_grad() def sample_and_analyze_given_pocket(self, n_samples, dataset=None, batch_size=None): print(f'Analyzing sampled molecules given pockets at epoch ' f'{self.current_epoch}...') batch_size = self.batch_size if batch_size is None else batch_size batch_size = min(batch_size, n_samples) # each item in molecules is a tuple (position, atom_type_encoded) molecules = [] atom_types = [] aa_types = [] receptors = [] for i in range(math.ceil(n_samples / batch_size)): n_samples_batch = min(batch_size, n_samples - len(molecules)) # Create a batch batch = dataset.collate_fn( [dataset[(i * batch_size + j) % len(dataset)] for j in range(n_samples_batch)] ) ligand, pocket = self.get_ligand_and_pocket(batch) receptors.extend([self.get_full_path(x) for x in batch['receptors']]) if self.virtual_nodes: num_nodes_lig = self.max_num_nodes else: num_nodes_lig = self.ddpm.size_distribution.sample_conditional( n1=None, n2=pocket['size']) xh_lig, xh_pocket, lig_mask, _ = self.ddpm.sample_given_pocket( pocket, num_nodes_lig) x = xh_lig[:, :self.x_dims].detach().cpu() atom_type = xh_lig[:, self.x_dims:].argmax(1).detach().cpu() lig_mask = lig_mask.cpu() if self.virtual_nodes: # Remove virtual nodes for analysis vnode_mask = (atom_type == self.virtual_atom) x = x[~vnode_mask, :] atom_type = atom_type[~vnode_mask] lig_mask = lig_mask[~vnode_mask] molecules.extend(list( zip(utils.batch_to_list(x, lig_mask), utils.batch_to_list(atom_type, lig_mask)) )) atom_types.extend(atom_type.tolist()) aa_types.extend( xh_pocket[:, self.x_dims:].argmax(1).detach().cpu().tolist()) return self.analyze_sample(molecules, atom_types, aa_types, receptors=receptors) def sample_and_save(self, n_samples): num_nodes_lig, num_nodes_pocket = \ self.ddpm.size_distribution.sample(n_samples) xh_lig, xh_pocket, lig_mask, pocket_mask = \ self.ddpm.sample(n_samples, num_nodes_lig, num_nodes_pocket, device=self.device) if self.pocket_representation == 'CA': # convert residues into atom representation for visualization x_pocket, one_hot_pocket = utils.residues_to_atoms( xh_pocket[:, :self.x_dims], self.lig_type_encoder) else: x_pocket, one_hot_pocket = \ xh_pocket[:, :self.x_dims], xh_pocket[:, self.x_dims:] x = torch.cat((xh_lig[:, :self.x_dims], x_pocket), dim=0) one_hot = torch.cat((xh_lig[:, self.x_dims:], one_hot_pocket), dim=0) outdir = Path(self.outdir, f'epoch_{self.current_epoch}') save_xyz_file(str(outdir) + '/', one_hot, x, self.lig_type_decoder, name='molecule', batch_mask=torch.cat((lig_mask, pocket_mask))) # visualize(str(outdir), dataset_info=self.dataset_info, wandb=wandb) visualize(str(outdir), dataset_info=self.dataset_info, wandb=None) def sample_and_save_given_pocket(self, n_samples): batch = self.val_dataset.collate_fn( [self.val_dataset[i] for i in torch.randint(len(self.val_dataset), size=(n_samples,))] ) ligand, pocket = self.get_ligand_and_pocket(batch) if self.virtual_nodes: num_nodes_lig = self.max_num_nodes else: num_nodes_lig = self.ddpm.size_distribution.sample_conditional( n1=None, n2=pocket['size']) xh_lig, xh_pocket, lig_mask, pocket_mask = \ self.ddpm.sample_given_pocket(pocket, num_nodes_lig) if self.pocket_representation == 'CA': # convert residues into atom representation for visualization x_pocket, one_hot_pocket = utils.residues_to_atoms( xh_pocket[:, :self.x_dims], self.lig_type_encoder) else: x_pocket, one_hot_pocket = \ xh_pocket[:, :self.x_dims], xh_pocket[:, self.x_dims:] x = torch.cat((xh_lig[:, :self.x_dims], x_pocket), dim=0) one_hot = torch.cat((xh_lig[:, self.x_dims:], one_hot_pocket), dim=0) outdir = Path(self.outdir, f'epoch_{self.current_epoch}') save_xyz_file(str(outdir) + '/', one_hot, x, self.lig_type_decoder, name='molecule', batch_mask=torch.cat((lig_mask, pocket_mask))) # visualize(str(outdir), dataset_info=self.dataset_info, wandb=wandb) visualize(str(outdir), dataset_info=self.dataset_info, wandb=None) def sample_chain_and_save(self, keep_frames): n_samples = 1 num_nodes_lig, num_nodes_pocket = \ self.ddpm.size_distribution.sample(n_samples) chain_lig, chain_pocket, _, _ = self.ddpm.sample( n_samples, num_nodes_lig, num_nodes_pocket, return_frames=keep_frames, device=self.device) chain_lig = utils.reverse_tensor(chain_lig) chain_pocket = utils.reverse_tensor(chain_pocket) # Repeat last frame to see final sample better. chain_lig = torch.cat([chain_lig, chain_lig[-1:].repeat(10, 1, 1)], dim=0) chain_pocket = torch.cat( [chain_pocket, chain_pocket[-1:].repeat(10, 1, 1)], dim=0) # Prepare entire chain. x_lig = chain_lig[:, :, :self.x_dims] one_hot_lig = chain_lig[:, :, self.x_dims:] one_hot_lig = F.one_hot( torch.argmax(one_hot_lig, dim=2), num_classes=len(self.lig_type_decoder)) x_pocket = chain_pocket[:, :, :self.x_dims] one_hot_pocket = chain_pocket[:, :, self.x_dims:] one_hot_pocket = F.one_hot( torch.argmax(one_hot_pocket, dim=2), num_classes=len(self.pocket_type_decoder)) if self.pocket_representation == 'CA': # convert residues into atom representation for visualization x_pocket, one_hot_pocket = utils.residues_to_atoms( x_pocket, self.lig_type_encoder) x = torch.cat((x_lig, x_pocket), dim=1) one_hot = torch.cat((one_hot_lig, one_hot_pocket), dim=1) # flatten (treat frame (chain dimension) as batch for visualization) x_flat = x.view(-1, x.size(-1)) one_hot_flat = one_hot.view(-1, one_hot.size(-1)) mask_flat = torch.arange(x.size(0)).repeat_interleave(x.size(1)) outdir = Path(self.outdir, f'epoch_{self.current_epoch}', 'chain') save_xyz_file(str(outdir), one_hot_flat, x_flat, self.lig_type_decoder, name='/chain', batch_mask=mask_flat) visualize_chain(str(outdir), self.dataset_info, wandb=wandb) def sample_chain_and_save_given_pocket(self, keep_frames): n_samples = 1 batch = self.val_dataset.collate_fn([ self.val_dataset[torch.randint(len(self.val_dataset), size=(1,))] ]) ligand, pocket = self.get_ligand_and_pocket(batch) if self.virtual_nodes: num_nodes_lig = self.max_num_nodes else: num_nodes_lig = self.ddpm.size_distribution.sample_conditional( n1=None, n2=pocket['size']) chain_lig, chain_pocket, _, _ = self.ddpm.sample_given_pocket( pocket, num_nodes_lig, return_frames=keep_frames) chain_lig = utils.reverse_tensor(chain_lig) chain_pocket = utils.reverse_tensor(chain_pocket) # Repeat last frame to see final sample better. chain_lig = torch.cat([chain_lig, chain_lig[-1:].repeat(10, 1, 1)], dim=0) chain_pocket = torch.cat( [chain_pocket, chain_pocket[-1:].repeat(10, 1, 1)], dim=0) # Prepare entire chain. x_lig = chain_lig[:, :, :self.x_dims] one_hot_lig = chain_lig[:, :, self.x_dims:] one_hot_lig = F.one_hot( torch.argmax(one_hot_lig, dim=2), num_classes=len(self.lig_type_decoder)) x_pocket = chain_pocket[:, :, :3] one_hot_pocket = chain_pocket[:, :, 3:] one_hot_pocket = F.one_hot( torch.argmax(one_hot_pocket, dim=2), num_classes=len(self.pocket_type_decoder)) if self.pocket_representation == 'CA': # convert residues into atom representation for visualization x_pocket, one_hot_pocket = utils.residues_to_atoms( x_pocket, self.lig_type_encoder) x = torch.cat((x_lig, x_pocket), dim=1) one_hot = torch.cat((one_hot_lig, one_hot_pocket), dim=1) # flatten (treat frame (chain dimension) as batch for visualization) x_flat = x.view(-1, x.size(-1)) one_hot_flat = one_hot.view(-1, one_hot.size(-1)) mask_flat = torch.arange(x.size(0)).repeat_interleave(x.size(1)) outdir = Path(self.outdir, f'epoch_{self.current_epoch}', 'chain') save_xyz_file(str(outdir), one_hot_flat, x_flat, self.lig_type_decoder, name='/chain', batch_mask=mask_flat) visualize_chain(str(outdir), self.dataset_info, wandb=wandb) def prepare_pocket(self, biopython_residues, repeats=1): if self.pocket_representation == 'CA': pocket_coord = torch.tensor(np.array( [res['CA'].get_coord() for res in biopython_residues]), device=self.device, dtype=FLOAT_TYPE) pocket_types = torch.tensor( [self.pocket_type_encoder[three_to_one(res.get_resname())] for res in biopython_residues], device=self.device) else: pocket_atoms = [a for res in biopython_residues for a in res.get_atoms() if (a.element.capitalize() in self.pocket_type_encoder or a.element != 'H')] pocket_coord = torch.tensor(np.array( [a.get_coord() for a in pocket_atoms]), device=self.device, dtype=FLOAT_TYPE) pocket_types = torch.tensor( [self.pocket_type_encoder[a.element.capitalize()] for a in pocket_atoms], device=self.device) pocket_one_hot = F.one_hot( pocket_types, num_classes=len(self.pocket_type_encoder) ) pocket_size = torch.tensor([len(pocket_coord)] * repeats, device=self.device, dtype=INT_TYPE) pocket_mask = torch.repeat_interleave( torch.arange(repeats, device=self.device, dtype=INT_TYPE), len(pocket_coord) ) pocket = { 'x': pocket_coord.repeat(repeats, 1), 'one_hot': pocket_one_hot.repeat(repeats, 1), 'size': pocket_size, 'mask': pocket_mask } return pocket def generate_ligands(self, pdb_file, n_samples, pocket_ids=None, ref_ligand=None, num_nodes_lig=None, sanitize=False, largest_frag=False, relax_iter=0, timesteps=None, n_nodes_bias=0, n_nodes_min=0, **kwargs): """ Generate ligands given a pocket Args: pdb_file: PDB filename n_samples: number of samples pocket_ids: list of pocket residues in : format ref_ligand: alternative way of defining the pocket based on a reference ligand given in : format if the ligand is contained in the PDB file, or path to an SDF file that contains the ligand num_nodes_lig: number of ligand nodes for each sample (list of integers), 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 n_nodes_bias: added to the sampled (or provided) number of nodes n_nodes_min: lower bound on the number of sampled nodes kwargs: additional inpainting parameters Returns: list of molecules """ assert (pocket_ids is None) ^ (ref_ligand is None) self.ddpm.eval() # Load PDB pdb_struct = PDBParser(QUIET=True).get_structure('', pdb_file)[0] if pocket_ids is not None: # define pocket with list of residues residues = [ pdb_struct[x.split(':')[0]][(' ', int(x.split(':')[1]), ' ')] for x in pocket_ids] else: # define pocket with reference ligand residues = utils.get_pocket_from_ligand(pdb_struct, ref_ligand) pocket = self.prepare_pocket(residues, repeats=n_samples) # Pocket's center of mass pocket_com_before = scatter_mean(pocket['x'], pocket['mask'], dim=0) # Create dummy ligands if num_nodes_lig is None: num_nodes_lig = self.ddpm.size_distribution.sample_conditional( n1=None, n2=pocket['size']) # Add bias num_nodes_lig = num_nodes_lig + n_nodes_bias # Apply minimum ligand size num_nodes_lig = torch.clamp(num_nodes_lig, min=n_nodes_min) # Use inpainting if type(self.ddpm) == EnVariationalDiffusion: lig_mask = utils.num_nodes_to_batch_mask( len(num_nodes_lig), num_nodes_lig, self.device) ligand = { 'x': torch.zeros((len(lig_mask), self.x_dims), device=self.device, dtype=FLOAT_TYPE), 'one_hot': torch.zeros((len(lig_mask), self.atom_nf), device=self.device, dtype=FLOAT_TYPE), 'size': num_nodes_lig, 'mask': lig_mask } # Fix all pocket nodes but sample lig_mask_fixed = torch.zeros(len(lig_mask), device=self.device) pocket_mask_fixed = torch.ones(len(pocket['mask']), device=self.device) xh_lig, xh_pocket, lig_mask, pocket_mask = self.ddpm.inpaint( ligand, pocket, lig_mask_fixed, pocket_mask_fixed, timesteps=timesteps, **kwargs) # Use conditional generation elif type(self.ddpm) == ConditionalDDPM: xh_lig, xh_pocket, lig_mask, pocket_mask = \ self.ddpm.sample_given_pocket(pocket, num_nodes_lig, timesteps=timesteps) else: raise NotImplementedError # Move generated molecule back to the original pocket position pocket_com_after = scatter_mean( xh_pocket[:, :self.x_dims], pocket_mask, dim=0) xh_pocket[:, :self.x_dims] += \ (pocket_com_before - pocket_com_after)[pocket_mask] xh_lig[:, :self.x_dims] += \ (pocket_com_before - pocket_com_after)[lig_mask] # Build mol objects x = xh_lig[:, :self.x_dims].detach().cpu() atom_type = xh_lig[:, self.x_dims:].argmax(1).detach().cpu() lig_mask = lig_mask.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, self.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 def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_val, gradient_clip_algorithm): if not self.clip_grad: return # Allow gradient norm to be 150% + 2 * stdev of the recent history. max_grad_norm = 1.5 * self.gradnorm_queue.mean() + \ 2 * self.gradnorm_queue.std() # Get current grad_norm params = [p for g in optimizer.param_groups for p in g['params']] grad_norm = utils.get_grad_norm(params) # Lightning will handle the gradient clipping self.clip_gradients(optimizer, gradient_clip_val=max_grad_norm, gradient_clip_algorithm='norm') if float(grad_norm) > max_grad_norm: self.gradnorm_queue.add(float(max_grad_norm)) else: self.gradnorm_queue.add(float(grad_norm)) if float(grad_norm) > max_grad_norm: print(f'Clipped gradient with value {grad_norm:.1f} ' f'while allowed {max_grad_norm:.1f}') class WeightSchedule: def __init__(self, T, max_weight, mode='linear'): if mode == 'linear': self.weights = torch.linspace(max_weight, 0, T + 1) elif mode == 'constant': self.weights = max_weight * torch.ones(T + 1) else: raise NotImplementedError(f'{mode} weight schedule is not ' f'available.') def __call__(self, t_array): """ all values in t_array are assumed to be integers in [0, T] """ return self.weights[t_array].to(t_array.device)