aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--.gitignore1
-rw-r--r--schroedinger/__init__.py2
-rw-r--r--schroedinger/schrodinger_solve.py30
-rw-r--r--schroedinger/schroedinger.py29
-rw-r--r--test/__init__.py1
-rw-r--r--test/infinite.inp7
-rw-r--r--test/test_infinite.py45
7 files changed, 86 insertions, 29 deletions
diff --git a/.gitignore b/.gitignore
index 251d000..8b15847 100644
--- a/.gitignore
+++ b/.gitignore
@@ -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))