1# This file is part of QuTiP: Quantum Toolbox in Python. 2# 3# Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson. 4# All rights reserved. 5# 6# Redistribution and use in source and binary forms, with or without 7# modification, are permitted provided that the following conditions are 8# met: 9# 10# 1. Redistributions of source code must retain the above copyright notice, 11# this list of conditions and the following disclaimer. 12# 13# 2. Redistributions in binary form must reproduce the above copyright 14# notice, this list of conditions and the following disclaimer in the 15# documentation and/or other materials provided with the distribution. 16# 17# 3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names 18# of its contributors may be used to endorse or promote products derived 19# from this software without specific prior written permission. 20# 21# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 22# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 23# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A 24# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 25# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 26# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 27# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 28# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 29# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 30# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 31# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 32############################################################################### 33 34__all__ = ['orbital'] 35 36import numpy as np 37from scipy.special import factorial 38 39 40def orbital(theta, phi, *args): 41 r"""Calculates an angular wave function on a sphere. 42 ``psi = orbital(theta,phi,ket1,ket2,...)`` calculates 43 the angular wave function on a sphere at the mesh of points 44 defined by theta and phi which is 45 :math:`\sum_{lm} c_{lm} Y_{lm}(theta,phi)` where :math:`C_{lm}` are the 46 coefficients specified by the list of kets. Each ket has 2l+1 components 47 for some integer l. 48 49 Parameters 50 ---------- 51 theta : list/array 52 Polar angles 53 54 phi : list/array 55 Azimuthal angles 56 57 args : list/array 58 ``list`` of ket vectors. 59 60 Returns 61 ------- 62 ``array`` for angular wave function 63 64 """ 65 psi = 0.0 66 if isinstance(args[0], list): 67 # use the list in args[0] 68 args = args[0] 69 70 for k in range(len(args)): 71 ket = args[k] 72 if not ket.type == 'ket': 73 raise TypeError('Invalid input ket in orbital') 74 sk = ket.shape 75 nchk = (sk[0] - 1) / 2.0 76 if nchk != np.floor(nchk): 77 raise ValueError( 78 'Kets must have odd number of components in orbital') 79 l = int((sk[0] - 1) / 2) 80 if l == 0: 81 SPlm = np.sqrt(2) * np.ones((np.size(theta), 1), dtype=complex) 82 else: 83 SPlm = _sch_lpmv(l, np.cos(theta)) 84 fac = np.sqrt((2.0 * l + 1) / (8 * np.pi)) 85 kf = ket.full() 86 psi += np.sqrt(2) * fac * kf[l, 0] * np.ones((np.size(phi), 87 np.size(theta)), 88 dtype=complex) * SPlm[0] 89 for m in range(1, l + 1): 90 psi += ((-1.0) ** m * fac * kf[l - m, 0]) * \ 91 np.array([np.exp(1.0j * 1 * phi)]).T * \ 92 np.ones((np.size(phi), np.size(theta)), 93 dtype=complex) * SPlm[1] 94 for m in range(-l, 0): 95 psi = psi + (fac * kf[l - m, 0]) * \ 96 np.array([np.exp(1.0j * 1 * phi)]).T * \ 97 np.ones((np.size(phi), np.size(theta)), dtype=complex) * \ 98 SPlm[abs(m)] 99 return psi 100 101 102# Schmidt Semi-normalized Associated Legendre Functions 103def _sch_lpmv(n, x): 104 ''' 105 Outputs array of Schmidt Seminormalized Associated Legendre Functions 106 S_{n}^{m} for m<=n. 107 108 Parameters 109 ---------- 110 n : int 111 Degree of polynomial. 112 113 x : float 114 Point at which to evaluate 115 116 Returns 117 ------- 118 array of values for Legendre functions. 119 120 ''' 121 from scipy.special import lpmv 122 n = int(n) 123 sch = np.array([1.0]) 124 sch2 = np.array([(-1.0) ** m * np.sqrt( 125 (2.0 * factorial(n - m)) / factorial(n + m)) for m in range(1, n + 1)]) 126 sch = np.append(sch, sch2) 127 if isinstance(x, float) or len(x) == 1: 128 leg = lpmv(np.arange(0, n + 1), n, x) 129 return np.array([sch * leg]).T 130 else: 131 for j in range(0, len(x)): 132 leg = lpmv(range(0, n + 1), n, x[j]) 133 if j == 0: 134 out = np.array([sch * leg]).T 135 else: 136 out = np.append(out, np.array([sch * leg]).T, axis=1) 137 return out 138