Add function to more quickly compute exp(1j*x)

This is a major time sink in the calculation of the control matrix using
calculate_control_matrix_from_scratch(). Computing real and imaginary
part separately and writing to a preallocated array should make the
whole calculation faster by a factor of two.

Author: thangleiter

filter_functions/numeric.py

@@ -56,7 +56,7 @@
 from .basis import Basis, ggm_expand
 from .plotting import plot_infidelity_convergence
 from .types import Coefficients, Operator
-from .util import (abs2, get_indices_from_identifiers, progressbar,
+from .util import (abs2, cexp, get_indices_from_identifiers, progressbar,
                    symmetrize_spectrum)
 
 __all__ = ['diagonalize', 'liouville_representation',
@@ -108,8 +108,7 @@ def diagonalize(H: ndarray, dt: Coefficients) -> Sequence[ndarray]:
     # P = np.empty((500, 4, 4), dtype=complex)
     # for l, (V, D) in enumerate(zip(HV, np.exp(-1j*dt*HD.T).T)):
     #     P[l] = (V * D) @ V.conj().T
-    P = np.einsum('lij,jl,lkj->lik', HV, np.exp(-1j*np.asarray(dt)*HD.T),
-                  HV.conj())
+    P = np.einsum('lij,jl,lkj->lik', HV, cexp(-np.asarray(dt)*HD.T), HV.conj())
     # The cumulative propagator Q with the identity operator as first
     # element (Q_0 = P_0 = I), i.e.
     # Q = [Q_0, Q_1, ..., Q_n] = [P_0, P_1 @ P_0, ..., P_n @ ... @ P_0]
@@ -250,7 +249,7 @@ def progressbar_range(*args):
         # Mask the integral to avoid convergence problems
         mask_small = np.abs(EdE*dt[l]) <= 1e-7
         integral[mask_small] = dt[l]
-        integral[~mask_small] = (np.exp(1j*EdE[~mask_small]*dt[l]) - 1) /\
+        integral[~mask_small] = (cexp(EdE[~mask_small]*dt[l]) - 1) /\
                                 (1j*(EdE[~mask_small]))
         """
         # Test convergence of integral as argument of exponential -> 0
@@ -263,7 +262,7 @@ def progressbar_range(*args):
         # Faster for d = 2 to also contract over the time dimension instead of
         # loop, but for readability we don't distinguish.
         R += np.einsum('o,j,jmn,omn,knm->jko',
-                       np.exp(1j*E*t[l]), n_coeffs[:, l], B[:, l], integral,
+                       cexp(E*t[l]), n_coeffs[:, l], B[:, l], integral,
                        QdagV[l].conj().T @ basis @ QdagV[l],
                        optimize=R_path)
 

filter_functions/pulse_sequence.py

@@ -57,10 +57,11 @@
                       calculate_control_matrix_from_scratch,
                       calculate_control_matrix_periodic,
                       calculate_filter_function,
-                      calculate_pulse_correlation_filter_function, diagonalize,
+                      calculate_pulse_correlation_filter_function,
+                      diagonalize,
                       liouville_representation)
 from .types import Coefficients, Hamiltonian, Operator, PulseMapping
-from .util import (CalculationError, all_array_equal,
+from .util import (CalculationError, all_array_equal, cexp,
                    get_indices_from_identifiers, hash_array_along_axis, mdot,
                    tensor, tensor_insert, tensor_merge, tensor_transpose)
 
@@ -659,7 +660,7 @@ def cache_total_phases(self, omega: Coefficients,
             they are computed.
         """
         if total_phases is None:
-            total_phases = np.exp(1j*np.asarray(omega)*self.t[-1])
+            total_phases = cexp(np.asarray(omega)*self.t[-1])
 
         self.omega = omega
         self._total_phases = total_phases
@@ -814,7 +815,7 @@ def propagator_at_arb_t(self, t: Coefficients) -> ndarray:
         idx[idx < 0] = 0
         Q_prev = self.Q[idx]
         U_curr = np.einsum('lij,jl,lkj->lik', self.HV[idx],
-                           np.exp(-1j*(t - self.t[idx])*self.HD[idx].T),
+                           cexp((self.t[idx] - t)*self.HD[idx].T),
                            self.HV[idx].conj())
 
         return U_curr @ Q_prev

filter_functions/util.py

@@ -176,7 +176,7 @@ def _get_notebook_name() -> str:
 P_np = [P.full() for P in P_qt]
 
 
-def abs2(x):
+def abs2(x: ndarray) -> ndarray:
     r"""
     Fast function to calculate the absolute value squared,
 
@@ -191,6 +191,31 @@ def abs2(x):
     return x.real**2 + x.imag**2
 
 
+def cexp(x: ndarray) -> ndarray:
+    r"""Fast complex exponential.
+
+    Parameters
+    ----------
+    x : ndarray
+        Argument of the complex exponential :math:`\exp(i x)`.
+
+    Returns
+    -------
+    y : ndarray
+        Complex exponential :math:`y = \exp(i x)`.
+
+    References
+    ----------
+    https://software.intel.com/en-us/forums/intel-distribution-for-python/topic/758148  # noqa
+    """
+    df_exp = np.empty(x.shape, dtype=np.complex128)
+    trig_buf = np.cos(x)
+    df_exp.real[:] = trig_buf
+    np.sin(x, out=trig_buf)
+    df_exp.imag[:] = trig_buf
+    return df_exp
+
+
 def _tensor_product_shape(shape_A: Sequence[int], shape_B: Sequence[int],
                           rank: int):
     """Get shape of the tensor product between A and B of rank rank"""

tests/test_util.py

@@ -35,6 +35,22 @@ def test_abs2(self):
         x = randn(20, 100) + 1j*randn(20, 100)
         self.assertArrayAlmostEqual(np.abs(x)**2, util.abs2(x))
 
+    def test_cexp(self):
+        """Fast complex exponential."""
+        x = randn(50, 100)
+        a = util.cexp(x)
+        b = np.exp(1j*x)
+        self.assertArrayAlmostEqual(a, b)
+
+        a = util.cexp(-x)
+        b = np.exp(-1j*x)
+        self.assertArrayAlmostEqual(a, b)
+
+        x = randn(50, 100)*(1 + 1j)
+        a = util.cexp(x)
+        b = np.exp(1j*x)
+        self.assertArrayAlmostEqual(a, b)
+
     def test_get_indices_from_identifiers(self):
         pulse = PulseSequence(
             [[util.P_np[3], [2], 'Z'],