File size: 2,783 Bytes
3a8e4d3 | 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 | #
# For licensing see accompanying LICENSE file.
# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
#
import torch
from tqdm import tqdm
from einops import repeat
from onescience.utils.simplefold.boltz_utils import center_random_augmentation
class EMSampler():
"""
A Euler-Maruyama solver for SDEs.
"""
def __init__(
self,
num_timesteps=500,
t_start=1e-4,
tau=0.3,
log_timesteps=False,
w_cutoff=0.99,
):
self.num_timesteps = num_timesteps
self.log_timesteps = log_timesteps
self.t_start = t_start
self.tau = tau
self.w_cutoff = w_cutoff
if self.log_timesteps:
t = 1.0 - torch.logspace(-2, 0, self.num_timesteps + 1).flip(0)
t = t - torch.min(t)
t = t / torch.max(t)
self.steps = t.clamp(min=self.t_start, max=1.0)
else:
self.steps = torch.linspace(
self.t_start, 1.0, steps=self.num_timesteps + 1
)
def diffusion_coefficient(self, t, eps=0.01):
# determine diffusion coefficient
w = (1.0 - t) / (t + eps)
if t >= self.w_cutoff:
w = 0.0
return w
@torch.no_grad()
def euler_maruyama_step(
self,
model_fn,
flow,
y,
t,
t_next,
batch,
):
dt = t_next - t
eps = torch.randn_like(y).to(y)
y = center_random_augmentation(
y,
batch["atom_pad_mask"],
augmentation=False,
centering=True,
)
batched_t = repeat(t, " -> b", b=y.shape[0])
velocity = model_fn(
noised_pos=y,
t=batched_t,
feats=batch,
)['predict_velocity']
score = flow.compute_score_from_velocity(velocity, y, t)
diff_coeff = self.diffusion_coefficient(t)
drift = velocity + diff_coeff * score
mean_y = y + drift * dt
y_sample = mean_y + torch.sqrt(2.0 * dt * diff_coeff * self.tau) * eps
return y_sample
@torch.no_grad()
def sample(self, model_fn, flow, noise, batch):
sampling_timesteps = self.num_timesteps
steps = self.steps.to(noise.device)
y_sampled = noise
feats = batch
for i in tqdm(
range(sampling_timesteps),
desc="Sampling",
total=sampling_timesteps,
):
t = steps[i]
t_next = steps[i + 1]
y_sampled = self.euler_maruyama_step(
model_fn,
flow,
y_sampled,
t,
t_next,
feats,
)
return {
"denoised_coords": y_sampled
}
|