Add test for infinite potential well

1 files changed, 3 insertions, 27 deletions

diff --git a/schroedinger/schrodinger_solve.py b/schroedinger/schrodinger_solve.py
index 618d4ce..8f9f515 100644
--- a/schroedinger/schrodinger_solve.py
+++ b/schroedinger/schrodinger_solve.py
@@ -1,35 +1,11 @@
 import argparse
 
 import numpy as np
-from scipy.linalg import eigh_tridiagonal
 
-from schroedinger import Config, potential_interp
-
-def build_potential(config: Config):
-    start, end, steps = config.interval
-    potential = np.zeros((steps, 2))
-    potential[:, 0] = np.linspace(start, end, steps)
-    delta = np.abs(potential[1, 0] - potential[0, 0])
-    interp = potential_interp(config.interpolation, config.points)
-    potential[:, 1] = interp(potential[:, 0])
-    return potential, delta
-
-
-def solve_schroedinger(mass, potential, delta, eig_interval=None):
-    '''
-    returns eigen values and wave functions specified by eig_interval
-    '''
-    n = potential.shape[0]
-    a = 1 / mass / delta**2
-    w, v = eigh_tridiagonal(a + potential,
-                            -a * np.ones(n - 1, dtype=np.float_) / 2.0,
-                            select='i', select_range=eig_interval)
-    # Normalize eigenfunctions
-    for i in range(w.shape[0]):
-        v[:, i] /= np.sqrt(delta * np.sum(np.abs(v[:, i])**2))
-
-    return w, v
+from schroedinger import (
+    Config, potential_interp, build_potential, solve_schroedinger
 
+)
 
 def save_wavefuncs(filename, x, v):
     wavefuncs = np.zeros((x.shape[0], v.shape[1] + 1))