Fix measurement

examples/PauliSumOp/pauli_sum_op.py

@@ -44,17 +44,16 @@ def __init__(self):
         self.rz0 = tq.RZ(has_params=True, trainable=True)
         self.crx0 = tq.CRX(has_params=True, trainable=True)
 
-        self.measure = tq.MeasureMultiPauliSum(
+        self.measure = tq.MeasureMultiQubitPauliSum(
             obs_list=[
+                {'coefficient': [0.2, 0.5]},
                 {
                     "wires": [0, 2, 3, 1],
                     "observables": ["x", "y", "z", "i"],
-                    "coefficient": [1, 0.5, 0.4, 0.3],
                 },
                 {
                     "wires": [0, 2, 3, 1],
                     "observables": ["x", "x", "z", "i"],
-                    "coefficient": [1, 0.5, 0.4, 0.3],
                 },
             ]
         )

test/measurement/test_measuremultipletimes.py

@@ -0,0 +1,57 @@
+import numpy as np
+import torchquantum as tq
+
+
+def test_non_trivial_pauli_expectation():
+    class TestCircuit(tq.QuantumModule):
+        def __init__(self):
+            super().__init__()
+
+            self.meas = tq.measurement.MeasureMultipleTimes([
+                {'wires': range(2), 'observables': 'xx'},
+                {'wires': range(2), 'observables': 'yy'},
+                {'wires': range(2), 'observables': 'zz'},
+            ])
+
+        def forward(self, qdev: tq.QuantumDevice):
+            """
+                Prepare and measure the expexctation value of the state
+                    exp(-i pi/8)/sqrt(2) * (cos pi/12,
+                                            -i sin pi/12,
+                                            -i sin pi/12 * exp(i pi/4),
+                                            -i sin pi/12 * exp(i pi/4))
+            """
+            # prepare bell state
+            tq.h(qdev, 0)
+            tq.cnot(qdev, [0, 1])
+
+            # add some phases
+            tq.rz(qdev, wires=0, params=np.pi / 4)
+            tq.rx(qdev, wires=1, params=np.pi / 6)
+            return self.meas(qdev)
+
+    test_circuit = TestCircuit()
+    qdev = tq.QuantumDevice(bsz=1, n_wires=2)  # Batch size 1 for testing
+
+    # Run the circuit
+    meas_results = test_circuit(qdev)[0]
+
+    # analytical results for XX, YY, ZZ expval respectively
+    expval_xx = np.cos(np.pi / 4)
+    expval_yy = -np.cos(np.pi / 4) * np.cos(np.pi / 6)
+    expval_zz = np.cos(np.pi / 6)
+
+    atol = 1e-6
+
+    assert np.isclose(meas_results[0].item(), expval_xx, atol=atol), \
+        f"Expected {expval_xx}, got {meas_results[0].item()}"
+    assert np.isclose(meas_results[1].item(), expval_yy, atol=atol), \
+        f"Expected {expval_yy}, got {meas_results[1].item()}"
+    assert np.isclose(meas_results[2].item(), expval_zz, atol=atol), \
+        f"Expected {expval_zz}, got {meas_results[2].item()}"
+
+    print("Test passed!")
+
+
+if __name__ == "__main__":
+    test_non_trivial_pauli_expectation()

torchquantum/measurement/measurements.py

@@ -23,7 +23,6 @@
     "expval",
     "MeasureAll",
     "MeasureMultipleTimes",
-    "MeasureMultiPauliSum",
     "MeasureMultiQubitPauliSum",
     "gen_bitstrings",
     "measure",
@@ -367,15 +366,14 @@ def forward(self, qdev: tq.QuantumDevice):
 
             observables = []
             for wire in range(qdev.n_wires):
-                observables.append(tq.I())
+                observables.append("I")
 
             for wire, observable in zip(layer["wires"], layer["observables"]):
-                observables[wire] = tq.op_name_dict[observable]()
+                observables[wire] = observable
 
-            res = expval(
+            res = expval_joint_analytical(
                 qdev_new,
-                wires=list(range(qdev.n_wires)),
-                observables=observables,
+                observable="".join(observables),
             )
 
             if self.v_c_reg_mapping is not None:
@@ -390,42 +388,13 @@ def forward(self, qdev: tq.QuantumDevice):
                 res = res[:, perm]
             res_all.append(res)
 
-        return torch.cat(res_all)
+
+        return torch.stack(res_all, dim=-1)
 
     def set_v_c_reg_mapping(self, mapping):
         self.v_c_reg_mapping = mapping
 
 
-class MeasureMultiPauliSum(tq.QuantumModule):
-    """
-    similar to qiskit.opflow PauliSumOp
-    obs list:
-    list of dict: example
-    [{'wires': [0, 2, 3, 1],
-    'observables': ['x', 'y', 'z', 'i'],
-    'coefficient': [1, 0.5, 0.4, 0.3]
-    },
-    {'wires': [0, 2, 3, 1],
-    'observables': ['x', 'y', 'z', 'i'],
-    'coefficient': [1, 0.5, 0.4, 0.3]
-    },
-    ]
-    """
-
-    def __init__(self, obs_list, v_c_reg_mapping=None):
-        super().__init__()
-        self.obs_list = obs_list
-        self.v_c_reg_mapping = v_c_reg_mapping
-        self.measure_multiple_times = MeasureMultipleTimes(
-            obs_list=obs_list, v_c_reg_mapping=v_c_reg_mapping
-        )
-
-    def forward(self, qdev: tq.QuantumDevice):
-        res_all = self.measure_multiple_times(qdev).prod(-1)
-
-        return res_all.sum(-1)
-
-
 class MeasureMultiQubitPauliSum(tq.QuantumModule):
     """obs list:
     list of dict: example
@@ -449,8 +418,9 @@ def __init__(self, obs_list, v_c_reg_mapping=None):
         )
 
     def forward(self, qdev: tq.QuantumDevice):
-        res_all = self.measure_multiple_times(qdev).prod(-1)
-
+        # returns batch x len(obs_list) object, return sum times coefficient
+        res_all = self.measure_multiple_times(qdev)
+
         return (res_all * torch.tensor(self.obs_list[0]["coefficient"])).sum(-1)