21.5.2. Schrodinger model - Gallery

Show code cell source Hide code cell source

#> Libraries
import numpy as np

import matplotlib.pyplot as plt
import plotly.graph_objects as go
from plotly.subplots import make_subplots

import plotly.io as pio
from IPython.display import display, HTML

pio.renderers.default = 'iframe'  

from scipy.special import eval_genlaguerre, sph_harm_y, factorial
from scipy.interpolate import interpn

from skimage import measure             # requires scikit-image

21.5.2.1. Stationary states

\[\psi_{n,\ell,m_\ell}(r, \theta, \phi) = R_{n\ell}(r) \, Y_\ell^{m_\ell}(\theta, \phi)\]

Radial Wavefunctions, \(R_{n\ell}(r)\). Using the dimensionless variable \(\rho = \frac{2r}{n a_0}\),

\[R_{n\ell}(r) = -\sqrt{\left(\frac{2}{n a_0}\right)^3 \frac{(n-\ell-1)!}{2n [(n+\ell)!]^3}} e^{-\rho/2} \rho^\ell L_{n-\ell-1}^{2\ell+1}(\rho) \ ,\]

where \(a_0 = \frac{4\pi \varepsilon_0 \hbar^2}{m_e e^2}\) is the Bohr radius, and \(L^k_p(\rho)\) the solution of the associated Laguerre equation.

Spherical Harmonics, \(Y_\ell^{m_\ell}(\theta, \phi) = \Theta_{\ell, m_{\ell}}(\theta) \Phi_{m_{\ell}}(\phi)\).

\[Y_\ell^{m_\ell}(\theta, \phi) = (-1)^{m_\ell} \sqrt{\frac{(2\ell+1)}{4\pi} \frac{(\ell-m_\ell)!}{(\ell+m_\ell)!}} P_\ell^{m_\ell}(\cos\theta) \, e^{i m_\ell \phi}\]

with \(P^{m_{\ell}}_{\ell}\) the associated Legendre polynomials. The polar functions are real-valued and by definition \(\Theta_{\ell,-m_\ell}(\theta) = (-1)^{m_\ell} \Theta_{\ell,m_\ell}(\theta)\).

#> Spatial resolution and default parameters

Show code cell source Hide code cell source

# Build 3D spatial grid
box_size = 18.0
grid_pts = 50
x = np.linspace(-box_size, box_size, grid_pts)
y = np.linspace(-box_size, box_size, grid_pts)
z = np.linspace(-box_size, box_size, grid_pts)
X_def, Y_def, Z_def = np.meshgrid(x, y, z, indexing='ij')

#> Functions

#> Evaluate stationary states

Show code cell source Hide code cell source

def hydro_psi_complex(n, l, m, X, Y, Z):
    """Computes complex eigenfunction psi_{n,l,m}(x,y,z) in atomic units a0=1."""
    R = np.sqrt(X**2 + Y**2 + Z**2)
    R[R == 0] = 1e-10  # Avoid division by zero

    Theta = np.arccos(Z / R)       # Polar angle [0, pi]
    Phi = np.arctan2(Y, X)        # Azimuthal angle [-pi, pi]

    # Dimensionless variable rho
    rho = 2.0 * R / n
    degree = n - l - 1
    alpha = 2 * l + 1

    # Radial wavefunction R_{n,l}
    norm_R = (2.0 / n)**1.5 * np.sqrt(factorial(n - l - 1) / (2.0 * n * factorial(n + l)))
    lag_poly = eval_genlaguerre(degree, alpha, rho)
    R_nl = norm_R * np.exp(-0.5 * rho) * (rho**l) * lag_poly

    # Spherical harmonic Y_l^m(theta, phi)
    # SciPy's sph_harm takes arguments: sph_harm(m, l, azimuth_phi, polar_theta)
    Y_lm = sph_harm_y(l, m, Theta, Phi)

    return R_nl * Y_lm
    

#> Plot functions

#> Isosurface of probability density of stationary states, colored with phase

Show code cell source Hide code cell source

def plot_complex_eigenfunction_colored_by_phase(n, l, m, X=X_def, Y=Y_def, Z=Z_def, iso_val=0.001):
    psi = hydro_psi_complex(n, l, m, X, Y, Z)
    prob_density = np.abs(psi)**2
    phase = np.angle(psi)

    # 1. Extract marching cubes mesh
    dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
    verts, faces, normals, values = measure.marching_cubes(
        volume=prob_density, level=iso_val, spacing=(dx, dy, dz)
    )

    verts_x = x[0] + verts[:, 0]
    verts_y = y[0] + verts[:, 1]
    verts_z = z[0] + verts[:, 2]

    # 2. Compute the exact 3D center (centroid) of every triangular face
    face_centers_x = verts_x[faces].mean(axis=1)
    face_centers_y = verts_y[faces].mean(axis=1)
    face_centers_z = verts_z[faces].mean(axis=1)
    face_centroids = np.column_stack([face_centers_x, face_centers_y, face_centers_z])

    # 3. Interpolate cos and sin at the face centroids (smooth, non-jumping fields)
    interp_cos = interpn((x, y, z), np.cos(phase), face_centroids, method='linear')
    interp_sin = interpn((x, y, z), np.sin(phase), face_centroids, method='linear')

    # 4. Reconstruct phase angle for each face center
    face_phase = np.arctan2(interp_sin, interp_cos)

    # 5. Map phase directly to RGB strings for EACH FACE
    norm_phase = (face_phase + np.pi) / (2 * np.pi)
    cmap = plt.get_cmap('twilight')
    face_colors = [
        f'rgb({int(r*255)}, {int(g*255)}, {int(b*255)})'
        for r, g, b, _ in cmap(norm_phase)
    ]

    # 6. Pass facecolor to Plotly Mesh3d (renders flat per-face colors without interpolation)
    fig = go.Figure(data=[
        go.Mesh3d(
            x=verts_x,
            y=verts_y,
            z=verts_z,
            i=faces[:, 0],
            j=faces[:, 1],
            k=faces[:, 2],
            facecolor=face_colors,  # Colors each triangle uniformly, preventing face interpolation
            opacity=0.9
        )
    ])

    fig.update_layout(
        title=dict(
            text=f"Complex Eigenfunction ψ<sub>{n},{l},{m}</sub> (|ψ|² Isosurface colored by Phase)",
            x=0.5,
            xanchor='center'
        ),
        scene=dict(
            xaxis=dict(
                title='X [a₀]',
                # titlefont=dict(size=14, color='black'),
                # showgrid=True,
                # showbackground=True,
                # backgroundcolor="rgb(240, 240, 240)"
            ),
            yaxis=dict(
                title='Y [a₀]',
                # titlefont=dict(size=14, color='black'),
                # showgrid=True,
                # showbackground=True,
                # backgroundcolor="rgb(240, 240, 240)"
            ),
            zaxis=dict(
                title='Z [a₀]',
                # titlefont=dict(size=14, color='black'),
                # showgrid=True,
                # showbackground=True,
                # backgroundcolor="rgb(240, 240, 240)"
            ),
            # aspectmode='data'
        ),
        # margin=dict(l=50, r=50, b=50, t=60),  # Increase margins so labels aren't cut off
        width=720,
        height=720
    )

    return fig

Show code cell source Hide code cell source

def plot_real_orbitals(n, l, m, X=X_def, Y=Y_def, Z=Z_def, iso_val=0.001):

    figs = []
    
    if ( m == 0 ):
        figs += [ plot_complex_eigenfunction_colored_by_phase( \
            n, l, m, X=X, Y=Y, Z=Z, iso_val=iso_val ) ]

    else:
        
        psi_p = hydro_psi_complex(n, l, m, X, Y, Z)
        psi_m = hydro_psi_complex(n, l,-m, X, Y, Z)

        psi_list = [
            np.real(      ( psi_m + (-1)**m * psi_p ) / np.sqrt(2) ),
            np.real( 1j * ( psi_m - (-1)**m * psi_p ) / np.sqrt(2) )
        ]

        for psi in psi_list:
    
            prob_density = np.abs(psi)**2
            phase = np.angle(psi)
        
            # 1. Extract marching cubes mesh
            dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
            verts, faces, normals, values = measure.marching_cubes(
                volume=prob_density, level=iso_val, spacing=(dx, dy, dz)
            )
        
            verts_x = x[0] + verts[:, 0]
            verts_y = y[0] + verts[:, 1]
            verts_z = z[0] + verts[:, 2]
        
            # 2. Compute the exact 3D center (centroid) of every triangular face
            face_centers_x = verts_x[faces].mean(axis=1)
            face_centers_y = verts_y[faces].mean(axis=1)
            face_centers_z = verts_z[faces].mean(axis=1)
            face_centroids = np.column_stack([face_centers_x, face_centers_y, face_centers_z])
        
            # 3. Interpolate cos and sin at the face centroids (smooth, non-jumping fields)
            interp_cos = interpn((x, y, z), np.cos(phase), face_centroids, method='linear')
            interp_sin = interpn((x, y, z), np.sin(phase), face_centroids, method='linear')
        
            # 4. Reconstruct phase angle for each face center
            face_phase = np.arctan2(interp_sin, interp_cos)
        
            # 5. Map phase directly to RGB strings for EACH FACE
            norm_phase = (face_phase + np.pi) / (2 * np.pi)
            cmap = plt.get_cmap('twilight')
            face_colors = [
                f'rgb({int(r*255)}, {int(g*255)}, {int(b*255)})'
                for r, g, b, _ in cmap(norm_phase)
            ]
        
            # 6. Pass facecolor to Plotly Mesh3d (renders flat per-face colors without interpolation)
            fig = go.Figure(data=[
                go.Mesh3d(
                    x=verts_x,
                    y=verts_y,
                    z=verts_z,
                    i=faces[:, 0],
                    j=faces[:, 1],
                    k=faces[:, 2],
                    facecolor=face_colors,  # Colors each triangle uniformly, preventing face interpolation
                    opacity=0.9
                )
            ])
        
            fig.update_layout(
                title=dict(
                    # text=f"Complex Eigenfunction ψ<sub>{n},{l},{m}</sub> (|ψ|² Isosurface colored by Phase)",
                    x=0.5,
                    xanchor='center'
                ),
                scene=dict(
                    xaxis=dict(
                        title='X [a₀]',
                        # titlefont=dict(size=14, color='black'),
                        # showgrid=True,
                        # showbackground=True,
                        # backgroundcolor="rgb(240, 240, 240)"
                    ),
                    yaxis=dict(
                        title='Y [a₀]',
                        # titlefont=dict(size=14, color='black'),
                        # showgrid=True,
                        # showbackground=True,
                        # backgroundcolor="rgb(240, 240, 240)"
                    ),
                    zaxis=dict(
                        title='Z [a₀]',
                        # titlefont=dict(size=14, color='black'),
                        # showgrid=True,
                        # showbackground=True,
                        # backgroundcolor="rgb(240, 240, 240)"
                    ),
                    # aspectmode='data'
                ),
                # margin=dict(l=50, r=50, b=50, t=60),  # Increase margins so labels aren't cut off
                width=720,
                height=720
            )

            figs += [ fig ]

    return figs
    

#> Plot a set of stationary states

Show code cell source Hide code cell source

nmin, nmax = 1, 4

for n in range(1, nmax + 1):
    for ell in range(0, n):
        ml_vals = np.arange(-ell, ell + 1)
        num_m = len(ml_vals)

        # Create 1 row with (2*ell + 1) 3D subplots
        fig = go.Figure()
        fig = make_subplots(
            rows=1,
            cols=num_m,
            subplot_titles=[f"m = {ml}" for ml in ml_vals],
            specs=[[{"type": "scene"} for _ in range(num_m)]]
        )

        for idx, ml in enumerate(ml_vals, start=1):
            # Compute/get traces for this specific orbital
            # Assuming your helper function can return traces or a figure
            if ( n == 4 ):
                iso_val = 0.00005
            else:
                iso_val = 0.0001
            sub_fig = plot_complex_eigenfunction_colored_by_phase(
                n=n, l=ell, m=ml, iso_val=iso_val
            )

            # Add all 3D traces from sub_fig into the corresponding grid cell
            for trace in sub_fig.data:
                fig.add_trace(trace, row=1, col=idx)

        # Update layout title and dimensions
        fig.update_layout(
            title_text=f"Hydrogen Orbitals for n = {n}, ℓ = {ell}",
            height=450,
            width=450 * num_m,
            showlegend=False
        )

        # 1. Force a unique filename for every (n, l) pair so they don't overwrite each other
        unique_id = f"orbital_n{n}_l{ell}"
        
        # 2. Render explicitly using display() in each iteration
        # Passing a unique div_id prevents Plotly JS from overriding previous DOM containers
        html_str = pio.to_html(fig, include_plotlyjs='cdn', full_html=False, div_id=unique_id)
        display(HTML(html_str))

21.5.2.2. Real orbitals

#> Plot a set of real orbitals

Show code cell source Hide code cell source

nmin, nmax = 1, 4

for n in range(1, nmax + 1):
    for ell in range(0, n):
        ml_vals = np.arange(0, ell + 1)
        num_m = len(ml_vals) * 2 - 1

        idx = 1

        # Create 1 row with (2*ell + 1) 3D subplots
        fig = go.Figure()
        subplot_titles = num_m * [""]
        subplot_titles[0] = f"m = 0"
        if ( ell > 0 ):
            subplot_titles[1::2] = [f"m = {ml}" for ml in ml_vals[1:]]
        
        fig = make_subplots(
            rows=1,
            cols=num_m,
            subplot_titles=subplot_titles,
            specs=[[{"type": "scene"} for _ in range(num_m)]]
        )

        for ml in ml_vals:
            # Compute/get traces for this specific orbital
            # Assuming your helper function can return traces or a figure
            sub_figs = plot_real_orbitals(
                n=n, l=ell, m=ml, iso_val=0.0001 if n < 4 else 0.00005
            )

            # Add all 3D traces from sub_fig into the corresponding grid cell
            # print(sub_figs)
            for sub_fig in sub_figs:
                for trace in sub_fig.data:
                    fig.add_trace(trace, row=1, col=idx)
                    idx += 1

        # Update layout title and dimensions
        fig.update_layout(
            title_text=f"Real Orbitals for n = {n}, ℓ = {ell}",
            height=450,
            width=450 * num_m,
            showlegend=False
        )

        # 1. Force a unique filename for every (n, l) pair so they don't overwrite each other
        unique_id = f"real_orbital_n{n}_l{ell}"
        
        # 2. Render explicitly using display() in each iteration
        # Passing a unique div_id prevents Plotly JS from overriding previous DOM containers
        html_str = pio.to_html(fig, include_plotlyjs='cdn', full_html=False, div_id=unique_id)
        display(HTML(html_str))

21.5.2.3. Some details

As \(\Theta_{\ell, -m_\ell} = (-1)^{m_{\ell}} \Theta_{\ell, m_\ell}\), it follows that

\[\begin{split}\begin{cases} Y_{\ell}^{m_\ell} + (-1)^{m_{ell}} Y_{\ell}^{-m_\ell} = \Theta_{\ell, m_{\ell}}(\theta) \left[ e^{i m_{\ell} \phi} + e^{-i m_{\ell} \phi} \right] \propto \cos( m_{\ell} \theta ) \\ Y_{\ell}^{m_\ell} - (-1)^{m_{ell}} Y_{\ell}^{-m_\ell} = \Theta_{\ell, m_{\ell}}(\theta) \left[ e^{i m_{\ell} \phi} - e^{-i m_{\ell} \phi} \right] \propto i \sin( m_{\ell} \theta ) \ , \end{cases}\end{split}\]

and thus the real orbitals as the linear combinations

\[\begin{split}\psi_{n,\ell,m_\ell, \text{real}} = \begin{cases} \psi_{n,\ell,0} & \text{for } m_\ell = 0 \\ \frac{1}{\sqrt{2}} \left( \psi_{n,\ell,-\vert{}m_\ell\vert{}} + (-1)^{m_\ell} \psi_{n,\ell,\vert{}m_\ell\vert{}} \right) \propto \cos(\vert{}m_\ell\vert{}\phi) & \text{for } \text{real orbital } (x, xy, x^2-y^2, \dots) \\ \frac{i}{\sqrt{2}} \left( \psi_{n,\ell,-\vert{}m_\ell\vert{}} - (-1)^{m_\ell} \psi_{n,\ell,\vert{}m_\ell\vert{}} \right) \propto \sin(\vert{}m_\ell\vert{}\phi) & \text{for } \text{real orbital } (y, yz, \dots) \end{cases}\end{split}\]

21.5.2.3.1. Radial and spherical factors

\[\psi_{n,\ell,m_\ell}(r, \theta, \phi) = R_{n\ell}(r) \, Y_\ell^{m_\ell}(\theta, \phi)\]

21.5.2.3.1.1. Radial wavefunctions, \(R_{n, \ell}(r)\)

Radial Wavefunctions, \(R_{n\ell}(r)\). Using the dimensionless variable \(\rho = \frac{2r}{n a_0}\),

\[R_{n\ell}(r) = -\sqrt{\left(\frac{2}{n a_0}\right)^3 \frac{(n-\ell-1)!}{2n [(n+\ell)!]^3}} e^{-\rho/2} \rho^\ell L_{n-\ell-1}^{2\ell+1}(\rho) \ ,\]

where \(a_0 = \frac{4\pi \varepsilon_0 \hbar^2}{m_e e^2}\) is the Bohr radius, and \(L^k_p(\rho)\) the solution of the associated Laguerre equation.

21.5.2.3.1.2. Polar and azimuthal wavefunctions, \(Y^{m_\ell}_{\ell}(\theta,\phi) = \Theta_{\ell, m_\ell}(\theta) \Phi_{m_\ell}(\phi)\)

\[Y_\ell^{m_\ell}(\theta, \phi) = (-1)^{m_\ell} \sqrt{\frac{(2\ell+1)}{4\pi} \frac{(\ell-m_\ell)!}{(\ell+m_\ell)!}} P_\ell^{m_\ell}(\cos\theta) \, e^{i m_\ell \phi}\]

with \(P^{m_{\ell}}_{\ell}\) the associated Legendre polynomials. The first associated Legendre polynomials are

\(\ell\)	\(m_\ell\)	\(P^{m_{\ell}}_{\ell}(x)\)
\(0\)	\(0\)	\(1\)
\(1\)	\(1\)	\(-(1-x^2)^{\frac{1}{2}}\)
	\(0\)	\(x\)
	\(-1\)	\(-\frac{1}{2}P_1^1(x)\)
\(2\)	\(2\)	\( 3 ( 1 -x^2 )\)
	\(1\)	\(- 3 x ( 1 - x^2 )^{\frac{1}{2}}\)
	\(0\)	\(\frac{1}{2}( 3 x^2 - 1)\)
	\(-1\)	\(- \frac{1}{6} P_{2}^{1}(x)\)
	\(-2\)	\(\frac{1}{24} P_2^2(x)\)