Add test for the add_dd_and_align_right in 2D GHZ state generation

cirq-core/cirq/experiments/ghz_2d.py

@@ -50,7 +50,11 @@ def _transform_circuit(circuit: circuits.Circuit) -> circuits.Circuit:
     circuit = circuit + circuits.Circuit(ops.measure(*qubits, key="m"))
     circuit = transformers.align_right(transformers.merge_single_qubit_gates_to_phxz(circuit))
     circuit = transformers.stratified_circuit(
-        circuit[::-1], categories=[lambda op: protocols.num_qubits(op) == k for k in (1, 2)]
+        circuit[::-1],
+        categories=[
+            lambda op: protocols.num_qubits(op) == 1,
+            lambda op: protocols.num_qubits(op) == 2,
+        ],
     )[::-1]
     circuit = transformers.add_dynamical_decoupling(circuit)
     circuit = circuits.Circuit(circuit[:-1])

cirq-core/cirq/experiments/ghz_2d_test.py

@@ -153,3 +153,67 @@ def test_ghz_invalid_inputs():
         ghz_2d.generate_2d_ghz_circuit(
             center_qubit, graph, num_qubits=len(graph.nodes) + 1  # invalid
         )
+
+
+allowed_single_qubit_gates_in_transformed = (
+    cirq.PhasedXZGate,
+    cirq.XPowGate,
+    cirq.YPowGate,
+    cirq.ZPowGate,
+    cirq.IdentityGate,
+)
+
+allowed_single_qubit_gates_in_base = (cirq.HPowGate,)
+
+
+def _is_allowed_single_qubit_gate(op: cirq.Operation, allowed_gates: tuple) -> bool:
+    """Checks if a single-qubit operation is in the allowed gate tuple"""
+
+    return isinstance(op.gate, allowed_gates)
+
+
+def test_dd_is_applied_by_structural_check():
+    """Verifies that DD is applied by checking for an increase in circuit length
+    and the presence of the standard DD gate in the transformed circuit.
+    """
+
+    # Define the graph
+    q0 = cirq.GridQubit(0, 0)
+    q1 = cirq.GridQubit(0, 1)
+    q2 = cirq.GridQubit(0, 2)
+    center_qubit = q0
+
+    graph_3q = nx.Graph([(q0, q1), (q1, q2)])
+
+    # Create the base circuit
+    base_circuit = ghz_2d.generate_2d_ghz_circuit(
+        center_qubit, graph_3q, num_qubits=3, add_dd_and_align_right=False
+    )
+
+    # Create the transformed circuit
+    transformed_circuit = ghz_2d.generate_2d_ghz_circuit(
+        center_qubit, graph_3q, num_qubits=3, add_dd_and_align_right=True
+    )
+
+    # Find the single qubit operations in base circuit
+    single_qubit_ops_in_base = list(
+        base_circuit.findall_operations(lambda op: cirq.num_qubits(op) == 1)
+    )
+
+    # Find the single qubit operations in transformed circuit
+    single_qubit_ops_in_transformed = list(
+        transformed_circuit.findall_operations(lambda op: cirq.num_qubits(op) == 1)
+    )
+
+    # Assert the transformed circuit has more single qubit operations than the base circuit
+    assert len(single_qubit_ops_in_transformed) > len(single_qubit_ops_in_base)
+
+    # The base circuit should only have a Hadamard as a single qubit gate
+    for op in single_qubit_ops_in_base:
+        assert _is_allowed_single_qubit_gate(op[1], allowed_single_qubit_gates_in_base)
+
+    # The transformed circuit should have DD gates and PhasedXZ gates as single qubit gates
+    for op in single_qubit_ops_in_transformed:
+        assert _is_allowed_single_qubit_gate(op[1], allowed_single_qubit_gates_in_transformed)
+
+    assert transformed_circuit != base_circuit