diff options
Diffstat (limited to 'schroedinger/schrodinger_solve.py')
-rw-r--r-- | schroedinger/schrodinger_solve.py | 48 |
1 files changed, 47 insertions, 1 deletions
diff --git a/schroedinger/schrodinger_solve.py b/schroedinger/schrodinger_solve.py index 9369e2e..8f9f515 100644 --- a/schroedinger/schrodinger_solve.py +++ b/schroedinger/schrodinger_solve.py @@ -1,5 +1,51 @@ +import argparse + +import numpy as np + +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)) + 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(): - pass + parser = argparse.ArgumentParser( + prog='schrodinger_solve', + description='a', + epilog='a') + + parser.add_argument('filename') + args = parser.parse_args() + conf = Config(args.filename) + + potential, delta = build_potential(conf) + np.savetxt('potential.dat', potential) + + e, v = 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() |