def _init_pswf_func2d(self, c, eps):
"""
Initialize the whole set of PSWF functions with the input bandlimit and error
:param c: bandlimit (>0) can be estimated by beta * pi * rcut
:param eps: error tolerance
:return:
alpha_all (list of arrays):
alpha = alpha_all[i] contains all the eigenvalues for N=i such that
lambda > eps, where lambda is the normalized alpha values (i.e.
lambda is between 0 and 1), given by lambda=sqrt(c*np.absolute(alpha)/(2*pi)).
d_vec_all (list of 2D lists): the corresponding eigenvectors for alpha_all.
n_order_length_vec (list of ints): n_order_length_vec[i] = len(alpha_all[i])
"""
d_vec_all = []
alpha_all = []
n_order_length_vec = []
m = 0
n = int(np.ceil(2 * c / np.pi))
r, w = lgwt(n, 0, 1)
cons = c / 2 / np.pi
while True:
alpha, d_vec, a = self.pswf_func2d(m, n, c, eps, r, w)
lambda_var = np.sqrt(cons * np.absolute(alpha))
n_end = np.where(lambda_var <= eps)[0]
if len(n_end) != 0:
n_end = n_end[0]
if n_end == 0:
break
n_order_length_vec.append(n_end)
alpha_all.append(alpha[:n_end])
d_vec_all.append(d_vec[:, :n_end])
m += 1
n = n_end + 1
else:
n *= 2
r, w = lgwt(n, 0, 1)
return d_vec_all, alpha_all, n_order_length_vec
def filter_to_fb_mat(h_fun, fbasis):
"""
Convert a nonradial function in k space into a basis representation.
:param h_fun: The function form in k space.
:param fbasis: The basis object for expanding.
:return: a BlkDiagMatrix instance representation using the
`fbasis` expansion.
"""
# These form a circular dependence, import locally until time to clean up.
from aspire.basis import FFBBasis2D
from aspire.basis.basis_utils import lgwt
if not isinstance(fbasis, FFBBasis2D):
raise NotImplementedError("Currently only fast FB method is supported")
# Set same dimensions as basis object
n_k = fbasis.n_r
n_theta = fbasis.n_theta
radial = fbasis.get_radial()
# get 2D grid in polar coordinate
k_vals, wts = lgwt(n_k, 0, 0.5, dtype=fbasis.dtype)
k, theta = np.meshgrid(k_vals,
np.arange(n_theta) * 2 * np.pi / (2 * n_theta),
indexing="ij")
# Get function values in polar 2D grid and average out angle contribution
omegax = k * np.cos(theta)
omegay = k * np.sin(theta)
omega = 2 * np.pi * np.vstack((omegax.flatten("C"), omegay.flatten("C")))
h_vals2d = h_fun(omega).reshape(n_k, n_theta).astype(fbasis.dtype)
h_vals = np.sum(h_vals2d, axis=1) / n_theta
# Represent 1D function values in fbasis
h_fb = BlkDiagMatrix.empty(2 * fbasis.ell_max + 1, dtype=fbasis.dtype)
ind_ell = 0
for ell in range(0, fbasis.ell_max + 1):
k_max = fbasis.k_max[ell]
rmat = 2 * k_vals.reshape(n_k, 1) * fbasis.r0[0:k_max, ell].T
fb_vals = np.zeros_like(rmat)
ind_radial = np.sum(fbasis.k_max[0:ell])
fb_vals[:, 0:k_max] = radial[ind_radial:ind_radial + k_max].T
h_fb_vals = fb_vals * h_vals.reshape(n_k, 1)
h_fb_ell = fb_vals.T @ (h_fb_vals * k_vals.reshape(n_k, 1) *
wts.reshape(n_k, 1))
h_fb[ind_ell] = h_fb_ell
ind_ell += 1
if ell > 0:
h_fb[ind_ell] = h_fb[ind_ell - 1]
ind_ell += 1
return h_fb
def filter_to_fb_mat(h_fun, fbasis):
"""
Convert a nonradial function in k space into a basis representation.
:param h_fun: The function form in k space.
:param fbasis: The basis object for expanding.
:return: a BlkDiagMatrix instance representation using the
`fbasis` expansion.
"""
if not isinstance(fbasis, FFBBasis2D):
raise NotImplementedError('Currently only fast FB method is supported')
# Set same dimensions as basis object
n_k = int(np.ceil(4 * fbasis.rcut * fbasis.kcut))
n_theta = np.ceil(16 * fbasis.kcut * fbasis.rcut)
n_theta = int((n_theta + np.mod(n_theta, 2)) / 2)
# get 2D grid in polar coordinate
k_vals, wts = lgwt(n_k, 0, 0.5)
k, theta = np.meshgrid(
k_vals, np.arange(n_theta) * 2 * np.pi / (2 * n_theta), indexing='ij')
# Get function values in polar 2D grid and average out angle contribution
omegax = k * np.cos(theta)
omegay = k * np.sin(theta)
omega = 2 * np.pi * np.vstack((omegax.flatten('C'), omegay.flatten('C')))
h_vals2d = h_fun(omega).reshape(n_k, n_theta)
h_vals = np.sum(h_vals2d, axis=1)/n_theta
# Represent 1D function values in fbasis
h_fb = BlkDiagMatrix.empty(2 * fbasis.ell_max + 1, dtype=fbasis.dtype)
ind = 0
for ell in range(0, fbasis.ell_max+1):
k_max = fbasis.k_max[ell]
rmat = 2*k_vals.reshape(n_k, 1)*fbasis.r0[0:k_max, ell].T
fb_vals = np.zeros_like(rmat)
for ik in range(0, k_max):
fb_vals[:, ik] = jv(ell, rmat[:, ik])
fb_nrms = 1/np.sqrt(2)*abs(jv(ell+1, fbasis.r0[0:k_max, ell].T))/2
fb_vals = fb_vals/fb_nrms
h_fb_vals = fb_vals*h_vals.reshape(n_k, 1)
h_fb_ell = fb_vals.T @ (
h_fb_vals * k_vals.reshape(n_k, 1) * wts.reshape(n_k, 1))
h_fb[ind] = h_fb_ell
ind += 1
if ell > 0:
h_fb[ind] = h_fb[ind-1]
ind += 1
return h_fb
def _precomp(self):
"""
Precomute the basis functions on a polar Fourier grid
Gaussian quadrature points and weights are also generated.
The sampling criterion requires n_r=4*c*R and n_theta= 16*c*R.
"""
n_r = self.n_r
n_theta = self.n_theta
r, w = lgwt(n_r, 0.0, self.kcut, dtype=self.dtype)
radial = np.zeros(shape=(np.sum(self.k_max), n_r), dtype=self.dtype)
ind_radial = 0
for ell in range(0, self.ell_max + 1):
for k in range(1, self.k_max[ell] + 1):
radial[ind_radial] = jv(ell,
self.r0[k - 1, ell] * r / self.kcut)
# NOTE: We need to remove the factor due to the discretization here
# since it is already included in our quadrature weights
# Only normalized by the radial part of basis function
nrm = 1 / (np.sqrt(np.prod(
self.sz))) * self.radial_norms[ind_radial]
radial[ind_radial] /= nrm
ind_radial += 1
# Only calculate "positive" frequencies in one half-plane.
freqs_x = m_reshape(r, (n_r, 1)) @ m_reshape(
np.cos(
np.arange(n_theta, dtype=self.dtype) * 2 * pi / (2 * n_theta)),
(1, n_theta),
)
freqs_y = m_reshape(r, (n_r, 1)) @ m_reshape(
np.sin(
np.arange(n_theta, dtype=self.dtype) * 2 * pi / (2 * n_theta)),
(1, n_theta),
)
freqs = np.vstack((freqs_y[np.newaxis, ...], freqs_x[np.newaxis, ...]))
return {
"gl_nodes": r,
"gl_weights": w,
"radial": radial,
"freqs": freqs
}
def testLGQuad(self):
resx, resw = lgwt(
ndeg=10, # degree
a=0.0, # start x
b=0.5 # end x
)
self.assertTrue(np.allclose(
resx,
[0.00652337, 0.03373416, 0.08014761, 0.14165115, 0.21278142,
0.28721858, 0.35834885, 0.41985239, 0.46626584, 0.49347663
]
))
self.assertTrue(np.allclose(
resw,
[0.01666784, 0.03736284, 0.05477159, 0.06731668, 0.07388106,
0.07388106, 0.06731668, 0.05477159, 0.03736284, 0.01666784
]
))
def _precomp(self):
"""
Precomute the basis functions on a polar Fourier grid
Gaussian quadrature points and weights are also generated.
The sampling criterion requires n_r=4*c*R and n_theta= 16*c*R.
"""
n_r = int(np.ceil(4 * self.rcut * self.kcut))
r, w = lgwt(n_r, 0.0, self.kcut)
radial = np.zeros(shape=(n_r, np.sum(self.k_max)))
ind_radial = 0
for ell in range(0, self.ell_max + 1):
for k in range(1, self.k_max[ell] + 1):
radial[:, ind_radial] = jv(ell, self.r0[k - 1, ell] * r / self.kcut)
# NOTE: We need to remove the factor due to the discretization here
# since it is already included in our quadrature weights
nrm = 1 / (np.sqrt(np.prod(self.sz))) * self.basis_norm_2d(ell, k)
radial[:, ind_radial] /= nrm
ind_radial += 1
n_theta = np.ceil(16 * self.kcut * self.rcut)
n_theta = int((n_theta + np.mod(n_theta, 2)) / 2)
# Only calculate "positive" frequencies in one half-plane.
freqs_x = m_reshape(r, (n_r, 1)) @ m_reshape(
np.cos(np.arange(n_theta) * 2 * pi / (2 * n_theta)), (1, n_theta))
freqs_y = m_reshape(r, (n_r, 1)) @ m_reshape(
np.sin(np.arange(n_theta) * 2 * pi / (2 * n_theta)), (1, n_theta))
freqs = np.vstack((freqs_x[np.newaxis, ...], freqs_y[np.newaxis, ...]))
return {
'gl_nodes': r,
'gl_weights': w,
'radial': radial,
'freqs': freqs
}
def _generate_pswf_radial_quad(self, n, bandlimit, phi_approximate_error,
lambda_max):
"""
Generate Gaussian quadrature points and weights for the radical parts of 2D PSWFs
"""
x, w = lgwt(20 * n, 0, 1)
big_n = 0
x_as_mat = x.reshape((len(x), 1))
alpha_n, d_vec, approx_length = self.pswf_func2d(
big_n, n, bandlimit, phi_approximate_error, x, w)
cut_indices = np.where(
bandlimit / 2 / np.pi * np.absolute(alpha_n) < lambda_max)[0]
if len(cut_indices) == 0:
k = len(alpha_n)
else:
k = cut_indices[0]
if k % 2 == 0:
k = k + 1
range_array = np.arange(approx_length).reshape((1, approx_length))
idx_for_quad_nodes = int((k + 1) / 2)
num_quad_pts = idx_for_quad_nodes - 1
phi_zeros = self._find_initial_nodes(x, n, bandlimit / 2,
phi_approximate_error,
idx_for_quad_nodes)
def phi_for_quad_weights(t):
return np.dot(t_x_mat_dot(t, big_n, range_array, approx_length),
d_vec[:, :k - 1])
b = np.dot(w * np.sqrt(x), phi_for_quad_weights(x_as_mat))
a = phi_for_quad_weights(phi_zeros.reshape(
(len(phi_zeros), 1))).transpose() * np.sqrt(phi_zeros)
init_quad_weights = lstsq(a, b, rcond=None)
init_quad_weights = init_quad_weights[0]
tolerance = np.spacing(1)
def obj_func(quad_rule):
q = quad_rule.reshape((len(quad_rule), 1))
temp = np.dot((phi_for_quad_weights(q[:num_quad_pts]) *
np.sqrt(q[:num_quad_pts])).transpose(),
q[num_quad_pts:])
temp = temp.reshape(temp.shape[0])
return temp - b
arr_to_send = np.concatenate((phi_zeros, init_quad_weights))
quad_rule_final = least_squares(obj_func,
arr_to_send,
xtol=tolerance,
ftol=tolerance,
max_nfev=1000)
quad_rule_final = quad_rule_final.x
quad_rule_pts = quad_rule_final[:num_quad_pts]
quad_rule_weights = quad_rule_final[num_quad_pts:]
return quad_rule_pts, quad_rule_weights
def _precomp(self):
"""
Precomute the basis functions on a polar Fourier 3D grid
Gaussian quadrature points and weights are also generated
in radical and phi dimensions.
"""
n_r = int(self.ell_max + 1)
n_theta = int(2 * self.sz[0])
n_phi = int(self.ell_max + 1)
r, wt_r = lgwt(n_r, 0.0, self.kcut)
z, wt_z = lgwt(n_phi, -1, 1)
r = m_reshape(r, (n_r, 1))
wt_r = m_reshape(wt_r, (n_r, 1))
z = m_reshape(z, (n_phi, 1))
wt_z = m_reshape(wt_z, (n_phi, 1))
phi = np.arccos(z)
wt_phi = wt_z
theta = 2 * pi * np.arange(n_theta).T / (2 * n_theta)
theta = m_reshape(theta, (n_theta, 1))
# evaluate basis function in the radial dimension
radial_wtd = np.zeros(shape=(n_r, np.max(self.k_max),
self.ell_max + 1))
for ell in range(0, self.ell_max + 1):
k_max_ell = self.k_max[ell]
rmat = r * self.r0[0:k_max_ell, ell].T / self.kcut
radial_ell = np.zeros_like(rmat)
for ik in range(0, k_max_ell):
radial_ell[:, ik] = sph_bessel(ell, rmat[:, ik])
nrm = np.abs(sph_bessel(ell + 1, self.r0[0:k_max_ell, ell].T) / 4)
radial_ell = radial_ell / nrm
radial_ell_wtd = r**2 * wt_r * radial_ell
radial_wtd[:, 0:k_max_ell, ell] = radial_ell_wtd
# evaluate basis function in the phi dimension
ang_phi_wtd_even = []
ang_phi_wtd_odd = []
for m in range(0, self.ell_max + 1):
n_even_ell = int(
np.floor((self.ell_max - m) / 2) + 1 -
np.mod(self.ell_max, 2) * np.mod(m, 2))
n_odd_ell = int(self.ell_max - m + 1 - n_even_ell)
phi_wtd_m_even = np.zeros((n_phi, n_even_ell), dtype=phi.dtype)
phi_wtd_m_odd = np.zeros((n_phi, n_odd_ell), dtype=phi.dtype)
ind_even = 0
ind_odd = 0
for ell in range(m, self.ell_max + 1):
phi_m_ell = norm_assoc_legendre(ell, m, z)
nrm_inv = np.sqrt(0.5 / pi)
phi_m_ell = nrm_inv * phi_m_ell
phi_wtd_m_ell = wt_phi * phi_m_ell
if np.mod(ell, 2) == 0:
phi_wtd_m_even[:, ind_even] = phi_wtd_m_ell[:, 0]
ind_even = ind_even + 1
else:
phi_wtd_m_odd[:, ind_odd] = phi_wtd_m_ell[:, 0]
ind_odd = ind_odd + 1
ang_phi_wtd_even.append(phi_wtd_m_even)
ang_phi_wtd_odd.append(phi_wtd_m_odd)
# evaluate basis function in the theta dimension
ang_theta = np.zeros((n_theta, 2 * self.ell_max + 1),
dtype=theta.dtype)
ang_theta[:, 0:self.ell_max] = np.sqrt(2) * np.sin(
theta @ m_reshape(np.arange(self.ell_max, 0, -1),
(1, self.ell_max)))
ang_theta[:, self.ell_max] = np.ones(n_theta, dtype=theta.dtype)
ang_theta[:,
self.ell_max + 1:2 * self.ell_max + 1] = np.sqrt(2) * np.cos(
theta @ m_reshape(np.arange(1, self.ell_max + 1),
(1, self.ell_max)))
ang_theta_wtd = (2 * pi / n_theta) * ang_theta
theta_grid, phi_grid, r_grid = np.meshgrid(theta,
phi,
r,
sparse=False,
indexing='ij')
fourier_x = m_flatten(r_grid * np.cos(theta_grid) * np.sin(phi_grid))
fourier_y = m_flatten(r_grid * np.sin(theta_grid) * np.sin(phi_grid))
fourier_z = m_flatten(r_grid * np.cos(phi_grid))
fourier_pts = 2 * pi * np.vstack(
(fourier_x[np.newaxis, ...], fourier_y[np.newaxis, ...],
fourier_z[np.newaxis, ...]))
return {
'radial_wtd': radial_wtd,
'ang_phi_wtd_even': ang_phi_wtd_even,
'ang_phi_wtd_odd': ang_phi_wtd_odd,
'ang_theta_wtd': ang_theta_wtd,
'fourier_pts': fourier_pts
}