Add a new sweep metadata from/to proto approach (#6925)

Continue #6869

This PR added a new path to encode and decode metadata to this
corresponding proto. Notice this is a side path if the metadata in sweep
is not encoded as `Metadata` class, the code path is still old and same
for `from_proto` case. If the proto does not contain `metadata` proto
field, the code still convert back as the old style

Author: BichengYing

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

@@ -136,7 +136,10 @@ message Metadata {
     optional string label = 2;
 
     // If true, store this sweep as parameters instead of the independent axes.
-    optional bool as_parameter = 3;
+    optional bool is_const = 3;
+
+    // A temporary solution that we store the unit information here.
+    optional string unit = 4;
 }
 
 // A bundle of multiple DeviceParameters and their values.

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

@@ -20,7 +20,7 @@
 import cirq
 from cirq.study import sweeps
 from cirq_google.api.v2 import run_context_pb2
-from cirq_google.study.device_parameter import DeviceParameter
+from cirq_google.study.device_parameter import DeviceParameter, Metadata
 
 
 def _build_sweep_const(value: Any) -> run_context_pb2.ConstValue:
@@ -116,13 +116,17 @@ 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
+        # Encode the metadata if present
+        if isinstance(sweep.metadata, Metadata):
+            out.single_sweep.metadata.MergeFrom(metadata_to_proto(sweep.metadata))
+        else:
+            # 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):
         sweep = cast(cirq.Points, sweep_transformer(sweep))
         out.single_sweep.parameter_key = sweep.key
@@ -135,13 +139,17 @@ def sweep_to_proto(
                 unit.to_proto(out.single_sweep.points.unit)
             else:
                 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
+        # Encode the metadata if present
+        if isinstance(sweep.metadata, Metadata):
+            out.single_sweep.metadata.MergeFrom(metadata_to_proto(sweep.metadata))
+        else:
+            # 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:
@@ -190,6 +198,7 @@ def sweep_from_proto(
         raise ValueError(f'invalid sweep function type: {func_type}')
     if which == 'single_sweep':
         key = msg.single_sweep.parameter_key
+        metadata: DeviceParameter | Metadata | None
         if msg.single_sweep.HasField("parameter"):
             metadata = DeviceParameter(
                 path=msg.single_sweep.parameter.path,
@@ -204,6 +213,8 @@ def sweep_from_proto(
                     else None
                 ),
             )
+        elif msg.single_sweep.HasField("metadata"):
+            metadata = metadata_from_proto(msg.single_sweep.metadata)
         else:
             metadata = None
 
@@ -245,6 +256,38 @@ def sweep_from_proto(
     raise ValueError(f'sweep type not set: {msg}')  # pragma: no cover
 
 
+def metadata_to_proto(metadata: Metadata) -> run_context_pb2.Metadata:
+    """Convert the metadata dataclass to the metadata proto."""
+    device_parameters: list[run_context_pb2.DeviceParameter] = []
+    if params := getattr(metadata, "device_parameters", None):
+        for param in params:
+            path = getattr(param, "path", None)
+            idx = getattr(param, "idx", None)
+            device_parameters.append(run_context_pb2.DeviceParameter(path=path, idx=idx))
+
+    return run_context_pb2.Metadata(
+        device_parameters=device_parameters or None,  # If empty set this field as None.
+        label=metadata.label,
+        is_const=metadata.is_const,
+        unit=metadata.unit,
+    )
+
+
+def metadata_from_proto(metadata_pb: run_context_pb2.Metadata) -> Metadata:
+    """Convert the metadata proto to the metadata dataclass."""
+    device_parameters: list[DeviceParameter] = []
+    for param in metadata_pb.device_parameters:
+        device_parameters.append(
+            DeviceParameter(path=param.path, idx=param.idx if param.HasField("idx") else None)
+        )
+    return Metadata(
+        device_parameters=device_parameters or None,
+        label=metadata_pb.label if metadata_pb.HasField("label") else None,
+        is_const=metadata_pb.is_const,
+        unit=metadata_pb.unit if metadata_pb.HasField("unit") else None,
+    )
+
+
 def run_context_to_proto(
     sweepable: cirq.Sweepable, repetitions: int, *, out: Optional[run_context_pb2.RunContext] = None
 ) -> run_context_pb2.RunContext:

cirq-google/cirq_google/api/v2/run_context_pb2.pyi

@@ -22,7 +22,7 @@
 
 import cirq
 from cirq.study import sweeps
-from cirq_google.study import DeviceParameter
+from cirq_google.study import DeviceParameter, Metadata
 from cirq_google.api import v2
 
 
@@ -55,6 +55,27 @@ def _values(self) -> Iterator[float]:
             [1, 1.5, 2, 2.5, 3],
             metadata=DeviceParameter(path=['path', 'to', 'parameter'], idx=2, units='GHz'),
         ),
+        cirq.Points(
+            'a',
+            [1, 1.5, 2, 2.5, 3],
+            metadata=Metadata(
+                device_parameters=[DeviceParameter(path=['path', 'to', 'parameter'], idx=2)],
+                label="bb",
+            ),
+        ),
+        cirq.Points(
+            'a',
+            [1],
+            metadata=Metadata(
+                device_parameters=[
+                    DeviceParameter(path=['path', 'to', 'parameter']),
+                    DeviceParameter(path=['path', 'to', 'parameter2']),
+                ],
+                label="bb",
+                is_const=True,
+            ),
+        ),
+        cirq.Linspace('a', 0, 10, 100, metadata=Metadata(is_const=True)),
         cirq.Points(
             'b',
             [1, 1.5, 2, 2.5, 3],

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

@@ -78,5 +78,6 @@ def _old_xmon(*args, **kwargs):
         'cirq.google.GridDevice': cirq_google.GridDevice,
         'cirq.google.GoogleCZTargetGateset': cirq_google.GoogleCZTargetGateset,
         'cirq.google.DeviceParameter': cirq_google.study.device_parameter.DeviceParameter,
+        'cirq.google.Metadata': cirq_google.study.device_parameter.Metadata,
         'InternalGate': cirq_google.InternalGate,
     }

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

@@ -0,0 +1,18 @@
+{
+    "cirq_type":"cirq.google.Metadata",
+    "device_parameters":[
+       {
+          "cirq_type":"cirq.google.DeviceParameter",
+          "path":[
+             "test",
+             "subdir"
+          ],
+          "idx":null,
+          "value":null,
+          "units":null
+       }
+    ],
+    "is_const":true,
+    "label":"fake_label",
+    "unit":null
+}