Use v2 protos in cirq.google.Engine

cirq/google/api/v2/results.py

@@ -1,6 +1,7 @@
 from typing import (Dict, Iterable, Iterator, List, NamedTuple, Optional, Set,
                     Union, cast)
 
+from collections import OrderedDict
 import numpy as np
 
 from cirq.api.google.v2 import result_pb2
@@ -138,10 +139,9 @@ def results_to_proto(
                     qmr.results = pack_bits(m_data[:, i])
     return out
 
-
 def results_from_proto(
         msg: result_pb2.Result,
-        measurements: List[MeasureInfo],
+        measurements: List[MeasureInfo] = None,
 ) -> List[List[study.TrialResult]]:
     """Converts a v2 result proto into List of list of trial results.
 
@@ -152,7 +152,8 @@ def results_from_proto(
     Returns:
         A list containing a list of trial results for each sweep.
     """
-    measure_map = {m.key: m for m in measurements}
+
+    measure_map = {m.key: m for m in measurements} if measurements else None
     return [
         _trial_sweep_from_proto(sweep_result, measure_map)
         for sweep_result in msg.sweep_results
@@ -161,20 +162,23 @@ def results_from_proto(
 
 def _trial_sweep_from_proto(
         msg: result_pb2.SweepResult,
-        measurements: Dict[str, MeasureInfo],
+        measurements: Dict[str, MeasureInfo] = None,
 ) -> List[study.TrialResult]:
     trial_sweep: List[study.TrialResult] = []
     for pr in msg.parameterized_results:
         m_data: Dict[str, np.ndarray] = {}
         for mr in pr.measurement_results:
-            m = measurements[mr.key]
-            qubit_results: Dict[devices.GridQubit, np.ndarray] = {}
+            qubit_results: OrderedDict[devices.GridQubit, np.ndarray] = {}
             for qmr in mr.qubit_measurement_results:
                 qubit = devices.GridQubit.from_proto_id(qmr.qubit.id)
                 if qubit in qubit_results:
                     raise ValueError('qubit already exists: {}'.format(qubit))
                 qubit_results[qubit] = unpack_bits(qmr.results, msg.repetitions)
-            ordered_results = [qubit_results[qubit] for qubit in m.qubits]
+            if measurements:
+                ordered_results = [qubit_results[qubit]
+                                   for qubit in measurements[mr.key].qubits]
+            else:
+                ordered_results = list(qubit_results.values())
             m_data[mr.key] = np.array(ordered_results).transpose()
         trial_sweep.append(
             study.TrialResult(

cirq/google/api/v2/results_test.py

@@ -221,3 +221,39 @@ def test_results_from_proto_duplicate_qubit():
         qmr.results = bytes([results])
     with pytest.raises(ValueError, match='qubit already exists'):
         v2.results_from_proto(proto, measurements)
+
+
+def test_results_from_proto_default_ordering():
+    proto = result_pb2.Result()
+    sr = proto.sweep_results.add()
+    sr.repetitions = 8
+    pr = sr.parameterized_results.add()
+    pr.params.assignments.update({'i': 1})
+    mr = pr.measurement_results.add()
+    mr.key = 'foo'
+    for qubit, results in [
+        (q(0, 1), 0b1100_1100),
+        (q(1, 1), 0b1010_1010),
+        (q(0, 0), 0b1111_0000),
+    ]:
+        qmr = mr.qubit_measurement_results.add()
+        qmr.qubit.id = qubit.proto_id()
+        qmr.results = bytes([results])
+
+    trial_results = v2.results_from_proto(proto)
+    trial = trial_results[0][0]
+    assert trial.params == cirq.ParamResolver({'i': 1})
+    assert trial.repetitions == 8
+    np.testing.assert_array_equal(
+        trial.measurements['foo'],
+        np.array([
+            [0, 0, 0],
+            [0, 1, 0],
+            [1, 0, 0],
+            [1, 1, 0],
+            [0, 0, 1],
+            [0, 1, 1],
+            [1, 0, 1],
+            [1, 1, 1],
+        ],
+                 dtype=bool))

cirq/google/engine/engine.py

@@ -24,31 +24,45 @@
 """
 
 import base64
+import enum
 import random
 import re
 import string
 import time
 import urllib.parse
 from collections import Iterable
-from typing import cast, Dict, List, Optional, Sequence, TYPE_CHECKING, Union
+from typing import Any, cast, Dict, List, Optional, Sequence, Tuple, Union
+
+import google.protobuf as gp
 from apiclient import discovery
+from google.protobuf import any_pb2
 
 from cirq import optimizers, circuits
+from cirq.api.google import v1, v2
+from cirq.google.api import v2 as api_v2
 from cirq.google.convert_to_xmon_gates import ConvertToXmonGates
 from cirq.google.params import sweep_to_proto_dict
 from cirq.google.programs import schedule_to_proto_dicts, unpack_results
+from cirq.google.serializable_gate_set import SerializableGateSet
 from cirq.schedules import Schedule, moment_by_moment_schedule
 from cirq.study import ParamResolver, Sweep, Sweepable, TrialResult
 from cirq.study.sweeps import Points, UnitSweep, Zip
 
-if TYPE_CHECKING:
-    # pylint: disable=unused-import
-    from typing import Any
-
 gcs_prefix_pattern = re.compile('gs://[a-z0-9._/-]+')
 TERMINAL_STATES = ['SUCCESS', 'FAILURE', 'CANCELLED']
 
 
+class ProtoVersion(enum.Enum):
+    """Protocol buffer version to use for requests to the quantum engine."""
+    UNDEFINED = 0
+    V1 = 1
+    V2 = 2
+
+
+# Quantum programs to run can be specified as circuits or schedules.
+Program = Union[circuits.Circuit, Schedule]
+
+
 class JobConfig:
     """Configuration for a program and job to run on the Quantum Engine API.
 
@@ -111,11 +125,11 @@ def __repr__(self):
                 'gcs_prefix={!r}, '
                 'gcs_program={!r}, '
                 'gcs_results={!r})').format(self.project_id,
-                                             self.program_id,
-                                             self.job_id,
-                                             self.gcs_prefix,
-                                             self.gcs_program,
-                                             self.gcs_results)
+                                            self.program_id,
+                                            self.job_id,
+                                            self.gcs_prefix,
+                                            self.gcs_program,
+                                            self.gcs_results)
 
 
 class Engine:
@@ -149,8 +163,9 @@ def __init__(self,
                  default_project_id: Optional[str] = None,
                  discovery_url: Optional[str] = None,
                  default_gcs_prefix: Optional[str] = None,
-                 **kwargs
-                 ) -> None:
+                 proto_version: ProtoVersion = ProtoVersion.V1,
+                 gate_set: Optional[SerializableGateSet] = None,
+                 **kwargs) -> None:
         """Engine service client.
 
         Args:
@@ -174,11 +189,13 @@ def __init__(self,
                                                '$discovery/rest'
                                                '?version={apiVersion}')
         self.default_gcs_prefix = default_gcs_prefix
+        self.proto_version = proto_version
+        self.gate_set = gate_set
 
         discovery_service_url = (
             self.discovery_url if self.api_key is None else (
                 "%s&key=%s" % (self.discovery_url, urllib.parse.quote_plus(
-                               self.api_key))))
+                    self.api_key))))
         self.service = discovery.build(
             self.api,
             self.version,
@@ -188,7 +205,7 @@ def __init__(self,
     def run(
             self,
             *,  # Force keyword args.
-            program: Union[circuits.Circuit, Schedule],
+            program: Program,
             job_config: Optional[JobConfig] = None,
             param_resolver: ParamResolver = ParamResolver({}),
             repetitions: int = 1,
@@ -298,9 +315,7 @@ def implied_job_config(self, job_config: Optional[JobConfig]) -> JobConfig:
 
         return implied_job_config
 
-    def program_as_schedule(self,
-                            program: Union[circuits.Circuit,
-                                           Schedule]) -> Schedule:
+    def program_as_schedule(self, program: Program) -> Schedule:
         if isinstance(program, circuits.Circuit):
             device = program.device
             circuit_copy = program.copy()
@@ -317,7 +332,7 @@ def program_as_schedule(self,
     def run_sweep(
             self,
             *,  # Force keyword args.
-            program: Union[circuits.Circuit, Schedule],
+            program: Program,
             job_config: Optional[JobConfig] = None,
             params: Sweepable = None,
             repetitions: int = 1,
@@ -343,26 +358,23 @@ def run_sweep(
         """
 
         job_config = self.implied_job_config(job_config)
-        schedule = self.program_as_schedule(program)
 
         # Check program to run and program parameters.
         if not 0 <= priority < 1000:
             raise ValueError('priority must be between 0 and 1000')
 
-        schedule.device.validate_schedule(schedule)
-
-        # Create program.
         sweeps = _sweepable_to_sweeps(params or ParamResolver({}))
-        program_dict = {}  # type: Dict[str, Any]
+        if self.proto_version == ProtoVersion.V1:
+            code, run_context = self._serialize_program_v1(
+                program, sweeps, repetitions)
+        elif self.proto_version == ProtoVersion.V2:
+            code, run_context = self._serialize_program_v2(
+                program, sweeps, repetitions)
+        else:
+            raise ValueError('invalid proto version: {}'.format(
+                self.proto_version))
 
-        program_dict['parameter_sweeps'] = [
-            sweep_to_proto_dict(sweep, repetitions) for
-            sweep in sweeps]
-        program_dict['operations'] = [op for op in
-                                      schedule_to_proto_dicts(schedule)]
-        code = {
-            '@type': 'type.googleapis.com/cirq.api.google.v1.Program'}
-        code.update(program_dict)
+        # Create program.
         request = {
             'name': 'projects/%s/programs/%s' % (job_config.project_id,
                                                  job_config.program_id,),
@@ -392,11 +404,52 @@ def run_sweep(
                 }
             },
         }
+        if run_context is not None:
+            request['run_context'] = run_context
         response = self.service.projects().programs().jobs().create(
             parent=response['name'], body=request).execute()
 
         return EngineJob(job_config, response, self)
 
+    def _serialize_program_v1(
+            self, program: Program, sweeps: List[Sweep], repetitions: int
+    ) -> Tuple[Dict[str, Any], Optional[Dict[str, Any]]]:
+        schedule = self.program_as_schedule(program)
+        schedule.device.validate_schedule(schedule)
+
+        descriptor = v1.program_pb2.Program.DESCRIPTOR
+
+        program_dict = {}  # type: Dict[str, Any]
+        program_dict['@type'] = 'type.googleapis.com/' + descriptor.full_name
+        program_dict['parameter_sweeps'] = [
+            sweep_to_proto_dict(sweep, repetitions) for sweep in sweeps
+        ]
+        program_dict['operations'] = [
+            op for op in schedule_to_proto_dicts(schedule)
+        ]
+        return program_dict, None  # run context included in program
+
+    def _serialize_program_v2(
+            self, program: Program, sweeps: List[Sweep], repetitions: int
+    ) -> Tuple[Dict[str, Any], Optional[Dict[str, Any]]]:
+        if isinstance(program, Schedule):
+            program.device.validate_schedule(program)
+
+        program = self.gate_set.serialize(program)
+
+        run_context = v2.run_context_pb2.RunContext()
+        for sweep in sweeps:
+            sweep_proto = run_context.parameter_sweeps.add()
+            sweep_proto.repetitions = repetitions
+            api_v2.sweep_to_proto(sweep, out=sweep_proto.sweep)
+
+        def any_dict(message: gp.message.Message) -> Dict[str, Any]:
+            any_message = any_pb2.Any()
+            any_message.Pack(message)
+            return gp.json_format.MessageToDict(any_message)
+
+        return any_dict(program), any_dict(run_context)
+
     def get_program(self, program_resource_name: str) -> Dict:
         """Returns the previously created quantum program.
 
@@ -439,8 +492,16 @@ def get_job_results(self, job_resource_name: str) -> List[TrialResult]:
         """
         response = self.service.projects().programs().jobs().getResult(
             parent=job_resource_name).execute()
+        if self.proto_version == ProtoVersion.V1:
+            return self._get_job_results_v1(response['result'])
+        if self.proto_version == ProtoVersion.V2:
+            return self._get_job_results_v2(response['result'])
+        raise ValueError('invalid proto version: {}'.format(
+            self.proto_version))
+
+    def _get_job_results_v1(self, result: Dict[str, Any]) -> List[TrialResult]:
         trial_results = []
-        for sweep_result in response['result']['sweepResults']:
+        for sweep_result in result['sweepResults']:
             sweep_repetitions = sweep_result['repetitions']
             key_sizes = [(m['key'], len(m['qubits']))
                          for m in sweep_result['measurementKeys']]
@@ -449,13 +510,27 @@ def get_job_results(self, job_resource_name: str) -> List[TrialResult]:
                 measurements = unpack_results(data, sweep_repetitions,
                                               key_sizes)
 
-                trial_results.append(TrialResult(
-                    params=ParamResolver(
+                trial_results.append(
+                    TrialResult(params=ParamResolver(
                         result.get('params', {}).get('assignments', {})),
-                    repetitions=sweep_repetitions,
-                    measurements=measurements))
+                                repetitions=sweep_repetitions,
+                                measurements=measurements))
         return trial_results
 
+    def _get_job_results_v2(self,
+                            result_dict: Dict[str, Any]) -> List[TrialResult]:
+        result_any = any_pb2.Any()
+        gp.json_format.ParseDict(result_dict, result_any)
+        result = v2.result_pb2.Result()
+        result_any.Unpack(result)
+
+        sweep_results = api_v2.results_from_proto(result)
+        # Flatten to single list to match to sampler api.
+        return [
+            trial_result for sweep_result in sweep_results
+            for trial_result in sweep_result
+        ]
+
     def cancel_job(self, job_resource_name: str):
         """Cancels the given job.
 
@@ -581,8 +656,9 @@ def cancel(self):
         """Cancel the job."""
         self._engine.cancel_job(self.job_resource_name)
 
-    def results(self) -> List[TrialResult]:
-        """Returns the job results, blocking until the job is complete."""
+    def results(self)-> List[TrialResult]:
+        """Returns the job results, blocking until the job is complete.
+        """
         if not self._results:
             job = self._update_job()
             for _ in range(1000):
@@ -612,14 +688,12 @@ def _sweepable_to_sweeps(sweepable: Sweepable) -> List[Sweep]:
         if isinstance(next(iter(iterable)), Sweep):
             sweeps = iterable
             return list(sweeps)
-        else:
-            resolvers = iterable
-            return [_resolver_to_sweep(p) for p in resolvers]
-
+        resolvers = iterable
+        return [_resolver_to_sweep(p) for p in resolvers]
     raise TypeError('Unexpected Sweepable.')  # coverage: ignore
 
 
 def _resolver_to_sweep(resolver: ParamResolver) -> Sweep:
     return Zip(*[Points(key, [value]) for key, value in
                  resolver.param_dict.items()]) if len(
-        resolver.param_dict) else UnitSweep
+                     resolver.param_dict) else UnitSweep