1"""
2A base classes for quantum systems constructed from a contraction of
3coefficients with a set of operators.
4"""
5
6import numpy as np
7from qugrad import QuantumSystem, HilbertSpace
8
9def partial_flatten(array: np.ndarray, n: int) -> np.ndarray:
10 """
11 Flattens the first `n` dimensions of an `array`. If `n` is negative, it
12 flattens the last ``-n`` dimensions of the array.
13
14 Parameters
15 ----------
16 array : np.ndarray
17 The array to be partially flattened.
18 n : int
19 The number of dimensions to flatten. If `n` is negative, it flattens
20 the last ``-n`` dimensions of the ``array``.
21
22 Returns
23 -------
24 np.ndarray
25 The partially flattened array.
26 """
27 if array.size == 0:
28 if n >= 0:
29 return np.empty((np.prod(array.shape[:n], dtype=int),)+array.shape[n:], dtype=array.dtype)
30 return np.empty((array.shape[:n])+(np.prod(array.shape[n:], dtype=int),), dtype=array.dtype)
31 if n >= 0:
32 return array.reshape((-1,)+array.shape[n:])
33 return array.reshape((array.shape[:n])+(-1,))
34
35def contract_general(a: np.ndarray,
36 b: np.ndarray,
37 remaining_ndim_b: int
38 ) -> np.ndarray:
39 """
40 Contracts `a` and `b` such that only the last `remaining_ndim_b` axes of `b`
41 and any excess axes of `a` are left uncontracted.
42
43 Parameters
44 ----------
45 a : np.ndarray
46 The first array to be contracted.
47 b : np.ndarray
48 The second array to be contracted.
49 remaining_ndim_b : int
50 The number of axes of `b` (from the end) that should remain
51 uncontracted.
52
53 Returns
54 -------
55 np.ndarray
56 The result of the contraction.
57 """
58 ndims_to_contract = b.ndim-remaining_ndim_b
59 b = partial_flatten(b, ndims_to_contract)
60 if ndims_to_contract == 0:
61 a = np.expand_dims(a, axis=-1)
62 else:
63 a = partial_flatten(a, -ndims_to_contract)
64 if remaining_ndim_b == 0:
65 return np.einsum("...i,i->...", a, b)
66 a_pre_shape = a.shape[:-1]
67 b_post_shape = b.shape[1:]
68 shape = a_pre_shape + b_post_shape
69 b = partial_flatten(b, -remaining_ndim_b)
70 output = np.einsum("...i,ij->...j", a, b)
71 return output.reshape(shape)
72
73def contract_skeleton(coefficients: np.ndarray,
74 skeleton:np.ndarray
75 ) -> np.ndarray:
76 """
77 Contracts coefficients with an operator skeleton.
78
79 Parameters
80 ----------
81 coefficients : np.ndarray
82 The coefficients to be contracted with the skeleton.
83 skeleton : np.ndarray
84 The operator skeleton to be contracted with the coefficients.
85
86 Returns
87 -------
88 np.ndarray
89 The result of the contraction.
90 """
91 return contract_general(coefficients, skeleton, 2)
92
93def contract_skeletons(coefficients: list[np.ndarray],
94 skeletons: list[np.ndarray]
95 ) -> np.ndarray:
96 """
97 Contracts a list of coefficients with a list of operator skeletons and then
98 sums the results.
99
100 Parameters
101 ----------
102 coefficients : list[np.ndarray]
103 The coefficients to be contracted with the skeletons.
104 skeletons : list[np.ndarray]
105 The operator skeletons to be contracted with the coefficients.
106
107 Returns
108 -------
109 np.ndarray
110 The result of the contractions.
111 """
112 return sum(map(contract_skeleton, coefficients, skeletons))
113
114def flatten_skeleton(skeleton: np.ndarray) -> np.ndarray:
115 """
116 Flattens a skeleton into an array of operators.
117
118 Parameters
119 ----------
120 skeleton : np.ndarray
121 The operator skeleton to be flattened.
122
123 Returns
124 -------
125 np.ndarray
126 The flattened skeleton.
127 """
128 return partial_flatten(skeleton, skeleton.ndim - 2)
129
130def get_Hs(coefficients: list[np.ndarray],
131 skeletons: list[np.ndarray]
132 ) -> np.ndarray:
133 """
134
135 """
136 Hs = map(contract_skeleton, coefficients, skeletons)
137 return np.concatenate(list(map(flatten_skeleton, Hs)))
138
[docs]
139class SkeletalSystem(QuantumSystem):
140 """
141 A quantum system with a Hamiltonian built upon an operator skeleton.
142 """
[docs]
143 def __init__(self,
144 drift_coefficients: list[np.ndarray],
145 drift_skeletons: list[np.ndarray],
146 ctrl_coefficients: list[np.ndarray],
147 ctrl_skeletons: list[np.ndarray],
148 hilbert_space: HilbertSpace,
149 use_graph: bool = True):
150 """
151
152 Parameters
153 ----------
154 drift_coefficients : list[np.ndarray]
155 The coefficients of the drift Hamiltonian
156 drift_skeletons : list[np.ndarray]
157 The operator skeletons for the drift Hamiltonian
158 coefficients : list[np.ndarray]
159 The coefficients of the control Hamiltonians
160 skeletons : list[np.ndarray]
161 The operator skeletons for the control Hamiltonians
162 hilbert_space : HilbertSpace
163 The Hilbert space of the system
164 use_graph : bool
165 Whether to use `TensorFlow <https://www.tensorflow.org>`__ graphs
166 during computation, by default ``True``
167 """
168 Hs = get_Hs(ctrl_coefficients, ctrl_skeletons)
169 del ctrl_skeletons
170 H0 = contract_skeletons(drift_coefficients, drift_skeletons)
171 del drift_skeletons
172 super().__init__(H0, Hs, hilbert_space, use_graph)