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()