def _setup_so3_fft(b, nl, weighted): from lie_learn.representations.SO3.wigner_d import wigner_d_matrix import lie_learn.spaces.S3 as S3 import numpy as np import logging betas = (np.arange(2 * b) + 0.5) / (2 * b) * np.pi w = S3.quadrature_weights(b) assert len(w) == len(betas) logging.getLogger("trainer").info("Compute Wigner: b=%d nbeta=%d nl=%d nspec=%d", b, len(betas), nl, nl ** 2) dss = [] for b, beta in enumerate(betas): ds = [] for l in range(nl): d = wigner_d_matrix(l, beta, field='complex', normalization='quantum', order='centered', condon_shortley='cs') d = d.reshape(((2 * l + 1) ** 2,)) if weighted: d *= w[b] else: d *= 2 * l + 1 # d # [m * n] ds.append(d) ds = np.concatenate(ds) # [l * m * n] dss.append(ds) dss = np.stack(dss) # [beta, l * m * n] return dss
def setup_d_transform(b, L2_normalized, field='complex', normalization='quantum', order='centered', condon_shortley='cs'): """ Precompute arrays of samples from the Wigner-d function, for use in the Wigner-d transform. Specifically, the samples that are required are: d^l_mn(beta_k) for: l = 0, ..., b - 1 -l <= m, n <= l k = 0, ..., 2b - 1 (where beta_k = pi (2 b + 1) / 4b) This data is returned as a list d indexed by l (of length b), where each element of the list is an array d[l] of shape (2l+1, 2b, 2l+1) indexed by (m, k, n) In the Wigner-d transform, for each l, we reduce an array d[l] of shape (2l+1, 2b, 2l+1) against a data array of the same shape, along the beta axis (axis 1 of length 2b). :param b: bandwidth of the transform :param L2_normalized: whether to use L2_normalized versions of the Wigner-d functions. :param field, normalization, order, condon_shortley: the basis and normalization convention (see irrep_bases.py) :return a list d of length b, where d[l] is an array of shape (2l+1, 2b, 2l+1) """ # Compute array of beta values as described in SOFT 2.0 documentation beta = np.pi * (2 * np.arange(0, 2 * b) + 1) / (4. * b) # For each l=0, ..., b-1, we compute a 3D tensor of shape (2l+1, 2b, 2l+1) for axes (m, beta, n) # Together, these indices (l, m, beta, n) identify d^l_mn(beta) convention = {'field': field, 'normalization': normalization, 'order': order, 'condon_shortley': condon_shortley} d = [np.array([wigner_d_matrix(l, bt, **convention) for bt in beta]).transpose(1, 0, 2) for l in range(b)] if L2_normalized: # TODO: this should be integrated in the normalization spec above, no? # The Unitary matrix elements have norm: # | U^\lambda_mn |^2 = 1/(2l+1) # where the 2-norm is defined in terms of normalized Haar measure. # So T = sqrt(2l + 1) U are L2-normalized functions d = [d[l] * np.sqrt(2 * l + 1) for l in range(len(d))] # We want the L2 normalized functions: # d = [d[l] * np.sqrt(l + 0.5) for l in range(len(d))] return d