| |
| |
| |
| |
|
|
| 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): |
| |
| 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 |
| } |
|
|
