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
