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

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 if len(operation.gate.KM) == 0: 

159 continue 

160 

161 indices_order += operation.indices 

162 

163 if len(KM) == 0: 

164 KM = deepcopy(operation.gate.KM) 

165 else: 

166 KM = KM ^ operation.gate.KM 

167 

168 if len(KM) == 0: 

169 return KM 

170 

171 register = self.operations[0].register 

172 missing_indices = [ 

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

174 ] 

175 

176 for j in missing_indices: 

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

178 

179 combined_indices = indices_order + missing_indices 

180 sorted_ind = list(argsort(combined_indices)) 

181 

182 KM = KM.transpose(sorted_ind) 

183 

184 return KM 

185 

186 def gen_c_ops(self): 

187 c_ops = Qarray.from_list([]) 

188 

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

190 return c_ops 

191 

192 for operation in self.operations: 

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

194 continue 

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

196 c_ops = concatenate([c_ops, promoted_c_ops]) 

197 

198 return c_ops 

199 

200 def gen_ts(self): 

201 ts = None 

202 

203 for operation in self.operations: 

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

205 if ts is None: 

206 ts = operation.gate.ts 

207 else: 

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

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

210 ) 

211 return ts 

212 

213@struct.dataclass 

214class Circuit: 

215 register: Register 

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

217 

218 @classmethod 

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

220 if layers is None: 

221 layers = [] 

222 

223 return Circuit( 

224 register=register, 

225 layers=layers, 

226 ) 

227 

228 def append_layer(self, layer: Layer): 

229 self.layers.append(layer) 

230 

231 def append_operation( 

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

233 ): 

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

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

236 ) 

237 

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

239 

240 if new_layer: 

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

242 self.append_layer( 

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

244 ) 

245 else: 

246 if default_simulate_mode is not None: 

247 assert ( 

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

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

250 

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

252 

253 def append( 

254 self, 

255 gate: Gate, 

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

257 default_simulate_mode: Optional[SimulateMode] = None, 

258 new_layer: bool = True, 

259 ): 

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

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