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