diff options
Diffstat (limited to 'schroedinger/schrodinger_solve.py')
| -rw-r--r-- | schroedinger/schrodinger_solve.py | 16 |
1 files changed, 12 insertions, 4 deletions
diff --git a/schroedinger/schrodinger_solve.py b/schroedinger/schrodinger_solve.py index 6885100..2eb83e8 100644 --- a/schroedinger/schrodinger_solve.py +++ b/schroedinger/schrodinger_solve.py @@ -3,15 +3,19 @@ import argparse from pathlib import Path import numpy as np +from numpy.typing import NDArray from schroedinger import ( Config, potential_interp, build_potential, solve_schroedinger - ) DESCRIPTION='Solve time independent Schrödinger\'s equation for a given system.' -def save_wavefuncs(filename, x, v): +def save_wavefuncs( + filename: Path, + x: NDArray[np.float64], + v: NDArray[np.float64] +) -> None: wavefuncs = np.zeros((x.shape[0], v.shape[1] + 1)) wavefuncs[:, 0] = x for i in range(v.shape[1]): @@ -19,7 +23,11 @@ def save_wavefuncs(filename, x, v): np.savetxt(filename, wavefuncs) -def save_expvalues(filename, x, v): +def save_expvalues( + filename: Path, + x: NDArray[np.float64], + v: NDArray[np.float64] +) -> None: n = v.shape[1] delta = np.abs(x[1] - x[0]) expvalues = np.zeros((n, 2)) @@ -31,7 +39,7 @@ def save_expvalues(filename, x, v): np.savetxt(filename, expvalues) -def main(): +def main() -> None: parser = argparse.ArgumentParser( prog='schrodinger_solve', description=DESCRIPTION, |