Spaces:
Sleeping
Sleeping
File size: 2,538 Bytes
4a3f787 |
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 |
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_)
score_ = grad(x, sigma[:, None], N=100) / p_
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]
|