Coverage for jaxquantum/circuits/library/qubit.py: 87%

1"""qubit gates.""" 

2 

3from jaxquantum.core.operators import ( 

4 identity, 

5 sigmax, 

6 sigmay, 

7 sigmaz, 

8 basis, 

9 hadamard, 

10 qubit_rotation, 

11) 

12from jaxquantum.circuits.gates import Gate 

13from jaxquantum.circuits.channels import apply_elementwise_channel 

14from jaxquantum.core.qarray import Qarray 

15import jax.numpy as jnp 

16 

17 

18def _reset_apply(rho, params): 

19 output = jnp.zeros_like(rho) 

20 return output.at[..., 0, 0].set(jnp.trace(rho, axis1=-2, axis2=-1)) 

21 

22 

23def _imperfect_reset_apply(rho, params): 

24 p_eg = jnp.asarray(params["p_eg"])[..., None] 

25 p_ee = jnp.asarray(params["p_ge"])[..., None] 

26 ground = (1 - p_eg) * rho[..., 0, 0] + (1 - p_ee) * rho[..., 1, 1] 

27 excited = p_eg * rho[..., 0, 0] + p_ee * rho[..., 1, 1] 

28 zero = jnp.zeros_like(ground) 

29 return jnp.stack( 

30 (jnp.stack((ground, zero), -1), jnp.stack((zero, excited), -1)), 

31 -2, 

32 ) 

33 

34 

35def _thermal_qubit_apply(rho, params): 

36 probability = jnp.asarray(params["err_prob"])[..., None] 

37 n_bar = jnp.asarray(params["n_bar"])[..., None] 

38 p0 = (n_bar + 1) / (2 * n_bar + 1) 

39 p1 = n_bar / (2 * n_bar + 1) 

40 coherence = jnp.sqrt(1 - probability) 

41 ground = (1 - p1 * probability) * rho[..., 0, 0] + p0 * probability * rho[..., 1, 1] 

42 excited = ( 

43 p1 * probability * rho[..., 0, 0] + (1 - p0 * probability) * rho[..., 1, 1] 

44 ) 

45 return jnp.stack( 

46 ( 

47 jnp.stack((ground, coherence * rho[..., 0, 1]), -1), 

48 jnp.stack((coherence * rho[..., 1, 0], excited), -1), 

49 ), 

50 -2, 

51 ) 

52 

53 

54def _measure_x_apply(rho, params): 

55 sign = params["_sign"] 

56 weight = 0.25 * ( 

57 rho[..., 0, 0] + rho[..., 1, 1] + sign * (rho[..., 0, 1] + rho[..., 1, 0]) 

58 ) 

59 return jnp.stack( 

60 ( 

61 jnp.stack((weight, sign * weight), -1), 

62 jnp.stack((sign * weight, weight), -1), 

63 ), 

64 -2, 

65 ) 

66 

67 

68def _dephase_x_apply(rho, params): 

69 return 0.5 * (rho + rho[..., ::-1, ::-1]) 

70 

71 

72def _dephase_z_apply(rho, params): 

73 factor = (1 - 2 * jnp.asarray(params["err_prob"]))[..., None] 

74 one = jnp.ones_like(factor) 

75 return jnp.stack( 

76 ( 

77 jnp.stack((one * rho[..., 0, 0], factor * rho[..., 0, 1]), -1), 

78 jnp.stack((factor * rho[..., 1, 0], one * rho[..., 1, 1]), -1), 

79 ), 

80 -2, 

81 ) 

82 

83 

84def X(): 

85 return Gate.create(2, name="X", gen_U=lambda params: sigmax(), num_modes=1) 

86 

87 

88def Y(): 

89 return Gate.create(2, name="Y", gen_U=lambda params: sigmay(), num_modes=1) 

90 

91 

92def Z(): 

93 return Gate.create(2, name="Z", gen_U=lambda params: sigmaz(), num_modes=1) 

94 

95 

96def H(): 

97 return Gate.create(2, name="H", gen_U=lambda params: hadamard(), num_modes=1) 

98 

99 

100def Rx(theta, ts=None): 

101 gen_Ht = None 

102 if ts is not None: 

103 delta_t = ts[-1] - ts[0] 

104 amp = theta / delta_t 

105 gen_Ht = lambda params: (lambda t: amp / 2 * sigmax()) 

106 

107 return Gate.create( 

108 2, 

109 name="Rx", 

110 params={"theta": theta}, 

111 gen_U=lambda params: qubit_rotation(params["theta"], 1, 0, 0), 

112 gen_Ht=gen_Ht, 

113 ts=ts, 

114 num_modes=1, 

115 ) 

116 

117 

118def Ry(theta, ts=None): 

119 gen_Ht = None 

120 if ts is not None: 

121 delta_t = ts[-1] - ts[0] 

122 amp = theta / delta_t 

123 gen_Ht = lambda params: (lambda t: amp / 2 * sigmay()) 

124 return Gate.create( 

125 2, 

126 name="Ry", 

127 params={"theta": theta}, 

128 gen_U=lambda params: qubit_rotation(params["theta"], 0, 1, 0), 

129 gen_Ht=gen_Ht, 

130 ts=ts, 

131 num_modes=1, 

132 ) 

133 

134 

135def Rz(theta, ts=None): 

136 gen_Ht = None 

137 if ts is not None: 

138 delta_t = ts[-1] - ts[0] 

139 amp = theta / delta_t 

140 gen_Ht = lambda params: (lambda t: amp / 2 * sigmaz()) 

141 return Gate.create( 

142 2, 

143 name="Rz", 

144 params={"theta": theta}, 

145 gen_U=lambda params: qubit_rotation(params["theta"], 0, 0, 1), 

146 gen_Ht=gen_Ht, 

147 ts=ts, 

148 num_modes=1, 

149 ) 

150 

151 

152def MZ(measure=None): 

153 if measure is None: 

154 gate_name = "MZ" 

155 factor = jnp.eye(2) 

156 elif measure == +1: 

157 gate_name = "MZ_plus" 

158 factor = jnp.array([[1, 0], [0, 0]]) 

159 elif measure == -1: 

160 gate_name = "MZ_minus" 

161 factor = jnp.array([[0, 0], [0, 1]]) 

162 else: 

163 raise ValueError("measure should be None, +1 or -1") 

164 

165 def kmap(params): 

166 g, e = basis(2, 0), basis(2, 1) 

167 if measure is None: 

168 return Qarray.from_list([g @ g.dag(), e @ e.dag()]) 

169 state = g if measure == 1 else e 

170 return Qarray.from_list([state @ state.dag()]) 

171 

172 return Gate.create( 

173 2, 

174 name=gate_name, 

175 params={"_factor": factor}, 

176 gen_KM=kmap, 

177 channel_apply=apply_elementwise_channel, 

178 lazy_kraus=True, 

179 num_modes=1, 

180 ) 

181 

182 

183def MX(measure=None): 

184 if measure is None: 

185 gate_name = "MX" 

186 channel_apply = _dephase_x_apply 

187 sign = 0 

188 elif measure == +1: 

189 gate_name = "MX_plus" 

190 channel_apply = _measure_x_apply 

191 sign = 1 

192 elif measure == -1: 

193 gate_name = "MX_minus" 

194 channel_apply = _measure_x_apply 

195 sign = -1 

196 else: 

197 raise ValueError("measure should be None, +1 or -1") 

198 

199 def kmap(params): 

200 g, e = basis(2, 0), basis(2, 1) 

201 plus, minus = (g + e).unit(), (g - e).unit() 

202 if measure is None: 

203 return Qarray.from_list([plus @ plus.dag(), minus @ minus.dag()]) 

204 state = plus if measure == 1 else minus 

205 return Qarray.from_list([state @ state.dag()]) 

206 

207 return Gate.create( 

208 2, 

209 name=gate_name, 

210 params={"_sign": sign}, 

211 gen_KM=kmap, 

212 channel_apply=channel_apply, 

213 lazy_kraus=True, 

214 num_modes=1, 

215 ) 

216 

217 

218def Reset(): 

219 def kmap(params): 

220 g, e = basis(2, 0), basis(2, 1) 

221 return Qarray.from_list([g @ g.dag(), g @ e.dag()]) 

222 

223 return Gate.create( 

224 2, 

225 name="Reset", 

226 gen_KM=kmap, 

227 channel_apply=_reset_apply, 

228 lazy_kraus=True, 

229 num_modes=1, 

230 ) 

231 

232 

233def IP_Reset(p_eg, p_ee): 

234 def kmap(params): 

235 g, e = basis(2, 0), basis(2, 1) 

236 gg, ge = g @ g.dag(), g @ e.dag() 

237 eg, ee = e @ g.dag(), e @ e.dag() 

238 return Qarray.from_list( 

239 [ 

240 jnp.sqrt(1 - p_eg) * gg, 

241 jnp.sqrt(p_ee) * ee, 

242 jnp.sqrt(p_eg) * eg, 

243 jnp.sqrt(1 - p_ee) * ge, 

244 ] 

245 ) 

246 

247 return Gate.create( 

248 2, 

249 name="IP_Reset", 

250 params={"p_eg": p_eg, "p_ge": p_ee}, 

251 gen_KM=kmap, 

252 channel_apply=_imperfect_reset_apply, 

253 lazy_kraus=True, 

254 num_modes=1, 

255 ) 

256 

257 

258def CX(): 

259 g = basis(2, 0) 

260 e = basis(2, 1) 

261 

262 gg = g @ g.dag() 

263 ee = e @ e.dag() 

264 

265 op = (gg ^ identity(2)) + (ee ^ sigmax()) 

266 

267 return Gate.create([2, 2], name="CX", gen_U=lambda params: op, num_modes=2) 

268 

269 

270def _Thermal_Kraus_Ops_Qb(err_prob, n_bar): 

271 """Return generalized amplitude-damping Kraus operators for a qubit.""" 

272 err_prob = jnp.asarray(err_prob) 

273 zero = jnp.zeros_like(err_prob) 

274 one = jnp.ones_like(err_prob) 

275 normalization = 2 * n_bar + 1 

276 sqrt_p0 = jnp.sqrt((n_bar + 1) / normalization)[..., None, None] 

277 sqrt_p1 = jnp.sqrt(n_bar / normalization)[..., None, None] 

278 retained = jnp.sqrt(1 - err_prob) 

279 lost = jnp.sqrt(err_prob) 

280 

281 def matrix(a, b, c, d): 

282 return jnp.stack( 

283 (jnp.stack((a, b), -1), jnp.stack((c, d), -1)), 

284 -2, 

285 ) 

286 

287 return [ 

288 Qarray.create(sqrt_p0 * matrix(one, zero, zero, retained)), 

289 Qarray.create(sqrt_p0 * matrix(zero, lost, zero, zero)), 

290 Qarray.create(sqrt_p1 * matrix(zero, zero, lost, zero)), 

291 Qarray.create(sqrt_p1 * matrix(retained, zero, zero, one)), 

292 ] 

293 

294 

295def Thermal_Ch_Qb(err_prob, n_bar): 

296 kmap = lambda params: Qarray.from_list(_Thermal_Kraus_Ops_Qb(err_prob, n_bar)) 

297 return Gate.create( 

298 2, 

299 name="Thermal_Ch_Qb", 

300 params={"err_prob": err_prob, "n_bar": n_bar}, 

301 gen_KM=kmap, 

302 channel_apply=_thermal_qubit_apply, 

303 lazy_kraus=True, 

304 num_modes=1, 

305 ) 

306 

307 

308def _Pure_Dephasing_Ops_Qb(err_prob): 

309 """ " Returns the Kraus Operators for a thermal channel with probability 

310 err_prob and average photon number n_bar in a Hilbert Space of size 2""" 

311 return [jnp.sqrt(1 - err_prob) * identity(2), jnp.sqrt(err_prob) * sigmaz()] 

312 

313 

314def Dephasing_Ch_Qb(err_prob): 

315 kmap = lambda params: Qarray.from_list(_Pure_Dephasing_Ops_Qb(err_prob)) 

316 return Gate.create( 

317 2, 

318 name="Dephasing_Ch_Qb", 

319 params={"err_prob": err_prob}, 

320 gen_KM=kmap, 

321 channel_apply=_dephase_z_apply, 

322 lazy_kraus=True, 

323 num_modes=1, 

324 )