aboutsummaryrefslogtreecommitdiff
path: root/test
diff options
context:
space:
mode:
authorThomas Albers Raviola <thomas@thomaslabs.org>2024-07-19 16:11:50 +0200
committerThomas Albers Raviola <thomas@thomaslabs.org>2024-07-19 16:11:50 +0200
commit17faed5a64ca7c989b81dc3b27467e5c0aa14733 (patch)
tree4ea3e866cf5e41db3deff4676fde6fafc56757a3 /test
parent9ff2d4dd949f74b0052d21ecd83e398ab818b02e (diff)
Add type hinting
Diffstat (limited to 'test')
-rw-r--r--test/test_graphic.py10
-rw-r--r--test/test_infinite.py9
2 files changed, 8 insertions, 11 deletions
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)