Source code for qugrad.systems._systems

"""
Defines classes for quantum systems
"""

from typing import Callable, Optional, Union, Any

import numpy as np
import tensorflow as tf

from py_ste.evolvers import UnitaryEvolver

from .._hilbert_space import HilbertSpace

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.
    """
    ...

# 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
        """
        ...
    @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>`__.
        """
        ...

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
        """
        ...
    @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>`__.
        """
        ...



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
    ----------
    ctrl_amp : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
        The envolope control amplitudes
    initial_state : NDArray[Shape[:attr:`state_shape`], complex]
        The initial state for the integrator
    dt : float
        The itegration time step
    frequencies : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
        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
    -------
    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 = True
            Whether to use TensorFlow graphs during computation.
        """
        ...

            
    def __del__(self):
        ...
    @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.
        """
        ...
    @using_graph.setter
    def using_graph(self, value: bool):
        ...
    @property
    def hilbert_space(self) -> HilbertSpace:
        "The Hilbert space of the system"
        ...
    @property
    def H0(self) -> np.ndarray[complex]:
        """
        The systems drift Hamiltonian as a :attr:`dim` x :attr:`dim` matrix.

        See Also
        --------
        :attr:`Hs`
        """
        ...
    @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`
        """
        ...
    @property
    def dim(self) -> int:
        """
        The dimension of states in the quantum system.

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

        See Also
        --------
        :attr:`dim`
        """
        ...
    @property
    def n_ctrl(self) -> int:
        """
        The number of control Hamiltonians.
        """
        ...
    @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()`.
        """
        ...


    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.
        """
        ...

    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).
        """
        ...


    def H(self,
          ctrl_amp: Union[np.ndarray[float], np.ndarray[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]]]
            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, the control
            amplitudes can be passed as
            ``np.ndarray[Callable[[float], np.ndarray[float]]]`` where each
            element is a function of time. This 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).
        """
        ...



    def _pre_processing(self,
                        ctrl_amp : np.ndarray[complex],
                        initial_state : np.ndarray[complex],
                        dt: float,
                        frequencies: np.ndarray[complex],
                        number_channels: list[int]
                       ) -> tuple:
        """
        When calling any evolution method (listed in the
        :ref:`See also section <pre_processing_see_also>`) :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.

        :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
        ----------
        ctrl_amp : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        dt : float
            The itegration time step
        frequencies : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            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
        -------
        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.

        Warning
        -------
        Keyword arguments are not supported.

        .. _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()`
        """
        ...



    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()`
        """
        ...



    def propagate(self,
                  ctrl_amp : np.ndarray[complex],
                  initial_state : np.ndarray[complex],
                  dt: float,
                  frequencies: np.ndarray[complex],
                  number_channels: list[int]
                 ) -> 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
        ----------
        ctrl_amp : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        dt : float
            The itegration time step
        frequencies : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            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.

        Warning
        -------
        Keyword arguments are not supported.

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

        See Also
        --------
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        * :meth:`get_evolution()`
        """
        ...



    def propagate_collection(self,
                             ctrl_amp : np.ndarray[complex],
                             initial_states : np.ndarray[complex],
                             dt: float,
                             frequencies: np.ndarray[complex],
                             number_channels: list[int]
                            ) -> 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
        ----------
        ctrl_amp : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            The envolope control amplitudes
        initial_states : NDArray[Shape[n_states, :attr:`state_shape`], complex]
            The initial state for the integrator
        dt : float
            The itegration time step
        frequencies : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            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.

        Warning
        -------
        Keyword arguments are not supported.

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

        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_all()`
        * :meth:`get_evolution()`
        """
        ...



    def propagate_all(self,
                      ctrl_amp : np.ndarray[complex],
                      initial_states : np.ndarray[complex],
                      dt: float,
                      frequencies: np.ndarray[complex],
                      number_channels: list[int]
                     ) -> 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
        ----------
        ctrl_amp : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        dt : float
            The itegration time step
        frequencies : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            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.

        Warning
        -------
        Keyword arguments are not supported.

        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()`
        """
        ...



    def evolved_expectation_value(self,
                                  ctrl_amp : np.ndarray[complex],
                                  initial_state : np.ndarray[complex],
                                  dt: float,
                                  frequencies: np.ndarray[complex],
                                  number_channels: list[int],
                                  observable : np.ndarray[complex]
                                 ) -> 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
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        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.
        observable : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The observable to take the expectation value of.

        Warning
        -------
        Keyword arguments are not supported.

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

        See Also
        --------
        * :meth:`evolved_expectation_value_all()`
        * :meth:`gradient()`
        """
        ...



    def evolved_expectation_value_all(self,
                                      ctrl_amp : np.ndarray[complex],
                                      initial_state : np.ndarray[complex],
                                      dt: float,
                                      frequencies: np.ndarray[complex],
                                      number_channels: list[int],
                                      observable : np.ndarray[complex]
                                     ) -> 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
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        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.
        observable : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The observable to take the expectation value of.

        Warning
        -------
        Keyword arguments are not supported.

        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()`
        """
        ...



    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
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        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.
        fixed_vector : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The fixed vector to calculate the inner product with.

        Warning
        -------
        Keyword arguments are not supported.

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

        See Also
        --------
        * :meth:`evolved_inner_product_all()`
        """
        ...



    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
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        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.
        fixed_vector : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The fixed vector to calculate the inner product with.

        Warning
        -------
        Keyword arguments are not supported.

        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()`
        """
        ...



    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
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            This is a redundant parameter for this function.
        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.

        Warning
        -------
        Keyword arguments are not supported.

        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()`
        """
        ...



    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
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            This is a redundant parameter for this function.
        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.
        target : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The target gate to calculate the gate infidelity with respect to.

        Warning
        -------
        Keyword arguments are not supported.
        
        Returns
        -------
        float
            The gate infidelity

        See Also
        --------
        * :meth:`gate_gradient()`
        """
        ...



    def get_driving_pulses(self,
                           ctrl_amp : np.ndarray[complex],
                           initial_states : np.ndarray[complex],
                           dt: float,
                           frequencies: np.ndarray[complex],
                           number_channels: list[int]
                          ) -> 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
        ----------
        ctrl_amp : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        dt : float
            The itegration time step
        frequencies : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            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.

        Warning
        -------
        Keyword arguments are not supported.

        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()`
        """
        ...

    def _eager_processing(self,
                          ctrl_amp : np.ndarray[complex],
                          initial_states : np.ndarray[complex],
                          dt: float,
                          frequencies: np.ndarray[complex],
                          number_channels: list[int]
                         ) -> 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
        ----------
        ctrl_amp : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        dt : float
            The itegration time step
        frequencies : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            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.

        Warning
        -------
        Keyword arguments are not supported.

        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
        """
        ...
    def _traceable_eager_processing(self,
                                   ctrl_amp : np.ndarray[complex],
                                   initial_states : np.ndarray[complex],
                                   dt: float,
                                   frequencies: np.ndarray[complex],
                                   number_channels: list[int]
                                  ) -> 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
        ----------
        ctrl_amp : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        dt : float
            The itegration time step
        frequencies : NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex]
            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.

        Warning
        -------
        Keyword arguments are not supported.

        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
        """
        ...


    def gradient(self,
                 ctrl_amp : np.ndarray[complex],
                 initial_state : np.ndarray[complex],
                 dt: float,
                 frequencies: np.ndarray[complex],
                 number_channels: list[int],
                 observable : np.ndarray[complex]
                ) -> 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
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            The initial state for the integrator
        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.
        observable : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The observable to take the expectation value of.

        Warning
        -------
        Keyword arguments are not supported.

        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()`
        """
        ...



    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
        ----------
        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
            The envolope control amplitudes
        initial_state : NDArray[Shape[:attr:`state_shape`], complex]
            This is a redundant parameter for this function.
        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.
        observable : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
            The observable to take the expectation value of.

        Warning
        -------
        Keyword arguments are not supported.
        
        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()`
        """
        ...



    def pulse_form(self,
                   pulse_function: Callable,
                   path: Optional[str] = None
                  ) -> "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`
        """
        ...





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
        """
        ...

    @property
    def original_system(self) -> QuantumSystem:
        "The system that was transformed into this 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`.
        """
        ...


    def _pre_processing(self, *args):
        """
        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. The actual parameters will be the same
            as ``original_system._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()`
        """
        ...



    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. The actual parameters will be the same
            as ``original_system._pre_processing()``.

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

        See Also
        --------
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        """
        ...



    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. The actual parameters will be the same
            as ``original_system._pre_processing()``.

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

        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_all()`
        """
        ...



    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. The actual parameters will be the same
            as ``original_system._pre_processing()``.

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

        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        """
        ...



    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. The actual parameters will be the same
            as ``original_system._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()`
        """
        ...



    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. The actual parameters will be the same
            as ``original_system._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 state at each integrator time step (including the initial
            state).

         See Also
        --------
        * :meth:`evolved_expectation_value()`
        * :meth:`gradient()`
        """
        ...



    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. The actual parameters will be the same
            as ``original_system._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:`evolved_expectation_value()`
        * :meth:`evolved_expectation_value_all()`
        * :meth:`gradient()`
        """
        ...



    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.DenseUnitaryEvoler.switching_function()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. The actual parameters will be the same
            as ``original_system._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()`
        """
        ...





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


    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``.
        """
        ...



    def _pre_processing(self, *args):
        """
        When calling any evolution method (listed in the
        :ref:`See also section <PulseForm_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(*pulse_function())``
        (or ``pulse_function(*original_system._pre_processing())`` if
        attr:`appended` is ``True``).

        Parameters
        ----------
        *args
            The placeholder parameters. The actual parameters will be the same
            as :attr:`pulse_function` if ``append`` is ``False`` or
            ``original_system._pre_processing()`` if ``append`` is ``True``.

        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.


        .. _PulseForm_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()`
        """
        ...



    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. The actual parameters will be the same
            as :attr:`pulse_function` if ``append`` is ``False`` or
            ``original_system._pre_processing()`` if ``append`` is ``True``.

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

        See Also
        --------
        * :meth:`propagate_collection()`
        * :meth:`propagate_all()`
        """
        ...



    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. The actual parameters will be the same
            as :attr:`pulse_function` if ``append`` is ``False`` or
            ``original_system._pre_processing()`` if ``append`` is ``True``.

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

        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_all()`
        """
        ...



    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. The actual parameters will be the same
            as :attr:`pulse_function` if ``append`` is ``False`` or
            ``original_system._pre_processing()`` if ``append`` is ``True``.

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

        See Also
        --------
        * :meth:`propagate()`
        * :meth:`propagate_collection()`
        """
        ...



    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. The actual parameters will be the same
            as :attr:`pulse_function` if ``append`` is ``False`` or
            ``original_system._pre_processing()`` if ``append`` is ``True``.
        ``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()`
        """
        ...



    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. The actual parameters will be the same
            as :attr:`pulse_function` if ``append`` is ``False`` or
            ``original_system._pre_processing()`` if ``append`` is ``True``.
        ``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 state at each integrator time step (including the initial
            state).

         See Also
        --------
        * :meth:`evolved_expectation_value()`
        * :meth:`gradient()`
        """
        ...



    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. The actual parameters will be the same
            as :attr:`pulse_function` if ``append`` is ``False`` or
            ``original_system._pre_processing()`` if ``append`` is ``True``.

        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:`evolved_expectation_value()`
        * :meth:`evolved_expectation_value_all()`
        * :meth:`gradient()`
        """
        ...



    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.DenseUnitaryEvoler.switching_function()`
        from `PySTE <https://PySTE.readthedocs.io>`__.

        Parameters
        ----------
        ``*args[:-1]``
            The placeholder parameters. The actual parameters will be the same
            as :attr:`pulse_function` if ``append`` is ``False`` or
            ``original_system._pre_processing()`` if ``append`` is ``True``.
        ``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()`
        """
        ...