Coverage for jaxquantum/circuits/gates.py: 90%

1"""Gates.""" 

2 

3from copy import deepcopy 

4from flax import struct 

5from jax import Array, config 

6from typing import List, Dict, Any, Optional, Callable, Union 

7import jax.numpy as jnp 

8 

9 

10from jaxquantum.core.qarray import Qarray, concatenate 

11 

12config.update("jax_enable_x64", True) 

13 

14 

15@struct.dataclass 

16class Gate: 

17 dims: List[int] = struct.field(pytree_node=False) 

18 _U: Optional[Array] # Unitary 

19 _Ht: Optional[Array] # Hamiltonian 

20 _KM: Optional[Qarray] # Kraus map 

21 _c_ops: Optional[Qarray] 

22 _params: Dict[str, Any] 

23 _ts: Array 

24 _name: str = struct.field(pytree_node=False) 

25 num_modes: int = struct.field(pytree_node=False) 

26 _gen_KM: Optional[Callable] = struct.field(pytree_node=False, default=None) 

27 _channel_apply: Optional[Callable] = struct.field( 

28 pytree_node=False, default=None 

29 ) 

30 

31 @classmethod 

32 def create( 

33 cls, 

34 dims: Union[int, List[int]], 

35 name: str = "Gate", 

36 params: Optional[Dict[str, Any]] = None, 

37 ts: Optional[Array] = None, 

38 gen_U: Optional[Callable[[Dict[str, Any]], Qarray]] = None, 

39 gen_Ht: Optional[Callable[[Dict[str, Any]], Qarray]] = None, 

40 gen_c_ops: Optional[Callable[[Dict[str, Any]], Qarray]] = None, 

41 gen_KM: Optional[Callable[[Dict[str, Any]], List[Qarray]]] = None, 

42 channel_apply: Optional[Callable[[Array, Dict[str, Any]], Array]] = None, 

43 lazy_kraus: bool = False, 

44 num_modes: int = 1, 

45 ): 

46 """Create a gate. 

47 

48 Args: 

49 dims: Dimensions of the gate. 

50 name: Name of the gate. 

51 params: Parameters of the gate. 

52 ts: Times of the gate. 

53 gen_U: Function to generate the unitary of the gate. 

54 gen_Ht: Function to generate a function Ht(t) that takes in a time t and outputs a Hamiltonian Qarray. 

55 gen_KM: Function to generate the Kraus map of the gate. 

56 channel_apply: Optional direct density-matrix channel kernel. 

57 lazy_kraus: Generate Kraus operators only when ``KM`` is accessed. 

58 num_modes: Number of modes of the gate. 

59 """ 

60 

61 # TODO: add params to device? 

62 

63 if isinstance(dims, int): 

64 dims = [dims] 

65 

66 assert len(dims) == num_modes, ( 

67 "Number of dimensions must match number of modes." 

68 ) 

69 

70 # Unitary 

71 _U = gen_U(params) if gen_U is not None else None 

72 _Ht = gen_Ht(params) if gen_Ht is not None else None 

73 _c_ops = gen_c_ops(params) if gen_c_ops is not None else Qarray.from_list([]) 

74 

75 _KM = gen_KM(params) if gen_KM is not None and not lazy_kraus else None 

76 

77 return Gate( 

78 dims = dims, 

79 _U = _U, 

80 _Ht = _Ht, 

81 _KM = _KM, 

82 _c_ops = _c_ops, 

83 _gen_KM=gen_KM if lazy_kraus else None, 

84 _channel_apply=channel_apply, 

85 _params = params if params is not None else {}, 

86 _ts=ts if ts is not None else jnp.array([]), 

87 _name=name, 

88 num_modes=num_modes, 

89 ) 

90 

91 def __str__(self): 

92 return self._name 

93 

94 def __repr__(self): 

95 return self._name 

96 

97 @property 

98 def name(self): 

99 return self._name 

100 

101 @property 

102 def U(self): 

103 return self._U 

104 

105 @property 

106 def Ht(self): 

107 return self._Ht 

108 

109 @property 

110 def KM(self): 

111 if self._KM is not None: 

112 return self._KM 

113 if self._gen_KM is not None: 

114 return self._gen_KM(self.params) 

115 if self._U is None: 

116 return Qarray.from_list([]) 

117 impl = type(self._U._impl).from_data(self._U.data[None]) 

118 return Qarray._from_impl(impl, self._U._qdims) 

119 

120 @property 

121 def c_ops(self): 

122 return self._c_ops 

123 

124 @property 

125 def channel_apply(self): 

126 """Direct density-matrix kernel, if defined.""" 

127 return self._channel_apply 

128 

129 @property 

130 def params(self): 

131 return self._params 

132 

133 @property 

134 def ts(self): 

135 return self._ts 

136 

137 def add_Ht(self, Ht: Callable[[float], Qarray]): 

138 """Add a Hamiltonian function to the gate.""" 

139 def new_Ht(t): 

140 return Ht(t) + self.Ht(t) if self.Ht is not None else Ht(t) 

141 

142 return Gate( 

143 dims = self.dims, 

144 _U = self.U, 

145 _Ht = new_Ht, 

146 _KM = self._KM, 

147 _c_ops = self.c_ops, 

148 _gen_KM=self._gen_KM, 

149 _channel_apply=self._channel_apply, 

150 _params = self.params, 

151 _ts = self.ts, 

152 _name = self.name, 

153 num_modes = self.num_modes, 

154 ) 

155 

156 def add_c_ops(self, c_ops: Qarray): 

157 """Add a c_ops to the gate.""" 

158 return Gate( 

159 dims = self.dims, 

160 _U = self.U, 

161 _Ht = self.Ht, 

162 _KM = self._KM, 

163 _c_ops = concatenate([self.c_ops, c_ops]), 

164 _gen_KM=self._gen_KM, 

165 _channel_apply=self._channel_apply, 

166 _params = self.params, 

167 _ts = self.ts, 

168 _name = self.name, 

169 num_modes = self.num_modes, 

170 ) 

171 

172 def copy(self): 

173 """Return a copy of the gate.""" 

174 return Gate( 

175 dims = deepcopy(self.dims), 

176 _U = self.U, 

177 _Ht = deepcopy(self.Ht), 

178 _KM = self._KM, 

179 _c_ops = self.c_ops, 

180 _gen_KM=self._gen_KM, 

181 _channel_apply=self._channel_apply, 

182 _params = deepcopy(self.params), 

183 _ts = self.ts, 

184 _name = self.name, 

185 num_modes = self.num_modes, 

186 )