Add JSON serialization for Calibration metrics (#3508)

* Add JSON serialization for Calibration metrics

- Add json_dict and from_json_dict for serialization
of Calibration objects using the underlying proto representation.
- Add a valid repr for Calibration objects.

Author: dstrain115

cirq/google/engine/calibration.py

@@ -18,6 +18,7 @@
 
 from typing import Any, Dict, Iterator, Optional, Tuple, TYPE_CHECKING
 
+import google.protobuf.json_format as json_format
 from cirq import devices, vis
 from cirq.google.api import v2
 
@@ -45,14 +46,25 @@ class Calibration(abc.Mapping):
 
         `calibration['t1']`
 
+    This class can be instantiated either from a `MetricsSnapshot` proto
+    or from a dictionary of metric values.
+
     Attributes:
         timestamp: The time that this calibration was run, in milliseconds since
             the epoch.
     """
 
-    def __init__(self, calibration: v2.metrics_pb2.MetricsSnapshot) -> None:
+    def __init__(self,
+                 calibration: v2.metrics_pb2.MetricsSnapshot = v2.metrics_pb2.
+                 MetricsSnapshot(),
+                 metrics: Optional[
+                     Dict[str, Dict[Tuple['cirq.GridQubit', ...], Any]]] = None
+                ) -> None:
         self.timestamp = calibration.timestamp_ms
-        self._metric_dict = self._compute_metric_dict(calibration.metrics)
+        if metrics is None:
+            self._metric_dict = self._compute_metric_dict(calibration.metrics)
+        else:
+            self._metric_dict = metrics
 
     def _compute_metric_dict(
             self, metrics: v2.metrics_pb2.MetricsSnapshot
@@ -103,8 +115,47 @@ def __len__(self) -> int:
         return len(self._metric_dict)
 
     def __str__(self) -> str:
-
-        return 'Calibration(keys={})'.format(list(sorted(self.keys())))
+        return f'Calibration(keys={list(sorted(self.keys()))})'
+
+    def __repr__(self) -> str:
+        return ('cirq.google.Calibration(metrics='
+                f'{repr(dict(self._metric_dict))})')
+
+    def to_proto(self) -> v2.metrics_pb2.MetricsSnapshot:
+        """Reconstruct the protobuf message represented by this class."""
+        proto = v2.metrics_pb2.MetricsSnapshot()
+        for key in self._metric_dict:
+            for target, value_list in self._metric_dict[key].items():
+                current_metric = proto.metrics.add()
+                current_metric.name = key
+                current_metric.targets.extend(
+                    [v2.qubit_to_proto_id(q) for q in target])
+                for value in value_list:
+                    current_value = current_metric.values.add()
+                    if isinstance(value, float):
+                        current_value.double_val = value
+                    elif isinstance(value, int):
+                        current_value.int64_val = value
+                    elif isinstance(value, str):
+                        current_value.str_val = value
+                    else:
+                        raise ValueError(f'Unsupported metric value {value}. '
+                                         'Must be int, float, or str to '
+                                         'convert to proto.')
+        return proto
+
+    @classmethod
+    def _from_json_dict_(cls, metrics: str, **kwargs) -> 'Calibration':
+        """Magic method for the JSON serialization protocol."""
+        metric_proto = v2.metrics_pb2.MetricsSnapshot()
+        return cls(json_format.ParseDict(metrics, metric_proto))
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        """Magic method for the JSON serialization protocol."""
+        return {
+            'cirq_type': 'Calibration',
+            'metrics': json_format.MessageToDict(self.to_proto())
+        }
 
     def timestamp_str(self,
                       tz: Optional[datetime.tzinfo] = None,

cirq/google/engine/calibration_test.py

@@ -94,6 +94,11 @@ def test_calibration_str():
                                 "'xeb'])")
 
 
+def test_calibration_repr():
+    calibration = cg.Calibration(_CALIBRATION_DATA)
+    cirq.testing.assert_equivalent_repr(calibration)
+
+
 def test_calibration_timestamp_str():
     calibration = cg.Calibration(_CALIBRATION_DATA)
     assert (calibration.timestamp_str(
@@ -103,6 +108,17 @@ def test_calibration_timestamp_str():
             hours=1))) == '2019-07-08 01:00:00.021021+01:00')
 
 
+def test_to_proto():
+    calibration = cg.Calibration(_CALIBRATION_DATA)
+    assert calibration == cg.Calibration(calibration.to_proto())
+    invalid_value = cg.Calibration(
+        metrics={'metric': {
+            (cirq.GridQubit(1, 1),): [1.1, {}]
+        }})
+    with pytest.raises(ValueError, match='Unsupported metric value'):
+        invalid_value.to_proto()
+
+
 def test_calibration_heatmap():
     calibration = cg.Calibration(_CALIBRATION_DATA)
 

cirq/protocols/json_serialization.py

@@ -86,6 +86,7 @@ def two_qubit_matrix_gate(matrix):
                 'CCXPowGate': cirq.CCXPowGate,
                 'CCZPowGate': cirq.CCZPowGate,
                 'CNotPowGate': cirq.CNotPowGate,
+                'Calibration': cirq.google.Calibration,
                 'CalibrationLayer': cirq.google.CalibrationLayer,
                 'CalibrationResult': cirq.google.CalibrationResult,
                 'CalibrationTag': cirq.google.CalibrationTag,

cirq/protocols/json_serialization_test.py

@@ -276,7 +276,6 @@ def test_mutually_exclusive_blacklist():
 NOT_YET_SERIALIZABLE = [
     'AsymmetricDepolarizingChannel',
     'AxisAngleDecomposition',
-    'Calibration',
     'CalibrationLayer',
     'CalibrationResult',
     'CircuitDag',

cirq/protocols/json_test_data/Calibration.json

@@ -0,0 +1,72 @@
+{
+  "cirq_type": "Calibration",
+  "metrics": {
+    "metrics": [
+      {
+        "name": "xeb",
+        "targets": [
+          "0_0",
+          "0_1"
+        ],
+        "values": [
+          {
+            "doubleVal": 0.9999
+          }
+        ]
+      },
+      {
+        "name": "xeb",
+        "targets": [
+          "0_0",
+          "1_0"
+        ],
+        "values": [
+          {
+            "doubleVal": 0.9998
+          }
+        ]
+      },
+      {
+        "name": "t1",
+        "targets": [
+          "0_0"
+        ],
+        "values": [
+          {
+            "int64Val": "321"
+          }
+        ]
+      },
+      {
+        "name": "t1",
+        "targets": [
+          "0_1"
+        ],
+        "values": [
+          {
+            "int64Val": "911"
+          }
+        ]
+      },
+      {
+        "name": "t1",
+        "targets": [
+          "1_0"
+        ],
+        "values": [
+          {
+            "int64Val": "505"
+          }
+        ]
+      },
+      {
+        "name": "globalMetric",
+        "values": [
+          {
+            "strVal": "abcd"
+          }
+        ]
+      }
+    ]
+  }
+}

cirq/protocols/json_test_data/Calibration.repr

@@ -0,0 +1,13 @@
+cirq.google.Calibration(
+    metrics={
+        'xeb': {
+            (cirq.GridQubit(0, 0), cirq.GridQubit(0, 1)): [0.9999],
+            (cirq.GridQubit(0, 0), cirq.GridQubit(1, 0)): [0.9998]
+            },
+        't1': {
+            (cirq.GridQubit(0, 0),): [321],
+            (cirq.GridQubit(0, 1),): [911],
+            (cirq.GridQubit(1, 0),): [505]},
+        'globalMetric': {(): ['abcd']}
+    }
+)