Coverage for jaxquantum/circuits/simulate.py: 0%

1"""Circuit simulation methods.""" 

2 

3from flax import struct 

4from jax import config 

5from typing import List 

6 

7 

8from jaxquantum.core.qarray import Qarray, ket2dm 

9from jaxquantum.circuits.circuits import Circuit, Layer 

10from jaxquantum.circuits.constants import SimulateMode 

11 

12config.update("jax_enable_x64", True) 

13 

14 

15@struct.dataclass 

16class Results: 

17 results: List[Qarray] = struct.field(pytree_node=False) 

18 

19 @classmethod 

20 def create(cls, results: List[Qarray]): 

21 return Results(results=results) 

22 

23 def __getitem__(self, j: int): 

24 return self.results[j] 

25 

26 def __str__(self): 

27 return self.__repr__() 

28 

29 def __repr__(self): 

30 return str(self.results) 

31 

32 def append(self, result: Qarray): 

33 self.results.append(result) 

34 

35 def __len__(self): 

36 return len(self.results) 

37 

38 

39def simulate( 

40 circuit: Circuit, initial_state: Qarray, mode: SimulateMode = SimulateMode.UNITARY 

41) -> Results: 

42 """ 

43 Simulates the evolution of a quantum state through a given quantum circuit. 

44 

45 Args: 

46 circuit (Circuit): The quantum circuit to simulate. The circuit is composed of layers, 

47 each of which can generate unitary or Kraus operators. 

48 initial_state (Qarray): The initial quantum state to be evolved. This can be a state vector 

49 or a density matrix. 

50 mode (SimulateMode, optional): The mode of simulation. It can be either SimulateMode.UNITARY 

51 for unitary evolution or SimulateMode.KRAUS for Kraus operator 

52 evolution. Defaults to SimulateMode.UNITARY. 

53 

54 Returns: 

55 Results: An object containing the results of the simulation, which includes the quantum states 

56 at each step of the circuit. 

57 """ 

58 

59 results = Results.create([]) 

60 state = initial_state 

61 results.append(Qarray.from_list([state])) 

62 

63 for layer in circuit.layers: 

64 result = simulate_layer(layer, state, mode=mode) 

65 results.append(result) 

66 state = result[-1] 

67 

68 return results 

69 

70 

71def simulate_layer( 

72 layer: Layer, initial_state: Qarray, mode: SimulateMode = SimulateMode.UNITARY 

73) -> Qarray: 

74 """ 

75 Simulates the evolution of a quantum state through a given layer. 

76 

77 Args: 

78 layer (Layer): The layer through which the quantum state evolves. 

79 This layer should have methods to generate unitary (gen_U) 

80 and Kraus (gen_KM) operators. 

81 initial_state (Qarray): The initial quantum state to be evolved. 

82 This can be a state vector or a density matrix. 

83 mode (SimulateMode, optional): The mode of simulation. It can be either 

84 SimulateMode.UNITARY for unitary evolution 

85 or SimulateMode.KRAUS for Kraus operator evolution 

86 or SimulateMode.DEFAULT to use the default simulate mode in the layer. 

87 Defaults to SimulateMode.UNITARY. 

88 Returns: 

89 Qarray: The result of the simulation containing the evolved quantum state. 

90 """ 

91 

92 state = initial_state 

93 

94 if mode == SimulateMode.DEFAULT: 

95 mode = layer._default_simulate_mode 

96 

97 if mode == SimulateMode.UNITARY: 

98 U = layer.gen_U() 

99 if state.is_dm(): 

100 state = U @ state @ U.dag() 

101 else: 

102 state = U @ state 

103 

104 result = Qarray.from_list([state]) 

105 

106 elif mode == SimulateMode.KRAUS: 

107 KM = layer.gen_KM() 

108 

109 state = ket2dm(state) 

110 state = (KM @ state @ KM.dag()).collapse() 

111 

112 # new_state = 0 

113 # for op_j in range(len(KM)): 

114 # op = KM[op_j] 

115 # new_state += op @ state @ op.dag() 

116 # state = new_state 

117 

118 result = Qarray.from_list([state]) 

119 

120 return result