赛前稳定版

src/qibotn/backends/vidal_tebd.py

@@ -50,10 +50,10 @@ def _transpose(xp, tensor, axes):
     return np.transpose(tensor, axes) if xp is np else tensor.permute(*axes)
 
 
-def _vdot_real(xp, left, right):
+def _vdot(xp, left, right):
     if xp is np:
-        return np.vdot(left.reshape(-1), right.reshape(-1)).real
-    return xp.vdot(left.reshape(-1), right.reshape(-1)).real
+        return np.vdot(left.reshape(-1), right.reshape(-1))
+    return xp.vdot(left.reshape(-1), right.reshape(-1))
 
 
 def _to_float(x):
@@ -62,6 +62,19 @@ def _to_float(x):
     return float(x)
 
 
+def _to_scalar(x):
+    if hasattr(x, "detach"):
+        return x.detach().cpu().item()
+    if isinstance(x, np.ndarray):
+        return x.item()
+    return x
+
+
+def _real_if_close(x, tol=1000):
+    value = np.real_if_close(x, tol=tol)
+    return value.item() if isinstance(value, np.ndarray) else value
+
+
 def _to_numpy(tensor):
     if hasattr(tensor, "detach"):
         return tensor.detach().cpu().numpy()
@@ -154,8 +167,8 @@ def _safe_inverse(xp, values):
 @dataclass
 class VidalTEBDExecutor:
     nqubits: int
-    max_bond: int
-    cut_ratio: float = 1e-12
+    max_bond: int | None
+    cut_ratio: float | None = 1e-12
     tensor_module: str = "torch"
 
     def __post_init__(self):
@@ -175,6 +188,7 @@ class VidalTEBDExecutor:
             _ones(self.xp, 1, self.dtype, self.device) for _ in range(self.nqubits + 1)
         ]
         self._accumulated_truncation_error = 0.0
+        self._max_truncation_error = 0.0
 
     def run_circuit(self, circuit, compile_circuit=True):
         gates = circuit.queue
@@ -189,6 +203,10 @@ class VidalTEBDExecutor:
     def truncation_error(self):
         return self._accumulated_truncation_error
 
+    @property
+    def max_truncation_error(self):
+        return self._max_truncation_error
+
     def _apply_gate(self, gate):
         sites = _gate_sites(gate)
         matrix = _asarray(self.xp, gate.matrix(), self.dtype)
@@ -232,6 +250,10 @@ class VidalTEBDExecutor:
         update = _make_two_site_update(item, umat, singvals, vh,
                                        self.max_bond, self.cut_ratio, self.xp)
         self._accumulated_truncation_error += update["truncation_error"]
+        self._max_truncation_error = max(
+            self._max_truncation_error,
+            update["truncation_error"],
+        )
         i = update["site"]
         self.gammas[i] = update["left"]
         self.gammas[i + 1] = update["right"]
@@ -252,25 +274,37 @@ class VidalTEBDExecutor:
             "Y": _asarray(self.xp, [[0, -1j], [1j, 0]], self.dtype),
             "Z": _asarray(self.xp, [[1, 0], [0, -1]], self.dtype),
         }
-        value = 0.0
+        operator_terms = [
+            (
+                coeff,
+                tuple((site, paulis[name.upper()]) for name, site in ops),
+            )
+            for coeff, ops in terms
+        ]
+        return self.expectation_operator_sum(operator_terms)
+
+    def expectation_operator_sum(self, terms):
+        value = 0.0 + 0.0j
         norm = self.norm()
         for coeff, ops in terms:
-            ops = tuple((name.upper(), site) for name, site in ops)
+            operators = {
+                int(site): _asarray(self.xp, matrix, self.dtype)
+                for site, matrix in ops
+            }
             if len(ops) == 0:
                 term_value = norm
-            elif len(ops) == 1:
-                name, site = ops[0]
-                term_value = _to_float(self._expect_one_site(site, paulis[name]))
-            elif len(ops) == 2 and abs(ops[0][1] - ops[1][1]) == 1:
-                (name0, site0), (name1, site1) = sorted(ops, key=lambda item: item[1])
-                term_value = _to_float(
-                    self._expect_adjacent(site0, paulis[name0], paulis[name1])
+            elif len(operators) == 1:
+                site, matrix = next(iter(operators.items()))
+                term_value = _to_scalar(self._expect_one_site(site, matrix))
+            elif len(operators) == 2 and abs(max(operators) - min(operators)) == 1:
+                site0, site1 = sorted(operators)
+                term_value = _to_scalar(
+                    self._expect_adjacent(site0, operators[site0], operators[site1])
                 )
             else:
-                operators = {site: paulis[name] for name, site in ops}
-                term_value = _to_float(self.expect_product_operators(operators))
-            value += float(np.real(coeff)) * term_value
-        return value / norm
+                term_value = _to_scalar(self.expect_product_operators(operators))
+            value += complex(coeff) * complex(term_value)
+        return _real_if_close(value / norm)
 
     def _expect_one_site(self, site, op):
         theta = self.xp.einsum(
@@ -280,7 +314,7 @@ class VidalTEBDExecutor:
             self.lambdas[site + 1],
         )
         op_theta = self.xp.einsum("us,asb->aub", op, theta)
-        return _vdot_real(self.xp, theta, op_theta)
+        return _vdot(self.xp, theta, op_theta)
 
     def _expect_adjacent(self, site, op_left, op_right):
         theta = self.xp.einsum(
@@ -292,7 +326,7 @@ class VidalTEBDExecutor:
             self.lambdas[site + 2],
         )
         op_theta = self.xp.einsum("us,vt,astc->auvc", op_left, op_right, theta)
-        return _vdot_real(self.xp, theta, op_theta)
+        return _vdot(self.xp, theta, op_theta)
 
     def expect_product_operators(self, operators):
         env = _asarray(self.xp, [[1.0 + 0.0j]], self.dtype)
@@ -303,10 +337,38 @@ class VidalTEBDExecutor:
             env = self.xp.einsum(
                 "xy,xsb,st,ytd->bd", env, _conj(self.xp, tensor), op, tensor
             )
-        return env.reshape(-1)[0].real
+        return env.reshape(-1)[0]
 
     def norm(self):
-        return _to_float(self.expect_product_operators({}))
+        return float(np.real(_to_scalar(self.expect_product_operators({}))))
+
+    def expectation_mpo(self, mpo_tensors):
+        """Compute ``<psi|MPO|psi> / <psi|psi>``.
+
+        MPO tensors are expected in ``(left_bond, phys_out, phys_in, right_bond)``
+        order, with physical dimension 2 on every site.
+        """
+        if len(mpo_tensors) != self.nqubits:
+            raise ValueError(
+                f"Expected {self.nqubits} MPO tensors, got {len(mpo_tensors)}."
+            )
+        env = _asarray(self.xp, [[[1.0 + 0.0j]]], self.dtype)
+        for site, raw_mpo in enumerate(mpo_tensors):
+            mpo = _asarray(self.xp, raw_mpo, self.dtype)
+            if mpo.ndim != 4 or mpo.shape[1:3] != (2, 2):
+                raise ValueError(
+                    "Each MPO tensor must have shape "
+                    "(left_bond, 2, 2, right_bond)."
+                )
+            tensor = self.gammas[site] * self.lambdas[site + 1].reshape(1, 1, -1)
+            env = self.xp.einsum(
+                "xlc,xub,lutr,ctd->brd",
+                env,
+                _conj(self.xp, tensor),
+                mpo,
+                tensor,
+            )
+        return _real_if_close(_to_scalar(env.reshape(-1)[0]) / self.norm())
 
 
 def _build_theta_svd_matrix(op, xp, lam_left, lam_mid, lam_right, gamma_left, gamma_right):
@@ -328,8 +390,8 @@ def _build_theta_svd_matrix(op, xp, lam_left, lam_mid, lam_right, gamma_left, ga
 
 def _choose_bond(singvals, max_bond, cut_ratio, xp):
     max_possible = int(singvals.shape[0])
-    keep = min(max_possible, max_bond)
-    if cut_ratio > 0 and max_possible > 0:
+    keep = max_possible if max_bond is None else min(max_possible, int(max_bond))
+    if cut_ratio is not None and cut_ratio > 0 and max_possible > 0:
         threshold = singvals[0] * cut_ratio
         if xp is np:
             ratio_keep = int(np.count_nonzero(singvals > threshold))
@@ -415,8 +477,8 @@ class _RoutedTwoQubitGate:
     name = "routed_two_qubit"
     control_qubits = ()
 
-    def __init__(self, original_gate, left_phys, right_phys):
-        self.target_qubits = (left_phys, right_phys)
+    def __init__(self, original_gate, physical_sites):
+        self.target_qubits = tuple(physical_sites)
         self._matrix = original_gate.matrix()
 
     def matrix(self):
@@ -447,8 +509,10 @@ def _route_non_adjacent_gates(gates, nqubits):
         for pos in range(right - 1, left, -1):
             routed.append(_SWAPGate(pos, pos + 1))
 
-        # Apply the original gate at the adjacent physical positions
-        routed.append(_RoutedTwoQubitGate(gate, left, left + 1))
+        # Apply the original gate in its original qubit order.  For gates like
+        # CNOT(5, 0), sorting the routed sites would swap control and target.
+        physical_map = {left: left, right: left + 1}
+        routed.append(_RoutedTwoQubitGate(gate, [physical_map[site] for site in sites]))
 
         # Reverse SWAPs to restore original ordering
         for pos in range(left + 1, right):