feat(core): implement Circuit.append with corresponding unit tests (#59)

src/tsim/circuit.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Iterable, Literal, cast, overload
+from typing import Any, Iterable, Literal, Sequence, Union, cast, overload
 
 import pyzx_param as zx
 import stim
@@ -65,6 +65,102 @@ def append_from_stim_program_text(self, stim_program_text: str) -> None:
         )
         self._stim_circ = self._stim_circ.flattened()
 
+    @overload
+    def append(
+        self,
+        name: str,
+        targets: Union[
+            int,
+            stim.GateTarget,
+            stim.PauliString,
+            Iterable[Union[int, stim.GateTarget, stim.PauliString]],
+        ] = (),
+        arg: Union[float, Iterable[float], None] = None,
+        *,
+        tag: str = "",
+    ) -> None: ...
+
+    @overload
+    def append(
+        self,
+        name: Union[stim.CircuitInstruction, stim.CircuitRepeatBlock, stim.Circuit],
+    ) -> None: ...
+
+    def append(
+        self,
+        name: Union[
+            str, stim.CircuitInstruction, stim.CircuitRepeatBlock, stim.Circuit
+        ],
+        targets: Union[
+            int,
+            stim.GateTarget,
+            stim.PauliString,
+            Iterable[Union[int, stim.GateTarget, stim.PauliString]],
+        ] = (),
+        arg: Union[float, Iterable[float], None] = None,
+        *,
+        tag: str = "",
+    ) -> None:
+        """Append an operation into the circuit.
+
+        Args:
+            name: The name of the operation's gate (e.g. "H" or "M" or "CNOT").
+
+                This argument can also be set to a `stim.CircuitInstruction` or
+                `stim.CircuitInstructionBlock`, which results in the instruction or
+                block being appended to the circuit. The other arguments (targets
+                and arg) can't be specified when doing so.
+
+            targets: The objects operated on by the gate. This can be either a
+                single target or an iterable of multiple targets.
+
+                Each target can be:
+                    An int: The index of a targeted qubit.
+                    A `stim.GateTarget`: Could be a variety of things. Methods like
+                        `stim.target_rec`, `stim.target_sweet`, `stim.target_x`, and
+                        `stim.CircuitInstruction.__getitem__` all return this type.
+                    A `stim.PauliString`: This will automatically be expanded into
+                        a product of pauli targets like `X1*Y2*Z3`.
+            arg: The "parens arguments" for the gate, such as the probability for a
+                noise operation. A double or list of doubles parameterizing the
+                gate. Different gates take different parens arguments. For example,
+                X_ERROR takes a probability, OBSERVABLE_INCLUDE takes an observable
+                index, and PAULI_CHANNEL_1 takes three disjoint probabilities.
+
+            tag: A customizable string attached to the instruction.
+
+        """
+        if isinstance(name, str):
+            if name == "T":
+                name = "S"
+                tag = "T"
+            elif name == "T_DAG":
+                name = "S_DAG"
+                tag = "T"
+            elif name in ("R_X", "R_Y", "R_Z"):
+                assert arg is not None, f"For {name} gates, an angle must be provided."
+                args = list(arg) if isinstance(arg, Iterable) else [arg]
+                assert (
+                    len(args) == 1
+                ), f"For {name} gates, a single angle must be provided."
+                tag = f"{name}(theta={args[0]}*pi)"
+                name = "I"
+                arg = None
+            elif name == "U3":
+                assert arg is not None and (
+                    isinstance(arg, Iterable) and len(list(arg)) == 3
+                ), f"For U3 gates, three rotation angles must be provided."
+                theta, phi, lam = list(arg)
+                tag = f"U3(theta={theta}*pi, phi={phi}*pi, lambda={lam}*pi)"
+                name = "I"
+                arg = None
+
+            self._stim_circ.append(name=name, targets=targets, arg=arg, tag=tag)  # type: ignore
+        else:
+            self._stim_circ.append(name=name)
+            if isinstance(name, stim.CircuitRepeatBlock):
+                self._stim_circ = self._stim_circ.flattened()
+
     @classmethod
     def from_file(cls, filename: str) -> Circuit:
         """Create a Circuit from a file.

test/unit/test_circuit.py

@@ -628,3 +628,43 @@ def test_diagram_pyzx_scale_horizontally(
         g = c.diagram(type=type, scale_horizontally=2)
         mock_draw.assert_called_once()
     assert hasattr(g, "vertices")
+
+
+def test_circuit_append():
+    c = Circuit()
+    c.append("T", [0, 1])
+    assert str(c) == "T 0 1"
+
+    c.append("T_DAG", 2)
+    assert "T_DAG 2" in str(c)
+
+    c.append("R_Z", 0, arg=0.25)
+    assert "R_Z(0.25) 0" in str(c)
+
+    c.append("R_X", 1, arg=[0.1])
+    assert "R_X(0.1) 1" in str(c)
+
+    c.append("U3", 0, arg=(0.3, 0.24, 0.49))
+    assert "U3(0.3, 0.24, 0.49) 0" in str(c)
+
+
+def test_circuit_append_circuit_instruction():
+    c = Circuit()
+    c.append(stim.CircuitInstruction("H", [0]))
+    assert str(c) == "H 0"
+
+
+def test_circuit_append_circuit_repeat_block():
+    c = Circuit()
+    block = stim.CircuitRepeatBlock(3, stim.Circuit("H 0"))
+    c.append(block)
+    # Should be flattened
+    assert str(c) == "H 0 0 0"
+
+
+def test_circuit_append_circuit():
+    c = Circuit()
+    sub_c = stim.Circuit("H 0\nCNOT 0 1")
+    c.append(sub_c)
+    assert "H 0" in str(c)
+    assert "CX 0 1" in str(c) or "CNOT 0 1" in str(c)