import math import os import numpy as np import torch from e3nn import o3 from e3nn.o3 import Linear from torch import nn from torch.nn import functional as F from torch_cluster import radius, radius_graph from torch_scatter import scatter, scatter_mean from onescience.datapipes.diffdock.process_mols import ( lig_feature_dims, rec_atom_feature_dims, rec_residue_feature_dims, ) from onescience.utils.diffdock import so3, torus from .layers import AtomEncoder, GaussianSmearing from .tensor_layers import TensorProductConvLayer, get_irrep_seq def _env_flag(name): return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"} try: from esm.pretrained import load_model_and_alphabet except ImportError: load_model_and_alphabet = None class CGModel(torch.nn.Module): def __init__( self, t_to_sigma, device, timestep_emb_func, in_lig_edge_features=4, sigma_embed_dim=32, sh_lmax=2, ns=16, nv=4, num_conv_layers=2, lig_max_radius=5, rec_max_radius=30, cross_max_distance=250, center_max_distance=30, distance_embed_dim=32, cross_distance_embed_dim=32, no_torsion=False, scale_by_sigma=True, norm_by_sigma=True, use_second_order_repr=False, batch_norm=True, dynamic_max_cross=False, dropout=0.0, smooth_edges=False, odd_parity=False, separate_noise_schedule=False, lm_embedding_type=None, confidence_mode=False, confidence_dropout=0, confidence_no_batchnorm=False, asyncronous_noise_schedule=False, affinity_prediction=False, parallel=1, parallel_aggregators="mean max min std", num_confidence_outputs=1, atom_num_confidence_outputs=1, fixed_center_conv=False, no_aminoacid_identities=False, include_miscellaneous_atoms=False, differentiate_convolutions=True, tp_weights_layers=2, num_prot_emb_layers=0, reduce_pseudoscalars=False, embed_also_ligand=False, atom_confidence=False, sidechain_pred=False, depthwise_convolution=False, ): super(CGModel, self).__init__() assert parallel == 1, "not implemented" assert (not no_aminoacid_identities) or ( lm_embedding_type is None ), "no language model emb without identities" self.t_to_sigma = t_to_sigma self.in_lig_edge_features = in_lig_edge_features sigma_embed_dim *= 3 if separate_noise_schedule else 1 self.sigma_embed_dim = sigma_embed_dim self.lig_max_radius = lig_max_radius self.rec_max_radius = rec_max_radius self.include_miscellaneous_atoms = include_miscellaneous_atoms self.cross_max_distance = cross_max_distance self.dynamic_max_cross = dynamic_max_cross self.center_max_distance = center_max_distance self.distance_embed_dim = distance_embed_dim self.cross_distance_embed_dim = cross_distance_embed_dim self.sh_irreps = o3.Irreps.spherical_harmonics(lmax=sh_lmax) self.ns, self.nv = ns, nv self.scale_by_sigma = scale_by_sigma self.norm_by_sigma = norm_by_sigma self.device = device self.no_torsion = no_torsion self.smooth_edges = smooth_edges self.odd_parity = odd_parity self.timestep_emb_func = timestep_emb_func self.separate_noise_schedule = separate_noise_schedule self.confidence_mode = confidence_mode self.num_conv_layers = num_conv_layers self.num_prot_emb_layers = num_prot_emb_layers self.asyncronous_noise_schedule = asyncronous_noise_schedule self.affinity_prediction = affinity_prediction self.fixed_center_conv = fixed_center_conv self.no_aminoacid_identities = no_aminoacid_identities self.differentiate_convolutions = differentiate_convolutions self.reduce_pseudoscalars = reduce_pseudoscalars self.atom_confidence = atom_confidence self.atom_num_confidence_outputs = atom_num_confidence_outputs self.sidechain_pred = sidechain_pred self.lm_embedding_type = lm_embedding_type if lm_embedding_type is None: lm_embedding_dim = 0 elif lm_embedding_type == "precomputed": lm_embedding_dim = 1280 else: if load_model_and_alphabet is None: raise ImportError( "esm is required when lm_embedding_type is not None and not 'precomputed'." ) lm, alphabet = load_model_and_alphabet(lm_embedding_type) self.batch_converter = alphabet.get_batch_converter() lm.lm_head = torch.nn.Identity() lm.contact_head = torch.nn.Identity() lm_embedding_dim = lm.embed_dim self.lm = lm atom_encoder_class = AtomEncoder self.lig_node_embedding = atom_encoder_class( emb_dim=ns, feature_dims=lig_feature_dims, sigma_embed_dim=sigma_embed_dim, ) self.lig_edge_embedding = nn.Sequential( nn.Linear(in_lig_edge_features + sigma_embed_dim + distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) self.rec_node_embedding = atom_encoder_class( emb_dim=ns, feature_dims=rec_residue_feature_dims, sigma_embed_dim=0, lm_embedding_dim=lm_embedding_dim, ) self.rec_edge_embedding = nn.Sequential( nn.Linear(distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) self.rec_sigma_embedding = nn.Sequential( nn.Linear(sigma_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) if self.include_miscellaneous_atoms: self.misc_atom_node_embedding = atom_encoder_class( emb_dim=ns, feature_dims=rec_atom_feature_dims, sigma_embed_dim=sigma_embed_dim, ) self.misc_atom_edge_embedding = nn.Sequential( nn.Linear(sigma_embed_dim + distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) self.ar_edge_embedding = nn.Sequential( nn.Linear(sigma_embed_dim + distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) self.la_edge_embedding = nn.Sequential( nn.Linear(sigma_embed_dim + cross_distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) self.cross_edge_embedding = nn.Sequential( nn.Linear(sigma_embed_dim + cross_distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) self.lig_distance_expansion = GaussianSmearing(0.0, lig_max_radius, distance_embed_dim) self.rec_distance_expansion = GaussianSmearing(0.0, rec_max_radius, distance_embed_dim) self.cross_distance_expansion = GaussianSmearing( 0.0, cross_max_distance, cross_distance_embed_dim ) irrep_seq = get_irrep_seq(ns, nv, use_second_order_repr, reduce_pseudoscalars) assert not self.include_miscellaneous_atoms, "currently not supported" rec_emb_layers = [] for i in range(num_prot_emb_layers): in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)] out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)] layer = TensorProductConvLayer( in_irreps=in_irreps, sh_irreps=self.sh_irreps, out_irreps=out_irreps, n_edge_features=3 * ns, hidden_features=3 * ns, residual=True, batch_norm=batch_norm, dropout=dropout, faster=sh_lmax == 1 and not use_second_order_repr, tp_weights_layers=tp_weights_layers, edge_groups=1, depthwise=depthwise_convolution, ) rec_emb_layers.append(layer) self.rec_emb_layers = nn.ModuleList(rec_emb_layers) self.embed_also_ligand = embed_also_ligand if embed_also_ligand: lig_emb_layers = [] for i in range(num_prot_emb_layers): in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)] out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)] layer = TensorProductConvLayer( in_irreps=in_irreps, sh_irreps=self.sh_irreps, out_irreps=out_irreps, n_edge_features=3 * ns, hidden_features=3 * ns, residual=True, batch_norm=batch_norm, dropout=dropout, faster=sh_lmax == 1 and not use_second_order_repr, tp_weights_layers=tp_weights_layers, edge_groups=1, depthwise=depthwise_convolution, ) lig_emb_layers.append(layer) self.lig_emb_layers = nn.ModuleList(lig_emb_layers) conv_layers = [] for i in range(num_prot_emb_layers, num_prot_emb_layers + num_conv_layers): in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)] out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)] layer = TensorProductConvLayer( in_irreps=in_irreps, sh_irreps=self.sh_irreps, out_irreps=out_irreps, n_edge_features=3 * ns, hidden_features=3 * ns, residual=True, batch_norm=batch_norm, dropout=dropout, faster=sh_lmax == 1 and not use_second_order_repr, tp_weights_layers=tp_weights_layers, edge_groups=( 1 if not differentiate_convolutions else (2 if i == num_prot_emb_layers + num_conv_layers - 1 else 4) ), depthwise=depthwise_convolution, ) conv_layers.append(layer) self.conv_layers = nn.ModuleList(conv_layers) if sidechain_pred: self.sidechain_predictor = Linear( irreps_in=irrep_seq[min(num_prot_emb_layers + num_conv_layers, len(irrep_seq) - 1)], irreps_out="4x0e + 2x1e + 4x0o + 2x1o", internal_weights=True, shared_weights=True, ) if self.confidence_mode: input_size = ( ns + (nv if reduce_pseudoscalars else ns) if num_conv_layers + num_prot_emb_layers >= 3 else ns ) if self.atom_confidence: self.atom_confidence_predictor = nn.Sequential( nn.Linear(input_size, ns), nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(), nn.ReLU(), nn.Dropout(confidence_dropout), nn.Linear(ns, ns), nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(), nn.ReLU(), nn.Dropout(confidence_dropout), nn.Linear(ns, atom_num_confidence_outputs + ns), ) input_size = ns self.confidence_predictor = nn.Sequential( nn.Linear(input_size, ns), nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(), nn.ReLU(), nn.Dropout(confidence_dropout), nn.Linear(ns, ns), nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(), nn.ReLU(), nn.Dropout(confidence_dropout), nn.Linear(ns, num_confidence_outputs + (1 if self.affinity_prediction else 0)), ) else: self.center_distance_expansion = GaussianSmearing( 0.0, center_max_distance, distance_embed_dim ) self.center_edge_embedding = nn.Sequential( nn.Linear(distance_embed_dim + sigma_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) self.final_conv = TensorProductConvLayer( in_irreps=self.conv_layers[-1].out_irreps, sh_irreps=self.sh_irreps, out_irreps="2x1o + 2x1e" if not self.odd_parity else "1x1o + 1x1e", n_edge_features=2 * ns, residual=False, dropout=dropout, batch_norm=batch_norm, ) self.tr_final_layer = nn.Sequential( nn.Linear(1 + sigma_embed_dim, ns), nn.Dropout(dropout), nn.ReLU(), nn.Linear(ns, 1), ) self.rot_final_layer = nn.Sequential( nn.Linear(1 + sigma_embed_dim, ns), nn.Dropout(dropout), nn.ReLU(), nn.Linear(ns, 1), ) if not no_torsion: self.final_edge_embedding = nn.Sequential( nn.Linear(distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns), ) self.final_tp_tor = o3.FullTensorProduct(self.sh_irreps, "2e") self.tor_bond_conv = TensorProductConvLayer( in_irreps=self.conv_layers[-1].out_irreps, sh_irreps=self.final_tp_tor.irreps_out, out_irreps=f"{ns}x0o + {ns}x0e" if not self.odd_parity else f"{ns}x0o", n_edge_features=3 * ns, residual=False, dropout=dropout, batch_norm=batch_norm, ) self.tor_final_layer = nn.Sequential( nn.Linear(2 * ns if not self.odd_parity else ns, ns, bias=False), nn.Tanh(), nn.Dropout(dropout), nn.Linear(ns, 1, bias=False), ) @staticmethod def _resolve_edge_store(data, primary_key, fallback_key): edge_types = getattr(data, "edge_types", ()) if primary_key in edge_types: return data[primary_key] if fallback_key in edge_types: return data[fallback_key] try: return data[primary_key] except Exception: return data[fallback_key] def _ligand_edge_store(self, data): return self._resolve_edge_store( data, ("ligand", "ligand"), ("ligand", "lig_bond", "ligand"), ) def _receptor_edge_store(self, data): return self._resolve_edge_store( data, ("receptor", "receptor"), ("receptor", "rec_contact", "receptor"), ) def ligand_embedding(self, data): lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight = self.build_lig_conv_graph( data ) lig_node_attr = self.lig_node_embedding(lig_node_attr) lig_edge_attr = self.lig_edge_embedding(lig_edge_attr) assert self.embed_also_ligand, "otherwise reimplement padding" for l in range(len(self.lig_emb_layers)): edge_attr_ = torch.cat( [ lig_edge_attr, lig_node_attr[lig_edge_index[0], : self.ns], lig_node_attr[lig_edge_index[1], : self.ns], ], -1, ) lig_node_attr = self.lig_emb_layers[l]( lig_node_attr, lig_edge_index, edge_attr_, lig_edge_sh, edge_weight=lig_edge_weight, ) return lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight def embedding(self, data): receptor_edge_store = self._receptor_edge_store(data) if not hasattr(data["receptor"], "rec_node_attr"): if self.lm_embedding_type not in [None, "precomputed"]: sequences = [s for l in data["receptor"].sequence for s in l] if isinstance(sequences[0], list): sequences = [s for l in sequences for s in l] sequences = [(i, s) for i, s in enumerate(sequences)] batch_labels, batch_strs, batch_tokens = self.batch_converter(sequences) out = self.lm( batch_tokens.to(data["receptor"].x.device), repr_layers=[self.lm.num_layers], return_contacts=False, ) rec_lm_emb = torch.cat( [ t[: len(sequences[i][1])] for i, t in enumerate(out["representations"][self.lm.num_layers]) ], dim=0, ) data["receptor"].x = torch.cat([data["receptor"].x, rec_lm_emb], dim=-1) rec_node_attr, rec_edge_attr, rec_edge_sh, rec_edge_weight = self.build_rec_conv_graph(data) rec_node_attr = self.rec_node_embedding(rec_node_attr) rec_edge_attr = self.rec_edge_embedding(rec_edge_attr) for l in range(len(self.rec_emb_layers)): edge_attr_ = torch.cat( [ rec_edge_attr, rec_node_attr[receptor_edge_store.edge_index[0], : self.ns], rec_node_attr[receptor_edge_store.edge_index[1], : self.ns], ], -1, ) rec_node_attr = self.rec_emb_layers[l]( rec_node_attr, receptor_edge_store.edge_index, edge_attr_, rec_edge_sh, edge_weight=rec_edge_weight, ) data["receptor"].rec_node_attr = rec_node_attr receptor_edge_store.rec_edge_attr = rec_edge_attr receptor_edge_store.edge_sh = rec_edge_sh receptor_edge_store.edge_weight = rec_edge_weight rec_sigma_emb = self.rec_sigma_embedding(self.timestep_emb_func(data.complex_t["tr"])) rec_node_attr = data["receptor"].rec_node_attr + 0 rec_node_attr[:, : self.ns] = rec_node_attr[:, : self.ns] + rec_sigma_emb[data["receptor"].batch] rec_edge_attr = receptor_edge_store.rec_edge_attr + rec_sigma_emb[ data["receptor"].batch[receptor_edge_store.edge_index[0]] ] lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight = self.ligand_embedding( data ) return ( lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight, rec_node_attr, receptor_edge_store.edge_index, rec_edge_attr, receptor_edge_store.edge_sh, receptor_edge_store.edge_weight, ) def forward(self, data): if self.no_aminoacid_identities: data["receptor"].x = data["receptor"].x * 0 if not self.confidence_mode: tr_sigma, rot_sigma, tor_sigma = self.t_to_sigma( *[data.complex_t[noise_type] for noise_type in ["tr", "rot", "tor"]] ) else: tr_sigma, rot_sigma, tor_sigma = [ data.complex_t[noise_type] for noise_type in ["tr", "rot", "tor"] ] ( lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight, rec_node_attr, rec_edge_index, rec_edge_attr, rec_edge_sh, rec_edge_weight, ) = self.embedding(data) if self.dynamic_max_cross: cross_cutoff = (tr_sigma * 3 + 20).unsqueeze(1) else: cross_cutoff = self.cross_max_distance lr_edge_index, lr_edge_attr, lr_edge_sh, rev_lr_edge_sh, lr_edge_weight = self.build_cross_conv_graph( data, cross_cutoff ) lr_edge_attr = self.cross_edge_embedding(lr_edge_attr) node_attr = torch.cat([lig_node_attr, rec_node_attr], dim=0) lr_edge_index[1] = lr_edge_index[1] + len(lig_node_attr) edge_index = torch.cat( [ lig_edge_index, lr_edge_index, rec_edge_index + len(lig_node_attr), torch.flip(lr_edge_index, dims=[0]), ], dim=1, ) edge_attr = torch.cat([lig_edge_attr, lr_edge_attr, rec_edge_attr, lr_edge_attr], dim=0) edge_sh = torch.cat([lig_edge_sh, lr_edge_sh, rec_edge_sh, rev_lr_edge_sh], dim=0) edge_weight = ( torch.cat([lig_edge_weight, lr_edge_weight, rec_edge_weight, lr_edge_weight], dim=0) if torch.is_tensor(lig_edge_weight) else torch.ones((len(edge_index[0]), 1), device=edge_index.device) ) s1, s2, s3 = ( len(lig_edge_index[0]), len(lig_edge_index[0]) + len(lr_edge_index[0]), len(lig_edge_index[0]) + len(lr_edge_index[0]) + len(rec_edge_index[0]), ) for l in range(len(self.conv_layers)): if l < len(self.conv_layers) - 1: edge_attr_ = torch.cat( [ edge_attr, node_attr[edge_index[0], : self.ns], node_attr[edge_index[1], : self.ns], ], -1, ) if self.differentiate_convolutions: edge_attr_ = [edge_attr_[:s1], edge_attr_[s1:s2], edge_attr_[s2:s3], edge_attr_[s3:]] node_attr = self.conv_layers[l]( node_attr, edge_index, edge_attr_, edge_sh, edge_weight=edge_weight, ) else: edge_attr_ = torch.cat( [ edge_attr[:s2], node_attr[edge_index[0, :s2], : self.ns], node_attr[edge_index[1, :s2], : self.ns], ], -1, ) if self.differentiate_convolutions: edge_attr_ = [edge_attr_[:s1], edge_attr_[s1:s2]] node_attr = self.conv_layers[l]( node_attr, edge_index[:, :s2], edge_attr_, edge_sh[:s2], edge_weight=edge_weight[:s2], ) lig_node_attr = node_attr[: len(lig_node_attr)] if self.confidence_mode: scalar_lig_attr = ( torch.cat( [ lig_node_attr[:, : self.ns], lig_node_attr[:, -(self.nv if self.reduce_pseudoscalars else self.ns) :], ], dim=1, ) if self.num_conv_layers + self.num_prot_emb_layers >= 3 else lig_node_attr[:, : self.ns] ) if self.atom_confidence: scalar_lig_attr = self.atom_confidence_predictor(scalar_lig_attr) atom_confidence = scalar_lig_attr[:, : self.atom_num_confidence_outputs] scalar_lig_attr = scalar_lig_attr[:, self.atom_num_confidence_outputs :] else: atom_confidence = torch.zeros((len(lig_node_attr),), device=lig_node_attr.device) confidence = self.confidence_predictor( scatter_mean(scalar_lig_attr, data["ligand"].batch, dim=0) ).squeeze(dim=-1) return confidence, atom_confidence center_edge_index, center_edge_attr, center_edge_sh = self.build_center_conv_graph(data) center_edge_attr = self.center_edge_embedding(center_edge_attr) if self.fixed_center_conv: center_edge_attr = torch.cat( [center_edge_attr, lig_node_attr[center_edge_index[1], : self.ns]], -1 ) else: center_edge_attr = torch.cat( [center_edge_attr, lig_node_attr[center_edge_index[0], : self.ns]], -1 ) global_pred = self.final_conv( lig_node_attr, center_edge_index, center_edge_attr, center_edge_sh, out_nodes=data.num_graphs, ) tr_pred = global_pred[:, :3] + (global_pred[:, 6:9] if not self.odd_parity else 0) rot_pred = global_pred[:, 3:6] + (global_pred[:, 9:] if not self.odd_parity else 0) if self.separate_noise_schedule: data.graph_sigma_emb = torch.cat( [self.timestep_emb_func(data.complex_t[noise_type]) for noise_type in ["tr", "rot", "tor"]], dim=1, ) elif self.asyncronous_noise_schedule: data.graph_sigma_emb = self.timestep_emb_func(data.complex_t["t"]) else: data.graph_sigma_emb = self.timestep_emb_func(data.complex_t["tr"]) tr_norm = torch.linalg.vector_norm(tr_pred, dim=1).unsqueeze(1) tr_pred = tr_pred / tr_norm * self.tr_final_layer( torch.cat([tr_norm, data.graph_sigma_emb], dim=1) ) rot_norm = torch.linalg.vector_norm(rot_pred, dim=1).unsqueeze(1) rot_pred = rot_pred / rot_norm * self.rot_final_layer( torch.cat([rot_norm, data.graph_sigma_emb], dim=1) ) if self.scale_by_sigma: tr_pred = tr_pred / tr_sigma.unsqueeze(1) rot_pred = rot_pred * so3.score_norm(rot_sigma.cpu()).unsqueeze(1).to( data["ligand"].x.device ) sidechain_pred = None if self.sidechain_pred: rec_node_attr = node_attr[len(lig_node_attr) :] sidechain_pred = self.sidechain_predictor(rec_node_attr) sidechain_pred = sidechain_pred[:, :10] + sidechain_pred[:, 10:] if self.no_torsion or data["ligand"].edge_mask.sum() == 0: return tr_pred, rot_pred, torch.empty(0, device=self.device), sidechain_pred tor_bonds, tor_edge_index, tor_edge_attr, tor_edge_sh, tor_edge_weight = self.build_bond_conv_graph( data ) tor_bond_vec = data["ligand"].pos[tor_bonds[1]] - data["ligand"].pos[tor_bonds[0]] tor_bond_attr = lig_node_attr[tor_bonds[0]] + lig_node_attr[tor_bonds[1]] tor_bonds_sh = o3.spherical_harmonics( "2e", tor_bond_vec, normalize=True, normalization="component" ) tor_edge_sh = self.final_tp_tor(tor_edge_sh, tor_bonds_sh[tor_edge_index[0]]) tor_edge_attr = torch.cat( [ tor_edge_attr, lig_node_attr[tor_edge_index[1], : self.ns], tor_bond_attr[tor_edge_index[0], : self.ns], ], -1, ) tor_pred = self.tor_bond_conv( lig_node_attr, tor_edge_index, tor_edge_attr, tor_edge_sh, out_nodes=data["ligand"].edge_mask.sum(), reduce="mean", edge_weight=tor_edge_weight, ) tor_pred = self.tor_final_layer(tor_pred).squeeze(1) ligand_edge_store = self._ligand_edge_store(data) edge_sigma = tor_sigma[data["ligand"].batch][ligand_edge_store.edge_index[0]][ data["ligand"].edge_mask ] if self.scale_by_sigma: tor_pred = tor_pred * torch.sqrt( torch.tensor(torus.score_norm(edge_sigma.cpu().numpy())).float().to( data["ligand"].x.device ) ) return tr_pred, rot_pred, tor_pred, sidechain_pred def torsional_forward(self, data): tor_sigma = self.t_to_sigma(data.complex_t["tor"]) lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight = self.ligand_embedding( data ) if self.separate_noise_schedule: data.graph_sigma_emb = torch.cat( [self.timestep_emb_func(data.complex_t[noise_type]) for noise_type in ["tr", "rot", "tor"]], dim=1, ) elif self.asyncronous_noise_schedule: data.graph_sigma_emb = self.timestep_emb_func(data.complex_t["t"]) else: data.graph_sigma_emb = self.timestep_emb_func(data.complex_t["tr"]) tor_bonds, tor_edge_index, tor_edge_attr, tor_edge_sh, tor_edge_weight = self.build_bond_conv_graph( data ) tor_bond_vec = data["ligand"].pos[tor_bonds[1]] - data["ligand"].pos[tor_bonds[0]] tor_bond_attr = lig_node_attr[tor_bonds[0]] + lig_node_attr[tor_bonds[1]] tor_bonds_sh = o3.spherical_harmonics( "2e", tor_bond_vec, normalize=True, normalization="component" ) tor_edge_sh = self.final_tp_tor(tor_edge_sh, tor_bonds_sh[tor_edge_index[0]]) tor_edge_attr = torch.cat( [ tor_edge_attr, lig_node_attr[tor_edge_index[1], : self.ns], tor_bond_attr[tor_edge_index[0], : self.ns], ], -1, ) tor_pred = self.tor_bond_conv( lig_node_attr, tor_edge_index, tor_edge_attr, tor_edge_sh, out_nodes=data["ligand"].edge_mask.sum(), reduce="mean", edge_weight=tor_edge_weight, ) tor_pred = self.tor_final_layer(tor_pred).squeeze(1) ligand_edge_store = self._ligand_edge_store(data) edge_sigma = tor_sigma[data["ligand"].batch][ligand_edge_store.edge_index[0]][ data["ligand"].edge_mask ] if self.scale_by_sigma: tor_pred = tor_pred * torch.sqrt( torch.tensor(torus.score_norm(edge_sigma.cpu().numpy())).float().to( data["ligand"].x.device ) ) return 0, 0, tor_pred, 0 def get_edge_weight(self, edge_vec, max_norm): if self.smooth_edges: normalised_norm = torch.clip(edge_vec.norm(dim=-1) * np.pi / max_norm, max=np.pi) return 0.5 * (torch.cos(normalised_norm) + 1.0).unsqueeze(-1) return 1.0 def build_lig_conv_graph(self, data): if self.separate_noise_schedule: data["ligand"].node_sigma_emb = torch.cat( [self.timestep_emb_func(data["ligand"].node_t[noise_type]) for noise_type in ["tr", "rot", "tor"]], dim=1, ) elif self.asyncronous_noise_schedule: data["ligand"].node_sigma_emb = self.timestep_emb_func(data["ligand"].node_t["t"]) else: data["ligand"].node_sigma_emb = self.timestep_emb_func(data["ligand"].node_t["tr"]) ligand_edge_store = self._ligand_edge_store(data) radius_edges = radius_graph(data["ligand"].pos, self.lig_max_radius, data["ligand"].batch) edge_index = torch.cat([ligand_edge_store.edge_index, radius_edges], 1).long() edge_attr = torch.cat( [ ligand_edge_store.edge_attr, torch.zeros( radius_edges.shape[-1], self.in_lig_edge_features, device=data["ligand"].x.device, ), ], 0, ) edge_sigma_emb = data["ligand"].node_sigma_emb[edge_index[0].long()] edge_attr = torch.cat([edge_attr, edge_sigma_emb], 1) node_attr = torch.cat([data["ligand"].x, data["ligand"].node_sigma_emb], 1) src, dst = edge_index edge_vec = data["ligand"].pos[dst.long()] - data["ligand"].pos[src.long()] edge_length_emb = self.lig_distance_expansion(edge_vec.norm(dim=-1)) edge_attr = torch.cat([edge_attr, edge_length_emb], 1) edge_sh = o3.spherical_harmonics( self.sh_irreps, edge_vec, normalize=True, normalization="component" ) edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius) return node_attr, edge_index, edge_attr, edge_sh, edge_weight def build_rec_conv_graph(self, data): assert not self.separate_noise_schedule or self.asyncronous_noise_schedule, "removed support in this function" node_attr = data["receptor"].x receptor_edge_store = self._receptor_edge_store(data) edge_index = receptor_edge_store.edge_index src, dst = edge_index edge_vec = data["receptor"].pos[dst.long()] - data["receptor"].pos[src.long()] edge_length_emb = self.rec_distance_expansion(edge_vec.norm(dim=-1)) edge_attr = edge_length_emb edge_sh = o3.spherical_harmonics( self.sh_irreps, edge_vec, normalize=True, normalization="component" ) edge_weight = self.get_edge_weight(edge_vec, self.rec_max_radius) return node_attr, edge_attr, edge_sh, edge_weight def build_misc_atom_conv_graph(self, data): if self.separate_noise_schedule: data["misc_atom"].node_sigma_emb = torch.cat( [self.timestep_emb_func(data["misc_atom"].node_t[noise_type]) for noise_type in ["tr", "rot", "tor"]], dim=1, ) elif self.asyncronous_noise_schedule: data["misc_atom"].node_sigma_emb = self.timestep_emb_func(data["misc_atom"].node_t["t"]) else: data["misc_atom"].node_sigma_emb = self.timestep_emb_func(data["misc_atom"].node_t["tr"]) node_attr = torch.cat([data["misc_atom"].x, data["misc_atom"].node_sigma_emb], 1) edge_index = data["misc_atom", "misc_atom"].edge_index src, dst = edge_index edge_vec = data["misc_atom"].pos[dst.long()] - data["misc_atom"].pos[src.long()] edge_length_emb = self.lig_distance_expansion(edge_vec.norm(dim=-1)) edge_sigma_emb = data["misc_atom"].node_sigma_emb[edge_index[0].long()] edge_attr = torch.cat([edge_sigma_emb, edge_length_emb], 1) edge_sh = o3.spherical_harmonics( self.sh_irreps, edge_vec, normalize=True, normalization="component" ) edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius) return node_attr, edge_index, edge_attr, edge_sh, edge_weight def build_cross_conv_graph(self, data, cross_distance_cutoff): radius_on_cpu = _env_flag("DIFFDOCK_RADIUS_ON_CPU") if torch.is_tensor(cross_distance_cutoff): receptor_pos = data["receptor"].pos ligand_pos = data["ligand"].pos receptor_batch = data["receptor"].batch ligand_batch = data["ligand"].batch cutoff = cross_distance_cutoff if radius_on_cpu: receptor_pos = receptor_pos.detach().cpu() ligand_pos = ligand_pos.detach().cpu() receptor_batch = receptor_batch.detach().cpu() ligand_batch = ligand_batch.detach().cpu() cutoff = cutoff.detach().cpu() edge_index = radius( receptor_pos / cutoff[receptor_batch], ligand_pos / cutoff[ligand_batch], 1, receptor_batch, ligand_batch, max_num_neighbors=10000, ) else: receptor_pos = data["receptor"].pos ligand_pos = data["ligand"].pos receptor_batch = data["receptor"].batch ligand_batch = data["ligand"].batch if radius_on_cpu: receptor_pos = receptor_pos.detach().cpu() ligand_pos = ligand_pos.detach().cpu() receptor_batch = receptor_batch.detach().cpu() ligand_batch = ligand_batch.detach().cpu() edge_index = radius( receptor_pos, ligand_pos, cross_distance_cutoff, receptor_batch, ligand_batch, max_num_neighbors=10000, ) if radius_on_cpu: edge_index = edge_index.to(data["ligand"].pos.device) src, dst = edge_index edge_vec = data["receptor"].pos[dst.long()] - data["ligand"].pos[src.long()] edge_length_emb = self.cross_distance_expansion(edge_vec.norm(dim=-1)) edge_sigma_emb = data["ligand"].node_sigma_emb[src.long()] edge_attr = torch.cat([edge_sigma_emb, edge_length_emb], 1) edge_sh = o3.spherical_harmonics( self.sh_irreps, edge_vec, normalize=True, normalization="component" ) rev_edge_sh = o3.spherical_harmonics( self.sh_irreps, -edge_vec, normalize=True, normalization="component" ) cutoff_d = ( cross_distance_cutoff[data["ligand"].batch[src]].squeeze() if torch.is_tensor(cross_distance_cutoff) else cross_distance_cutoff ) edge_weight = self.get_edge_weight(edge_vec, cutoff_d) return edge_index, edge_attr, edge_sh, rev_edge_sh, edge_weight def build_misc_cross_conv_graph(self, data, lr_cross_distance_cutoff): if torch.is_tensor(lr_cross_distance_cutoff): lr_edge_index = radius( data["receptor"].pos / lr_cross_distance_cutoff[data["receptor"].batch], data["ligand"].pos / lr_cross_distance_cutoff[data["ligand"].batch], 1, data["receptor"].batch, data["ligand"].batch, max_num_neighbors=10000, ) else: lr_edge_index = radius( data["receptor"].pos, data["ligand"].pos, lr_cross_distance_cutoff, data["receptor"].batch, data["ligand"].batch, max_num_neighbors=10000, ) lr_edge_vec = data["receptor"].pos[lr_edge_index[1].long()] - data["ligand"].pos[lr_edge_index[0].long()] lr_edge_length_emb = self.cross_distance_expansion(lr_edge_vec.norm(dim=-1)) lr_edge_sigma_emb = data["ligand"].node_sigma_emb[lr_edge_index[0].long()] lr_edge_attr = torch.cat([lr_edge_sigma_emb, lr_edge_length_emb], 1) lr_edge_sh = o3.spherical_harmonics( self.sh_irreps, lr_edge_vec, normalize=True, normalization="component" ) cutoff_d = ( lr_cross_distance_cutoff[data["ligand"].batch[lr_edge_index[0]]].squeeze() if torch.is_tensor(lr_cross_distance_cutoff) else lr_cross_distance_cutoff ) lr_edge_weight = self.get_edge_weight(lr_edge_vec, cutoff_d) la_edge_index = radius( data["misc_atom"].pos, data["ligand"].pos, self.lig_max_radius, data["misc_atom"].batch, data["ligand"].batch, max_num_neighbors=10000, ) la_edge_vec = data["misc_atom"].pos[la_edge_index[1].long()] - data["ligand"].pos[la_edge_index[0].long()] la_edge_length_emb = self.cross_distance_expansion(la_edge_vec.norm(dim=-1)) la_edge_sigma_emb = data["ligand"].node_sigma_emb[la_edge_index[0].long()] la_edge_attr = torch.cat([la_edge_sigma_emb, la_edge_length_emb], 1) la_edge_sh = o3.spherical_harmonics( self.sh_irreps, la_edge_vec, normalize=True, normalization="component" ) la_edge_weight = self.get_edge_weight(la_edge_vec, self.lig_max_radius) ar_edge_index = data["misc_atom", "receptor"].edge_index ar_edge_vec = data["receptor"].pos[ar_edge_index[1].long()] - data["misc_atom"].pos[ar_edge_index[0].long()] ar_edge_length_emb = self.rec_distance_expansion(ar_edge_vec.norm(dim=-1)) ar_edge_sigma_emb = data["misc_atom"].node_sigma_emb[ar_edge_index[0].long()] ar_edge_attr = torch.cat([ar_edge_sigma_emb, ar_edge_length_emb], 1) ar_edge_sh = o3.spherical_harmonics( self.sh_irreps, ar_edge_vec, normalize=True, normalization="component" ) ar_edge_weight = 1 return ( lr_edge_index, lr_edge_attr, lr_edge_sh, lr_edge_weight, la_edge_index, la_edge_attr, la_edge_sh, la_edge_weight, ar_edge_index, ar_edge_attr, ar_edge_sh, ar_edge_weight, ) def build_center_conv_graph(self, data): edge_index = torch.cat( [ data["ligand"].batch.unsqueeze(0), torch.arange(len(data["ligand"].batch)).to(data["ligand"].x.device).unsqueeze(0), ], dim=0, ) center_pos, count = torch.zeros((data.num_graphs, 3)).to(data["ligand"].x.device), torch.zeros( (data.num_graphs, 3) ).to(data["ligand"].x.device) center_pos.index_add_(0, index=data["ligand"].batch, source=data["ligand"].pos) center_pos = center_pos / torch.bincount(data["ligand"].batch).unsqueeze(1) edge_vec = data["ligand"].pos[edge_index[1]] - center_pos[edge_index[0]] edge_attr = self.center_distance_expansion(edge_vec.norm(dim=-1)) edge_sigma_emb = data["ligand"].node_sigma_emb[edge_index[1].long()] edge_attr = torch.cat([edge_attr, edge_sigma_emb], 1) edge_sh = o3.spherical_harmonics( self.sh_irreps, edge_vec, normalize=True, normalization="component" ) return edge_index, edge_attr, edge_sh def build_bond_conv_graph(self, data): ligand_edge_store = self._ligand_edge_store(data) bonds = ligand_edge_store.edge_index[:, data["ligand"].edge_mask].long() bond_pos = (data["ligand"].pos[bonds[0]] + data["ligand"].pos[bonds[1]]) / 2 bond_batch = data["ligand"].batch[bonds[0]] edge_index = radius( data["ligand"].pos, bond_pos, self.lig_max_radius, batch_x=data["ligand"].batch, batch_y=bond_batch, ) edge_vec = data["ligand"].pos[edge_index[1]] - bond_pos[edge_index[0]] edge_attr = self.lig_distance_expansion(edge_vec.norm(dim=-1)) edge_attr = self.final_edge_embedding(edge_attr) edge_sh = o3.spherical_harmonics( self.sh_irreps, edge_vec, normalize=True, normalization="component" ) edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius) return bonds, edge_index, edge_attr, edge_sh, edge_weight