Sequence sampler (#345)

* A first gist of sequence sampler

plain dataframe of samples

extract samples in the nested dict form

A note: basis that are unused by declared channels are arrays full of
zeros, while it is simply not initilized in the dict in the simulation
module.

add test to compare with extraction from simulation module

Add various docstrings

type annotations

sample per qubit, test passes

add support for LocalNoises

add seed to the amplitude_noise generator

fix the doppler noise to use a normal distribution

Write test for doppler noise. It fails.

It fails probably because of the different seeds in the random
generators ?

refactor a bit

Rename as module sampler

Refactor by distinguishing the local and global extraction

Remove unused import

Refactor the noises to their own modules

Add modulation, but no implementation

Revert "Refactor the noises to their own modules"

This reverts commit 37093330c57a9152604f239e79ec6940a213b432.

Use strategies for _TimeSlot extraction

Remove functions made useless by refactor to strategy pattern

Add doctrings and make some functions private

Extend test to local channels

Remove dead code

Add modulation

Refactor a bit: docstrings, order of functions, etc.

Fix _TimeSlot grouping by handling delays correctly

Tidy the comments

refactor the extraction strategies

Docstring

* Refactor sampler in independent module

* Consistent support for noises, SLM and modulation of Global and Local 

Deprecate temporarily the amplitude_noise as it misses the point

It urges for a correct implementation

Refactor sample writing

Allow local decay with on SLM

Add support for noises on global channel

Add TODO for global modulation

Fix usage of noises.apply()

Fix modulation and apply SLM on local channels as well

Fix imports

* Refactor the noises module

Also fix imports for mypy compliance

* Improve docstrings and small refactoring

* Small typos in comments

* Fix the decay of global channels to local ones

* Fix the modulation feature

* Refactor

Adds a post_init check for same length of quantities arrays

Add documentation, testing and refactoring

Refactor the dict construction with modulation

Simple comment to clarify the design of sampler.sample()

Remove unnecessary value in enum _GroupType

Add a diagram for _TimeSlot grouping

Fix embedded NoiseModel by using copies

Update docstrings with args

Expose the sample() function from the __init__

Test the amplitude noise

Remove unused commented import

Add usage note on doppler noise

Add a docstring for the amplitude noise

Remove empty line

Create a better test for doppler noise

Small refactor

Mark the doppler test as expected to fail

Refactor the control flow of sampler.sample()

Refactor the dict exporting

Refactor

Fix typo

Add a docstring

Refactor the TimeSlot grouping and strategy

* Simplify the flow of samples.samples()

Global and local channels are sampled identically.
The export to a dict form like the simulation one handles the
distinctions.
It introduces an additional dict keeping track of the addressing of each
channel: Global, Decayed, Local.

Trim dead code, refactor, defensive coding

Write global channel from QubitSamples

Rename modulation function

* Improve test and get total coverage

Table based tests

Add a modulation test

I don't like it. It is testing the sampling for sure, but still I would
like to test against a known output of the modulation.

Remove unnecessary if branch

functools.reduce got us covered already

Extend to digital

Fix mypy error because of weird scope of variables

Test sequence in XY and fix a typo in the relevant code

Test for corner cases and omit defensive coding with pragma no cover

* Simplify the sampling by removing the Samples dataclass

* Fix a docstring

* Fix docstrings formatting and content

* Remove unnecessary nonzero check in amplitude noise

* Fix the noise seed default value: defaults to None

* Improve docstrings of helper functions in noise module

* Apply noises before the SLM masking

* Rename misnamed variables

* Change the scope of _key_func() and remove the _GroupType class

* Fix dictionary construction with defaultdict

* Fix misnamed keys in tests

* Add a match to pytest.raises

* Move the noises to the simulation module

Import NoiseModel as noises.NoiseModel to avoid a weird circular import

* Improve the simulation.noises docstring

* Remove noisy helper functions

* Write test case for _write_dict exception

It removes the # pragma: no cover.

* Reorder tests

* Add a more fundamental test

* Reword and add docstrings

* Fix other misnamed keys in tests

* Fix SLM masking: consider only the seq._mask_times

Sequence._slm_mask_time is not a list of times, but a pair with a
start and end time of the SLM masking.

* Change for defaultdict in new_qdict

* Add a note about performance for sampling of global channels

For global channels, we hold the same data in N copies, N the number of
qubits of the register. It scales poorly with the size of the register,
but it remains manageable for current size of registers.

If needed, we should patch this, at the expense of a more complex
control flow in the sample() function at least.

* Simply testing of sequence

* Refactor the testing of sequence with a helper function

* Fix the testing of blackman modulation

* Move noise-related tests to an independent testing files

* Fix the unneeded # pragma: no cover in _sample_slots

* Add a Failing test for the SLM, for discussion

* Move the comment on performance

* Make the SLM detectino more idiomatic

* Simplify the assertions for sequence in XY

* Remove superfluous print statements

* Test SLM sampling alongside the check against simulation

* Refactor using a helper for nested dicts

* Rearrange the test file

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

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

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."
+            )