add set_noise_model method

amazon-braket · Feb 22, 2024 · 6696e02 · 6696e02
1 parent c19120e
commit 6696e02
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 12 deletions.
diff --git a/src/braket/aws/aws_device.py b/src/braket/aws/aws_device.py
@@ -113,9 +113,7 @@ def __init__(
         self._aws_session = self._get_session_and_initialize(aws_session or AwsSession())
         self._ports = None
         self._frames = None
-        if noise_model:
-            self._validate_device_noise_model_support(noise_model)
-        self._noise_model = noise_model
+        self.set_noise_model(noise_model)
 
     def run(
         self,

diff --git a/src/braket/devices/device.py b/src/braket/devices/device.py
@@ -116,6 +116,18 @@ def status(self) -> str:
         """
         return self._status
 
+    def set_noise_model(self, noise_model: NoiseModel) -> None:
+        """Set the noise model of the device.
+
+        Args:
+            noise_model (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.
+        """
+        if noise_model:
+            self._validate_device_noise_model_support(noise_model)
+        self._noise_model = noise_model
+
     def _validate_device_noise_model_support(self, noise_model: NoiseModel) -> None:
         supported_noises = set(
             SUPPORTED_NOISE_PRAGMA_TO_NOISE[pragma].__name__

diff --git a/src/braket/devices/local_simulator.py b/src/braket/devices/local_simulator.py
@@ -72,9 +72,7 @@ def __init__(
             status="AVAILABLE",
         )
         self._delegate = delegate
-        if noise_model:
-            self._validate_device_noise_model_support(noise_model)
-        self._noise_model = noise_model
+        self.set_noise_model(noise_model)
 
     def run(
         self,

diff --git a/test/unit_tests/braket/devices/test_local_simulator.py b/test/unit_tests/braket/devices/test_local_simulator.py
@@ -632,24 +632,36 @@ def noise_model():
 
 
 @pytest.mark.parametrize("backend", ["dummy_oq3_dm"])
-def test_valid_local_device_for_noise_model(backend, noise_model):
-    device = LocalSimulator(backend, noise_model=noise_model)
+def test_set_noise_model(backend, noise_model):
+    device = LocalSimulator(backend)
+    device.set_noise_model(noise_model)
     assert device._noise_model.instructions == [
         NoiseModelInstruction(Noise.BitFlip(0.05), GateCriteria(Gate.H)),
         NoiseModelInstruction(Noise.TwoQubitDepolarizing(0.10), GateCriteria(Gate.CNot)),
     ]
 
 
 @pytest.mark.parametrize("backend", ["dummy_oq3"])
-def test_invalid_local_device_for_noise_model(backend, noise_model):
+def test_set_noise_model_invalid_device(backend, noise_model):
     with pytest.raises(ValueError):
-        _ = LocalSimulator(backend, noise_model=noise_model)
+        device = LocalSimulator(backend)
+        device.set_noise_model(noise_model)
 
 
 @pytest.mark.parametrize("backend", ["dummy_oq3_dm"])
-def test_local_device_with_invalid_noise_model(backend, noise_model):
+def test_set_noise_model_invalid_noise_model(backend, noise_model):
     with pytest.raises(TypeError):
-        _ = LocalSimulator(backend, noise_model=Mock())
+        device = LocalSimulator(backend)
+        device.set_noise_model(Mock())
+
+
+@pytest.mark.parametrize("backend", ["dummy_oq3_dm"])
+def test_valid_local_device_for_noise_model(backend, noise_model):
+    device = LocalSimulator(backend, noise_model=noise_model)
+    assert device._noise_model.instructions == [
+        NoiseModelInstruction(Noise.BitFlip(0.05), GateCriteria(Gate.H)),
+        NoiseModelInstruction(Noise.TwoQubitDepolarizing(0.10), GateCriteria(Gate.CNot)),
+    ]
 
 
 @patch.object(DummyProgramDensityMatrixSimulator, "run")