From 5a58264f7b639440db3cdf7cd48ae5ebaf8672c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Henrique=20Silv=C3=A9rio?= <29920212+HGSilveri@users.noreply.github.com> Date: Mon, 2 Dec 2024 18:27:58 +0100 Subject: [PATCH] Fixes to the amplitude noise implementation (#768) * Include the beam's propagation direction in Channel * Apply amplitude fluctuations from run to run instead of pulse to pulse * Correct emulation of finite-waist effects on amplitude * Fix handling or laser_waist to/from SimConfig * UTs for the noise * Include propagation_dir in the JSON schema * Fix error in conversion between SimConfig and NoiseModel * Add the pulser version to serialized devices * Missing UT for Channel * Adopting review suggestions --- pulser-core/pulser/channels/base_channel.py | 25 +- pulser-core/pulser/devices/_device_datacls.py | 9 +- .../abstract_repr/schemas/device-schema.json | 224 ++++++++++++++++++ pulser-core/pulser/noise_model.py | 9 +- .../pulser_simulation/hamiltonian.py | 61 +++-- .../pulser_simulation/simconfig.py | 13 +- tests/test_abstract_repr.py | 4 + tests/test_channels.py | 7 +- tests/test_simconfig.py | 12 +- tests/test_simresults.py | 6 +- tests/test_simulation.py | 27 ++- 11 files changed, 367 insertions(+), 30 deletions(-) diff --git a/pulser-core/pulser/channels/base_channel.py b/pulser-core/pulser/channels/base_channel.py index bc0e4d62..54f27bfe 100644 --- a/pulser-core/pulser/channels/base_channel.py +++ b/pulser-core/pulser/channels/base_channel.py @@ -34,7 +34,7 @@ ChannelType = TypeVar("ChannelType", bound="Channel") -OPTIONAL_ABSTR_CH_FIELDS = ("min_avg_amp",) +OPTIONAL_ABSTR_CH_FIELDS = ("min_avg_amp", "propagation_dir") # States ranked in decreasing order of their associated eigenenergy States = Literal["u", "d", "r", "g", "h", "x"] @@ -78,6 +78,8 @@ class Channel(ABC): min_avg_amp: The minimum average amplitude of a pulse (when not zero). mod_bandwidth: The modulation bandwidth at -3dB (50% reduction), in MHz. + propagation_dir: The propagation direction of the beam associated with + the channel, given as a vector in 3D space. Example: To create a channel targeting the 'ground-rydberg' transition globally, @@ -96,6 +98,7 @@ class Channel(ABC): min_avg_amp: float = 0 mod_bandwidth: Optional[float] = None # MHz eom_config: Optional[BaseEOM] = field(init=False, default=None) + propagation_dir: tuple[float, float, float] | None = None @property def name(self) -> str: @@ -196,7 +199,12 @@ def __post_init__(self) -> None: getattr(self, p) is None ), f"'{p}' must be left as None in a Global channel." else: + assert self.addressing == "Local" parameters += local_only + if self.propagation_dir is not None: + raise NotImplementedError( + "'propagation_dir' must be left as None in Local channels." + ) for param in parameters: value = getattr(self, param) @@ -244,6 +252,18 @@ def __post_init__(self) -> None: "modulation bandwidth." ) + if self.propagation_dir is not None: + dir_vector = np.array(self.propagation_dir, dtype=float) + if dir_vector.size != 3 or np.sum(dir_vector) == 0.0: + raise ValueError( + "'propagation_dir' must be given as a non-zero 3D vector;" + f" got {self.propagation_dir} instead." + ) + # Make sure it's stored as a tuple + object.__setattr__( + self, "propagation_dir", tuple(self.propagation_dir) + ) + @property def rise_time(self) -> int: """The rise time (in ns). @@ -340,6 +360,7 @@ def Global( cls: Type[ChannelType], max_abs_detuning: Optional[float], max_amp: Optional[float], + # TODO: Impose a default propagation_dir in pulser-core 1.3 **kwargs: Any, ) -> ChannelType: """Initializes the channel with global addressing. @@ -361,6 +382,8 @@ def Global( bandwidth at -3dB (50% reduction), in MHz. min_avg_amp: The minimum average amplitude of a pulse (when not zero). + propagation_dir: The propagation direction of the beam associated + with the channel, given as a vector in 3D space. """ # Can't initialize a channel whose addressing is determined internally for cls_field in fields(cls): diff --git a/pulser-core/pulser/devices/_device_datacls.py b/pulser-core/pulser/devices/_device_datacls.py index 3a6872a3..0a248125 100644 --- a/pulser-core/pulser/devices/_device_datacls.py +++ b/pulser-core/pulser/devices/_device_datacls.py @@ -24,6 +24,7 @@ import numpy as np from scipy.spatial.distance import squareform +import pulser import pulser.json.abstract_repr as pulser_abstract_repr import pulser.math as pm from pulser.channels.base_channel import Channel, States, get_states_from_bases @@ -565,7 +566,13 @@ def _to_abstract_repr(self) -> dict[str, Any]: for ch_name, ch_obj in self.channels.items(): ch_list.append(ch_obj._to_abstract_repr(ch_name)) # Add version and channels to params - params.update({"version": "1", "channels": ch_list}) + params.update( + { + "version": "1", + "pulser_version": pulser.__version__, + "channels": ch_list, + } + ) dmm_list = [] for dmm_name, dmm_obj in self.dmm_channels.items(): dmm_list.append(dmm_obj._to_abstract_repr(dmm_name)) diff --git a/pulser-core/pulser/json/abstract_repr/schemas/device-schema.json b/pulser-core/pulser/json/abstract_repr/schemas/device-schema.json index 6d410851..7a97fbd3 100644 --- a/pulser-core/pulser/json/abstract_repr/schemas/device-schema.json +++ b/pulser-core/pulser/json/abstract_repr/schemas/device-schema.json @@ -82,6 +82,22 @@ "null" ] }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." + }, "total_bottom_detuning": { "description": "Minimum possible detuning of the whole DMM channel (in rad/µs), must be below zero.", "type": [ @@ -520,6 +536,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -610,6 +642,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -700,6 +748,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -857,6 +921,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -950,6 +1030,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -1043,6 +1139,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -1193,6 +1305,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -1274,6 +1402,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -1355,6 +1499,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -1500,6 +1660,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -1581,6 +1757,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -1662,6 +1854,22 @@ "number", "null" ] + }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." } }, "required": [ @@ -1751,6 +1959,22 @@ "null" ] }, + "propagation_dir": { + "anyOf": [ + { + "items": { + "type": "number" + }, + "maxItems": 3, + "minItems": 3, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "The propagation direction of the beam leaving the channel." + }, "total_bottom_detuning": { "description": "Minimum possible detuning of the whole DMM channel (in rad/µs), must be below zero.", "type": [ diff --git a/pulser-core/pulser/noise_model.py b/pulser-core/pulser/noise_model.py index e46a9533..a79f25e0 100644 --- a/pulser-core/pulser/noise_model.py +++ b/pulser-core/pulser/noise_model.py @@ -144,10 +144,11 @@ class NoiseModel: p_false_neg: Probability of measuring a false negative. temperature: Temperature, set in µK, of the atoms in the array. Also sets the standard deviation of the speed of the atoms. - laser_waist: Waist of the gaussian laser, set in µm, for global - pulses. - amp_sigma: Dictates the fluctuations in amplitude as a standard - deviation of a normal distribution centered in 1. + laser_waist: Waist of the gaussian lasers, set in µm, for global + pulses. Assumed to be the same for all global channels. + amp_sigma: Dictates the fluctuations in amplitude of global pulses + from run to run as a standard deviation of a normal distribution + centered in 1. Assumed to be the same for all global channels. relaxation_rate: The rate of relaxation from the Rydberg to the ground state (in 1/µs). Corresponds to 1/T1. dephasing_rate: The rate of a dephasing occuring (in 1/µs) in a diff --git a/pulser-simulation/pulser_simulation/hamiltonian.py b/pulser-simulation/pulser_simulation/hamiltonian.py index 13a83f62..c17decf9 100644 --- a/pulser-simulation/pulser_simulation/hamiltonian.py +++ b/pulser-simulation/pulser_simulation/hamiltonian.py @@ -15,6 +15,7 @@ from __future__ import annotations +import functools import itertools from collections import defaultdict from collections.abc import Mapping @@ -221,6 +222,30 @@ def set_config(self, cfg: NoiseModel) -> None: # Noise, samples and Hamiltonian update routine self._construct_hamiltonian() + @staticmethod + @functools.cache + def _finite_waist_amp_fraction( + coords: tuple[float, ...], + propagation_dir: tuple[float, float, float], + laser_waist: float, + ) -> float: + pos_vec = np.zeros(3, dtype=float) + pos_vec[: len(coords)] = np.array(coords, dtype=float) + u_vec = np.array(propagation_dir, dtype=float) + u_vec = u_vec / np.linalg.norm(u_vec) + # Given a line crossing the origin with normalized direction vector + # u_vec, the closest point along said line and an arbitrary point + # pos_vec is at k*u_vec, where + k = np.dot(pos_vec, u_vec) + # The distance between pos_vec and the line is then given by + dist = np.linalg.norm(pos_vec - k * u_vec) + # We assume the Rayleigh length of the gaussian beam is very large, + # resulting in a negligble drop in amplitude along the propagation + # direction. Therefore, the drop in amplitude at a given position + # is dictated solely by its distance to the optical axis (as dictated + # by the propagation direction), ie + return float(np.exp(-((dist / laser_waist) ** 2))) + def _extract_samples(self) -> None: """Populates samples dictionary with every pulse in the sequence.""" local_noises = True @@ -244,15 +269,14 @@ def add_noise( slot: _PulseTargetSlot, samples_dict: Mapping[QubitId, dict[str, np.ndarray]], is_global_pulse: bool, + amp_fluctuation: float, + propagation_dir: tuple | None, ) -> None: """Builds hamiltonian coefficients. Taking into account, if necessary, noise effects, which are local and depend on the qubit's id qid. """ - noise_amp_base = max( - 0, np.random.normal(1.0, self.config.amp_sigma) - ) for qid in slot.targets: if "doppler" in self.config.noise_types: noise_det = self._doppler_detune[qid] @@ -260,22 +284,31 @@ def add_noise( # Gaussian beam loss in amplitude for global pulses only # Noise is drawn at random for each pulse if "amplitude" in self.config.noise_types and is_global_pulse: - amp_fraction = 1.0 + amp_fraction = amp_fluctuation if self.config.laser_waist is not None: - position = self._qdict[qid] - r = np.linalg.norm(position) - w0 = self.config.laser_waist - amp_fraction = np.exp(-((r / w0) ** 2)) - noise_amp = noise_amp_base * amp_fraction - samples_dict[qid]["amp"][slot.ti : slot.tf] *= noise_amp + # Default to an optical axis along y + prop_dir = propagation_dir or (0.0, 1.0, 0.0) + amp_fraction *= self._finite_waist_amp_fraction( + tuple(self._qdict[qid]), + tuple(prop_dir), + self.config.laser_waist, + ) + samples_dict[qid]["amp"][slot.ti : slot.tf] *= amp_fraction if local_noises: for ch, ch_samples in self.samples_obj.channel_samples.items(): - addr = self.samples_obj._ch_objs[ch].addressing - basis = self.samples_obj._ch_objs[ch].basis - samples_dict = samples["Local"][basis] + _ch_obj = self.samples_obj._ch_objs[ch] + samples_dict = samples["Local"][_ch_obj.basis] for slot in ch_samples.slots: - add_noise(slot, samples_dict, addr == "Global") + add_noise( + slot, + samples_dict, + _ch_obj.addressing == "Global", + amp_fluctuation=max( + 0, np.random.normal(1.0, self.config.amp_sigma) + ), + propagation_dir=_ch_obj.propagation_dir, + ) # Delete samples for badly prepared atoms for basis in samples["Local"]: for qid in samples["Local"][basis]: diff --git a/pulser-simulation/pulser_simulation/simconfig.py b/pulser-simulation/pulser_simulation/simconfig.py index 5811495a..cf229f4b 100644 --- a/pulser-simulation/pulser_simulation/simconfig.py +++ b/pulser-simulation/pulser_simulation/simconfig.py @@ -15,8 +15,8 @@ from __future__ import annotations +import math from dataclasses import dataclass, field, fields -from math import sqrt from typing import Any, Optional, Tuple, Type, TypeVar, Union, cast import qutip @@ -58,7 +58,7 @@ def doppler_sigma(temperature: float) -> float: Arg: temperature: The temperature in K. """ - return KEFF * sqrt(KB * temperature / MASS) + return KEFF * math.sqrt(KB * temperature / MASS) @dataclass(frozen=True) @@ -139,16 +139,23 @@ def from_noise_model(cls: Type[T], noise_model: NoiseModel) -> T: kwargs[_DIFF_NOISE_PARAMS.get(param, param)] = getattr( noise_model, param ) + # When laser_waist is None, it should be given as inf instead + # Otherwise, the legacy default laser_waist value will be taken + if "amplitude" in noise_model.noise_types: + kwargs.setdefault("laser_waist", float("inf")) kwargs.pop("with_leakage", None) return cls(**kwargs) def to_noise_model(self) -> NoiseModel: """Creates a NoiseModel from the SimConfig.""" + laser_waist_ = ( + None if math.isinf(self.laser_waist) else self.laser_waist + ) relevant_params = NoiseModel._find_relevant_params( cast(Tuple[NoiseTypes, ...], self.noise), self.eta, self.amp_sigma, - self.laser_waist, + laser_waist_, ) kwargs = {} for param in relevant_params: diff --git a/tests/test_abstract_repr.py b/tests/test_abstract_repr.py index 95ca4305..0d8874e4 100644 --- a/tests/test_abstract_repr.py +++ b/tests/test_abstract_repr.py @@ -232,6 +232,9 @@ def abstract_device(self, request): def test_device_schema(self, abstract_device): validate_abstract_repr(json.dumps(abstract_device), "device") + def test_pulser_version(self, abstract_device): + assert abstract_device["pulser_version"] == pulser.__version__ + def test_roundtrip(self, abstract_device): def _roundtrip(abstract_device): device = deserialize_device(json.dumps(abstract_device)) @@ -509,6 +512,7 @@ def test_optional_device_fields(self, og_device, field, value): "ch_obj", [ Rydberg.Global(None, None, min_avg_amp=1), + Rydberg.Global(None, None, propagation_dir=(1, 0, 0)), Rydberg.Global( None, None, diff --git a/tests/test_channels.py b/tests/test_channels.py index bbf1a321..5f2e2cfe 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -36,6 +36,8 @@ ("mod_bandwidth", 0), ("mod_bandwidth", MODBW_TO_TR * 1e3 + 1), ("min_avg_amp", -1e-3), + ("propagation_dir", (0, 0, 0)), + ("propagation_dir", [1, 0]), ], ) def test_bad_init_global_channel(bad_param, bad_value): @@ -63,12 +65,15 @@ def test_bad_init_global_channel(bad_param, bad_value): ("mod_bandwidth", -1e4), ("mod_bandwidth", MODBW_TO_TR * 1e3 + 1), ("min_avg_amp", -1e-3), + ("propagation_dir", (1, 0, 0)), ], ) def test_bad_init_local_channel(bad_param, bad_value): kwargs = dict(max_abs_detuning=None, max_amp=None) kwargs[bad_param] = bad_value - if bad_param == "mod_bandwidth" and bad_value > 1: + if ( + bad_param == "mod_bandwidth" and bad_value > 1 + ) or bad_param == "propagation_dir": error_type = NotImplementedError else: error_type = ValueError diff --git a/tests/test_simconfig.py b/tests/test_simconfig.py index 14c879cf..2b414dd0 100644 --- a/tests/test_simconfig.py +++ b/tests/test_simconfig.py @@ -138,9 +138,19 @@ def test_noise_model_conversion(): noise_model = NoiseModel( p_false_neg=0.4, p_false_pos=0.1, + amp_sigma=1e-3, + runs=10, + samples_per_run=1, ) expected_simconfig = SimConfig( - noise="SPAM", epsilon=0.1, epsilon_prime=0.4, eta=0.0 + noise=("SPAM", "amplitude"), + epsilon=0.1, + epsilon_prime=0.4, + eta=0.0, + amp_sigma=1e-3, + laser_waist=float("inf"), + runs=10, + samples_per_run=1, ) assert SimConfig.from_noise_model(noise_model) == expected_simconfig assert expected_simconfig.to_noise_model() == noise_model diff --git a/tests/test_simresults.py b/tests/test_simresults.py index e366fa5a..927a3756 100644 --- a/tests/test_simresults.py +++ b/tests/test_simresults.py @@ -309,7 +309,7 @@ def test_expect_noisy(results_noisy): with pytest.raises(ValueError, match="non-diagonal"): results_noisy.expect([bad_op]) op = qutip.tensor([qutip.qeye(2), qutip.basis(2, 0).proj()]) - assert np.isclose(results_noisy.expect([op])[0][-1], 0.7333333333333334) + assert np.isclose(results_noisy.expect([op])[0][-1], 0.7466666666666667) def test_plot(results_noisy, results): @@ -326,7 +326,7 @@ def test_sim_without_measurement(seq_no_meas): ) results_no_meas = sim_no_meas.run() assert results_no_meas.sample_final_state() == Counter( - {"11": 587, "10": 165, "01": 164, "00": 84} + {"11": 580, "10": 173, "01": 167, "00": 80} ) @@ -358,7 +358,7 @@ def test_sample_final_state_three_level(seq_no_meas, pi_pulse): def test_sample_final_state_noisy(seq_no_meas, results_noisy): np.random.seed(123) assert results_noisy.sample_final_state(N_samples=1234) == Counter( - {"11": 725, "10": 265, "01": 192, "00": 52} + {"11": 676, "10": 244, "01": 218, "00": 96} ) res_3level = QutipEmulator.from_sequence( seq_no_meas, config=SimConfig(noise=("SPAM", "doppler"), runs=10) diff --git a/tests/test_simulation.py b/tests/test_simulation.py index 5921e55a..cd4650e9 100644 --- a/tests/test_simulation.py +++ b/tests/test_simulation.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import dataclasses from collections import Counter from unittest.mock import patch @@ -1535,7 +1536,19 @@ def test_effective_size_disjoint(channel_type): ) -def test_simulation_with_modulation(mod_device, reg, patch_plt_show): +@pytest.mark.parametrize( + "propagation_dir", (None, (1, 0, 0), (0, 1, 0), (0, 0, 1)) +) +def test_simulation_with_modulation( + mod_device, reg, propagation_dir, patch_plt_show +): + channels = mod_device.channels + channels["rydberg_global"] = dataclasses.replace( + channels["rydberg_global"], propagation_dir=propagation_dir + ) + mod_device = dataclasses.replace( + mod_device, channel_objects=tuple(channels.values()), channel_ids=None + ) seq = Sequence(reg, mod_device) seq.declare_channel("ch0", "rydberg_global") seq.config_slm_mask({"control1"}) @@ -1589,7 +1602,17 @@ def test_simulation_with_modulation(mod_device, reg, patch_plt_show): ) def pos_factor(qid): - r = np.linalg.norm(reg.qubits[qid].as_array()) + if propagation_dir is None or propagation_dir == (0, 1, 0): + # Optical axis long y, x dicates the distance + r = reg.qubits[qid].as_array()[0] + elif propagation_dir == (1, 0, 0): + # Optical axis long x, y dicates the distance + r = reg.qubits[qid].as_array()[1] + else: + # Optical axis long z, use distance to origin + assert propagation_dir == (0, 0, 1) + r = np.linalg.norm(reg.qubits[qid].as_array()) + w0 = sim_config.laser_waist return np.exp(-((r / w0) ** 2))