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 spherical_harmonics(order, alpha, beta, dtype=None):
"""
spherical harmonics
:param order: int or list
:param alpha: float or tensor of shape [A]
:param beta: float or tensor of shape [A]
:return: tensor of shape [m, A]
- compatible with irr_repr and compose
"""
from lie_learn.representations.SO3.spherical_harmonics import sh # real valued by default
import numpy as np
if not isinstance(order, list):
order = [order]
if not torch.is_tensor(alpha):
alpha = torch.tensor(alpha, dtype=torch.float64)
if not torch.is_tensor(beta):
beta = torch.tensor(beta, dtype=torch.float64)
Js = np.concatenate([J * np.ones(2 * J + 1) for J in order], 0)
Ms = np.concatenate([np.arange(-J, J + 1, 1) for J in order], 0)
Js = Js.reshape(-1, *[1] * alpha.dim())
Ms = Ms.reshape(-1, *[1] * alpha.dim())
alpha = alpha.view(1, *alpha.size())
beta = beta.view(1, *beta.size())
Y = sh(Js, Ms, math.pi - beta.cpu().numpy(), alpha.cpu().numpy())
Y = torch.tensor(
Y, dtype=torch.get_default_dtype() if dtype is None else dtype)
return Y
def spherical_harmonics(order,
alpha,
beta,
sph_last=False,
dtype=None,
device=None):
"""
spherical harmonics
:param order: int or list
:param alpha: float or tensor of shape [...]
:param beta: float or tensor of shape [...]
:param sph_last: return the spherical harmonics in the last channel
:param dtype:
:param device:
:return: tensor of shape [m, ...] (or [..., m] if sph_last)
- compatible with irr_repr and compose
"""
from lie_learn.representations.SO3.spherical_harmonics import sh # real valued by default
import numpy as np
try:
order = list(order)
except TypeError:
order = [order]
if dtype is None and torch.is_tensor(alpha):
dtype = alpha.dtype
if dtype is None and torch.is_tensor(beta):
dtype = beta.dtype
if dtype is None:
dtype = torch.get_default_dtype()
if device is None and torch.is_tensor(alpha):
device = alpha.device
if device is None and torch.is_tensor(beta):
device = beta.device
if not torch.is_tensor(alpha):
alpha = torch.tensor(alpha, dtype=torch.float64)
if not torch.is_tensor(beta):
beta = torch.tensor(beta, dtype=torch.float64)
Js = np.concatenate([J * np.ones(2 * J + 1) for J in order], 0)
Ms = np.concatenate([np.arange(-J, J + 1, 1) for J in order], 0)
Js = Js.reshape(-1, *[1] * alpha.dim())
Ms = Ms.reshape(-1, *[1] * alpha.dim())
alpha = alpha.unsqueeze(0)
beta = beta.unsqueeze(0)
Y = sh(Js, Ms, math.pi - beta.cpu().numpy(), alpha.cpu().numpy())
if sph_last:
rank = len(Y.shape)
return torch.tensor(Y, dtype=dtype,
device=device).permute(*range(1, rank),
0).contiguous()
else:
return torch.tensor(Y, dtype=dtype, device=device)
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 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 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_harmonics(order, alpha, beta, dtype=None):
"""
spherical harmonics
- compatible with irr_repr and compose
"""
from lie_learn.representations.SO3.spherical_harmonics import sh # real valued by default
Y = torch.tensor(
[
sh(order, m, math.pi - beta, alpha)
for m in range(-order, order + 1)
],
dtype=torch.get_default_dtype() if dtype is None else dtype)
# if order == 1:
# # change of basis to have vector_field[x, y, z] = [vx, vy, vz]
# A = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]])
# return A @ Y
return Y
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_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.)
phi = 0.1*torch.randn(bs,1024,10, dtype=dtype)
cu_theta = theta.to(device)
cu_phi = phi.to(device)
s0 = s1 = s2 = 0
max_error = -1.
sph_har = SphericalHarmonics()
for l in range(10):
for m in range(l, -l-1, -1):
start = time.time()
#y = tesseral_harmonics(l, m, theta, phi)
y = sph_har.get_element(l, m, cu_theta, cu_phi).type(torch.float32)
#y = sph_har.lpmv(l, m, phi)
s0 += time.time() - start
start = time.time()
z = sh(l, m, theta, phi)
#z = lpmv_scipy(m, l, phi).numpy()
s1 += time.time() - start
error = np.mean(np.abs((y.cpu().numpy() - z) / z))
max_error = max(max_error, error)
print(f"l: {l}, m: {m} ", error)
#start = time.time()
#sph_har.get(l, theta, phi)
#s2 += time.time() - start
print('#################')
print(f"Max error: {max_error}")
print(f"Time diff: {s0/s1}")
