完善mps的vidal机制,多节点并行;补充tn搜索时dask集群搜索的方式

src/qibotn/backends/quimb.py

@@ -37,6 +37,8 @@ GATE_MAP = {
     "measure": "measure",
 }
 
+PAULI_DENSE_MAX_QUBITS = 8
+
 
 def _torch_cpu_array(data, dtype=None):
     """Convert array-like data to a contiguous CPU torch tensor."""
@@ -68,6 +70,81 @@ def _arrays_to_backend(arrays, backend, engine):
     return [engine.asarray(array) for array in arrays]
 
 
+def _pauli_term_to_dense_operator(factors):
+    op = None
+    where = []
+    for qubit, gate_name in factors:
+        pauli = qu.pauli(gate_name.lower())
+        op = pauli if op is None else op & pauli
+        where.append(qubit)
+    return op, tuple(where)
+
+
+def pauli_product_expectation_tn(
+    quimb_circuit,
+    factors,
+    simplify_sequence="ADCRS",
+    simplify_atol=1e-12,
+    simplify_equalize_norms=True,
+):
+    """Build the scalar TN for ``<psi|P|psi>`` without dense Pauli strings."""
+    import numpy as np
+
+    op_by_site = {
+        int(qubit): qu.pauli(str(gate_name).lower())
+        for qubit, gate_name in factors
+        if str(gate_name).upper() != "I"
+    }
+    ket = quimb_circuit.get_psi_simplified(
+        seq=simplify_sequence,
+        atol=simplify_atol,
+        equalize_norms=simplify_equalize_norms,
+    )
+    bra = ket.conj().reindex(
+        {
+            quimb_circuit.ket_site_ind(qubit): quimb_circuit.bra_site_ind(qubit)
+            for qubit in range(quimb_circuit.N)
+        }
+    )
+
+    tn = bra | ket
+    identity = np.eye(2, dtype=complex)
+    for qubit in range(quimb_circuit.N):
+        data = op_by_site.get(qubit, identity)
+        tn |= qtn.Tensor(
+            data=data,
+            inds=(
+                quimb_circuit.bra_site_ind(qubit),
+                quimb_circuit.ket_site_ind(qubit),
+            ),
+        )
+
+    tn.full_simplify_(
+        output_inds=(),
+        seq=simplify_sequence,
+        atol=simplify_atol,
+        equalize_norms=simplify_equalize_norms,
+    )
+    return tn
+
+
+def pauli_product_expectation(
+    quimb_circuit,
+    factors,
+    backend,
+    optimize,
+    simplify_sequence="ADCRS",
+    simplify_atol=1e-12,
+):
+    tn = pauli_product_expectation_tn(
+        quimb_circuit,
+        factors,
+        simplify_sequence=simplify_sequence,
+        simplify_atol=simplify_atol,
+    )
+    return tn.contract(all, output_inds=(), optimize=optimize, backend=backend)
+
+
 def __init__(self, quimb_backend="numpy", contraction_optimizer="auto-hq"):
     super(self.__class__, self).__init__()
 
@@ -340,6 +417,9 @@ def _qibo_circuit_to_quimb(
             circ.apply_gate("CNOT", c, t)
             continue
         if quimb_gate_name is None:
+            if hasattr(gate, "matrix"):
+                circ.apply_gate_raw(gate.matrix(), getattr(gate, "qubits", ()))
+                continue
             raise_error(ValueError, f"Gate {gate_name} not supported in Quimb backend.")
 
         params = getattr(gate, "parameters", ())
@@ -426,20 +506,24 @@ def expectation(self, circuit, observable, parallel=None, parallel_opts=None):
 
         exp_val = 0.0
         for coeff, factors in all_terms:
-            op = None
-            where = []
-            for qubit, gate_name in factors:
-                p = qu.pauli(gate_name.lower())
-                op = p if op is None else op & p
-                where.append(qubit)
-
-            val = qc.local_expectation(
-                op, tuple(where),
-                backend=self.backend,
-                optimize=self.contractions_optimizer,
-                simplify_sequence="ADCRS",
-                simplify_atol=1e-12,
-            )
+            if len(factors) > PAULI_DENSE_MAX_QUBITS:
+                val = pauli_product_expectation(
+                    qc,
+                    factors,
+                    backend=self.backend,
+                    optimize=self.contractions_optimizer,
+                    simplify_sequence="ADCRS",
+                    simplify_atol=1e-12,
+                )
+            else:
+                op, where = _pauli_term_to_dense_operator(factors)
+                val = qc.local_expectation(
+                    op, where,
+                    backend=self.backend,
+                    optimize=self.contractions_optimizer,
+                    simplify_sequence="ADCRS",
+                    simplify_atol=1e-12,
+                )
             exp_val += coeff * val
 
         return self.real(exp_val)
@@ -487,14 +571,11 @@ def _expectation_parallel(self, circuit, observable, method, opts):
 
     my_exp = 0.0
     for coeff, factors in my_terms:
-        op = None
-        where = []
-        for qubit, gate_name in factors:
-            p = qu.pauli(gate_name.lower())
-            op = p if op is None else op & p
-            where.append(qubit)
-
-        tn = qc.local_expectation(op, tuple(where), rehearse='tn')
+        if len(factors) > PAULI_DENSE_MAX_QUBITS:
+            tn = pauli_product_expectation_tn(qc, factors)
+        else:
+            op, where = _pauli_term_to_dense_operator(factors)
+            tn = qc.local_expectation(op, where, rehearse='tn')
 
         tree = parallel_path_search(
             tn, tn.outer_inds(),