File size: 4,682 Bytes
ca7299e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """Inspired by https://github.com/jasonkyuyim/se3_diffusion/blob/master/data/r3_diffuser.py"""
from math import sqrt
import torch
from src.utils.tensor_utils import inflate_array_like
class R3Diffuser:
"""VPSDE diffusion module."""
def __init__(
self,
min_b: float = 0.1,
max_b: float = 20.0,
coordinate_scaling: float = 1.0,
):
self.min_b = min_b
self.max_b = max_b
self.coordinate_scaling = coordinate_scaling
def scale(self, x):
return x * self.coordinate_scaling
def unscale(self, x):
return x / self.coordinate_scaling
def b_t(self, t: torch.Tensor):
if torch.any(t < 0) or torch.any(t > 1):
raise ValueError(f'Invalid t={t}')
return self.min_b + t * (self.max_b - self.min_b)
def diffusion_coef(self, t):
return torch.sqrt(self.b_t(t))
def drift_coef(self, x, t):
return -0.5 * self.b_t(t) * x
def sample_prior(self, shape, device=None):
return torch.randn(size=shape, device=device)
def marginal_b_t(self, t):
return t*self.min_b + 0.5*(t**2)*(self.max_b-self.min_b)
def calc_trans_0(self, score_t, x_t, t):
beta_t = self.marginal_b_t(t)
beta_t = beta_t[..., None, None]
cond_var = 1 - torch.exp(-beta_t)
return (score_t * cond_var + x_t) / torch.exp(-0.5*beta_t)
def forward_marginal(
self,
x_0: torch.Tensor,
t: torch.Tensor
):
"""Samples marginal p(x(t) | x(0)).
Args:
x_0: [..., n, 3] initial positions in Angstroms.
t: continuous time in [0, 1].
Returns:
x_t: [..., n, 3] positions at time t in Angstroms.
score_t: [..., n, 3] score at time t in scaled Angstroms.
"""
t = inflate_array_like(t, x_0)
x_0 = self.scale(x_0)
loc = torch.exp(-0.5 * self.marginal_b_t(t)) * x_0
scale = torch.sqrt(1 - torch.exp(-self.marginal_b_t(t)))
z = torch.randn_like(x_0)
x_t = z * scale + loc
score_t = self.score(x_t, x_0, t)
x_t = self.unscale(x_t)
return x_t, score_t
def score_scaling(self, t: torch.Tensor):
return 1.0 / torch.sqrt(self.conditional_var(t))
def reverse(
self,
x_t: torch.Tensor,
score_t: torch.Tensor,
t: torch.Tensor,
dt: float,
mask: torch.Tensor = None,
center: bool = True,
noise_scale: float = 1.0,
probability_flow: bool = True,
):
"""Simulates the reverse SDE for 1 step
Args:
x_t: [..., 3] current positions at time t in angstroms.
score_t: [..., 3] rotation score at time t.
t: continuous time in [0, 1].
dt: continuous step size in [0, 1].
mask: True indicates which residues to diffuse.
probability_flow: whether to use probability flow ODE.
Returns:
[..., 3] positions at next step t-1.
"""
t = inflate_array_like(t, x_t)
x_t = self.scale(x_t)
f_t = self.drift_coef(x_t, t)
g_t = self.diffusion_coef(t)
z = noise_scale * torch.randn_like(score_t)
rev_drift = (f_t - g_t ** 2 * score_t) * dt * (0.5 if probability_flow else 1.)
rev_diffusion = 0. if probability_flow else (g_t * sqrt(dt) * z)
perturb = rev_drift + rev_diffusion
if mask is not None:
perturb *= mask[..., None]
else:
mask = torch.ones_like(x_t[..., 0])
x_t_1 = x_t - perturb # reverse in time
if center:
com = torch.sum(x_t_1, dim=-2) / torch.sum(mask, dim=-1)[..., None] # reduce length dim
x_t_1 -= com[..., None, :]
x_t_1 = self.unscale(x_t_1)
return x_t_1
def conditional_var(self, t, use_torch=False):
"""Conditional variance of p(xt|x0).
Var[x_t|x_0] = conditional_var(t) * I
"""
return 1.0 - torch.exp(-self.marginal_b_t(t))
def score(self, x_t, x_0, t, scale=False):
t = inflate_array_like(t, x_t)
if scale:
x_t, x_0 = self.scale(x_t), self.scale(x_0)
return -(x_t - torch.exp(-0.5 * self.marginal_b_t(t)) * x_0) / self.conditional_var(t)
def distribution(self, x_t, score_t, t, mask, dt):
x_t = self.scale(x_t)
f_t = self.drift_coef(x_t, t)
g_t = self.diffusion_coef(t)
std = g_t * sqrt(dt)
mu = x_t - (f_t - g_t**2 * score_t) * dt
if mask is not None:
mu *= mask[..., None]
return mu, std |