diff --git a/filter_functions/numeric.py b/filter_functions/numeric.py index 49d57bc..cbe50b3 100644 --- a/filter_functions/numeric.py +++ b/filter_functions/numeric.py @@ -170,7 +170,7 @@ def _first_order_integral(E: ndarray, eigvals: ndarray, dt: float, def _second_order_integral(E: ndarray, eigvals: ndarray, dt: float, int_buf: ndarray, frc_bufs: Tuple[ndarray, ndarray], dE_bufs: Tuple[ndarray, ndarray, ndarray], - msk_bufs: Tuple[ndarray, ndarray, ndarray, ndarray]) -> ndarray: + msk_bufs: Tuple[ndarray, ndarray]) -> ndarray: r"""Calculate the nested integral of second order Magnus expansion. The integral is evaluated as @@ -200,32 +200,21 @@ def _second_order_integral(E: ndarray, eigvals: ndarray, dt: float, # frc_buf1 has shape (len(E), *dE.shape), frc_buf2 has shape dE.shape*2 frc_buf1, frc_buf2 = frc_bufs dEdE, EdE, dEE = dE_bufs - mask_nE_dE_dE, mask_nE_ndE_dE, mask_nE_dE_ndE, mask_nEdE_ndEE = msk_bufs + mask_nEdE_dEE, mask_nEdE_ndEE = msk_bufs dE = np.subtract.outer(eigvals, eigvals) dEdE = np.add.outer(dE, dE, out=dEdE) EdE = np.add.outer(E, dE, out=EdE) dEE = np.add.outer(-E, dE, out=dEE) - mask_E = E != 0 - mask_dE = dE != 0 mask_dEE = dEE != 0 - mask_EdE = EdE != 0 mask_dEdE = dEdE != 0 - mask_dE_dE = mask_dE[..., None, None] & mask_dE - mask_ndE_dE = ~mask_dE[..., None, None] & mask_dE - mask_dE_ndE = mask_dE[..., None, None] & ~mask_dE - mask_ndE_ndE = ~mask_dE[..., None, None] & ~mask_dE - mask_nE_dE_dE = np.logical_and(~mask_E[:, None, None, None, None], mask_dE_dE, - out=mask_nE_dE_dE) - mask_nE_ndE_dE = np.logical_and(~mask_E[:, None, None, None, None], mask_ndE_dE, - out=mask_nE_ndE_dE) - mask_nE_dE_ndE = np.logical_and(~mask_E[:, None, None, None, None], mask_dE_ndE, - out=mask_nE_dE_ndE) - mask_nEdE_ndEE = np.logical_and(~mask_E[:, None, None, None, None], mask_ndE_ndE, - out=mask_nEdE_ndEE) - mask_EdE_dEE = np.broadcast_to(mask_EdE[:, None, None], int_buf.shape) + mask_EdE = np.broadcast_to((EdE != 0)[:, None, None], int_buf.shape) + mask_nEdE_dEE = np.logical_and(~mask_EdE, mask_dEE[..., None, None], out=mask_nEdE_dEE) + mask_nEdE_ndEE = np.logical_and(~mask_EdE, ~mask_dEE[..., None, None], out=mask_nEdE_ndEE) + # ω + Ω_mn != 0 + # This buffer holds the term (exp(Ω_ij - ω) - 1)/(Ω_ij - ω), which we reuse later. frc_buf1 = util.cexpm1(dEE*dt, where=mask_dEE, out=frc_buf1) frc_buf1 = np.divide(frc_buf1, dEE, where=mask_dEE, out=frc_buf1) frc_buf1[~mask_dEE] = 1j*dt @@ -235,23 +224,18 @@ def _second_order_integral(E: ndarray, eigvals: ndarray, dt: float, frc_buf2[~mask_dEdE] = 1j*dt int_buf = np.subtract(frc_buf1[..., None, None], frc_buf2[None, ...], - out=int_buf, where=mask_EdE_dEE) - int_buf = np.divide(int_buf, EdE[:, None, None], out=int_buf, where=mask_EdE_dEE) - - # Limiting cases for ω -> 0 - dEt = dE*dt - # Ω_ij == 0, Ω_mn != 0 - int_buf = np.divide(1j*dt - np.divide(util.cexpm1(dEt), dE, where=mask_dE), dE, - where=mask_nE_ndE_dE, out=int_buf) - # Ω_ij != 0, Ω_mn == 0 - int_buf = np.divide((np.divide(util.cexpm1(dEt), dE, where=mask_dE) - - 1j*dt*util.cexp(dEt))[..., None, None], - dE[..., None, None], where=mask_nE_dE_ndE, out=int_buf) - # Ω_ij != 0, Ω_mn != 0 - int_buf = np.subtract(np.divide(util.cexpm1(dEt), dE, where=mask_dE)[..., None, None], - frc_buf2, where=mask_nE_dE_dE, out=int_buf) - int_buf = np.divide(int_buf, dE, where=mask_nE_dE_dE, out=int_buf) - # Ω_ij == 0, Ω_mn == 0 + out=int_buf, where=mask_EdE) + int_buf = np.divide(int_buf, EdE[:, None, None], out=int_buf, where=mask_EdE) + + # Limiting cases + # ω + Ω_mn = 0, Ω_ij - ω != 0 + frc_buf1 = np.subtract(frc_buf1, 1j*dt*util.cexp(dEE*dt, where=mask_dEE), + where=mask_dEE, out=frc_buf1) + frc_buf1 = np.divide(frc_buf1, dEE, where=mask_dEE, out=frc_buf1) + int_buf[mask_nEdE_dEE] = np.broadcast_to(frc_buf1[..., None, None], int_buf.shape)[ + mask_nEdE_dEE + ] + # ω + Ω_mn = 0, Ω_ij - ω = 0 int_buf[mask_nEdE_ndEE] = dt**2 / 2 return int_buf @@ -1586,7 +1570,7 @@ def calculate_second_order_filter_function_from_scratch( np.empty((n_omega, d, d), dtype=float), np.empty((n_omega, d, d), dtype=float)) frc_bufs = (np.empty((n_omega, d, d), dtype=complex), np.empty((d, d, d, d), dtype=complex)) - msk_bufs = np.empty((4, n_omega, d, d, d, d), dtype=bool) + msk_bufs = np.empty((2, n_omega, d, d, d, d), dtype=bool) n_opers_basis_buf = np.empty((n_nops, n_basis, d, d), dtype=complex) complete_steps = np.zeros((n_nops, n_nops, n_basis, n_basis, n_omega), dtype=complex) incomplete_steps = np.zeros_like(complete_steps) diff --git a/tests/test_precision.py b/tests/test_precision.py index 1999be8..711e56d 100644 --- a/tests/test_precision.py +++ b/tests/test_precision.py @@ -54,7 +54,7 @@ def _get_integrals_second_order(d, E, eigval, dt, t0): frc_bufs = (np.empty((len(E), d, d), dtype=complex), np.empty((d, d, d, d), dtype=complex)) int_buf = np.empty((len(E), d, d, d, d), dtype=complex) - msk_bufs = np.empty((4, len(E), d, d, d, d), dtype=bool) + msk_bufs = np.empty((2, len(E), d, d, d, d), dtype=bool) tspace = np.linspace(0, dt, 1001) + t0 dE = np.subtract.outer(eigval, eigval) @@ -492,6 +492,22 @@ def test_integration(self): integral, integral_numeric = _get_integrals_second_order(d, E, eigval, dt, t) self.assertArrayAlmostEqual(integral, integral_numeric, atol=1e-4) + def test_integration_edge_cases(self): + pulse = testutil.rand_pulse_sequence(testutil.rng.integers(2, 10), + testutil.rng.integers(1, 5)) + for i, (eigval, dt, t) in enumerate(zip(pulse.eigvals, pulse.dt, pulse.t)): + # \Omega_ij = \omega or \Omega_mn = -\omega + omega = np.repeat(testutil.rng.choice([-1, 1]) + * np.diff(testutil.rng.choice(eigval, 2, replace=False)), + 3) + omega[0] -= 1e-10 + omega[2] += 1e-10 + integral, integral_numeric = _get_integrals_first_order(pulse.d, omega, eigval, dt, t) + self.assertArrayAlmostEqual(integral, integral_numeric, atol=1e-4) + + integral, integral_numeric = _get_integrals_second_order(pulse.d, omega, eigval, dt, t) + self.assertArrayAlmostEqual(integral, integral_numeric, atol=1e-4) + def test_infidelity(self): """Benchmark infidelity results against previous version's results""" rng = np.random.default_rng(seed=123456789)