From 17faed5a64ca7c989b81dc3b27467e5c0aa14733 Mon Sep 17 00:00:00 2001 From: Thomas Albers Raviola Date: Fri, 19 Jul 2024 16:11:50 +0200 Subject: Add type hinting --- test/test_infinite.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) (limited to 'test/test_infinite.py') 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) -- cgit v1.2.3