Coverage for jaxquantum/codes/cat.py: 97%

1""" 

2Cat Code Qubit 

3""" 

4 

5from typing import ClassVar, Tuple 

6 

7from jaxquantum.codes.base import BosonicQubit 

8import jaxquantum as jqt 

9import jax.numpy as jnp 

10 

11from jax import config 

12 

13config.update("jax_enable_x64", True) 

14 

15 

16class CatQubit(BosonicQubit): 

17 """ 

18 Cat Qubit Class. 

19 """ 

20 

21 PARAMETERS: ClassVar[list[str]] = ["alpha", "delta"] 

22 

23 name = "cat" 

24 

25 @property 

26 def _non_device_params(self): 

27 param_list = super()._non_device_params 

28 param_list.extend(["alpha", "delta"]) 

29 return param_list 

30 

31 def _params_validation(self): 

32 super()._params_validation() 

33 if "alpha" not in self.params: 

34 self.params["alpha"] = 2 

35 if "delta" not in self.params: 

36 self.params["delta"] = 1.0 

37 if not 0 < self.params["delta"] <= 1: 

38 raise ValueError("delta must satisfy 0 < delta <= 1") 

39 

40 @classmethod 

41 def displaced_squeezed_state(cls, N, alpha, delta): 

42 """Return D(alpha) S(-log(delta))|0>.""" 

43 return ( 

44 jqt.displace(N, alpha) @ jqt.squeeze(N, -jnp.log(delta)) @ jqt.basis(N, 0) 

45 ) 

46 

47 @classmethod 

48 def cat_state(cls, N, alpha, delta, parity="even"): 

49 """Return the even or odd squeezed cat from the paper.""" 

50 if parity not in ("even", "odd"): 

51 raise ValueError("parity must be 'even' or 'odd'") 

52 plus = cls.displaced_squeezed_state(N, alpha, delta) 

53 minus = cls.displaced_squeezed_state(N, -alpha, delta) 

54 return jqt.unit(plus + (1 if parity == "even" else -1) * minus) 

55 

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

57 """Return the displaced squeezed states at +/- alpha.""" 

58 N = self.params["N"] 

59 alpha = self.params["alpha"] 

60 delta = self.params["delta"] 

61 return ( 

62 self.displaced_squeezed_state(N, alpha, delta), 

63 self.displaced_squeezed_state(N, -alpha, delta), 

64 )