Add solve_schroedinger

2 files changed, 87 insertions, 14 deletions

diff --git a/schroedinger/schrodinger_solve.py b/schroedinger/schrodinger_solve.py
index a759c3a..d301498 100644
--- a/schroedinger/schrodinger_solve.py
+++ b/schroedinger/schrodinger_solve.py
@@ -1,6 +1,13 @@
-import numpy as np
-from pathlib import Path
 import argparse
+from pathlib import Path
+
+import numpy as np
+from numpy.polynomial import Polynomial
+from numpy.typing import NDArray
+
+from scipy.linalg import eigh_tridiagonal
+import scipy.interpolate as interp
+
 
 class Config:
     def __init__(self, path):
@@ -24,7 +31,8 @@ class Config:
         with open(path, 'r') as fd:
             try:
                 self.mass = float(next_parameter(fd))
-                self.interval = [float(attr) for attr in next_parameter(fd).split()]
+                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.interpolation = next_parameter(fd)
 
@@ -37,19 +45,83 @@ class Config:
                 print('Syntax error in \'{}\' line {}'.format(path.name, current_line))
 
 
+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])
+
+    if config.interpolation == 'linear':
+        potential[:, 1] = np.interp(potential[:, 0], config.points[:, 0],
+                                    config.points[:, 1])
+    elif config.interpolation == 'polynomial':
+        p = Polynomial.fit(config.points[:, 0], config.points[:, 1],
+                           config.points.shape[0] - 1)
+        potential[:, 1] = p(potential[:, 0])
+    elif config.interpolation == 'cspline':
+        cs = CubicSpline(config.points[:, 0], config.points[:, 1])
+        potential[:, 1] = cs(potential[:, 0])
+    else:
+        raise ValueError()
+
+    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
+
+
+def save_wavefuncs(filename, x, v):
+    wavefuncs = np.zeros((x.shape[0], v.shape[1] + 1))
+    wavefuncs[:, 0] = x
+    for i in range(v.shape[1]):
+        wavefuncs[:, 1 + i] = v[:, i]
+    np.savetxt(filename, wavefuncs)
+
+
+def save_expvalues(filename, x, v):
+    n = v.shape[1]
+    delta = np.abs(x[1] - x[0])
+    expvalues = np.zeros((n, 2))
+    for i in range(n):
+        exp_x = delta * np.sum(v[:, i] * x * v[:, i])
+        exp_xsq = delta * np.sum(v[:, i] * x**2 * v[:, i])
+        expvalues[i, 0] = exp_x
+        expvalues[i, 1] = np.sqrt(exp_xsq - exp_x ** 2)
+    np.savetxt(filename, expvalues)
+
+
 def main():
     parser = argparse.ArgumentParser(
         prog='schrodinger_solve',
         description='a',
         epilog='a')
+
     parser.add_argument('filename')
     args = parser.parse_args()
     conf = Config(args.filename)
-    print(conf.mass)
-    print(conf.interval)
-    print(conf.eig_interval)
-    print(conf.interpolation)
-    print(conf.points)
+
+    potential, delta = build_potential(conf)
+    np.savetxt('potential.dat', potential)
+
+    e, f = solve_schroedinger(conf.mass, potential[:, 1], delta, conf.eig_interval)
+
+    np.savetxt('energies.dat', e)
+    save_wavefuncs('wavefuncs.dat', potential[:, 0], v)
+    save_expvalues('expvalues.dat', potential[:, 0], v)
+
 
 if __name__ == '__main__':
     main()
diff --git a/test/schrodinger.inp b/test/schrodinger.inp
index 8c9a16b..98bcbe2 100644
--- a/test/schrodinger.inp
+++ b/test/schrodinger.inp
@@ -1,7 +1,8 @@
-2.0                    # mass
--2.0    2.0     1999   # xMin xMax nPoint
+4.0                    # mass
+-5.0    5.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
+polynomial             # interpolation type
+3                      # nr. of interpolation points and xy declarations
+-1.0 0.5
+0.0 0.0
+1.0 0.5