Source code for py_ste._wrapper
1import numpy as np
2from numpy.typing import ArrayLike
3# import os
4# os.environ["OMP_NUM_THREADS"] = "2"
5from . import evolvers
6
[docs]
7def get_unitary_evolver(drift_hamiltonian: ArrayLike,
8 control_hamiltonians: ArrayLike,
9 sparse: bool = False,
10 force_dynamic: bool = False
11 ) -> evolvers.UnitaryEvolver:
12 """
13 Initialises a class to store the diagonalised drift and control
14 Hamiltonians. On initialisation the Hamiltonians are diagonalised and the
15 eigenvectors and values stored. This initial diagonalisation may be slow and
16 takes
17 $O(\\textrm{dim}^3)$
18 time for a
19 $\\textrm{dim}\\times \\textrm{dim}$
20 Hamiltonian. However, it allows each step of the Suzuki-Trotter expansion
21 to be implimented in
22 $O(\\textrm{dim}^2)$
23 time with matrix multiplication and only scalar exponentiation opposed to
24 matrix exponentiation which takes
25 $O(\\textrm{dim}^3)$
26 time.
27
28 Parameters
29 ----------
30 drift_hamiltonian : ArrayLike
31 The drift Hamiltonian. Must be a square matrix of dimension ``dim``.
32 control_hamiltonians : ArrayLike
33 The control Hamiltonians. Can either by a 3D array with shape
34 ``(n_ctrl, dim, dim)`` representing a stack of control Hamiltonians
35 indexed by the first axis. Alternatively a 3D array with shape
36 ``(n_ctrl * dim, dim)`` can be passed with the control Hamiltonians
37 being concatenated along the first axis.
38 sparse : bool, optional
39 Whether to use sparse matrices for the evolution. The default is False.
40 force_dynamic : bool, optional
41 Whether to force the use of dynamically sized matrices for the
42 evolution. The default is False.
43
44 Returns
45 -------
46 evolvers.UnitaryEvolver
47 An instance of a child class of
48 :class:`evolvers.UnitaryEvolver <py_ste.evolvers.UnitaryEvolver>`.
49 If ``sparse == False`` then the returned instance will be a child class
50 of
51 :class:`evolvers.DenseUnitaryEvolver <py_ste.evolvers.DenseUnitaryEvolver>`
52 else
53 :class:`evolvers.SparseUnitaryEvolver <py_ste.evolvers.SparseUnitaryEvolver>`
54 is returned. Both
55 :class:`evolvers.DenseUnitaryEvolver <py_ste.evolvers.DenseUnitaryEvolver>`
56 and
57 :class:`evolvers.SparseUnitaryEvolver <py_ste.evolvers.SparseUnitaryEvolver>`
58 are dynamic evolvers: evolvers for which the number of control
59 Hamiltonians and the vector space dimension are determined at runtime
60 based on the shapes of ``drift_hamiltonian`` and ``control_hamiltonians``. If possible
61 ``get_unitary_evolver()`` will return an instance of
62 :class:`evolvers.DenseUnitaryEvolver_nctrl_dim <py_ste.evolvers.DenseUnitaryEvolver_nctrl_dim>`
63 or
64 :class:`evolvers.SparseUnitaryEvolver_nctrl_dim <py_ste.evolvers.SparseUnitaryEvolver_nctrl_dim>`,
65 where ``nctrl`` and ``dim`` are substituted for their corresponding
66 values.
67 These are fixed evolvers where the number of controls (``nctrl``) and
68 the vector space dimension (``dim``) are know at the C++ compile time
69 allowing for more efficient C++ methods to be compiled.
70
71 Note
72 ----
73 An instance of
74 :class:`evolvers.UnitaryEvolver <py_ste.evolvers.UnitaryEvolver>`
75 itself will never be returned.
76 :class:`evolvers.UnitaryEvolver <py_ste.evolvers.UnitaryEvolver>`
77 is simply a base class for all evolvers allowing for checks such as::
78
79 isinstance(evolver, evolvers.UnitaryEvolver)
80
81 Raises
82 ------
83 ValueError
84 ``drift_hamiltonian`` must be 2D.
85 ValueError
86 ``drift_hamiltonian`` must be square.
87 ValueError
88 The control Hamiltonians (``control_hamiltonians``) must have the same number of columns as
89 ``drift_hamiltonian``.
90 ValueError
91 ``control_hamiltonians`` must be 2D or 3D.
92 ValueError
93 Each control Hamiltonian in ``control_hamiltonians`` must have the same dimension as ``drift_hamiltonian``.
94 """
95 drift_hamiltonian = np.array(drift_hamiltonian, dtype=np.complex128)
96 control_hamiltonians = np.array(control_hamiltonians, dtype=np.complex128)
97
98 if drift_hamiltonian.ndim != 2:
99 raise ValueError("``drift_hamiltonian`` must be 2D.")
100
101 dim: int = drift_hamiltonian.shape[0]
102 if drift_hamiltonian.shape[1] != dim:
103 raise ValueError("``drift_hamiltonian`` must be square.")
104
105 if control_hamiltonians.ndim == 3:
106 control_hamiltonians = control_hamiltonians.reshape((-1, dim))
107 elif control_hamiltonians.ndim == 2:
108 if control_hamiltonians.shape[1] != dim:
109 raise ValueError("The control Hamiltonians (``control_hamiltonians``) must have the same number of columns as ``drift_hamiltonian``.")
110 elif control_hamiltonians.size == 0:
111 control_hamiltonians = np.zeros((0, dim), dtype=np.complex128)
112 else:
113 raise ValueError("``control_hamiltonians`` must be 2D or 3D.")
114
115 if control_hamiltonians.shape[0] % dim != 0:
116 raise ValueError("Each control Hamiltonian in control_hamiltonians must have the same dimension as ``drift_hamiltonian``.")
117
118 n_ctrl: int = control_hamiltonians.shape[0]//dim
119 name: str = ("Sparse" if sparse else "Dense") + "UnitaryEvolver"
120 if force_dynamic:
121 Evolver = getattr(evolvers, name)
122 else:
123 try:
124 Evolver = getattr(evolvers, name+f"_{n_ctrl}_{dim}")
125 except AttributeError:
126 try:
127 Evolver = getattr(evolvers, name+f"_Dynamic_{dim}")
128 except AttributeError:
129 try:
130 Evolver = getattr(evolvers, name+f"_{n_ctrl}_Dynamic")
131 except AttributeError:
132 Evolver = getattr(evolvers, name)
133 return Evolver(drift_hamiltonian, control_hamiltonians)