aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorThomas Albers Raviola <thomas@thomaslabs.org>2024-07-19 17:25:19 +0200
committerThomas Albers Raviola <thomas@thomaslabs.org>2024-07-19 17:25:19 +0200
commit5ee9ca1de9df21be6b3a749193481dac8844c90e (patch)
tree44b24c387bfdb6abae446814a11ef9e1a752939f
parent0d6bd534753b6cf80352e66d4f71471b9d9ff445 (diff)
Improve pylint score
-rw-r--r--schroedinger/schrodinger_solve.py61
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?')