Add metadata to sweeps (#6099)

* Add metadata to sweeps

- Adds optional metadata attribute to Linspace and Points.
- This will be used to adjust device parameters via sweep.

See https://tinyurl.com/cirq-sweep-metadata-public for RFC.

Co-authored-by: Matthew Neeley <mneeley@gmail.com>

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

@@ -1,7 +1,20 @@
-{
-  "cirq_type": "Linspace",
-  "key": "a",
-  "start": 0,
-  "stop": 1,
-  "length": 4
-}
+[
+    {
+      "cirq_type": "Linspace",
+      "key": "a",
+      "start": 0,
+      "stop": 1,
+      "length": 4
+    },
+    {
+      "cirq_type": "Linspace",
+      "key": "b",
+      "start": 1,
+      "stop": 2,
+      "length": 3,
+      "metadata": {
+              "cirq_type": "NamedQubit",
+              "name": "b"
+      }
+    }
+]

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

@@ -1 +1,4 @@
-cirq.Linspace("a", start=0, stop=1, length=4)
+[
+    cirq.Linspace("a", start=0, stop=1, length=4),
+    cirq.Linspace("b", start=1, stop=2, length=3, metadata=cirq.NamedQubit("b"))
+]

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

@@ -1,8 +1,19 @@
-{
-  "cirq_type": "Points",
-  "key": "a",
-  "points": [
-    0,
-    0.4
-  ]
-}
+[
+    {
+      "cirq_type": "Points",
+      "key": "a",
+      "points": [
+        0,
+        0.4
+      ]
+    },
+    {
+      "cirq_type": "Points",
+      "key": "amp",
+      "points": [
+        0,
+        1
+      ],
+      "metadata": "used to turn amp on or off"
+    }
+]

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

@@ -1 +1,4 @@
-cirq.Points('a', [0, 0.4])
+[
+    cirq.Points('a', [0, 0.4]),
+    cirq.Points('amp', [0, 1], metadata='used to turn amp on or off'),
+]

cirq-core/cirq/study/sweeps.py

@@ -18,6 +18,7 @@
     Iterable,
     Iterator,
     List,
+    Optional,
     overload,
     Sequence,
     TYPE_CHECKING,
@@ -421,9 +422,23 @@ def _values(self) -> Iterator[float]:
 class Points(SingleSweep):
     """A simple sweep with explicitly supplied values."""
 
-    def __init__(self, key: 'cirq.TParamKey', points: Sequence[float]) -> None:
-        super(Points, self).__init__(key)
+    def __init__(
+        self, key: 'cirq.TParamKey', points: Sequence[float], metadata: Optional[Any] = None
+    ) -> None:
+        """Creates a sweep on a variable with supplied values.
+
+        Args:
+            key: sympy.Symbol or equivalent to sweep across.
+            points: sequence of floating point values that represent
+                the values to sweep across.  The length of the sweep
+                will be equivalent to the length of this sequence.
+            metadata: Optional metadata to attach to the sweep to
+                annotate the sweep or its variable.
+
+        """
+        super().__init__(key)
         self.points = points
+        self.metadata = metadata
 
     def _tuple(self) -> Tuple[Union[str, sympy.Expr], Sequence[float]]:
         return self.key, tuple(self.points)
@@ -435,25 +450,44 @@ def _values(self) -> Iterator[float]:
         return iter(self.points)
 
     def __repr__(self) -> str:
-        return f'cirq.Points({self.key!r}, {self.points!r})'
+        metadata_repr = f', metadata={self.metadata!r}' if self.metadata is not None else ""
+        return f'cirq.Points({self.key!r}, {self.points!r}{metadata_repr})'
 
     def _json_dict_(self) -> Dict[str, Any]:
+        if self.metadata is not None:
+            return protocols.obj_to_dict_helper(self, ["key", "points", "metadata"])
         return protocols.obj_to_dict_helper(self, ["key", "points"])
 
 
 class Linspace(SingleSweep):
     """A simple sweep over linearly-spaced values."""
 
-    def __init__(self, key: 'cirq.TParamKey', start: float, stop: float, length: int) -> None:
+    def __init__(
+        self,
+        key: 'cirq.TParamKey',
+        start: float,
+        stop: float,
+        length: int,
+        metadata: Optional[Any] = None,
+    ) -> None:
         """Creates a linear-spaced sweep for a given key.
 
         For the given args, assigns to the list of values
             start, start + (stop - start) / (length - 1), ..., stop
+
+        Args:
+            key: sympy.Symbol or equivalent to sweep across.
+            start: minimum value of linear sweep.
+            stop: maximum value of linear sweep.
+            length: number of points in the sweep.
+            metadata: Optional metadata to attach to the sweep to
+                annotate the sweep or its variable.
         """
-        super(Linspace, self).__init__(key)
+        super().__init__(key)
         self.start = start
         self.stop = stop
         self.length = length
+        self.metadata = metadata
 
     def _tuple(self) -> Tuple[Union[str, sympy.Expr], float, float, int]:
         return (self.key, self.start, self.stop, self.length)
@@ -470,12 +504,17 @@ def _values(self) -> Iterator[float]:
                 yield self.start * (1 - p) + self.stop * p
 
     def __repr__(self) -> str:
+        metadata_repr = f', metadata={self.metadata!r}' if self.metadata is not None else ""
         return (
             f'cirq.Linspace({self.key!r}, start={self.start!r}, '
-            f'stop={self.stop!r}, length={self.length!r})'
+            f'stop={self.stop!r}, length={self.length!r}{metadata_repr})'
         )
 
     def _json_dict_(self) -> Dict[str, Any]:
+        if self.metadata is not None:
+            return protocols.obj_to_dict_helper(
+                self, ["key", "start", "stop", "length", "metadata"]
+            )
         return protocols.obj_to_dict_helper(self, ["key", "start", "stop", "length"])
 
 

cirq-core/cirq/study/sweeps_test.py

@@ -273,6 +273,12 @@ def test_repr():
     )
     cirq.testing.assert_equivalent_repr(cirq.Points('zero&pi', [0, 3.14159]))
     cirq.testing.assert_equivalent_repr(cirq.Linspace('I/10', 0, 1, 10))
+    cirq.testing.assert_equivalent_repr(
+        cirq.Points('zero&pi', [0, 3.14159], metadata='example str')
+    )
+    cirq.testing.assert_equivalent_repr(
+        cirq.Linspace('for_q0', 0, 1, 10, metadata=cirq.LineQubit(0))
+    )
 
 
 def test_zip_product_str():