Sequence sampler

pulser/sampler/sampler.py

@@ -1,6 +1,9 @@
 """Exposes the sample() functions.
 
 It contains many helpers.
+
+NOTE(perf): it not very efficient to hold the same data for all qubit in global
+channels, but i remains manageable for register # with less than 100 qubits.
 """
 from __future__ import annotations
 
@@ -61,12 +64,10 @@ def sample(
 
         ch_noises = list(common_noises)
 
+        slm_on = len(seq._slm_mask_targets) > 0 and len(seq._slm_mask_time) > 0
+
         if addr == "Global":
-            decay = (
-                len(seq._slm_mask_targets) > 0
-                or len(global_noises) > 0
-                or len(common_noises) > 0
-            )
+            decay = slm_on or len(global_noises) > 0 or len(common_noises) > 0
             if decay:
                 addr = "Decayed"
                 ch_noises.extend(global_noises)
@@ -80,8 +81,12 @@ def sample(
 
         s = noises.apply(s, ch_noises)
 
-        unmasked_qubits = seq._qids - seq._slm_mask_targets
-        s = [x for x in s if x.qubit in unmasked_qubits]  # SLM
+        if slm_on:  # Update the samples of masked qubits during SLM on times
+            for i, _ in enumerate(s):
+                if s[i].qubit in seq._slm_mask_targets:
+                    ti, tf = seq._slm_mask_time[0], seq._slm_mask_time[1]
+                    s[i].amp[ti:tf] = 0.0
+                    # apply only on amp since it's just a shutter
 
         samples[ch_name] = s
 
@@ -105,7 +110,7 @@ def new_qty_dict() -> dict:
         }
 
     def new_qdict() -> dict:
-        return {qubit: new_qty_dict() for qubit in seq._qids}
+        return defaultdict(new_qty_dict)
 
     if seq._in_xy:
         return {
@@ -180,9 +185,8 @@ def _sample_slots(N: int, *slots: _TimeSlot) -> list[QubitSamples]:
     # weird, it's not enforced by the structure,bad design?)
     qubits = slots[0].targets
     amp, det, phase = np.zeros(N), np.zeros(N), np.zeros(N)
-    for s in slots:
-        if type(s.type) is str:  # pragma: no cover
-            continue
+    pulse_slots = [s for s in slots if isinstance(s.type, Pulse)]
+    for s in pulse_slots:
         pulse = cast(Pulse, s.type)
         amp[s.ti : s.tf] += pulse.amplitude.samples
         det[s.ti : s.tf] += pulse.detuning.samples

pulser/tests/test_sampling_noises.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import pulser
+import pulser.simulation.noises as noises
+from pulser.devices import MockDevice
+from pulser.pulse import Pulse
+from pulser.sampler import sample
+from pulser.waveforms import ConstantWaveform
+
+
+def test_amplitude_noise():
+    """Test the noise related to the amplitude profile of global pulses."""
+    N = 100
+    amplitude = 1.0
+    waist_width = 2.0  # µm
+
+    coords = np.array([[-2.0, 0.0], [0.0, 0.0], [2.0, 0.0]])
+    reg = pulser.Register.from_coordinates(coords, prefix="q")
+    seq = pulser.Sequence(reg, MockDevice)
+    seq.declare_channel("ch0", "rydberg_global")
+    seq.add(
+        Pulse.ConstantAmplitude(amplitude, ConstantWaveform(N, 0.0), 0.0),
+        "ch0",
+    )
+    seq.measure()
+
+    def expected_samples(vec: np.ndarray) -> np.ndarray:
+        """Defines the non-noisy effect of a gaussian amplitude profile."""
+        r = np.linalg.norm(vec)
+        a = np.ones(N)
+        a *= amplitude
+        a *= np.exp(-((r / waist_width) ** 2))
+        return a
+
+    s = sample(
+        seq, global_noises=[noises.amplitude(reg, waist_width, random=False)]
+    )
+
+    for q, coords in reg.qubits.items():
+        got = s["Local"]["ground-rydberg"][q]["amp"]
+        want = expected_samples(coords)
+        np.testing.assert_equal(got, want)
+
+
+@pytest.mark.xfail(
+    reason="Test a different doppler effect than the one implemented; "
+    "no surprise it fails."
+)
+def test_doppler_noise():
+    """What is exactly the doppler noise here?
+
+    A constant detuning shift per pulse seems weird. A global shift seems more
+    reasonable, but how can it be constant during the all sequence? It is not
+    clear to me here, I find the current implementation in the simulation
+    module to be unsatisfactory.
+
+    No surprise I make it fail on purpose right now 😅
+    """
+    N = 100
+    det_value = np.pi
+
+    reg = pulser.Register.from_coordinates(np.array([[0.0, 0.0]]), prefix="q")
+    seq = pulser.Sequence(reg, MockDevice)
+    seq.declare_channel("ch0", "rydberg_global")
+    for _ in range(3):
+        seq.add(
+            Pulse.ConstantDetuning(ConstantWaveform(N, 1.0), det_value, 0.0),
+            "ch0",
+        )
+        seq.delay(100, "ch0")
+    seq.measure()
+
+    MASS = 1.45e-25  # kg
+    KB = 1.38e-23  # J/K
+    KEFF = 8.7  # µm^-1
+    doppler_sigma = KEFF * np.sqrt(KB * 50.0e-6 / MASS)
+    seed = 42
+    rng = np.random.default_rng(seed)
+
+    shifts = rng.normal(0, doppler_sigma, 3)
+    want = np.zeros(6 * N)
+    want[0:100] = det_value + shifts[0]
+    want[200:300] = det_value + shifts[1]
+    want[400:500] = det_value + shifts[2]
+
+    local_noises = [noises.doppler(reg, doppler_sigma, seed=seed)]
+    samples = sample(seq, common_noises=local_noises)
+    got = samples["Local"]["ground-rydberg"]["q0"]["det"]
+
+    np.testing.assert_array_equal(got, want)