diff options
| -rw-r--r-- | schroedinger/schrodinger_solve.py | 88 | ||||
| -rw-r--r-- | test/schrodinger.inp | 13 |
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 |