3D spherical harmonic visualization capable superposing multiple user-defined Ylm functions.

sph_harm_viz.py

import numpy as np
from mpl_toolkits.mplot3d import Axes3D 
from matplotlib import cm, pyplot as plt
from scipy.special import sph_harm # (very convenient)

theta = np.linspace(0, 2*np.pi, 361) # define independent vars
phi   = np.linspace(0, np.pi, 181)

THETA, PHI = np.meshgrid(theta, phi) # 2D versions of the independent variables
XYZ = np.array([np.sin(PHI) * np.sin(THETA), np.sin(PHI) * np.cos(THETA), np.cos(PHI)]) 


def plot_sph_harm(*params): # pass a tuple of (c,l,m) for each harmonic to be summed
    orbitals = "spdf"
    fig = plt.figure(figsize=(7,7)) # init figure
    ax = fig.add_subplot( 111 , projection='3d') # init 3D axis
    Y = 0 # init function value
    titlestr = "" # init plot title
    
    for term in params:
        (c,l,m) = term
        if type(l) != int:
            l = orbitals.index(l)
        Y_lm = c*sph_harm(m, l, THETA, PHI) # calculate the spherical harmonic solution for this m and l combo
        Y += Y_lm # add it to the full thing
        titlestr += "+%.2fY$_{%i}^{%i}$" %(c,l,m)
    
    [X,Y,Z] = np.abs(Y.real)*XYZ # take the real part's absolute value and map it back into cartesian
    colormap = cm.ScalarMappable(cmap=plt.get_cmap("Wistia_r"))
    axislim = 0.25
    ax.plot_surface(X, Y, Z, facecolors=colormap.to_rgba(np.sqrt(X*X+Y*Y+Z*Z)), rstride=1, cstride=1) # plot ittttttt
    ax.set_xlim(-axislim,axislim)
    ax.set_ylim(-axislim,axislim)
    ax.set_zlim(-axislim,axislim)
    ax.set_axis_off() # turn off the grid and axes to make it pretty
    fig.suptitle("Laplace spherical harmonic ($\mathbb{R}^3$): "+titlestr[1:])
    plt.show()
    
plot_sph_harm((0.8,0,0),(0.1,9,6),(0.05,6,0)) # pass a 3-tuple of (c,l,m) for each harmonic (of magnitude c, degree l, and order m) to be summed

visualization.md