def test_S2FFT_NFFT():
"""
Testing S2FFT NFFT
"""
b = 8
convention = 'Gauss-Legendre'
#convention = 'Clenshaw-Curtis'
x = S2.meshgrid(b=b, grid_type=convention)
print(x[0].shape, x[1].shape)
x = np.c_[x[0][..., None], x[1][..., None]]#.reshape(-1, 2)
print(x.shape)
x = x.reshape(-1, 2)
w = S2.quadrature_weights(b=b, grid_type=convention).flatten()
F = S2FFT_NFFT(L_max=b, x=x, w=w)
for l in range(0, b):
for m in range(-l, l + 1):
#l = b; m = b
f = sh(l, m, x[..., 0], x[..., 1], field='real', normalization='quantum', condon_shortley=True)
#f2 = np.random.randn(*f.shape)
print(f)
f_hat = F.analyze(f)
print(np.round(f_hat, 3))
f_reconst = F.synthesize(f_hat)
#print np.round(f, 3)
print(np.round(f_reconst, 3))
#print np.round(f/f_reconst, 3)
print(np.abs(f-f_reconst).sum())
assert np.isclose(np.abs(f-f_reconst).sum(), 0.)
print(np.round(f_hat, 3))
assert np.isclose(f_hat[l ** 2 + l + m], 1.)
#assert False
def make_sgrid(b):
theta, phi = S2.meshgrid(b=b, grid_type='Driscoll-Healy')
sgrid = S2.change_coordinates(np.c_[theta[..., None], phi[..., None]],
p_from='S',
p_to='C')
sgrid = sgrid.reshape((-1, 3))
return sgrid
def make_sgrid(b):
from lie_learn.spaces import S2
theta, phi = S2.meshgrid(b=b, grid_type='SOFT')
sgrid = S2.change_coordinates(np.c_[theta[..., None], phi[..., None]],
p_from='S',
p_to='C')
sgrid = sgrid.reshape((-1, 3))
return (theta, phi), sgrid
def make_sgrid(b, alpha, beta, gamma, grid_type):
theta, phi = S2.meshgrid(b=b, grid_type=grid_type)
sgrid = S2.change_coordinates(np.c_[theta[..., None], phi[..., None]],
p_from='S',
p_to='C')
sgrid = sgrid.reshape((-1, 3))
R = mesh_op.rotmat(alpha, beta, gamma, hom_coord=False)
sgrid = np.einsum('ij,nj->ni', R, sgrid)
return sgrid
def make_sgrid(b, alpha, beta, gamma):
from lie_learn.spaces import S2
theta, phi = S2.meshgrid(b=b, grid_type='SOFT')
sgrid = S2.change_coordinates(np.c_[theta[..., None], phi[..., None]], p_from='S', p_to='C')
sgrid = sgrid.reshape((-1, 3))
R = rotmat(alpha, beta, gamma, hom_coord=False)
sgrid = np.einsum('ij,nj->ni', R, sgrid)
return sgrid
def get_projection_grid(b, grid_type="Driscoll-Healy"):
"""
returns the spherical grid in euclidean
coordinates, where the sphere's center is moved
to (0, 0, 1)
"""
theta, phi = S2.meshgrid(b=b, grid_type=grid_type)
grid = S2.change_coordinates(np.c_[theta[..., None], phi[..., None]],
p_from='S',
p_to='C')
grid = grid.reshape((-1, 3)).astype(np.float32)
return grid
def test_spherical_quadrature():
"""
Testing spherical quadrature rule versus numerical integration.
"""
b = 8 # 10
# Create grids on the sphere
x_gl = S2.meshgrid(b=b, grid_type='Gauss-Legendre')
x_cc = S2.meshgrid(b=b, grid_type='Clenshaw-Curtis')
x_soft = S2.meshgrid(b=b, grid_type='SOFT')
x_gl = np.c_[x_gl[0][..., None], x_gl[1][..., None]]
x_cc = np.c_[x_cc[0][..., None], x_cc[1][..., None]]
x_soft = np.c_[x_soft[0][..., None], x_soft[1][..., None]]
# Compute quadrature weights
w_gl = S2.quadrature_weights(b=b, grid_type='Gauss-Legendre')
w_cc = S2.quadrature_weights(b=b, grid_type='Clenshaw-Curtis')
w_soft = S2.quadrature_weights(b=b, grid_type='SOFT')
# Define a polynomial function, to be evaluated at one point or at an array of points
def f1a(xs):
xc = S2.change_coordinates(coords=xs, p_from='S', p_to='C')
return xc[..., 0]**2 * xc[..., 1] - 1.4 * xc[..., 2] * xc[
..., 1]**3 + xc[..., 1] - xc[..., 2]**2 + 2.
def f1(theta, phi):
xs = np.array([theta, phi])
return f1a(xs)
# Obtain the "true" value of the integral of the function over the sphere, using scipy's numerical integration
# routines
i1 = S2.integrate(f1, normalize=False)
# Compute the integral using the quadrature formulae
# i1_gl_w = (w_gl * f1a(x_gl)).sum()
i1_gl_w = S2.integrate_quad(f1a(x_gl),
grid_type='Gauss-Legendre',
normalize=False,
w=w_gl)
print(i1_gl_w, i1, 'diff:', np.abs(i1_gl_w - i1))
assert np.isclose(np.abs(i1_gl_w - i1), 0.0)
# i1_cc_w = (w_cc * f1a(x_cc)).sum()
i1_cc_w = S2.integrate_quad(f1a(x_cc),
grid_type='Clenshaw-Curtis',
normalize=False,
w=w_cc)
print(i1_cc_w, i1, 'diff:', np.abs(i1_cc_w - i1))
assert np.isclose(np.abs(i1_cc_w - i1), 0.0)
i1_soft_w = (w_soft * f1a(x_soft)).sum()
print(i1_soft_w, i1, 'diff:', np.abs(i1_soft_w - i1))
print(i1_soft_w)
print(i1)
def get_projection_grid(b, grid_type="Driscoll-Healy"): theta, phi = S2.meshgrid(b=b, grid_type=grid_type) x_ = np.sin(theta) * np.cos(phi) y_ = np.sin(theta) * np.sin(phi) z_ = np.cos(theta) #pdb.set_trace() return x_, y_, z_
def naive_conv(l1=1, m1=1, l2=1, m2=1, g_parameterization='EA313'):
f1 = lambda t, p: sh(l=l1,
m=m1,
theta=t,
phi=p,
field='real',
normalization='quantum',
condon_shortley=True)
f2 = lambda t, p: sh(l=l2,
m=m2,
theta=t,
phi=p,
field='real',
normalization='quantum',
condon_shortley=True)
theta, phi = S2.meshgrid(b=3, grid_type='Gauss-Legendre')
f1_grid = f1(theta, phi)
f2_grid = f2(theta, phi)
alpha, beta, gamma = S3.meshgrid(b=3,
grid_type='SOFT') # TODO check convention
f12_grid = np.zeros_like(alpha)
for i in range(alpha.shape[0]):
for j in range(alpha.shape[1]):
for k in range(alpha.shape[2]):
f12_grid[i, j, k] = naive_S2_conv_v2(f1, f2, alpha[i, j, k],
beta[i, j, k], gamma[i, j,
k],
g_parameterization)
print(i, j, k, f12_grid[i, j, k])
return f1_grid, f2_grid, f12_grid
def get_grid(b, radius, grid_type="Driscoll-Healy"):
"""
returns the spherical grid in euclidean coordinates, which, to be specify,
for each image in range(train_size):
for each point in range(num_points):
generate the 2b * 2b S2 points , each is (x, y, z)
therefore returns tensor (train_size * num_points, 2b * 2b, 3)
:param b: the number of grids on the sphere
:param radius: the radius of each sphere
:param grid_type: "Driscoll-Healy"
:return: tensor (batch_size, num_points, 4 * b * b, 3)
"""
# theta in shape (2b, 2b), range [0, pi]; phi range [0, 2 * pi]
theta, phi = S2.meshgrid(b=b, grid_type=grid_type)
theta = torch.from_numpy(theta).cuda()
phi = torch.from_numpy(phi).cuda()
x_ = radius * torch.sin(theta) * torch.cos(phi)
"""x will be reshaped to have one dimension of 1, then can broadcast
look this link for more information: https://pytorch.org/docs/stable/notes/broadcasting.html
"""
x = x_.reshape((1, 4 * b * b)) # tensor (1, 4 * b * b)
# same for y and z
y_ = radius * torch.sin(theta) * torch.sin(phi)
y = y_.reshape((1, 4 * b * b))
z_ = radius * torch.cos(theta)
z = z_.reshape((1, 4 * b * b))
grid = torch.cat((x, y, z), dim=0) # (3, 4 * b * b)
assert grid.shape == torch.Size([3, 4 * b * b])
# grid = grid.reshape((1, 4 * b * b, 3))
return grid
def inverse_render_model(points: np.ndarray, sgrid: np.ndarray, lib):
# wait for implementing
#print("InverseRenderModel")
try:
x = points[:, 0]
except IndexError:
print(points.shape[0], points.shape[1])
y = points[:, 1]
z = points[:, 2]
# Fistly, we get the centroid, then translate the points
centroid_x = np.sum(x) / points.shape[0]
centroid_y = np.sum(y) / points.shape[0]
centroid_z = np.sum(z) / points.shape[0]
centroid = np.array([centroid_x, centroid_y, centroid_z])
points = points.astype(np.float)
points -= centroid
# After normalization, compute the distance between the sphere and points
radius = np.sqrt(points[:, 0]**2 + points[:, 1]**2 + points[:, 2]**2)
# dist = 1 - (1 / np.max(radius)) * radius
# Projection
from lie_learn.spaces import S2
radius = np.repeat(radius, 3).reshape(-1, 3)
points_on_sphere = points / radius
# ssgrid = sgrid.reshape(-1, 3)
# phi, theta = S2.change_coordinates(ssgrid, p_from='C', p_to='S')
out = S2.change_coordinates(points_on_sphere, p_from='C', p_to='S')
phi = out[..., 0]
theta = out[..., 1]
phi = phi
theta = theta % (np.pi * 2)
# Interpolate
b = sgrid.shape[0] / 2 # bandwidth
# By computing the m,n, we can find
# the neighbours on the sphere
m = np.trunc((phi - np.pi / (4 * b)) / (np.pi / (2 * b)))
m = m.astype(int)
n = np.trunc(theta / (np.pi / b))
n = n.astype(int)
dist_im, center_grid, east_grid, south_grid, southeast_grid = interpolate(
m=m,
n=n,
sgrid=sgrid,
points_on_sphere=points_on_sphere,
radius=radius)
dot_img, cross_img = angle(m=m, n=n, sgrid=sgrid, dist_im=dist_im, lib=lib)
#dist_im=dist_im.reshape(-1,1)
#wait for validation
im = np.stack((dist_im, dot_img, cross_img), axis=0)
return im
def get_projection_grid(b, grid_type="Driscoll-Healy"):
''' returns the spherical grid in euclidean
coordinates, where the sphere's center is moved
to (0, 0, 1)'''
theta, phi = S2.meshgrid(b=b, grid_type=grid_type)
x_ = np.sin(theta) * np.cos(phi)
y_ = np.sin(theta) * np.sin(phi)
z_ = np.cos(theta)
return x_, y_, z_
def make_sgrid_(b):
theta = np.linspace(0, m.pi, num=b)
phi = np.linspace(0, 2 * m.pi, num=b)
theta_m, phi_m = np.meshgrid(theta, phi)
sgrid = S2.change_coordinates(np.c_[theta_m[..., None], phi_m[..., None]],
p_from='S',
p_to='C')
sgrid = sgrid.reshape((-1, 3))
return sgrid
def __init__(self,
L_max,
grid_type='Gauss-Legendre',
field='real',
normalization='quantum',
condon_shortley='cs'):
super().__init__()
self.b = L_max + 1
# Compute a grid of spatial sampling points and associated quadrature weights
beta, alpha = S2.meshgrid(b=self.b, grid_type=grid_type)
self.w = S2.quadrature_weights(b=self.b, grid_type=grid_type)
self.spatial_grid_shape = beta.shape
self.num_spatial_points = beta.size
# Determine for which degree and order we want the spherical harmonics
irreps = np.arange(
self.b
) # TODO find out upper limit for exact integration for each grid type
ls = [[ls] * (2 * ls + 1) for ls in irreps]
ls = np.array([ll for sublist in ls
for ll in sublist]) # 0, 1, 1, 1, 2, 2, 2, 2, 2, ...
ms = [list(range(-ls, ls + 1)) for ls in irreps]
ms = np.array([mm for sublist in ms
for mm in sublist]) # 0, -1, 0, 1, -2, -1, 0, 1, 2, ...
self.num_spectral_points = ms.size # This equals sum_{l=0}^{b-1} 2l+1 = b^2
# In one shot, sample the spherical harmonics at all spectral (l, m) and spatial (beta, alpha) coordinates
self.Y = sh(ls[None, None, :],
ms[None, None, :],
beta[:, :, None],
alpha[:, :, None],
field=field,
normalization=normalization,
condon_shortley=condon_shortley == 'cs')
# Convert to a matrix
self.Ymat = self.Y.reshape(self.num_spatial_points,
self.num_spectral_points)
def test_S2_FT_Naive():
L_max = 6
for grid_type in ('Gauss-Legendre', 'Clenshaw-Curtis'):
theta, phi = S2.meshgrid(b=L_max + 1, grid_type=grid_type)
for field in ('real', 'complex'):
for normalization in (
'quantum', 'seismology'
): # TODO Others should work but are not normalized
for condon_shortley in ('cs', 'nocs'):
fft = S2_FT_Naive(L_max,
grid_type=grid_type,
field=field,
normalization=normalization,
condon_shortley=condon_shortley)
for l in range(L_max):
for m in range(-l, l + 1):
y_true = sh(
l,
m,
theta,
phi,
field=field,
normalization=normalization,
condon_shortley=condon_shortley == 'cs')
y_hat = fft.analyze(y_true)
# The flat index for (l, m) is l^2 + l + m
# Before the harmonics of degree l, there are this many harmonics:
# sum_{i=0}^{l-1} 2i+1 = l^2
# There are 2l+1 harmonics of degree l, with order m=0 at the center,
# so the m-th harmonic of degree is at l + m within the block of degree l.
y_hat_true = np.zeros_like(y_hat)
y_hat_true[l**2 + l + m] = 1
y = fft.synthesize(y_hat_true)
diff = np.sum(np.abs(y_hat - y_hat_true))
print(grid_type, field, normalization,
condon_shortley, l, m, diff)
assert np.isclose(diff, 0.)
diff = np.sum(np.abs(y - y_true))
print(grid_type, field, normalization,
condon_shortley, l, m, diff)
assert np.isclose(diff, 0.)
def naive_S2_conv_v2(f1, f2, alpha, beta, gamma, g_parameterization='EA323'):
"""
Compute int_S^2 f1(x) f2(g^{-1} x)* dx,
where x = (theta, phi) is a point on the sphere S^2,
and g = (alpha, beta, gamma) is a point in SO(3) in Euler angle parameterization
:param f1, f2: functions to be convolved
:param alpha, beta, gamma: the rotation at which to evaluate the result of convolution
:return:
"""
theta, phi = S2.meshgrid(b=3, grid_type='Gauss-Legendre')
w = S2.quadrature_weights(b=3, grid_type='Gauss-Legendre')
print(theta.shape, phi.shape)
s2_coords = np.c_[theta[..., None], phi[..., None]]
print(s2_coords.shape)
r3_coords = np.c_[theta[..., None], phi[..., None],
np.ones_like(theta)[..., None]]
# g_inv = SO3.invert((alpha, beta, gamma), parameterization=g_parameterization)
# g_inv = (-gamma, -beta, -alpha)
g_inv = (alpha, beta, gamma) # wrong
ginvx = SO3.transform_r3(g=g_inv,
x=r3_coords,
g_parameterization=g_parameterization,
x_parameterization='S')
print(ginvx.shape)
g_inv_theta = ginvx[..., 0]
g_inv_phi = ginvx[..., 1]
g_inv_r = ginvx[..., 2]
print(g_inv_theta, g_inv_phi, g_inv_r)
f1_grid = f1(theta, phi)
f2_grid = f2(g_inv_theta, g_inv_phi)
print(f1_grid.shape, f2_grid.shape, w.shape)
return np.sum(f1_grid * f2_grid * w)
def get_grids(b,
num_grids,
base_radius=1,
center=[0, 0, 0],
grid_type="Driscoll-Healy"):
"""
:param b: the number of grids on the sphere
:param base_radius: the radius of each sphere
:param grid_type: "Driscoll-Healy"
:param num_grids: number of grids
:return: [(radius, tensor([2b, 2b, 3])) * num_grids]
"""
grids = list()
radiuses = [
round(i, 2)
for i in list(np.linspace(0, base_radius, num_grids + 1))[1:]
]
# Each grid has differet radius, the radiuses are distributed uniformly based on number
for radius in radiuses:
# theta in shape (2b, 2b), range [0, pi]; phi range [0, 2 * pi]
theta, phi = S2.meshgrid(b=b, grid_type=grid_type)
theta = torch.from_numpy(theta)
phi = torch.from_numpy(phi)
# x will be reshaped to have one dimension of 1, then can broadcast
# look this link for more information: https://pytorch.org/docs/stable/notes/broadcasting.html
x_ = radius * torch.sin(theta) * torch.cos(phi)
x = x_.reshape((1, 4 * b * b)) # tensor -> [1, 4 * b * b]
x = x + center[0]
y_ = radius * torch.sin(theta) * torch.sin(phi)
y = y_.reshape((1, 4 * b * b))
y = y + center[1]
z_ = radius * torch.cos(theta)
z = z_.reshape((1, 4 * b * b))
z = z + center[2]
grid = torch.cat((x, y, z), dim=0) # -> [3, 4b^2]
grid = grid.transpose(0, 1) # -> [4b^2, 3]
grid = grid.view(2 * b, 2 * b, 3) # -> [2b, 2b, 3]
grid = grid.float().cuda()
grids.append((radius, grid))
assert len(grids) == num_grids
return grids
def setup_legendre_transform(b):
"""
Compute a set of matrices containing coefficients to be used in a discrete Legendre transform.
The discrete Legendre transform of a data vector s[k] (k=0, ..., 2b-1) is defined as
s_hat(l, m) = sum_{k=0}^{2b-1} P_l^m(cos(beta_k)) s[k]
for l = 0, ..., b-1 and -l <= m <= l,
where P_l^m is the associated Legendre function of degree l and order m,
beta_k = ...
Computing Fourier Transforms and Convolutions on the 2-Sphere
J.R. Driscoll, D.M. Healy
FFTs for the 2-Sphere–Improvements and Variations
D.M. Healy, Jr., D.N. Rockmore, P.J. Kostelec, and S. Moore
:param b: bandwidth of the transform
:return: lt, an array of shape (N, 2b), containing samples of the Legendre functions,
where N is the number of spectral points for a signal of bandwidth b.
"""
dim = np.sum(np.arange(b) * 2 + 1)
lt = np.empty((2 * b, dim))
beta, _ = S2.linspace(b, grid_type='Driscoll-Healy')
sample_points = np.cos(beta)
# TODO move quadrature weight computation to S2.py
weights = [(1. / b) * np.sin(np.pi * j * 0.5 / b) * np.sum([
1. / (2 * l + 1) * np.sin((2 * l + 1) * np.pi * j * 0.5 / b)
for l in range(b)
]) for j in range(2 * b)]
weights = np.array(weights)
zeros = np.zeros_like(sample_points)
i = 0
for l in range(b):
for m in range(-l, l + 1):
# Z = np.sqrt(((2 * l + 1) * factorial(l - m)) / float(4 * np.pi * factorial(l + m))) * np.pi / 2
# lt[i, :] = lpmv(m, l, sample_points) * weights * Z
# The spherical harmonics code appears to be more stable than the (unnormalized) associated Legendre
# function code.
# (Note: the spherical harmonics evaluated at alpha=0 is the associated Legendre function))
lt[:,
i] = csh(l, m, beta, zeros,
normalization='seismology').real * weights * np.pi / 2
i += 1
return lt
def get_projection_grid(b, images, radius, grid_type="Driscoll-Healy"):
"""
returns the spherical grid in euclidean coordinates, which, to be specify,
for each image in range(train_size):
for each point in range(num_points):
generate the 2b * 2b S2 points , each is (x, y, z)
therefore returns tensor (train_size * num_points, 2b * 2b, 3)
:param b: the number of grids on the sphere
:param images: tensor (batch_size, num_points, 3)
:param radius: the radius of each sphere
:param grid_type: "Driscoll-Healy"
:return: tensor (batch_size, num_points, 4 * b * b, 3)
"""
assert type(images) == torch.Tensor
assert len(images.shape) == 3
assert images.shape[-1] == 3
batch_size = images.shape[0]
num_points = images.shape[1]
images = images.reshape((-1, 3)) # -> (B * 512, 3)
# theta in shape (2b, 2b), range [0, pi]; phi range [0, 2 * pi]
theta, phi = S2.meshgrid(b=b, grid_type=grid_type)
theta = torch.from_numpy(theta).cuda()
phi = torch.from_numpy(phi).cuda()
x_ = radius * torch.sin(theta) * torch.cos(phi)
"""x will be reshaped to have one dimension of 1, then can broadcast
look this link for more information: https://pytorch.org/docs/stable/notes/broadcasting.html
"""
x = x_.reshape((1, 4 * b * b)) # tensor (1, 4 * b * b)
px = images[:, 0].reshape((-1, 1)) # tensor (batch_size * 512, 1)
x = x + px # (batch_size * num_points, 4 * b * b)
# same for y and z
y_ = radius * torch.sin(theta) * torch.sin(phi)
y = y_.reshape((1, 4 * b * b))
py = images[:, 1].reshape((-1, 1))
y = y + py
z_ = radius * torch.cos(theta)
z = z_.reshape((1, 4 * b * b))
pz = images[:, 2].reshape((-1, 1))
z = z + pz
# give x, y, z extra dimension, so that it can concat by that dimension
x = torch.unsqueeze(x, 2) # (B * 512, 4 * b * b, 1)
y = torch.unsqueeze(y, 2)
z = torch.unsqueeze(z, 2)
grid = torch.cat((x, y, z), 2) # (B * 512, 4 * b * b, 3)
grid = grid.reshape((batch_size, num_points, 4 * b * b, 3))
return grid
def check_orthogonality(L_max=3,
grid_type='Gauss-Legendre',
field='real',
normalization='quantum',
condon_shortley=True):
theta, phi = S2.meshgrid(b=L_max + 1, grid_type=grid_type)
w = S2.quadrature_weights(b=L_max + 1, grid_type=grid_type)
for l in range(L_max):
for m in range(-l, l + 1):
for l2 in range(L_max):
for m2 in range(-l2, l2 + 1):
Ylm = sh(l, m, theta, phi, field, normalization,
condon_shortley)
Ylm2 = sh(l2, m2, theta, phi, field, normalization,
condon_shortley)
dot_numerical = S2.integrate_quad(Ylm * Ylm2.conj(),
grid_type=grid_type,
normalize=False,
w=w)
dot_numerical2 = S2.integrate(
lambda t, p: sh(l, m, t, p, field, normalization, condon_shortley) * \
sh(l2, m2, t, p, field, normalization, condon_shortley).conj(), normalize=False)
sqnorm_analytical = sh_squared_norm(l,
normalization,
normalized_haar=False)
dot_analytical = sqnorm_analytical * (l == l2 and m == m2)
print(l, m, l2, m2, field, normalization, condon_shortley,
dot_analytical, dot_numerical, dot_numerical2)
assert np.isclose(dot_numerical, dot_analytical)
assert np.isclose(dot_numerical2, dot_analytical)
def spherical_voxel_optimized(points: np.ndarray, size_bandwidth: int, size_radial_divisions: int,
radius_support: float, do_random_sampling: bool, num_random_points: int) \
-> Tuple[np.ndarray, np.ndarray]:
"""Compute spherical voxel using the C++ code.
Compute Spherical Voxel signal as defined in:
Pointwise Rotation-Invariant Network withAdaptive Sampling and 3D Spherical Voxel Convolution.
Yang You, Yujing Lou, Qi Liu, Yu-Wing Tai, Weiming Wang, Lizhuang Ma and Cewu Lu.
AAAI 2020.
:param points: the points to convert.
:param size_bandwidth: alpha and beta bandwidth.
:param size_radial_divisions: the number of bins along radial dimension.
:param radius_support: the radius used to compute the points in the support.
:param do_random_sampling: if true a subset of random points will be used to compute the spherical voxel.
:param num_random_points: the number of points to keep if do_random_sampling is true.
:return: A tuple containing:
The spherical voxel, shape(size_radial_divisions, 2 * size_bandwidth, 2 * size_bandwidth).
The points used to compute the signal normalized according the the farthest point.
"""
if do_random_sampling:
min_limit = 1 if points.shape[0] > 1 else 0
indices_random = np.random.randint(min_limit, points.shape[0], num_random_points)
points = points[indices_random]
pts_norm = np.linalg.norm(points, axis=1)
# Scale points to fit unit sphere
pts_normed = points / pts_norm[:, None]
pts_normed = np.clip(pts_normed, -1, 1)
pts_s2_coord = S2.change_coordinates(pts_normed, p_from='C', p_to='S')
# Convert to spherical voxel indices
pts_s2_coord[:, 0] *= 2 * size_bandwidth / np.pi # [0, pi]
pts_s2_coord[:, 1] *= size_bandwidth / np.pi
pts_s2_coord[:, 1][pts_s2_coord[:, 1] < 0] += 2 * size_bandwidth
# Adaptive sampling factor
daas_weights = np.sin(np.pi * (2 * np.arange(2 * size_bandwidth) + 1) / 4 / size_bandwidth).astype(np.float32)
voxel = np.asarray(sv.compute(pts_on_s2=pts_s2_coord,
pts_norm=pts_norm,
size_bandwidth=size_bandwidth,
size_radial_divisions=size_radial_divisions,
radius_support=radius_support,
daas_weights=daas_weights))
pts_normed = points / np.max(pts_norm)
return voxel.astype(np.float32), pts_normed.astype(np.float32)
def __getitem__(self, index):
b = self.bw
pts = np.array(self.pts[index])
# randomly sample points
sub_idx = np.random.randint(0, pts.shape[0], 2048)
pts = pts[sub_idx]
if self.aug:
rot = rnd_rot()
pts = np.einsum('ij,nj->ni', rot, pts)
pts += np.random.rand(3)[None, :] * 0.05
pts = np.einsum('ij,nj->ni', rot.T, pts)
segs = np.array(self.segs[index])
segs = segs[sub_idx]
labels = self.labels[index]
pts_norm = np.linalg.norm(pts, axis=1)
pts_normed = pts / pts_norm[:, None]
rand_rot = rnd_rot() if self.rand_rot else np.eye(3)
rotated_pts_normed = np.clip(pts_normed @ rand_rot, -1, 1)
pts_s2 = S2.change_coordinates(rotated_pts_normed,
p_from='C',
p_to='S')
pts_s2[:, 0] *= 2 * b / np.pi # [0, pi]
pts_s2[:, 1] *= b / np.pi
pts_s2[:, 1][pts_s2[:, 1] < 0] += 2 * b
pts_s2_float = pts_s2
# N * 3
pts_so3 = np.stack([
pts_norm * 2 - 1, pts_s2_float[:, 1] /
(2 * b - 1) * 2 - 1, pts_s2_float[:, 0] / (2 * b - 1) * 2 - 1
],
axis=1)
pts_so3 = np.clip(pts_so3, -1, 1)
features = np.asarray(
compute(pts_s2_float, np.linalg.norm(pts, axis=1), 2 * b, b,
np.sin(np.pi * (2 * np.arange(2 * b) + 1) / 4 / b)))
features = np.moveaxis(features, [0, 1, 2], [2, 0, 1])[None]
return features.astype(np.float32), pts_so3.astype(
np.float32), segs.astype(np.int64), pts @ rand_rot, labels.astype(
np.int64)
def __getitem__(self, index):
b = hyper.BANDWIDTH_IN
pts = np.array(self.pts[index])
# randomly sample points
sub_idx = np.random.randint(0, pts.shape[0], hyper.N_PTCLOUD)
pts = pts[sub_idx]
if self.aug:
rot = rnd_rot()
pts = np.einsum('ij,nj->ni', rot, pts)
pts += np.random.rand(3)[None, :] * 0.05
pts = np.einsum('ij,nj->ni', rot.T, pts)
segs = np.array(self.segs[index])
segs = segs[sub_idx]
labels = self.labels[index]
pts_norm = np.linalg.norm(pts, axis=1)
pts_normed = pts / pts_norm[:, None]
rand_rot = rnd_rot() if self.rand_rot else np.eye(3)
rotated_pts_normed = np.clip(pts_normed @ rand_rot, -1, 1)
pts_s2 = S2.change_coordinates(rotated_pts_normed, p_from='C', p_to='S')
pts_s2[:, 0] *= 2 * b / np.pi # [0, pi]
pts_s2[:, 1] *= b / np.pi
pts_s2[:, 1][pts_s2[:, 1] < 0] += 2 * b
pts_s2_float = pts_s2
# N * 3
pts_so3 = np.stack([pts_norm * 2 - 1, pts_s2_float[:, 1] / (2 * b - 1) * 2 - 1, pts_s2_float[:, 0] / (2 * b - 1) * 2 - 1], axis=1)
pts_so3 = np.clip(pts_so3, -1, 1)
# one hundred times speed up !
features = np.asarray(compute(pts_s2_float, np.linalg.norm(pts, axis=1), hyper.R_IN, b, np.sin(np.pi * (2 * np.arange(2 * b) + 1) / 4 / b)))
print('SSS', type(pts), pts.shape, pts_so3.shape, features.shape) # 2048 x 3, 2048 x 3, 64 x 64 x 64
print('Mins/maxs/avg pts', pts.min(0), pts.max(0), pts.mean(0))
print('Mins/maxs p so3ts', pts_so3.min(0), pts_so3.max(0))
return features.astype(np.float32), pts_so3.astype(np.float32), segs.astype(np.int64), pts @ rand_rot, labels.astype(np.int64)
def compare_naive_and_spectral_conv():
f1 = lambda t, p: sh(l=2,
m=1,
theta=t,
phi=p,
field='real',
normalization='quantum',
condon_shortley=True)
f2 = lambda t, p: sh(l=2,
m=1,
theta=t,
phi=p,
field='real',
normalization='quantum',
condon_shortley=True)
theta, phi = S2.meshgrid(b=4, grid_type='Gauss-Legendre')
f1_grid = f1(theta, phi)
f2_grid = f2(theta, phi)
alpha, beta, gamma = S3.meshgrid(b=4,
grid_type='SOFT') # TODO check convention
f12_grid_spectral = spectral_S2_conv(f1_grid,
f2_grid,
s2_fft=None,
so3_fft=None)
f12_grid = np.zeros_like(alpha)
for i in range(alpha.shape[0]):
for j in range(alpha.shape[1]):
for k in range(alpha.shape[2]):
f12_grid[i, j, k] = naive_S2_conv(f1, f2, alpha[i, j, k],
beta[i, j, k], gamma[i, j,
k])
print(i, j, k, f12_grid[i, j, k])
return f1_grid, f2_grid, f12_grid, f12_grid_spectral
def naive_S2_conv(f1, f2, alpha, beta, gamma, g_parameterization='EA323'):
"""
Compute int_S^2 f1(x) f2(g^{-1} x)* dx,
where x = (theta, phi) is a point on the sphere S^2,
and g = (alpha, beta, gamma) is a point in SO(3) in Euler angle parameterization
:param f1, f2: functions to be convolved
:param alpha, beta, gamma: the rotation at which to evaluate the result of convolution
:return:
"""
# This fails
def integrand(theta, phi):
g_inv = SO3.invert((alpha, beta, gamma),
parameterization=g_parameterization)
g_inv_theta, g_inv_phi, _ = SO3.transform_r3(
g=g_inv,
x=(theta, phi, 1.),
g_parameterization=g_parameterization,
x_parameterization='S')
return f1(theta, phi) * f2(g_inv_theta, g_inv_phi).conj()
return S2.integrate(f=integrand, normalize=True)
def linspace(b, grid_type='SOFT'):
"""
Compute a linspace on the 3-sphere.
Since S3 is ismorphic to SO(3), we use the grid grid_type from:
FFTs on the Rotation Group
Peter J. Kostelec and Daniel N. Rockmore
http://www.cs.dartmouth.edu/~geelong/soft/03-11-060.pdf
:param b:
:return:
"""
# alpha = 2 * np.pi * np.arange(2 * b) / (2. * b)
# beta = np.pi * (2 * np.arange(2 * b) + 1) / (4. * b)
# gamma = 2 * np.pi * np.arange(2 * b) / (2. * b)
beta, alpha = S2.linspace(b, grid_type)
# According to this paper:
# "Sampling sets and quadrature formulae on the rotation group"
# We can just tack a sampling grid for S^1 to a sampling grid for S^2 to get a sampling grid for SO(3).
gamma = 2 * np.pi * np.arange(2 * b) / (2. * b)
return alpha, beta, gamma
def test_S2FFT():
L_max = 10
beta, alpha = S2.meshgrid(b=L_max + 1, grid_type='Driscoll-Healy')
lt = setup_legendre_transform(b=L_max + 1)
lti = setup_legendre_transform_indices(b=L_max + 1)
for l in range(L_max):
for m in range(-l, l + 1):
Y = sh(l,
m,
beta,
alpha,
field='complex',
normalization='seismology',
condon_shortley=True)
y_hat = sphere_fft(Y, lt, lti)
# The flat index for (l, m) is l^2 + l + m
# Before the harmonics of degree l, there are this many harmonics: sum_{i=0}^{l-1} 2i+1 = l^2
# There are 2l+1 harmonics of degree l, with order m=0 at the center,
# so the m-th harmonic of degree is at l + m within the block of degree l.
y_hat_true = np.zeros_like(y_hat)
y_hat_true[l**2 + l + m] = 1
diff = np.sum(np.abs(y_hat - y_hat_true))
nz = 1. - np.isclose(y_hat, 0.)
diff_nz = np.sum(np.abs(nz - y_hat_true))
print(l, m, diff, diff_nz)
print(np.round(y_hat, 4))
print(y_hat_true)
# assert np.isclose(diff, 0.) # TODO make this work
print(nz)
assert np.isclose(diff_nz, 0.)
def __getitem__(self, index):
b = hyper.BANDWIDTH_IN
pts = np.array(self.pts[index])
# randomly sample points
sub_idx = np.random.randint(0, pts.shape[0], hyper.N_PTCLOUD)
pts = pts[sub_idx]
if self.aug:
rot = rnd_rot()
np.einsum('ij,nj->ni', rot, pts)
pts += np.random.rand(3)[None, :] * 0.05
np.einsum('ij,nj->ni', rot.T, pts)
segs = np.array(self.segs[index])
segs = segs[sub_idx]
labels = self.labels[index]
pts_norm = np.linalg.norm(pts, axis=1)
pts_normed = pts / pts_norm[:, None]
rand_rot = rnd_rot() if self.rand_rot else np.eye(3)
rotated_pts_normed = np.clip(pts_normed @ rand_rot, -1, 1)
pts_s2 = S2.change_coordinates(rotated_pts_normed,
p_from='C',
p_to='S')
pts_s2[:, 0] *= 2 * b / np.pi # [0, pi]
pts_s2[:, 1] *= b / np.pi
pts_s2[:, 1][pts_s2[:, 1] < 0] += 2 * b
pts_s2_float = pts_s2
pts_s2 = (pts_s2 + 0.5).astype(np.int)
pts_s2[:, 0] = np.clip(pts_s2[:, 0], 0, 2 * b - 1)
pts_s2[:, 1] = np.clip(pts_s2[:, 1], 0, 2 * b - 1) # [0, 2pi]
# N * 3
pts_so3 = np.stack([
pts_norm * 2 - 1, pts_s2_float[:, 1] /
(2 * b - 1) * 2 - 1, pts_s2_float[:, 0] / (2 * b - 1) * 2 - 1
],
axis=1)
pts_so3 = np.clip(pts_so3, -1, 1)
# cache data
try:
if self.cache_dir is None:
raise FileNotFoundError
features = np.load(
os.path.join(self.cache_dir, 'features%d.npy' % index))
except (OSError, FileNotFoundError):
features = []
interval = 1. / hyper.R_IN
dist = np.linalg.norm(pts, axis=1)
# adaptive sampling
wt = np.sin(np.pi * (2 * np.arange(2 * b) + 1) / 4 / b)
# TODO: rewrite this in an efficient way
for i in range(hyper.R_IN):
idx = (dist < (i + 2) * interval) & (dist > i * interval)
pts_idx = pts_s2[idx]
pts_idx_float = pts_s2_float[idx]
im = np.zeros([2 * b, 2 * b], np.float32)
for beta in range(2 * b):
filt = (pts_idx[:, 0] == beta)
for alpha in range(2 * b):
filt_alpha = filt & (
pts_idx_float[:, 1] > alpha - 1. / 2 / wt[beta]
) & (pts_idx_float[:, 1] < alpha + 1. / 2 / wt[beta])
if not np.any(filt_alpha):
continue
im[beta,
alpha] = (1 - np.abs(dist[idx][filt_alpha] -
(i + 1) * interval) / interval
).sum() / np.count_nonzero(filt_alpha)
features.append(im)
features = np.stack(features, axis=0)
if self.cache_dir is not None:
np.save(os.path.join(self.cache_dir, 'features%d.npy' % index),
features)
return features, pts_so3.astype(np.float32), segs.astype(
np.int64), pts @ rand_rot, labels.astype(np.int64)
def inverse_render_model(points: np.ndarray, sgrid: np.ndarray):
# wait for implementing
print("Aloha")
x = points[:, 0]
y = points[:, 1]
z = points[:, 2]
# Fistly, we get the centroid, then translate the points
centroid_x = np.sum(x) / points.shape[0]
centroid_y = np.sum(y) / points.shape[0]
centroid_z = np.sum(z) / points.shape[0]
centroid = np.array([centroid_x, centroid_y, centroid_z])
points = points.astype(np.float)
points -= centroid
# After normalization, compute the distance between the sphere and points
radius = np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2 + points[:, 2] ** 2)
# dist = 1 - (1 / np.max(radius)) * radius
# Projection
from lie_learn.spaces import S2
radius = np.repeat(radius, 3).reshape(-1, 3)
points_on_sphere = points / radius
# ssgrid = sgrid.reshape(-1, 3)
# phi, theta = S2.change_coordinates(ssgrid, p_from='C', p_to='S')
out = S2.change_coordinates(points_on_sphere, p_from='C', p_to='S')
phi = out[..., 0]
theta = out[..., 1]
phi = phi
theta = theta % (np.pi * 2)
# Interpolate
b = sgrid.shape[0] / 2 # bandwidth
# By computing the m,n, we can find
# the neighbours on the sphere
m = np.trunc((phi - np.pi / (4 * b)) / (np.pi / (2 * b)))
m = m.astype(int)
n = np.trunc(theta / (np.pi / b))
n = n.astype(int)
dist_im, center_grid, east_grid, south_grid, southeast_grid = interpolate(m=m, n=n, sgrid=sgrid,
points_on_sphere=points_on_sphere,
radius=radius)
coef, intercept,center_points =angle(m=m,n=n,sgrid=sgrid,dist_im=dist_im)
# utilizing linear regression to create a plane
# use a mask to avoid the index out of the boundary
"""
mask_m = m - 1 >= 0
mask = mask_m
m = m[mask]
n = n[mask]
"""
# =======================================================================================
fig = plt.figure()
grid = make_sgrid(bandwidth, 0, 0, 0)
grid = grid.reshape((-1, 3))
xx = grid[:, 0]
yy = grid[:, 1]
zz = grid[:, 2]
xx = xx.reshape(-1, 1)
yy = yy.reshape(-1, 1)
zz = zz.reshape(-1, 1)
ax = Axes3D(fig)
ax.scatter(0, 0, 0)
ax.scatter(points[:, 0], points[:, 1], points[:, 2])
ax.scatter(points_on_sphere[:, 0], points_on_sphere[:, 1], points_on_sphere[:, 2])
ax.scatter(center_grid[:, 0], center_grid[:, 1], center_grid[:, 2])
ax.scatter(east_grid[:, 0], east_grid[:, 1], east_grid[:, 2])
ax.scatter(south_grid[:, 0], south_grid[:, 1], south_grid[:, 2])
ax.scatter(southeast_grid[:, 0], southeast_grid[:, 1], southeast_grid[:, 2])
# ax.scatter(xx, yy, zz)
# plt.legend()
# draw line
ax = fig.gca(projection='3d')
zero = np.zeros(points_on_sphere.shape[0])
ray_x = np.stack((zero, points_on_sphere[:, 0]), axis=1).reshape(-1, 2)
ray_y = np.stack((zero, points_on_sphere[:, 1]), axis=1).reshape(-1, 2)
ray_z = np.stack((zero, points_on_sphere[:, 2]), axis=1).reshape(-1, 2)
for index in range(points_on_sphere.shape[0]):
ax.plot(ray_x[index], ray_y[index], ray_z[index])
# draw plane
for index in range(center_points.shape[0]):
X = np.arange(center_points[index, 0] - 0.05, center_points[index, 0] + 0.05, 0.01)
Y = np.arange(center_points[index, 1] - 0.05, center_points[index, 1] + 0.05, 0.01)
X, Y = np.meshgrid(X, Y)
a1 = coef[index, 0]
a2 = coef[index, 1]
b = intercept[index]
Z = a1 * X + a2 * Y + b
surf = ax.plot_surface(X, Y, Z)
plt.show()
im = dist_im
return im
def f1a(xs):
xc = S2.change_coordinates(coords=xs, p_from='S', p_to='C')
return xc[..., 0]**2 * xc[..., 1] - 1.4 * xc[..., 2] * xc[
..., 1]**3 + xc[..., 1] - xc[..., 2]**2 + 2.
