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