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}")