Add states labels to Channels and SequenceSamples (#705)

* Associate states to bases

* Add states to conventions

* Use states in Hamiltonian

* Add Channel.eigenstates, corespondance btw eigenstates and labels

* Fixing widths in table

* Revert changes to convetions, make table in Channels docstring

* Add r"""

* Fix indentation

* Fix table in eigenstates docstring

* Fix typo

* Add multiple_bases_states, check for eigenstates

* Sort imports

* Move test on EIGENSTATES to unit tests

* Change name of multiple_bases_states

* Fix typo

* Fix import of Collection

pulser-core/pulser/channels/base_channel.py

@@ -17,8 +17,9 @@
 
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Collection
 from dataclasses import MISSING, dataclass, field, fields
-from typing import Any, Literal, Optional, Type, TypeVar, cast
+from typing import Any, Literal, Optional, Type, TypeVar, cast, get_args
 
 import numpy as np
 from numpy.typing import ArrayLike
@@ -35,6 +36,23 @@
 
 OPTIONAL_ABSTR_CH_FIELDS = ("min_avg_amp",)
 
+# States ranked in decreasing order of their associated eigenenergy
+States = Literal["u", "d", "r", "g", "h"]  # TODO: add "x" for leakage
+
+STATES_RANK = get_args(States)
+
+EIGENSTATES: dict[str, list[States]] = {
+    "ground-rydberg": ["r", "g"],
+    "digital": ["g", "h"],
+    "XY": ["u", "d"],
+}
+
+
+def get_states_from_bases(bases: Collection[str]) -> list[States]:
+    """The states associated to a list of bases, ranked by their energies."""
+    all_states = set().union(*(set(EIGENSTATES[basis]) for basis in bases))
+    return [state for state in STATES_RANK if state in all_states]
+
 
 @dataclass(init=True, repr=False, frozen=True)
 class Channel(ABC):
@@ -90,12 +108,45 @@ def basis(self) -> str:
         """The addressed basis name."""
         pass
 
+    @property
+    def eigenstates(self) -> list[States]:
+        r"""The eigenstates associated with the basis.
+
+        Returns a tuple of labels, ranked in decreasing order
+        of their associated eigenenergy, as such:
+
+        .. list-table::
+            :align: center
+            :widths: 50 35 35
+            :header-rows: 1
+
+            * - Name
+              - Eigenstate (see :doc:`/conventions`)
+              - Associated label
+            * - Up state
+              - :math:`|0\rangle`
+              - ``"u"``
+            * - Down state
+              - :math:`|1\rangle`
+              - ``"d"``
+            * - Rydberg state
+              - :math:`|r\rangle`
+              - ``"r"``
+            * - Ground state
+              - :math:`|g\rangle`
+              - ``"g"``
+            * - Hyperfine state
+              - :math:`|h\rangle`
+              - ``"h"``
+        """
+        return EIGENSTATES[self.basis]
+
     @property
     def _internal_param_valid_options(self) -> dict[str, tuple[str, ...]]:
         """Internal parameters and their valid options."""
         return dict(
             name=("Rydberg", "Raman", "Microwave", "DMM"),
-            basis=("ground-rydberg", "digital", "XY"),
+            basis=tuple(EIGENSTATES.keys()),
             addressing=("Local", "Global"),
         )
 

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.base_channel import Channel
+from pulser.channels.base_channel import Channel, States, get_states_from_bases
 from pulser.channels.dmm import DMM
 from pulser.devices.interaction_coefficients import c6_dict
 from pulser.json.abstract_repr.serializer import AbstractReprEncoder
@@ -270,6 +270,11 @@ def supported_bases(self) -> set[str]:
         """Available electronic transitions for control and measurement."""
         return {ch.basis for ch in self.channel_objects}
 
+    @property
+    def supported_states(self) -> list[States]:
+        """Available states ranked by their energy levels (highest first)."""
+        return get_states_from_bases(self.supported_bases)
+
     @property
     def interaction_coeff(self) -> float:
         r"""The interaction coefficient for the chosen Rydberg level.

pulser-core/pulser/sampler/samples.py

@@ -9,7 +9,12 @@
 
 import numpy as np
 
-from pulser.channels.base_channel import Channel
+from pulser.channels.base_channel import (
+    EIGENSTATES,
+    Channel,
+    States,
+    get_states_from_bases,
+)
 from pulser.channels.eom import BaseEOM
 from pulser.register import QubitId
 from pulser.register.weight_maps import DetuningMap
@@ -468,6 +473,13 @@ def used_bases(self) -> set[str]:
             if not ch_samples.is_empty()
         }
 
+    @property
+    def eigenbasis(self) -> list[States]:
+        """The basis of eigenstates used for simulation."""
+        if len(self.used_bases) == 0:
+            return EIGENSTATES["XY" if self._in_xy else "ground-rydberg"]
+        return get_states_from_bases(self.used_bases)
+
     @property
     def _in_xy(self) -> bool:
         """Checks if the sequence is in XY mode."""

pulser-core/pulser/sequence/sequence.py

@@ -41,7 +41,7 @@
 import pulser
 import pulser.devices as devices
 import pulser.sequence._decorators as seq_decorators
-from pulser.channels.base_channel import Channel
+from pulser.channels.base_channel import Channel, States, get_states_from_bases
 from pulser.channels.dmm import DMM, _dmm_id_from_name, _get_dmm_name
 from pulser.channels.eom import RydbergEOM
 from pulser.devices._device_datacls import BaseDevice
@@ -460,6 +460,10 @@ def get_addressed_bases(self) -> tuple[str, ...]:
         """Returns the bases addressed by the declared channels."""
         return tuple(self._basis_ref)
 
+    def get_addressed_states(self) -> list[States]:
+        """Returns the states addressed by the declared channels."""
+        return get_states_from_bases(self.get_addressed_bases())
+
     @seq_decorators.screen
     def current_phase_ref(
         self, qubit: QubitId, basis: str = "digital"

pulser-simulation/pulser_simulation/hamiltonian.py

@@ -23,6 +23,7 @@
 import numpy as np
 import qutip
 
+from pulser.channels.base_channel import STATES_RANK
 from pulser.devices._device_datacls import BaseDevice
 from pulser.noise_model import NoiseModel
 from pulser.register.base_register import QubitId
@@ -315,35 +316,34 @@ def _update_noise(self) -> None:
 
     def _build_basis_and_op_matrices(self) -> None:
         """Determine dimension, basis and projector operators."""
-        if self._interaction == "XY":
-            self.basis_name = "XY"
-            self.dim = 2
-            basis = ["u", "d"]
-            projectors = ["uu", "du", "ud", "dd"]
-        else:
-            if "digital" not in self.samples_obj.used_bases:
-                self.basis_name = "ground-rydberg"
-                self.dim = 2
-                basis = ["r", "g"]
-                projectors = ["gr", "rr", "gg"]
-            elif "ground-rydberg" not in self.samples_obj.used_bases:
-                self.basis_name = "digital"
-                self.dim = 2
-                basis = ["g", "h"]
-                projectors = ["hg", "hh", "gg"]
+        if len(self.samples_obj.used_bases) == 0:
+            if self.samples_obj._in_xy:
+                self.basis_name = "XY"
             else:
-                self.basis_name = "all"  # All three states
-                self.dim = 3
-                basis = ["r", "g", "h"]
-                projectors = ["gr", "hg", "rr", "gg", "hh"]
+                self.basis_name = "ground-rydberg"
+        elif len(self.samples_obj.used_bases) == 1:
+            self.basis_name = list(self.samples_obj.used_bases)[0]
+        else:
+            self.basis_name = "all"  # All three rydberg states
+        eigenbasis = self.samples_obj.eigenbasis
 
-        self.basis = {b: qutip.basis(self.dim, i) for i, b in enumerate(basis)}
-        self.op_matrix = {"I": qutip.qeye(self.dim)}
+        # TODO: Add leakage
 
-        for proj in projectors:
-            self.op_matrix["sigma_" + proj] = (
-                self.basis[proj[0]] * self.basis[proj[1]].dag()
-            )
+        self.eigenbasis = [
+            state for state in STATES_RANK if state in eigenbasis
+        ]
+
+        self.dim = len(self.eigenbasis)
+        self.basis = {
+            b: qutip.basis(self.dim, i) for i, b in enumerate(self.eigenbasis)
+        }
+        self.op_matrix = {"I": qutip.qeye(self.dim)}
+        for proj0 in self.eigenbasis:
+            for proj1 in self.eigenbasis:
+                proj_name = "sigma_" + proj0 + proj1
+                self.op_matrix[proj_name] = (
+                    self.basis[proj0] * self.basis[proj1].dag()
+                )
 
     def _construct_hamiltonian(self, update: bool = True) -> None:
         """Constructs the hamiltonian from the sampled Sequence and noise.

tests/test_channels.py

@@ -20,6 +20,7 @@
 import pulser
 from pulser import Pulse
 from pulser.channels import Microwave, Raman, Rydberg
+from pulser.channels.base_channel import EIGENSTATES, STATES_RANK
 from pulser.channels.eom import MODBW_TO_TR, BaseEOM, RydbergBeam, RydbergEOM
 from pulser.waveforms import BlackmanWaveform, ConstantWaveform
 
@@ -140,6 +141,21 @@ def test_device_channels():
                 assert ch.max_targets == int(ch.max_targets)
 
 
+def test_eigenstates():
+    for _, states in EIGENSTATES.items():
+        idx_0, idx_1 = STATES_RANK.index(states[0]), STATES_RANK.index(
+            states[1]
+        )
+        assert idx_0 != -1 and idx_1 != -1, f"States must be in {STATES_RANK}."
+        assert (
+            idx_0 < idx_1
+        ), "Eigenstates must be ranked with highest energy first."
+
+    assert Raman.Global(None, None).eigenstates == ["g", "h"]
+    assert Rydberg.Global(None, None).eigenstates == ["r", "g"]
+    assert Microwave.Global(None, None).eigenstates == ["u", "d"]
+
+
 def test_validate_duration():
     ch = Rydberg.Local(20, 10, min_duration=16, max_duration=1000)
     with pytest.raises(TypeError, match="castable to an int"):

tests/test_devices.py

@@ -22,7 +22,12 @@
 import pulser
 from pulser.channels import Microwave, Raman, Rydberg
 from pulser.channels.dmm import DMM
-from pulser.devices import Device, DigitalAnalogDevice, VirtualDevice
+from pulser.devices import (
+    Device,
+    DigitalAnalogDevice,
+    MockDevice,
+    VirtualDevice,
+)
 from pulser.register import Register, Register3D
 from pulser.register.register_layout import RegisterLayout
 from pulser.register.special_layouts import (
@@ -188,6 +193,30 @@ def test_default_channel_ids(test_params):
     )
 
 
+@pytest.mark.parametrize(
+    "channels, states",
+    [
+        ((Rydberg.Local(None, None),), ["r", "g"]),
+        ((Raman.Local(None, None),), ["g", "h"]),
+        (DigitalAnalogDevice.channel_objects, ["r", "g", "h"]),
+        (
+            (
+                Microwave.Global(None, None),
+                Raman.Global(None, None),
+            ),
+            ["u", "d", "g", "h"],
+        ),
+        ((Microwave.Global(None, None),), ["u", "d"]),
+        (MockDevice.channel_objects, ["u", "d", "r", "g", "h"]),
+    ],
+)
+def test_eigenstates(test_params, channels, states):
+    test_params["interaction_coeff_xy"] = 10000.0
+    test_params["channel_objects"] = channels
+    dev = VirtualDevice(**test_params)
+    assert dev.supported_states == states
+
+
 def test_tuple_conversion(test_params):
     test_params["channel_objects"] = [Rydberg.Global(None, None)]
     test_params["channel_ids"] = ["custom_channel"]

tests/test_sequence.py

@@ -85,12 +85,15 @@ def test_channel_declaration(reg, device):
     seq = Sequence(reg, device)
     available_channels = set(seq.available_channels)
     assert seq.get_addressed_bases() == ()
+    assert seq.get_addressed_states() == []
     with pytest.raises(ValueError, match="Name starting by 'dmm_'"):
         seq.declare_channel("dmm_1_2", "raman")
     seq.declare_channel("ch0", "rydberg_global")
     assert seq.get_addressed_bases() == ("ground-rydberg",)
+    assert seq.get_addressed_states() == ["r", "g"]
     seq.declare_channel("ch1", "raman_local")
     assert seq.get_addressed_bases() == ("ground-rydberg", "digital")
+    assert seq.get_addressed_states() == ["r", "g", "h"]
     with pytest.raises(ValueError, match="No channel"):
         seq.declare_channel("ch2", "raman")
     with pytest.raises(ValueError, match="not available"):
@@ -129,6 +132,8 @@ def test_channel_declaration(reg, device):
         match="cannot work simultaneously with the declared 'Microwave'",
     ):
         seq2.declare_channel("ch3", "rydberg_global")
+    assert seq2.get_addressed_bases() == ("XY",)
+    assert seq2.get_addressed_states() == ["u", "d"]
 
 
 def test_dmm_declaration(reg, device, det_map):