| 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 |
| SIGMA_MIN, SIGMA_MAX, SIGMA_N = 3e-3, 2, 5000 |
|
|
| 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 |
|
|
| tmp_path = os.environ['precomputed_arrays'] |
| print(tmp_path) |
| if os.path.exists(os.path.join(tmp_path,'.p.npy')): |
| p_ = np.load(os.path.join(tmp_path,'.p.npy')) |
| score_ = np.load(os.path.join(tmp_path,'.score.npy')) |
| else: |
| print("Precomputing and saving to cache torus distribution table") |
| p_ = p(x, sigma[:, None], N=100) |
| np.save(os.path.join(tmp_path,'.p.npy'), p_) |
|
|
| score_ = grad(x, sigma[:, None], N=100) / p_ |
| np.save(os.path.join(tmp_path,'.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] |
|
|
