From 768259679955564f74c8d5c7dd15ce1d88da2309 Mon Sep 17 00:00:00 2001 From: Tim Date: Thu, 22 Feb 2024 22:07:25 -0500 Subject: [PATCH] feat: support Braket noise model --- .../providers/braket_backend.py | 17 ++++++++++++++++- .../providers/braket_provider.py | 8 +++++++- tests/providers/mocks.py | 10 ++++++++++ tests/providers/test_braket_backend.py | 18 ++++++++++++++++++ tests/providers/test_braket_provider.py | 19 ++++++++++++++++++- 5 files changed, 69 insertions(+), 3 deletions(-) diff --git a/qiskit_braket_provider/providers/braket_backend.py b/qiskit_braket_provider/providers/braket_backend.py index 1da02bee..25c3baa0 100644 --- a/qiskit_braket_provider/providers/braket_backend.py +++ b/qiskit_braket_provider/providers/braket_backend.py @@ -10,6 +10,7 @@ from braket.aws import AwsDevice, AwsQuantumTask, AwsQuantumTaskBatch from braket.aws.queue_information import QueueDepthInfo from braket.circuits import Circuit +from braket.circuits.noise_model import NoiseModel as BraketNoiseModel from braket.device_schema import DeviceActionType from braket.devices import Device, LocalSimulator from braket.tasks.local_quantum_task import LocalQuantumTask @@ -59,7 +60,12 @@ def _get_gateset(self) -> Optional[set[str]]: class BraketLocalBackend(BraketBackend): """BraketLocalBackend.""" - def __init__(self, name: str = "default", **fields): + def __init__( + self, + name: str = "default", + noise_model: Optional[BraketNoiseModel] = None, + **fields, + ): """AWSBraketLocalBackend for local execution of circuits. Example: @@ -71,11 +77,15 @@ def __init__(self, name: str = "default", **fields): Args: name: name of backend + noise_model (Optional[NoiseModel]): The Braket noise model to apply to the circuit + before execution. Noise model can only be added to the devices that support + noise simulation. **fields: extra fields """ super().__init__(name=name, **fields) self.backend_name = name self._local_device = LocalSimulator(backend=self.backend_name) + self._local_device.set_noise_model(noise_model) self._target = local_simulator_to_target(self._local_device) self.status = self._local_device.status @@ -178,6 +188,7 @@ def __init__( # pylint: disable=too-many-arguments description: str = None, online_date: datetime.datetime = None, backend_version: str = None, + noise_model: Optional[BraketNoiseModel] = None, **fields, ): """AWSBraketBackend for execution circuits against AWS Braket devices. @@ -196,6 +207,9 @@ def __init__( # pylint: disable=too-many-arguments description: description of backend online_date: online date backend_version: backend version + noise_model (Optional[NoiseModel]): The Braket noise model to apply to the circuit + before execution. Noise model can only be added to the devices that support + noise simulation. **fields: other arguments """ super().__init__( @@ -208,6 +222,7 @@ def __init__( # pylint: disable=too-many-arguments ) user_agent = f"QiskitBraketProvider/{version.__version__}" device.aws_session.add_braket_user_agent(user_agent) + device.set_noise_model(noise_model) self._aws_device = device self._target = aws_device_to_target(device=device) diff --git a/qiskit_braket_provider/providers/braket_provider.py b/qiskit_braket_provider/providers/braket_provider.py index ef5ae132..59b4459d 100644 --- a/qiskit_braket_provider/providers/braket_provider.py +++ b/qiskit_braket_provider/providers/braket_provider.py @@ -1,6 +1,9 @@ """AWS Braket provider.""" +from typing import Optional + from braket.aws import AwsDevice +from braket.circuits.noise_model import NoiseModel as BraketNoiseModel from braket.device_schema.dwave import DwaveDeviceCapabilities from braket.device_schema.quera import QueraDeviceCapabilities from braket.device_schema.xanadu import XanaduDeviceCapabilities @@ -28,7 +31,9 @@ class AWSBraketProvider(ProviderV1): BraketBackend[dm1]] """ - def backends(self, name=None, **kwargs): + def backends( + self, name=None, noise_model: Optional[BraketNoiseModel] = None, **kwargs + ): if kwargs.get("local"): return [BraketLocalBackend(name="default")] names = [name] if name else None @@ -57,6 +62,7 @@ def backends(self, name=None, **kwargs): description=f"AWS Device: {device.provider_name} {device.name}.", online_date=device.properties.service.updatedAt, backend_version="2", + noise_model=noise_model, ) ) return backends diff --git a/tests/providers/mocks.py b/tests/providers/mocks.py index ae716bb3..cd3606f4 100644 --- a/tests/providers/mocks.py +++ b/tests/providers/mocks.py @@ -15,6 +15,7 @@ RIGETTI_ARN = "arn:aws:braket:::device/qpu/rigetti/Aspen-10" RIGETTI_ASPEN_ARN = "arn:aws:braket:::device/qpu/rigetti/Aspen-M-3" SV1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/sv1" +DM1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/dm1" TN1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/tn1" RIGETTI_REGION = "us-west-1" SIMULATOR_REGION = "us-west-1" @@ -136,6 +137,15 @@ "deviceCapabilities": MOCK_GATE_MODEL_SIMULATOR_CAPABILITIES.json(), } +MOCK_GATE_MODEL_SIMULATOR_DM = { + "deviceName": "dm1", + "deviceType": "SIMULATOR", + "providerName": "provider1", + "deviceStatus": "ONLINE", + "deviceArn": DM1_ARN, + "deviceCapabilities": MOCK_GATE_MODEL_SIMULATOR_CAPABILITIES.json(), +} + MOCK_GATE_MODEL_SIMULATOR_TN = { "deviceName": "tn1", "deviceType": "SIMULATOR", diff --git a/tests/providers/test_braket_backend.py b/tests/providers/test_braket_backend.py index 2287e089..ad061968 100644 --- a/tests/providers/test_braket_backend.py +++ b/tests/providers/test_braket_backend.py @@ -22,6 +22,7 @@ from qiskit_braket_provider.providers import AWSBraketBackend, BraketLocalBackend from qiskit_braket_provider.providers.adapter import aws_device_to_target from tests.providers.mocks import ( + MOCK_GATE_MODEL_SIMULATOR_CAPABILITIES, RIGETTI_MOCK_GATE_MODEL_QPU_CAPABILITIES, RIGETTI_MOCK_M_3_QPU_CAPABILITIES, ) @@ -75,6 +76,16 @@ def test_device_backend(self): with self.assertRaises(NotImplementedError): backend.control_channel([0, 1]) + def test_aws_backend_noise_model(self): + """Tests local backend.""" + mock_noise_model = Mock() + mock_set_noise_model = Mock() + device = Mock() + device.properties = MOCK_GATE_MODEL_SIMULATOR_CAPABILITIES + device.set_noise_model = mock_set_noise_model + _ = AWSBraketBackend(device, noise_model=mock_noise_model) + self.assertEqual(mock_set_noise_model.call_args[0][0], mock_noise_model) + def test_local_backend(self): """Tests local backend.""" backend = BraketLocalBackend(name="default") @@ -90,6 +101,13 @@ def test_local_backend(self): with self.assertRaises(NotImplementedError): backend.control_channel([0, 1]) + @patch("braket.devices.LocalSimulator.set_noise_model") + def test_local_backend_noise_model(self, mock_set_noise_model): + """Tests local backend.""" + mock_noise_model = Mock() + _ = BraketLocalBackend(name="default", noise_model=mock_noise_model) + self.assertEqual(mock_set_noise_model.call_args[0][0], mock_noise_model) + def test_local_backend_output(self): """Test local backend output""" first_backend = BraketLocalBackend(name="braket_dm") diff --git a/tests/providers/test_braket_provider.py b/tests/providers/test_braket_provider.py index 81502c5b..f939f233 100644 --- a/tests/providers/test_braket_provider.py +++ b/tests/providers/test_braket_provider.py @@ -17,6 +17,7 @@ BraketBackend, ) from tests.providers.mocks import ( + MOCK_GATE_MODEL_SIMULATOR_DM, MOCK_GATE_MODEL_SIMULATOR_SV, MOCK_GATE_MODEL_SIMULATOR_TN, MOCK_RIGETTI_GATE_MODEL_M_3_QPU, @@ -30,7 +31,11 @@ class TestAWSBraketProvider(TestCase): def setUp(self): self.mock_session = Mock() - simulators = [MOCK_GATE_MODEL_SIMULATOR_SV, MOCK_GATE_MODEL_SIMULATOR_TN] + simulators = [ + MOCK_GATE_MODEL_SIMULATOR_SV, + MOCK_GATE_MODEL_SIMULATOR_DM, + MOCK_GATE_MODEL_SIMULATOR_TN, + ] self.mock_session.get_device.side_effect = simulators self.mock_session.region = SIMULATOR_REGION self.mock_session.boto_session.region_name = SIMULATOR_REGION @@ -71,6 +76,18 @@ def test_real_devices(self): with self.subTest(f"{backend.name}"): self.assertIsInstance(backend, AWSBraketBackend) + @patch("qiskit_braket_provider.providers.braket_provider.AWSBraketBackend") + @patch("qiskit_braket_provider.providers.braket_backend.AwsDevice.get_devices") + def test_get_backend_noise_model(self, mock_get_devices, mock_aws_braket_backend): + """Tests qiskit circuit transpilation.""" + mock_noise_model = Mock() + mock_get_devices.return_value = [Mock()] + provider = AWSBraketProvider() + _ = provider.get_backend(name="dm1", noise_model=mock_noise_model) + self.assertEqual( + mock_aws_braket_backend.call_args[1]["noise_model"], mock_noise_model + ) + @patch("qiskit_braket_provider.providers.braket_backend.AWSBraketBackend") @patch("qiskit_braket_provider.providers.braket_backend.AwsDevice.get_devices") def test_qiskit_circuit_transpilation_run(