diff --git a/src/braket/aws/aws_device.py b/src/braket/aws/aws_device.py index 70fc3c078..338f9c17b 100644 --- a/src/braket/aws/aws_device.py +++ b/src/braket/aws/aws_device.py @@ -113,8 +113,6 @@ 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 def run( diff --git a/src/braket/devices/device.py b/src/braket/devices/device.py index dbc7b6b35..827b8e626 100644 --- a/src/braket/devices/device.py +++ b/src/braket/devices/device.py @@ -116,6 +116,17 @@ 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. + """ + 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 = { SUPPORTED_NOISE_PRAGMA_TO_NOISE[pragma].__name__ diff --git a/src/braket/devices/local_simulator.py b/src/braket/devices/local_simulator.py index 15ec904de..1c4ab99d8 100644 --- a/src/braket/devices/local_simulator.py +++ b/src/braket/devices/local_simulator.py @@ -74,8 +74,6 @@ def __init__( status="AVAILABLE", ) self._delegate = delegate - if noise_model: - self._validate_device_noise_model_support(noise_model) self._noise_model = noise_model def run( diff --git a/test/unit_tests/braket/devices/test_local_simulator.py b/test/unit_tests/braket/devices/test_local_simulator.py index 216d161c7..2d36e2c58 100644 --- a/test/unit_tests/braket/devices/test_local_simulator.py +++ b/test/unit_tests/braket/devices/test_local_simulator.py @@ -670,8 +670,9 @@ 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)), @@ -679,15 +680,26 @@ def test_valid_local_device_for_noise_model(backend, noise_model): @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")