outputs the value of the angular wave function at (theta, phi) for any given (l, m)

angular_wave_function.py

"""
outputs the value of the associated Legendre function P_l,m(x) at a given point.
l and m must be non-negative integers.
"""
import numpy as np
from math import pi, sqrt, cos, factorial, e

def legendre(l):
    """returns a legendre polynomial in terms of cos(theta)"""
    if l == 0:
        return lambda x: 1 # return a funciton that returns 1
    p = np.poly1d([1, 0, -1]).__pow__(l) #create polynomial
    p_deriv = p.deriv(l) #get degree-l derivative of (x^2 - 1)^l
    #legendre_pol = np.poly1d([i/(2**l * math.factorial(l)) for i in p_deriv])
    legendre_pol = p_deriv/(2**l * factorial(l))

    def getpol(theta):
        """returns the value of the legendre polynomial
        at the point x = cos(theta)"""
        #only required if passing the output of legendre() as a function
        pol = 0
        for power, coeff in enumerate(legendre_pol):
            pol += coeff*(cos(theta)**power)
        return pol
    
    return legendre_pol #returns a poly1d object
    
def assoc_legendre(m,l):
    """returns a function in terms of cos(theta) that is 
    the associated Legendre function for m,l"""
    if l == 0 and m == 0:
        return lambda x: 1 #returns a function that returns 1

    elif abs(m) <= l:
        legendre_pol = legendre(l)
        legendre_pol_deriv = legendre_pol.deriv(abs(m)) if m != 0 else legendre_pol

        def getassoc(theta):
            """returns the associated Legendre function at
            the point x = cos(theta)"""
            return legendre_pol_deriv(cos(theta))*((1 - cos(theta)**2)**abs(m/2.0))
        return getassoc #returns a function in terms of x

    else:
        return lambda x: "m cannot be greater than l" 
        #returns a null-returning function if |m| > l


def angular_wave_func(m, l, theta, phi):
    eps = -1**m if m > 0 else 1
#    print "\nepsilon = " + str(eps) + '\n'
    expr = sqrt((((2*l) + 1) * factorial(l - abs(m))) / ((4.0*pi) * factorial(l + abs(m))))
#    print "expr value for for m = %d, l = %d: \n %.3f \n" % (m, l, expr)
    p = assoc_legendre(m, l)(theta)
#    print "assoc_legendre value for m = %d, l = %d at x = %.2f: %.3f \n" %(m,l,cos(theta),p)
    e_part = (e**(m*phi*1j))
#    print "Exponential part:"
#    print str(round(e_part.real, 5)) + "   " + str(round(e_part.imag, 5)) + 'j\n'
    y = eps * expr * e_part * p
    return np.round(y, 5)

print legendre(1)
print assoc_legendre(1, 1)(pi/6)

print 'angular_wave_func(0,0,0,0)'
ans=angular_wave_func(0,0,0,0)
print ans
print 'angular_wave_func (0,1,c.pi ,0)'
ans=angular_wave_func (0,1,pi ,0)
print ans
print 'angular_wave_func(1,1,c.pi/2,c.pi)'
ans=angular_wave_func(1,1,pi/2,pi)
print ans
print 'angular_wave_func (0,2,c.pi ,0)'
ans=angular_wave_func (0,2,pi ,0)
print ans