def expand_sph_harm_order(self, A_ell, ell):
"""
Evaluate a single order term in the expansion of a scalar field in
spherical harmonics. Specifically, if field phi(r) has been expanded
phi(r) = sum sum { A_{l,m}(|r|) * Y_lm(Omega) }
l m
then this function takes A_{ell,m} for one value of ell and returns
phi_ell(r) = sum { A_{l,m}(|r|) * Y_lm(Omega) }
m
evaluated at points on a regular square grid.
Parameters
----------
A_ell : np.ndarray, float
A shape ( 2`ell`+1, len(`radii`) ) array of spherical harmonic
coefficients for a single ell band.
ell : int
The value of l for which to compute the term.
Returns
-------
grid_term : np.ndarray, complex128
One term, representing an-ell band in the spherical harmonic
expansion of an object.
See Also
--------
expand_sph_harm : method
Evaluate the expansion for all terms {ell} simultaneously.
"""
l = ell # rename
if not A_ell.shape == (2*l+1, len(self.radii)):
raise ValueError('`A_ell` has an invalid shape. Must be a (2l+1) x'
' len(radii) shape array. Got '
'(%s)' % str(A_ell.shape))
grid_term = np.zeros(self._grid_dimensions, dtype=np.complex128)
for ir,r in enumerate(self._radii):
S, u = self._compute_shell(r)
w_grid = np.zeros_like(grid)
w_grid[S] = np.power(1 - np.power(np.abs(u)[S], 3), 3)
for m in range(-l, l+1):
Y_lm = sph_harm(l, m, self._theta, self._phi)
grid_term += coefficients[ir,l,m] * w_grid * Y_lm
return grid_term
def expand_sph_harm(self, coefficients):
"""
Evaluate a field phi(r) on a regular grid by expanding a series of
spherial harmonic coefficients:
phi(r) = sum sum { A_{l,m}(r) * Y_lm(Omega) }
l m
Parameters
----------
coefficients : np.ndarray
A three dimensional array of coefficients, indexed by (radii, l, m).
Note that the 'm' values wrap around, so that for m < 0 the
index is 2l+1-m.
Returns
-------
grid : np.ndarray, complex128
The expansion evaluated on the regular grid.
See Also
--------
expand_sph_harm_order : method
Evaluate the expansion for only a single order {ell}.
"""
if not coefficients.shape[0] == len(self.radii):
raise ValueError('`coefficients` first dimension must be the same '
'length as self.radii')
l_max = coefficients.shape[1]
grid = np.zeros(self._grid_dimensions, dtype=np.complex128)
for ir,r in enumerate(self._radii):
S, u = self._compute_shell(r)
w_grid = np.zeros_like(grid)
w_grid[S] = np.power(1 - np.power(np.abs(u)[S], 3), 3)
for l in range(self.L_max):
for m in range(-l, l+1):
Y_lm = sph_harm(l, m, self._theta, self._phi)
grid += coefficients[ir,l,m] * w_grid * Y_lm
return grid
def test1():
"""
Create a spherical harmonic function, and test to see how it gets
decomposed.
July 10: this seems to be working OK, just based on visual inspection
of coefficient magnitudes...
"""
grid_dimesions = (30, 30, 30)
grid_center = np.array(grid_dimesions) / 2.0
radii = np.arange(2.0, 10.0)
test_r_index = 5
test_r = radii[test_r_index]
print 'testing at r=%f' % test_r
L_max = 7
l = 3
m = 1
shg = SphHarmGrid(grid_dimesions, grid_center, radii, L_max)
u = np.abs(shg._grid_distances - test_r)
S = ( u < shg._delta )
test_grid = np.zeros(grid_dimesions)
test_grid[S] = sph_harm(l, m, shg._theta[S].flatten(),
shg._phi[S].flatten())
coeff = shg(test_grid)
# index is radius, l, m
print '\nl_slice'
print coeff[test_r_index,3,:]
print '\nm_slice'
print coeff[test_r_index,:,1]
print '\nall'
print coeff[test_r_index,:,:]
return
def test_sph_hrm():
phi = np.linspace(0.01, 2*np.pi-.01, 50)
theta = np.linspace(0.01, np.pi-.01, 50)
for l in range(2,5):
for m in range(-l, l+1):
ref = special.sph_harm(m, l, phi, theta)
Ylm = math2.sph_harm(l, m, theta, phi)
#ref /= ref[0]
#Ylm /= Ylm[0]
# print l, m, (ref-Ylm) / ref
assert_allclose(np.real(ref), np.real(Ylm), atol=1e-6, rtol=1e-4)
assert_allclose(np.imag(ref), np.imag(Ylm), atol=1e-6, rtol=1e-4)
def radial_sph_harm(n, l, m, u, theta, phi, Rtype='legendre', **kwargs):
"""
Return the radially augmented spherical harmonic Y_nlm evaluated at
(r, theta, phi).
These functions are the regurlar spherical harmonics with a Legendre
radial component:
Y_nlm = R_n * Y_lm
R_n(u) = u^n
Where P_n are the Legendre polynomials
"""
if not (theta.shape == phi.shape) and (theta.shape == u.shape):
raise ValueError('`u`, `theta`, and `phi` must all have the same shape!')
Ylm = sph_harm(l, m, theta, phi)
if Rtype == 'monomial':
Rn = np.power(u, n)
elif Rtype == 'legendre':
if not 'delta' in kwargs.keys():
raise TypeError('with Rtype=legendre, you must provide a `delta` '
'kwarg parameter.')
if not 'radius' in kwargs.keys():
raise TypeError('with Rtype=legendre, you must provide a `radius` '
'kwarg parameter.')
d = kwargs['delta']
r0 = kwargs['radius']
f = special.legendre(n)
Rn = np.sqrt( (2.0*n+1.0) / (2.0*d) ) / (d * u + r0) * f(u)
else:
raise NotImplementedError('No known implementation for radial basis '
'function type: `%s`' % Rtype)
return Rn * Ylm
def _compute_projections(self):
"""
Compute the spherical harmonic projection vectors for the specific grid.
"""
for ir, r in enumerate(self._radii):
logger.debug('Computing projections for radius: %f' % r)
S, u = self._compute_shell(r)
# determine weights, tricube rule
w = np.power(1 - np.power(np.abs(u)[S], 3), 3)
W = np.diag( w.astype(np.complex128) )
# compute Y, the N x M matrix where each column is the function
# Y_nlm evaluated at all the points S in the shell
N = np.sum(S)
M = len(self.n_values) * self._L_max ** 2
Y = np.zeros((N, M), dtype=np.complex128)
logger.debug('Computing %d x %d (points x functions) design matrix '
'Y' % (N, M))
for n in self.n_values:
for l in range(self._L_max):
for m in range(-l, l+1):
i = self._projection_index(n, l, m)
logger.debug('computing projection for %d-th function, '
'(n l m) = (%d %d %d)' % (i, n, l, m))
if self.Rtype == 'flat':
if n != 0:
raise ValueError('n must be 0 for Rtype `flat` (got n=%d)' % n)
Y[:,i] = sph_harm(l, m, self._theta[S], self._phi[S])
elif self.Rtype == 'monomial':
Y[:,i] = radial_sph_harm(n, l, m,
u[S],
self._theta[S],
self._phi[S],
Rtype='monomial')
elif self.Rtype == 'legendre':
Y[:,i] = radial_sph_harm(n, l, m,
u[S],
self._theta[S],
self._phi[S],
Rtype='legendre',
delta=self._delta,
radius=r)
else:
raise ValueError('Unknown Rtype: `%s`' % Rtype)
# perform weighted LSQ
G = np.dot(np.conjugate(Y).T, np.dot(W, Y))
B = np.dot(np.conjugate(Y).T, W)
assert G.shape == (M, M)
assert B.shape == (M, N)
P = np.linalg.solve(G, B)
assert P.shape == (M, N)
# store the projection for later use
self._projections[r] = P
return
def sph_harm_coefficients(trajectory, q_values, weights=None,
num_coefficients=10):
"""
Numerically evaluates the coefficients of the projection of a structure's
fourier transform onto the three-dimensional spherical harmonic basis. Can
be used to directly compare a proposed structure to the same coefficients
computed from correlations in experimentally observed scattering profiles.
Parameters
----------
trajectory : mdtraj.trajectory
A trajectory object representing a Boltzmann ensemble.
q_values : ndarray, float
A list of the reciprocal space magnitudes at which to evaluate
coefficients.
weights : ndarray, float
A list of weights, for how to weight each snapshot in the trajectory.
If not provided, treats each snapshot with equal weight.
num_coefficients : int
The order at which to truncate the spherical harmonic expansion
Returns
-------
sph_coefficients : ndarray, float
A 3-dimensional array of coefficients. The first dimension indexes the
order of the spherical harmonic. The second two dimensions index the
array `q_values`.
References
----------
.[1] Kam, Z. Determination of macromolecular structure in solution by
spatial correlation of scattering fluctuations. Macromolecules 10, 927-934
(1977).
"""
logger.debug('Projecting image into spherical harmonic basis...')
# first, deal with weights
if weights == None:
weights = np.ones(trajectory.n_frames)
else:
if not len(weights) == trajectory.n_frames:
raise ValueError('length of `weights` array must be the same as the'
'number of snapshots in `trajectory`')
weights /= weights.sum()
# initialize the q_values array
q_values = np.array(q_values).flatten()
num_q_mags = len(q_values)
# don't do odd values of ell
l_vals = range(0, 2*num_coefficients, 2)
# initialize spherical harmonic coefficient array
# note that it's 4* num_coeff - 3 b/c we're skipping odd l's -- (2l+1)
Slm = np.zeros(( num_coefficients, 4*num_coefficients-3, num_q_mags),
dtype=np.complex128 )
# get the quadrature vectors we'll use, a 900 x 4 array : [q_x, q_y, q_z, w]
# from thor.refdata import sph_quad_900
q_phi = arctan3(sph_quad_900[:,1], sph_quad_900[:,0])
q_theta = np.arccos(sph_quad_900[:,2])
# iterate over all snapshots in the trajectory
for i in range(trajectory.n_frames):
for iq,q in enumerate(q_values):
logger.info('Computing coefficients for q=%f\t(%d/%d)' % (q, iq+1, num_q_mags))
# compute S, the single molecule scattering intensity
S_q = simulate_shot(trajectory[i], 1, q * sph_quad_900[:,:3],
force_no_gpu=True)
# project S onto the spherical harmonics using spherical quadrature
for il,l in enumerate(l_vals):
for m in range(0, l+1):
logger.debug('Projecting onto Ylm, l=%d/m=%d' % (l, m))
# compute a spherical harmonic, turns out this is
# unexpectedly annoying...
# -----------
# option (1) : scipy (slow & incorrect?)
# scipy switched the convention for theta/phi in this fxn
#Ylm = special.sph_harm(m, l, q_phi, q_theta)
# -----------
# option (2) : roll your own
Ylm = sph_harm(l, m, q_theta, q_phi)
# -----------
# NOTE: we're going to use the fact that negative array
# indices wrap around here -- the value of m can be
# negative, but those values just end up at the *end*
# of the array
r = np.sum( S_q * Ylm * sph_quad_900[:,3] )
Slm[il, m, iq] = r
Slm[il, -m, iq] = ((-1) ** m) * np.conjugate(r)
# now, reduce the Slm solution to C_l(q1, q2)
sph_coefficients = np.zeros((num_coefficients, num_q_mags, num_q_mags))
for iq1, q1 in enumerate(q_values):
for iq2, q2 in enumerate(q_values):
for il, l in enumerate(l_vals):
ip = np.sum( Slm[il,:,iq1] * np.conjugate(Slm[il,:,iq2]) )
if not np.imag(ip) < 1e-6:
logger.warning('C_l coefficient has non-zero imaginary'
' component (%f) -- this is theoretically'
' forbidden and usually a sign something '
'went wrong numerically' % np.imag(ip))
sph_coefficients[il, iq1, iq2] += weights[i] * np.real(ip)
return sph_coefficients
