separable_eigenstates3d.py

import numpy as np
from numpy.linalg import eigh
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from skimage import measure


n: int = 256
length: float = 64
x: np.ndarray = np.linspace(-length/2, length/2-length/(n), n,
                            dtype=np.float)
dx: float = x[1] - x[0]
V: np.ndarray = 50.0*(x/length)**2
hbar: float = 1.0
m: float = 1.0


def get_separable_hamiltonian(Vx: np.ndarray) -> np.ndarray:
    H: np.ndarray = np.zeros((n, n), dtype=np.float)
    for i in range(n):
        if i-1 >= 0: H[i-1, i] = (-hbar**2/(2*m))/dx**2
        H[i, i] = Vx[i] - 2.0*(-hbar**2/(2*m))/dx**2
        if i+1 < n: H[i+1, i] = (-hbar**2/(2*m))/dx**2
    return H

isw = [np.zeros([n]) for _ in range(3)]
sho = [a*(x/length)**2 for a in [40.0, 40.0, 40.0]]
pot1 = [4.0*(1.0 - np.cos(2.0*np.pi*x/length)), 40.0*(x/length)**2,
        2*abs(x/length)]

Hx, Hy, Hz = [get_separable_hamiltonian(V) for V in sho]
ex, vx = eigh(Hx)
ey, vy = eigh(Hy)
ez, vz = eigh(Hz)

ex0 = np.multiply.outer(vx.T[0], np.outer(vy.T[1], vz.T[0]))**2/np.sqrt(2.0)
ex0 += np.multiply.outer(vx.T[1], np.outer(vy.T[0], vz.T[0]))**2/np.sqrt(2.0)
# ex0 += np.multiply.outer(vx.T[0], np.outer(vy.T[0], vz.T[1]))**2/np.sqrt(3.0)
ex_04 = ex[0:4]
ey_04 = ey[0:4]
ez_04 = ez[0:4]
energies = np.multiply.outer(ex_04, np.outer(ey_04, ez_04))
# print(energies[1, 0, 1])

verts, faces, normals, values = measure.marching_cubes(1000000.0*ex0 - 4.0, 0.0,
                                                       spacing=(1.0, 1.0, 1.0))
print(verts)
print(ex0.shape)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
mesh = Poly3DCollection(verts[faces])
ax.add_collection3d(mesh)
view_range: 'List[int]' = [4*n//10, 6*n//10]
ax.set_xlim(*view_range)
ax.set_ylim(*view_range)
ax.set_zlim(*view_range)
plt.tight_layout()
plt.show()
plt.close()