<jupyter_start><jupyter_text>2D  Simple Harmonic Oscillator<jupyter_code>import numpy as np
from scipy.constants import physical_constants as pc
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
import plotly
import plotly.graph_objects as go
amu_to_au = 1.0/(pc['kilogram-atomic mass unit relationship'][0]*pc['atomic unit of mass'][0])
            #1822.888479031408
hartree_to_cm1 = pc['hartree-hertz relationship'][0]/pc['speed of light in vacuum'][0]/100.0 
            #2.194746313705e5
sec = pc['atomic unit of time'][0]  #  2.418884326505e-17
cee = 100.0*pc['speed of light in vacuum'][0] #  2.99792458e10 cm/s
bohr_to_angstroms = pc['atomic unit of length'][0]*1e10
mumass = 1.0
we = 400.00   # Harmonic frequency in wavenumbers

nmu = mumass*amu_to_au
nfreq = we*sec*cee # angular harmonic frequency in au 
nkeq = ((2.0*np.pi*nfreq)**2)*nmu  # force constant in au
k = -1.0/(2.0*nmu)   # Use this if you want to change the units (don't)
def potential(point, k):
    x, y = point
    r = np.sqrt(x*x + y*y)
    u = 1/2 * k * (r) ** 2
    return u

def gen_V(x, u):
    """
    Assemble the matrix representation of the potential energy
    """
    V = np.zeros((NUM*NUM, NUM*NUM)).astype(np.complex)
    for m in range(NUM*NUM):
        V[m, m] = u[m]
    return(V)


def gen_T(xy, k):
    dx = xy[1][0] - xy[0][0]
    dy = xy[NUM][1] - xy[0][1]
    
    T = np.zeros((NUM*NUM, NUM*NUM)).astype(np.complex)
    row_len = int(np.sqrt(NUM*NUM))

    for m in range(0, NUM*NUM):
        for n in [m-NUM, m-1, m, m+1, m+NUM]:
            if (n >= 0 and n < NUM*NUM):
                factor = laplace_factor(m, n, row_len, dx, dy)
                T[m, n] = k * factor
    return(T)


def convert(xy, row_len):
    return xy//row_len, xy%row_len


def laplace_factor(idxy1, idxy2, row_len, dx, dy):    

    idx1, idy1 = convert(idxy1, row_len)
    idx2, idy2 = convert(idxy2, row_len)
    iddx = np.abs(idx2 - idx1) 
    iddy = np.abs(idy2 - idy1)
    
    if (iddx == 1 and iddy == 0):
        return 1 / (dx ** 2) 
    elif (iddx == 0 and iddy == 1):
        return 1 / (dy ** 2) 
    elif (iddx == 0 and iddy == 0):
        return -2 / (dy ** 2) -2 / (dx ** 2) 
    else:
        return 0


def solve_eigenproblem(H):
    """
    Solve an eigenproblem and return the eigenvalues and eigenvectors.
    """
    vals, vecs = np.linalg.eigh(H)
    idx = np.real(vals).argsort()
    vals = vals[idx]
    vecs = vecs.T[idx]
    return(vals, vecs)
NUM = 60
half_r = 2

x = np.linspace(-half_r, half_r, NUM)
y = np.linspace(-half_r, half_r, NUM)
x, y = np.meshgrid(x, y)
xy = np.stack((x, y), axis=-1)

xy = np.reshape(xy, (-1, 2))
xy.shape
%%time
u = np.zeros(xy.shape[0])

for i, point in enumerate(xy):
    u[i] = potential(point, nkeq)
    
V = gen_V(xy, u)
# plot potential
fig = go.Figure(data=[go.Surface(z=u.reshape(NUM, NUM))])

fig.update_layout(title='Potential Surface', autosize=False,
                  width=500, height=500,
                  margin=dict(l=65, r=50, b=65, t=90))

fig.show()
%%time
T = gen_T(xy, k)
T.shape
plt.imshow(T.real)
plt.imshow(T[:int(1.5*NUM), :int(1.5*NUM)].real)
%%time
H = T + V
evals, evecs = solve_eigenproblem(H)
evals.imag  #  Good  real eigenvalues
plt.plot(evals.real)
xy = xy.reshape(NUM, NUM, 2)

wave = evecs[0].reshape([NUM, NUM])

fig = go.Figure(data=[go.Surface(z=wave.real)])

fig.update_layout(title='Wavefunction', autosize=False,
                  width=500, height=500,
                  margin=dict(l=65, r=50, b=65, t=90))

fig.show()
wave = evecs[1].reshape([NUM, NUM])

fig = go.Figure(data=[go.Surface(z=wave.real)])

fig.update_layout(title='Wavefunction', autosize=False,
                  width=500, height=500,
                  margin=dict(l=65, r=50, b=65, t=90))

fig.show()
wave = evecs[15].reshape([NUM, NUM])

fig = go.Figure(data=[go.Surface(z=wave.real)])

fig.update_layout(title='Wavefunction', autosize=False,
                  width=500, height=500,
                  margin=dict(l=65, r=50, b=65, t=90))

fig.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1641938674579624, 'fraction_numerical': 0.050939663699307615, 'mean_word_length': 2.9082125603864735, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, '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': '620623', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 1787, 'n_tokens_neox': 1642, 'n_words': 384, 'path': 'notebooks/pchem/SHO-2D.ipynb', 'repo_name': 'yueyericardo/data_visualization_chem', 'size': '4044'}