Wavefunc plto

mas/pyploy7/main.py

@@ -0,0 +1,46 @@
+#!/usr/bin/python3
+import matplotlib.pyplot as plt
+import numpy as np
+
+fig = plt.figure()
+ax = plt.axes(projection='3d')
+
+# constants
+L = 100 # box length
+h = 1 # plancks constant
+
+D = np.sqrt(2/L) # norm factor
+m = 1 # mass
+
+# total energy function
+def E(n):
+    return (((h*n)/L)**2)*(1/(8*m))
+
+def k(n):
+    return np.sqrt( (8 * ((np.pi)**2) * m * E(n))/(h**2) )
+
+# wave function
+def psi(x: float, n: int=1) -> float:
+    return D * np.sin( k(n) * x ) 
+
+# probability density function
+def prob(x: float, n: int=1) -> float:
+    return (psi(x, n))**2
+
+
+x = np.arange(0, L, 0.1)
+n = np.arange(1, 10, 1)
+
+X, N = np.meshgrid(x, n)
+P = prob(X, N)
+
+fig = plt.figure()
+ax = plt.axes(projection='3d')
+ax.plot_surface(X, N, P, rstride=1, cstride=1, cmap='viridis', edgecolor='none')
+ax.set_xlabel('x')
+ax.set_ylabel('n')
+ax.set_zlabel('p');
+
+plt.show()
+
+