Address qubits with indices

pulser/register/mappable_reg.py

@@ -22,6 +22,7 @@
 if TYPE_CHECKING:  # pragma: no cover
     from pulser.register.base_register import BaseRegister, QubitId
     from pulser.register.register_layout import RegisterLayout
+from typing import Sequence as abcSequence
 
 
 class MappableRegister:
@@ -67,9 +68,50 @@ def build_register(self, qubits: Mapping[QubitId, int]) -> BaseRegister:
             raise ValueError(
                 "All qubits must be labeled with pre-declared qubit IDs."
             )
+        trap_ordered_qubits = {
+            id: qubits[id] for id in self._qubit_ids if id in chosen_ids
+        }
         return self._layout.define_register(
-            *tuple(qubits.values()), qubit_ids=chosen_ids
+            *tuple(trap_ordered_qubits.values()),
+            qubit_ids=tuple(trap_ordered_qubits.keys()),
         )
 
+    def find_indices(
+        self, chosen_ids: set[QubitId], id_list: abcSequence[QubitId]
+    ) -> list[int]:
+        """Computes indices of qubits according to a register mapping.
+
+        This can especially be useful when building a Pulser Sequence
+        with a parameter denoting qubits.
+
+        Example:
+            ``
+            mapp_reg = TriangularLatticeLayout(50, 5).make_mappable_register(5)
+            qubit_map = {"q0": 1, "q2": 4, "q4": 2}
+            indices = mapp_reg.find_indices(
+                qubit_map,
+                ["q4", "q2", "q1", "q2"])
+            seq.build(qubits=qubit_map, qubit_indices=indices)
+            ``
+
+        Args:
+            chosen_ids: IDs of the qubits that are chosen to map the
+                MappableRegister
+            id_list: IDs of the qubits to denote
+
+        Returns:
+            Indices of the qubits to denote, only valid for the given mapping.
+        """
+        if not chosen_ids <= set(self._qubit_ids):
+            raise ValueError(
+                "Chosen IDs must be selected among pre-declared qubit IDs."
+            )
+        if not set(id_list) <= chosen_ids:
+            raise ValueError(
+                "The IDs list must be selected among the chosen IDs."
+            )
+        ordered_ids = [id for id in self.qubit_ids if id in chosen_ids]
+        return [ordered_ids.index(id) for id in id_list]
+
     def _to_dict(self) -> dict[str, Any]:
         return obj_to_dict(self, self._layout, *self._qubit_ids)

pulser/sequence.py

@@ -20,7 +20,7 @@
 import os
 import warnings
 from collections import namedtuple
-from collections.abc import Callable, Generator, Iterable, Mapping
+from collections.abc import Callable, Generator, Iterable, Mapping, Set
 from functools import wraps
 from itertools import chain
 from sys import version_info
@@ -98,38 +98,53 @@ def wrapper(self: Sequence, *args: Any, **kwargs: Any) -> Any:
     return cast(F, wrapper)
 
 
-def _store(func: F) -> F:
-    """Stores any Sequence building call for deferred execution."""
+def _verify_variable(seq: Sequence, x: Any) -> None:
+    if isinstance(x, Parametrized):
+        # If not already, the sequence becomes parametrized
+        seq._building = False
+        for name, var in x.variables.items():
+            if name not in seq._variables:
+                raise ValueError(f"Unknown variable '{name}'.")
+            elif seq._variables[name] is not var:
+                raise ValueError(
+                    f"{x} has variables that don't come from this "
+                    "Sequence. Use only what's returned by this"
+                    "Sequence's 'declare_variable' method as your"
+                    "variables."
+                )
+    elif isinstance(x, Iterable) and not isinstance(x, str):
+        # Recursively look for parametrized objs inside the arguments
+        for y in x:
+            _verify_variable(seq, y)
+
+
+def _verify_parametrization(func: F) -> F:
+    """Checks and updates the sequence status' consistency with the call.
+
+    - Checks the sequence can still be modified.
+    - Checks if all Parametrized inputs stem from declared variables.
+    """
 
     @wraps(func)
     def wrapper(self: Sequence, *args: Any, **kwargs: Any) -> Any:
-        def verify_variable(x: Any) -> None:
-            if isinstance(x, Parametrized):
-                # If not already, the sequence becomes parametrized
-                self._building = False
-                for name, var in x.variables.items():
-                    if name not in self._variables:
-                        raise ValueError(f"Unknown variable '{name}'.")
-                    elif self._variables[name] is not var:
-                        raise ValueError(
-                            f"{x} has variables that don't come from this "
-                            "Sequence. Use only what's returned by this"
-                            "Sequence's 'declare_variable' method as your"
-                            "variables."
-                        )
-            elif isinstance(x, Iterable) and not isinstance(x, str):
-                # Recursively look for parametrized objs inside the arguments
-                for y in x:
-                    verify_variable(y)
-
         if self._is_measured and self.is_parametrized():
             raise RuntimeError(
                 "The sequence has been measured, no further "
                 "changes are allowed."
             )
-        # Check if all Parametrized inputs stem from declared variables
         for x in chain(args, kwargs.values()):
-            verify_variable(x)
+            _verify_variable(self, x)
+        func(self, *args, **kwargs)
+
+    return cast(F, wrapper)
+
+
+def _store(func: F) -> F:
+    """Checks and stores the call to call it when building the Sequence."""
+
+    @wraps(func)
+    @_verify_parametrization
+    def wrapper(self: Sequence, *args: Any, **kwargs: Any) -> Any:
         storage = self._calls if self._building else self._to_build_calls
         func(self, *args, **kwargs)
         storage.append(_Call(func.__name__, args, kwargs))
@@ -579,73 +594,6 @@ def declare_channel(
             )
         )
 
-    @overload
-    def declare_variable(
-        self,
-        name: str,
-        *,
-        dtype: Union[type[int], type[float], type[str]] = float,
-    ) -> VariableItem:
-        pass
-
-    @overload
-    def declare_variable(
-        self,
-        name: str,
-        *,
-        size: int,
-        dtype: Union[type[int], type[float], type[str]] = float,
-    ) -> Variable:
-        pass
-
-    def declare_variable(
-        self,
-        name: str,
-        size: Optional[int] = None,
-        dtype: Union[type[int], type[float], type[str]] = float,
-    ) -> Union[Variable, VariableItem]:
-        """Declare a new variable within this Sequence.
-
-        The declared variables can be used to create parametrized versions of
-        ``Waveform`` and ``Pulse`` objects, which in turn can be added to the
-        ``Sequence``. Additionally, simple arithmetic operations involving
-        variables are also supported and will return parametrized objects that
-        are dependent on the involved variables.
-
-        Args:
-            name (str): The name for the variable. Must be unique within a
-                Sequence.
-
-        Keyword Args:
-            size (int=1): The number of entries stored in the variable.
-            dtype (default=float): The type of the data that will be assigned
-                to the variable. Must be ``float``, ``int`` or ``str``.
-
-        Returns:
-            Variable: The declared Variable instance.
-
-        Note:
-            To avoid confusion, it is recommended to store the returned
-            Variable instance in a Python variable with the same name.
-        """
-        if name == "qubits":
-            # Necessary because 'qubits' is a keyword arg in self.build()
-            raise ValueError(
-                "'qubits' is a protected name. Please choose a different name "
-                "for the variable."
-            )
-
-        if name in self._variables:
-            raise ValueError("Name for variable is already being used.")
-
-        if size is None:
-            var = self.declare_variable(name, size=1, dtype=dtype)
-            return var[0]
-        else:
-            var = Variable(name, dtype, size=size)
-            self._variables[name] = var
-            return var
-
     @_store
     def add(
         self,
@@ -833,6 +781,33 @@ def target(
 
         self._target(qubits, channel)
 
+    @_verify_parametrization
+    def target_index(
+        self,
+        qubits: Union[int, Iterable[int], Parametrized],
+        channel: Union[str, Parametrized],
+    ) -> None:
+        """Changes the target qubit of a 'Local' channel.
+
+        Args:
+            qubits (Union[int, str, Iterable]): The new target for this
+                channel. Must correspond to a qubit index (a number between 0
+                 and the number of qubits, which is determined at build time
+                 when using a MappableRegister) or an iterable
+                of qubit indices, when multi-qubit addressing is possible.
+            channel (str): The channel's name provided when declared. Must be
+                a channel with 'Local' addressing.
+        """
+        if not self.is_parametrized():
+            raise RuntimeError(
+                f"Sequence.{self.target_index.__name__} can't be called in"
+                " non parametrized sequences."
+            )
+
+        qubits = cast(int, qubits)
+        channel = cast(str, channel)
+        self._target_index(qubits, channel)
+
     @_store
     def delay(
         self,
@@ -1186,9 +1161,88 @@ def draw(
             fig.savefig(fig_name, **kwargs_savefig)
         plt.show()
 
-    def _target(
-        self, qubits: Union[Iterable[QubitId], QubitId], channel: str
-    ) -> None:
+    @overload
+    def declare_variable(
+        self,
+        name: str,
+        *,
+        dtype: Union[type[int], type[float], type[str]] = float,
+    ) -> VariableItem:
+        pass
+
+    @overload
+    def declare_variable(
+        self,
+        name: str,
+        *,
+        size: int,
+        dtype: Union[type[int], type[float], type[str]] = float,
+    ) -> Variable:
+        pass
+
+    def declare_variable(
+        self,
+        name: str,
+        size: Optional[int] = None,
+        dtype: Union[type[int], type[float], type[str]] = float,
+    ) -> Union[Variable, VariableItem]:
+        """Declare a new variable within this Sequence.
+
+        The declared variables can be used to create parametrized versions of
+        ``Waveform`` and ``Pulse`` objects, which in turn can be added to the
+        ``Sequence``. Additionally, simple arithmetic operations involving
+        variables are also supported and will return parametrized objects that
+        are dependent on the involved variables.
+
+        Args:
+            name (str): The name for the variable. Must be unique within a
+                Sequence.
+
+        Keyword Args:
+            size (int=1): The number of entries stored in the variable.
+            dtype (default=float): The type of the data that will be assigned
+                to the variable. Must be ``float``, ``int`` or ``str``.
+
+        Returns:
+            Variable: The declared Variable instance.
+
+        Note:
+            To avoid confusion, it is recommended to store the returned
+            Variable instance in a Python variable with the same name.
+        """
+        if name == "qubits":
+            # Necessary because 'qubits' is a keyword arg in self.build()
+            raise ValueError(
+                "'qubits' is a protected name. Please choose a different name "
+                "for the variable."
+            )
+
+        if name in self._variables:
+            raise ValueError("Name for variable is already being used.")
+
+        if size is None:
+            var = self.declare_variable(name, size=1, dtype=dtype)
+            return var[0]
+        else:
+            var = Variable(name, dtype, size=size)
+            self._variables[name] = var
+            return var
+
+    @overload
+    def _precheck_target_qubits_set(
+            self, qubits: Union[Iterable[int], int], channel: str
+    ) -> Union[Set[int]]:
+        pass
+
+    @overload
+    def _precheck_target_qubits_set(
+            self, qubits: Union[Iterable[QubitId], QubitId], channel: str
+    ) -> Union[Set[QubitId], Set[int]]:
+        pass
+
+    def _precheck_target_qubits_set(
+            self, qubits: Union[Iterable[QubitId], QubitId], channel: str
+    ) -> Union[Set[QubitId], Set[int]]:
         self._validate_channel(channel)
         channel_obj = self._channels[channel]
         try:
@@ -1214,11 +1268,38 @@ def _target(
                     raise ValueError(
                         "All non-variable qubits must belong to the register."
                     )
-            return
 
-        elif not qubits_set.issubset(self._qids):
-            raise ValueError("All given qubits must belong to the register.")
+        return qubits_set
+
+    def _target(
+        self, qubits: Union[Iterable[QubitId], QubitId], channel: str
+    ) -> None:
+        qubits_set = self._precheck_target_qubits_set(qubits, channel)
+        if not self.is_parametrized():
+            self._perform_target_non_parametrized(qubits_set, channel)
+
+    @_store
+    def _target_index(
+        self, qubits: Union[Iterable[int], int], channel: str
+    ) -> None:
+
+        qubits_set = self._precheck_target_qubits_set(qubits, channel)
+        if not self.is_parametrized():
+            qubit_ids_set = {
+                self.register.qubit_ids[index] for index in qubits_set
+            }  # TODO try except
+            self._perform_target_non_parametrized(qubit_ids_set, channel)
+
+    def _perform_target_non_parametrized(
+        self, qubits_set: Set[QubitId], channel: str
+    ) -> None:
+        for qubit in qubits_set:
+            if qubit not in self._qids:
+                raise ValueError(
+                    f"The qubit ID '{qubit}' does not belong to the register."
+                )
 
+        channel_obj = self._channels[channel]
         basis = channel_obj.basis
         phase_refs = {self._phase_ref[basis][q].last_phase for q in qubits_set}
         if len(phase_refs) != 1:
@@ -1261,7 +1342,9 @@ def _target(
             tf = 0
 
         self._last_target[channel] = tf
-        self._add_to_schedule(channel, _TimeSlot("target", ti, tf, qubits_set))
+        self._add_to_schedule(
+            channel, _TimeSlot("target", ti, tf, set(qubits_set))
+        )
 
     def _delay(self, duration: int, channel: str) -> None:
         self._validate_channel(channel)