import numpy import pytest import dataclasses from typing import List, Tuple from numpy.testing import assert_allclose, assert_array_equal from fdfd_tools import fdtd prng = numpy.random.RandomState(12345) def assert_fields_close(a, b, *args, **kwargs): numpy.testing.assert_allclose(a, b, verbose=False, err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(a, -1), numpy.rollaxis(b, -1)), *args, **kwargs) def assert_close(a, b, *args, **kwargs): numpy.testing.assert_allclose(a, b, *args, **kwargs) def test_initial_fields(sim): # Make sure initial fields didn't change e0 = sim.es[0] h0 = sim.hs[0] j0 = sim.js[0] mask = (j0 != 0) assert_fields_close(e0[mask], j0[mask] / sim.epsilon[mask]) assert not e0[~mask].any() assert not h0.any() def test_initial_energy(sim): """ Assumes fields start at 0 before J0 is added """ j0 = sim.js[0] e0 = sim.es[0] h0 = sim.hs[0] h1 = sim.hs[1] mask = (j0 != 0) dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0) u0 = (j0 * j0.conj() / sim.epsilon * dV).sum(axis=0) args = {'dxes': sim.dxes, 'epsilon': sim.epsilon} # Make sure initial energy and E dot J are correct energy0 = fdtd.energy_estep(h0=h0, e1=e0, h2=h1, **args) e0_dot_j0 = fdtd.delta_energy_j(j0=j0, e1=e0, dxes=sim.dxes) assert_fields_close(energy0, u0) assert_fields_close(e0_dot_j0, u0) def test_energy_conservation(sim): """ Assumes fields start at 0 before J0 is added """ e0 = sim.es[0] j0 = sim.js[0] u = fdtd.delta_energy_j(j0=j0, e1=e0, dxes=sim.dxes).sum() args = {'dxes': sim.dxes, 'epsilon': sim.epsilon} for ii in range(1, 8): u_hstep = fdtd.energy_hstep(e0=sim.es[ii-1], h1=sim.hs[ii], e2=sim.es[ii], **args) u_estep = fdtd.energy_estep(h0=sim.hs[ii], e1=sim.es[ii], h2=sim.hs[ii + 1], **args) delta_j_A = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii-1], dxes=sim.dxes) delta_j_B = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii], dxes=sim.dxes) u += delta_j_A.sum() assert_close(u_hstep.sum(), u) u += delta_j_B.sum() assert_close(u_estep.sum(), u) def test_poynting_divergence(sim): args = {'dxes': sim.dxes, 'epsilon': sim.epsilon} dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0) u_eprev = None for ii in range(1, 8): u_hstep = fdtd.energy_hstep(e0=sim.es[ii-1], h1=sim.hs[ii], e2=sim.es[ii], **args) u_estep = fdtd.energy_estep(h0=sim.hs[ii], e1=sim.es[ii], h2=sim.hs[ii + 1], **args) delta_j_B = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii], dxes=sim.dxes) du_half_h2e = u_estep - u_hstep - delta_j_B div_s_h2e = sim.dt * fdtd.poynting_divergence(e=sim.es[ii], h=sim.hs[ii], dxes=sim.dxes) * dV assert_fields_close(du_half_h2e, -div_s_h2e) if u_eprev is None: u_eprev = u_estep continue # previous half-step delta_j_A = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii-1], dxes=sim.dxes) du_half_e2h = u_hstep - u_eprev - delta_j_A div_s_e2h = sim.dt * fdtd.poynting_divergence(e=sim.es[ii-1], h=sim.hs[ii], dxes=sim.dxes) * dV assert_fields_close(du_half_e2h, -div_s_e2h) u_eprev = u_estep def test_poynting_planes(sim): mask = (sim.js[0] != 0) if mask.sum() > 1: pytest.skip('test_poynting_planes can only test single point sources') args = {'dxes': sim.dxes, 'epsilon': sim.epsilon} dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0) mx = numpy.roll(mask, (-1, -1), axis=(0, 1)) my = numpy.roll(mask, -1, axis=2) mz = numpy.roll(mask, (+1, -1), axis=(0, 3)) px = numpy.roll(mask, -1, axis=0) py = mask.copy() pz = numpy.roll(mask, +1, axis=0) u_eprev = None for ii in range(1, 8): u_hstep = fdtd.energy_hstep(e0=sim.es[ii-1], h1=sim.hs[ii], e2=sim.es[ii], **args) u_estep = fdtd.energy_estep(h0=sim.hs[ii], e1=sim.es[ii], h2=sim.hs[ii + 1], **args) s_h2e = -fdtd.poynting(e=sim.es[ii], h=sim.hs[ii]) * sim.dt s_h2e[0] *= sim.dxes[0][1][None, :, None] * sim.dxes[0][2][None, None, :] s_h2e[1] *= sim.dxes[0][0][:, None, None] * sim.dxes[0][2][None, None, :] s_h2e[2] *= sim.dxes[0][0][:, None, None] * sim.dxes[0][1][None, :, None] planes = [s_h2e[px].sum(), -s_h2e[mx].sum(), s_h2e[py].sum(), -s_h2e[my].sum(), s_h2e[pz].sum(), -s_h2e[mz].sum()] assert_close(sum(planes), (u_estep - u_hstep).sum()) if u_eprev is None: u_eprev = u_estep continue s_e2h = -fdtd.poynting(e=sim.es[ii - 1], h=sim.hs[ii]) * sim.dt s_e2h[0] *= sim.dxes[0][1][None, :, None] * sim.dxes[0][2][None, None, :] s_e2h[1] *= sim.dxes[0][0][:, None, None] * sim.dxes[0][2][None, None, :] s_e2h[2] *= sim.dxes[0][0][:, None, None] * sim.dxes[0][1][None, :, None] planes = [s_e2h[px].sum(), -s_e2h[mx].sum(), s_e2h[py].sum(), -s_e2h[my].sum(), s_e2h[pz].sum(), -s_e2h[mz].sum()] assert_close(sum(planes), (u_hstep - u_eprev).sum()) # previous half-step u_eprev = u_estep ##################################### # Test fixtures ##################################### @pytest.fixture(scope='module', params=[(5, 5, 1), (5, 1, 5), (5, 5, 5), # (7, 7, 7), ]) def shape(request): yield (3, *request.param) @pytest.fixture(scope='module', params=[0.3]) def dt(request): yield request.param @pytest.fixture(scope='module', params=[1.0, 1.5]) def epsilon_bg(request): yield request.param @pytest.fixture(scope='module', params=[1.0, 2.5]) def epsilon_fg(request): yield request.param @pytest.fixture(scope='module', params=['center', '000', 'random']) def epsilon(request, shape, epsilon_bg, epsilon_fg): is3d = (numpy.array(shape) == 1).sum() == 0 if is3d: if request.param == '000': pytest.skip('Skipping 000 epsilon because test is 3D (for speed)') if epsilon_bg != 1: pytest.skip('Skipping epsilon_bg != 1 because test is 3D (for speed)') if epsilon_fg not in (1.0, 2.0): pytest.skip('Skipping epsilon_fg not in (1, 2) because test is 3D (for speed)') epsilon = numpy.full(shape, epsilon_bg, dtype=float) if request.param == 'center': epsilon[:, shape[1]//2, shape[2]//2, shape[3]//2] = epsilon_fg elif request.param == '000': epsilon[:, 0, 0, 0] = epsilon_fg elif request.param == 'random': epsilon[:] = prng.uniform(low=min(epsilon_bg, epsilon_fg), high=max(epsilon_bg, epsilon_fg), size=shape) yield epsilon @pytest.fixture(scope='module', params=[1.0])#, 1.5]) def j_mag(request): yield request.param @pytest.fixture(scope='module', params=['center', 'random']) def j_distribution(request, shape, j_mag): j = numpy.zeros(shape) if request.param == 'center': j[:, shape[1]//2, shape[2]//2, shape[3]//2] = j_mag elif request.param == '000': j[:, 0, 0, 0] = j_mag elif request.param == 'random': j[:] = prng.uniform(low=-j_mag, high=j_mag, size=shape) yield j @pytest.fixture(scope='module', params=[1.0, 1.5]) def dx(request): yield request.param @pytest.fixture(scope='module', params=['uniform']) def dxes(request, shape, dx): if request.param == 'uniform': dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)] yield dxes @pytest.fixture(scope='module', params=[(0,), (0, 4, 8), ] ) def j_steps(request): yield request.param @dataclasses.dataclass() class SimResult: shape: Tuple[int] dt: float dxes: List[List[numpy.ndarray]] epsilon: numpy.ndarray j_distribution: numpy.ndarray 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(scope='module') def sim(request, shape, epsilon, dxes, dt, j_distribution, j_steps): is3d = (numpy.array(shape) == 1).sum() == 0 if is3d: if dt != 0.3: pytest.skip('Skipping dt != 0.3 because test is 3D (for speed)') sim = SimResult( shape=shape, dt=dt, dxes=dxes, epsilon=epsilon, j_distribution=j_distribution, j_steps=j_steps, ) e = numpy.zeros_like(epsilon) h = numpy.zeros_like(epsilon) assert 0 in j_steps j_zeros = numpy.zeros_like(j_distribution) eh2h = fdtd.maxwell_h(dt=dt, dxes=dxes) eh2e = fdtd.maxwell_e(dt=dt, dxes=dxes) for tt in range(10): e = e.copy() h = h.copy() eh2h(e, h) eh2e(e, h, epsilon) if tt in j_steps: e += j_distribution / epsilon sim.js.append(j_distribution) else: sim.js.append(j_zeros) sim.es.append(e) sim.hs.append(h) return sim