Repository: pasqal-io/Pulser

Issue:
Register ids type change after serialization and deserialization
When using numbers as atom ids, they are converted to strings when going through the serialization and deserialization. The consequence is that the outcome will differ from the original value.

For example:

```python
from pulser.register import Register
reg = Register(dict(enumerate([(2,3),(5,1),(10,0)])))
j=json.dumps(reg, cls=PulserEncoder)
r=json.loads(j, cls=PulserDecoder)
assert reg == r # raises because ids do not have the same types
```

This is due to json converting object ids to strings. This might lead to undesired behaviors, because we never know what consequences we might meet by having strings instead of integers for the ids.

I can find two solutions:

- either fix this type change by replacing the json objects with json lists, where the first element would represent the id: this way, the type will be kept
- or systmetically convert int ids to strings, even when not using serialization: the behavior with or without encoding/decoding will be identical

Discussion:

Patch:
diff --git a/pulser-core/pulser/json/supported.py b/pulser-core/pulser/json/supported.py
index 611a730..c38cedf 100644
--- a/pulser-core/pulser/json/supported.py
+++ b/pulser-core/pulser/json/supported.py
@@ -48,7 +48,13 @@ SUPPORTED_NUMPY = (
     "tan",
 )
 
-SUPPORTS_SUBMODULE = ("Pulse", "BlackmanWaveform", "KaiserWaveform")
+SUPPORTS_SUBMODULE = (
+    "Pulse",
+    "BlackmanWaveform",
+    "KaiserWaveform",
+    "Register",
+    "Register3D",
+)
 
 SUPPORTED_MODULES = {
     "builtins": SUPPORTED_BUILTINS,
diff --git a/pulser-core/pulser/register/base_register.py b/pulser-core/pulser/register/base_register.py
index 1276189..1506cff 100644
--- a/pulser-core/pulser/register/base_register.py
+++ b/pulser-core/pulser/register/base_register.py
@@ -67,6 +67,9 @@ class BaseRegister(ABC):
         self._coords = [np.array(v, dtype=float) for v in qubits.values()]
         self._dim = self._coords[0].size
         self._layout_info: Optional[_LayoutInfo] = None
+        self._init_kwargs(**kwargs)
+
+    def _init_kwargs(self, **kwargs: Any) -> None:
         if kwargs:
             if kwargs.keys() != {"layout", "trap_ids"}:
                 raise ValueError(
@@ -133,6 +136,7 @@ class BaseRegister(ABC):
         center: bool = True,
         prefix: Optional[str] = None,
         labels: Optional[abcSequence[QubitId]] = None,
+        **kwargs: Any,
     ) -> T:
         """Creates the register from an array of coordinates.
 
@@ -172,7 +176,7 @@ class BaseRegister(ABC):
             qubits = dict(zip(cast(Iterable, labels), coords))
         else:
             qubits = dict(cast(Iterable, enumerate(coords)))
-        return cls(qubits)
+        return cls(qubits, **kwargs)
 
     def _validate_layout(
         self, register_layout: RegisterLayout, trap_ids: tuple[int, ...]
@@ -202,16 +206,41 @@ class BaseRegister(ABC):
 
     @abstractmethod
     def _to_dict(self) -> dict[str, Any]:
-        qs = dict(zip(self._ids, map(np.ndarray.tolist, self._coords)))
-        if self._layout_info is not None:
-            return obj_to_dict(self, qs, **(self._layout_info._asdict()))
-        return obj_to_dict(self, qs)
+        """Serializes the object.
+
+        During deserialization, it will be reconstructed using
+        'from_coordinates', so that it uses lists instead of a dictionary
+        (in JSON, lists elements keep their types, but dictionaries keys do
+        not).
+        """
+        cls_dict = obj_to_dict(
+            None,
+            _build=False,
+            _name=self.__class__.__name__,
+            _module=self.__class__.__module__,
+        )
+
+        kwargs = (
+            {} if self._layout_info is None else self._layout_info._asdict()
+        )
+
+        return obj_to_dict(
+            self,
+            cls_dict,
+            [np.ndarray.tolist(qubit_coords) for qubit_coords in self._coords],
+            False,
+            None,
+            self._ids,
+            **kwargs,
+            _submodule=self.__class__.__name__,
+            _name="from_coordinates",
+        )
 
     def __eq__(self, other: Any) -> bool:
         if type(other) is not type(self):
             return False
 
-        return set(self._ids) == set(other._ids) and all(
+        return list(self._ids) == list(other._ids) and all(
             (
                 np.allclose(  # Accounts for rounding errors
                     self._coords[i],


Test patch:
diff --git a/tests/test_json.py b/tests/test_json.py
index c46b219..b1c2405 100644
--- a/tests/test_json.py
+++ b/tests/test_json.py
@@ -90,6 +90,20 @@ def test_register_from_layout():
     assert new_reg._layout_info.trap_ids == (1, 0)
 
 
+@pytest.mark.parametrize(
+    "reg",
+    [
+        Register(dict(enumerate([(2, 3), (5, 1), (10, 0)]))),
+        Register3D({3: (2, 3, 4), 4: (3, 4, 5), 2: (4, 5, 7)}),
+    ],
+)
+def test_register_numbered_keys(reg):
+    j = json.dumps(reg, cls=PulserEncoder)
+    decoded_reg = json.loads(j, cls=PulserDecoder)
+    assert reg == decoded_reg
+    assert all([type(i) == int for i in decoded_reg.qubit_ids])
+
+
 def test_mappable_register():
     layout = RegisterLayout([[0, 0], [1, 1], [1, 0], [0, 1]])
     mapp_reg = layout.make_mappable_register(2)
diff --git a/tests/test_register.py b/tests/test_register.py
index 251eac0..67413f3 100644
--- a/tests/test_register.py
+++ b/tests/test_register.py
@@ -405,7 +405,8 @@ def assert_ineq(left, right):
 def test_equality_function():
     reg1 = Register({"c": (1, 2), "d": (8, 4)})
     assert_eq(reg1, reg1)
-    assert_eq(reg1, Register({"d": (8, 4), "c": (1, 2)}))
+    assert_eq(reg1, Register({"c": (1, 2), "d": (8, 4)}))
+    assert_ineq(reg1, Register({"d": (8, 4), "c": (1, 2)}))
     assert_ineq(reg1, Register({"c": (8, 4), "d": (1, 2)}))
     assert_ineq(reg1, Register({"c": (1, 2), "d": (8, 4), "e": (8, 4)}))
     assert_ineq(reg1, 10)
@@ -413,6 +414,8 @@ def test_equality_function():
     reg2 = Register3D({"a": (1, 2, 3), "b": (8, 5, 6)})
     assert_eq(reg2, reg2)
     assert_eq(reg2, Register3D({"a": (1, 2, 3), "b": (8, 5, 6)}))
+    assert_eq(reg2, Register3D({"a": (1, 2, 3), "b": (8, 5, 6)}))
+    assert_ineq(reg2, Register3D({"b": (8, 5, 6), "a": (1, 2, 3)}))
     assert_ineq(reg2, Register3D({"b": (1, 2, 3), "a": (8, 5, 6)}))
     assert_ineq(
         reg2, Register3D({"a": (1, 2, 3), "b": (8, 5, 6), "e": (8, 5, 6)})