mirror of
https://github.com/gcorso/DiffDock.git
synced 2026-06-04 18:04:23 +08:00
561f70a897
* Ensure we calculate rotatable bonds on the version of the ligand with no hydrogens. Also fix spelling of rotable -> rotatable. Closes GH-220 (@Nobody-Zhang) * Vectorize SO3 calculations. Closes PR GH-218 (@tornikeo) * Pin pytorch-lightning version. Closes GH-193 (@mikael-h-christensen) * Guard against divide by zero in torus.py. Closes GH-161 (@amorehead) * Update e3nn version to 0.5.1. Closes GH-155 (@amorehead) * Add a little more info on docker container to README.md
84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
import numpy as np
|
|
import tqdm
|
|
import os
|
|
|
|
"""
|
|
Preprocessing for the SO(2)/torus sampling and score computations, truncated infinite series are computed and then
|
|
cached to memory, therefore the precomputation is only run the first time the repository is run on a machine
|
|
"""
|
|
|
|
|
|
def p(x, sigma, N=10):
|
|
p_ = 0
|
|
for i in tqdm.trange(-N, N + 1):
|
|
p_ += np.exp(-(x + 2 * np.pi * i) ** 2 / 2 / sigma ** 2)
|
|
return p_
|
|
|
|
|
|
def grad(x, sigma, N=10):
|
|
p_ = 0
|
|
for i in tqdm.trange(-N, N + 1):
|
|
p_ += (x + 2 * np.pi * i) / sigma ** 2 * np.exp(-(x + 2 * np.pi * i) ** 2 / 2 / sigma ** 2)
|
|
return p_
|
|
|
|
|
|
X_MIN, X_N = 1e-5, 5000 # relative to pi
|
|
SIGMA_MIN, SIGMA_MAX, SIGMA_N = 3e-3, 2, 5000 # relative to pi
|
|
|
|
x = 10 ** np.linspace(np.log10(X_MIN), 0, X_N + 1) * np.pi
|
|
sigma = 10 ** np.linspace(np.log10(SIGMA_MIN), np.log10(SIGMA_MAX), SIGMA_N + 1) * np.pi
|
|
|
|
if os.path.exists('.p.npy'):
|
|
p_ = np.load('.p.npy')
|
|
score_ = np.load('.score.npy')
|
|
else:
|
|
p_ = p(x, sigma[:, None], N=100)
|
|
np.save('.p.npy', p_)
|
|
|
|
eps = np.finfo(p_.dtype).eps
|
|
score_ = grad(x, sigma[:, None], N=100) / (p_ + eps)
|
|
np.save('.score.npy', score_)
|
|
|
|
|
|
def score(x, sigma):
|
|
x = (x + np.pi) % (2 * np.pi) - np.pi
|
|
sign = np.sign(x)
|
|
x = np.log(np.abs(x) / np.pi)
|
|
x = (x - np.log(X_MIN)) / (0 - np.log(X_MIN)) * X_N
|
|
x = np.round(np.clip(x, 0, X_N)).astype(int)
|
|
sigma = np.log(sigma / np.pi)
|
|
sigma = (sigma - np.log(SIGMA_MIN)) / (np.log(SIGMA_MAX) - np.log(SIGMA_MIN)) * SIGMA_N
|
|
sigma = np.round(np.clip(sigma, 0, SIGMA_N)).astype(int)
|
|
return -sign * score_[sigma, x]
|
|
|
|
|
|
def p(x, sigma):
|
|
x = (x + np.pi) % (2 * np.pi) - np.pi
|
|
x = np.log(np.abs(x) / np.pi)
|
|
x = (x - np.log(X_MIN)) / (0 - np.log(X_MIN)) * X_N
|
|
x = np.round(np.clip(x, 0, X_N)).astype(int)
|
|
sigma = np.log(sigma / np.pi)
|
|
sigma = (sigma - np.log(SIGMA_MIN)) / (np.log(SIGMA_MAX) - np.log(SIGMA_MIN)) * SIGMA_N
|
|
sigma = np.round(np.clip(sigma, 0, SIGMA_N)).astype(int)
|
|
return p_[sigma, x]
|
|
|
|
|
|
def sample(sigma):
|
|
out = sigma * np.random.randn(*sigma.shape)
|
|
out = (out + np.pi) % (2 * np.pi) - np.pi
|
|
return out
|
|
|
|
|
|
score_norm_ = score(
|
|
sample(sigma[None].repeat(10000, 0).flatten()),
|
|
sigma[None].repeat(10000, 0).flatten()
|
|
).reshape(10000, -1)
|
|
score_norm_ = (score_norm_ ** 2).mean(0)
|
|
|
|
|
|
def score_norm(sigma):
|
|
sigma = np.log(sigma / np.pi)
|
|
sigma = (sigma - np.log(SIGMA_MIN)) / (np.log(SIGMA_MAX) - np.log(SIGMA_MIN)) * SIGMA_N
|
|
sigma = np.round(np.clip(sigma, 0, SIGMA_N)).astype(int)
|
|
return score_norm_[sigma]
|