Preserve identical signature with bitsize information during Cirq interop

[tanujkhattar]

This is built on top of quantumlib/Cirq#6286 and quantumlib/Cirq#6298 and should address parts of #360

This is the last big change to support seamless Cirq-FT interop. See example below for updated usage.

bloq = CirqGateAsBloq(cirq_ft.SwapWithZeroGate(2, 3, 4))
cbloq = bloq.decompose_bloq()
fig, ax = draw_musical_score(get_musical_score_data(cbloq))
fig.set_size_inches(11, 4)

cc @fdmalone @mpharrigan

Note: Failing tests will be fixed once quantumlib/Cirq#6298 is merged.

[mpharrigan]

Super cool stuff. The logic -- while complex -- looks to be organized quite well in this PR. And of course the users will see a seamless api!

I have two requests

• Asserts should be used such that if one is violated it means the developer has made a mistake. Some of these assertions made sense as assertions when we were always making quregs ourselves, for example. But now there are places where user-provided input can trigger an assertion failure including those without accompanying messages. Can you replace data input validation insertions with if not x: raise ValueError(f"Helpful error message")?
• The closures (ie nested function definitions) inside this large and complex function are too large and complex. It is not apparent what is being closed over (which is bb). And since we already have a ton of variable names which are all variants of "qvar" or "qureg" there's a non-zero chance of accidentally shadowing the name of an outer scope variable inside the closure accidentally. Can you please move these to free functions.

When I was reviewing I also re-worded the docstring and did some of the value error things, fyc below

From 654e8654a6fbd62bd714ed60a6f2bcd8b2f47eea Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Mon, 25 Sep 2023 17:14:08 -0700
Subject: [PATCH 2/3] docstring

---
 qualtran/cirq_interop/_cirq_interop.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/qualtran/cirq_interop/_cirq_interop.py b/qualtran/cirq_interop/_cirq_interop.py
index d1b92ad..208364d 100644
--- a/qualtran/cirq_interop/_cirq_interop.py
+++ b/qualtran/cirq_interop/_cirq_interop.py
@@ -180,21 +180,21 @@ def cirq_optree_to_cbloq(
     """Convert a Cirq OP-TREE into a `CompositeBloq` with signature `signature`.
     
     Each `cirq.Operation` will be wrapped into a `CirqGateAsBloq` wrapper.
+    The signature of the resultant CompositeBloq is `signature`, if provided. Otherwise, use
+    one thru-register named "qubits" of shape `(n_qubits,)`.
 
-    If `signature` is not None, the signature of the resultant CompositeBloq is `signature`. For
-    multi-dimensional registers and registers with > 1 bitsize, this function automatically
-    splits the input soquets into a flat list and joins the output soquets into the correct shape
-    to ensure compatibility with the flat API expected by Cirq. When specifying a signature, users
-    must also specify the `cirq_quregs` argument, which is a mapping of cirq qubits used in the
-    OP-TREE corresponding to the `signature`. If `signature` has registers with entry
+    For multi-dimensional registers and registers with bitsize>1, this function automatically
+    splits the input soquets and joins the output soquets to ensure compatibility with the
+    flat-list-of-qubits expected by Cirq.
+
+    When specifying a signature, users must also specify the `cirq_quregs` argument, which is a
+    mapping of cirq qubits used in the OP-TREE corresponding to the `signature`. If `signature`
+    has registers with entry
         - `Register('x', bitsize=2, shape=(3, 4))` and
         - `Register('y', bitsize=1, shape=(10, 20))`
     then `cirq_quregs` should have one entry corresponding to each register as follows:
-        - key='x'; value=`np.array(cirq_qubits_used_in_optree, shape=(3, 4, 2))` and
-        - key='y'; value=`np.array(cirq_qubits_used_in_optree, shape=(10, 20, 1))`.
-
-    If `signature` is None, the resultant composite bloq will have one thru-register named "qubits"
-    of shape `(n_qubits,)`.
+        - key='x'; value=`np.array(cirq_qubits_used_for_x, shape=(3, 4, 2))` and
+        - key='y'; value=`np.array(cirq_qubits_used_for_y, shape=(10, 20, 1))`.
     """
     circuit = cirq.Circuit(optree)
     # "qubits" means cirq qubits | "qvars" means bloq Soquets
-- 
2.42.0.515.g380fc7ccd1-goog

From f43cbcb138314dadcdff4b2126baea98bb834c5f Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Mon, 25 Sep 2023 17:14:42 -0700
Subject: [PATCH 3/3] value error 2

---
 qualtran/cirq_interop/_cirq_interop.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/qualtran/cirq_interop/_cirq_interop.py b/qualtran/cirq_interop/_cirq_interop.py
index 208364d..1d54dee 100644
--- a/qualtran/cirq_interop/_cirq_interop.py
+++ b/qualtran/cirq_interop/_cirq_interop.py
@@ -197,17 +197,17 @@ def cirq_optree_to_cbloq(
         - key='y'; value=`np.array(cirq_qubits_used_for_y, shape=(10, 20, 1))`.
     """
     circuit = cirq.Circuit(optree)
-    # "qubits" means cirq qubits | "qvars" means bloq Soquets
     if signature is None:
-        assert cirq_quregs is None
+        if cirq_quregs is not None:
+            raise ValueError("`cirq_quregs` requires specifying `signature`.")
         all_qubits = sorted(circuit.all_qubits())
         signature = Signature([Register('qubits', 1, shape=(len(all_qubits),))])
         cirq_quregs = {'qubits': np.array(all_qubits).reshape(len(all_qubits), 1)}
+    else:
+        if cirq_quregs is None:
+            raise ValueError("`signature` requires specifying `cirq_quregs`.")
 
     cirq_quregs = {k: np.apply_along_axis(QReg, -1, v) for k, v in cirq_quregs.items()}
-
-    assert signature is not None and cirq_quregs is not None
-
     bb, initial_soqs = BloqBuilder.from_signature(signature, add_registers_allowed=False)
         
     # 0. Helper functions
-- 
2.42.0.515.g380fc7ccd1-goog

[mpharrigan]

please replace with a helpfully worded ValueError

[mpharrigan]

go 4 it

[tanujkhattar]

@mpharrigan Address all comments.