add more type hints
This commit is contained in:
parent
d13a3796a9
commit
bc428f5e8e
17 changed files with 170 additions and 106 deletions
|
|
@ -3,6 +3,7 @@
|
|||
Test fixtures
|
||||
|
||||
"""
|
||||
from typing import Tuple, Iterable, List
|
||||
import numpy # type: ignore
|
||||
import pytest # type: ignore
|
||||
|
||||
|
|
@ -14,22 +15,26 @@ from .utils import PRNG
|
|||
(5, 5, 5),
|
||||
# (7, 7, 7),
|
||||
])
|
||||
def shape(request):
|
||||
def shape(request: pytest.FixtureRequest) -> Iterable[Tuple[int, ...]]:
|
||||
yield (3, *request.param)
|
||||
|
||||
|
||||
@pytest.fixture(scope='module', params=[1.0, 1.5])
|
||||
def epsilon_bg(request):
|
||||
def epsilon_bg(request: pytest.FixtureRequest) -> Iterable[float]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope='module', params=[1.0, 2.5])
|
||||
def epsilon_fg(request):
|
||||
def epsilon_fg(request: pytest.FixtureRequest) -> Iterable[float]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope='module', params=['center', '000', 'random'])
|
||||
def epsilon(request, shape, epsilon_bg, epsilon_fg):
|
||||
def epsilon(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
epsilon_bg: float,
|
||||
epsilon_fg: float,
|
||||
) -> Iterable[numpy.ndarray]:
|
||||
is3d = (numpy.array(shape) == 1).sum() == 0
|
||||
if is3d:
|
||||
if request.param == '000':
|
||||
|
|
@ -53,17 +58,20 @@ def epsilon(request, shape, epsilon_bg, epsilon_fg):
|
|||
|
||||
|
||||
@pytest.fixture(scope='module', params=[1.0]) # 1.5
|
||||
def j_mag(request):
|
||||
def j_mag(request: pytest.FixtureRequest) -> Iterable[float]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope='module', params=[1.0, 1.5])
|
||||
def dx(request):
|
||||
def dx(request: pytest.FixtureRequest) -> Iterable[float]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope='module', params=['uniform', 'centerbig'])
|
||||
def dxes(request, shape, dx):
|
||||
def dxes(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
dx: float,
|
||||
) -> Iterable[List[List[numpy.ndarray]]]:
|
||||
if request.param == 'uniform':
|
||||
dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
|
||||
elif request.param == 'centerbig':
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import List, Tuple
|
||||
from typing import List, Tuple, Iterable, Optional
|
||||
import dataclasses
|
||||
import pytest # type: ignore
|
||||
import numpy # type: ignore
|
||||
|
|
@ -9,14 +9,14 @@ from ..fdmath import vec, unvec
|
|||
from .utils import assert_close # , assert_fields_close
|
||||
|
||||
|
||||
def test_residual(sim):
|
||||
def test_residual(sim: 'FDResult') -> None:
|
||||
A = fdfd.operators.e_full(sim.omega, sim.dxes, vec(sim.epsilon)).tocsr()
|
||||
b = -1j * sim.omega * vec(sim.j)
|
||||
residual = A @ vec(sim.e) - b
|
||||
assert numpy.linalg.norm(residual) < 1e-10
|
||||
|
||||
|
||||
def test_poynting_planes(sim):
|
||||
def test_poynting_planes(sim: 'FDResult') -> None:
|
||||
mask = (sim.j != 0).any(axis=0)
|
||||
if mask.sum() != 2:
|
||||
pytest.skip(f'test_poynting_planes will only test 2-point sources, got {mask.sum()}')
|
||||
|
|
@ -53,17 +53,17 @@ def test_poynting_planes(sim):
|
|||
# Also see conftest.py
|
||||
|
||||
@pytest.fixture(params=[1 / 1500])
|
||||
def omega(request):
|
||||
def omega(request: pytest.FixtureRequest) -> Iterable[float]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(params=[None])
|
||||
def pec(request):
|
||||
def pec(request: pytest.FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(params=[None])
|
||||
def pmc(request):
|
||||
def pmc(request: pytest.FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
|
||||
yield request.param
|
||||
|
||||
|
||||
|
|
@ -74,7 +74,10 @@ def pmc(request):
|
|||
|
||||
|
||||
@pytest.fixture(params=['diag']) # 'center'
|
||||
def j_distribution(request, shape, j_mag):
|
||||
def j_distribution(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
j_mag: float,
|
||||
) -> Iterable[numpy.ndarray]:
|
||||
j = numpy.zeros(shape, dtype=complex)
|
||||
center_mask = numpy.zeros(shape, dtype=bool)
|
||||
center_mask[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = True
|
||||
|
|
@ -89,7 +92,7 @@ def j_distribution(request, shape, j_mag):
|
|||
|
||||
@dataclasses.dataclass()
|
||||
class FDResult:
|
||||
shape: Tuple[int]
|
||||
shape: Tuple[int, ...]
|
||||
dxes: List[List[numpy.ndarray]]
|
||||
epsilon: numpy.ndarray
|
||||
omega: complex
|
||||
|
|
@ -100,7 +103,15 @@ class FDResult:
|
|||
|
||||
|
||||
@pytest.fixture()
|
||||
def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc):
|
||||
def sim(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
epsilon: numpy.ndarray,
|
||||
dxes: List[List[numpy.ndarray]],
|
||||
j_distribution: numpy.ndarray,
|
||||
omega: float,
|
||||
pec: Optional[numpy.ndarray],
|
||||
pmc: Optional[numpy.ndarray],
|
||||
) -> FDResult:
|
||||
"""
|
||||
Build simulation from parts
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,15 +1,15 @@
|
|||
#####################################
|
||||
from typing import Optional, Tuple, Iterable, List
|
||||
import pytest # type: ignore
|
||||
import numpy # type: ignore
|
||||
from numpy.testing import assert_allclose # type: ignore
|
||||
|
||||
from .. import fdfd
|
||||
from ..fdmath import vec, unvec
|
||||
from ..fdmath import vec, unvec, dx_lists_mut
|
||||
#from .utils import assert_close, assert_fields_close
|
||||
from .test_fdfd import FDResult
|
||||
|
||||
|
||||
def test_pml(sim, src_polarity):
|
||||
def test_pml(sim: FDResult, src_polarity: int) -> None:
|
||||
e_sqr = numpy.squeeze((sim.e.conj() * sim.e).sum(axis=0))
|
||||
|
||||
# from matplotlib import pyplot
|
||||
|
|
@ -42,34 +42,40 @@ def test_pml(sim, src_polarity):
|
|||
# Also see conftest.py
|
||||
|
||||
@pytest.fixture(params=[1 / 1500])
|
||||
def omega(request):
|
||||
def omega(request: pytest.FixtureRequest) -> Iterable[float]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(params=[None])
|
||||
def pec(request):
|
||||
def pec(request: pytest.FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(params=[None])
|
||||
def pmc(request):
|
||||
def pmc(request: pytest.FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(params=[(30, 1, 1),
|
||||
(1, 30, 1),
|
||||
(1, 1, 30)])
|
||||
def shape(request):
|
||||
def shape(request: pytest.FixtureRequest) -> Iterable[Tuple[int, ...]]:
|
||||
yield (3, *request.param)
|
||||
|
||||
|
||||
@pytest.fixture(params=[+1, -1])
|
||||
def src_polarity(request):
|
||||
def src_polarity(request: pytest.FixtureRequest) -> Iterable[int]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def j_distribution(request, shape, epsilon, dxes, omega, src_polarity):
|
||||
def j_distribution(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
epsilon: numpy.ndarray,
|
||||
dxes: dx_lists_mut,
|
||||
omega: float,
|
||||
src_polarity: int,
|
||||
) -> Iterable[numpy.ndarray]:
|
||||
j = numpy.zeros(shape, dtype=complex)
|
||||
|
||||
dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis
|
||||
|
|
@ -101,13 +107,22 @@ def j_distribution(request, shape, epsilon, dxes, omega, src_polarity):
|
|||
|
||||
|
||||
@pytest.fixture()
|
||||
def epsilon(request, shape, epsilon_bg, epsilon_fg):
|
||||
def epsilon(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
epsilon_bg: float,
|
||||
epsilon_fg: float,
|
||||
) -> Iterable[numpy.ndarray]:
|
||||
epsilon = numpy.full(shape, epsilon_fg, dtype=float)
|
||||
yield epsilon
|
||||
|
||||
|
||||
@pytest.fixture(params=['uniform'])
|
||||
def dxes(request, shape, dx, omega, epsilon_fg):
|
||||
def dxes(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
dx: float,
|
||||
omega: float,
|
||||
epsilon_fg: float,
|
||||
) -> Iterable[List[List[numpy.ndarray]]]:
|
||||
if request.param == 'uniform':
|
||||
dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
|
||||
dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis
|
||||
|
|
@ -120,7 +135,15 @@ def dxes(request, shape, dx, omega, epsilon_fg):
|
|||
|
||||
|
||||
@pytest.fixture()
|
||||
def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc):
|
||||
def sim(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
epsilon: numpy.ndarray,
|
||||
dxes: dx_lists_mut,
|
||||
j_distribution: numpy.ndarray,
|
||||
omega: float,
|
||||
pec: Optional[numpy.ndarray],
|
||||
pmc: Optional[numpy.ndarray],
|
||||
) -> FDResult:
|
||||
j_vec = vec(j_distribution)
|
||||
eps_vec = vec(epsilon)
|
||||
e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec,
|
||||
|
|
@ -129,7 +152,7 @@ def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc):
|
|||
|
||||
sim = FDResult(
|
||||
shape=shape,
|
||||
dxes=dxes,
|
||||
dxes=[list(d) for d in dxes],
|
||||
epsilon=epsilon,
|
||||
j=j_distribution,
|
||||
e=e,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import List, Tuple
|
||||
from typing import List, Tuple, Iterable
|
||||
import dataclasses
|
||||
import pytest # type: ignore
|
||||
import numpy # type: ignore
|
||||
|
|
@ -8,7 +8,7 @@ from .. import fdtd
|
|||
from .utils import assert_close, assert_fields_close, PRNG
|
||||
|
||||
|
||||
def test_initial_fields(sim):
|
||||
def test_initial_fields(sim: 'TDResult') -> None:
|
||||
# Make sure initial fields didn't change
|
||||
e0 = sim.es[0]
|
||||
h0 = sim.hs[0]
|
||||
|
|
@ -20,7 +20,7 @@ def test_initial_fields(sim):
|
|||
assert not h0.any()
|
||||
|
||||
|
||||
def test_initial_energy(sim):
|
||||
def test_initial_energy(sim: 'TDResult') -> None:
|
||||
"""
|
||||
Assumes fields start at 0 before J0 is added
|
||||
"""
|
||||
|
|
@ -41,7 +41,7 @@ def test_initial_energy(sim):
|
|||
assert_fields_close(e0_dot_j0, u0)
|
||||
|
||||
|
||||
def test_energy_conservation(sim):
|
||||
def test_energy_conservation(sim: 'TDResult') -> None:
|
||||
"""
|
||||
Assumes fields start at 0 before J0 is added
|
||||
"""
|
||||
|
|
@ -63,7 +63,7 @@ def test_energy_conservation(sim):
|
|||
assert_close(u_estep.sum(), u)
|
||||
|
||||
|
||||
def test_poynting_divergence(sim):
|
||||
def test_poynting_divergence(sim: 'TDResult') -> None:
|
||||
args = {'dxes': sim.dxes,
|
||||
'epsilon': sim.epsilon}
|
||||
|
||||
|
|
@ -90,7 +90,7 @@ def test_poynting_divergence(sim):
|
|||
u_eprev = u_estep
|
||||
|
||||
|
||||
def test_poynting_planes(sim):
|
||||
def test_poynting_planes(sim: 'TDResult') -> None:
|
||||
mask = (sim.js[0] != 0).any(axis=0)
|
||||
if mask.sum() > 1:
|
||||
pytest.skip('test_poynting_planes can only test single point sources, got {}'.format(mask.sum()))
|
||||
|
|
@ -140,30 +140,33 @@ def test_poynting_planes(sim):
|
|||
|
||||
|
||||
@pytest.fixture(params=[0.3])
|
||||
def dt(request):
|
||||
def dt(request: pytest.FixtureRequest) -> Iterable[float]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@dataclasses.dataclass()
|
||||
class TDResult:
|
||||
shape: Tuple[int]
|
||||
shape: Tuple[int, ...]
|
||||
dt: float
|
||||
dxes: List[List[numpy.ndarray]]
|
||||
epsilon: numpy.ndarray
|
||||
j_distribution: numpy.ndarray
|
||||
j_steps: Tuple[int]
|
||||
j_steps: Tuple[int, ...]
|
||||
es: List[numpy.ndarray] = dataclasses.field(default_factory=list)
|
||||
hs: List[numpy.ndarray] = dataclasses.field(default_factory=list)
|
||||
js: List[numpy.ndarray] = dataclasses.field(default_factory=list)
|
||||
|
||||
|
||||
@pytest.fixture(params=[(0, 4, 8)]) # (0,)
|
||||
def j_steps(request):
|
||||
def j_steps(request: pytest.fixtureRequest) -> Iterable[Tuple[int, ...]]:
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(params=['center', 'random'])
|
||||
def j_distribution(request, shape, j_mag):
|
||||
def j_distribution(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
j_mag: float,
|
||||
) -> Iterable[numpy.ndarray]:
|
||||
j = numpy.zeros(shape)
|
||||
if request.param == 'center':
|
||||
j[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = j_mag
|
||||
|
|
@ -175,7 +178,14 @@ def j_distribution(request, shape, j_mag):
|
|||
|
||||
|
||||
@pytest.fixture()
|
||||
def sim(request, shape, epsilon, dxes, dt, j_distribution, j_steps):
|
||||
def sim(request: pytest.FixtureRequest,
|
||||
shape: Tuple[int, ...],
|
||||
epsilon: numpy.ndarray,
|
||||
dxes: List[List[numpy.ndarray]],
|
||||
dt: float,
|
||||
j_distribution: numpy.ndarray,
|
||||
j_steps: Tuple[int, ...],
|
||||
) -> TDResult:
|
||||
is3d = (numpy.array(shape) == 1).sum() == 0
|
||||
if is3d:
|
||||
if dt != 0.3:
|
||||
|
|
|
|||
|
|
@ -1,13 +1,22 @@
|
|||
from typing import Any
|
||||
import numpy # type: ignore
|
||||
|
||||
PRNG = numpy.random.RandomState(12345)
|
||||
|
||||
def assert_fields_close(x, y, *args, **kwargs):
|
||||
def assert_fields_close(x: numpy.ndarray,
|
||||
y: numpy.ndarray,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
numpy.testing.assert_allclose(
|
||||
x, y, verbose=False,
|
||||
err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(x, -1),
|
||||
numpy.rollaxis(y, -1)), *args, **kwargs)
|
||||
|
||||
def assert_close(x, y, *args, **kwargs):
|
||||
def assert_close(x: numpy.ndarray,
|
||||
y: numpy.ndarray,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
numpy.testing.assert_allclose(x, y, *args, **kwargs)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue