File size: 2,965 Bytes
b2cb4a0 | 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 | #
# For licensing see accompanying LICENSE file.
# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
#
import mlx.core as mx
from tqdm import tqdm
from einops.array_api import repeat
from onescience.utils.simplefold.mlx_utils import center_random_augmentation
def logspace(start, end, steps, base=10.0, dtype=mx.float32):
# create a linear space between start and end
lin = mx.linspace(start, end, steps, dtype=dtype)
# raise base to that power
return mx.power(mx.array(base, dtype=dtype), lin)
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 - logspace(-2, 0, steps=self.num_timesteps + 1)[::-1, ...]
t = t - mx.min(t)
t = t / mx.max(t)
self.steps = mx.clip(t, a_min=self.t_start, a_max=1.0)
else:
self.steps = mx.linspace(self.t_start, 1.0, num=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
def euler_maruyama_step(
self,
model_fn,
flow,
y,
t,
t_next,
batch,
):
dt = t_next - t
eps = mx.random.normal(y.shape)
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 + mx.sqrt(2.0 * dt * diff_coeff * self.tau) * eps
return y_sample
def sample(self, model_fn, flow, noise, batch):
sampling_timesteps = self.num_timesteps
steps = self.steps
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,
)
mx.eval(y_sampled)
return {"denoised_coords": y_sampled}
|