Small improvements

filter_functions/basis.py

@@ -553,8 +553,8 @@ def _full_from_partial(elems: Sequence, traceless: bool, labels: Sequence[str])
         traceless = elems.istraceless
     else:
         if traceless and not elems.istraceless:
-            raise ValueError("The basis elements are not traceless (up to an identity element) " +
-                             "but a traceless basis was requested!")
+            raise ValueError("The basis elements are not traceless (up to an identity element) "
+                             + "but a traceless basis was requested!")
 
     if labels is not None and len(labels) not in (len(elems), elems.d**2):
         raise ValueError(f'Got {len(labels)} labels but expected {len(elems)} or {elems.d**2}')
@@ -677,7 +677,8 @@ def expand(M: Union[ndarray, Basis], basis: Union[ndarray, Basis],
 
     """
 
-    def cast(arr): return arr.real if hermitian and basis.isherm else arr
+    def cast(arr):
+        return arr.real if hermitian and basis.isherm else arr
 
     coefficients = cast(np.tensordot(M, basis, axes=[(-2, -1), (-1, -2)]))
     if not normalized:
@@ -724,7 +725,8 @@ def ggm_expand(M: Union[ndarray, Basis], traceless: bool = False,
     if M.shape[-1] != M.shape[-2]:
         raise ValueError('M should be square in its last two axes')
 
-    def cast(arr): return arr.real if hermitian else arr
+    def cast(arr):
+        return arr.real if hermitian else arr
 
     # Squeeze out an extra dimension to be shape agnostic
     square = M.ndim < 3

filter_functions/gradient.py

@@ -55,8 +55,8 @@
 
 import numpy as np
 import opt_einsum as oe
-from opt_einsum.contract import ContractExpression
 from numpy import ndarray
+from opt_einsum.contract import ContractExpression
 
 from . import numeric, superoperator, util
 from .basis import Basis
@@ -656,13 +656,13 @@ def infidelity_derivative(
         Section A: General, Atomic and Solid State Physics, 303(4), 249–252.
         https://doi.org/10.1016/S0375-9601(02)01272-0
     """
-    spectrum = numeric._parse_spectrum(spectrum, omega, range(len(pulse.n_opers)))
+    spectrum = util.parse_spectrum(spectrum, omega, range(len(pulse.n_opers)))
     filter_function_deriv = pulse.get_filter_function_derivative(omega,
                                                                  control_identifiers,
                                                                  n_oper_identifiers,
                                                                  n_coeffs_deriv)
 
-    integrand = np.einsum('ao,atho->atho', spectrum, filter_function_deriv)
+    integrand = np.einsum('...o,...tho->...tho', spectrum, filter_function_deriv)
     infid_deriv = util.integrate(integrand, omega) / (2*np.pi*pulse.d)
 
     return infid_deriv

filter_functions/numeric.py

@@ -249,24 +249,6 @@ def _second_order_integral(E: ndarray, eigvals: ndarray, dt: float,
     return int_buf
 
 
-def _parse_spectrum(spectrum: Sequence, omega: Sequence, idx: Sequence) -> ndarray:
-    spectrum = np.asanyarray(spectrum)
-    error = 'Spectrum should be of shape {}, not {}.'
-    shape = (len(idx),)*(spectrum.ndim - 1) + (len(omega),)
-    if spectrum.shape != shape and spectrum.ndim <= 3:
-        raise ValueError(error.format(shape, spectrum.shape))
-
-    if spectrum.ndim == 1:
-        # As we broadcast over the noise operators
-        spectrum = spectrum[None, ...]
-    if spectrum.ndim == 3 and not np.allclose(spectrum, spectrum.conj().swapaxes(0, 1)):
-        raise ValueError('Cross-spectra given but not Hermitian along first two axes')
-    elif spectrum.ndim > 3:
-        raise ValueError(f'Expected spectrum to have < 4 dimensions, not {spectrum.ndim}')
-
-    return spectrum
-
-
 def _get_integrand(
         spectrum: ndarray,
         omega: ndarray,
@@ -330,7 +312,7 @@ def _get_integrand(
             # Everything simpler if noise operators always on 2nd-to-last axes
             filter_function = np.moveaxis(filter_function, source=[-5, -4], destination=[-3, -2])
 
-    spectrum = _parse_spectrum(spectrum, omega, idx)
+    spectrum = util.parse_spectrum(spectrum, omega, idx)
     if spectrum.ndim in (1, 2):
         if filter_function is not None:
             integrand = (filter_function[..., tuple(idx), tuple(idx), :]*spectrum)
@@ -342,17 +324,17 @@ def _get_integrand(
             # R is not None
             if which_pulse == 'correlations':
                 if which_FF == 'fidelity':
-                    einsum_str = 'gako,ao,hako->ghao'
+                    einsum_str = 'g...ko,...o,h...ko->gh...o'
                 else:
                     # which_FF == 'generalized'
-                    einsum_str = 'gako,ao,halo->ghaklo'
+                    einsum_str = 'g...ko,...o,h...lo->gh...klo'
             else:
                 # which_pulse == 'total'
                 if which_FF == 'fidelity':
-                    einsum_str = 'ako,ao,ako->ao'
+                    einsum_str = '...ko,...o,...ko->...o'
                 else:
                     # which_FF == 'generalized'
-                    einsum_str = 'ako,ao,alo->aklo'
+                    einsum_str = '...ko,...o,...lo->...klo'
 
             integrand = np.einsum(einsum_str,
                                   ctrl_left[..., idx, :, :], spectrum, ctrl_right[..., idx, :, :])
@@ -690,7 +672,8 @@ def calculate_control_matrix_from_atomic(
         control_matrix = np.zeros(control_matrix_atomic.shape, dtype=complex)
         for g in util.progressbar_range(n, show_progressbar=show_progressbar,
                                         desc='Calculating control matrix'):
-            control_matrix[g] = expr(phases[g]*control_matrix_atomic[g], propagators_liouville[g])
+            control_matrix[g] = expr(phases[g]*control_matrix_atomic[g], propagators_liouville[g],
+                                     out=control_matrix[g])
 
     return control_matrix
 
@@ -1068,8 +1051,8 @@ def calculate_cumulant_function(
     N, d = pulse.basis.shape[:2]
     if spectrum is None and omega is None:
         if decay_amplitudes is None or (frequency_shifts is None and second_order):
-            raise ValueError('Require either spectrum and frequencies or precomputed ' +
-                             'decay amplitudes (frequency shifts)')
+            raise ValueError('Require either spectrum and frequencies or precomputed '
+                             + 'decay amplitudes (frequency shifts)')
 
     if which == 'correlations' and second_order:
         raise ValueError('Cannot compute correlation cumulant function for second order terms')
@@ -1242,8 +1225,8 @@ def calculate_decay_amplitudes(
         # which == 'correlations'
         if pulse.is_cached('omega'):
             if not np.array_equal(pulse.omega, omega):
-                raise ValueError('Pulse correlation decay amplitudes requested but omega not ' +
-                                 'equal to cached frequencies.')
+                raise ValueError('Pulse correlation decay amplitudes requested but omega not '
+                                 + 'equal to cached frequencies.')
 
         if pulse.is_cached('filter_function_pc_gen'):
             control_matrix = None
@@ -1810,8 +1793,8 @@ def error_transfer_matrix(
     """
     if cumulant_function is None:
         if pulse is None or spectrum is None or omega is None:
-            raise ValueError('Require either precomputed cumulant function ' +
-                             'or pulse, spectrum, and omega as arguments.')
+            raise ValueError('Require either precomputed cumulant function '
+                             + 'or pulse, spectrum, and omega as arguments.')
 
         cumulant_function = calculate_cumulant_function(pulse, spectrum, omega,
                                                         n_oper_identifiers, 'total', second_order,
@@ -2011,8 +1994,8 @@ def infidelity(
         try:
             omega_IR = omega.get('omega_IR', 2*np.pi/pulse.tau*1e-2)
         except AttributeError:
-            raise TypeError('omega should be dictionary with parameters ' +
-                            'when test_convergence == True.')
+            raise TypeError('omega should be dictionary with parameters '
+                            + 'when test_convergence == True.')
 
         omega_UV = omega.get('omega_UV', 2*np.pi/pulse.tau*1e+2)
         spacing = omega.get('spacing', 'linear')
@@ -2049,8 +2032,8 @@ def infidelity(
             # but trace tensor plays a role, cf eq. (39). For traceless bases,
             # the trace tensor term reduces to delta_ij.
             traces = pulse.basis.four_element_traces
-            traces_diag = (sparse.diagonal(traces, axis1=2, axis2=3).sum(-1) -
-                           sparse.diagonal(traces, axis1=1, axis2=3).sum(-1)).todense()
+            traces_diag = (sparse.diagonal(traces, axis1=2, axis2=3).sum(-1)
+                           - sparse.diagonal(traces, axis1=1, axis2=3).sum(-1)).todense()
 
             control_matrix = pulse.get_control_matrix(omega, show_progressbar, cache_intermediates)
             filter_function = np.einsum('ako,blo,kl->abo',
@@ -2061,14 +2044,9 @@ def infidelity(
                                                         cache_intermediates=cache_intermediates)
     else:
         # which == 'correlations'
-        if not pulse.basis.istraceless:
-            warn('Calculating pulse correlation fidelities with non-' +
-                 'traceless basis. The results will be off.')
-
-        if pulse.is_cached('omega'):
-            if not np.array_equal(pulse.omega, omega):
-                raise ValueError('Pulse correlation infidelities requested ' +
-                                 'but omega not equal to cached frequencies.')
+        if pulse.is_cached('omega') and not np.array_equal(pulse.omega, omega):
+            raise ValueError('Pulse correlation infidelities requested '
+                             + 'but omega not equal to cached frequencies.')
 
         filter_function = pulse.get_pulse_correlation_filter_function()
 
@@ -2078,8 +2056,8 @@ def infidelity(
 
     if return_smallness:
         if spectrum.ndim > 2:
-            raise NotImplementedError('Smallness parameter only implemented ' +
-                                      'for uncorrelated noise sources')
+            raise NotImplementedError('Smallness parameter only implemented '
+                                      + 'for uncorrelated noise sources')
 
         T1 = util.integrate(spectrum, omega)/(2*np.pi)
         T2 = (pulse.dt*pulse.n_coeffs[idx]).sum(axis=-1)**2

filter_functions/plotting.py

@@ -55,8 +55,8 @@
 from numpy import ndarray
 
 from . import numeric, util
-from .types import (Axes, Coefficients, Colormap, Figure, FigureAxes, FigureAxesLegend, FigureGrid,
-                    Grid, Operator, State)
+from .types import (Axes, Coefficients, Colormap, Cycler, Figure, FigureAxes, FigureAxesLegend,
+                    FigureGrid, Grid, Operator, State)
 
 __all__ = ['plot_cumulant_function', 'plot_infidelity_convergence', 'plot_filter_function',
            'plot_pulse_correlation_filter_function', 'plot_pulse_train']
@@ -129,9 +129,7 @@ def init_bloch_sphere(**bloch_kwargs) -> qt.Bloch:
     return b
 
 
-@util.parse_optional_parameters(prop=('total', 'piecewise'))
-def get_states_from_prop(U: Sequence[Operator], psi0: Optional[State] = None,
-                         prop: str = 'total') -> ndarray:
+def get_states_from_prop(U: Sequence[Operator], psi0: Optional[State] = None) -> ndarray:
     r"""
     Get the the quantum state at time t from the propagator and the
     inital state:
@@ -140,31 +138,18 @@ def get_states_from_prop(U: Sequence[Operator], psi0: Optional[State] = None,
 
         |\psi(t)\rangle = U(t, 0)|\psi(0)\rangle
 
-    If *prop* is 'piecewise', then it is assumed that *U* is the
-    propagator of a piecewise-constant control:
-
-    .. math::
-        |\psi(t)\rangle = \prod_{l=1}^n U(t_l, t_{l-1})|\psi(0)\rangle
-
-    with :math:`t_0\equiv 0` and :math:`t_n\equiv t`.
-
     """
     if psi0 is None:
-        psi0 = np.c_[1:-1:-1]  # |0>
-
-    psi0 = psi0.full() if hasattr(psi0, 'full') else psi0  # qutip.Qobj
-    d = max(psi0.shape)
-    states = np.empty((len(U), d, 1), dtype=complex)
-    if prop == 'total':
-        for j in range(len(U)):
-            states[j] = U[j] @ psi0
-    else:
-        # prop == 'piecewise'
-        states[0] = U[0] @ psi0
-        for j in range(1, len(U)):
-            states[j] = U[j] @ states[j-1]
+        # default to |0>
+        psi0 = np.c_[1:-1:-1]
+    elif hasattr(psi0, 'full'):
+        # qutip.Qobj
+        psi0 = psi0.full()
+
+    if psi0.shape[-2:] != (2, 1):
+        raise ValueError('Initial state should be shape (..., 2, 1)')
 
-    return states
+    return U @ psi0
 
 
 def plot_bloch_vector_evolution(
@@ -289,7 +274,7 @@ def plot_pulse_train(
         c_oper_identifiers: Optional[Sequence[int]] = None,
         fig: Optional[Figure] = None,
         axes: Optional[Axes] = None,
-        cycler: Optional['cycler.Cycler'] = None,
+        cycler: Optional[Cycler] = None,
         plot_kw: Optional[dict] = {},
         subplot_kw: Optional[dict] = None,
         gridspec_kw: Optional[dict] = None,
@@ -380,7 +365,7 @@ def plot_filter_function(
         xscale: str = 'log',
         yscale: str = 'linear',
         omega_in_units_of_tau: bool = True,
-        cycler: Optional['cycler.Cycler'] = None,
+        cycler: Optional[Cycler] = None,
         plot_kw: dict = {},
         subplot_kw: Optional[dict] = None,
         gridspec_kw: Optional[dict] = None,
@@ -510,7 +495,7 @@ def plot_pulse_correlation_filter_function(
         xscale: str = 'log',
         yscale: str = 'linear',
         omega_in_units_of_tau: bool = True,
-        cycler: Optional['cycler.Cycler'] = None,
+        cycler: Optional[Cycler] = None,
         plot_kw: dict = {},
         subplot_kw: Optional[dict] = None,
         gridspec_kw: Optional[dict] = None,
@@ -731,7 +716,7 @@ def plot_cumulant_function(
 
     Parameters
     ----------
-    pulse: 'PulseSequence'
+    pulse: PulseSequence
         The pulse sequence.
     spectrum: ndarray
         The two-sided noise spectrum.
@@ -801,13 +786,15 @@ def plot_cumulant_function(
                 n_oper_identifiers = [f'$B_{{{i}}}$' for i in range(len(n_oper_inds))]
         else:
             if len(n_oper_identifiers) != len(K):
-                raise ValueError('Both precomputed cumulant function and n_oper_identifiers ' +
-                                 f'given but not same len: {len(K)} != {len(n_oper_identifiers)}')
+                raise ValueError(
+                    'Both precomputed cumulant function and n_oper_identifiers '
+                    + f'given but not same len: {len(K)} != {len(n_oper_identifiers)}'
+                )
 
     else:
         if pulse is None or spectrum is None or omega is None:
-            raise ValueError('Require either precomputed cumulant function ' +
-                             'or pulse, spectrum, and omega as arguments.')
+            raise ValueError('Require either precomputed cumulant function '
+                             + 'or pulse, spectrum, and omega as arguments.')
 
         n_oper_inds = util.get_indices_from_identifiers(pulse.n_oper_identifiers,
                                                         n_oper_identifiers)
@@ -852,8 +839,8 @@ def plot_cumulant_function(
         grid = axes_grid1.ImageGrid(fig, **grid_kw)
     else:
         if len(grid) != len(n_oper_inds):
-            raise ValueError('Size of supplied ImageGrid instance does not ' +
-                             'match the number of n_oper_identifiers given!')
+            raise ValueError('Size of supplied ImageGrid instance does not '
+                             + 'match the number of n_oper_identifiers given!')
 
         fig = grid[0].get_figure()