Add Device Parameter metadata serialization to cirq_google

cirq-google/cirq_google/api/v2/run_context.proto

@@ -84,6 +84,22 @@ message SweepFunction {
   repeated Sweep sweeps = 2;
 }
 
+message DeviceParameter {
+
+    // Path to the parameter key
+    repeated string path = 1;
+
+    // If the value is an array, the index of the array to change.
+    int64 idx = 2;
+
+    // String representation of the units, if any.
+    // Examples: "GHz", "ns", etc.
+    string units = 3;
+
+    // Note that the device parameter values will be populated
+    // by the sweep values themselves.
+}
+
 
 // A set of values to loop over for a particular parameter.
 message SingleSweep {
@@ -98,6 +114,10 @@ message SingleSweep {
     // Uniformly-spaced sampling over a range.
     Linspace linspace = 3;
   }
+
+  // Optional arguments for if this is a device parameter.
+  // (as opposed to a circuit symbol)
+  DeviceParameter parameter = 4;
 }
 
 // A list of explicit values.

cirq-google/cirq_google/api/v2/sweeps.py

@@ -19,6 +19,7 @@
 import cirq
 from cirq_google.api.v2 import batch_pb2
 from cirq_google.api.v2 import run_context_pb2
+from cirq_google.study.device_parameter import DeviceParameter
 
 
 def sweep_to_proto(
@@ -54,9 +55,23 @@ def sweep_to_proto(
         out.single_sweep.linspace.first_point = sweep.start
         out.single_sweep.linspace.last_point = sweep.stop
         out.single_sweep.linspace.num_points = sweep.length
+        # Use duck-typing to support google-internal Parameter objects
+        if sweep.metadata and getattr(sweep.metadata, 'path', None):
+            out.single_sweep.parameter.path.extend(sweep.metadata.path)
+        if sweep.metadata and getattr(sweep.metadata, 'idx', None):
+            out.single_sweep.parameter.idx = sweep.metadata.idx
+        if sweep.metadata and getattr(sweep.metadata, 'units', None):
+            out.single_sweep.parameter.units = sweep.metadata.units
     elif isinstance(sweep, cirq.Points) and not isinstance(sweep.key, sympy.Expr):
         out.single_sweep.parameter_key = sweep.key
         out.single_sweep.points.points.extend(sweep.points)
+        # Use duck-typing to support google-internal Parameter objects
+        if sweep.metadata and getattr(sweep.metadata, 'path', None):
+            out.single_sweep.parameter.path.extend(sweep.metadata.path)
+        if sweep.metadata and getattr(sweep.metadata, 'idx', None):
+            out.single_sweep.parameter.idx = sweep.metadata.idx
+        if sweep.metadata and getattr(sweep.metadata, 'units', None):
+            out.single_sweep.parameter.units = sweep.metadata.units
     elif isinstance(sweep, cirq.ListSweep):
         sweep_dict: Dict[str, List[float]] = {}
         for param_resolver in sweep:
@@ -88,15 +103,22 @@ def sweep_from_proto(msg: run_context_pb2.Sweep) -> cirq.Sweep:
         raise ValueError(f'invalid sweep function type: {func_type}')
     if which == 'single_sweep':
         key = msg.single_sweep.parameter_key
+        if msg.single_sweep.HasField("parameter"):
+            metadata = DeviceParameter(
+                path=msg.single_sweep.parameter.path, idx=msg.single_sweep.parameter.idx
+            )
+        else:
+            metadata = None
         if msg.single_sweep.WhichOneof('sweep') == 'linspace':
             return cirq.Linspace(
                 key=key,
                 start=msg.single_sweep.linspace.first_point,
                 stop=msg.single_sweep.linspace.last_point,
                 length=msg.single_sweep.linspace.num_points,
+                metadata=metadata,
             )
         if msg.single_sweep.WhichOneof('sweep') == 'points':
-            return cirq.Points(key=key, points=msg.single_sweep.points.points)
+            return cirq.Points(key=key, points=msg.single_sweep.points.points, metadata=metadata)
 
         raise ValueError(f'single sweep type not set: {msg}')
 

cirq-google/cirq_google/api/v2/sweeps_test.py

@@ -19,6 +19,7 @@
 
 import cirq
 from cirq.study import sweeps
+from cirq_google.study import DeviceParameter
 from cirq_google.api import v2
 
 
@@ -41,7 +42,19 @@ def _values(self) -> Iterator[float]:
     [
         cirq.UnitSweep,
         cirq.Linspace('a', 0, 10, 100),
+        cirq.Linspace(
+            'a',
+            0,
+            10,
+            100,
+            metadata=DeviceParameter(path=['path', 'to', 'parameter'], idx=2, units='ns'),
+        ),
         cirq.Points('b', [1, 1.5, 2, 2.5, 3]),
+        cirq.Points(
+            'b',
+            [1, 1.5, 2, 2.5, 3],
+            metadata=DeviceParameter(path=['path', 'to', 'parameter'], idx=2, units='GHz'),
+        ),
         cirq.Linspace('a', 0, 1, 5) * cirq.Linspace('b', 0, 1, 5),
         cirq.Points('a', [1, 2, 3]) + cirq.Linspace('b', 0, 1, 3),
         (
@@ -62,14 +75,23 @@ def test_sweep_to_proto_roundtrip(sweep):
 
 
 def test_sweep_to_proto_linspace():
-    proto = v2.sweep_to_proto(cirq.Linspace('foo', 0, 1, 20))
+    proto = v2.sweep_to_proto(
+        cirq.Linspace(
+            'foo', 0, 1, 20, metadata=DeviceParameter(path=['path', 'to', 'parameter'], idx=2)
+        )
+    )
     assert isinstance(proto, v2.run_context_pb2.Sweep)
     assert proto.HasField('single_sweep')
     assert proto.single_sweep.parameter_key == 'foo'
     assert proto.single_sweep.WhichOneof('sweep') == 'linspace'
     assert proto.single_sweep.linspace.first_point == 0
     assert proto.single_sweep.linspace.last_point == 1
     assert proto.single_sweep.linspace.num_points == 20
+    assert proto.single_sweep.parameter.path == ['path', 'to', 'parameter']
+    assert proto.single_sweep.parameter.idx == 2
+    assert v2.sweep_from_proto(proto).metadata == DeviceParameter(
+        path=['path', 'to', 'parameter'], idx=2
+    )
 
 
 def test_list_sweep_bad_expression():