Add the classes for EOM mode configuration

pulser-core/pulser/channels/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2022 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.
+"""The various hardware channel types."""
+
+from pulser.channels.channels import Microwave, Raman, Rydberg

pulser-core/pulser/channels.py → pulser-core/pulser/channels/channels.py

@@ -11,20 +11,21 @@
 # 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.
-"""The various hardware channel types."""
+"""Defines the Channel ABC and its subclasses."""
 
 from __future__ import annotations
 
 import warnings
 from abc import ABC, abstractmethod
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass, field, fields
 from sys import version_info
 from typing import Any, Optional, cast
 
 import numpy as np
 from numpy.typing import ArrayLike
 from scipy.fft import fft, fftfreq, ifft
 
+from pulser.channels.eom import MODBW_TO_TR, BaseEOM, RydbergEOM
 from pulser.json.utils import obj_to_dict
 from pulser.pulse import Pulse
 
@@ -43,10 +44,6 @@
 # Warnings of adjusted waveform duration appear just once
 warnings.filterwarnings("once", "A duration of")
 
-# Conversion factor from modulation bandwith to rise time
-# For more info, see https://tinyurl.com/bdeumc8k
-MODBW_TO_TR = 0.48
-
 ADDRESSING = Literal["Global", "Local"]
 CH_TYPE = Literal["Rydberg", "Raman", "Microwave"]
 BASIS = Literal["ground-rydberg", "digital", "XY"]
@@ -94,6 +91,7 @@ class Channel(ABC):
     min_duration: int = 1  # ns
     max_duration: Optional[int] = int(1e8)  # ns
     mod_bandwidth: Optional[float] = None  # MHz
+    eom_config: Optional[BaseEOM] = field(init=False, default=None)
 
     @property
     def name(self) -> CH_TYPE:
@@ -210,6 +208,10 @@ def is_virtual(self) -> bool:
         """Whether the channel is virtual (i.e. partially defined)."""
         return bool(self._undefined_fields())
 
+    def supports_eom(self) -> bool:
+        """Whether the channel supports EOM mode operation."""
+        return hasattr(self, "eom_config") and self.eom_config is not None
+
     def _undefined_fields(self) -> list[str]:
         optional = [
             "max_amp",
@@ -366,48 +368,62 @@ def validate_pulse(self, pulse: Pulse) -> None:
             )
 
     def modulate(
-        self, input_samples: np.ndarray, keep_ends: bool = False
+        self,
+        input_samples: np.ndarray,
+        keep_ends: bool = False,
+        eom: bool = False,
     ) -> np.ndarray:
         """Modulates the input according to the channel's modulation bandwidth.
 
         Args:
             input_samples: The samples to modulate.
             keep_ends: Assume the end values of the samples were kept
                 constant (i.e. there is no ramp from zero on the ends).
+            eom: Whether to calculate the modulation using the EOM
+                bandwidth.
 
         Returns:
             The modulated output signal.
         """
-        if not self.mod_bandwidth:
+        if eom:
+            if not self.supports_eom():
+                raise TypeError(f"The channel {self} does not have an EOM.")
+            eom_config = cast(BaseEOM, self.eom_config)
+            mod_bandwidth = eom_config.mod_bandwidth
+            rise_time = eom_config.rise_time
+
+        elif not self.mod_bandwidth:
             warnings.warn(
                 f"No modulation bandwidth defined for channel '{self}',"
                 " 'Channel.modulate()' returns the 'input_samples' unchanged.",
                 stacklevel=2,
             )
             return input_samples
+        else:
+            mod_bandwidth = self.mod_bandwidth
+            rise_time = self.rise_time
 
         # The cutoff frequency (fc) and the modulation transfer function
         # are defined in https://tinyurl.com/bdeumc8k
-        fc = self.mod_bandwidth * 1e-3 / np.sqrt(np.log(2))
+        fc = mod_bandwidth * 1e-3 / np.sqrt(np.log(2))
         if keep_ends:
-            samples = np.pad(input_samples, 2 * self.rise_time, mode="edge")
+            samples = np.pad(input_samples, 2 * rise_time, mode="edge")
         else:
-            samples = np.pad(input_samples, self.rise_time)
+            samples = np.pad(input_samples, rise_time)
         freqs = fftfreq(samples.size)
         modulation = np.exp(-(freqs**2) / fc**2)
         mod_samples = ifft(fft(samples) * modulation).real
         if keep_ends:
             # Cut off the extra ends
-            return cast(
-                np.ndarray, mod_samples[self.rise_time : -self.rise_time]
-            )
+            return cast(np.ndarray, mod_samples[rise_time:-rise_time])
         return cast(np.ndarray, mod_samples)
 
     def calc_modulation_buffer(
         self,
         input_samples: ArrayLike,
         mod_samples: ArrayLike,
         max_allowed_diff: float = 1e-2,
+        eom: bool = False,
     ) -> tuple[int, int]:
         """Calculates the minimal buffers needed around a modulated waveform.
 
@@ -417,17 +433,23 @@ def calc_modulation_buffer(
                 ``len(input_samples) + 2 * self.rise_time``.
             max_allowed_diff: The maximum allowed difference between
                 the input and modulated samples at the end points.
+            eom: Whether to calculate the modulation buffers with the EOM
+                bandwidth.
 
         Returns:
             The minimum buffer times at the start and end of
             the samples, in ns.
         """
-        if not self.mod_bandwidth:
-            raise TypeError(
-                f"The channel {self} doesn't have a modulation bandwidth."
-            )
-
-        tr = self.rise_time
+        if eom:
+            if not self.supports_eom():
+                raise TypeError(f"The channel {self} does not have an EOM.")
+            tr = cast(BaseEOM, self.eom_config).rise_time
+        else:
+            if not self.mod_bandwidth:
+                raise TypeError(
+                    f"The channel {self} doesn't have a modulation bandwidth."
+                )
+            tr = self.rise_time
         samples = np.pad(input_samples, tr)
         diffs = np.abs(samples - mod_samples) <= max_allowed_diff
         try:
@@ -475,7 +497,10 @@ def __repr__(self) -> str:
         return self.name + config
 
     def _to_dict(self) -> dict[str, Any]:
-        return obj_to_dict(self, **asdict(self))
+        params = {
+            f.name: getattr(self, f.name) for f in fields(self) if f.init
+        }
+        return obj_to_dict(self, _module="pulser.channels", **params)
 
 
 @dataclass(init=True, repr=False, frozen=True)
@@ -500,6 +525,22 @@ class Rydberg(Channel):
     thus enconding the 'ground-rydberg' basis. See base class.
     """
 
+    eom_config: Optional[RydbergEOM] = None
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.eom_config is not None:
+            if not isinstance(self.eom_config, RydbergEOM):
+                raise TypeError(
+                    "When defined, 'eom_config' must be a valid 'RydbergEOM'"
+                    f" instance, not {type(self.eom_config)}."
+                )
+            if self.mod_bandwidth is None:
+                raise ValueError(
+                    "'eom_config' can't be defined in a Channel without a "
+                    "modulation bandwidth."
+                )
+
     @property
     def basis(self) -> Literal["ground-rydberg"]:
         """The addressed basis name."""

pulser-core/pulser/channels/eom.py

@@ -0,0 +1,170 @@
+# Copyright 2022 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.
+"""Configuration parameters for the a channel's EOM."""
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+from enum import Flag
+from itertools import chain
+from typing import Any, List, cast
+
+import numpy as np
+
+from pulser.json.utils import obj_to_dict
+
+# Conversion factor from modulation bandwith to rise time
+# For more info, see https://tinyurl.com/bdeumc8k
+MODBW_TO_TR = 0.48
+
+
+class RydbergBeam(Flag):
+    """The beams that make up a Rydberg channel."""
+
+    BLUE = 1
+    RED = 2
+
+    def _to_dict(self) -> dict[str, Any]:
+        return obj_to_dict(self, self.value)
+
+
+@dataclass(frozen=True)
+class BaseEOM:
+    """A base class for the EOM configuration."""
+
+    mod_bandwidth: float  # MHz
+
+    def __post_init__(self) -> None:
+        if self.mod_bandwidth <= 0.0:
+            raise ValueError(
+                "'mod_bandwidth' must be greater than zero, not"
+                f" {self.mod_bandwidth}."
+            )
+
+    @property
+    def rise_time(self) -> int:
+        """The rise time (in ns).
+
+        Defined as the time taken to go from 10% to 90% output in response to
+        a step change in the input.
+        """
+        return int(MODBW_TO_TR / self.mod_bandwidth * 1e3)
+
+    def _to_dict(self) -> dict[str, Any]:
+        params = {
+            f.name: getattr(self, f.name) for f in fields(self) if f.init
+        }
+        return obj_to_dict(self, **params)
+
+
+@dataclass(frozen=True)
+class RydbergEOM(BaseEOM):
+    """The EOM configuration for a Rydberg channel.
+
+    Attributes:
+        mod_bandwidth: The EOM modulation bandwidth at -3dB (50% reduction),
+            in MHz.
+        limiting_beam: The beam with the smallest amplitude range.
+        max_limiting_beam: The maximum amplitude the limiting beam can reach,
+            in rad/µs.
+        intermediate_detuning: The detuning between the two beams, in rad/µs.
+        controlled_beams: The beams that can be switched on/off with an EOM.
+    """
+
+    limiting_beam: RydbergBeam
+    max_limiting_amp: float  # rad/µs
+    intermediate_detuning: float  # rad/µs
+    controlled_beams: tuple[RydbergBeam, ...]
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        for param in ["max_limiting_amp", "intermediate_detuning"]:
+            value = getattr(self, param)
+            if value <= 0.0:
+                raise ValueError(
+                    f"'{param}' must be greater than zero, not {value}."
+                )
+        if not isinstance(self.controlled_beams, tuple):
+            if not isinstance(self.controlled_beams, list):
+                raise TypeError(
+                    "The 'controlled_beams' must be provided as a tuple "
+                    "or list."
+                )
+            # Convert list to tuple to keep RydbergEOM hashable
+            object.__setattr__(
+                self, "controlled_beams", tuple(self.controlled_beams)
+            )
+        if not self.controlled_beams:
+            raise ValueError(
+                "There must be at least one beam in 'controlled_beams'."
+            )
+        for beam in chain((self.limiting_beam,), self.controlled_beams):
+            if not (isinstance(beam, RydbergBeam) and beam in RydbergBeam):
+                raise TypeError(
+                    "Every beam must be one of options of the `RydbergBeam`"
+                    f" enumeration, not {self.limiting_beam}."
+                )
+
+    def detuning_off_options(
+        self, rabi_frequency: float, detuning_on: float
+    ) -> list[float]:
+        """Calculates the possible detuning values when the amplitude is off.
+
+        Args:
+            rabi_frequency: The Rabi frequency when executing a pulse,
+                in rad/µs.
+            detuning_on: The detuning when executing a pulse, in rad/µs.
+
+        Returns:
+            The possible detuning values when in between pulses.
+        """
+        offset = detuning_on - self._lightshift(rabi_frequency, *RydbergBeam)
+        if len(self.controlled_beams) == 1:
+            lightshifts = [
+                self._lightshift(rabi_frequency, ~self.controlled_beams[0])
+            ]
+
+        else:
+            lightshifts = [
+                self._lightshift(rabi_frequency, beam)
+                for beam in self.controlled_beams
+            ]
+            # Extra case where both beams are off
+            lightshifts.append(0.0)
+        return cast(List[float], (offset + np.array(lightshifts)).tolist())
+
+    def _lightshift(
+        self, rabi_frequency: float, *beams_on: RydbergBeam
+    ) -> float:
+        rabi_freqs = self._rabi_freq_per_beam(rabi_frequency)
+        bias = {RydbergBeam.RED: -1, RydbergBeam.BLUE: 1}
+        return sum(bias[beam] * rabi_freqs[beam] ** 2 for beam in beams_on) / (
+            4 * self.intermediate_detuning
+        )
+
+    def _rabi_freq_per_beam(
+        self, rabi_frequency: float
+    ) -> dict[RydbergBeam, float]:
+        limit_rabi_freq = self.max_limiting_amp**2 / (
+            2 * self.intermediate_detuning
+        )
+        if rabi_frequency <= limit_rabi_freq:
+            beam_amp = np.sqrt(2 * rabi_frequency * self.intermediate_detuning)
+            return {beam: beam_amp for beam in RydbergBeam}
+        return {
+            self.limiting_beam: self.max_limiting_amp,
+            ~self.limiting_beam: 2
+            * self.intermediate_detuning
+            * rabi_frequency
+            / self.max_limiting_amp,
+        }

pulser-core/pulser/devices/_device_datacls.py

@@ -23,7 +23,7 @@
 import numpy as np
 from scipy.spatial.distance import pdist, squareform
 
-from pulser.channels import Channel
+from pulser.channels.channels import Channel
 from pulser.devices.interaction_coefficients import c6_dict
 from pulser.json.utils import obj_to_dict
 from pulser.register.base_register import BaseRegister, QubitId