Spaces:
Sleeping
Sleeping
File size: 3,589 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 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import os
import numpy as np
import torch
from scipy.spatial.transform import Rotation
MIN_EPS, MAX_EPS, N_EPS = 0.01, 2, 1000
X_N = 2000
"""
Preprocessing for the SO(3) 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
"""
omegas = np.linspace(0, np.pi, X_N + 1)[1:]
def _compose(r1, r2): # R1 @ R2 but for Euler vecs
return Rotation.from_matrix(Rotation.from_rotvec(r1).as_matrix() @ Rotation.from_rotvec(r2).as_matrix()).as_rotvec()
def _expansion(omega, eps, L=2000): # the summation term only
p = 0
for l in range(L):
p += (2 * l + 1) * np.exp(-l * (l + 1) * eps**2) * np.sin(omega * (l + 1 / 2)) / np.sin(omega / 2)
return p
def _density(expansion, omega, marginal=True): # if marginal, density over [0, pi], else over SO(3)
if marginal:
return expansion * (1 - np.cos(omega)) / np.pi
else:
return expansion / 8 / np.pi ** 2 # the constant factor doesn't affect any actual calculations though
def _score(exp, omega, eps, L=2000): # score of density over SO(3)
dSigma = 0
for l in range(L):
hi = np.sin(omega * (l + 1 / 2))
dhi = (l + 1 / 2) * np.cos(omega * (l + 1 / 2))
lo = np.sin(omega / 2)
dlo = 1 / 2 * np.cos(omega / 2)
dSigma += (2 * l + 1) * np.exp(-l * (l + 1) * eps**2) * (lo * dhi - hi * dlo) / lo ** 2
return dSigma / exp
if os.path.exists('.so3_omegas_array2.npy'):
_omegas_array = np.load('.so3_omegas_array2.npy')
_cdf_vals = np.load('.so3_cdf_vals2.npy')
_score_norms = np.load('.so3_score_norms2.npy')
_exp_score_norms = np.load('.so3_exp_score_norms2.npy')
else:
_eps_array = 10 ** np.linspace(np.log10(MIN_EPS), np.log10(MAX_EPS), N_EPS)
_omegas_array = np.linspace(0, np.pi, X_N + 1)[1:]
_exp_vals = np.asarray([_expansion(_omegas_array, eps) for eps in _eps_array])
_pdf_vals = np.asarray([_density(_exp, _omegas_array, marginal=True) for _exp in _exp_vals])
_cdf_vals = np.asarray([_pdf.cumsum() / X_N * np.pi for _pdf in _pdf_vals])
_score_norms = np.asarray([_score(_exp_vals[i], _omegas_array, _eps_array[i]) for i in range(len(_eps_array))])
_exp_score_norms = np.sqrt(np.sum(_score_norms**2 * _pdf_vals, axis=1) / np.sum(_pdf_vals, axis=1) / np.pi)
np.save('.so3_omegas_array2.npy', _omegas_array)
np.save('.so3_cdf_vals2.npy', _cdf_vals)
np.save('.so3_score_norms2.npy', _score_norms)
np.save('.so3_exp_score_norms2.npy', _exp_score_norms)
def sample(eps):
eps_idx = (np.log10(eps) - np.log10(MIN_EPS)) / (np.log10(MAX_EPS) - np.log10(MIN_EPS)) * N_EPS
eps_idx = np.clip(np.around(eps_idx).astype(int), a_min=0, a_max=N_EPS - 1)
x = np.random.rand()
return np.interp(x, _cdf_vals[eps_idx], _omegas_array)
def sample_vec(eps):
x = np.random.randn(3)
x /= np.linalg.norm(x)
return x * sample(eps)
def score_vec(eps, vec):
eps_idx = (np.log10(eps) - np.log10(MIN_EPS)) / (np.log10(MAX_EPS) - np.log10(MIN_EPS)) * N_EPS
eps_idx = np.clip(np.around(eps_idx).astype(int), a_min=0, a_max=N_EPS - 1)
om = np.linalg.norm(vec)
return np.interp(om, _omegas_array, _score_norms[eps_idx]) * vec / om
def score_norm(eps):
eps = eps.numpy()
eps_idx = (np.log10(eps) - np.log10(MIN_EPS)) / (np.log10(MAX_EPS) - np.log10(MIN_EPS)) * N_EPS
eps_idx = np.clip(np.around(eps_idx).astype(int), a_min=0, a_max=N_EPS-1)
return torch.from_numpy(_exp_score_norms[eps_idx]).float()
|