diff options
| -rw-r--r-- | schroedinger/schrodinger_solve.py | 95 | ||||
| -rw-r--r-- | schroedinger/schroedinger.py | 118 | ||||
| -rw-r--r-- | test/test_graphic.py | 10 | ||||
| -rw-r--r-- | test/test_infinite.py | 9 |
4 files changed, 158 insertions, 74 deletions
diff --git a/schroedinger/schrodinger_solve.py b/schroedinger/schrodinger_solve.py index 8f9f515..aec63e1 100644 --- a/schroedinger/schrodinger_solve.py +++ b/schroedinger/schrodinger_solve.py @@ -1,50 +1,93 @@ +'''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 import numpy as np +from numpy.typing import NDArray from schroedinger import ( - Config, potential_interp, build_potential, solve_schroedinger - + Config, 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]) +DESCRIPTION='Solve time independent Schrödinger\'s equation for a given system.' + +def save_wavefuncs( + filename: Path, + pos: NDArray[np.float64], + wavefuncs: NDArray[np.float64] +) -> None: + '''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, + pos: NDArray[np.float64], + wavefuncs: NDArray[np.float64] +) -> None: + '''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(): +def main() -> None: + '''Program entry point + ''' parser = argparse.ArgumentParser( prog='schrodinger_solve', - description='a', - epilog='a') + description=DESCRIPTION, + epilog='') + + parser.add_argument('filename', + help='File describing the system to solve') + parser.add_argument('-o', '--output-dir', + help='Output directory for the results') - parser.add_argument('filename') args = parser.parse_args() + conf = Config(args.filename) + output_path = Path(args.output_dir) if args.output_dir else Path.cwd() potential, delta = build_potential(conf) - np.savetxt('potential.dat', potential) - e, v = solve_schroedinger(conf.mass, potential[:, 1], delta, conf.eig_interval) + eig, vec = 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) + try: + np.savetxt(output_path / 'potential.dat', potential) + 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?') + sys.exit(1) if __name__ == '__main__': diff --git a/schroedinger/schroedinger.py b/schroedinger/schroedinger.py index 4c5d1f7..0c26041 100644 --- a/schroedinger/schroedinger.py +++ b/schroedinger/schroedinger.py @@ -1,12 +1,28 @@ +'''schroedinger.py + +Collection of method common to solve and tests +''' + +import sys from pathlib import Path +from typing import TextIO, Type import numpy as np from numpy.polynomial import Polynomial from numpy.typing import NDArray + from scipy.interpolate import CubicSpline from scipy.linalg import eigh_tridiagonal + +Interpolator = Type[callable] | Type[Polynomial] | Type[CubicSpline] +# type Interpolator = callable | Polynomial | CubicSpline + + class Config: + '''Wrapper for reading, parsing and storing the contents of a configuration + file + ''' def __init__(self, path): current_line = 0 @@ -14,53 +30,70 @@ class Config: if path is not Path: path = Path(path) - def next_parameter(fd): + def next_parameter(file: TextIO): '''Read next parameter, ignoring comments or empty lines''' content = None nonlocal current_line while not content: - str = fd.readline() + line = file.readline() current_line += 1 - index = str.find('#') - content = str[0:index].strip() + index = line.find('#') + content = line[0:index].strip() return content - with open(path, 'r') as fd: + with open(path, 'r', encoding='utf8') as file: try: - self.mass = float(next_parameter(fd)) - start, end, steps = next_parameter(fd).split() + self.mass = float(next_parameter(file)) + start, end, steps = next_parameter(file).split() self.interval = [float(start), float(end), int(steps)] - self.eig_interval = [int(attr) - 1 for attr in next_parameter(fd).split()] - self.interpolation = next_parameter(fd) + self.eig_interval = [int(attr) - 1 for attr in next_parameter(file).split()] + self.interpolation = next_parameter(file) - npoints = int(next_parameter(fd)) + npoints = int(next_parameter(file)) self.points = np.zeros((npoints, 2)) for i in range(npoints): - line = next_parameter(fd) + line = next_parameter(file) self.points[i] = np.array([float(comp) for comp in line.split()]) - # TODO: don't be a moron, catch only relevant exceptions - except: - print('Syntax error in \'{}\' line {}'.format(path.name, current_line)) - raise ValueError() + except ValueError: + print(f'Syntax error in \'{path.name}\' line {current_line}') + sys.exit(1) -def potential_interp(interpolation, points): +def potential_interp( + interpolation: str, + points: NDArray[np.float64] +) -> Interpolator: + '''Create an interpolator for a set of predefined potential points + + :param interpolation: Kind of interpolation + :param points: Points to interpolate within + :return: Interpolating object + ''' + interpolator = None + if interpolation == 'linear': - def line(x): - return np.interp(x, points[:, 0], points[:, 1]) - return line + def line(pos): + return np.interp(pos, points[:, 0], points[:, 1]) + interpolator = line elif interpolation == 'polynomial': - poly = Polynomial.fit(points[:, 0], points[:, 1], - points.shape[0] - 1) - return poly + interpolator = Polynomial.fit(points[:, 0], points[:, 1], + points.shape[0] - 1) elif interpolation == 'cspline': - cs = CubicSpline(points[:, 0], points[:, 1]) - return cs + interpolator = CubicSpline(points[:, 0], points[:, 1]) + else: + raise ValueError('Invalid interpolator kind. Use any of: linear, polynomial or cspline') + + return interpolator - raise ValueError() +def build_potential( + config: Config +) -> tuple[NDArray[np.float64], np.float64]: + '''Build a potential based on the options inside a configuration file -def build_potential(config: Config): + :param config: System parameters + :return: Potential and distance between samples + ''' start, end, steps = config.interval potential = np.zeros((steps, 2)) potential[:, 0] = np.linspace(start, end, steps) @@ -70,17 +103,28 @@ def build_potential(config: Config): return potential, delta -def solve_schroedinger(mass, potential, delta, eig_interval=None): - ''' - returns eigen values and wave functions specified by eig_interval +def solve_schroedinger( + mass: float, + potential: NDArray[np.float64], + delta: float, + eig_interval: tuple[int, int]=None +) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + '''Solve the one dimensional, time independent Schrödinger's equation for a + particle of given mass inside a discretized potential + + :param mass: Mass of the particle + :param potential: Discretized potential + :param delta: Distance between points in the potential + :param eig_interval: Interval of quantum numbers for which the states are to be calculated + :return: Eigenvalues and normalized wave functions ''' - n = potential.shape[0] - a = 1 / mass / delta**2 - w, v = eigh_tridiagonal(a + potential, - -a * np.ones(n - 1, dtype=np.float64) / 2.0, - select='i', select_range=eig_interval) + qnumber = potential.shape[0] + coeff = 1 / mass / delta**2 + eig, vec = eigh_tridiagonal(coeff + potential, + -coeff * np.ones(qnumber - 1, dtype=np.float64) / 2.0, + select='i', select_range=eig_interval) # Normalize eigenfunctions - for i in range(w.shape[0]): - v[:, i] /= np.sqrt(delta * np.sum(np.abs(v[:, i])**2)) + for i in range(eig.shape[0]): + vec[:, i] /= np.sqrt(delta * np.sum(np.abs(vec[:, i])**2)) - return w, v + return eig, vec diff --git a/test/test_graphic.py b/test/test_graphic.py index 73d8b82..8433cae 100644 --- a/test/test_graphic.py +++ b/test/test_graphic.py @@ -1,7 +1,3 @@ -from pathlib import Path -import sys -sys.path.insert(0, str(Path.cwd().parent/'schroedinger')) - import pytest import numpy as np @@ -18,19 +14,19 @@ v = {} FORMS = ['asymmetric', 'double_linear', 'double_cubic'] @pytest.mark.parametrize('form', FORMS) -def test_potential(form): +def test_potential(form: str) -> None: potential_prec = np.loadtxt(f'test/{form}/potential.dat') assert np.allclose(potential[form], potential_prec, rtol=1e-2, atol=1e-2) @pytest.mark.parametrize('form', FORMS) -def test_e(form): +def test_e(form: str) -> None: e_prec = np.loadtxt(f'test/{form}/energies.dat') assert np.allclose(e[form], e_prec, rtol=1e-2, atol=1e-2) @pytest.mark.parametrize('form', FORMS) -def test_v(form): +def test_v(form: str) -> None: v_prec = np.loadtxt(f'test/{form}/wavefuncs.dat')[:, 1:] assert np.allclose(v[form], v_prec, rtol=1e-2, atol=1e-2) diff --git a/test/test_infinite.py b/test/test_infinite.py index aa85ae5..bca21b6 100644 --- a/test/test_infinite.py +++ b/test/test_infinite.py @@ -1,5 +1,6 @@ import pytest import numpy as np +from numpy.typing import NDArray import matplotlib.pyplot as plt @@ -8,17 +9,17 @@ from schroedinger import ( ) -def psi(x, n, a): +def psi(x: NDArray[np.float64], n: int, a: float) -> NDArray[np.float64]: n += 1 # Index starting from 0 x = -x + a / 2 # Reflect x and move to the left return np.sqrt(2 / a) * np.sin(n * np.pi * x / a) -def energy(n, a, m): - return (n + 1)**2 * np.pi**2 / m / a**2 / 2.0 +def energy(n: int, a: float, mass: float) -> float: + return (n + 1)**2 * np.pi**2 / mass / a**2 / 2.0 -def test_infinite(): +def test_infinite() -> None: conf = Config('test/infinite.inp') potential, delta = build_potential(conf) |