| 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 |
|
|
