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 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 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, 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):
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 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 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 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 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 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 __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 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 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 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 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 get_projection_grid(bandwidth, grid_type="Driscoll-Healy"):
theta, phi = S2.meshgrid(b=bandwidth, grid_type=grid_type)
x_ = np.sin(theta) * np.cos(phi)
y_ = np.sin(theta) * np.sin(phi)
z_ = np.cos(theta)
return np.array((x_, y_, z_))