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