Add initial Engine client for Calibration API (#3474)

* Add initial Engine client for Calibration API

- Adds run_calibration to Engine and EngineProgram.
- Adds calibration_results() to EngineJob to retrieve results.
- Adds CalibrationLayer and CalibrationResult container data classes
for sending and storing results, respectively.

cirq/google/__init__.py

@@ -28,6 +28,8 @@
 
 from cirq.google.engine import (
     Calibration,
+    CalibrationLayer,
+    CalibrationResult,
     Engine,
     engine_from_environment,
     EngineJob,

cirq/google/api/v2/__init__.py

@@ -15,6 +15,7 @@
 
 from cirq.google.api.v2 import (
     batch_pb2,
+    calibration_pb2,
     device_pb2,
     metrics_pb2,
     program_pb2,

cirq/google/arg_func_langs.py

@@ -85,11 +85,11 @@ def _function_languages_from_arg(arg_proto: v2.program_pb2.Arg
                 yield from _function_languages_from_arg(a)
 
 
-def _arg_to_proto(value: ARG_LIKE,
-                  *,
-                  arg_function_language: Optional[str],
-                  out: Optional[v2.program_pb2.Arg] = None
-                 ) -> v2.program_pb2.Arg:
+def arg_to_proto(value: ARG_LIKE,
+                 *,
+                 arg_function_language: Optional[str] = None,
+                 out: Optional[v2.program_pb2.Arg] = None
+                ) -> v2.program_pb2.Arg:
     """Writes an argument value into an Arg proto.
 
     Args:
@@ -133,21 +133,21 @@ def check_support(func_type: str) -> str:
     elif isinstance(value, sympy.Add):
         msg.func.type = check_support('add')
         for arg in value.args:
-            _arg_to_proto(arg,
-                          arg_function_language=arg_function_language,
-                          out=msg.func.args.add())
+            arg_to_proto(arg,
+                         arg_function_language=arg_function_language,
+                         out=msg.func.args.add())
     elif isinstance(value, sympy.Mul):
         msg.func.type = check_support('mul')
         for arg in value.args:
-            _arg_to_proto(arg,
-                          arg_function_language=arg_function_language,
-                          out=msg.func.args.add())
+            arg_to_proto(arg,
+                         arg_function_language=arg_function_language,
+                         out=msg.func.args.add())
     elif isinstance(value, sympy.Pow):
         msg.func.type = check_support('pow')
         for arg in value.args:
-            _arg_to_proto(arg,
-                          arg_function_language=arg_function_language,
-                          out=msg.func.args.add())
+            arg_to_proto(arg,
+                         arg_function_language=arg_function_language,
+                         out=msg.func.args.add())
     else:
         raise ValueError(f'Unrecognized arg type: {type(value)}')
 

cirq/google/arg_func_langs_test.py

@@ -21,7 +21,7 @@
 import cirq
 from cirq.google.arg_func_langs import (
     arg_from_proto,
-    _arg_to_proto,
+    arg_to_proto,
     ARG_LIKE,
     LANGUAGE_ORDER,
 )
@@ -94,13 +94,13 @@ def test_correspondence(min_lang: str, value: ARG_LIKE,
         if i < min_i:
             with pytest.raises(ValueError,
                                match='not supported by arg_function_language'):
-                _ = _arg_to_proto(value, arg_function_language=lang)
+                _ = arg_to_proto(value, arg_function_language=lang)
             with pytest.raises(ValueError, match='Unrecognized function type'):
                 _ = arg_from_proto(msg, arg_function_language=lang)
         else:
             parsed = arg_from_proto(msg, arg_function_language=lang)
             packed = json_format.MessageToDict(
-                _arg_to_proto(value, arg_function_language=lang),
+                arg_to_proto(value, arg_function_language=lang),
                 including_default_value_fields=True,
                 preserving_proto_field_name=True,
                 use_integers_for_enums=True)
@@ -121,7 +121,7 @@ def test_double_value():
 
 
 def test_serialize_sympy_constants():
-    proto = _arg_to_proto(sympy.pi, arg_function_language='')
+    proto = arg_to_proto(sympy.pi, arg_function_language='')
     packed = json_format.MessageToDict(proto,
                                        including_default_value_fields=True,
                                        preserving_proto_field_name=True,
@@ -136,7 +136,7 @@ def test_serialize_sympy_constants():
 
 def test_unsupported_function_language():
     with pytest.raises(ValueError, match='Unrecognized arg_function_language'):
-        _ = _arg_to_proto(1, arg_function_language='NEVER GONNAH APPEN')
+        _ = arg_to_proto(1, arg_function_language='NEVER GONNAH APPEN')
     with pytest.raises(ValueError, match='Unrecognized arg_function_language'):
         _ = arg_from_proto(None, arg_function_language='NEVER GONNAH APPEN')
 
@@ -160,8 +160,8 @@ def test_unsupported_function_language():
 def test_serialize_conversion(value: ARG_LIKE, proto: v2.program_pb2.Arg):
     msg = v2.program_pb2.Arg()
     json_format.ParseDict(proto, msg)
-    packed = json_format.MessageToDict(_arg_to_proto(value,
-                                                     arg_function_language=''),
+    packed = json_format.MessageToDict(arg_to_proto(value,
+                                                    arg_function_language=''),
                                        including_default_value_fields=True,
                                        preserving_proto_field_name=True,
                                        use_integers_for_enums=True)

cirq/google/engine/__init__.py

@@ -18,6 +18,10 @@
 from cirq.google.engine.calibration import (
     Calibration,)
 
+from cirq.google.engine.calibration_layer import (
+    CalibrationLayer,)
+from cirq.google.engine.calibration_result import (
+    CalibrationResult,)
 from cirq.google.engine.engine import (
     Engine,
     get_engine,

cirq/google/engine/calibration_layer.py

@@ -0,0 +1,29 @@
+# Copyright 2020 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import dataclasses
+from typing import Dict, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import cirq
+
+
+@dataclasses.dataclass
+class CalibrationLayer:
+    """Python implementation of the proto found in
+    cirq.google.api.v2.calibration_pb2.FocusedCalibrationLayer for use
+    in Engine calls."""
+    calibration_type: str
+    program: 'cirq.Circuit'
+    args: Dict[str, Union[str, float]]

cirq/google/engine/calibration_result.py

@@ -0,0 +1,39 @@
+# Copyright 2020 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import dataclasses
+import datetime
+from typing import Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import cirq
+    import cirq.google.api.v2.calibration_pb2 as calibration_pb2
+
+
+@dataclasses.dataclass
+class CalibrationResult:
+    """Python implementation of the proto found in
+    cirq.google.api.v2.calibration_pb2.CalibrationLayerResult for use
+    in Engine calls.
+
+    Note that, if these fields are not filled out by the calibration API,
+    they will be set to the default values in the proto, as defined here:
+    https://developers.google.com/protocol-buffers/docs/proto3#default
+    These defaults will converted to `None` by the API client.
+    """
+    code: 'calibration_pb2.CalibrationLayerCode'
+    error_message: Optional[str]
+    token: Optional[str]
+    valid_until: Optional[datetime.datetime]
+    metrics: 'cirq.google.Calibration'

cirq/google/engine/engine.py

@@ -33,9 +33,11 @@
 
 from google.protobuf import any_pb2
 from cirq.google.engine.client import quantum
+from cirq.google.engine.result_type import ResultType
 from cirq import circuits, study, value
 from cirq.google import serializable_gate_set as sgs
 from cirq.google.api import v2
+from cirq.google.arg_func_langs import arg_to_proto
 from cirq.google.engine import (engine_client, engine_program, engine_job,
                                 engine_processor, engine_sampler)
 
@@ -360,6 +362,77 @@ def run_batch(
                                         description=job_description,
                                         labels=job_labels)
 
+    def run_calibration(
+            self,
+            layers: List['cirq.google.CalibrationLayer'],
+            program_id: Optional[str] = None,
+            job_id: Optional[str] = None,
+            processor_id: str = None,
+            processor_ids: Sequence[str] = (),
+            gate_set: Optional[sgs.SerializableGateSet] = None,
+            program_description: Optional[str] = None,
+            program_labels: Optional[Dict[str, str]] = None,
+            job_description: Optional[str] = None,
+            job_labels: Optional[Dict[str, str]] = None,
+    ) -> engine_job.EngineJob:
+        """Runs the specified calibrations via the Calibration API.
+
+        Each calibration will be specified by a `CalibrationLayer`
+        that contains the type of the calibrations to run, a `Circuit`
+        to optimize, and any arguments needed by the calibration routine.
+
+        Arguments and circuits needed for each layer will vary based on the
+        calibration type.  However, the typical calibration routine may
+        require a single moment defining the gates to optimize, for example.
+
+        Note: this is an experimental API and is not yet fully supported
+        for all users.
+
+        Args:
+            layers: The layers of calibration to execute as a batch.
+            program_id: A user-provided identifier for the program. This must
+                be unique within the Google Cloud project being used. If this
+                parameter is not provided, a random id of the format
+                'calibration-################YYMMDD' will be generated,
+                where # is alphanumeric and YYMMDD is the current year, month,
+                and day.
+            job_id: Job identifier to use. If this is not provided, a random id
+                of the format 'calibration-################YYMMDD' will be
+                generated, where # is alphanumeric and YYMMDD is the current
+                year, month, and day.
+            processor_id: The engine processor that should run the calibration.
+                If this is specified, processor_ids should not be specified.
+            processor_ids: The engine processors that should be candidates
+                to run the program. Only one of these will be scheduled for
+                execution.
+            gate_set: The gate set used to serialize the circuit. The gate set
+                must be supported by the selected processor.
+            program_description: An optional description to set on the program.
+            program_labels: Optional set of labels to set on the program.
+            job_description: An optional description to set on the job.
+            job_labels: Optional set of labels to set on the job.  By defauly,
+                this will add a 'calibration' label to the job.
+
+        Returns:
+            An EngineJob whose results can be retrieved by calling
+            calibration_results().
+        """
+        if processor_id and processor_ids:
+            raise ValueError('Only one of processor_id and processor_ids '
+                             'can be specified.')
+        if not processor_ids and not processor_id:
+            raise ValueError('Processor id must be specified.')
+        if processor_id:
+            processor_ids = [processor_id]
+        if job_labels is None:
+            job_labels = {'calibration': ''}
+        engine_program = self.create_calibration_program(
+            layers, program_id, gate_set, program_description, program_labels)
+        return engine_program.run_calibration(job_id=job_id,
+                                              processor_ids=processor_ids,
+                                              description=job_description,
+                                              labels=job_labels)
+
     def create_program(
             self,
             program: 'cirq.Circuit',
@@ -446,7 +519,62 @@ def create_batch_program(
                                             new_program_id,
                                             self.context,
                                             new_program,
-                                            batch_mode=True)
+                                            result_type=ResultType.Batch)
+
+    def create_calibration_program(
+            self,
+            layers: List['cirq.google.CalibrationLayer'],
+            program_id: Optional[str] = None,
+            gate_set: Optional[sgs.SerializableGateSet] = None,
+            description: Optional[str] = None,
+            labels: Optional[Dict[str, str]] = None,
+    ) -> engine_program.EngineProgram:
+        """Wraps a list of calibration layers into an Any for Quantum Engine.
+
+        Args:
+            layers: The calibration routines to execute.  All layers will be
+                executed within the same API call in the order specified,
+                though some layers may be interleaved together using
+                hardware-specific batching.
+            program_id: A user-provided identifier for the program. This must be
+                unique within the Google Cloud project being used. If this
+                parameter is not provided, a random id of the format
+                'calibration-################YYMMDD' will be generated,
+                where # is alphanumeric and YYMMDD is the current year, month,
+                and day.
+            gate_set: The gate set used to serialize the circuits in each
+                layer.  The gate set must be supported by the processor.
+            description: An optional description to set on the program.
+            labels: Optional set of labels to set on the program.
+
+        Returns:
+            A EngineProgram for the newly created program.
+        """
+        if not gate_set:
+            raise ValueError('Gate set must be specified.')
+        if not program_id:
+            program_id = _make_random_id('calibration-')
+
+        calibration = v2.calibration_pb2.FocusedCalibration()
+        for layer in layers:
+            new_layer = calibration.layers.add()
+            new_layer.calibration_type = layer.calibration_type
+            for arg in layer.args:
+                arg_to_proto(layer.args[arg], out=new_layer.args[arg])
+            gate_set.serialize(layer.program, msg=new_layer.layer)
+
+        new_program_id, new_program = self.context.client.create_program(
+            self.project_id,
+            program_id,
+            code=self._pack_any(calibration),
+            description=description,
+            labels=labels)
+
+        return engine_program.EngineProgram(self.project_id,
+                                            new_program_id,
+                                            self.context,
+                                            new_program,
+                                            result_type=ResultType.Calibration)
 
     def _serialize_program(self, program: 'cirq.Circuit',
                            gate_set: sgs.SerializableGateSet) -> any_pb2.Any: