diff options
| -rw-r--r-- | schroedinger/schrodinger_solve.py | 61 |
1 files changed, 40 insertions, 21 deletions
diff --git a/schroedinger/schrodinger_solve.py b/schroedinger/schrodinger_solve.py index 2eb83e8..aec63e1 100644 --- a/schroedinger/schrodinger_solve.py +++ b/schroedinger/schrodinger_solve.py @@ -1,3 +1,8 @@ +'''schrodinger_solve + +Solve one dimensional, time independent Schrödinger's equation for an user +provided system description ''' + import sys import argparse from pathlib import Path @@ -6,40 +11,54 @@ import numpy as np from numpy.typing import NDArray from schroedinger import ( - Config, potential_interp, build_potential, solve_schroedinger + Config, build_potential, solve_schroedinger ) DESCRIPTION='Solve time independent Schrödinger\'s equation for a given system.' def save_wavefuncs( filename: Path, - x: NDArray[np.float64], - v: NDArray[np.float64] + pos: NDArray[np.float64], + wavefuncs: NDArray[np.float64] ) -> None: - 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) + '''Save wave functions to file + + :param filename: Filename where to output wave functions + :param pos: x Axis + :param wavefuncs: Wave functions + ''' + content = np.zeros((pos.shape[0], wavefuncs.shape[1] + 1)) + content[:, 0] = pos + for i in range(wavefuncs.shape[1]): + content[:, 1 + i] = wavefuncs[:, i] + np.savetxt(filename, content) def save_expvalues( filename: Path, - x: NDArray[np.float64], - v: NDArray[np.float64] + pos: NDArray[np.float64], + wavefuncs: NDArray[np.float64] ) -> None: - 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]) + '''Save expected values to file + + :param filename: Filename where to output expected values + :param pos: x Axis + :param wavefuncs: Wave functions + ''' + nfuncs = wavefuncs.shape[1] + delta = np.abs(pos[1] - pos[0]) + expvalues = np.zeros((nfuncs, 2)) + for i in range(nfuncs): + exp_x = delta * np.sum(wavefuncs[:, i] * pos * wavefuncs[:, i]) + exp_xsq = delta * np.sum(wavefuncs[:, i] * pos**2 * wavefuncs[:, i]) expvalues[i, 0] = exp_x expvalues[i, 1] = np.sqrt(exp_xsq - exp_x ** 2) np.savetxt(filename, expvalues) def main() -> None: + '''Program entry point + ''' parser = argparse.ArgumentParser( prog='schrodinger_solve', description=DESCRIPTION, @@ -57,14 +76,14 @@ def main() -> None: potential, delta = build_potential(conf) - e, v = solve_schroedinger(conf.mass, potential[:, 1], delta, - conf.eig_interval) + eig, vec = solve_schroedinger(conf.mass, potential[:, 1], delta, + conf.eig_interval) try: np.savetxt(output_path / 'potential.dat', potential) - np.savetxt(output_path / 'energies.dat', e) - save_wavefuncs(output_path / 'wavefuncs.dat', potential[:, 0], v) - save_expvalues(output_path / 'expvalues.dat', potential[:, 0], v) + np.savetxt(output_path / 'energies.dat', eig) + save_wavefuncs(output_path / 'wavefuncs.dat', potential[:, 0], vec) + save_expvalues(output_path / 'expvalues.dat', potential[:, 0], vec) except FileNotFoundError: print('Output files could not be saved.' ' Are you sure the output directory exists?') |