Adding tags to Moments (#7467)

- This PR adds the ability to add tags to Moment objects.
- tags in Moments function similar to tags in operations, in that they
can be any hashable object and are lost if the circuit undergoes
transformation, such as by a transformer.

Author: dstrain115

cirq-core/cirq/circuits/moment.py

@@ -24,6 +24,7 @@
     Any,
     Callable,
     cast,
+    Hashable,
     Iterable,
     Iterator,
     Mapping,
@@ -77,7 +78,12 @@ class Moment:
             are no such operations, returns an empty Moment.
     """
 
-    def __init__(self, *contents: cirq.OP_TREE, _flatten_contents: bool = True) -> None:
+    def __init__(
+        self,
+        *contents: cirq.OP_TREE,
+        _flatten_contents: bool = True,
+        tags: tuple[Hashable, ...] = (),
+    ) -> None:
         """Constructs a moment with the given operations.
 
         Args:
@@ -88,6 +94,12 @@ def __init__(self, *contents: cirq.OP_TREE, _flatten_contents: bool = True) -> N
                 we skip flattening and assume that contents already consists
                 of individual operations. This is used internally by helper
                 methods to avoid unnecessary validation.
+            tags:  Optional tags to denote specific Moment objects with meta-data.
+                These are a tuple of any Hashable object.  Typically, a class
+                will be passed.  Tags apply only to this specific set of operations
+                and will be lost on any transformation of the
+                Moment.  For instance, if operations are added to the Moment, tags
+                will be dropped unless explicitly added back in by the user.
 
         Raises:
             ValueError: A qubit appears more than once.
@@ -110,9 +122,10 @@ def __init__(self, *contents: cirq.OP_TREE, _flatten_contents: bool = True) -> N
 
         self._measurement_key_objs: frozenset[cirq.MeasurementKey] | None = None
         self._control_keys: frozenset[cirq.MeasurementKey] | None = None
+        self._tags = tags
 
     @classmethod
-    def from_ops(cls, *ops: cirq.Operation) -> cirq.Moment:
+    def from_ops(cls, *ops: cirq.Operation, tags: tuple[Hashable, ...] = ()) -> cirq.Moment:
         """Construct a Moment from the given operations.
 
         This avoids calling `flatten_to_ops` in the moment constructor, which
@@ -122,8 +135,11 @@ def from_ops(cls, *ops: cirq.Operation) -> cirq.Moment:
 
         Args:
             *ops: Operations to include in the Moment.
+            tags:  Optional tags to denote specific Moment objects with meta-data.
+                These are a tuple of any Hashable object.  Tags will be dropped if
+                the operations in the Moment are modified or transformed.
         """
-        return cls(*ops, _flatten_contents=False)
+        return cls(*ops, _flatten_contents=False, tags=tags)
 
     @property
     def operations(self) -> tuple[cirq.Operation, ...]:
@@ -133,6 +149,34 @@ def operations(self) -> tuple[cirq.Operation, ...]:
     def qubits(self) -> frozenset[cirq.Qid]:
         return frozenset(self._qubit_to_op)
 
+    @property
+    def tags(self) -> tuple[Hashable, ...]:
+        """Returns a tuple of the operation's tags."""
+        return self._tags
+
+    def with_tags(self, *new_tags: Hashable) -> cirq.Moment:
+        """Creates a new Moment with the current ops and the specified tags.
+
+        If the moment already has tags, this will add the new_tags to the
+        preexisting tags.
+
+        This method can be used to attach meta-data to moments
+        without affecting their functionality.  The intended usage is to
+        attach classes intended for this purpose or strings to mark operations
+        for specific usage that will be recognized by consumers.
+
+        Tags can be a list of any type of object that is useful to identify
+        this operation as long as the type is hashable.  If you wish the
+        resulting operation to be eventually serialized into JSON, you should
+        also restrict the operation to be JSON serializable.
+
+        Please note that tags should be instantiated if classes are
+        used.  Raw types are not allowed.
+        """
+        if not new_tags:
+            return self
+        return Moment(*self._operations, _flatten_contents=False, tags=(*self._tags, *new_tags))
+
     def operates_on_single_qubit(self, qubit: cirq.Qid) -> bool:
         """Determines if the moment has operations touching the given qubit.
         Args:
@@ -170,6 +214,8 @@ def operation_at(self, qubit: raw_types.Qid) -> cirq.Operation | None:
     def with_operation(self, operation: cirq.Operation) -> cirq.Moment:
         """Returns an equal moment, but with the given op added.
 
+        Any tags on the Moment will be dropped.
+
         Args:
             operation: The operation to append.
 
@@ -198,6 +244,9 @@ def with_operation(self, operation: cirq.Operation) -> cirq.Moment:
     def with_operations(self, *contents: cirq.OP_TREE) -> cirq.Moment:
         """Returns a new moment with the given contents added.
 
+        Any tags on the original Moment object are dropped if the Moment
+        is changed.
+
         Args:
             *contents: New operations to add to this moment.
 
@@ -235,6 +284,9 @@ def with_operations(self, *contents: cirq.OP_TREE) -> cirq.Moment:
     def without_operations_touching(self, qubits: Iterable[cirq.Qid]) -> cirq.Moment:
         """Returns an equal moment, but without ops on the given qubits.
 
+        Any tags on the original Moment object are dropped if the Moment
+        is changed.
+
         Args:
             qubits: Operations that touch these will be removed.
 
@@ -510,11 +562,13 @@ def _superoperator_(self) -> np.ndarray:
         return qis.kraus_to_superoperator(self._kraus_())
 
     def _json_dict_(self) -> dict[str, Any]:
-        return protocols.obj_to_dict_helper(self, ['operations'])
+        # For backwards compatibility, only output tags if they exist.
+        args = ['operations', 'tags'] if self._tags else ['operations']
+        return protocols.obj_to_dict_helper(self, args)
 
     @classmethod
-    def _from_json_dict_(cls, operations, **kwargs):
-        return cls.from_ops(*operations)
+    def _from_json_dict_(cls, operations, tags=(), **kwargs):
+        return cls(*operations, tags=tags)
 
     def __add__(self, other: cirq.OP_TREE) -> cirq.Moment:
         if isinstance(other, circuit.AbstractCircuit):

cirq-core/cirq/circuits/moment_test.py

@@ -939,3 +939,68 @@ def test_superoperator() -> None:
     assert m._has_superoperator_()
     s = m._superoperator_()
     assert np.allclose(s, np.array([[1, 0, 0, 1], [0, 0, 0, 0], [0, 0, 0, 0], [1, 0, 0, 1]]) / 2)
+
+
+def test_moment_with_tags() -> None:
+    q0 = cirq.LineQubit(0)
+    q1 = cirq.LineQubit(1)
+    op1 = cirq.X(q0)
+    op2 = cirq.Y(q1)
+
+    # Test initialization with no tags
+    moment_no_tags = cirq.Moment(op1)
+    assert moment_no_tags.tags == ()
+
+    # Test initialization with tags
+    moment_with_tags = cirq.Moment(op1, op2, tags=("initial_tag_1", "initial_tag_2"))
+    assert moment_with_tags.tags == ("initial_tag_1", "initial_tag_2")
+
+    # Test with_tags method to add new tags
+    new_moment = moment_with_tags.with_tags("new_tag_1", "new_tag_2")
+
+    # Ensure the original moment's tags are unchanged
+    assert moment_with_tags.tags == ("initial_tag_1", "initial_tag_2")
+
+    # Ensure the new moment has both old and new tags
+    assert new_moment.tags == ("initial_tag_1", "initial_tag_2", "new_tag_1", "new_tag_2")
+
+    # Test with_tags on a moment that initially had no tags
+    new_moment_from_no_tags = moment_no_tags.with_tags("single_new_tag")
+    assert new_moment_from_no_tags.tags == ("single_new_tag",)
+
+    # Test adding no new tags
+    same_moment_tags = moment_with_tags.with_tags()
+    assert same_moment_tags.tags == ("initial_tag_1", "initial_tag_2")
+
+    class CustomTag:
+        """Example Hashable Tag"""
+
+        def __init__(self, value):
+            self.value = value
+
+        def __hash__(self):
+            return hash(self.value)  # pragma: nocover
+
+        def __eq__(self, other):
+            return isinstance(other, CustomTag) and self.value == other.value  # pragma: nocover
+
+        def __repr__(self):
+            return f"CustomTag({self.value})"  # pragma: nocover
+
+    tag_obj = CustomTag("complex_tag")
+    moment_with_custom_tag = cirq.Moment(op1, tags=("string_tag", 123, tag_obj))
+    assert moment_with_custom_tag.tags == ("string_tag", 123, tag_obj)
+
+    new_moment_with_custom_tag = moment_with_custom_tag.with_tags(456)
+    assert new_moment_with_custom_tag.tags == ("string_tag", 123, tag_obj, 456)
+
+    # Test that tags are dropped if the Moment is changed.
+    moment = cirq.Moment.from_ops(op1, tags=(tag_obj,))
+    assert moment.tags == (tag_obj,)
+    assert moment.with_operation(op2).tags == ()
+    assert moment.with_operations(op2).tags == ()
+    assert moment.without_operations_touching([q0]).tags == ()
+
+    # Test that tags are retained if the Moment is unchanged.
+    assert moment.with_operations().tags == (tag_obj,)
+    assert moment.without_operations_touching([q1]).tags == (tag_obj,)

cirq-core/cirq/protocols/json_test_data/Moment.json

@@ -39,5 +39,28 @@
         }
       }
     ]
+  },
+  {
+	  "cirq_type": "Moment",
+	  "operations": [
+	      	  {
+		    	  "cirq_type": "SingleQubitPauliStringGateOperation",
+		    	  "pauli": {
+			  	  "cirq_type": "_PauliX",
+			  	  "exponent": 1.0,
+			  	  "global_shift": 0.0
+		    	  },
+		    	  "qubit": {
+			  	  "cirq_type": "LineQubit",
+			  	  "x": 0
+		    	  }
+	      	  }
+	  ],
+	  "tags": [
+	      {
+	      	  "cirq_type": "Duration",
+	      	  "picos": 25000
+	      }
+	  ]
   }
-]
+]

cirq-core/cirq/protocols/json_test_data/Moment.repr

@@ -2,4 +2,9 @@
     cirq.X(cirq.LineQubit(0)),
     cirq.Y(cirq.LineQubit(1)),
     cirq.Z(cirq.LineQubit(2)),
-)]
+),
+cirq.Moment(
+    cirq.X(cirq.LineQubit(0)),
+    tags=(cirq.Duration(nanos=25),)
+)
+]