diff options
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | schroedinger/__init__.py | 2 | ||||
| -rw-r--r-- | schroedinger/schrodinger_solve.py | 30 | ||||
| -rw-r--r-- | schroedinger/schroedinger.py | 29 | ||||
| -rw-r--r-- | test/__init__.py | 1 | ||||
| -rw-r--r-- | test/infinite.inp | 7 | ||||
| -rw-r--r-- | test/test_infinite.py | 45 |
7 files changed, 86 insertions, 29 deletions
@@ -2,3 +2,4 @@ docs/_build .idea/ schroedinger/__pycache__ notes.org +**/__pycache__ diff --git a/schroedinger/__init__.py b/schroedinger/__init__.py index 1de4510..e952971 100644 --- a/schroedinger/__init__.py +++ b/schroedinger/__init__.py @@ -1 +1 @@ -from .schroedinger import Config +from .schroedinger import * 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)) 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 diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/test/__init__.py @@ -0,0 +1 @@ + diff --git a/test/infinite.inp b/test/infinite.inp new file mode 100644 index 0000000..a61ffbe --- /dev/null +++ b/test/infinite.inp @@ -0,0 +1,7 @@ +2.0 # mass +-2.0 2.0 1999 # xMin xMax nPoint +1 5 # first and last eigenvalue to print +linear # interpolation type +2 # nr. of interpolation points and xy declarations +-2.0 0.0 +2.0 0.0 diff --git a/test/test_infinite.py b/test/test_infinite.py new file mode 100644 index 0000000..aa85ae5 --- /dev/null +++ b/test/test_infinite.py @@ -0,0 +1,45 @@ +import pytest +import numpy as np + +import matplotlib.pyplot as plt + +from schroedinger import ( + Config, potential_interp, build_potential, solve_schroedinger +) + + +def psi(x, n, a): + n += 1 # Index starting from 0 + x = -x + a / 2 # Reflect x and move to the left + return np.sqrt(2 / a) * np.sin(n * np.pi * x / a) + + +def energy(n, a, m): + return (n + 1)**2 * np.pi**2 / m / a**2 / 2.0 + + +def test_infinite(): + conf = Config('test/infinite.inp') + potential, delta = build_potential(conf) + + e, v = solve_schroedinger(conf.mass, potential[:, 1], delta, conf.eig_interval) + # Account for -v also being an eigenvector if v is one + v = np.abs(v) + + a = np.abs(conf.points[1][0] - conf.points[0][0]) + e_theory = np.zeros(e.shape) + v_theory = np.zeros(v.shape) + + + for n in range(e.shape[0]): + e_theory[n] = energy(n, a, conf.mass) + v_theory[:, n] = np.abs(psi(potential[:, 0], n, a)) + + # for n in range(e.shape[0]): + # plt.plot(potential[:, 0], np.abs(v[:, n]), label='Num{}'.format(n)) + # plt.plot(potential[:, 0], np.abs(v_theory[:, n]), ls='--', label='Theory{}'.format(n)) + # plt.legend() + # plt.savefig('test.pdf') + + assert (np.allclose(e, e_theory, rtol=1e-2, atol=1e-2) + and np.allclose(v, v_theory, rtol=1e-2, atol=1e-2)) |