代码封装

.venv/lib/python3.12/site-packages/quimb/tensor/circuit.py

@@ -1573,6 +1573,23 @@ def _combine_1q_gate_run(gates, array_fn=None):
     return Gate.from_raw(G, gates[0].qubits)
 
 
+def _combine_2q_gate_run(gates, array_fn=None):
+    """Combine a run of two qubit gates in application order."""
+    gates = tuple(gate for _, gate in gates)
+    G = gates[0].array
+    if array_fn is not None:
+        G = array_fn(G)
+    G = reshape(G, (4, 4))
+
+    for gate in gates[1:]:
+        Gi = gate.array
+        if array_fn is not None:
+            Gi = array_fn(Gi)
+        G = reshape(Gi, (4, 4)) @ G
+
+    return Gate.from_raw(reshape(G, (2, 2, 2, 2)), gates[0].qubits)
+
+
 def _can_merge_1q_gate(gate):
     return (
         (gate.controls is None)
@@ -1583,48 +1600,96 @@ def _can_merge_1q_gate(gate):
     )
 
 
-def _iter_gates_with_merged_1q_runs(gates):
+def _can_merge_2q_gate(gate):
+    return (
+        (gate.controls is None)
+        and (not gate.special)
+        and (not gate.parametrize)
+        and (gate.qubits is not None)
+        and (len(gate.qubits) == 2)
+    )
+
+
+def _iter_gates_with_merged_runs(gates, merge_1q=True, merge_2q=True):
     """Yield ``(gate_to_apply, gates_to_record)``, merging adjacent runs of
-    single qubit gates that are not interrupted by any operation touching the
-    same qubit.
+    local gates that are not interrupted by any operation touching the same
+    qubits.
     """
-    pending = {}
+    pending_1q = {}
+    pending_2q = {}
 
     def flush_qubit(q):
-        run = pending.pop(q, None)
+        run = pending_1q.pop(q, None)
         if run is None:
             return
         if len(run) == 1:
             return run[0][1], run
         return None, run
 
+    def flush_pair(pair):
+        run = pending_2q.pop(pair, None)
+        if run is None:
+            return
+        if len(run) == 1:
+            return run[0][1], run
+        return None, run
+
+    def flush_touched(touched, keep_qubit=None, keep_pair=None):
+        for q in tuple(pending_1q):
+            if q == keep_qubit:
+                continue
+            if q in touched:
+                item = flush_qubit(q)
+                if item is not None:
+                    yield item
+
+        for pair in tuple(pending_2q):
+            if pair == keep_pair:
+                continue
+            if touched.intersection(pair):
+                item = flush_pair(pair)
+                if item is not None:
+                    yield item
+
     def flush_all():
-        for q in tuple(pending):
+        for q in tuple(pending_1q):
             item = flush_qubit(q)
             if item is not None:
                 yield item
+        for pair in tuple(pending_2q):
+            item = flush_pair(pair)
+            if item is not None:
+                yield item
 
     for i, gate in enumerate(gates):
-        if _can_merge_1q_gate(gate):
+        if merge_1q and _can_merge_1q_gate(gate):
             (q,) = gate.qubits
-            pending.setdefault(q, []).append((i, gate))
+            yield from flush_touched({q}, keep_qubit=q)
+            pending_1q.setdefault(q, []).append((i, gate))
+            continue
+
+        if merge_2q and _can_merge_2q_gate(gate):
+            pair = gate.qubits
+            yield from flush_touched(set(pair), keep_pair=pair)
+            pending_2q.setdefault(pair, []).append((i, gate))
             continue
 
         touched = set(gate.qubits or ())
         if gate.controls:
             touched.update(gate.controls)
 
-        for q in tuple(pending):
-            if q in touched:
-                item = flush_qubit(q)
-                if item is not None:
-                    yield item
+        yield from flush_touched(touched)
 
         yield gate, ((i, gate),)
 
     yield from flush_all()
 
 
+_iter_gates_with_merged_1q_runs = functools.partial(
+    _iter_gates_with_merged_runs, merge_1q=True, merge_2q=False
+)
+
+
 # --------------------------- main circuit class ---------------------------- #
 
 
@@ -2103,6 +2168,24 @@ class Circuit:
 
         self._psi.gate_(G, gates[0][1].qubits, tags=tags, **opts)
 
+    def _apply_merged_2q_gate_run(self, gates, gate_number_offset=0, **gate_opts):
+        tags = tags_to_oset(gate_opts.pop("tags", None))
+        for i, gate in gates:
+            tags |= self._gate_tags_for_record(
+                gate, gate_number=gate_number_offset + i
+            )
+
+        opts = {**self.gate_opts, **gate_opts}
+
+        if self.convert_eager:
+            G = _combine_2q_gate_run(
+                gates, array_fn=self._maybe_convert_gate_array
+            ).array
+        else:
+            G = _combine_2q_gate_run(gates).array
+
+        self._psi.gate_(G, gates[0][1].qubits, tags=tags, **opts)
+
     def apply_gate(
         self,
         gate_id,
@@ -2178,11 +2261,14 @@ class Circuit:
             Supplied to :meth:`~quimb.tensor.circuit.Circuit.apply_gate`.
         """
         merge_1q = gate_opts.pop("merge_1q", "auto")
+        merge_2q = gate_opts.pop("merge_2q", "auto")
 
         if merge_1q == "auto":
             merge_1q = True
+        if merge_2q == "auto":
+            merge_2q = True
 
-        if merge_1q:
+        if merge_1q or merge_2q:
             gates = tuple(
                 gate if isinstance(gate, Gate) else parse_to_gate(gate)
                 for gate in gates
@@ -2195,15 +2281,22 @@ class Circuit:
                 pbar = _progbar(total=len(gates))
 
             gate_number_offset = len(self._gates)
-            for gate, gates_to_record in _iter_gates_with_merged_1q_runs(
-                gates
+            for gate, gates_to_record in _iter_gates_with_merged_runs(
+                gates, merge_1q=merge_1q, merge_2q=merge_2q
             ):
                 if gate is None:
-                    self._apply_merged_1q_gate_run(
-                        gates_to_record,
-                        gate_number_offset=gate_number_offset,
-                        **gate_opts,
-                    )
+                    if len(gates_to_record[0][1].qubits) == 1:
+                        self._apply_merged_1q_gate_run(
+                            gates_to_record,
+                            gate_number_offset=gate_number_offset,
+                            **gate_opts,
+                        )
+                    else:
+                        self._apply_merged_2q_gate_run(
+                            gates_to_record,
+                            gate_number_offset=gate_number_offset,
+                            **gate_opts,
+                        )
                 else:
                     self._apply_gate(
                         gate,
@@ -4892,11 +4985,16 @@ class CircuitMPS(Circuit):
 
     def apply_gates(self, gates, progbar=False, **gate_opts):
         merge_1q = gate_opts.pop("merge_1q", "auto")
+        merge_2q = gate_opts.pop("merge_2q", "auto")
 
         if merge_1q == "auto":
             merge_1q = True
+        if merge_2q == "auto":
+            # MPS truncation semantics are sensitive to when a 2q gate is
+            # materialized, so keep the default conservative here.
+            merge_2q = False
 
-        if merge_1q:
+        if merge_1q or merge_2q:
             gates = tuple(
                 gate if isinstance(gate, Gate) else parse_to_gate(gate)
                 for gate in gates
@@ -4913,15 +5011,22 @@ class CircuitMPS(Circuit):
                 )
 
             gate_number_offset = len(self._gates)
-            for gate, gates_to_record in _iter_gates_with_merged_1q_runs(
-                gates
+            for gate, gates_to_record in _iter_gates_with_merged_runs(
+                gates, merge_1q=merge_1q, merge_2q=merge_2q
             ):
                 if gate is None:
-                    self._apply_merged_1q_gate_run(
-                        gates_to_record,
-                        gate_number_offset=gate_number_offset,
-                        **gate_opts,
-                    )
+                    if len(gates_to_record[0][1].qubits) == 1:
+                        self._apply_merged_1q_gate_run(
+                            gates_to_record,
+                            gate_number_offset=gate_number_offset,
+                            **gate_opts,
+                        )
+                    else:
+                        self._apply_merged_2q_gate_run(
+                            gates_to_record,
+                            gate_number_offset=gate_number_offset,
+                            **gate_opts,
+                        )
                     gate_for_progress = gates_to_record[-1][1]
                 else:
                     self._apply_gate(

.venv/lib/python3.12/site-packages/quimb/tensor/tn1d/core.py

@@ -5050,8 +5050,6 @@ class TNLinearOperator1D(spla.LinearOperator):
 
         if self.is_conj:
             T = T.conj()
-        print(T)
-        assert(0)
         return T.to_dense(self.left_inds, self.right_inds)
 
     def toarray(self):