Coverage for jaxquantum/codes/base.py: 85%

1""" 

2Base Bosonic Qubit Class 

3""" 

4 

5from typing import ClassVar, Dict, Optional, Tuple 

6from abc import abstractmethod, ABCMeta 

7 

8from jaxquantum.utils.utils import device_put_params 

9import jaxquantum as jqt 

10 

11from jax import config 

12import numpy as np 

13import jax.numpy as jnp 

14import matplotlib.pyplot as plt 

15 

16config.update("jax_enable_x64", True) 

17 

18 

19class BosonicQubit(metaclass=ABCMeta): 

20 """ 

21 Base class for Bosonic Qubits. 

22 """ 

23 

24 BASE_PARAMETERS: ClassVar[list[str]] = ["N"] 

25 PARAMETERS: ClassVar[list[str]] = [] 

26 

27 name = "bqubit" 

28 

29 @property 

30 def _non_device_params(self): 

31 """ 

32 Can be overriden in child classes. 

33 """ 

34 return ["N"] 

35 

36 def __init__( 

37 self, params: Optional[Dict[str, float]] = None, name: Optional[str] = None 

38 ): 

39 if name is not None: 

40 self.name = name 

41 

42 self.params = params if params else {} 

43 self._params_validation() 

44 

45 self.params = device_put_params(self.params, self._non_device_params) 

46 

47 self.common_gates: Dict[str, jqt.Qarray] = {} 

48 self._gen_common_gates() 

49 

50 self.wigner_pts = jnp.linspace(-4.5, 4.5, 61) 

51 

52 self.basis = self._get_basis_states() 

53 

54 for basis_state in ["+x", "-x", "+y", "-y", "+z", "-z"]: 

55 assert basis_state in self.basis, ( 

56 f"Please set the {basis_state} basis state." 

57 ) 

58 

59 def _params_validation(self): 

60 """ 

61 Override this method to add additional validation to params. 

62 

63 E.g. 

64 if "N" not in self.params: 

65 self.params["N"] = 50 

66 """ 

67 

68 for key in self.params: 

69 if key not in self.BASE_PARAMETERS + self.PARAMETERS: 

70 raise ValueError( 

71 f"Invalid parameter {key}. Allowed parameters are {self.BASE_PARAMETERS + self.PARAMETERS}" 

72 ) 

73 

74 if "N" not in self.params: 

75 self.params["N"] = 50 

76 

77 def _gen_common_gates(self): 

78 """ 

79 Override this method to add additional common gates. 

80 

81 E.g. 

82 if "N" not in self.params: 

83 self.params["N"] = 50 

84 """ 

85 N = self.params["N"] 

86 self.common_gates["a_dag"] = jqt.create(N) 

87 self.common_gates["a"] = jqt.destroy(N) 

88 

89 @abstractmethod 

90 def _get_basis_z(self) -> Tuple[jqt.Qarray, jqt.Qarray]: 

91 """ 

92 Returns: 

93 plus_z (jqt.Qarray), minus_z (jqt.Qarray): z basis states 

94 """ 

95 

96 def _get_basis_states(self) -> Dict[str, jqt.Qarray]: 

97 """ 

98 Construct basis states |+-x>, |+-y>, |+-z> 

99 """ 

100 plus_z, minus_z = self._get_basis_z() 

101 return self._gen_basis_states_from_z(plus_z, minus_z) 

102 

103 def _gen_basis_states_from_z( 

104 self, plus_z: jqt.Qarray, minus_z: jqt.Qarray 

105 ) -> Dict[str, jqt.Qarray]: 

106 """ 

107 Construct basis states |+-x>, |+-y>, |+-z> from |+-z> 

108 """ 

109 basis: Dict[str, jqt.Qarray] = {} 

110 

111 # import to make sure that each basis state is a column vec 

112 # otherwise, transposing a 1D vector will do nothing 

113 

114 basis["+z"] = plus_z 

115 basis["-z"] = minus_z 

116 

117 basis["+x"] = jqt.unit(basis["+z"] + basis["-z"]) 

118 basis["-x"] = jqt.unit(basis["+z"] - basis["-z"]) 

119 basis["+y"] = jqt.unit(basis["+z"] + 1j * basis["-z"]) 

120 basis["-y"] = jqt.unit(basis["+z"] - 1j * basis["-z"]) 

121 return basis 

122 

123 def jqt2qt(self, state): 

124 return jqt.jqt2qt(state) 

125 

126 # gates 

127 # ====================================================== 

128 # @abstractmethod 

129 # def stabilize(self) -> None: 

130 # """ 

131 # Stabilizing/measuring syndromes. 

132 # """ 

133 

134 @property 

135 def x_U(self) -> jqt.Qarray: 

136 """ 

137 Logical X unitary gate. 

138 """ 

139 return self._gen_pauli_U("x") 

140 

141 @property 

142 def x_H(self) -> Optional[jqt.Qarray]: 

143 """ 

144 Logical X hamiltonian. 

145 """ 

146 return None 

147 

148 @property 

149 def y_U(self) -> jqt.Qarray: 

150 """ 

151 Logical Y unitary gate. 

152 """ 

153 return self._gen_pauli_U("y") 

154 

155 @property 

156 def y_H(self) -> Optional[jqt.Qarray]: 

157 """ 

158 Logical Y hamiltonian. 

159 """ 

160 return None 

161 

162 @property 

163 def z_U(self) -> jqt.Qarray: 

164 """ 

165 Logical Z unitary gate. 

166 """ 

167 return self._gen_pauli_U("z") 

168 

169 @property 

170 def z_H(self) -> Optional[jqt.Qarray]: 

171 """ 

172 Logical Z hamiltonian. 

173 """ 

174 return None 

175 

176 @property 

177 def h_H(self) -> Optional[jqt.Qarray]: 

178 """ 

179 Logical Hadamard hamiltonian. 

180 """ 

181 return None 

182 

183 @property 

184 def h_U(self) -> jqt.Qarray: 

185 """ 

186 Logical Hadamard unitary gate. 

187 """ 

188 return ( 

189 self.basis["+x"] @ self.basis["+z"].dag() 

190 + self.basis["-x"] @ self.basis["-z"].dag() 

191 ) 

192 

193 def _gen_pauli_U(self, basis_state: str) -> jqt.Qarray: 

194 """ 

195 Generates unitary for Pauli X, Y, Z. 

196 

197 Args: 

198 basis_state (str): "x", "y", "z" 

199 

200 Returns: 

201 U (jqt.Qarray): Pauli unitary 

202 """ 

203 H = getattr(self, basis_state + "_H") 

204 if H is not None: 

205 return jqt.expm(1.0j * H) 

206 

207 gate = ( 

208 self.basis["+" + basis_state] @ self.basis["+" + basis_state].dag() 

209 - self.basis["-" + basis_state] @ self.basis["-" + basis_state].dag() 

210 ) 

211 

212 return gate 

213 

214 @property 

215 def projector(self): 

216 return ( 

217 self.basis["+z"] @ self.basis["+z"].dag() 

218 + self.basis["-z"] @ self.basis["-z"].dag() 

219 ) 

220 

221 @property 

222 def maximally_mixed_state(self): 

223 return (1 / 2.0) * self.projector 

224 

225 # Plotting 

226 # ====================================================== 

227 def _prepare_state_plot(self, state): 

228 """ 

229 Can be overriden. 

230 

231 E.g. in the case of cavity x transmon system 

232 return qt.ptrace(state, 0) 

233 """ 

234 return state 

235 

236 def plot(self, state, ax=None, qp_type=jqt.WIGNER, **kwargs) -> None: 

237 if ax is None: 

238 fig, ax = plt.subplots(1, figsize=(4, 3), dpi=200) 

239 fig = ax.get_figure() 

240 

241 if qp_type == jqt.WIGNER: 

242 vmin = -1 

243 vmax = 1 

244 elif qp_type == jqt.QFUNC: 

245 vmin = 0 

246 vmax = 1 

247 

248 w_plt = self._plot_single(state, ax=ax, qp_type=qp_type, **kwargs) 

249 

250 ax.set_title(qp_type.capitalize() + " Quasi-Probability Dist.") 

251 ticks = np.linspace(vmin, vmax, 5) 

252 fig.colorbar(w_plt, ax=ax, ticks=ticks) 

253 ax.set_xlabel(r"Re$(\alpha)$") 

254 ax.set_ylabel(r"Im$(\alpha)$") 

255 fig.tight_layout() 

256 

257 plt.show() 

258 

259 def _plot_single(self, state, ax=None, contour=True, qp_type=jqt.WIGNER): 

260 """ 

261 Assumes state has same dims as initial_state. 

262 """ 

263 

264 if ax is None: 

265 _, ax = plt.subplots(1, figsize=(4, 3), dpi=200) 

266 

267 return jqt.plot_qp( 

268 state, self.wigner_pts, axs=ax, contour=contour, qp_type=qp_type 

269 ) 

270 

271 def plot_code_states(self, qp_type: str = jqt.WIGNER, **kwargs): 

272 """ 

273 Plot |±x⟩, |±y⟩, |±z⟩ code states. 

274 

275 Args: 

276 qp_type (str): 

277 WIGNER or QFUNC 

278 

279 Return: 

280 axs: Axes 

281 """ 

282 fig, axs = plt.subplots(2, 3, figsize=(12, 6), dpi=200) 

283 

284 for i, label in enumerate(["+z", "+x", "+y", "-z", "-x", "-y"]): 

285 state = self._prepare_state_plot(self.basis[label]) 

286 pos = (i // 3, i % 3) 

287 ax = axs[pos] 

288 self._plot_single(state, ax=ax, qp_type=qp_type, **kwargs) 

289 ax.set_title(f"|{label}" + r"$\rangle$") 

290 ax.set_xlabel(r"Re[$\alpha$]") 

291 ax.set_ylabel(r"Im[$\alpha$]") 

292 

293 fig.suptitle(self.name) 

294 fig.align_xlabels(axs) 

295 fig.align_ylabels(axs)