Developer API Reference

The documentation presented for qugrad.QuantumSystem and its child classes is demonstrates their usage by specifying method arguments. The underlying methods actually accept *args to allow developers to extend the functionality of QuGrad. To aid in extending QuGrad’s functionality we include source code for qugrad.QuantumSystem and its child classes below. The documentation in this source code is focused on extending the functionality oposed to utilising it.

"""
Defines classes for quantum systems
"""

from typing import Callable, Optional, Union, Any

import numpy as np
import tensorflow as tf

from py_ste import get_unitary_evolver
from py_ste.evolvers import UnitaryEvolver

from .._hilbert_space import HilbertSpace
from ..pulses import compose_unpack

def generate_channel_couplings(number_channels: list[int]) -> np.ndarray[bool]:
    """Generates a boolean array indicating which channels couple to which
    drives.

    Parameters
    ----------
    number_chanels : list[int]
        A list with the nth entry indicating the number of channels that
        correspond the nth drive.

    Returns
    -------
    NDArray[Shape[``len(number_channels)``, total_number_channels], bool]
        An entry is `True` if the channel corresponds to the drive.
    """
    number_channels: np.ndarray[int] = np.array(number_channels)
    if len(number_channels) == 0:
        return np.empty((0, 0), dtype=bool)
    number_channels = number_channels.cumsum()
    channel_couplings: np.ndarray[bool] = np.zeros((len(number_channels),
                                                    number_channels[-1]),
                                                   dtype=bool)
    c_old = 0
    for i, c in enumerate(number_channels):
        channel_couplings[i, c_old: (c_old := c)] = True
    return channel_couplings

# Defining classes
class ExpValCustom:
    """
    A class implimenting :meth:`QuantumSystem.evolved_expectation_value()` with
    a `TensorFlow <https://www.tensorflow.org>`__ gradient.
    """

    system: "QuantumSystem"
    "The system in which to take the expectation value"
    
    initial_state: np.ndarray[complex]
    "The initial state for the integrator"
    
    dt: float
    "The integrator time step"
    
    observable: np.ndarray[complex]
    "The observable to take the expectation value of"

    switching_function: Any
    """
    The value of the switching function multiplied by :attr:`dt` evaluated at
    each integration timestep
    """
    
    def __init__(self,
                 system: "QuantumSystem",
                 initial_state: np.ndarray[complex],
                 dt: float,
                 observable: np.ndarray[complex]):
        """
        Initialises the class with the `system` in which to take the expectation
        value and the `observable` in the `system` to take the expectation value
        of. Additionally, the initial state before evolution and
        the integrator time step are specified.

        Parameters
        ----------
        system : QuantumSystem
            The system in which to take the expectation value
        initial_state : NDArray[Shape[``system.state_shape``], complex]
            The initial state for the integrator
        dt : float
            The integrator time step
        observable : NDArray[Shape[``system.dim``, ``system.dim``], complex])
            The observable to take the expectation value of
        """
        self.system = system
        self.initial_state = initial_state
        self.dt = dt
        self.observable = observable
    @tf.custom_gradient
    def run(self, ctrl_amp):
        """
        Computes the expectation value of the :attr:`observable` in the
        :attr:`system` with respect to the :attr:`initial_state` evolved under
        the Hamiltonian generated by the specified control amplitudes.

        Parameters
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, ``system.n_ctrl``], tf.complex128]
            $\left(a_{ij}\right)$ The control amplitudes at each time step
            expressed as an $N\times\textrm{length}$ matrix where the element
            $a_{ij}$ corresponds to the control amplitude of the $j$th control
            Hamiltonian at the $i$th time step.

        Returns
        -------
        tf.Tensor[Shape[], tf.complex128]
            The expectation value

        Note
        ----
        This function is differentiable using TensorFlow's ``tf.GradientTape``.

        Note
        ----
        This function uses
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.switching_function()`
        from `PySTE <https://PySTE.readthedocs.io>`__.
        """
        E, self.switching_function = self.system.evolver.switching_function(ctrl_amp.numpy().copy(), self.initial_state, self.dt, self.observable)
        self.switching_function = tf.constant(self.switching_function*self.dt, dtype=tf.complex128)
        def grad(upstream): return self.switching_function
        return tf.math.real(tf.constant(E, dtype=tf.complex128)), grad # @tf.custom_gradient removes the second returned variable

class GateInfidelityCustom():
    """
    A class implimenting :meth:`QuantumSystem.evolved_gate_infidelity()` with
    a `TensorFlow <https://www.tensorflow.org>`__ gradient.
    """
    
    system: "QuantumSystem"
    "The system in which to calculate the gate infidelity"

    dt: float
    "The integrator time step"

    target: np.ndarray[complex]
    "The target gate to calculate the gate infidelity with respect to."

    switching_function: Any
    """
    The value of the switching function multiplied by :attr:`dt` evaluated at
    each integration timestep
    """
    
    def __init__(self,
                 system: "QuantumSystem",
                 dt: float,
                 target: np.ndarray[complex]):
        """
        Initialises the class with the `system` in which to calculate the
        gate infidelity and the `target` gate calculate the gate infidelity
        from. Additionally, the integrator time step is specified.

        Parameters
        ----------
        system : QuantumSystem
            The system in which to calculate the gate infidelity
        dt : float
            The integrator time step
        target : NDArray[Shape[``system.dim``, ``system.dim``], complex])
            The target gate to calculate the gate infidelity with respect to
        """
        self.system = system
        self.dt = dt
        self.target = target
    @tf.custom_gradient
    def run(self, ctrl_amp):
        """
        Computes the gate infidelity from :attr:`target` of the evolution due to
        the :attr:`system` Hamiltonian generated by the specified control
        amplitudes.

        Parameters
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, ``system.n_ctrl``], tf.complex128]
            $\left(a_{ij}\right)$ The control amplitudes at each time step
            expressed as an $N\times\textrm{length}$ matrix where the element
            $a_{ij}$ corresponds to the control amplitude of the $j$th control
            Hamiltonian at the $i$th time step.

        Returns
        -------
        tf.Tensor[Shape[], tf.complex128]
            The gate infidelity

        Note
        ----
        This function is differentiable using TensorFlow's ``tf.GradientTape``.

        Note
        ----
        This function uses
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.gate_switching_function()`
        from `PySTE <https://PySTE.readthedocs.io>`__.
        """
        infidelity, self.switching_function = self.system.evolver.gate_switching_function(ctrl_amp.numpy().copy(), self.dt, self.target)
        self.switching_function = tf.constant(self.switching_function*self.dt, dtype=tf.complex128)
        def grad(upstream): return self.switching_function
        return tf.constant(infidelity, dtype=tf.float64), grad # @tf.custom_gradient removes the second returned variable

class QuantumSystem:
    """
    A class storing the properties of a quantum system.
    """
    _evolver: Optional[UnitaryEvolver] = None
    "The integrator used for time evolutions of the system."
    
    _hilbert_space: HilbertSpace
    "The Hilbert space of the system"
    
    _H0: np.ndarray
    "The systems drift Hamiltonian as a :attr:`dim` x :attr:`dim` matrix."
    
    _Hs: np.ndarray
    """
    An array of the system's control Hamiltonians with shape
    (:attr:`n_ctrl`, :attr:`dim`, :attr:`dim`).
    """
    
    _graph_processing: Callable[..., tuple]
    "A Tensorflow graph of :attr:`_processing()`"
    
    _processing: Callable[..., tuple]
    """
    Executes :meth:`_pre_processing()` followed by
    :meth:`_envolope_processing()` eagerly (i.e. without using a TensorFlow
    graph). Nonetheless, :meth:`_eager_processing()` is still auto
    differentiable.

    Parameters
    ----------
    *args
        The placeholder parameters. See ``_systems.pyi`` for actual parameters.
        Each child class that implements a new :meth:`_pre_processing()` should
        implement a ``.pyi`` file to document the parameters for this function:
        the same parameters as passed to :meth:`_pre_processing()`.

    Returns
    -------
    tuple[tf.Tensor[Shape[n_time_steps, :attr:`n_ctrl`], complex], tf.Tensor[Shape[:attr:`state_shape`], complex], tf.Tensor[Shape[], float]]
        A tuple of:
        1. Control amplitudes
        2. Initial state
        3. Integrator time step
    """
    
    _using_graph: bool
    """
    Whether to use TensorFlow graphs during computation. Using a TensorFlow
    graph will increase the speed of computation. However, you have to be
    careful that function parameters have not been baked into the graph leading
    to unexpected behaviour.
    """
    
    def __init__(self,
                 H0: np.ndarray[complex],
                 Hs: np.ndarray[complex],
                 hilbert_space: HilbertSpace,
                 use_graph: bool = True):
        """
        Initialises a new :class:`QuantumSystem`.

        Parameters
        ----------
        H0 : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The systems drift Hamiltonian
        Hs : NDArray[Shape[:attr:`n_ctrl`, :attr:`dim`, :attr:`dim`], complex] | NDArray[Shape[:attr:`n_ctrl` * :attr:`dim`, :attr:`dim`], complex]
            The systems control Hamiltonians either as an array of control
            Hamiltonians or the control Hamiltonians stacked along the first
            axis.
        hilbert_space : HilbertSpace
            The Hilbert space of the system
        use_graph : bool
            Whether to use `TensorFlow <https://www.tensorflow.org>`__ graphs
            during computation, by default ``True``
        """
        self._hilbert_space: HilbertSpace = hilbert_space
        self._H0 = np.array(H0)
        self._H0.flags.writeable = False
        Hs = np.array(Hs)
        self._Hs = Hs if Hs.ndim == 2 else Hs.reshape(-1, self.dim)
        self._Hs.flags.writeable = False
        self._graph_processing = tf.function(self._traceable_eager_processing, autograph=False)
        self.using_graph = use_graph
            
    def __del__(self):
        # to force clear up of tracing
        del self._graph_processing
        del self._processing
    @property
    def using_graph(self) -> bool:
        """
        Whether to use `TensorFlow <https://www.tensorflow.org>`__ graphs during
        computation. Using a `TensorFlow <https://www.tensorflow.org>`__ graph
        will increase the speed of computation. However, you have to be careful
        that function parameters have not been baked into the graph leading to
        unexpected behaviour.
        """
        return self._using_graph
    @using_graph.setter
    def using_graph(self, value: bool):
        self._using_graph = value
        self._processing = self._graph_processing if value else self._eager_processing
    @property
    def hilbert_space(self) -> HilbertSpace:
        "The Hilbert space of the system"
        return self._hilbert_space
    @property
    def H0(self) -> np.ndarray[complex]:
        """
        The systems drift Hamiltonian as a :attr:`dim` x :attr:`dim` matrix.

        See Also
        --------
        :attr:`Hs`
        """
        return self._H0
    @property
    def Hs(self) -> np.ndarray[complex]:
        """
        An array of the system's control Hamiltonians with shape
        (:attr:`n_ctrl`, :attr:`dim`, :attr:`dim`).

        See Also
        --------
        :attr:`H0`
        """
        return self._Hs.reshape((-1, self.dim, self.dim))
    @property
    def dim(self) -> int:
        """
        The dimension of states in the quantum system.

        See Also
        --------
        :attr:`state_shape`
        """
        return self._hilbert_space.dim
    @property
    def state_shape(self) -> tuple[int]:
        """
        The shape of the states in the system.

        See Also
        --------
        :attr:`dim`
        """
        return (self.dim,)
    @property
    def n_ctrl(self) -> int:
        """
        The number of control Hamiltonians.
        """
        return len(self.Hs)
    @property
    def evolver(self) -> UnitaryEvolver:
        """
        The integrator used for time evolutions of the system.

        Note
        ----
        The `evolver` can take a while to initialise and so is not initialised
        until `evolver` is is first used or when :meth:`initialise_evolver()` is
        called. Using `evolver` before calling :meth:`initialise_evolver()`
        initialises the `evolver` with the default parameters of
        :meth:`initialise_evolver()`.
        """
        if self._evolver is None:
            self.initialise_evolver()
        return self._evolver
    def initialise_evolver(self,
                           sparse: bool = False,
                           force_dynamic: bool = False):
        """
        Initialises :attr:`evolver` with an evolver from
        `PySTE <https://PySTE.readthedocs.io>`__.
        `PySTE <https://PySTE.readthedocs.io>`__ is Python
        wrapper around the C++ header-only library
        `Suzuki-Trotter-Evolver <https://Suzuki-Trotter-Evolver.readthedocs.io>`__:
        a fast Schrödinger solver utilising the first-order Suzuki-Trotter
        expansion.

        Warning
        -------
        This can take a very long time to execute, especially for large Hilbert
        space dimensions. If you plan to evolve the same quantum system many
        times we recommended pickling the :attr:`evolver`.

        Parameters
        ----------
        sparse : bool
            Whether to use sparse or dense matrices during integration.
            To make a decision on whether sparse or dense matrices are likely to
            lead to faster integration you can consult the benchmarks at
            https://PySTE.readthedocs.io/en/latest/benchmarks.
        force_dynamic : bool
            Whether to force `PySTE <https://PySTE.readthedocs.io>`__ to use a
            dynamic evolver.
            
            Note
            ----
            `PySTE <https://PySTE.readthedocs.io>`__ has precompiled evolvers
            for specific Hilbert space dimensions and numbers of control
            Hamiltonians. When these cannot be found
            `PySTE <https://PySTE.readthedocs.io>`__ uses less efficient
            evolvers with the Hilbert space dimension and the number of controls
            determined dynamically at runtime.
        """
        self._evolver = get_unitary_evolver(self.H0, self._Hs, sparse, force_dynamic)
    def _H(self, ctrl_amp: np.ndarray[float]) -> np.ndarray[complex]:
        """
        Computes the system Hamiltonian for the specified control amplitudes.

        Parameters
        ----------
        ctrl_amp : NDArray[Shape[s := Any_Shape, :attr:`n_ctrl`], float]
            The control amplitudes (stored in the last axis). The prior axes
            allow for multiple sets of control amplitudes to be passed and the
            Hamiltonian for each computed.

        Returns
        -------
        NDArray[Shape[s, :attr:`dim`, :attr:`dim`], complex]
            The system's Hamiltonian (stored in the last two axes).
        """
        return self._H0 + np.einsum("...i,ijk->...jk", ctrl_amp, self.Hs)
    def H(self,
          ctrl_amp: Union[np.ndarray[float], np.ndarray[Callable[[float], np.ndarray[float]]], Callable[[float], np.ndarray[float]]]
         ) -> Union[np.ndarray[complex], Callable[[float], np.ndarray[complex]]]:
        """
        Computes the system Hamiltonian for the specified control amplitudes.

        Parameters
        ----------
        ctrl_amp : NDArray[Shape[s := Any_Shape, :attr:`n_ctrl`], float | Callable[[float], np.ndarray[float]]] | Callable[[float], NDArray[Shape[:attr:`n_ctrl`], float]
            The control amplitudes (stored in the last axis). The prior axes
            allow for multiple sets of control amplitudes to be passed and the
            Hamiltonian for each computed. The control amplitudes can be passed
            as ``np.ndarray[float]`` to compute the system Hamiltonian for a
            specific value of the control ampltiudes. Alternatively,
            time-dependent control amplitudes can be passed.
            ``np.ndarray[Callable[[float], np.ndarray[float]]]`` can be passed
            where each element is a function of time. Alternatively, a function
            of time  that returns the control amplitudes can be passed as
            ``Callable[[float], NDArray[Shape[:attr:`n_ctrl`], float]``. These
            will generate a time-dependent Hamiltonian: a function that takes a
            single parameter (time) and returns the Hamiltonian at this time.

        Returns
        -------
        NDArray[Shape[s, :attr:`dim`, :attr:`dim`], complex] | NDArray[Shape[s], Callable[[float], np.ndarray[complex]]]]
            Either the systems Hamiltonian stored in the last two axes (if
            specific control amplitudes were passed) or a collection of
            time-dependent Hamiltonians (if time-dependent controls were
            passed).
        """
        if callable(ctrl_amp):
            return lambda t: self._H(ctrl_amp(t))
        ctrl_amp = np.array(ctrl_amp)
        if ctrl_amp.dtype == object:
            return lambda t: self._H([a(t) for a in ctrl_amp])
        return self._H(ctrl_amp)
    def _pre_processing(self, *args):
        """
        When calling any evolution method (listed in the
        :ref:`See also section <pre_processing_see_also>`)
        :meth:`_pre_processing()` is executed on the arguments before the
        control amplitudes are modulated by the frequencies (during
        :meth:`_envolope_processing()`) and then finally the modulated control
        amplitudes are used by the evolution method.

        :meth:`_pre_processing()` should be overridden to produce desired pulse
        shapes. You can either override :meth:`_pre_processing()` directly by
        creating a child class, or you can use :meth:`pulse_form()`.

        For :meth:`gradient()` to function correctly :meth:`_pre_processing()`
        should be written in `TensorFlow <https://www.tensorflow.org>`__.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        tuple[tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128], tf.Tensor[Shape[:attr:`state_shape`], tf.complex128], float, tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128], list[int]]
            A tuple of
            1. The control amplitude envolopes
            2. The initial state
            3. The integrator time step
            4. The frequencies to modulate the control amplitude envolopes with
            5. A list of the number of channels for each control Hamiltonian

            Warning
            -------
            The number of channels for each control Hamiltonian must be stored
            as a ``list`` and not an ``NDArray`` or a
            `TensorFlow <https://www.tensorflow.org>`__ tensor.


        .. _pre_processing_see_also:
        
        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        * :meth:`get_evolution()`
        * :meth:`evolved_expectation_value()`
        * :meth:`evolved_expectation_value_all()`
        * :meth:`evolved_inner_product()`
        * :meth:`evolved_inner_product_all()`
        * :meth:`evolved_gate_infidelity()`
        * :meth:`get_driving_pulses()`
        * :meth:`gradient()`
        * :meth:`gate_gradient()`
        """
        return args
    def _envolope_processing(self,
                             ctrl_amp,
                             dt: float,
                             frequencies,
                             number_channels: list[int]
                            ) -> tuple:
        """
        When calling any evolution method (listed in the
        :ref:`See also section <envolope_processing_see_also>` section)
        :meth:`_pre_processing()` is executed on the arguements before the
        control amplitudes are modulated by the frequencies during
        :meth:`_envolope_processing()` and then finally the modulated control
        amplitudes are used by the evolution method.

        Parameters
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        dt : float
            The itegration time step
        frequencies : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The frequencies to modulate the control amplitudes with
        number_channels : list[int]
            The number of channels associated with each control Hamiltonian

            Warning
            -------
            This must be a ``list`` and not an ``NDArray`` or a
            `TensorFlow <https://www.tensorflow.org>`__ tensor.

        Returns
        -------
        tf.Tensor[Shape[n_time_steps, :attr:`n_ctrl`], tf.complex128]
            The modulated control amplitudes


        .. _envolope_processing_see_also:
        
        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        * :meth:`get_evolution()`
        * :meth:`evolved_expectation_value()`
        * :meth:`evolved_expectation_value_all()`
        * :meth:`evolved_inner_product()`
        * :meth:`evolved_inner_product_all()`
        * :meth:`evolved_gate_infidelity()`
        * :meth:`get_driving_pulses()`
        * :meth:`gradient()`
        * :meth:`gate_gradient()`
        """
        dt = tf.cast(dt, dtype=tf.complex128)
        frequencies = tf.cast(frequencies, dtype=tf.complex128)
        ctrl_amp = tf.cast(ctrl_amp, dtype=tf.complex128)
        channel_couplings = tf.constant(generate_channel_couplings(number_channels), dtype=tf.complex128)
        x = tf.exp(tf.einsum("i,j->ij", dt*tf.cast(tf.range(tf.shape(ctrl_amp)[0]), dtype=tf.complex128), -1j*frequencies))
        ctrl_amp = tf.cast(tf.math.real(tf.einsum("tc,tc,dc->td", x, ctrl_amp, channel_couplings)), dtype=tf.complex128)
        return ctrl_amp
    def propagate(self, *args) -> np.ndarray[complex]:
        """
        Evolves a state vector under the time-dependent Hamiltonian defined by
        the control amplitudes using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.propagate()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        NDArray[Shape[:attr:`state_shape`], complex]
            The final state

        See Also
        --------
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        * :meth:`get_evolution()`
        """
        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args)
        return self.evolver.propagate(ctrl_amp, initial_state, dt)
    def propagate_collection(self, *args) -> np.ndarray[complex]:
        """
        Evolves a collection of state vectors under the time-dependent
        Hamiltonian defined by the control amplitudes using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.propagate_collection()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        NDArray[Shape[n_states, :attr:`state_shape`], complex]
            The final state

        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_all()`
        * :meth:`get_evolution()`
        """
        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args)
        return self.evolver.propagate_collection(ctrl_amp, initial_state, dt)
    def propagate_all(self, *args) -> np.ndarray[complex]:
        """
        Evolves a state vector under the time-dependent Hamiltonian defined by
        the control amplitudes using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.propagate_all()`
        from `PySTE <https://PySTE.readthedocs.io>`__ and returns the state at
        each time-step.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        NDArray[Shape[:attr:`state_shape`, n_time_steps+1], complex]
            The state at each integrator time step (including the initial
            state).

        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        * :meth:`get_evolution()`
        """
        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args)
        return self.evolver.propagate_all(ctrl_amp, initial_state, dt)
    def evolved_expectation_value(self, *args) -> complex:
        """
        Evolves a state vector under the time-dependent Hamiltonian defined by
        the control amplitudes and computes the expectation value of a specified
        observable with respect to the final state using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.evolved_expectation_value()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.
        ``args[-1]`` : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The observable to take the expectation value of.


        Returns
        -------
        complex
            The expectation value.

        See Also
        --------
        * :meth:`evolved_expectation_value_all()`
        * :meth:`gradient()`
        """
        observable: np.ndarray[complex] = args[-1]
        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args[:-1])
        return self.evolver.evolved_expectation_value(ctrl_amp,
                                                      initial_state,
                                                      dt,
                                                      observable)
    def evolved_expectation_value_all(self, *args) -> np.ndarray[complex]:
        """
        Evolves a state vector under the time-dependent Hamiltonian defined by
        the control amplitudes and computes the expectation value of a specified
        observable with respect to the state at each time-step using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.evolved_expectation_value_all()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.
        ``args[-1]`` : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The observable to take the expectation value of.

        Returns
        -------
        NDArray[Shape[n_time_steps+1], complex]
            The expectation value of the state at each integrator time step
            (including the initial state) with respect to the observable.

        See Also
        --------
        * :meth:`evolved_expectation_value()`
        * :meth:`gradient()`
        """
        observable: np.ndarray[complex] = args[-1]
        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args[:-1])
        return self.evolver.evolved_expectation_value_all(ctrl_amp,
                                                          initial_state,
                                                          dt,
                                                          observable)
    def evolved_inner_product(self, *args) -> complex:
        """
        Evolves a state vector under the time-dependent Hamiltonian defined by
        the control amplitudes and computes the inner product of the final state
        vector with a fixed vector using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.evolved_inner_product()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.
        ``args[-1]`` : NDArray[Shape[:attr:`dim`], complex]
            The fixed vector to calculate the inner product with.

        Returns
        -------
        complex
            The inner product of the evolved state vector with the fixed vector.

        See Also
        --------
        * :meth:`evolved_inner_product_all()`
        """
        fixed_vector: np.ndarray[complex] = args[-1]
        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args[:-1])
        return self.evolver.evolved_inner_product(ctrl_amp,
                                                  initial_state,
                                                  dt,
                                                  fixed_vector)
    def evolved_inner_product_all(self, *args) -> complex:
        """
        Evolves a state vector under the time-dependent Hamiltonian defined by
        the control amplitudes and computes the innper product of state at each
        time-step with a fixed vector using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.evolved_inner_product_all()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.
        ``args[-1]`` : NDArray[Shape[:attr:`dim`], complex]
            The fixed vector to calculate the inner product with.

        Returns
        -------
        NDArray[Shape[n_time_steps+1], complex]
            The inner product of state at each integrator time step (including
            the initial state) with the fixed vector.

        See Also
        --------
        * :meth:`evolved_inner_product()`
        """
        fixed_vector: np.ndarray[complex] = args[-1]
        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args[:-1])
        return self.evolver.evolved_inner_product_all(ctrl_amp,
                                                      initial_state,
                                                      dt,
                                                      fixed_vector)
    def get_evolution(self, *args) -> np.ndarray[complex]:
        """
        Computes the unitary evolution of the system under the time-dependent
        Hamiltonian defined by the control amplitudes using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.get_evolution()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The unitary corresponding to the evolution of the system.

        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        """
        ctrl_amp, _, dt = self.get_driving_pulses(*args)
        return self.evolver.get_evolution(ctrl_amp, dt)
    def evolved_gate_infidelity(self, *args) -> float:
        """
        Evolves the system under the time-dependent Hamiltonian defined by
        the control amplitudes and computes the gate infidelity to the target
        gate using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.evolved_gate_infidelity()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.
        ``args[-1]`` : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The target gate to calculate the gate infidelity with respect to.

        Returns
        -------
        float
            The gate infidelity

        See Also
        --------
        * :meth:`gate_gradient()`
        """
        target: np.ndarray[complex] = args[-1]
        ctrl_amp, _, dt = self.get_driving_pulses(*args[:-1])
        return self.evolver.evolved_gate_infidelity(ctrl_amp, dt, target)
    def get_driving_pulses(self, *args) -> tuple[np.ndarray[complex], np.ndarray[complex], float]:
        """
        When calling any evolution method (listed in the
        :ref:`See also section <get_driving_pulses_see_also>`) :meth:`get_driving_pulses()`
        is executed on the arguements before the evolution method.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        tuple[NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex], NDArray[Shape[:attr:`state_shape`], complex], float]
            A tuple of:
            1. Control amplitudes
            2. Initial state
            3. Integrator time step


        .. _get_driving_pulses_see_also:
        
        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        * :meth:`get_evolution()`
        * :meth:`evolved_expectation_value()`
        * :meth:`evolved_expectation_value_all()`
        * :meth:`evolved_inner_product()`
        * :meth:`evolved_inner_product_all()`
        * :meth:`evolved_gate_infidelity()`
        * :meth:`gradient()`
        * :meth:`gate_gradient()`
        """
        ctrl_amp, initial_state, dt = self._processing(*args)
        try: ctrl_amp: np.ndarray[complex] = ctrl_amp.numpy()
        except: pass
        try:  initial_state: np.ndarray[complex] = initial_state.numpy().flatten()
        except: pass
        try: dt = dt.numpy()
        except: pass
        dt = float(dt.real)
        return ctrl_amp, initial_state, dt
    def _eager_processing(self, *args) -> tuple:
        """
        Executes :meth:`_pre_processing()` followed by
        :meth:`_envolope_processing()` eagerly (i.e. without using a
        `TensorFlow <https://www.tensorflow.org>`__ graph). Nonetheless,
        :meth:`_eager_processing()` is still auto differentiable.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        tuple[tf.Tensor[Shape[n_time_steps, :attr:`n_ctrl`], complex], tf.Tensor[Shape[:attr:`state_shape`], complex], tf.Tensor[Shape[], float]]
            A tuple of:
            1. Control amplitudes
            2. Initial state
            3. Integrator time step
        """
        ctrl_amp, initial_state, dt, frequencies, number_channels = self._pre_processing(*args)
        return self._envolope_processing(ctrl_amp, dt, frequencies, list(number_channels)), initial_state, dt
    def _traceable_eager_processing(self, *args) -> tuple:
        """
        A function that will be traced by
        `TensorFlow <https://www.tensorflow.org>`__ to produce a graph of
        :meth:`_pre_processing()` followed by :meth:`_envolope_processing()`.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        tuple[tf.Tensor[Shape[n_time_steps, :attr:`n_ctrl`], complex], tf.Tensor[Shape[:attr:`state_shape`], complex], tf.Tensor[Shape[], float]]
            A tuple of:
            1. Control amplitudes
            2. Initial state
            3. Integrator time step
        """
        print("Tracing control amplitude graph.")
        return self._eager_processing(*args)
    def gradient(self, *args) -> tuple[float, np.ndarray[float]]:
        """
        Evolves a state vector under the time-dependent Hamiltonian defined by
        the control amplitudes and computes the expectation value of a specified
        observable with respect to the final state and then computes the
        gradient of the final state with respect to the first argument
        (``args[0]``) using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.switching_function()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.
        ``args[-1]`` : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The observable to take the expectation value of.

        Returns
        -------
        tuple[complex, NDArray[Shape[n_parameters], float]]
            A tuple of the expectation value and the gradient.

        See Also
        --------
        * :meth:`evolved_expectation_value()`
        * :meth:`evolved_expectation_value_all()`
        * :meth:`gate_gradient()`
        """
        cost: np.ndarray[complex] = args[-1]
        args = list(args[:-1])
        args[0] = tf.constant(args[0], dtype=tf.float64)             
        with tf.GradientTape(persistent=False,
                             watch_accessed_variables=False
                            ) as tape:
            tape.watch(args[0])
            ctrl_amp, initial_state, dt = self._processing(*args)
            try: initial_state: np.ndarray[complex] = initial_state.numpy().flatten()
            except: pass
            try: dt = dt.numpy()
            except: pass
            dt = float(dt.real)
            E = ExpValCustom(self, initial_state, dt, cost).run(ctrl_amp)
        grad = tape.gradient(E, args[0])
        del tape
        
        grad = tf.convert_to_tensor(grad)
        grad: np.ndarray[float] = grad.numpy()
        grad = grad.real
        return E.numpy(), grad
    def gate_gradient(self, *args) -> tuple[float, np.ndarray[float]]:
        """
        Evolves the system under the time-dependent Hamiltonian defined by
        the control amplitudes and computes the gate infidelity to the target
        gate and then computes the gradient with respect to the first argument
        (``args[0]``) using
        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.gate_switching_function()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.
        ``args[-1]`` : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The target gate to calculate the gate infidelity with respect to.

        Returns
        -------
        tuple[complex, NDArray[Shape[n_parameters], float]]
            A tuple of the gate infidelity value and the gradient.

        See Also
        --------
        * :meth:`evolved_gate_infidelity()`
        * :meth:`gradient()`
        """
        target: np.ndarray[complex] = args[-1]
        args = list(args[:-1])
        args[0] = tf.constant(args[0], dtype=tf.float64)             
        with tf.GradientTape(persistent=False,
                             watch_accessed_variables=False
                            ) as tape:
            tape.watch(args[0])
            ctrl_amp, _, dt = self._processing(*args)
            try: dt: float = dt.numpy()
            except: pass
            dt = float(dt.real)
            infidelity = GateInfidelityCustom(self, dt, target).run(ctrl_amp)
        grad = tape.gradient(infidelity, args[0])
        del tape
        
        grad = tf.convert_to_tensor(grad)
        grad: np.ndarray[float] = grad.numpy()
        grad = grad.real
        return infidelity.numpy(), grad
    def pulse_form(self,
                   pulse_function: Callable,
                   append: bool = False,                  
                  ) -> "PulseForm":
        """
        Initialises a new :class:`QuantumSystem` in which
        :meth:`_pre_processing()` corresponds to executing ``pulse_function()``
        and piping the output into the previous definition of
        :meth:`_pre_processing()`.

        Parameters
        ----------
        pulse_function : Callable
            The function to compose with :meth:`_pre_processing()`.
        append : bool, optional
            Whether to prepend
            (``self._pre_processing(*pulse_function())``) or append
            (``pulse_function(*self._pre_processing())``)
            ``pulse_function``. By default ``False``.

        Returns
        -------
        PulseForm
            The new :class:`QuantumSystem`
        """
        return PulseForm(self, pulse_function, append)

class TransformedSystem(QuantumSystem):
    """
    A base class for representing a transformation on a :class:`qugrad.QuantumSystem`.
    """

    _original_system: QuantumSystem
    "The system that was transformed into this system"
    
    _base_system: QuantumSystem
    """
    The system before any transformations were applied. That is `_base_system`
    is the recursive :attr:`original_system`
    (``original_system.original_system.original_system....``) until
    :attr:`original_system` is no longer a :class:`TransformedSystem`.
    """
    
    def __init__(self,
                 original_system: QuantumSystem,
                 H0: np.ndarray[complex],
                 Hs: Union[np.ndarray[complex], np.ndarray[complex]],
                 hilbert_space: HilbertSpace):
        """
        Performs a transformation on a :class:`qugrad.QuantumSystem`.

        Parameters
        ----------
        original_system: QuantumSystem
            The system to be transformed into this system
        H0: NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The new drift Hamiltonian
        Hs: NDArray[Shape[":attr:`n_ctrl`, :attr:`dim`, :attr:`dim`"], complex] | NDArray[Shape[:attr:`n_ctrl` * :attr:`dim`, :attr:`dim`], complex]
            The new control Hamiltonians either as an array of control
            Hamiltonians or the control Hamiltonians stacked along the first
            axis.
        hilbert_space: HilbertSpace
            The new Hilbert space of the system
        """
        self._original_system = original_system
        if isinstance(original_system, TransformedSystem):
            self._base_system = original_system._base_system
        else:
            self._base_system = original_system
        super().__init__(H0, Hs, hilbert_space, self._base_system.using_graph)
    @property
    def original_system(self) -> QuantumSystem:
        "The system that was transformed into this system"
        return self._original_system
    @property
    def base_system(self) -> QuantumSystem:
        """
        The system before any transformations were applied. That is
        :attr:`base_system` is the recursive :attr:`original_system`
        (``original_system.original_system.original_system....``) until
        :attr:`original_system` is no longer a :class:`TransformedSystem`.
        """
        return self._base_system
    def _pre_processing(self, *args) -> tuple:
        """
        When calling any evolution method (listed in the
        :ref:`See also section <TransformedSystem_pre_processing_see_also>` section)
        :meth:`_pre_processing()` is executed on the arguements before the
        control amplitudes are modulated by the frequencies (during
        :meth:`_envolope_processing()`) and then finally the modulated control
        amplitudes are used by the evolution method.

        This is a placeholder for ``original_system._pre_processing()``.

        Parameters
        ----------
        *args
            The placeholder parameters. See ``_systems.pyi`` for actual
            parameters. Each child class that implements a new
            :meth:`_pre_processing()` should implement a ``.pyi`` file to
            document the parameters for this function: the same parameters as
            passed to :meth:`_pre_processing()`.

        Returns
        -------
        tuple[tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128], tf.Tensor[Shape[:attr:`state_shape`], tf.complex128], float, tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128], list[int]]
            A tuple of
            1. The control amplitude envolopes
            2. The initial state
            3. The integrator time step
            4. The frequencies to modulate the control amplitude envolopes with
            5. A list of the number of channels for each control Hamiltonian

            Warning
            -------
            The number of channels for each control Hamiltonian must be stored
            as a ``list`` and not an ``NDArray`` or a
            `TensorFlow <https://www.tensorflow.org>`__ tensor.


        .. _TransformedSystem_pre_processing_see_also:
        
        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        * :meth:`evolved_expectation_value()`
        * :meth:`evolved_expectation_value_all()`
        * :meth:`get_driving_pulses()`
        * :meth:`gradient()`
        """
        return self._original_system._pre_processing(*args)
    def _envolope_processing(self,
                             ctrl_amp,
                             dt: float,
                             frequencies,
                             number_channels: list[int]):
        """
        When calling any evolution method (listed in the
        :ref:`See also section <envolope_processing_see_also>` section) :meth:`_pre_processing()`
        is executed on the arguements before the control amplitudes are
        modulated by the frequencies during :meth:`_envolope_processing()` and
        then finally the modulated control amplitudes are used by the evolution
        method.

        Parameters
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        dt : float
            The itegration time step
        frequencies : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The frequencies to modulate the control amplitudes with
        number_channels : list[int]
            The number of channels associated with each control Hamiltonian

            Warning
            -------
            This must be a ``list`` and not an ``NDArray`` or a
            `TensorFlow <https://www.tensorflow.org>`__ tensor.

        Returns
        -------
        tf.Tensor[Shape[n_time_steps,:attr:`n_ctrl`], tf.complex128]
            The modulated control amplitudes


        .. _envolope_processing_see_also:
        
        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        * :meth:`evolved_expectation_value()`
        * :meth:`evolved_expectation_value_all()`
        * :meth:`get_driving_pulses()`
        * :meth:`gradient()`
        """
        return self._original_system._envolope_processing(ctrl_amp,
                                                          dt,
                                                          frequencies,
                                                          number_channels)

class PulseForm(TransformedSystem):
    """
    A transformed :class:`qugrad.QuantumSystem` in which :meth:`_pre_processing()`
    has been composed with another pre processing function.
    """

    _pulse_function: Callable
    "The function composed with ``original_system._pre_processing()``"
    
    _appended: bool
    """
    Whether the :attr:`pulse_function` was prepended
    (``original_system._pre_processing(*pulse_function())``) or appended
    (``pulse_function(*original_system._pre_processing())``)
    """
        
    def __init__(self,
                 original_system: QuantumSystem,
                 pulse_function: Callable,
                 append: bool = False):
        """
        Initialises a new :class:`qugrad.QuantumSystem` in which :meth:`_pre_processing()`
        corresponds to running ``pulse_function()`` and piping the output into
        ``original_system._pre_processing()``.

        Parameters
        ----------
        original_system : QuantumSystem
            The system that was transformed into this system
        pulse_function : Callable
            The function to compose with :meth:`_pre_processing()`.
        append : bool
            Whether to prepend
            (``original_system._pre_processing(*pulse_function())``) or append
            (``pulse_function(*original_system._pre_processing())``)
            ``pulse_function``. By default ``False``.
        """
        super().__init__(original_system,
                         original_system.H0,
                         original_system.Hs,
                         original_system.hilbert_space)
        self._evolver = original_system._evolver
        if append:
            self._pre_processing = compose_unpack(
                pulse_function,
                original_system._pre_processing
            )
        else:
            self._pre_processing = compose_unpack(
                original_system._pre_processing,
                pulse_function
            )
        self._pulse_function = pulse_function
        self._appended = append
    @property
    def pulse_function(self) -> Callable:
        "The function composed with ``original_system._pre_processing()``"
        return self._pulse_function
    @property
    def appended(self) -> bool:
        """
        Whether the :attr:`pulse_function` was prepended
        (``original_system._pre_processing(*pulse_function())``) or appended
        (``pulse_function(*original_system._pre_processing())``)
        """
        return self._appended