Coverage for jaxquantum / core / operators.py: 87%

1"""States.""" 

2 

3from typing import List 

4from jax import config 

5from math import prod 

6 

7import jax.numpy as jnp 

8from jax.nn import one_hot 

9 

10from jaxquantum.core.qarray import Qarray, tensor 

11 

12config.update("jax_enable_x64", True) 

13 

14 

15def sigmax() -> Qarray: 

16 """σx 

17 

18 Returns: 

19 σx Pauli Operator 

20 """ 

21 return Qarray.create(jnp.array([[0.0, 1.0], [1.0, 0.0]])) 

22 

23 

24def sigmay() -> Qarray: 

25 """σy 

26 

27 Returns: 

28 σy Pauli Operator 

29 """ 

30 return Qarray.create(jnp.array([[0.0, -1.0j], [1.0j, 0.0]])) 

31 

32 

33def sigmaz() -> Qarray: 

34 """σz 

35 

36 Returns: 

37 σz Pauli Operator 

38 """ 

39 return Qarray.create(jnp.array([[1.0, 0.0], [0.0, -1.0]])) 

40 

41 

42def hadamard() -> Qarray: 

43 """H 

44 

45 Returns: 

46 H: Hadamard gate 

47 """ 

48 return Qarray.create(jnp.array([[1, 1], [1, -1]]) / jnp.sqrt(2)) 

49 

50 

51def sigmam() -> Qarray: 

52 """σ- 

53 

54 Returns: 

55 σ- Pauli Operator 

56 """ 

57 return Qarray.create(jnp.array([[0.0, 0.0], [1.0, 0.0]])) 

58 

59 

60def sigmap() -> Qarray: 

61 """σ+ 

62 

63 Returns: 

64 σ+ Pauli Operator 

65 """ 

66 return Qarray.create(jnp.array([[0.0, 1.0], [0.0, 0.0]])) 

67 

68 

69def qubit_rotation(theta: float, nx, ny, nz) -> Qarray: 

70 """Single qubit rotation. 

71 

72 Args: 

73 theta: rotation angle. 

74 nx: rotation axis x component. 

75 ny: rotation axis y component. 

76 nz: rotation axis z component. 

77 

78 Returns: 

79 Single qubit rotation operator. 

80 """ 

81 return jnp.cos(theta / 2) * identity(2) - 1j * jnp.sin(theta / 2) * ( 

82 nx * sigmax() + ny * sigmay() + nz * sigmaz() 

83 ) 

84 

85 

86def destroy(N) -> Qarray: 

87 """annihilation operator 

88 

89 Args: 

90 N: Hilbert space size 

91 

92 Returns: 

93 annilation operator in Hilber Space of size N 

94 """ 

95 return Qarray.create(jnp.diag(jnp.sqrt(jnp.arange(1, N)), k=1)) 

96 

97 

98def create(N) -> Qarray: 

99 """creation operator 

100 

101 Args: 

102 N: Hilbert space size 

103 

104 Returns: 

105 creation operator in Hilber Space of size N 

106 """ 

107 return Qarray.create(jnp.diag(jnp.sqrt(jnp.arange(1, N)), k=-1)) 

108 

109 

110def num(N) -> Qarray: 

111 """Number operator 

112 

113 Args: 

114 N: Hilbert Space size 

115 

116 Returns: 

117 number operator in Hilber Space of size N 

118 """ 

119 return Qarray.create(jnp.diag(jnp.arange(N))) 

120 

121 

122def identity(*args, **kwargs) -> Qarray: 

123 """Identity matrix. 

124 

125 Returns: 

126 Identity matrix. 

127 """ 

128 return Qarray.create(jnp.eye(*args, **kwargs)) 

129 

130 

131def identity_like(A) -> Qarray: 

132 """Identity matrix with the same shape as A. 

133 

134 Args: 

135 A: Matrix. 

136 

137 Returns: 

138 Identity matrix with the same shape as A. 

139 """ 

140 space_dims = A.space_dims 

141 total_dim = prod(space_dims) 

142 return Qarray.create(jnp.eye(total_dim, total_dim), dims=[space_dims, space_dims]) 

143 

144 

145def displace(N, α) -> Qarray: 

146 """Displacement operator 

147 

148 Args: 

149 N: Hilbert Space Size 

150 α: Phase space displacement 

151 

152 Returns: 

153 Displace operator D(α) 

154 """ 

155 a = destroy(N) 

156 return (α * a.dag() - jnp.conj(α) * a).expm() 

157 

158def squeeze(N, z): 

159 """Single-mode Squeezing operator. 

160 

161 

162 Args: 

163 N: Hilbert Space Size 

164 z: squeezing parameter 

165 

166 Returns: 

167 Sqeezing operator 

168 """ 

169 

170 a = destroy(N) 

171 op = (1 / 2.0) * jnp.conj(z) * (a @ a) - (1 / 2.0) * z * (a.dag() @ a.dag()) 

172 return op.expm() 

173 

174 

175def squeezing_linear_to_dB(z): 

176 return 20 * jnp.log10(jnp.exp(jnp.abs(z))) 

177 

178def squeezing_dB_to_linear(z_dB): 

179 return jnp.log(10**(z_dB/20)) 

180 

181# States --------------------------------------------------------------------- 

182 

183 

184def basis(N: int, k: int): 

185 """Creates a |k> (i.e. fock state) ket in a specified Hilbert Space. 

186 

187 Args: 

188 N: Hilbert space dimension 

189 k: fock number 

190 

191 Returns: 

192 Fock State |k> 

193 """ 

194 return Qarray.create(one_hot(k, N).reshape(N, 1)) 

195 

196def multi_mode_basis_set(Ns: List[int]) -> Qarray: 

197 """Creates a multi-mode basis set. 

198 

199 Args: 

200 Ns: List of Hilbert space dimensions for each mode. 

201 

202 Returns: 

203 Multi-mode basis set. 

204 """ 

205 data = jnp.eye(prod(Ns)) 

206 dims = (tuple(Ns), tuple([1 for _ in Ns])) 

207 return Qarray.create(data, dims=dims, bdims=(prod(Ns),)) 

208 

209 

210def coherent(N: int, α: complex) -> Qarray: 

211 """Coherent state. 

212 

213 Args: 

214 N: Hilbert Space Size. 

215 α: coherent state amplitude. 

216 

217 Return: 

218 Coherent state |α⟩. 

219 """ 

220 return displace(N, α) @ basis(N, 0) 

221 

222 

223def thermal_dm(N: int, n: float) -> Qarray: 

224 """Thermal state. 

225 

226 Args: 

227 N: Hilbert Space Size. 

228 n: average photon number. 

229 

230 Return: 

231 Thermal state. 

232 """ 

233 

234 beta = jnp.log(1 + 1 / n) 

235 

236 return Qarray.create( 

237 jnp.where( 

238 jnp.isposinf(beta), 

239 basis(N, 0).to_dm().data, 

240 jnp.diag(jnp.exp(-beta * jnp.linspace(0, N - 1, N))), 

241 ) 

242 ).unit() 

243 

244 

245def basis_like(A: Qarray, ks: List[int]) -> Qarray: 

246 """Creates a |k> (i.e. fock state) ket with the same space dims as A. 

247 

248 Args: 

249 A: state or operator. 

250 k: fock number. 

251 

252 Returns: 

253 Fock State |k> with the same space dims as A. 

254 """ 

255 space_dims = A.space_dims 

256 assert len(space_dims) == len(ks), "len(ks) must be equal to len(space_dims)" 

257 

258 kets = [] 

259 for j, k in enumerate(ks): 

260 kets.append(basis(space_dims[j], k)) 

261 return tensor(*kets)