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.

   1"""
   2Defines classes for quantum systems
   3"""
   4
   5from typing import Callable, Optional, Union
   6
   7import numpy as np
   8import tensorflow as tf
   9
  10from py_ste import get_unitary_evolver
  11from py_ste.evolvers import UnitaryEvolver
  12
  13from .._hilbert_space import HilbertSpace
  14from ..pulses import compose_unpack
  15
  16def generate_channel_couplings(number_channels: list[int]) -> np.ndarray[bool]:
  17    """Generates a boolean array indicating which channels couple to which
  18    drives.
  19
  20    Parameters
  21    ----------
  22    number_chanels : list[int]
  23        A list with the nth entry indicating the number of channels that
  24        correspond the nth drive.
  25
  26    Returns
  27    -------
  28    NDArray[Shape[``len(number_channels)``, total_number_channels], bool]
  29        An entry is `True` if the channel corresponds to the drive.
  30    """
  31    number_channels: np.ndarray[int] = np.array(number_channels)
  32    if len(number_channels) == 0:
  33        return np.empty((0, 0), dtype=bool)
  34    number_channels = number_channels.cumsum()
  35    channel_couplings: np.ndarray[bool] = np.zeros((len(number_channels),
  36                                                    number_channels[-1]),
  37                                                   dtype=bool)
  38    c_old = 0
  39    for i, c in enumerate(number_channels):
  40        channel_couplings[i, c_old: (c_old := c)] = True
  41    return channel_couplings
  42
  43# Defining classes
  44class ExpValCustom:
  45    """
  46    A class implimenting :meth:`QuantumSystem.evolved_expectation_value()` with
  47    a `TensorFlow <https://www.tensorflow.org>`__ gradient.
  48    """
  49
  50    system: "QuantumSystem"
  51    "The system in which to take the expectation value"
  52    
  53    initial_state: np.ndarray[complex]
  54    "The initial state for the integrator"
  55    
  56    dt: float
  57    "The integrator time step"
  58    
  59    observable: np.ndarray[complex]
  60    "The observable to take the expectation value of"
  61    
  62    def __init__(self,
  63                 system: "QuantumSystem",
  64                 initial_state: np.ndarray[complex],
  65                 dt: float,
  66                 observable: np.ndarray[complex]):
  67        """
  68        Initialises the class with the `system` in which to take the expectation
  69        value and the `observable` in the `system` to take the expectation value
  70        of. Additionally, the initial state before evolution and
  71        the integrator time step are specified.
  72
  73        Parameters
  74        ----------
  75        system : QuantumSystem
  76            The system in which to take the expectation value
  77        initial_state : NDArray[Shape[``system.state_shape``], complex]
  78            The initial state for the integrator
  79        dt : float
  80            The integrator time step
  81        observable : NDArray[Shape[``system.dim``, ``system.dim``], complex])
  82            The observable to take the expectation value of
  83        """
  84        self.system = system
  85        self.initial_state = initial_state
  86        self.dt = dt
  87        self.observable = observable
  88    @tf.custom_gradient
  89    def run(self, ctrl_amp):
  90        """
  91        Computes the expectation value of the :attr:`observable` in the
  92        :attr:`system` with respect to the :attr:`initial_state` evolved under
  93        the Hamiltonian generated by the specified control amplitudes.
  94
  95        Parameters
  96        ----------
  97        ctrl_amp : tf.Tensor[Shape[n_time_steps, ``system.n_ctrl``], tf.complex128]
  98            The control amplitudes
  99
 100        Returns
 101        -------
 102        tf.Tensor[Shape[], tf.complex128]
 103            The expectation value
 104
 105        Note
 106        ----
 107        This function is differentiable using TensorFlow's ``tf.GradientTape``.
 108        """
 109        E, self.switching_function = self.system.evolver.switching_function(ctrl_amp.numpy().copy(), self.initial_state, self.dt, self.observable)
 110        self.switching_function = tf.constant(self.switching_function*self.dt, dtype=tf.complex128)
 111        def grad(upstream) -> tf.Tensor: return self.switching_function
 112        return tf.math.real(tf.constant(E, dtype=tf.complex128)), grad # @tf.custom_gradient removes the second returned variable
 113
 114class QuantumSystem:
 115    """
 116    A class storing the properties of a quantum system.
 117    """
 118    _evolver: Optional[UnitaryEvolver] = None
 119    "The integrator used for time evolutions of the system."
 120    
 121    _hilbert_space: HilbertSpace
 122    "The Hilbert space of the system"
 123    
 124    _H0: np.ndarray
 125    "The systems drift Hamiltonian as a :attr:`dim` x :attr:`dim` matrix."
 126    
 127    _Hs: np.ndarray
 128    """
 129    An array of the system's control Hamiltonians with shape
 130    (:attr:`n_ctrl`, :attr:`dim`, :attr:`dim`).
 131    """
 132    
 133    _graph_processing: Callable[..., tuple]
 134    "A Tensorflow graph of :attr:`_processing()`"
 135    
 136    _processing: Callable[..., tuple]
 137    """
 138    Executes :meth:`_pre_processing()` followed by
 139    :meth:`_envolope_processing()` eagerly (i.e. without using a TensorFlow
 140    graph). Nonetheless, :meth:`_eager_processing()` is still auto
 141    differentiable.
 142
 143    Parameters
 144    ----------
 145    *args
 146        The placeholder parameters. See ``_systems.pyi`` for actual parameters.
 147        Each child class that implements a new :meth:`_pre_processing()` should
 148        implement a ``.pyi`` file to document the parameters for this function:
 149        the same parameters as passed to :meth:`_pre_processing()`.
 150
 151    Returns
 152    -------
 153    tuple[tf.Tensor[Shape[n_time_steps, :attr:`n_ctrl`], complex], tf.Tensor[Shape[:attr:`state_shape`], complex], tf.Tensor[Shape[], float]]
 154        A tuple of:
 155        1. Control amplitudes
 156        2. Initial state
 157        3. Integrator time step
 158    """
 159    
 160    _using_graph: bool
 161    """
 162    Whether to use TensorFlow graphs during computation. Using a TensorFlow
 163    graph will increase the speed of computation. However, you have to be
 164    careful that function parameters have not been baked into the graph leading
 165    to unexpected behaviour.
 166    """
 167    
 168    def __init__(self,
 169                 H0: np.ndarray[complex],
 170                 Hs: np.ndarray[complex],
 171                 hilbert_space: HilbertSpace,
 172                 use_graph: bool = True):
 173        """
 174        Initialises a new :class:`QuantumSystem`.
 175
 176        Parameters
 177        ----------
 178        H0 : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
 179            The systems drift Hamiltonian
 180        Hs : NDArray[Shape[:attr:`n_ctrl`, :attr:`dim`, :attr:`dim`], complex] | NDArray[Shape[:attr:`n_ctrl` * :attr:`dim`, :attr:`dim`], complex]
 181            The systems control Hamiltonians either as an array of control
 182            Hamiltonians or the control Hamiltonians stacked along the first
 183            axis.
 184        hilbert_space : HilbertSpace
 185            The Hilbert space of the system
 186        use_graph : bool
 187            Whether to use `TensorFlow <https://www.tensorflow.org>`__ graphs
 188            during computation, by default ``True``
 189        """
 190        self._hilbert_space: HilbertSpace = hilbert_space
 191        self._H0 = np.array(H0)
 192        self._H0.flags.writeable = False
 193        Hs = np.array(Hs)
 194        self._Hs = Hs if Hs.ndim == 2 else Hs.reshape(-1, self.dim)
 195        self._Hs.flags.writeable = False
 196        self._graph_processing = tf.function(self._traceable_eager_processing, autograph=False)
 197        self.using_graph = use_graph
 198            
 199    def __del__(self):
 200        # to force clear up of tracing
 201        del self._graph_processing
 202        del self._processing
 203    @property
 204    def using_graph(self) -> bool:
 205        """
 206        Whether to use `TensorFlow <https://www.tensorflow.org>`__ graphs during
 207        computation. Using a `TensorFlow <https://www.tensorflow.org>`__ graph
 208        will increase the speed of computation. However, you have to be careful
 209        that function parameters have not been baked into the graph leading to
 210        unexpected behaviour.
 211        """
 212        return self._using_graph
 213    @using_graph.setter
 214    def using_graph(self, value: bool):
 215        self._using_graph = value
 216        self._processing = self._graph_processing if value else self._eager_processing
 217    @property
 218    def hilbert_space(self) -> HilbertSpace:
 219        "The Hilbert space of the system"
 220        return self._hilbert_space
 221    @property
 222    def H0(self) -> np.ndarray[complex]:
 223        """
 224        The systems drift Hamiltonian as a :attr:`dim` x :attr:`dim` matrix.
 225
 226        See Also
 227        --------
 228        :attr:`Hs`
 229        """
 230        return self._H0
 231    @property
 232    def Hs(self) -> np.ndarray[complex]:
 233        """
 234        An array of the system's control Hamiltonians with shape
 235        (:attr:`n_ctrl`, :attr:`dim`, :attr:`dim`).
 236
 237        See Also
 238        --------
 239        :attr:`H0`
 240        """
 241        return self._Hs.reshape((-1, self.dim, self.dim))
 242    @property
 243    def dim(self) -> int:
 244        """
 245        The dimension of states in the quantum system.
 246
 247        See Also
 248        --------
 249        :attr:`state_shape`
 250        """
 251        return self._hilbert_space.dim
 252    @property
 253    def state_shape(self) -> tuple[int]:
 254        """
 255        The shape of the states in the system.
 256
 257        See Also
 258        --------
 259        :attr:`dim`
 260        """
 261        return (self.dim,)
 262    @property
 263    def n_ctrl(self) -> int:
 264        """
 265        The number of control Hamiltonians.
 266        """
 267        return len(self.Hs)
 268    @property
 269    def evolver(self) -> UnitaryEvolver:
 270        """
 271        The integrator used for time evolutions of the system.
 272
 273        Note
 274        ----
 275        The `evolver` can take a while to initialise and so is not initialised
 276        until `evolver` is is first used or when :meth:`initialise_evolver()` is
 277        called. Using `evolver` before calling :meth:`initialise_evolver()`
 278        initialises the `evolver` with the default parameters of
 279        :meth:`initialise_evolver()`.
 280        """
 281        if self._evolver is None:
 282            self.initialise_evolver()
 283        return self._evolver
 284    def initialise_evolver(self,
 285                           sparse: bool = False,
 286                           force_dynamic: bool = False):
 287        """
 288        Initialises :attr:`evolver` with an evolver from
 289        `PySTE <https://PySTE.readthedocs.io>`__.
 290        `PySTE <https://PySTE.readthedocs.io>`__ is Python
 291        wrapper around the C++ header-only library
 292        `Suzuki-Trotter-Evolver <https://Suzuki-Trotter-Evolver.readthedocs.io>`__:
 293        a fast Schrödinger solver utilising the first-order Suzuki-Trotter
 294        expansion.
 295
 296        Warning
 297        -------
 298        This can take a very long time to execute, especially for large Hilbert
 299        space dimensions. If you plan to evolve the same quantum system many
 300        times we recommended pickling the :attr:`evolver`.
 301
 302        Parameters
 303        ----------
 304        sparse : bool
 305            Whether to use sparse or dense matrices during integration.
 306            To make a decision on whether sparse or dense matrices are likely to
 307            lead to faster integration you can consult the benchmarks at
 308            https://PySTE.readthedocs.io/en/latest/benchmarks.
 309        force_dynamic : bool
 310            Whether to force `PySTE <https://PySTE.readthedocs.io>`__ to use a
 311            dynamic evolver.
 312            
 313            Note
 314            ----
 315            `PySTE <https://PySTE.readthedocs.io>`__ has precompiled evolvers
 316            for specific Hilbert space dimensions and numbers of control
 317            Hamiltonians. When these cannot be found
 318            `PySTE <https://PySTE.readthedocs.io>`__ uses less efficient
 319            evolvers with the Hilbert space dimension and the number of controls
 320            determined dynamically at runtime.
 321        """
 322        self._evolver = get_unitary_evolver(self.H0, self._Hs, sparse, force_dynamic)
 323    def _H(self, ctrl_amp: np.ndarray[float]) -> np.ndarray[complex]:
 324        """
 325        Computes the system Hamiltonian for the specified control amplitudes.
 326
 327        Parameters
 328        ----------
 329        ctrl_amp : NDArray[Shape[s := Any_Shape, :attr:`n_ctrl`], float]
 330            The control amplitudes (stored in the last axis). The prior axes
 331            allow for multiple sets of control amplitudes to be passed and the
 332            Hamiltonian for each computed.
 333
 334        Returns
 335        -------
 336        NDArray[Shape[s, :attr:`dim`, :attr:`dim`], complex]
 337            The system's Hamiltonian (stored in the last two axes).
 338        """
 339        return self._H0 + np.einsum("...i,ijk->...jk", ctrl_amp, self.Hs)
 340    def H(self,
 341          ctrl_amp: Union[np.ndarray[float], np.ndarray[Callable[[float], np.ndarray[float]]], Callable[[float], np.ndarray[float]]]
 342         ) -> Union[np.ndarray[complex], Callable[[float], np.ndarray[complex]]]:
 343        """
 344        Computes the system Hamiltonian for the specified control amplitudes.
 345
 346        Parameters
 347        ----------
 348        ctrl_amp : NDArray[Shape[s := Any_Shape, :attr:`n_ctrl`], float | Callable[[float], np.ndarray[float]]] | Callable[[float], NDArray[Shape[:attr:`n_ctrl`], float]
 349            The control amplitudes (stored in the last axis). The prior axes
 350            allow for multiple sets of control amplitudes to be passed and the
 351            Hamiltonian for each computed. The control amplitudes can be passed
 352            as ``np.ndarray[float]`` to compute the system Hamiltonian for a
 353            specific value of the control ampltiudes. Alternatively,
 354            time-dependent control amplitudes can be passed.
 355            ``np.ndarray[Callable[[float], np.ndarray[float]]]`` can be passed
 356            where each element is a function of time. Alternatively, a function
 357            of time  that returns the control amplitudes can be passed as
 358            ``Callable[[float], NDArray[Shape[:attr:`n_ctrl`], float]``. These
 359            will generate a time-dependent Hamiltonian: a function that takes a
 360            single parameter (time) and returns the Hamiltonian at this time.
 361
 362        Returns
 363        -------
 364        NDArray[Shape[s, :attr:`dim`, :attr:`dim`], complex] | NDArray[Shape[s], Callable[[float], np.ndarray[complex]]]]
 365            Either the systems Hamiltonian stored in the last two axes (if
 366            specific control amplitudes were passed) or a collection of
 367            time-dependent Hamiltonians (if time-dependent controls were
 368            passed).
 369        """
 370        if callable(ctrl_amp):
 371            return lambda t: self._H(ctrl_amp(t))
 372        ctrl_amp = np.array(ctrl_amp)
 373        if ctrl_amp.dtype == object:
 374            return lambda t: self._H([a(t) for a in ctrl_amp])
 375        return self._H(ctrl_amp)
 376    def _pre_processing(self, *args):
 377        """
 378        When calling any evolution method (listed in the
 379        :ref:`See also section <pre_processing_see_also>`)
 380        :meth:`_pre_processing()` is executed on the arguments before the
 381        control amplitudes are modulated by the frequencies (during
 382        :meth:`_envolope_processing()`) and then finally the modulated control
 383        amplitudes are used by the evolution method.
 384
 385        :meth:`_pre_processing()` should be overridden to produce desired pulse
 386        shapes. You can either override :meth:`_pre_processing()` directly by
 387        creating a child class, or you can use :meth:`pulse_form()`.
 388
 389        For :meth:`gradient()` to function correctly :meth:`_pre_processing()`
 390        should be written in `TensorFlow <https://www.tensorflow.org>`__.
 391
 392        Parameters
 393        ----------
 394        *args
 395            The placeholder parameters. See ``_systems.pyi`` for actual
 396            parameters. Each child class that implements a new
 397            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 398            document the parameters for this function: the same parameters as
 399            passed to :meth:`_pre_processing()`.
 400
 401        Returns
 402        -------
 403        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]]
 404            A tuple of
 405            1. The control amplitude envolopes
 406            2. The initial state
 407            3. The integrator time step
 408            4. The frequencies to modulate the control amplitude envolopes with
 409            5. A list of the number of channels for each control Hamiltonian
 410
 411            Warning
 412            -------
 413            The number of channels for each control Hamiltonian must be stored
 414            as a ``list`` and not an ``NDArray`` or a
 415            `TensorFlow <https://www.tensorflow.org>`__ tensor.
 416
 417
 418        .. _pre_processing_see_also:
 419        
 420        See Also
 421        --------
 422        * :meth:`propagate()`
 423        * :meth:`propagate_collection()`
 424        * :meth:`propagate_all()`
 425        * :meth:`evolved_expectation_value()`
 426        * :meth:`evolved_expectation_value_all()`
 427        * :meth:`get_driving_pulses()`
 428        * :meth:`gradient()`
 429        """
 430        return args
 431    def _envolope_processing(self,
 432                             ctrl_amp,
 433                             dt: float,
 434                             frequencies,
 435                             number_channels: list[int]
 436                            ) -> tuple:
 437        """
 438        When calling any evolution method (listed in the
 439        :ref:`See also section <envolope_processing_see_also>` section)
 440        :meth:`_pre_processing()` is executed on the arguements before the
 441        control amplitudes are modulated by the frequencies during
 442        :meth:`_envolope_processing()` and then finally the modulated control
 443        amplitudes are used by the evolution method.
 444
 445        Parameters
 446        ----------
 447        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
 448            The envolope control amplitudes
 449        dt : float
 450            The itegration time step
 451        frequencies : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
 452            The frequencies to modulate the control amplitudes with
 453        number_channels : list[int]
 454            The number of channels associated with each control Hamiltonian
 455
 456            Warning
 457            -------
 458            This must be a ``list`` and not an ``NDArray`` or a
 459            `TensorFlow <https://www.tensorflow.org>`__ tensor.
 460
 461        Returns
 462        -------
 463        tf.Tensor[Shape[n_time_steps, :attr:`n_ctrl`], tf.complex128]
 464            The modulated control amplitudes
 465
 466
 467        .. _envolope_processing_see_also:
 468        
 469        See Also
 470        --------
 471        * :meth:`propagate()`
 472        * :meth:`propagate_collection()`
 473        * :meth:`propagate_all()`
 474        * :meth:`evolved_expectation_value()`
 475        * :meth:`evolved_expectation_value_all()`
 476        * :meth:`get_driving_pulses()`
 477        * :meth:`gradient()`
 478        """
 479        dt = tf.cast(dt, dtype=tf.complex128)
 480        frequencies = tf.cast(frequencies, dtype=tf.complex128)
 481        ctrl_amp = tf.cast(ctrl_amp, dtype=tf.complex128)
 482        channel_couplings = tf.constant(generate_channel_couplings(number_channels), dtype=tf.complex128)
 483        x = tf.exp(tf.einsum("i,j->ij", dt*tf.cast(tf.range(tf.shape(ctrl_amp)[0]), dtype=tf.complex128), -1j*frequencies))
 484        ctrl_amp = tf.cast(tf.math.real(tf.einsum("tc,tc,dc->td", x, ctrl_amp, channel_couplings)), dtype=tf.complex128)
 485        return ctrl_amp
 486    def propagate(self, *args) -> np.ndarray[complex]:
 487        """
 488        Evolves a state vector under the time-dependent Hamiltonian defined by
 489        the control amplitudes using
 490        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.propagate()`
 491        from `PySTE <https://PySTE.readthedocs.io>`__.
 492
 493        Parameters
 494        ----------
 495        *args
 496            The placeholder parameters. See ``_systems.pyi`` for actual
 497            parameters. Each child class that implements a new
 498            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 499            document the parameters for this function: the same parameters as
 500            passed to :meth:`_pre_processing()`.
 501
 502        Returns
 503        -------
 504        NDArray[Shape[:attr:`state_shape`], complex]
 505            The final state
 506
 507        See Also
 508        --------
 509        * :meth:`propagate_collection()`
 510        * :meth:`propagate_all()`
 511        """
 512        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args)
 513        return self.evolver.propagate(ctrl_amp, initial_state, dt)
 514    def propagate_collection(self, *args) -> np.ndarray[complex]:
 515        """
 516        Evolves a collection of state vectors under the time-dependent
 517        Hamiltonian defined by the control amplitudes using
 518        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.propagate_collection()`
 519        from `PySTE <https://PySTE.readthedocs.io>`__.
 520
 521        Parameters
 522        ----------
 523        *args
 524            The placeholder parameters. See ``_systems.pyi`` for actual
 525            parameters. Each child class that implements a new
 526            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 527            document the parameters for this function: the same parameters as
 528            passed to :meth:`_pre_processing()`.
 529
 530        Returns
 531        -------
 532        NDArray[Shape[n_states, :attr:`state_shape`], complex]
 533            The final state
 534
 535        See Also
 536        --------
 537        * :meth:`propagate()`
 538        * :meth:`propagate_all()`
 539        """
 540        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args)
 541        return self.evolver.propagate_collection(ctrl_amp, initial_state, dt)
 542    def propagate_all(self, *args) -> np.ndarray[complex]:
 543        """
 544        Evolves a state vector under the time-dependent Hamiltonian defined by
 545        the control amplitudes using
 546        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.propagate_all()`
 547        from `PySTE <https://PySTE.readthedocs.io>`__ and returns the state at
 548        each time-step.
 549
 550        Parameters
 551        ----------
 552        *args
 553            The placeholder parameters. See ``_systems.pyi`` for actual
 554            parameters. Each child class that implements a new
 555            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 556            document the parameters for this function: the same parameters as
 557            passed to :meth:`_pre_processing()`.
 558
 559        Returns
 560        -------
 561        NDArray[Shape[:attr:`state_shape`, n_time_steps+1], complex]
 562            The state at each integrator time step (including the initial
 563            state).
 564
 565        See Also
 566        --------
 567        * :meth:`propagate()`
 568        * :meth:`propagate_collection()`
 569        """
 570        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args)
 571        return self.evolver.propagate_all(ctrl_amp, initial_state, dt)
 572    def evolved_expectation_value(self, *args) -> complex:
 573        """
 574        Evolves a state vector under the time-dependent Hamiltonian defined by
 575        the control amplitudes and computes the expectation value of a specified
 576        observable with respect to the final state using
 577        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.evolved_expectation_value()`
 578        from `PySTE <https://PySTE.readthedocs.io>`__.
 579
 580        Parameters
 581        ----------
 582        ``*args[:-1]``
 583            The placeholder parameters. See ``_systems.pyi`` for actual
 584            parameters. Each child class that implements a new
 585            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 586            document the parameters for this function: the same parameters as
 587            passed to :meth:`_pre_processing()`.
 588        ``args[-1]`` : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
 589            The observable to take the expectation value of.
 590
 591
 592        Returns
 593        -------
 594        complex
 595            The expectation value.
 596
 597        See Also
 598        --------
 599        * :meth:`evolved_expectation_value_all()`
 600        * :meth:`gradient()`
 601        """
 602        observable: np.ndarray[complex] = args[-1]
 603        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args[:-1])
 604        return self.evolver.evolved_expectation_value(ctrl_amp,
 605                                                      initial_state,
 606                                                      dt,
 607                                                      observable)
 608    def evolved_expectation_value_all(self, *args) -> np.ndarray[complex]:
 609        """
 610        Evolves a state vector under the time-dependent Hamiltonian defined by
 611        the control amplitudes and computes the expectation value of a specified
 612        observable with respect to the state at each time-step using
 613        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.evolved_expectation_value_all()`
 614        from `PySTE <https://PySTE.readthedocs.io>`__.
 615
 616        Parameters
 617        ----------
 618        ``*args[:-1]``
 619            The placeholder parameters. See ``_systems.pyi`` for actual
 620            parameters. Each child class that implements a new
 621            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 622            document the parameters for this function: the same parameters as
 623            passed to :meth:`_pre_processing()`.
 624        ``args[-1]`` : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
 625            The observable to take the expectation value of.
 626
 627        Returns
 628        -------
 629        NDArray[Shape[n_time_steps+1], complex]
 630            The state at each integrator time step (including the initial
 631            state).
 632
 633        See Also
 634        --------
 635        * :meth:`evolved_expectation_value()`
 636        * :meth:`gradient()`
 637        """
 638        observable: np.ndarray[complex] = args[-1]
 639        ctrl_amp, initial_state, dt = self.get_driving_pulses(*args[:-1])
 640        return self.evolver.evolved_expectation_value_all(ctrl_amp,
 641                                                          initial_state,
 642                                                          dt,
 643                                                          observable)
 644    def get_driving_pulses(self, *args) -> tuple[np.ndarray[complex], np.ndarray[complex], float]:
 645        """
 646        When calling any evolution method (listed in the
 647        :ref:`See also section <get_driving_pulses_see_also>`) :meth:`get_driving_pulses()`
 648        is executed on the arguements before the evolution method.
 649
 650        Parameters
 651        ----------
 652        *args
 653            The placeholder parameters. See ``_systems.pyi`` for actual
 654            parameters. Each child class that implements a new
 655            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 656            document the parameters for this function: the same parameters as
 657            passed to :meth:`_pre_processing()`.
 658
 659        Returns
 660        -------
 661        tuple[NDArray[Shape[n_time_steps, :attr:`n_ctrl`], complex], NDArray[Shape[:attr:`state_shape`], complex], float]
 662            A tuple of:
 663            1. Control amplitudes
 664            2. Initial state
 665            3. Integrator time step
 666
 667
 668        .. _get_driving_pulses_see_also:
 669        
 670        See Also
 671        --------
 672        * :meth:`propagate()`
 673        * :meth:`propagate_collection()`
 674        * :meth:`propagate_all()`
 675        * :meth:`evolved_expectation_value()`
 676        * :meth:`evolved_expectation_value_all()`
 677        * :meth:`gradient()`
 678        """
 679        ctrl_amp, initial_state, dt = self._processing(*args)
 680        try: ctrl_amp: np.ndarray[complex] = ctrl_amp.numpy()
 681        except: pass
 682        try:  initial_state: np.ndarray[complex] = initial_state.numpy().flatten()
 683        except: pass
 684        try: dt = dt.numpy()
 685        except: pass
 686        dt = float(dt.real)
 687        return ctrl_amp, initial_state, dt
 688    def _eager_processing(self, *args) -> tuple:
 689        """
 690        Executes :meth:`_pre_processing()` followed by
 691        :meth:`_envolope_processing()` eagerly (i.e. without using a
 692        `TensorFlow <https://www.tensorflow.org>`__ graph). Nonetheless,
 693        :meth:`_eager_processing()` is still auto differentiable.
 694
 695        Parameters
 696        ----------
 697        *args
 698            The placeholder parameters. See ``_systems.pyi`` for actual
 699            parameters. Each child class that implements a new
 700            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 701            document the parameters for this function: the same parameters as
 702            passed to :meth:`_pre_processing()`.
 703
 704        Returns
 705        -------
 706        tuple[tf.Tensor[Shape[n_time_steps, :attr:`n_ctrl`], complex], tf.Tensor[Shape[:attr:`state_shape`], complex], tf.Tensor[Shape[], float]]
 707            A tuple of:
 708            1. Control amplitudes
 709            2. Initial state
 710            3. Integrator time step
 711        """
 712        ctrl_amp, initial_state, dt, frequencies, number_channels = self._pre_processing(*args)
 713        return self._envolope_processing(ctrl_amp, dt, frequencies, list(number_channels)), initial_state, dt
 714    def _traceable_eager_processing(self, *args) -> tuple:
 715        """
 716        A function that will be traced by
 717        `TensorFlow <https://www.tensorflow.org>`__ to produce a graph of
 718        :meth:`_pre_processing()` followed by :meth:`_envolope_processing()`.
 719
 720        Parameters
 721        ----------
 722        *args
 723            The placeholder parameters. See ``_systems.pyi`` for actual
 724            parameters. Each child class that implements a new
 725            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 726            document the parameters for this function: the same parameters as
 727            passed to :meth:`_pre_processing()`.
 728
 729        Returns
 730        -------
 731        tuple[tf.Tensor[Shape[n_time_steps, :attr:`n_ctrl`], complex], tf.Tensor[Shape[:attr:`state_shape`], complex], tf.Tensor[Shape[], float]]
 732            A tuple of:
 733            1. Control amplitudes
 734            2. Initial state
 735            3. Integrator time step
 736        """
 737        print("Tracing control amplitude graph.")
 738        return self._eager_processing(*args)
 739    def gradient(self, *args) -> tuple[float, np.ndarray[float]]:
 740        """
 741        Evolves a state vector under the time-dependent Hamiltonian defined by
 742        the control amplitudes and computes the expectation value of a specified
 743        observable with respect to the final state and then computes the
 744        gradient of the final state with respect to the first argument
 745        (``args[0]``) using
 746        :meth:`~py_ste.evolvers.DenseUnitaryEvolver.switching_function()`
 747        from `PySTE <https://PySTE.readthedocs.io>`__.
 748
 749        Parameters
 750        ----------
 751        ``*args[:-1]``
 752            The placeholder parameters. See ``_systems.pyi`` for actual
 753            parameters. Each child class that implements a new
 754            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 755            document the parameters for this function: the same parameters as
 756            passed to :meth:`_pre_processing()`.
 757        ``args[-1]`` : NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
 758            The observable to take the expectation value of.
 759
 760        Returns
 761        -------
 762        tuple[complex, NDArray[Shape[n_parameters], float]]
 763            A tuple of the expectation value and the gradient.
 764
 765        See Also
 766        --------
 767        * :meth:`evolved_expectation_value()`
 768        * :meth:`evolved_expectation_value_all()`
 769        """
 770        cost: np.ndarray[complex] = args[-1]
 771        args = list(args[:-1])
 772        args[0] = tf.constant(args[0], dtype=tf.float64)             
 773        with tf.GradientTape(persistent=False,
 774                             watch_accessed_variables=False
 775                            ) as tape:
 776            tape.watch(args[0])
 777            ctrl_amp, initial_state, dt = self._processing(*args)
 778            try: initial_state: np.ndarray[complex] = initial_state.numpy().flatten()
 779            except: pass
 780            try: dt = dt.numpy()
 781            except: pass
 782            dt = float(dt.real)
 783            E = ExpValCustom(self, initial_state, dt, cost).run(ctrl_amp)
 784        grad = tape.gradient(E, args[0])
 785        del tape
 786        
 787        grad = tf.convert_to_tensor(grad)
 788        grad: np.ndarray[float] = grad.numpy()
 789        grad = grad.real
 790        return E.numpy(), grad
 791    def pulse_form(self,
 792                   pulse_function: Callable,
 793                   append: bool = False,                  
 794                  ) -> "PulseForm":
 795        """
 796        Initialises a new :class:`QuantumSystem` in which
 797        :meth:`_pre_processing()` corresponds to executing ``pulse_function()``
 798        and piping the output into the previous definition of
 799        :meth:`_pre_processing()`.
 800
 801        Parameters
 802        ----------
 803        pulse_function : Callable
 804            The function to compose with :meth:`_pre_processing()`.
 805
 806        Returns
 807        -------
 808        PulseForm
 809            The new :class:`QuantumSystem`
 810        """
 811        return PulseForm(self, pulse_function, append)
 812
 813class TransformedSystem(QuantumSystem):
 814    """
 815    A base class for representing a transformation on a :class:`qugrad.QuantumSystem`.
 816    """
 817
 818    _original_system: QuantumSystem
 819    "The system that was transformed into this system"
 820    
 821    _base_system: QuantumSystem
 822    """
 823    The system before any transformations were applied. That is `_base_system`
 824    is the recursive :attr:`original_system`
 825    (``original_system.original_system.original_system....``) until
 826    :attr:`original_system` is no longer a :class:`TransformedSystem`.
 827    """
 828    
 829    def __init__(self,
 830                 original_system: QuantumSystem,
 831                 H0: np.ndarray[complex],
 832                 Hs: Union[np.ndarray[complex], np.ndarray[complex]],
 833                 hilbert_space: HilbertSpace):
 834        """
 835        Performs a transformation on a :class:`qugrad.QuantumSystem`.
 836
 837        Parameters
 838        ----------
 839        original_system: QuantumSystem
 840            The system to be transformed into this system
 841        H0: NDArray[Shape[:attr:`dim`, :attr:`dim`], complex]
 842            The new drift Hamiltonian
 843        Hs: NDArray[Shape[":attr:`n_ctrl`, :attr:`dim`, :attr:`dim`"], complex] | NDArray[Shape[:attr:`n_ctrl` * :attr:`dim`, :attr:`dim`], complex]
 844            The new control Hamiltonians either as an array of control
 845            Hamiltonians or the control Hamiltonians stacked along the first
 846            axis.
 847        hilbert_space: HilbertSpace
 848            The new Hilbert space of the system
 849        """
 850        self._original_system = original_system
 851        if isinstance(original_system, TransformedSystem):
 852            self._base_system = original_system._base_system
 853        else:
 854            self._base_system = original_system
 855        super().__init__(H0, Hs, hilbert_space, self._base_system.using_graph)
 856    @property
 857    def original_system(self) -> QuantumSystem:
 858        "The system that was transformed into this system"
 859        return self._original_system
 860    @property
 861    def base_system(self) -> QuantumSystem:
 862        """
 863        The system before any transformations were applied. That is
 864        :attr:`base_system` is the recursive :attr:`original_system`
 865        (``original_system.original_system.original_system....``) until
 866        :attr:`original_system` is no longer a :class:`TransformedSystem`.
 867        """
 868        return self._base_system
 869    def _pre_processing(self, *args) -> tuple:
 870        """
 871        When calling any evolution method (listed in the
 872        :ref:`See also section <TransformedSystem_pre_processing_see_also>` section)
 873        :meth:`_pre_processing()` is executed on the arguements before the
 874        control amplitudes are modulated by the frequencies (during
 875        :meth:`_envolope_processing()`) and then finally the modulated control
 876        amplitudes are used by the evolution method.
 877
 878        This is a placeholder for ``original_system._pre_processing()``.
 879
 880        Parameters
 881        ----------
 882        *args
 883            The placeholder parameters. See ``_systems.pyi`` for actual
 884            parameters. Each child class that implements a new
 885            :meth:`_pre_processing()` should implement a ``.pyi`` file to
 886            document the parameters for this function: the same parameters as
 887            passed to :meth:`_pre_processing()`.
 888
 889        Returns
 890        -------
 891        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]]
 892            A tuple of
 893            1. The control amplitude envolopes
 894            2. The initial state
 895            3. The integrator time step
 896            4. The frequencies to modulate the control amplitude envolopes with
 897            5. A list of the number of channels for each control Hamiltonian
 898
 899            Warning
 900            -------
 901            The number of channels for each control Hamiltonian must be stored
 902            as a ``list`` and not an ``NDArray`` or a
 903            `TensorFlow <https://www.tensorflow.org>`__ tensor.
 904
 905
 906        .. _TransformedSystem_pre_processing_see_also:
 907        
 908        See Also
 909        --------
 910        * :meth:`propagate()`
 911        * :meth:`propagate_collection()`
 912        * :meth:`propagate_all()`
 913        * :meth:`evolved_expectation_value()`
 914        * :meth:`evolved_expectation_value_all()`
 915        * :meth:`get_driving_pulses()`
 916        * :meth:`gradient()`
 917        """
 918        return self._original_system._pre_processing(*args)
 919    def _envolope_processing(self,
 920                             ctrl_amp,
 921                             dt: float,
 922                             frequencies,
 923                             number_channels: list[int]):
 924        """
 925        When calling any evolution method (listed in the
 926        :ref:`See also section <envolope_processing_see_also>` section) :meth:`_pre_processing()`
 927        is executed on the arguements before the control amplitudes are
 928        modulated by the frequencies during :meth:`_envolope_processing()` and
 929        then finally the modulated control amplitudes are used by the evolution
 930        method.
 931
 932        Parameters
 933        ----------
 934        ctrl_amp : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
 935            The envolope control amplitudes
 936        dt : float
 937            The itegration time step
 938        frequencies : tf.Tensor[Shape[n_time_steps, total_n_channels], tf.complex128]
 939            The frequencies to modulate the control amplitudes with
 940        number_channels : list[int]
 941            The number of channels associated with each control Hamiltonian
 942
 943            Warning
 944            -------
 945            This must be a ``list`` and not an ``NDArray`` or a
 946            `TensorFlow <https://www.tensorflow.org>`__ tensor.
 947
 948        Returns
 949        -------
 950        tf.Tensor[Shape[n_time_steps,:attr:`n_ctrl`], tf.complex128]
 951            The modulated control amplitudes
 952
 953
 954        .. _envolope_processing_see_also:
 955        
 956        See Also
 957        --------
 958        * :meth:`propagate()`
 959        * :meth:`propagate_collection()`
 960        * :meth:`propagate_all()`
 961        * :meth:`evolved_expectation_value()`
 962        * :meth:`evolved_expectation_value_all()`
 963        * :meth:`get_driving_pulses()`
 964        * :meth:`gradient()`
 965        """
 966        return self._original_system._envolope_processing(ctrl_amp,
 967                                                          dt,
 968                                                          frequencies,
 969                                                          number_channels)
 970
 971class PulseForm(TransformedSystem):
 972    """
 973    A transformed :class:`qugrad.QuantumSystem` in which :meth:`_pre_processing()`
 974    has been composed with another pre processing function.
 975    """
 976
 977    _pulse_function: Callable
 978    "The function composed with ``original_system._pre_processing()``"
 979    
 980    _appended: bool
 981    """
 982    Whether the :attr:`pulse_function` was prepended
 983    (``original_system._pre_processing(*pulse_function())``) or appended
 984    (``pulse_function(*original_system._pre_processing())``)
 985    """
 986        
 987    def __init__(self,
 988                 original_system: QuantumSystem,
 989                 pulse_function: Callable,
 990                 append: bool = False):
 991        """
 992        Initialises a new :class:`qugrad.QuantumSystem` in which :meth:`_pre_processing()`
 993        corresponds to running ``pulse_function()`` and piping the output into
 994        ``original_system._pre_processing()``.
 995
 996        Parameters
 997        ----------
 998        original_system : QuantumSystem
 999            The system that was transformed into this system
1000        pulse_function : Callable
1001            The function to compose with :meth:`_pre_processing()`.
1002        append : bool
1003            Whether to prepend
1004            (``original_system._pre_processing(*pulse_function())``) or append
1005            (``pulse_function(*original_system._pre_processing())``)
1006            ``pulse_function``.
1007        """
1008        super().__init__(original_system,
1009                         original_system.H0,
1010                         original_system.Hs,
1011                         original_system.hilbert_space)
1012        self._evolver = original_system._evolver
1013        if append:
1014            self._pre_processing = compose_unpack(
1015                pulse_function,
1016                original_system._pre_processing
1017            )
1018        else:
1019            self._pre_processing = compose_unpack(
1020                original_system._pre_processing,
1021                pulse_function
1022            )
1023        self._pulse_function = pulse_function
1024        self._appended = append
1025    @property
1026    def pulse_function(self) -> Callable:
1027        "The function composed with ``original_system._pre_processing()``"
1028        return self._pulse_function
1029    @property
1030    def appended(self) -> bool:
1031        """
1032        Whether the :attr:`pulse_function` was prepended
1033        (``original_system._pre_processing(*pulse_function())``) or appended
1034        (``pulse_function(*original_system._pre_processing())``)
1035        """
1036        return self._appended