Improve pylint score

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?')