feature: Add support for executing OpenQasmPrograms (#116)

* feature: Add support for executing OpenQasmProgram * test updates * Update branch for checks * integ tests * Result types calculation for openqasm programs (#164) * fix: Result types calculation for openqasm * Add tests * Update assert * Reformat Co-authored-by: Cody Wang <[email protected]> Co-authored-by: ℂ𝓞𝔇𝚈 <[email protected]>
amazon-braket · Jan 12, 2022 · c12c432 · c12c432
1 parent 18ee1a1
commit c12c432
Show file tree

Hide file tree

Showing 10 changed files with 280 additions and 45 deletions.
diff --git a/src/braket/aws/aws_device.py b/src/braket/aws/aws_device.py
@@ -28,6 +28,7 @@
 from braket.device_schema import DeviceCapabilities, GateModelQpuParadigmProperties
 from braket.device_schema.dwave import DwaveProviderProperties
 from braket.devices.device import Device
+from braket.ir.openqasm import Program as OpenQasmProgram
 from braket.schema_common import BraketSchemaBase
 
 
@@ -78,7 +79,7 @@ def __init__(self, arn: str, aws_session: Optional[AwsSession] = None):
 
     def run(
         self,
-        task_specification: Union[Circuit, Problem],
+        task_specification: Union[Circuit, Problem, OpenQasmProgram],
         s3_destination_folder: Optional[AwsSession.S3DestinationFolder] = None,
         shots: Optional[int] = None,
         poll_timeout_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_TIMEOUT,
@@ -160,7 +161,7 @@ def run(
 
     def run_batch(
         self,
-        task_specifications: List[Union[Circuit, Problem]],
+        task_specifications: List[Union[Circuit, Problem, OpenQasmProgram]],
         s3_destination_folder: Optional[AwsSession.S3DestinationFolder] = None,
         shots: Optional[int] = None,
         max_parallel: Optional[int] = None,

diff --git a/src/braket/aws/aws_quantum_task.py b/src/braket/aws/aws_quantum_task.py
@@ -40,6 +40,7 @@
 from braket.device_schema.ionq import IonqDeviceParameters
 from braket.device_schema.rigetti import RigettiDeviceParameters
 from braket.device_schema.simulators import GateModelSimulatorDeviceParameters
+from braket.ir.openqasm import Program as OpenQasmProgram
 from braket.schema_common import BraketSchemaBase
 from braket.task_result import AnnealingTaskResult, GateModelTaskResult
 from braket.tasks import AnnealingQuantumTaskResult, GateModelQuantumTaskResult, QuantumTask
@@ -62,7 +63,7 @@ class AwsQuantumTask(QuantumTask):
     def create(
         aws_session: AwsSession,
         device_arn: str,
-        task_specification: Union[Circuit, Problem],
+        task_specification: Union[Circuit, Problem, OpenQasmProgram],
         s3_destination_folder: AwsSession.S3DestinationFolder,
         shots: int,
         device_parameters: Dict[str, Any] = None,
@@ -411,6 +412,22 @@ def _create_internal(
     raise TypeError("Invalid task specification type")
 
 
+@_create_internal.register
+def _(
+    open_qasm_program: OpenQasmProgram,
+    aws_session: AwsSession,
+    create_task_kwargs: Dict[str, Any],
+    device_arn: str,
+    _device_parameters: Union[dict, BraketSchemaBase],  # Not currently used for OpenQasmProgram
+    _disable_qubit_rewiring,
+    *args,
+    **kwargs,
+) -> AwsQuantumTask:
+    create_task_kwargs.update({"action": open_qasm_program.json()})
+    task_arn = aws_session.create_quantum_task(**create_task_kwargs)
+    return AwsQuantumTask(task_arn, aws_session, *args, **kwargs)
+
+
 @_create_internal.register
 def _(
     circuit: Circuit,
@@ -461,7 +478,7 @@ def _(
         DwaveAdvantageDeviceParameters,
         Dwave2000QDeviceParameters,
     ],
-    disable_qubit_rewiring,
+    _,
     *args,
     **kwargs,
 ) -> AwsQuantumTask:

diff --git a/src/braket/aws/aws_quantum_task_batch.py b/src/braket/aws/aws_quantum_task_batch.py
@@ -21,6 +21,7 @@
 from braket.aws.aws_quantum_task import AwsQuantumTask
 from braket.aws.aws_session import AwsSession
 from braket.circuits import Circuit
+from braket.ir.openqasm import Program as OpenQasmProgram
 
 
 class AwsQuantumTaskBatch:
@@ -41,7 +42,7 @@ def __init__(
         self,
         aws_session: AwsSession,
         device_arn: str,
-        task_specifications: List[Union[Circuit, Problem]],
+        task_specifications: List[Union[Circuit, Problem, OpenQasmProgram]],
         s3_destination_folder: AwsSession.S3DestinationFolder,
         shots: int,
         max_parallel: int,

diff --git a/src/braket/tasks/gate_model_quantum_task_result.py b/src/braket/tasks/gate_model_quantum_task_result.py
@@ -278,8 +278,12 @@ def _from_object_internal_computational_basis_sampling(cls, result: GateModelTas
                 f"Measured qubits {measured_qubits} is not equivalent to number of qubits "
                 + f"{measurements.shape[1]} in measurements"
             )
-        result_types = GateModelQuantumTaskResult._calculate_result_types(
-            additional_metadata.action.json(), measurements, measured_qubits
+        result_types = (
+            result.resultTypes
+            if result.resultTypes
+            else GateModelQuantumTaskResult._calculate_result_types(
+                additional_metadata.action.json(), measurements, measured_qubits
+            )
         )
         values = [rt.value for rt in result_types]
         return cls(

diff --git a/test/integ_tests/gate_model_device_testing_utils.py b/test/integ_tests/gate_model_device_testing_utils.py
@@ -13,14 +13,15 @@
 
 import concurrent.futures
 import math
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 import numpy as np
 
 from braket.aws import AwsDevice
 from braket.circuits import Circuit, Gate, Instruction, Observable, ResultType
 from braket.circuits.quantum_operator_helpers import get_pauli_eigenvalues
 from braket.devices import Device
+from braket.ir.openqasm import Program as OpenQasmProgram
 from braket.tasks import GateModelQuantumTaskResult
 
 
@@ -41,11 +42,14 @@ def qubit_ordering_testing(device: Device, run_kwargs: Dict[str, Any]):
 
 
 def no_result_types_testing(
-    circuit: Circuit, device: Device, run_kwargs: Dict[str, Any], expected: Dict[str, float]
+    program: Union[Circuit, OpenQasmProgram],
+    device: Device,
+    run_kwargs: Dict[str, Any],
+    expected: Dict[str, float],
 ):
     shots = run_kwargs["shots"]
     tol = get_tol(shots)
-    result = device.run(circuit, **run_kwargs).result()
+    result = device.run(program, **run_kwargs).result()
     probabilities = result.measurement_probabilities
     for bitstring in probabilities:
         assert np.allclose(probabilities[bitstring], expected[bitstring], **tol)
@@ -569,6 +573,71 @@ def batch_bell_pair_testing(device: AwsDevice, run_kwargs: Dict[str, Any]):
     assert [task.result() for task in batch.tasks] == results
 
 
+def bell_pair_openqasm_testing(device: AwsDevice, run_kwargs: Dict[str, Any]):
+    openqasm_string = (
+        "OPENQASM 3;"
+        "qubit[2] q;"
+        "bit[2] c;"
+        "h q[0];"
+        "cnot q[0], q[1];"
+        "c[0] = measure q[0];"
+        "c[1] = measure q[1];"
+    )
+    no_result_types_testing(
+        OpenQasmProgram(source=openqasm_string), device, run_kwargs, {"00": 0.5, "11": 0.5}
+    )
+
+
+def openqasm_noisy_circuit_1qubit_noise_full_probability(
+    device: Device, run_kwargs: Dict[str, Any]
+):
+    shots = run_kwargs["shots"]
+    tol = get_tol(shots)
+    openqasm_string = (
+        "OPENQASM 3;"
+        "qubit[2] q;"
+        "x q[0];"
+        "x q[1];"
+        "#pragma braket noise bit_flip(0.1) q[0]"
+        "#pragma braket result probability q[0], q[1]"
+    )
+    result = device.run(OpenQasmProgram(source=openqasm_string), **run_kwargs).result()
+    assert len(result.result_types) == 1
+    assert np.allclose(
+        result.get_value_by_result_type(ResultType.Probability(target=[0, 1])),
+        np.array([0.0, 0.1, 0, 0.9]),
+        **tol
+    )
+
+
+def openqasm_result_types_bell_pair_testing(device: Device, run_kwargs: Dict[str, Any]):
+    openqasm_string = (
+        "OPENQASM 3;"
+        "qubit[2] q;"
+        "h q[0];"
+        "cnot q[0], q[1];"
+        "#pragma braket result expectation h(q[0]) @ x(q[1])"
+        "#pragma braket result sample h(q[0]) @ x(q[1])"
+    )
+    result = device.run(OpenQasmProgram(source=openqasm_string), **run_kwargs).result()
+    assert len(result.result_types) == 2
+    assert (
+        0.6
+        < result.get_value_by_result_type(
+            ResultType.Expectation(observable=Observable.H() @ Observable.X(), target=[0, 1])
+        )
+        < 0.8
+    )
+    assert (
+        len(
+            result.get_value_by_result_type(
+                ResultType.Sample(observable=Observable.H() @ Observable.X(), target=[0, 1])
+            )
+        )
+        == run_kwargs["shots"]
+    )
+
+
 def many_layers(n_qubits: int, n_layers: int) -> Circuit:
     """
     Function to return circuit with many layers.

diff --git a/test/integ_tests/test_simulator_quantum_task.py b/test/integ_tests/test_simulator_quantum_task.py
@@ -16,10 +16,13 @@
 import pytest
 from gate_model_device_testing_utils import (
     batch_bell_pair_testing,
+    bell_pair_openqasm_testing,
     get_tol,
     many_layers,
     multithreaded_bell_pair_testing,
     no_result_types_bell_pair_testing,
+    openqasm_noisy_circuit_1qubit_noise_full_probability,
+    openqasm_result_types_bell_pair_testing,
     qubit_ordering_testing,
     result_types_all_selected_testing,
     result_types_bell_pair_full_probability_testing,
@@ -218,6 +221,28 @@ def test_batch_bell_pair(simulator_arn, aws_session, s3_destination_folder):
 
 
 @pytest.mark.parametrize("simulator_arn", SIMULATOR_ARNS)
+def test_bell_pair_openqasm(simulator_arn, aws_session, s3_destination_folder):
+    device = AwsDevice(simulator_arn, aws_session)
+    bell_pair_openqasm_testing(
+        device, {"shots": SHOTS, "s3_destination_folder": s3_destination_folder}
+    )
+
+
+@pytest.mark.parametrize("simulator_arn", SIMULATOR_ARNS)
+def test_bell_pair_openqasm_results(simulator_arn, aws_session, s3_destination_folder):
+    device = AwsDevice(simulator_arn, aws_session)
+    openqasm_result_types_bell_pair_testing(
+        device, {"shots": SHOTS, "s3_destination_folder": s3_destination_folder}
+    )
+
+
+def test_openqasm_probability_results(aws_session, s3_destination_folder):
+    device = AwsDevice("arn:aws:braket:::device/quantum-simulator/amazon/dm1", aws_session)
+    openqasm_noisy_circuit_1qubit_noise_full_probability(
+        device, {"shots": SHOTS, "s3_destination_folder": s3_destination_folder}
+    )
+
+
 @pytest.mark.parametrize("num_layers", [50, 100, 500, 1000])
 def test_many_layers(simulator_arn, num_layers, aws_session, s3_destination_folder):
     num_qubits = 10

diff --git a/test/integ_tests/test_tensor_network_simulator.py b/test/integ_tests/test_tensor_network_simulator.py
@@ -15,7 +15,7 @@
 import random
 
 import pytest
-from gate_model_device_testing_utils import no_result_types_testing
+from gate_model_device_testing_utils import bell_pair_openqasm_testing, no_result_types_testing
 
 from braket.aws import AwsDevice
 from braket.circuits import Circuit
@@ -52,6 +52,14 @@ def test_qft_iqft_h(simulator_arn, aws_session, s3_destination_folder):
     )
 
 
+@pytest.mark.parametrize("simulator_arn", SIMULATOR_ARNS)
+def test_bell_pair_openqasm(simulator_arn, aws_session, s3_destination_folder):
+    device = AwsDevice(simulator_arn, aws_session)
+    bell_pair_openqasm_testing(
+        device, {"shots": SHOTS, "s3_destination_folder": s3_destination_folder}
+    )
+
+
 def _ghz(num_qubits):
     circuit = Circuit()
     circuit.h(0)