Add custom operation conversion hook to cirq_optree_to_cbloq (#1834)

Adds a backward-compatible hook that lets callers provide a custom
Operation -> Bloq conversion function. This is needed by the FLASQ
library to perform custom analysis as we convert cirq optrees to bloqs.
For example, calculating the Manhattan distance between the two qubits
involved in a two-qubit gate.

The new parameter defaults to None, so all existing callers are
unaffected. When provided, it replaces the default extraction logic for
each operation in the circuit.

All existing cirq_interop tests pass unchanged and a new test was added
for this new functionality.

Co-authored-by: Matthew Harrigan <mpharrigan@google.com>

Author: wjhuggins

qualtran/cirq_interop/_cirq_to_bloq.py

@@ -18,7 +18,18 @@
 import itertools
 import warnings
 from functools import cached_property
-from typing import Any, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, TypeVar, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TYPE_CHECKING,
+    TypeVar,
+    Union,
+)
 
 import cirq
 import numpy as np
@@ -319,7 +330,7 @@ def _ensure_in_reg_exists(bb: BloqBuilder, in_reg: _QReg, qreg_to_qvar: Dict[_QR
                 soqs_to_join[qreg.qubits[0]] = soq
         elif len(in_reg_qubits) == 1 and qreg.qubits and qreg.qubits[0] in in_reg_qubits:
             # Cast single QBit registers to the appropriate single-bit register dtype.
-            err_msg = "Found non-QBit type register which shouldn't happen: " f"{soq}"
+            err_msg = f"Found non-QBit type register which shouldn't happen: {soq}"
             assert isinstance(soq.dtype, QBit), err_msg
             if not isinstance(in_reg.dtype, QBit):
                 qreg_to_qvar[in_reg] = bb.add(Cast(QBit(), in_reg.dtype), reg=soq)
@@ -465,6 +476,7 @@ def cirq_optree_to_cbloq(
     signature: Optional[Signature] = None,
     in_quregs: Optional[Dict[str, 'CirqQuregT']] = None,
     out_quregs: Optional[Dict[str, 'CirqQuregT']] = None,
+    op_conversion_method: Optional[Callable[[cirq.Operation], Bloq]] = None,
 ) -> CompositeBloq:
     """Convert a Cirq OP-TREE into a `CompositeBloq` with signature `signature`.
 
@@ -495,6 +507,17 @@ def cirq_optree_to_cbloq(
 
     Any qubit in `optree` which is not part of `in_quregs` and `out_quregs` is considered to be
     allocated & deallocated inside the CompositeBloq and does not show up in it's signature.
+
+    Args:
+        optree: A Cirq OP_TREE (e.g. a circuit or list of operations).
+        signature: The signature of the resulting CompositeBloq. If not provided, a default
+            signature with one thru-register named "qubits" is used.
+        in_quregs: Mapping from register names to arrays of cirq qubits for LEFT registers.
+        out_quregs: Mapping from register names to arrays of cirq qubits for RIGHT registers.
+        op_conversion_method: An optional callable that takes a ``cirq.Operation`` and returns
+            a ``Bloq``. If provided, this is used instead of the default ``_extract_bloq_from_op``
+            to convert each operation. This allows callers to attach custom metadata (e.g.
+            routing costs) to bloqs during conversion.
     """
     circuit = cirq.Circuit(optree)
     if signature is None:
@@ -533,7 +556,10 @@ def cirq_optree_to_cbloq(
 
     # 2. Add each operation to the composite Bloq.
     for op in circuit.all_operations():
-        bloq = _extract_bloq_from_op(op)
+        if op_conversion_method is not None:
+            bloq = op_conversion_method(op)
+        else:
+            bloq = _extract_bloq_from_op(op)
         if bloq.signature == Signature([]):
             bb.add(bloq)
             continue

qualtran/cirq_interop/_cirq_to_bloq_test.py

@@ -286,3 +286,33 @@ def test_cirq_gate_cost_via_decomp():
 
     gc_swappow = get_cost_value(swappow_bloq, QECGatesCost())
     assert gc_swappow == GateCounts(clifford=5, rotation=1, and_bloq=1, measurement=1)
+
+
+def test_cirq_optree_to_cbloq_op_conversion_method():
+    """Test the op_conversion_method parameter of cirq_optree_to_cbloq.
+
+    When provided, op_conversion_method should be called for each operation
+    instead of the default _extract_bloq_from_op.
+    """
+    qubits = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(cirq.H(qubits[0]), cirq.CNOT(qubits[0], qubits[1]), cirq.T(qubits[2]))
+
+    # Track which operations were converted
+    converted_ops: list[cirq.Operation] = []
+
+    def custom_converter(op: cirq.Operation) -> Bloq:
+        converted_ops.append(op)
+        # Fall back to the default behavior.
+        from qualtran.cirq_interop._cirq_to_bloq import _extract_bloq_from_op
+
+        return _extract_bloq_from_op(op)
+
+    cbloq = cirq_optree_to_cbloq(circuit, op_conversion_method=custom_converter)
+
+    # Verify the custom converter was called for each operation.
+    assert len(converted_ops) == 3
+
+    # The resulting CompositeBloq should still produce the correct unitary.
+    bloq_unitary = cbloq.tensor_contract()
+    cirq_unitary = circuit.unitary(qubits)
+    np.testing.assert_allclose(cirq_unitary, bloq_unitary, atol=1e-8)