<jupyter_start><jupyter_text>Hydrogen orbital isosurfaces© 2020 <NAME> and Department of Chemistry, University of California, Davis. All rights reserved.This Jupyter notebook was written for CHE 2AH, Fall 2020.  It uses Sympy to define the hydrogen atom wavefunctions and PyVista for plotting isosurfaces.How to use: 1. Evaluate the cell directly below this one to import the libraries and define functions.  If you're interested in coding, feel free to read through the code to get a better handle on what's going on.2. In the cell below the one containing the codes, set the *n*, *l*, and *m* quantum numbers (be sure that 0 < *l* < *n*, and -*l* < *m* < +*l*), and then evaluate it.  The three-dimensional plot should appear.  Click the "three bars" in the upper left to set the color map, opacity, change the viewpoint, etc.Note: In the interactive plot, the colors of the x, y, and z axes are red, yellow and green respectively.<jupyter_code>import numpy as np
from sympy.physics.hydrogen import Psi_nlm
from sympy import Symbol, lambdify
import pyvista as pv

pv.plotting.rcParams['cmap'] = 'jet'

def get_wavefunction(n, l, m):
    """
    Get functions to evaluate hydrogen atom wavefunction.

    Parameters
    ----------
    n, l, m : int
        Quantum numbers (must obey 0 < L < n, and -L < m < +L)
        (L is capitalized because otherwise it looks like the number 1).
        
    Returns
    -------
    f : function
        This function takes as input three array arguments containing
        r, theta, and phi values on a grid, and produce the hydrogen atom 
        wavefunction values corrresponding to n, l, m quantum numbers.
    """
    r=Symbol("r", real=True, positive=True)
    phi=Symbol("phi", real=True)
    theta=Symbol("theta", real=True)
    # Obtain the symbolic function.
    expr=Psi_nlm(n, l, m, r, phi, theta, 1)
    # The numerical function should take r, theta, phi in the expected order
    f = lambdify([r, theta, phi], expr)
    return f

def getOrbitalIsosurface(n, l, m, make_real=False, isovalue=None, resolution=101):
    """
    Get hydrogen orbital isosurfaces.
    
    Parameters
    ----------
    n, l, m : int
        Quantum numbers (must obey 0 < L < n, and -L < m < +L).
        
    make_real : bool
        When set to True: 
        Take the linear combination (e^imϕ + e^-imϕ)/√2 or (e^imϕ + e^-imϕ)/√2i 
        for positive and negative "m" to provide the "real" linear combinations 
        that are familiar to chemists. Both positive and negative isosurfaces
        will be given.
        
        When set to False:
        Do not take the linear combination and return the single orbital for
        the provided value of "m". Create a single isosurface corresponding to the norm
        and color it by the argument of the complex number.
        
    isovalue : float
        If provided, the isovalue for which to plot the functions.
        If not provided, a default value is used which is 0.05/2^n
        (i.e. 0.05, 0.025, 0.0125, 0.00625, 0.003125 for n=1, 2, 3, 4, 5)        
        
    resolution : int
        The number of points on the grid in each dimension. 
        Above 200, performance starts to get spotty.

    Returns
    ----------
    pl : pyvista.PlotterITK object
        Run pl.show(True) on the returned object to make the 3D plot
    """

    # If the isovalue is not provided, set it automatically.
    if isovalue is None:
        isovalue = 0.05/2**n
        
    # Set the plot limits (orbital size increases with n)
    if n == 1:
        limit = 5
    elif n == 2:
        limit = 10
    else:
        limit = 20*(n-2)

    # Create grid of x, y, z values.
    rj=resolution*1j
    x, y, z = np.mgrid[-limit:limit:rj, -limit:limit:rj, -limit:limit:rj]
    
    # Calculate r, theta, phi. The +1e-10 is to avoid singularities
    r = np.sqrt(x**2+y**2+z**2)
    theta = np.arccos(z/(r+1e-10))
    phi = np.arctan2(y, x+1e-10)

    # Create PlotterITK and StructuredGrid object
    pl = pv.PlotterITK()
    grid = pv.StructuredGrid(x, y, z)

    if make_real:
        # Evaluate +|m| and -|m| wavefunctions, and take linear combinations
        f_p = get_wavefunction(n, l, abs(m))
        f_m = get_wavefunction(n, l, -abs(m))
        if m == 0:
            f_array = f_p(r, theta, phi)
        elif m > 0:
            f_array = (f_p(r, theta, phi) + f_m(r, theta, phi))/np.sqrt(2)
        elif m < 0:
            f_array = (f_p(r, theta, phi) - f_m(r, theta, phi))/np.sqrt(2)
        # Take the real and imaginary parts; only one of them should be nonzero
        f_real = np.real(f_array)
        f_imag = np.imag(f_array)
        if np.linalg.norm(f_imag) > 1e-10 and np.linalg.norm(f_real) > 1e-10:
            print(np.linalg.norm(f_imag), np.linalg.norm(f_real))
            raise RuntimeError("At this point, only one of the real or imaginary parts should be nonzero")
        # Because only one part is nonzero, adding them just keeps the nonzero part
        f_data = f_real + f_imag
        # For some reason, flattening in the "normal" order reverses the axes ordering
        grid["vol"] = f_data.flatten(order='F')
        # Create two isosurfaces for positive and negative isovalues
        # (This is because I don't know how to set an initial colormap in the plotter window
        #  and the default colors don't look very good.)
        isosurface1 = grid.contour([isovalue])
        isosurface2 = grid.contour([-isovalue])
        isosurface3 = grid.contour([0])
        pl.add_mesh(isosurface1,color='orange',opacity=1.0)
        pl.add_mesh(isosurface2,color='white',opacity=1.0)
        pl.add_mesh(isosurface3,color='gray',opacity=0.4)
    else:
        # Evaluate wavefunction for the given m
        f = get_wavefunction(n, l, m)
        f_array = f(r, theta, phi)
        # Norm of the complex function
        f_abs = np.absolute(f_array)
        grid["vol"] = f_abs.flatten(order='F')
        isosurface = grid.contour([isovalue])
        # Evaluate the argument of the complex number
        # on the isosurface points and use that as the color map
        x, y, z = (isosurface.points[:,0], isosurface.points[:,1], isosurface.points[:,2])
        r = np.sqrt(x**2+y**2+z**2)
        theta = np.arccos(z/(r+1e-10))
        phi = np.arctan2(y, x+1e-10)
        f_mesh = f(r, theta, phi)
        f_arg = np.angle(f_mesh)
        pl.add_mesh(isosurface, scalars=f_arg)
        
    return pl
#===============================#
#| You may modify these values |#
#===============================#
n = 2
l = 1
m = 0
pl = getOrbitalIsosurface(n, l, m, make_real=True)
pl.show(True)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11353711790393013, 'fraction_numerical': 0.018822466496009636, 'mean_word_length': 2.523607427055703, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '211059', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 2284, 'n_tokens_neox': 2126, 'n_words': 771, 'path': 'Hydrogen_Orbitals.ipynb', 'repo_name': 'leeping/hydrogen-orbitals', 'size': '6648'}