improve light cone impl

Author: refraction-ray

AGENTS.md

@@ -22,6 +22,10 @@ TensorCircuit is a **Tensor Network-first**, **Multi-Backend** quantum computing
 - **Pattern**: Avoid Python control flow (if/else) that depends on tensor values. Use `tc.backend.cond` or `tc.backend.switch` if necessary, or structure code to be statically analyzable.
 - **Benefit**: This enables massive speedups on JAX and TensorFlow backends.
 
+### Known Issues and Intermittent Failures
+
+- **`tests/test_circuit.py::test_qiskit2tc`**: This test may fail intermittently (e.g., in Qiskit 0.46.3+) due to non-deterministic behavior in Qiskit's `UnitaryGate.control()` method during circuit translation. If this test fails while other tests pass, it is likely a framework-level "heisenbug" and not a regression in TensorCircuit core logic.
+
 ## Repository Structure
 
 - `tensorcircuit/`: Core package source code.
@@ -33,6 +37,12 @@ TensorCircuit is a **Tensor Network-first**, **Multi-Backend** quantum computing
 
 ## Configuration and Dependencies
 
+### Python Environment Management
+
+- **Rule**: Never install packages in the default system or user Python environment unless explicitly directed by the user.
+- **Protocol**: If a command fails due to a missing package or `ModuleNotFoundError`, **explicitly ask the user** for the correct Python environment to use (e.g., a specific Conda environment name like `2602`).
+- **Workflow**: Once an environment is identified, prefix all Python/pytest/pylint commands with `conda run -n <env_name>` or the appropriate environment activator.
+
 ### Core Dependencies
 
 - numpy

CHANGELOG.md

@@ -18,12 +18,17 @@
 
 - Add `bitwise_or`, `any`, `all` methods for backends.
 
+- Better lightcone cancellation.
+
 ### Fixed
 
 - Fix pauli propagation sign error for Y operator.
 
 - Fix jax backend QR gradient bug.
 
+- Lazy import for tensorflow and torch.
+
+
 ## v1.5.0
 
 ### Added
@@ -60,8 +65,6 @@
 
 - Generalize `scatter` method for backends.
 
-- Lazy import for tensorflow and torch.
-
 ## v1.4.0
 
 ### Added

examples/lightcone_simplify.py

@@ -5,7 +5,8 @@
 import numpy as np
 import tensorcircuit as tc
 
-K = tc.set_backend("tensorflow")
+K = tc.set_backend("jax")
+tc.set_contractor("cotengra")
 
 
 def brickwall_ansatz(c, params, gatename, nlayers):
@@ -34,8 +35,10 @@ def loss(params, n, nlayers, enable_lightcone):
 
 
 def efficiency():
-    for n in range(6, 40, 4):
-        for nlayers in range(2, 6, 2):
+    for n in range(10, 40, 4):
+        for nlayers in range(2, 6, 3):
+            if n > 20 and nlayers > 4:
+                continue
             print(n, nlayers)
             print("w lightcone")
             (v2, g2), _, _ = tc.utils.benchmark(
@@ -60,5 +63,5 @@ def correctness(n, nlayers):
 
 
 if __name__ == "__main__":
-    efficiency()
     correctness(7, 3)
+    efficiency()

tensorcircuit/basecircuit.py

@@ -102,6 +102,18 @@ def _tensors_to_nodes(
     def coloring_nodes(
         nodes: Sequence[tn.Node], is_dagger: bool = False, flag: str = "inputs"
     ) -> None:
+        r"""
+        Tag nodes with metadata used for casual lightcone simplification and tracing.
+
+        :param nodes: A sequence of tensornetwork nodes to tag.
+        :type nodes: Sequence[tn.Node]
+        :param is_dagger: Whether the nodes represent conjugate operations (U^\dagger),
+            defaults to False.
+        :type is_dagger: bool, optional
+        :param flag: A label for the node type (e.g., "gate", "inputs", "operator"),
+            defaults to "inputs".
+        :type flag: str, optional
+        """
         for node in nodes:
             node.is_dagger = is_dagger
             node.flag = flag
@@ -114,6 +126,19 @@ def coloring_copied_nodes(
         is_dagger: bool = True,
         flag: str = "inputs",
     ) -> None:
+        """
+        Tag copied nodes while preserving the original node's identity for lightcone cancellation.
+
+        :param nodes: A sequence of newly copied nodes.
+        :type nodes: Sequence[tn.Node]
+        :param nodes0: The sequence of original nodes from which `nodes` were copied.
+        :type nodes0: Sequence[tn.Node]
+        :param is_dagger: Whether the copied nodes represent conjugate operations,
+            defaults to True.
+        :type is_dagger: bool, optional
+        :param flag: A label for the node type, defaults to "inputs".
+        :type flag: str, optional
+        """
         for node, n0 in zip(nodes, nodes0):
             node.is_dagger = is_dagger
             node.flag = flag
@@ -238,7 +263,7 @@ def apply_general_gate(
 
             if applied is False:
                 gate.name = name
-                self.coloring_nodes([gate])
+                self.coloring_nodes([gate], flag="gate")
                 # gate.flag = "gate"
                 # gate.is_dagger = False
                 # gate.id = id(gate)
@@ -257,18 +282,16 @@ def apply_general_gate(
         elif mpo:  # gate in MPO format
             assert isinstance(gate, QuOperator)
             gatec = gate.copy()
+            self.coloring_nodes(gatec.nodes, flag="gate")
             for n in gatec.nodes:
-                n.flag = "gate"
-                n.is_dagger = False
                 n.id = id(gate)
                 n.name = name
             self._nodes += gatec.nodes
             if self.is_dm:
                 gateconj = gate.adjoint()
-                for n0, n in zip(gatec.nodes, gateconj.nodes):
-                    n.flag = "gate"
-                    n.is_dagger = True
-                    n.id = id(n0)
+                self.coloring_nodes(gateconj.nodes, flag="gate", is_dagger=True)
+                for _, n in zip(gatec.nodes, gateconj.nodes):
+                    n.id = id(gate)
                     n.name = name
                 self._nodes += gateconj.nodes
 
@@ -284,9 +307,8 @@ def apply_general_gate(
                 mps_nodes = [gate]
             else:
                 mps_nodes = gate.nodes
+            self.coloring_nodes(mps_nodes, flag="gate")
             for n in mps_nodes:
-                n.flag = "gate"
-                n.is_dagger = False
                 n.id = id(gate)
                 n.name = name
             self._nodes += mps_nodes
@@ -299,10 +321,9 @@ def apply_general_gate(
                 else:
                     gateconj = gate.adjoint()
                     gateconj_nodes = gateconj.nodes
-                for n0, n in zip(mps_nodes, gateconj_nodes):
-                    n.flag = "gate"
-                    n.is_dagger = True
-                    n.id = id(n0)
+                self.coloring_nodes(gateconj_nodes, flag="gate", is_dagger=True)
+                for _, n in zip(mps_nodes, gateconj_nodes):
+                    n.id = id(gate)
                     n.name = name
                 self._nodes += gateconj_nodes
 
@@ -463,15 +484,11 @@ def measure_jit(
                 else:
                     m = onehot_d_tensor(sample[i], d=self._d)
                 g1 = Gate(m)
-                g1.id = id(g1)
-                g1.is_dagger = False
-                g1.flag = "measurement"
+                self.coloring_nodes([g1], flag="measurement")
                 newnodes.append(g1)
                 g1.get_edge(0) ^ edge1[index[i]]
                 g2 = Gate(m)
-                g2.id = id(g2)
-                g2.is_dagger = True
-                g2.flag = "measurement"
+                self.coloring_nodes([g2], flag="measurement", is_dagger=True)
                 newnodes.append(g2)
                 g2.get_edge(0) ^ edge2[index[i]]
 
@@ -551,16 +568,14 @@ def amplitude_before(self, l: Union[str, Tensor]) -> List[Gate]:
             msconj = []
         for i in range(self._nqubits):
             n = tn.Node(endns[i])
-            n.flag = "measurement"
-            n.is_dagger = False
-            n.id = id(n)
+            self.coloring_nodes([n], flag="measurement")
             ms.append(n)
             d_edges[i] ^ n.get_edge(0)
             if self.is_dm:
                 nconj = tn.Node(endns[i])
-                nconj.flag = "measurement"
-                nconj.is_dagger = True
-                nconj.id = id(n)
+                self.coloring_copied_nodes(
+                    [nconj], [n], flag="measurement", is_dagger=True
+                )
                 msconj.append(nconj)
                 d_edges[i + self._nqubits] ^ nconj.get_edge(0)
         no.extend(ms)

tensorcircuit/circuit.py

@@ -117,7 +117,7 @@ def __init__(
         else:
             raise ValueError("No inputs provided")  # should not be reached
 
-        self.coloring_nodes(nodes)
+        self.coloring_nodes(nodes, flag="inputs")
         self._nodes = nodes
 
         self._start_index = len(nodes)
@@ -180,7 +180,7 @@ def replace_mps_inputs(self, mps_inputs: QuOperator) -> None:
                         new_front[j] ^ other[0][other[1]]
         j += 1
         self._front += new_front[j:]
-        self.coloring_nodes(new_nodes)
+        self.coloring_nodes(new_nodes, flag="inputs")
         self._nodes = new_nodes + self._nodes[self._start_index :]
         self._start_index = len(new_nodes)
 

tensorcircuit/simplify.py

@@ -203,79 +203,100 @@ def _full_rank_simplify(nodes: List[Any]) -> List[Any]:
 
 
 def _light_cone_cancel(nodes: List[Any]) -> Tuple[List[Any], bool]:
+    """
+    Scan the nodes and cancel pairs of U and U^dagger that are directly connected.
+    Assumes that for a gate node, the first half of its edges are 'output' (future-facing)
+    and the second half are 'input' (past-facing).
+    """
     is_changed = False
-    for ind in range(len(nodes) // 2, 0, -1):
-        n = nodes[ind]
-        if n.is_dagger is True:
+    nodes_to_remove = set()
+
+    # Identify the "future" side. Nodes in expectation are typically [ket_nodes, bra_nodes, ops]
+    # We scan backward through the nodes list.
+    for i in range(len(nodes) - 1, -1, -1):
+        n = nodes[i]
+        if n in nodes_to_remove:
+            continue
+        if getattr(n, "is_dagger", None) is True:
             continue
+
         noe = len(n.shape)
         if noe % 2 != 0:
             continue
-        e = n[0]
-        n1, n2 = e.node1, e.node2  # one of them is n itself
-        if n1 is None or n2 is None:
-            continue
-        if n1.is_dagger == n2.is_dagger:
-            continue
-        if n1.id != n2.id:
-            continue
-        if e.axis1 != e.axis2:
-            continue
-        for i in range(noe // 2):
-            e = n[i]
-            n3, n4 = e.node1, e.node2  # should also be n1 and n2
-            if n3 is None or n4 is None:
+
+        # Check if all output legs (0 to noe//2 - 1) are connected to the same conjugate node
+        match_node = None
+        for leg_idx in range(noe // 2):
+            e = n[leg_idx]
+            if e.is_dangling():
+                break
+            n1, n2 = e.node1, e.node2
+            other = n2 if n1 is n else n1
+
+            if getattr(other, "is_dagger", None) is not True:
                 break
-            if not ((n3 is n1 and n4 is n2) or (n3 is n2 and n4 is n1)):
+            if getattr(other, "id", None) != getattr(n, "id", -1):
                 break
             if e.axis1 != e.axis2:
                 break
+
+            if match_node is None:
+                match_node = other
+            elif match_node is not other:
+                break
         else:
-            if n1 is not n:
-                n1, n2 = n2, n1  # make sure n1 is n dagger is False
-
-            # contract
-            njs = [i for i, n in enumerate(nodes) if n is n1 or n is n2]
-            # new_node = tn.contract_between(e.node1, e.node2)
-            # contract(e) is not enough for multi edges between two tensors
-            for i in range(noe // 2):
-                e = n1[noe // 2 + i]
-                m1, m3 = e.node1, e.node2
-                i1, i3 = e.axis1, e.axis2
-                if m1 is not n1:
-                    m1, m3 = m3, m1  # m1 is n1, m3 is the one behind m1
-                    i1, i3 = i3, i1
-                e.disconnect()
-                e = n2[noe // 2 + i]
-                m2, m4 = e.node1, e.node2
-                i2, i4 = e.axis1, e.axis2
-                if m2 is not n2:
-                    m2, m4 = m4, m2  # m1 is n1, m3 is the one behind m1
-                    i2, i4 = i4, i2
-                e.disconnect()
-                m3[i3] ^ m4[i4]
-            nodes = _multi_remove(nodes, njs)
-
-            is_changed = True
-            break  # switch to the next node
-    return nodes, is_changed
+            if match_node is not None and match_node not in nodes_to_remove:
+                # Perform cancellation by bypass
+                for leg_idx in range(noe // 2, noe):
+                    e_n = n[leg_idx]
+                    e_m = match_node[leg_idx]
+
+                    m_n, i_n = (
+                        (e_n.node2, e_n.axis2)
+                        if e_n.node1 is n
+                        else (e_n.node1, e_n.axis1)
+                    )
+                    m_m, i_m = (
+                        (e_m.node2, e_m.axis2)
+                        if e_m.node1 is match_node
+                        else (e_m.node1, e_m.axis1)
+                    )
+
+                    e_n.disconnect()
+                    e_m.disconnect()
+                    m_n[i_n] ^ m_m[i_m]
+
+                nodes_to_remove.add(n)
+                nodes_to_remove.add(match_node)
+                is_changed = True
+
+    if is_changed:
+        new_nodes = [n for n in nodes if n not in nodes_to_remove]
+        return new_nodes, True
+    return nodes, False
 
 
 # TODO(@refraction-ray): better light cone cancellation in terms of MPO gates (three legs one)
+# MPO cancellation requires matching internal bond dimensions and identities of all nodes in the MPO.
 
 
 def _full_light_cone_cancel(nodes: List[Any]) -> List[Any]:
     """
     Simplify the list of tc.Nodes using casual lightcone structure.
 
-    :param nodes: _description_
+    :param nodes: List of nodes representing the tensor network.
     :type nodes: List[Any]
-    :return: _description_
+    :return: Simplified list of nodes.
     :rtype: List[Any]
     """
-    for n in nodes:
-        if getattr(n, "is_dagger", None) is None:
-            return nodes
+    if not nodes:
+        return nodes
+    # Check metadata availability
+    if any(getattr(n, "is_dagger", None) is None for n in nodes):
+        return nodes
+
+    # A single pass of _light_cone_cancel (optimized for O(N)) is often sufficient
+    # but we use a while loop to ensure all possible cancellations are resolved.
     nodes, is_changed = _light_cone_cancel(nodes)
     while is_changed:
         nodes, is_changed = _light_cone_cancel(nodes)