Repository: pasqal-io/Pulser (language: python)

Issue:
Sequence sampler component
Samples of `Sequence` instances are currently computed in place in the `Simulation` class and in the `_seq_drawer` function as well. This gives [two](https://github.com/pasqal-io/Pulser/blob/8d0c4296348be414da1793eb3ece174422d0158b/pulser/simulation/simulation.py#L408) [pieces](https://github.com/pasqal-io/Pulser/blob/8d0c4296348be414da1793eb3ece174422d0158b/pulser/_seq_drawer.py#L32) of code dedicated to the same task, which is bad. 

To 1) fix this 2) ease the construction of the samples dictionary and 3) decouple its use from its construction, it would be very beneficial to implement an independent Sequence sampler component.
Se the [original proposal](https://github.com/pasqal-io/Pulser/issues/337#issuecomment-1059064379_) in PR #337.

The sampler will receive a `Sequence` instance as well as the strategies to sample it: noise sources, qubits masked by an SLM, modulation effects from AOM/EOM. These strategies must play well with each others: they should be easy to add, remove or update independently.


Discussion / hints:

Fix patch:
diff --git a/pulser/sampler/__init__.py b/pulser/sampler/__init__.py
new file mode 100644
index 000000000..cbfbcc3fb
--- /dev/null
+++ b/pulser/sampler/__init__.py
@@ -0,0 +1,23 @@
+# Copyright 2020 Pulser Development Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module sampler enables the sampling of pulser sequences.
+
+Samples of a sequence are needed for plotting and simulation.
+
+  Typical usage:
+
+  sampler.sample(sequence)
+"""
+from pulser.sampler.sampler import sample
diff --git a/pulser/sampler/sampler.py b/pulser/sampler/sampler.py
new file mode 100644
index 000000000..aaf2243ef
--- /dev/null
+++ b/pulser/sampler/sampler.py
@@ -0,0 +1,284 @@
+"""Exposes the sample() functions.
+
+It contains many helpers.
+"""
+from __future__ import annotations
+
+import itertools
+from collections import defaultdict
+from typing import Callable, List, Optional, cast
+
+import numpy as np
+
+import pulser.simulation.noises as noises
+from pulser.channels import Channel
+from pulser.pulse import Pulse
+from pulser.sampler.samples import QubitSamples
+from pulser.sequence import Sequence, _TimeSlot
+
+
+def sample(
+    seq: Sequence,
+    modulation: bool = False,
+    common_noises: Optional[list[noises.NoiseModel]] = None,
+    global_noises: Optional[list[noises.NoiseModel]] = None,
+) -> dict:
+    """Samples the given Sequence and returns a nested dictionary.
+
+    It is intended to be used like the json.dumps() function.
+
+    Args:
+        seq (Sequence): A pulser.Sequence instance.
+        modulation (bool): Flag to account for the modulation of AOM/EOM
+            before sampling.
+        common_noises (Optional[list[LocalNoise]]): A list of the noise sources
+            for all channels.
+        global_noises (Optional[list[LocalNoise]]): A list of the noise sources
+            for global channels.
+
+    Returns:
+        A nested dictionnary of the samples of the amplitude, detuning and
+        phase at every nanoseconds for all channels.
+    """
+    if common_noises is None:
+        common_noises = []
+    if global_noises is None:
+        global_noises = []
+
+    # 1. determine if the global channel decay to a local one
+    # 2. extract samples
+    # 3. modulate
+    # 4. apply noises/SLM
+    # 5. write samples
+    #
+    # NOTE(perf): it not very efficient to hold copies of the same data for
+    # every qubits in a global channel, but it remains manageable for registers
+    # with less than 100 qubits.
+
+    samples: dict[str, list[QubitSamples]] = {}
+    addrs: dict[str, str] = {}
+
+    for ch_name, ch in seq.declared_channels.items():
+        s: list[QubitSamples]
+
+        addr = seq.declared_channels[ch_name].addressing
+
+        ch_noises = list(common_noises)
+
+        slm_on = seq._slm_mask_targets and seq._slm_mask_time
+
+        if addr == "Global":
+            decay = slm_on or len(global_noises) > 0 or len(common_noises) > 0
+            if decay:
+                addr = "Decayed"
+                ch_noises.extend(global_noises)
+
+        addrs[ch_name] = addr
+
+        strategy = _group_between_retargets if modulation else _regular
+        s = _sample_channel(seq, ch_name, strategy)
+        if modulation:
+            s = _modulate(ch, s)
+
+        s = noises.apply(s, ch_noises)
+
+        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
+
+    # format the samples in the simulation dict form
+    d = _write_dict(seq, samples, addrs)
+
+    return d
+
+
+def _prepare_dict(seq: Sequence, N: int) -> dict:
+    """Constructs empty dict of size N.
+
+    Usually N is the duration of seq.
+    """
+
+    def new_qty_dict() -> dict:
+        return {
+            "amp": np.zeros(N),
+            "det": np.zeros(N),
+            "phase": np.zeros(N),
+        }
+
+    def new_qdict() -> dict:
+        return defaultdict(new_qty_dict)
+
+    if seq._in_xy:
+        return {
+            "Global": {"XY": new_qty_dict()},
+            "Local": {"XY": new_qdict()},
+        }
+    else:
+        return {
+            "Global": defaultdict(new_qty_dict),
+            "Local": defaultdict(new_qdict),
+        }
+
+
+def _write_dict(
+    seq: Sequence,
+    samples: dict[str, list[QubitSamples]],
+    addrs: dict[str, str],
+) -> dict:
+    """Export the given samples to a nested dictionary."""
+    # Get the duration
+    if not _same_duration(samples):
+        raise ValueError("All the samples do not share the same duration.")
+    N = list(samples.values())[0][0].amp.size
+
+    d = _prepare_dict(seq, N)
+
+    for ch_name, some_samples in samples.items():
+        basis = seq.declared_channels[ch_name].basis
+        addr = addrs[ch_name]
+        if addr == "Global":
+            # Take samples on only one qubit and write them
+            a_qubit = next(iter(seq._qids))
+            to_write = [x for x in some_samples if x.qubit == a_qubit]
+            for s in to_write:
+                d["Global"][basis]["amp"] += s.amp
+                d["Global"][basis]["det"] += s.det
+                d["Global"][basis]["phase"] += s.phase
+        else:
+            for s in some_samples:
+                d["Local"][basis][s.qubit]["amp"] += s.amp
+                d["Local"][basis][s.qubit]["det"] += s.det
+                d["Local"][basis][s.qubit]["phase"] += s.phase
+    return d
+
+
+def _same_duration(samples: dict[str, list[QubitSamples]]) -> bool:
+    durations: list[int] = []
+    flatten_samples: list[QubitSamples] = []
+    for some_samples in samples.values():
+        flatten_samples.extend(some_samples)
+    for s in flatten_samples:
+        durations.extend((s.amp.size, s.det.size, s.phase.size))
+    return durations.count(durations[0]) == len(durations)
+
+
+def _sample_channel(
+    seq: Sequence, ch_name: str, strategy: TimeSlotExtractionStrategy
+) -> list[QubitSamples]:
+    """Compute a list of QubitSamples for a channel."""
+    qs: list[QubitSamples] = []
+    grouped_slots = strategy(seq._schedule[ch_name])
+
+    for group in grouped_slots:
+        ss = _sample_slots(seq.get_duration(), *group)
+        qs.extend(ss)
+    return qs
+
+
+def _sample_slots(N: int, *slots: _TimeSlot) -> list[QubitSamples]:
+    """Gather samples of a list of _TimeSlot in a single Samples instance."""
+    # Same target in one group, guaranteed by the strategy (this seems
+    # 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)
+    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
+        phase[s.ti : s.tf] += pulse.phase
+    qs = [
+        QubitSamples(
+            amp=amp.copy(), det=det.copy(), phase=phase.copy(), qubit=q
+        )
+        for q in qubits
+    ]
+    return qs
+
+
+TimeSlotExtractionStrategy = Callable[[List[_TimeSlot]], List[List[_TimeSlot]]]
+"""Extraction strategy of _TimeSlot's of a Channel.
+
+It's an alias for functions that returns a list of lists of _TimeSlots.
+_TimeSlots in the same group MUST share the same targets.
+
+NOTE:
+    This strategy type is used mostly for the necessity to extract samples
+    differently when taking into account the modulation of AOM/EOM. Despite
+    there are only two cases, whether it's necessary to modulate a local
+    channel or not, this pattern can accomodate for future needs.
+"""
+
+
+def _regular(ts: list[_TimeSlot]) -> list[list[_TimeSlot]]:
+    """No grouping performed, return only the pulses."""
+    return [[x] for x in ts if isinstance(x.type, Pulse)]
+
+
+def _group_between_retargets(
+    ts: list[_TimeSlot],
+) -> list[list[_TimeSlot]]:
+    """Filter and group _TimeSlots together.
+
+    Group the input slots by groups of successive Pulses and delays between
+    two target operations. Consider the following sequence consisting of pulses
+    A B C D E F, targeting different qubits:
+
+    .---A---B------.---C--D--E---.----F--
+    ^              ^             ^
+    |              |             |
+    target q0   target q1     target q0
+
+    It will group the pulses' _TimeSlot's in batches (A B), (C D E) and (F),
+    returning the following list of list of _TimeSlot instances:
+
+    [[A, B], [C, D, E], [F]]
+
+    Args:
+        ts (list[_TimeSlot]): A list of TimeSlot from a Sequence schedule.
+
+    Returns:
+        A list of list of _TimeSlot. _TimeSlot instances are successive and
+        share the same targets. They are of type either Pulse or "delay", all
+        "target" ones are discarded.
+    """
+    TO_KEEP = "pulses_and_delays"
+
+    def key_func(x: _TimeSlot) -> str:
+        if isinstance(x.type, Pulse) or x.type == "delay":
+            return TO_KEEP
+        else:
+            return "other"
+
+    grouped_slots: list[list[_TimeSlot]] = []
+
+    for key, group in itertools.groupby(ts, key_func):
+        g = list(group)
+        if key != TO_KEEP:
+            continue
+        grouped_slots.append(g)
+
+    return grouped_slots
+
+
+def _modulate(ch: Channel, samples: list[QubitSamples]) -> list[QubitSamples]:
+    """Modulate local samples according to the hardware specs.
+
+    Additional parameters will probably be needed (keep_end, etc).
+    """
+    modulated_samples: list[QubitSamples] = []
+    for s in samples:
+        modulated_samples.append(
+            QubitSamples(
+                amp=ch.modulate(s.amp),
+                det=ch.modulate(s.det),
+                phase=ch.modulate(s.phase),
+                qubit=s.qubit,
+            )
+        )
+    return modulated_samples
diff --git a/pulser/sampler/samples.py b/pulser/sampler/samples.py
new file mode 100644
index 000000000..4072ab064
--- /dev/null
+++ b/pulser/sampler/samples.py
@@ -0,0 +1,24 @@
+"""Defines samples dataclasses."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from pulser.sequence import QubitId
+
+
+@dataclass
+class QubitSamples:
+    """Gathers samples concerning a single qubit."""
+
+    amp: np.ndarray
+    det: np.ndarray
+    phase: np.ndarray
+    qubit: QubitId
+
+    def __post_init__(self) -> None:
+        if not len(self.amp) == len(self.det) == len(self.phase):
+            raise ValueError(
+                "ndarrays amp, det and phase must have the same length."
+            )
diff --git a/pulser/simulation/noises.py b/pulser/simulation/noises.py
new file mode 100644
index 000000000..4a537bdcb
--- /dev/null
+++ b/pulser/simulation/noises.py
@@ -0,0 +1,142 @@
+"""Contains noise models.
+
+For now, only the amplitude and doppler noises are implemented in the form a
+NoiseModel, which are the laser-atom interaction related noises relevant at
+sampling time.
+"""
+from __future__ import annotations
+
+import functools
+from typing import Callable, Optional
+
+import numpy as np
+
+from pulser.register import Register
+from pulser.sampler.samples import QubitSamples
+
+NoiseModel = Callable[[QubitSamples], QubitSamples]
+"""A function that apply some noise on a list of QubitSamples.
+
+A NoiseModel corresponds to a source of noises present in a device which is
+relevant when sampling the input pulses. Physical effects contributing to
+modifications of the shined amplitude, detuning and phase felt by qubits of the
+register are susceptible to be implemented by a NoiseModel.
+"""
+
+
+def amplitude(
+    reg: Register,
+    waist_width: float,
+    random: bool = True,
+    seed: Optional[int] = None,
+) -> NoiseModel:
+    """Generate a NoiseModel for the gaussian amplitude profile of laser beams.
+
+    The laser of a global channel has a non-constant amplitude profile in the
+    register plane. It makes global channels act differently on each qubit,
+    becoming local.
+
+    Args:
+        reg (Register): A Pulser register
+        waist_width (float): The laser waist_width in µm
+        random (bool): Adds an additional random noise on the amplitude
+        seed (int): Optional, seed for the numpy.random.Generator
+
+    Return:
+        NoiseModel: The function that applies the amplitude noise to some
+        QubitSamples.
+    """
+    rng = np.random.default_rng(seed)
+
+    def f(s: QubitSamples) -> QubitSamples:
+        r = np.linalg.norm(reg.qubits[s.qubit])
+
+        noise_amp = rng.normal(1.0, 1.0e-3) if random else 1.0
+        noise_amp *= np.exp(-((r / waist_width) ** 2))
+
+        amp = s.amp.copy()
+        amp *= noise_amp
+
+        return QubitSamples(
+            amp=amp,
+            det=s.det.copy(),
+            phase=s.phase.copy(),
+            qubit=s.qubit,
+        )
+
+    return f
+
+
+def doppler(reg: Register, std_dev: float, seed: Optional[int]) -> NoiseModel:
+    """Generate a NoiseModel for the Doppler effect detuning shifts.
+
+    Example usage:
+
+        MASS = 1.45e-25  # kg
+        KB = 1.38e-23  # J/K
+        KEFF = 8.7  # µm^-1
+        sigma = KEFF * np.sqrt(KB * 50.0e-6 / MASS)
+        doppler_noise = doppler(reg, sigma)
+        ...
+
+    Args:
+        reg (Register): A Pulser register
+        std_dev (float): The standard deviation of the normal distribution used
+            to sample the random detuning shifts
+        seed (int): Optional, seed for the numpy.random.Generator
+
+    Return:
+        NoiseModel: The function that applies the doppler noise to some
+        QubitSamples.
+    """
+    rng = np.random.default_rng(seed)
+    errs = rng.normal(0.0, std_dev, size=len(reg.qubit_ids))
+    detunings = dict(zip(reg.qubit_ids, errs))
+
+    def f(s: QubitSamples) -> QubitSamples:
+        det = s.det.copy()
+        det[np.nonzero(s.det)] += detunings[s.qubit]
+        return QubitSamples(
+            amp=s.amp.copy(),
+            det=det,
+            phase=s.phase.copy(),
+            qubit=s.qubit,
+        )
+
+    return f
+
+
+def compose_local_noises(*functions: NoiseModel) -> NoiseModel:
+    """Helper to compose multiple NoiseModel.
+
+    Args:
+        *functions: a list of functions
+
+    Returns:
+        The mathematical composition of *functions. The last element is applied
+        first. If *functions is [f, g, h], it returns f∘g∘h.
+    """
+    return functools.reduce(
+        lambda f, g: lambda x: f(g(x)), functions, lambda x: x
+    )
+
+
+def apply(
+    samples: list[QubitSamples], noises: list[NoiseModel]
+) -> list[QubitSamples]:
+    """Apply a list of NoiseModel on a list of QubitSamples.
+
+    The noises are composed using the compose_local_noises function, such that
+    the last element is applied first.
+
+    Args:
+        samples (list[QubitSamples]): A list of QubitSamples.
+        noises (list[NoiseModel]): A list of NoiseModel.
+
+    Return:
+        A list of QubitSamples on which each element of noises has been
+        applied.
+    """
+    tot_noise = compose_local_noises(*noises)
+
+    return [tot_noise(s) for s in samples]


Test patch:
diff --git a/pulser/tests/test_sampling_noises.py b/pulser/tests/test_sampling_noises.py
new file mode 100644
index 000000000..2822ea929
--- /dev/null
+++ b/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)
diff --git a/pulser/tests/test_sequence_sampler.py b/pulser/tests/test_sequence_sampler.py
new file mode 100644
index 000000000..c170e4da2
--- /dev/null
+++ b/pulser/tests/test_sequence_sampler.py
@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import textwrap
+from copy import deepcopy
+
+import numpy as np
+import pytest
+
+import pulser
+from pulser.channels import Rydberg
+from pulser.devices import Device, MockDevice
+from pulser.pulse import Pulse
+from pulser.sampler import sample
+from pulser.sampler.sampler import _write_dict
+from pulser.sampler.samples import QubitSamples
+from pulser.waveforms import BlackmanWaveform, RampWaveform
+
+# Helpers
+
+
+def assert_same_samples_as_sim(seq: pulser.Sequence) -> None:
+    """Check against the legacy sample extraction in the simulation module."""
+    got = sample(seq)
+    want = pulser.Simulation(seq).samples
+
+    assert_nested_dict_equality(got, want)
+
+
+def assert_nested_dict_equality(got, want: dict) -> None:
+    for basis in want["Global"]:
+        for qty in want["Global"][basis]:
+            np.testing.assert_array_equal(
+                got["Global"][basis][qty],
+                want["Global"][basis][qty],
+            )
+    for basis in want["Local"]:
+        for qubit in want["Local"][basis]:
+            for qty in want["Local"][basis][qubit]:
+                np.testing.assert_array_equal(
+                    got["Local"][basis][qubit][qty],
+                    want["Local"][basis][qubit][qty],
+                )
+
+
+# Tests
+
+
+def test_one_pulse_sampling():
+    """Test the sample function on a one-pulse sequence."""
+    reg = pulser.Register.square(1, prefix="q")
+    seq = pulser.Sequence(reg, MockDevice)
+    seq.declare_channel("ch0", "rydberg_global")
+    N = 1000
+    amp_wf = BlackmanWaveform(N, np.pi)
+    det_wf = RampWaveform(N, -np.pi / 2, np.pi / 2)
+    phase = 1.234
+    seq.add(Pulse(amp_wf, det_wf, phase), "ch0")
+    seq.measure()
+
+    got = sample(seq)["Global"]["ground-rydberg"]
+    want = (amp_wf.samples, det_wf.samples, np.ones(N) * phase)
+    for i, key in enumerate(["amp", "det", "phase"]):
+        np.testing.assert_array_equal(got[key], want[i])
+
+
+def test_table_sequence(seqs):
+    """A table-driven test designed to be extended easily."""
+    for seq in seqs:
+        assert_same_samples_as_sim(seq)
+
+
+def test_inXY() -> None:
+    """Test sequence in XY mode."""
+    pulse = Pulse(
+        BlackmanWaveform(200, np.pi / 2),
+        RampWaveform(200, -np.pi / 2, np.pi / 2),
+        0.0,
+    )
+    reg = pulser.Register.from_coordinates(np.array([[0.0, 0.0]]), prefix="q")
+    seq = pulser.Sequence(reg, MockDevice)
+    seq.declare_channel("ch0", "mw_global")
+    seq.add(pulse, "ch0")
+    seq.measure(basis="XY")
+
+    assert_same_samples_as_sim(seq)
+
+
+def test_modulation(mod_seq: pulser.Sequence) -> None:
+    """Test sampling for modulated channels."""
+    N = mod_seq.get_duration()
+    chan = mod_seq.declared_channels["ch0"]
+    blackman = np.clip(np.blackman(N), 0, np.inf)
+    input = (np.pi / 2) / (np.sum(blackman) / N) * blackman
+
+    want = chan.modulate(input)
+    got = sample(mod_seq, modulation=True)["Global"]["ground-rydberg"]["amp"]
+
+    np.testing.assert_array_equal(got, want)
+
+
+@pytest.fixture
+def seq_with_SLM() -> pulser.Sequence:
+    q_dict = {
+        "batman": np.array([-4.0, 0.0]),  # sometimes masked
+        "superman": np.array([4.0, 0.0]),  # always unmasked
+    }
+
+    reg = pulser.Register(q_dict)
+    seq = pulser.Sequence(reg, MockDevice)
+
+    seq.declare_channel("ch0", "rydberg_global")
+    seq.config_slm_mask(["batman"])
+
+    seq.add(
+        Pulse.ConstantDetuning(BlackmanWaveform(200, np.pi / 2), 0.0, 0.0),
+        "ch0",
+    )
+    seq.add(
+        Pulse.ConstantDetuning(BlackmanWaveform(200, np.pi / 2), 0.0, 0.0),
+        "ch0",
+    )
+    seq.measure()
+    return seq
+
+
+def test_SLM_samples(seq_with_SLM):
+    pulse = Pulse.ConstantDetuning(BlackmanWaveform(200, np.pi / 2), 0.0, 0.0)
+    a_samples = pulse.amplitude.samples
+
+    def z() -> np.ndarray:
+        return np.zeros(seq_with_SLM.get_duration())
+
+    want: dict = {
+        "Global": {},
+        "Local": {
+            "ground-rydberg": {
+                "batman": {"amp": z(), "det": z(), "phase": z()},
+                "superman": {"amp": z(), "det": z(), "phase": z()},
+            }
+        },
+    }
+    want["Local"]["ground-rydberg"]["batman"]["amp"][200:401] = a_samples
+    want["Local"]["ground-rydberg"]["superman"]["amp"][0:200] = a_samples
+    want["Local"]["ground-rydberg"]["superman"]["amp"][200:401] = a_samples
+
+    got = sample(seq_with_SLM)
+    assert_nested_dict_equality(got, want)
+
+
+slm_reason = textwrap.dedent(
+    """
+If the SLM is on, Global channels decay to local ones in the
+sampler, such that the Global key in the output dict is empty and
+all the samples are written in the Local dict. On the contrary, the
+simulation module use the Local dict only for the first pulse, and
+then write the remaining in the Global dict.
+"""
+)
+
+
+@pytest.mark.xfail(reason=slm_reason)
+def test_SLM_against_simulation(seq_with_SLM):
+    assert_same_samples_as_sim(seq_with_SLM)
+
+
+def test_corner_cases():
+    """Test corner cases of helper functions."""
+    with pytest.raises(
+        ValueError,
+        match="ndarrays amp, det and phase must have the same length.",
+    ):
+        _ = QubitSamples(
+            amp=np.array([1.0]),
+            det=np.array([1.0]),
+            phase=np.array([1.0, 1.0]),
+            qubit="q0",
+        )
+
+    reg = pulser.Register.square(1, prefix="q")
+    seq = pulser.Sequence(reg, MockDevice)
+    N, M = 10, 11
+    samples_dict = {
+        "a": [QubitSamples(np.zeros(N), np.zeros(N), np.zeros(N), "q0")],
+        "b": [QubitSamples(np.zeros(M), np.zeros(M), np.zeros(M), "q0")],
+    }
+    with pytest.raises(
+        ValueError, match="All the samples do not share the same duration."
+    ):
+        _write_dict(seq, samples_dict, {})
+
+
+# Fixtures
+
+
+@pytest.fixture
+def seqs(seq_rydberg) -> list[pulser.Sequence]:
+    seqs: list[pulser.Sequence] = []
+
+    pulse = Pulse(
+        BlackmanWaveform(200, np.pi / 2),
+        RampWaveform(200, -np.pi / 2, np.pi / 2),
+        0.0,
+    )
+
+    reg = pulser.Register.from_coordinates(np.array([[0.0, 0.0]]), prefix="q")
+    seq = pulser.Sequence(reg, MockDevice)
+    seq.declare_channel("ch0", "raman_global")
+    seq.add(pulse, "ch0")
+    seq.measure()
+    seqs.append(deepcopy(seq))
+
+    seqs.append(seq_rydberg)
+
+    return seqs
+
+
+@pytest.fixture
+def seq_rydberg() -> pulser.Sequence:
+    reg = pulser.Register.from_coordinates(
+        np.array([[0.0, 0.0], [2.0, 0.0]]), prefix="q"
+    )
+    seq = pulser.Sequence(reg, MockDevice)
+    seq.declare_channel("ch0", "rydberg_global")
+    seq.declare_channel("ch1", "rydberg_local", initial_target="q0")
+    seq.add(
+        Pulse.ConstantDetuning(BlackmanWaveform(100, np.pi / 8), 0.0, 0.0),
+        "ch0",
+    )
+    seq.delay(20, "ch0")
+    seq.add(
+        Pulse.ConstantAmplitude(0.0, BlackmanWaveform(100, np.pi / 8), 0.0),
+        "ch0",
+    )
+    seq.add(
+        Pulse.ConstantDetuning(BlackmanWaveform(100, np.pi / 8), 0.0, 0.0),
+        "ch1",
+    )
+    seq.target("q1", "ch1")
+    seq.add(
+        Pulse.ConstantAmplitude(1.0, BlackmanWaveform(100, np.pi / 8), 0.0),
+        "ch1",
+    )
+    seq.target(["q0", "q1"], "ch1")
+    seq.add(
+        Pulse.ConstantDetuning(BlackmanWaveform(100, np.pi / 8), 0.0, 0.0),
+        "ch1",
+    )
+    seq.measure()
+    return seq
+
+
+@pytest.fixture
+def mod_seq(mod_device: Device) -> pulser.Sequence:
+    reg = pulser.Register.from_coordinates(np.array([[0.0, 0.0]]), prefix="q")
+    seq = pulser.Sequence(reg, mod_device)
+    seq.declare_channel("ch0", "rydberg_global")
+    seq.add(
+        Pulse.ConstantDetuning(BlackmanWaveform(1000, np.pi / 2), 0.0, 0.0),
+        "ch0",
+    )
+    seq.measure()
+    return seq
+
+
+@pytest.fixture
+def mod_device() -> Device:
+    return Device(
+        name="ModDevice",
+        dimensions=3,
+        rydberg_level=70,
+        max_atom_num=2000,
+        max_radial_distance=1000,
+        min_atom_distance=1,
+        _channels=(
+            (
+                "rydberg_global",
+                Rydberg(
+                    "Global",
+                    1000,
+                    200,
+                    clock_period=1,
+                    min_duration=1,
+                    mod_bandwidth=4.0,  # MHz
+                ),
+            ),
+            (
+                "rydberg_local",
+                Rydberg(
+                    "Local",
+                    2 * np.pi * 20,
+                    2 * np.pi * 10,
+                    max_targets=2,
+                    phase_jump_time=0,
+                    fixed_retarget_t=0,
+                    min_retarget_interval=220,
+                    mod_bandwidth=4.0,
+                ),
+            ),
+        ),
+    )