Spaces:
Running
on
T4
Running
on
T4
| import math | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from scipy.stats import beta | |
| from utils.geometry import axis_angle_to_matrix, rigid_transform_Kabsch_3D_torch | |
| from utils.torsion import modify_conformer_torsion_angles | |
| def t_to_sigma(t_tr, t_rot, t_tor, args): | |
| tr_sigma = args.tr_sigma_min ** (1-t_tr) * args.tr_sigma_max ** t_tr | |
| rot_sigma = args.rot_sigma_min ** (1-t_rot) * args.rot_sigma_max ** t_rot | |
| tor_sigma = args.tor_sigma_min ** (1-t_tor) * args.tor_sigma_max ** t_tor | |
| return tr_sigma, rot_sigma, tor_sigma | |
| def modify_conformer(data, tr_update, rot_update, torsion_updates): | |
| lig_center = torch.mean(data['ligand'].pos, dim=0, keepdim=True) | |
| rot_mat = axis_angle_to_matrix(rot_update.squeeze()) | |
| rigid_new_pos = (data['ligand'].pos - lig_center) @ rot_mat.T + tr_update + lig_center | |
| if torsion_updates is not None: | |
| flexible_new_pos = modify_conformer_torsion_angles(rigid_new_pos, | |
| data['ligand', 'ligand'].edge_index.T[data['ligand'].edge_mask], | |
| data['ligand'].mask_rotate if isinstance(data['ligand'].mask_rotate, np.ndarray) else data['ligand'].mask_rotate[0], | |
| torsion_updates).to(rigid_new_pos.device) | |
| R, t = rigid_transform_Kabsch_3D_torch(flexible_new_pos.T, rigid_new_pos.T) | |
| aligned_flexible_pos = flexible_new_pos @ R.T + t.T | |
| data['ligand'].pos = aligned_flexible_pos | |
| else: | |
| data['ligand'].pos = rigid_new_pos | |
| return data | |
| def sinusoidal_embedding(timesteps, embedding_dim, max_positions=10000): | |
| """ from https://github.com/hojonathanho/diffusion/blob/master/diffusion_tf/nn.py """ | |
| assert len(timesteps.shape) == 1 | |
| half_dim = embedding_dim // 2 | |
| emb = math.log(max_positions) / (half_dim - 1) | |
| emb = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) * -emb) | |
| emb = timesteps.float()[:, None] * emb[None, :] | |
| emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) | |
| if embedding_dim % 2 == 1: # zero pad | |
| emb = F.pad(emb, (0, 1), mode='constant') | |
| assert emb.shape == (timesteps.shape[0], embedding_dim) | |
| return emb | |
| class GaussianFourierProjection(nn.Module): | |
| """Gaussian Fourier embeddings for noise levels. | |
| from https://github.com/yang-song/score_sde_pytorch/blob/1618ddea340f3e4a2ed7852a0694a809775cf8d0/models/layerspp.py#L32 | |
| """ | |
| def __init__(self, embedding_size=256, scale=1.0): | |
| super().__init__() | |
| self.W = nn.Parameter(torch.randn(embedding_size//2) * scale, requires_grad=False) | |
| def forward(self, x): | |
| x_proj = x[:, None] * self.W[None, :] * 2 * np.pi | |
| emb = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1) | |
| return emb | |
| def get_timestep_embedding(embedding_type, embedding_dim, embedding_scale=10000): | |
| if embedding_type == 'sinusoidal': | |
| emb_func = (lambda x : sinusoidal_embedding(embedding_scale * x, embedding_dim)) | |
| elif embedding_type == 'fourier': | |
| emb_func = GaussianFourierProjection(embedding_size=embedding_dim, scale=embedding_scale) | |
| else: | |
| raise NotImplemented | |
| return emb_func | |
| def get_t_schedule(inference_steps): | |
| return np.linspace(1, 0, inference_steps + 1)[:-1] | |
| def set_time(complex_graphs, t_tr, t_rot, t_tor, batchsize, all_atoms, device): | |
| complex_graphs['ligand'].node_t = { | |
| 'tr': t_tr * torch.ones(complex_graphs['ligand'].num_nodes).to(device), | |
| 'rot': t_rot * torch.ones(complex_graphs['ligand'].num_nodes).to(device), | |
| 'tor': t_tor * torch.ones(complex_graphs['ligand'].num_nodes).to(device)} | |
| complex_graphs['receptor'].node_t = { | |
| 'tr': t_tr * torch.ones(complex_graphs['receptor'].num_nodes).to(device), | |
| 'rot': t_rot * torch.ones(complex_graphs['receptor'].num_nodes).to(device), | |
| 'tor': t_tor * torch.ones(complex_graphs['receptor'].num_nodes).to(device)} | |
| complex_graphs.complex_t = {'tr': t_tr * torch.ones(batchsize).to(device), | |
| 'rot': t_rot * torch.ones(batchsize).to(device), | |
| 'tor': t_tor * torch.ones(batchsize).to(device)} | |
| if all_atoms: | |
| complex_graphs['atom'].node_t = { | |
| 'tr': t_tr * torch.ones(complex_graphs['atom'].num_nodes).to(device), | |
| 'rot': t_rot * torch.ones(complex_graphs['atom'].num_nodes).to(device), | |
| 'tor': t_tor * torch.ones(complex_graphs['atom'].num_nodes).to(device)} |