PyZX Interop: Converting bloqs to pyzx circuits

qualtran/_infra/bloq.py

@@ -21,6 +21,8 @@
 if TYPE_CHECKING:
     import cirq
     import networkx as nx
+    import numpy as np
+    import pyzx as zx
     import quimb.tensor as qtn
     import sympy
     from numpy.typing import NDArray
@@ -39,6 +41,7 @@
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.cirq_interop.t_complexity_protocol import TComplexity
     from qualtran.drawing import WireSymbol
+    from qualtran.pyzx_interop import ZXAncillaManager
     from qualtran.resource_counting import (
         BloqCountDictT,
         BloqCountT,
@@ -551,3 +554,8 @@ def wire_symbol(
 
     def __str__(self):
         return self.__class__.__name__
+
+    def as_zx_gates(
+        self, ancilla_manager: 'ZXAncillaManager', /, **qubits: 'NDArray[np.integer]'
+    ) -> tuple[list['zx.circuit.Gate'], dict[str, 'NDArray[np.integer]']]:
+        raise NotImplementedError(f"{self} does not declare a conversion to ZX gates.")

qualtran/bloqs/basic_gates/x_basis.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 from attrs import frozen
+from numpy.typing import NDArray
 
 from qualtran import (
     AddControlledT,
@@ -36,6 +37,7 @@
 
 if TYPE_CHECKING:
     import cirq
+    import pyzx as zx
     import quimb.tensor as qtn
 
     from qualtran.cirq_interop import CirqQuregT
@@ -263,3 +265,11 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym
             return Text('X')
 
         return ModPlus()
+
+    def as_zx_gates(
+        self, ancilla_manager, /, q: NDArray[np.integer]
+    ) -> tuple[list['zx.circuit.Gate'], dict[str, NDArray[np.integer]]]:
+        import pyzx as zx
+
+        (qubit,) = q
+        return [zx.circuit.NOT(qubit)], {'q': q}

qualtran/bloqs/basic_gates/z_basis.py

@@ -46,9 +46,11 @@
 
 if TYPE_CHECKING:
     import cirq
+    import pyzx as zx
     import quimb.tensor as qtn
 
     from qualtran.cirq_interop import CirqQuregT
+    from qualtran.pyzx_interop import ZXAncillaManager
     from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
 
 _ZERO = np.array([1, 0], dtype=np.complex128)
@@ -137,6 +139,28 @@ def wire_symbol(
         s = '1' if self.bit else '0'
         return directional_text_box(s, side=reg.side)
 
+    def as_zx_gates(
+        self, ancilla_manager: 'ZXAncillaManager', /, **qubits: NDArray[np.integer]
+    ) -> tuple[list['zx.circuit.Gate'], dict[str, NDArray[np.integer]]]:
+        import pyzx as zx
+
+        if self.state:
+            qubit = ancilla_manager.allocate()
+
+            gates = [zx.circuit.gates.InitAncilla(qubit), zx.circuit.gates.HAD(qubit)]
+            if self.bit:
+                gates = gates + [zx.circuit.gates.NOT(qubit)]
+
+            return gates, {'q': np.array([qubit])}
+        else:
+            (qubit,) = qubits.pop('q')
+
+            gates = [zx.circuit.gates.HAD(qubit), zx.circuit.gates.PostSelect(qubit)]
+            if self.bit:
+                gates = [zx.circuit.gates.NOT(qubit)] + gates
+
+            return gates, {}
+
 
 def _hide_base_fields(cls, fields):
     # for use in attrs `field_transformer`.
@@ -285,6 +309,14 @@ def wire_symbol(
 
         return TextBox('Z')
 
+    def as_zx_gates(
+        self, ancilla_manager, /, q: NDArray[np.integer]
+    ) -> tuple[list['zx.circuit.Gate'], dict[str, NDArray[np.integer]]]:
+        import pyzx as zx
+
+        (qubit,) = q
+        return [zx.circuit.Z(qubit)], {'q': q}
+
 
 @bloq_example
 def _zgate() -> ZGate:

qualtran/pyzx_interop/__init__.py

@@ -0,0 +1,14 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from .bloq_to_pyzx_circuit import bloq_to_pyzx_circuit, ZXAncillaManager

qualtran/pyzx_interop/bloq_to_pyzx_circuit.py

@@ -0,0 +1,179 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from typing import Iterable
+
+import numpy as np
+import pyzx as zx
+from attrs import define
+from numpy.typing import NDArray
+
+from qualtran import Bloq, CompositeBloq, LeftDangle, Register, RightDangle, Signature, Soquet
+
+ZXQubitMap = dict[str, NDArray[np.integer]]
+"""A mapping from register names to an NDArray of ZX qubits"""
+
+
+@define
+class ZXAncillaManager:
+    """A simple ancilla qubit manager for generating pyzx Circuits.
+
+    Attributes:
+        n: number of existing qubits in the starting circuit.
+    """
+
+    n: int
+
+    def allocate(self) -> int:
+        """Allocate an uninitialized ancilla qubit.
+
+        This returns the index of a new ancilla qubit for use.
+        Note that it must be manually initialized, e.g. using the `InitAncilla` gate.
+        """
+        idx = self.n
+        self.n += 1
+        return idx
+
+    def free(self, q: int):
+        """Free an ancilla qubit.
+
+        Discard an ancilla qubit. For now, this operation does nothing.
+        """
+
+
+def _empty_qubit_map_from_registers(registers: Iterable[Register]) -> ZXQubitMap:
+    """For each register, creates an empty NDArray of the appropriate shape to store the zx qubits."""
+    return {reg.name: np.empty(reg.shape + (reg.bitsize,), dtype=int) for reg in registers}
+
+
+def _initalize_zx_circuit_from_signature(
+    signature: Signature,
+) -> tuple[zx.Circuit, ZXQubitMap, ZXAncillaManager]:
+    """Initalize a pyzx circuit from a bloq signature.
+
+    This enumerates the qubits in the same order as the registers in the signature.
+
+    Args:
+        signature: the signature of the bloq.
+
+    Returns:
+        A tuple of the pyzx circuit, the mapping from register names to qubits, and an
+        ancilla manager for the circuit.
+    """
+
+    n_qubits: int = 0
+    qubit_d: ZXQubitMap = {}
+
+    for reg in signature.lefts():
+        n = reg.total_bits()
+        idxs = np.arange(n) + n_qubits
+        shape = reg.shape + (reg.bitsize,)
+        qubit_d[reg.name] = idxs.reshape(shape)
+        n_qubits += n
+
+    circ = zx.Circuit(qubit_amount=n_qubits)
+    return circ, qubit_d, ZXAncillaManager(n_qubits)
+
+
+def _add_bloq_to_pyzx_circuit(
+    circ: zx.Circuit, bloq: Bloq, ancilla_manager: ZXAncillaManager, in_qubits: ZXQubitMap
+) -> ZXQubitMap:
+    """Add a single bloq acting on the given input qubits to a pyzx circuit.
+
+    Args:
+        circ: the pyzx circuit.
+        bloq: the bloq to add.
+        ancilla_manager: the ancilla manager for `circ`.
+        in_qubits: the input qubits to the bloq.
+
+    Returns:
+        A mapping of output register names to output qubits.
+    """
+    try:
+        gates, out_qubits = bloq.as_zx_gates(ancilla_manager, **in_qubits)
+        for gate in gates:
+            circ.add_gate(gate)
+        return out_qubits
+    except NotImplementedError:
+        pass
+
+    cbloq = bloq.decompose_bloq()
+    return _add_cbloq_to_pyzx_circuit(circ, cbloq, ancilla_manager, in_qubits)
+
+
+def _add_cbloq_to_pyzx_circuit(
+    circ: zx.Circuit, cbloq: CompositeBloq, ancilla_manager: ZXAncillaManager, in_qubits: ZXQubitMap
+) -> ZXQubitMap:
+    """Add a composite bloq acting on the given input qubits to a pyzx circuit.
+
+    This iterates through the `cbloq` graph in topologically-sorted order, and adds each
+    bloq instance.
+
+    Args:
+        circ: the pyzx circuit.
+        cbloq: the composite bloq to add.
+        ancilla_manager: the ancilla manager for `circ`.
+        in_qubits: the input qubits to the bloq.
+
+    Returns:
+        A mapping of output register names to output qubits.
+    """
+    # initialize the soquets corresponding to the `cbloq` inputs
+    soq_map: dict[Soquet, NDArray[np.integer]] = {
+        soq: in_qubits[soq.reg.name][soq.idx]
+        for soq in cbloq.all_soquets
+        if soq.binst is LeftDangle
+    }
+
+    for binst, pred_cxns, succ_cxns in cbloq.iter_bloqnections():
+        bloq = binst.bloq
+
+        # compute the input qubits
+        bloq_in_qubits: ZXQubitMap = _empty_qubit_map_from_registers(bloq.signature.lefts())
+        for cxn in pred_cxns:
+            bloq_soq = cxn.right
+            bloq_in_qubits[bloq_soq.reg.name][bloq_soq.idx] = soq_map.pop(cxn.left)
+
+        out_qubits: ZXQubitMap = _add_bloq_to_pyzx_circuit(
+            circ, bloq, ancilla_manager, bloq_in_qubits
+        )
+
+        # forward the output qubits to their corresponding soqs
+        for cxn in succ_cxns:
+            bloq_soq = cxn.left
+            soq_map[bloq_soq] = out_qubits[bloq_soq.reg.name][bloq_soq.idx]
+
+    # forward the soqs to the cbloq output soqs
+    for cxn in cbloq.connections:
+        if cxn.right.binst is RightDangle:
+            soq_map[cxn.right] = soq_map.pop(cxn.left)
+
+    # get the output qubits
+    out_qubits = _empty_qubit_map_from_registers(cbloq.signature.rights())
+    for soq, qubits in soq_map.items():
+        out_qubits[soq.reg.name][soq.idx] = qubits
+    return out_qubits
+
+
+def bloq_to_pyzx_circuit(bloq: Bloq) -> zx.Circuit:
+    """Build a pyzx circuit of a bloq.
+
+    Args:
+        bloq: the bloq to convert.
+
+    Returns:
+        A pyzx circuit corresponding to `bloq`.
+    """
+    circ, in_qubits, ancilla_manager = _initalize_zx_circuit_from_signature(bloq.signature)
+    _ = _add_bloq_to_pyzx_circuit(circ, bloq, ancilla_manager, in_qubits)
+    return circ

qualtran/pyzx_interop/bloq_to_pyzx_circuit_test.py

@@ -0,0 +1,44 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import numpy as np
+import pyzx as zx
+
+from qualtran import Bloq, BloqBuilder, Signature, Soquet, SoquetT
+from qualtran.bloqs.basic_gates import OneEffect, XGate, ZeroState, ZGate
+from qualtran.pyzx_interop.bloq_to_pyzx_circuit import bloq_to_pyzx_circuit
+
+
+class TestBloq(Bloq):
+    @property
+    def signature(self) -> 'Signature':
+        return Signature.build(q=1)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', q: 'Soquet') -> dict[str, 'SoquetT']:
+        q = bb.add(ZGate(), q=q)
+        q = bb.add(XGate(), q=q)
+
+        a = bb.add(ZeroState())
+        a = bb.add(XGate(), q=a)
+        bb.add(OneEffect(), q=a)
+
+        return {'q': q}
+
+
+def test_bloq_to_pyzx_circuit():
+    bloq = TestBloq()
+    circ = bloq_to_pyzx_circuit(bloq)
+    tensor = bloq.tensor_contract()
+
+    assert zx.compare_tensors(circ, tensor)
+    np.testing.assert_allclose(np.imag(zx.find_scalar_correction(circ, tensor)), 0, atol=1e-7)

qualtran/simulation/tensor/_quimb.py

@@ -69,11 +69,8 @@ def cbloq_to_quimb(cbloq: CompositeBloq) -> qtn.TensorNetwork:
             # the tensor network. Add an identity tensor acting on this register to make sure the
             # tensor network has variables corresponding to all input / output registers.
 
-            n = cxn.left.reg.bitsize
             for j in range(cxn.left.reg.bitsize):
-
                 placeholder = Soquet(None, Register('simulation_placeholder', QBit()))  # type: ignore
-                Connection(cxn.left, placeholder)
                 tn.add(
                     qtn.Tensor(
                         data=np.eye(2),