Rectangular lattice register and layout (#665)

* Rectangular lattice and random layouts

* Precision on the lattice of Register.rectangle

* First batch of corrections

* Add RectangularLatticeLayout to supported.py

* Second batch of corrections, removed RandomLayout and TriangRectShape

* add json test for rectangular lattice layout

* Add test register rectangular lattice

pulser-core/pulser/json/supported.py

@@ -66,6 +66,7 @@
     "pulser.register.register3d": ("Register3D",),
     "pulser.register.register_layout": ("RegisterLayout",),
     "pulser.register.special_layouts": (
+        "RectangularLatticeLayout",
         "SquareLatticeLayout",
         "TriangularLatticeLayout",
     ),

pulser-core/pulser/register/__init__.py

@@ -20,6 +20,7 @@
 from pulser.register.special_layouts import (
     SquareLatticeLayout,
     TriangularLatticeLayout,
+    RectangularLatticeLayout,
 )
 
 __all__ = [
@@ -29,4 +30,5 @@
     "RegisterLayout",
     "SquareLatticeLayout",
     "TriangularLatticeLayout",
+    "RectangularLatticeLayout",
 ]

pulser-core/pulser/register/register.py

@@ -89,7 +89,7 @@ def rectangle(
         spacing: float = 4.0,
         prefix: Optional[str] = None,
     ) -> Register:
-        """Initializes the register with the qubits in a rectangular array.
+        """Creates a rectangular array of qubits on a square lattice.
 
         Args:
             rows: Number of rows.
@@ -102,6 +102,32 @@ def rectangle(
         Returns:
             A register with qubits placed in a rectangular array.
         """
+        return cls.rectangular_lattice(rows, columns, spacing, spacing, prefix)
+
+    @classmethod
+    def rectangular_lattice(
+        cls,
+        rows: int,
+        columns: int,
+        row_spacing: float = 4.0,
+        col_spacing: float = 2.0,
+        prefix: Optional[str] = None,
+    ) -> Register:
+        """Creates a rectangular array of qubits on a rectangular lattice.
+
+        Args:
+            rows: Number of rows.
+            columns: Number of columns.
+            row_spacing: The distance between rows in μm.
+            col_spacing: The distance between columns in μm.
+            prefix: The prefix for the qubit ids. If defined, each qubit
+                id starts with the prefix, followed by an int from 0 to N-1
+                (e.g. prefix='q' -> IDs: 'q0', 'q1', 'q2', ...)
+
+        Returns:
+            Register with qubits placed in a rectangular array on a
+            rectangular lattice.
+        """
         # Check rows
         if rows < 1:
             raise ValueError(
@@ -117,13 +143,12 @@ def rectangle(
             )
 
         # Check spacing
-        if spacing <= 0.0:
-            raise ValueError(
-                f"Spacing between atoms (`spacing` = {spacing})"
-                " must be greater than 0."
-            )
+        if row_spacing <= 0.0 or col_spacing <= 0.0:
+            raise ValueError("Spacing between atoms must be greater than 0.")
 
-        coords = patterns.square_rect(rows, columns) * spacing
+        coords = patterns.square_rect(rows, columns)
+        coords[:, 0] = coords[:, 0] * col_spacing
+        coords[:, 1] = coords[:, 1] * row_spacing
 
         return cls.from_coordinates(coords, center=True, prefix=prefix)
 

pulser-core/pulser/register/special_layouts.py

@@ -26,26 +26,33 @@
     from pulser.register import Register
 
 
-class SquareLatticeLayout(RegisterLayout):
-    """A RegisterLayout with a square lattice pattern in a rectangular shape.
+class RectangularLatticeLayout(RegisterLayout):
+    """RegisterLayout with rectangular lattice pattern in a rectangular shape.
 
     Args:
         rows: The number of rows of traps.
         columns: The number of columns of traps.
-        spacing: The distance between neighbouring traps (in µm).
+        col_spacing: Horizontal distance between neighbouring traps (in µm).
+        row_spacing: Vertical distance between neighbouring traps (in µm)
     """
 
-    def __init__(self, rows: int, columns: int, spacing: float):
-        """Initializes a SquareLatticeLayout."""
+    def __init__(
+        self, rows: int, columns: int, col_spacing: float, row_spacing: float
+    ):
+        """Initializes a RectangularLatticeLayout."""
         self._rows = int(rows)
         self._columns = int(columns)
-        self._spacing = float(spacing)
+        self._col_spacing = float(col_spacing)
+        self._row_spacing = float(row_spacing)
         slug = (
-            f"SquareLatticeLayout({self._rows}x{self._columns}, "
-            f"{self._spacing}µm)"
+            f"RectangularLatticeLayout({self._rows}x{self._columns}, "
+            f"{self._col_spacing}x{self._row_spacing}µm)"
         )
+        self._traps = patterns.square_rect(self._rows, self._columns)
+        self._traps[:, 0] = self._traps[:, 0] * self._col_spacing
+        self._traps[:, 1] = self._traps[:, 1] * self._row_spacing
         super().__init__(
-            patterns.square_rect(self._rows, self._columns) * self._spacing,
+            trap_coordinates=self._traps,
             slug=slug,
         )
 
@@ -84,16 +91,53 @@ def rectangular_register(
         if rows > self._rows or columns > self._columns:
             raise ValueError(
                 f"A '{rows}x{columns}' array doesn't fit a "
-                f"{self._rows}x{self._columns} SquareLatticeLayout."
+                f"{self._rows}x{self._columns} RectangularLatticeLayout."
             )
-        points = patterns.square_rect(rows, columns) * self._spacing
+        points = patterns.square_rect(rows, columns)
+        points[:, 0] = points[:, 0] * self._col_spacing
+        points[:, 1] = points[:, 1] * self._row_spacing
         trap_ids = self.get_traps_from_coordinates(*points)
         qubit_ids = [f"{prefix}{i}" for i in range(len(trap_ids))]
         return cast(
             pulser.Register,
             self.define_register(*trap_ids, qubit_ids=qubit_ids),
         )
 
+    def _to_dict(self) -> dict[str, Any]:
+        return obj_to_dict(
+            self,
+            self._rows,
+            self._columns,
+            self._col_spacing,
+            self._row_spacing,
+        )
+
+
+class SquareLatticeLayout(RectangularLatticeLayout):
+    """A RegisterLayout with a square lattice pattern in a rectangular shape.
+
+    Args:
+        rows: The number of rows of traps.
+        columns: The number of columns of traps.
+        spacing: The distance between neighbouring traps (in µm).
+    """
+
+    def __init__(self, rows: int, columns: int, spacing: float):
+        """Initializes a SquareLatticeLayout."""
+        self._rows = int(rows)
+        self._columns = int(columns)
+        self._spacing = float(spacing)
+        self._col_spacing = self._spacing
+        self._row_spacing = self._spacing
+        super().__init__(
+            self._rows, self._columns, self._spacing, self._spacing
+        )
+        slug = (
+            f"SquareLatticeLayout({self._rows}x{self._columns}, "
+            f"{self._spacing}µm)"
+        )
+        object.__setattr__(self, "slug", slug)
+
     def _to_dict(self) -> dict[str, Any]:
         return obj_to_dict(self, self._rows, self._columns, self._spacing)
 

tests/test_json.py

@@ -26,6 +26,7 @@
 from pulser.parametrized.decorators import parametrize
 from pulser.register.register_layout import RegisterLayout
 from pulser.register.special_layouts import (
+    RectangularLatticeLayout,
     SquareLatticeLayout,
     TriangularLatticeLayout,
 )
@@ -92,6 +93,11 @@ def test_layout():
     assert new_square_layout == square_layout
     assert type(new_square_layout) is SquareLatticeLayout
 
+    rectangular_layout = RectangularLatticeLayout(8, 10, 6, 5)
+    new_rectangular_layout = encode_decode(rectangular_layout)
+    assert new_rectangular_layout == rectangular_layout
+    assert type(new_rectangular_layout) is RectangularLatticeLayout
+
 
 def test_register_from_layout():
     layout = RegisterLayout([[0, 0], [1, 1], [1, 0], [0, 1]])

tests/test_register.py

@@ -84,6 +84,24 @@ def test_creation():
         Register(qubits, spacing=10, layout="square", trap_ids=(0, 1, 3))
 
 
+def test_rectangular_lattice():
+    # Check rows
+    with pytest.raises(ValueError, match="The number of rows"):
+        Register.rectangular_lattice(0, 2, 3, 4)
+
+    # Check columns
+    with pytest.raises(ValueError, match="The number of columns"):
+        Register.rectangular_lattice(2, 0, 3, 4)
+
+    # Check row spacing
+    with pytest.raises(ValueError, match="Spacing"):
+        Register.rectangular_lattice(2, 2, 0.0, 5)
+
+    # Check col spacing
+    with pytest.raises(ValueError, match="Spacing"):
+        Register.rectangular_lattice(2, 2, 3, 0.0)
+
+
 def test_rectangle():
     # Check rows
     with pytest.raises(ValueError, match="The number of rows"):

tests/test_register_layout.py

@@ -22,6 +22,7 @@
 from pulser.register import Register, Register3D
 from pulser.register.register_layout import RegisterLayout
 from pulser.register.special_layouts import (
+    RectangularLatticeLayout,
     SquareLatticeLayout,
     TriangularLatticeLayout,
 )
@@ -190,6 +191,22 @@ def test_square_lattice_layout():
         square.rectangular_register(10, 3)
 
 
+def test_rectangular_lattice_layout():
+    rectangle = RectangularLatticeLayout(9, 7, 2, 4)
+    assert str(rectangle) == "RectangularLatticeLayout(9x7, 2.0x4.0µm)"
+    assert rectangle.square_register(3) == Register.rectangular_lattice(
+        3, 3, col_spacing=2, row_spacing=4, prefix="q"
+    )
+    # An even number of atoms on the side won't align the center with an atom
+    assert rectangle.square_register(4) != Register.rectangular_lattice(
+        4, 4, col_spacing=2, row_spacing=4, prefix="q"
+    )
+    with pytest.raises(ValueError, match="'8x8' array doesn't fit"):
+        rectangle.square_register(8)
+    with pytest.raises(ValueError, match="'10x3' array doesn't fit"):
+        rectangle.rectangular_register(10, 3)
+
+
 def test_triangular_lattice_layout():
     tri = TriangularLatticeLayout(50, 5)
     assert str(tri) == "TriangularLatticeLayout(50, 5.0µm)"