from typing import List, Tuple import dataclasses import pytest import numpy from numpy.testing import assert_allclose, assert_array_equal from .. import fdfd, vec, unvec from .utils import assert_close, assert_fields_close, PRNG def test_residual(sim): 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): 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()}') points = numpy.where(mask) mask[points[0][0], points[1][0], points[2][0]] = 0 mx = numpy.roll(mask, -1, axis=0) my = numpy.roll(mask, -1, axis=1) mz = numpy.roll(mask, -1, axis=2) e2h = fdfd.operators.e2h(omega=sim.omega, dxes=sim.dxes, pmc=sim.pmc) ev = vec(sim.e) hv = e2h @ ev exh = fdfd.operators.poynting_e_cross(e=ev, dxes=sim.dxes) @ hv.conj() s = unvec(exh.real / 2, sim.shape[1:]) planes = [s[0, mask].sum(), -s[0, mx].sum(), s[1, mask].sum(), -s[1, my].sum(), s[2, mask].sum(), -s[2, mz].sum()] e_dot_j = sim.e * sim.j * sim.dxes[0][0][:, None, None] * sim.dxes[0][1][None, :, None] * sim.dxes[0][2][None, None, :] src_energy = -e_dot_j[:, mask].real / 2 assert_close(sum(planes), src_energy.sum(), rtol=2e-5) # TODO why is the tolerance so bad? ##################################### # Test fixtures ##################################### # Also see conftest.py @pytest.fixture(params=[1/1500]) def omega(request): yield request.param @pytest.fixture(params=[None]) def pec(request): yield request.param @pytest.fixture(params=[None]) def pmc(request): yield request.param @pytest.fixture(params=['center', 'diag']) def j_distribution(request, shape, j_mag): 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 if request.param == 'center': j[center_mask] = j_mag elif request.param == 'diag': j[numpy.roll(center_mask, [1, 1, 1], axis=(1, 2, 3))] = j_mag j[numpy.roll(center_mask, [-1, -1, -1], axis=(1, 2, 3))] = -1j * j_mag yield j @dataclasses.dataclass() class SimResult: shape: Tuple[int] dxes: List[List[numpy.ndarray]] epsilon: numpy.ndarray omega: complex j: numpy.ndarray e: numpy.ndarray pmc: numpy.ndarray pec: numpy.ndarray @pytest.fixture() def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc): # is3d = (numpy.array(shape) == 1).sum() == 0 # if is3d: # pytest.skip('Skipping dt != 0.3 because test is 3D (for speed)') j_vec = vec(j_distribution) eps_vec = vec(epsilon) e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec, matrix_solver_opts={'atol': 1e-12}) e = unvec(e_vec, shape[1:]) sim = SimResult( shape=shape, dxes=dxes, epsilon=epsilon, j=j_distribution, e=e, pec=pec, pmc=pmc, omega=omega, ) return sim