Fix BooleanHamiltonian.with_qubits() to work with not only NamedQubits (#4396)

The current code at head does not implement BooleanHamiltonian.with_qubits() correctly. It assumes that the qubits are named and then maps them to the variable name.

This is restrictive, does not follow the other classes' with_qubits(). It is also incorrect because it silently accepts a list of qubits with missing entries (e.g. there is a variable name 'a' but we don't provide a qubit named 'a'), and the error is thrown only when decompose_once() is called.

Instead, here the code accepts a list of qubits and replaces them in the map. It requires that there are as many qubits provided as there were when building the object. Also, it relies on the order of the qubit keys when building the BooleanHamiltonian object, but that is currently a restriction of the API.

Fixes #4390.

Author: tonybruguier

cirq-core/cirq/ops/boolean_hamiltonian.py

@@ -21,7 +21,7 @@
 [3] https://github.com/rsln-s/IEEE_QW_2020/blob/master/Slides.pdf
 """
 
-from typing import cast, Any, Dict, Generator, List, Sequence, Tuple
+from typing import Any, Dict, Generator, List, Sequence, Tuple
 
 import sympy.parsing.sympy_parser as sympy_parser
 
@@ -66,8 +66,14 @@ def __init__(
         self._theta: float = theta
 
     def with_qubits(self, *new_qubits: 'cirq.Qid') -> 'BooleanHamiltonian':
+        if len(self._qubit_map) != len(new_qubits):
+            raise ValueError('Length of replacement qubits must be the same')
+        new_qubit_map = {
+            variable_name: new_qubit
+            for variable_name, new_qubit in zip(self._qubit_map, new_qubits)
+        }
         return BooleanHamiltonian(
-            {cast(cirq.NamedQubit, q).name: q for q in new_qubits},
+            new_qubit_map,
             self._boolean_strs,
             self._theta,
         )

cirq-core/cirq/ops/boolean_hamiltonian_test.py

@@ -72,7 +72,6 @@ def test_circuit(boolean_str):
     hamiltonian_gate = cirq.BooleanHamiltonian(
         {q.name: q for q in qubits}, [boolean_str], 0.1 * math.pi
     )
-    assert hamiltonian_gate.with_qubits(*qubits) == hamiltonian_gate
 
     assert hamiltonian_gate.num_qubits() == n
 
@@ -83,3 +82,19 @@ def test_circuit(boolean_str):
 
     # Compare the two:
     np.testing.assert_array_equal(actual, expected)
+
+
+def test_with_custom_names():
+    q0, q1, q2, q3 = cirq.LineQubit.range(4)
+    original_op = cirq.BooleanHamiltonian(
+        {'a': q0, 'b': q1},
+        ['a'],
+        0.1,
+    )
+    assert cirq.decompose(original_op) == [cirq.Rz(rads=-0.05).on(q0)]
+
+    renamed_op = original_op.with_qubits(q2, q3)
+    assert cirq.decompose(renamed_op) == [cirq.Rz(rads=-0.05).on(q2)]
+
+    with pytest.raises(ValueError, match='Length of replacement qubits must be the same'):
+        original_op.with_qubits(q2)