diff options
Diffstat (limited to 'schroedinger/schroedinger.py')
| -rw-r--r-- | schroedinger/schroedinger.py | 29 |
1 files changed, 28 insertions, 1 deletions
diff --git a/schroedinger/schroedinger.py b/schroedinger/schroedinger.py index 10d994e..3ed2c3f 100644 --- a/schroedinger/schroedinger.py +++ b/schroedinger/schroedinger.py @@ -4,6 +4,7 @@ import numpy as np from numpy.polynomial import Polynomial from numpy.typing import NDArray from scipy.interpolate import CubicSpline +from scipy.linalg import eigh_tridiagonal class Config: def __init__(self, path): @@ -29,7 +30,7 @@ class Config: self.mass = float(next_parameter(fd)) start, end, steps = next_parameter(fd).split() self.interval = [float(start), float(end), int(steps)] - self.eig_interval = [int(attr) for attr in next_parameter(fd).split()] + self.eig_interval = [int(attr) - 1 for attr in next_parameter(fd).split()] self.interpolation = next_parameter(fd) npoints = int(next_parameter(fd)) @@ -57,3 +58,29 @@ def potential_interp(interpolation, points): return cs raise ValueError() + + +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 |