Coverage for jaxquantum/circuits/circuits.py: 89%

1"""Circuits. 

2 

3Inspired by a mix of Cirq and Qiskit circuits. 

4""" 

5 

6from flax import struct 

7from jax import config 

8from typing import List, Optional, Union 

9from copy import deepcopy 

10from numpy import argsort 

11import jax.numpy as jnp 

12 

13from jaxquantum.core.operators import identity 

14from jaxquantum.circuits.gates import Gate 

15from jaxquantum.circuits.constants import SimulateMode 

16from jaxquantum.core.qarray import Qarray, concatenate 

17 

18 

19config.update("jax_enable_x64", True) 

20 

21 

22@struct.dataclass 

23class Register: 

24 dims: List[int] = struct.field(pytree_node=False) 

25 

26 @classmethod 

27 def create(cls, dims: List[int]): 

28 return Register(dims=dims) 

29 

30 def __eq__(self, other): 

31 if isinstance(other, Register): 

32 return self.dims == other.dims 

33 return False 

34 

35 

36@struct.dataclass 

37class Operation: 

38 gate: Gate 

39 indices: List[int] = struct.field(pytree_node=False) 

40 register: Register 

41 

42 @classmethod 

43 def create(cls, gate: Gate, indices: Union[int, List[int]], register: Register): 

44 if isinstance(indices, int): 

45 indices = [indices] 

46 

47 assert gate.num_modes == len(indices), ( 

48 "Number of indices must match gate's num_modes." 

49 ) 

50 assert gate.dims == [register.dims[ind] for ind in indices], ( 

51 "Indices must match register dimensions." 

52 ) 

53 

54 if any( 

55 (0 > ind and ind >= len(register.dims)) or not isinstance(ind, int) 

56 for ind in indices 

57 ): 

58 raise ValueError("Indices must be integers within the register.") 

59 

60 return Operation(gate=gate, indices=indices, register=register) 

61 

62 

63 def promote(self, op: Qarray) -> Qarray: 

64 indices_order = self.indices 

65 missing_indices = [ 

66 i for i in range(len(self.register.dims)) if i not in indices_order 

67 ] 

68 for j in missing_indices: 

69 op = op ^ identity(self.register.dims[j]) 

70 combined_indices = indices_order + missing_indices 

71 sorted_ind = list(argsort(combined_indices)) 

72 op = op.transpose(sorted_ind) 

73 return op 

74 

75 

76@struct.dataclass 

77class Layer: 

78 operations: List[Operation] = struct.field(pytree_node=False) 

79 _unique_indices: List[int] = struct.field(pytree_node=False) 

80 _default_simulate_mode: SimulateMode = struct.field(pytree_node=False) 

81 

82 @classmethod 

83 def create( 

84 cls, operations: List[Operation], default_simulate_mode=SimulateMode.UNITARY 

85 ): 

86 all_indices = [ind for op in operations for ind in op.indices] 

87 unique_indices = list(set(all_indices)) 

88 

89 if ( 

90 default_simulate_mode != SimulateMode.HAMILTONIAN 

91 and len(all_indices) != len(unique_indices) 

92 ): 

93 raise ValueError("Operations must not have overlapping indices.") 

94 

95 return Layer( 

96 operations=operations, 

97 _unique_indices=unique_indices, 

98 _default_simulate_mode=default_simulate_mode, 

99 ) 

100 

101 def add(self, operation: Operation): 

102 if self._default_simulate_mode != SimulateMode.HAMILTONIAN and any( 

103 ind in self._unique_indices for ind in operation.indices 

104 ): 

105 raise ValueError("Operations must not have overlapping indices.") 

106 self.operations.append(operation) 

107 self._unique_indices.extend(operation.indices) 

108 

109 def gen_U(self): 

110 U = None 

111 

112 if len(self.operations) == 0: 

113 return None 

114 

115 indices_order = [] 

116 for operation in self.operations: 

117 indices_order += operation.indices 

118 

119 if U is None: 

120 U = operation.gate.U 

121 else: 

122 U = U ^ operation.gate.U 

123 

124 register = self.operations[0].register 

125 missing_indices = [ 

126 i for i in range(len(register.dims)) if i not in indices_order 

127 ] 

128 

129 for j in missing_indices: 

130 U = U ^ identity(register.dims[j]) 

131 

132 combined_indices = indices_order + missing_indices 

133 

134 sorted_ind = list(argsort(combined_indices)) 

135 U = U.transpose(sorted_ind) 

136 return U 

137 

138 def gen_Ht(self): 

139 Ht = lambda t: 0 

140 

141 if len(self.operations) == 0: 

142 return Ht 

143 

144 for operation in self.operations: 

145 def Ht(t, prev_Ht=Ht, prev_operation=operation): 

146 return prev_Ht(t) + prev_operation.promote(prev_operation.gate.Ht(t)) 

147 

148 return Ht 

149 

150 def gen_KM(self): 

151 KM = Qarray.from_list([]) 

152 

153 if len(self.operations) == 0: 

154 return KM 

155 

156 indices_order = [] 

157 for operation in self.operations: 

158 gate_map = operation.gate.KM 

159 if len(gate_map) == 0: 

160 continue 

161 

162 indices_order += operation.indices 

163 

164 if len(KM) == 0: 

165 KM = deepcopy(gate_map) 

166 else: 

167 KM = KM ^ gate_map 

168 

169 if len(KM) == 0: 

170 return KM 

171 

172 register = self.operations[0].register 

173 missing_indices = [ 

174 i for i in range(len(register.dims)) if i not in indices_order 

175 ] 

176 

177 for j in missing_indices: 

178 KM = KM ^ identity(register.dims[j]) 

179 

180 combined_indices = indices_order + missing_indices 

181 sorted_ind = list(argsort(combined_indices)) 

182 

183 KM = KM.transpose(sorted_ind) 

184 

185 return KM 

186 

187 def gen_c_ops(self): 

188 c_ops = Qarray.from_list([]) 

189 

190 if len(self.operations) == 0: 

191 return c_ops 

192 

193 for operation in self.operations: 

194 if len(operation.gate.c_ops) == 0: 

195 continue 

196 promoted_c_ops = operation.promote(operation.gate.c_ops) 

197 c_ops = concatenate([c_ops, promoted_c_ops]) 

198 

199 return c_ops 

200 

201 def gen_ts(self): 

202 ts = None 

203 

204 for operation in self.operations: 

205 if operation.gate.ts is not None and len(operation.gate.ts) > 0: 

206 if ts is None: 

207 ts = operation.gate.ts 

208 else: 

209 assert jnp.array_equal(ts, operation.gate.ts), ( 

210 "All operations in a layer must have the same specified time steps, but not all operations need to have time steps." 

211 ) 

212 return ts 

213 

214@struct.dataclass 

215class Circuit: 

216 register: Register 

217 layers: List[Layer] = struct.field(pytree_node=False) 

218 

219 @classmethod 

220 def create(cls, register: Register, layers: Optional[List[Layer]] = None): 

221 if layers is None: 

222 layers = [] 

223 

224 return Circuit( 

225 register=register, 

226 layers=layers, 

227 ) 

228 

229 def append_layer(self, layer: Layer): 

230 self.layers.append(layer) 

231 

232 def append_operation( 

233 self, operation: Operation, default_simulate_mode: Optional[SimulateMode] = None, new_layer: bool =True 

234 ): 

235 assert operation.register == self.register, ( 

236 f"Mismatch in operation register {operation.register} and circuit register {self.register}." 

237 ) 

238 

239 new_layer = new_layer or len(self.layers) == 0 

240 

241 if new_layer: 

242 default_simulate_mode = default_simulate_mode if default_simulate_mode is not None else SimulateMode.UNITARY 

243 self.append_layer( 

244 Layer.create([operation], default_simulate_mode=default_simulate_mode) 

245 ) 

246 else: 

247 if default_simulate_mode is not None: 

248 assert ( 

249 self.layers[-1]._default_simulate_mode == default_simulate_mode 

250 ), "Cannot append operation to last layer with different default simulate mode." 

251 

252 self.layers[-1].add(operation) 

253 

254 def append( 

255 self, 

256 gate: Gate, 

257 indices: Union[int, List[int]], 

258 default_simulate_mode: Optional[SimulateMode] = None, 

259 new_layer: bool = True, 

260 ): 

261 operation = Operation.create(gate, indices, self.register) 

262 self.append_operation(operation, default_simulate_mode=default_simulate_mode, new_layer=new_layer)