[Fix] Fix the dimension bug of expeval of density matrix.

Author: Zhuoyang Ye

examples/mnist/mnist_noise.py

@@ -82,7 +82,7 @@ def __init__(self):
         self.encoder = tq.GeneralEncoder(tq.encoder_op_list_name_dict["4x4_u3_h_rx"])
 
         self.q_layer = self.QLayer()
-        self.measure = tq.MeasureAll(tq.PauliZ)
+        self.measure = tq.MeasureAll_Density(tq.PauliZ)
 
     def forward(self, x, use_qiskit=False):
         qdev = tq.NoiseDevice(

torchquantum/device/noisedevices.py

@@ -73,12 +73,10 @@ def __init__(
         self.record_op = record_op
         self.op_history = []
 
-
     def reset_op_history(self):
         """Resets the all Operation of the quantum device"""
         self.op_history = []
 
-
     def print_2d(self, index):
         """Print the matrix value at the given index.
 
@@ -125,12 +123,16 @@ def get_probs_1d(self):
         """Return the states in a 1d tensor."""
         bsz = self.densities.shape[0]
         densities2d = torch.reshape(self.densities, [bsz, 2 ** self.n_wires, 2 ** self.n_wires])
-        return torch.diagonal(densities2d, offset=0, dim1=1, dim2=2)
+        return torch.abs(torch.diagonal(densities2d, offset=0, dim1=1, dim2=2))
 
     def get_prob_1d(self):
         """Return the state in a 1d tensor."""
         density2d = torch.reshape(self.density, [2 ** self.n_wires, 2 ** self.n_wires])
-        return torch.diagonal(density2d, offset=0, dim1=0, dim2=1)
+        return torch.abs(torch.diagonal(density2d, offset=0, dim1=0, dim2=1))
+
+    def clone_densities(self, existing_densities: torch.Tensor):
+        """Clone the densities of the other quantum device."""
+        self.densities = existing_densities.clone()
 
 
 for func_name, func in func_name_dict.items():

torchquantum/measurement/density_measurements.py

@@ -14,15 +14,16 @@
 import torchquantum.operator as op
 from copy import deepcopy
 import matplotlib.pyplot as plt
-from measurements import gen_bitstrings
-from measurements import find_observable_groups
+from .measurements import gen_bitstrings
+from .measurements import find_observable_groups
 
 __all__ = [
     "expval_joint_sampling_grouping",
     "expval_joint_analytical",
     "expval_joint_sampling",
     "expval",
     "measure",
+    "MeasureAll_Density"
 ]
 
 
@@ -194,7 +195,7 @@ def expval(
         wires: Union[int, List[int]],
         observables: Union[op.Observable, List[op.Observable]],
 ):
-    all_dims = np.arange(noisedev.densities.dim())
+    all_dims = np.arange(noisedev.n_wires+1)
     if isinstance(wires, int):
         wires = [wires]
         observables = [observables]
@@ -206,7 +207,8 @@ def expval(
 
     # compute magnitude
     state_mag = noisedev.get_probs_1d()
-
+    bsz = state_mag.shape[0]
+    state_mag = torch.reshape(state_mag, [bsz] + [2] * noisedev.n_wires)
     expectations = []
     for wire, observable in zip(wires, observables):
         # compute marginal magnitude
@@ -221,7 +223,7 @@ def expval(
     return torch.stack(expectations, dim=-1)
 
 
-class MeasureAll(tq.QuantumModule):
+class MeasureAll_Density(tq.QuantumModule):
     """Obtain the expectation value of all the qubits."""
 
     def __init__(self, obs, v_c_reg_mapping=None):
@@ -265,11 +267,7 @@ def set_v_c_reg_mapping(self, mapping):
     # measure the state on z basis
     print(tq.measure(qdev, n_shots=1024))
 
-
-
-    '''
     # obtain the expval on a observable
     expval = expval_joint_sampling(qdev, 'II', 100000)
-    expval_ana = expval_joint_analytical(qdev, 'II')
-    print(expval, expval_ana)
-    '''
+    # expval_ana = expval_joint_analytical(qdev, 'II')
+    print(expval)