Incorporate automatic layout generation in Register (#753)

* Adding the register layout generator

* Define Register.with_automatic_layout()

* Function scope modifications

* Unit tests

* Update tutorial

* Address review comments

* Account for when the Register requires more traps than max_traps

pulser-core/pulser/register/_layout_gen.py

@@ -0,0 +1,104 @@
+# Copyright 2024 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.
+from __future__ import annotations
+
+import numpy as np
+from scipy.spatial.distance import cdist
+
+
+def generate_trap_coordinates(
+    atom_coords: np.ndarray,
+    min_trap_dist: float,
+    max_radial_dist: int,
+    max_layout_filling: float,
+    optimal_layout_filling: float | None = None,
+    mesh_resolution: float = 1.0,
+    min_traps: int = 1,
+    max_traps: int | None = None,
+) -> list[np.ndarray]:
+    """Generates trap coordinates for a collection of atom coordinates.
+
+    Generates a mesh of resolution `mesh_resolution` covering a disk of radius
+    `max_radial_dist`. Deletes all the points of the mesh that are below a
+    radius `min_trap_dist` of any atoms or traps and iteratively selects from
+    the remaining points the necessary number of traps such that the ratio
+    number of atoms to number of traps is at most max_layout_filling and as
+    close as possible to optimal_layout_filling, while being above min_traps
+    and below max_traps.
+
+    Args:
+        atom_coords: The coordinates where atoms will be placed.
+        min_trap_dist: The minimum distance between traps, in µm.
+        max_radial_dist: The maximum distance from the origin, in µm.
+        max_layout_filling: The maximum ratio of atoms to traps.
+        optimal_layout_filling: An optional value for the optimal ratio of
+            atoms to traps. If not given, takes max_layout_filling.
+        mesh_resolution: The spacing between points in the mesh of candidate
+            coordinates, in µm.
+        min_traps: The minimum number of traps in the resulting layout.
+        max_traps: The maximum number of traps in the resulting layout.
+    """
+    optimal_layout_filling = optimal_layout_filling or max_layout_filling
+    assert optimal_layout_filling <= max_layout_filling
+    assert max_traps is None or min_traps <= max_traps
+
+    # Generate all coordinates where a trap can be placed
+    lx = 2 * max_radial_dist
+    side = np.linspace(0, lx, num=int(lx / mesh_resolution)) - max_radial_dist
+    x, y = np.meshgrid(side, side)
+    in_circle = x**2 + y**2 <= max_radial_dist**2
+    coords = np.c_[x[in_circle].ravel(), y[in_circle].ravel()]
+
+    # Get the atoms in the register (the "seeds")
+    seeds: list[np.ndarray] = list(atom_coords)
+    n_seeds = len(seeds)
+
+    # Record indices and distances between coords and seeds
+    c_indx = np.arange(len(coords))
+    all_dists = cdist(coords, seeds)
+
+    # Accounts for the case when the needed number is less than min_traps
+    min_traps = max(
+        np.ceil(n_seeds / max_layout_filling).astype(int), min_traps
+    )
+
+    # Use max() in case min_traps is larger than the optimal number
+    target_traps = max(
+        np.round(n_seeds / optimal_layout_filling).astype(int),
+        min_traps,
+    )
+    if max_traps:
+        target_traps = min(target_traps, max_traps)
+
+    # This is the region where we can still add traps
+    region_left = np.all(all_dists > min_trap_dist, axis=1)
+    # The traps start out as being just the seeds
+    traps = seeds.copy()
+    for _ in range(target_traps - n_seeds):
+        if not np.any(region_left):
+            break
+        # Select the point in the valid region that is closest to a seed
+        selected = c_indx[region_left][
+            np.argmin(np.min(all_dists[region_left][:, :n_seeds], axis=1))
+        ]
+        # Add the selected point to the traps
+        traps.append(coords[selected])
+        # Add the distances to the new trap
+        all_dists = np.append(all_dists, cdist(coords, [traps[-1]]), axis=1)
+        region_left *= all_dists[:, -1] > min_trap_dist
+    if len(traps) < min_traps:
+        raise RuntimeError(
+            f"Failed to find a site for {min_traps - len(traps)} traps."
+        )
+    return traps

pulser-core/pulser/register/register.py

@@ -17,7 +17,7 @@
 
 import warnings
 from collections.abc import Mapping
-from typing import Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -31,9 +31,13 @@
     deserialize_abstract_register,
 )
 from pulser.json.utils import stringify_qubit_ids
+from pulser.register._layout_gen import generate_trap_coordinates
 from pulser.register._reg_drawer import RegDrawer
 from pulser.register.base_register import BaseRegister, QubitId
 
+if TYPE_CHECKING:
+    from pulser.devices import Device
+
 
 class Register(BaseRegister, RegDrawer):
     """A 2D quantum register containing a set of qubits.
@@ -324,6 +328,44 @@ def max_connectivity(
 
         return cls.from_coordinates(coords, center=False, prefix=prefix)
 
+    def with_automatic_layout(
+        self,
+        device: Device,
+        layout_slug: str | None = None,
+    ) -> Register:
+        """Replicates the register with an automatically generated layout.
+
+        The generated `RegisterLayout` can be accessed via `Register.layout`.
+
+        Args:
+            device: The device constraints for the layout generation.
+            layout_slug: An optional slug for the generated layout.
+
+        Raises:
+            RuntimeError: If the automatic layout generation fails to meet
+                the device constraints.
+
+        Returns:
+            Register: A new register instance with identical qubit IDs and
+            coordinates but also the newly generated RegisterLayout.
+        """
+        if not isinstance(device, pulser.devices.Device):
+            raise TypeError(
+                f"'device' must be of type Device, not {type(device)}."
+            )
+        trap_coords = generate_trap_coordinates(
+            self.sorted_coords,
+            min_trap_dist=device.min_atom_distance,
+            max_radial_dist=device.max_radial_distance,
+            max_layout_filling=device.max_layout_filling,
+            optimal_layout_filling=device.optimal_layout_filling,
+            min_traps=device.min_layout_traps,
+            max_traps=device.max_layout_traps,
+        )
+        layout = pulser.register.RegisterLayout(trap_coords, slug=layout_slug)
+        trap_ids = layout.get_traps_from_coordinates(*self.sorted_coords)
+        return cast(Register, layout.define_register(*trap_ids))
+
     def rotated(self, degrees: float) -> Register:
         """Makes a new rotated register.
 

tests/test_register.py

@@ -13,13 +13,15 @@
 # limitations under the License.
 from __future__ import annotations
 
+import dataclasses
 from unittest.mock import patch
 
 import numpy as np
 import pytest
 
 from pulser import Register, Register3D
-from pulser.devices import DigitalAnalogDevice, MockDevice
+from pulser.devices import AnalogDevice, DigitalAnalogDevice, MockDevice
+from pulser.register import RegisterLayout
 
 
 def test_creation():
@@ -587,3 +589,71 @@ def test_register_recipes_torch(
     }
     reg = reg_classmethod(**kwargs)
     _assert_reg_requires_grad(reg, invert=not requires_grad)
+
+
+@pytest.mark.parametrize("optimal_filling", [None, 0.4, 0.1])
+def test_automatic_layout(optimal_filling):
+    reg = Register.square(4, spacing=5)
+    max_layout_filling = 0.5
+    min_traps = int(np.ceil(len(reg.qubits) / max_layout_filling))
+    optimal_traps = int(
+        np.ceil(len(reg.qubits) / (optimal_filling or max_layout_filling))
+    )
+    device = dataclasses.replace(
+        AnalogDevice,
+        max_atom_num=20,
+        max_layout_filling=max_layout_filling,
+        optimal_layout_filling=optimal_filling,
+        pre_calibrated_layouts=(),
+    )
+    device.validate_register(reg)
+
+    # On its own, it works
+    new_reg = reg.with_automatic_layout(device, layout_slug="foo")
+    assert isinstance(new_reg.layout, RegisterLayout)
+    assert str(new_reg.layout) == "foo"
+    trap_num = new_reg.layout.number_of_traps
+    assert min_traps <= trap_num <= optimal_traps
+    # To test the device limits on trap number are enforced
+    if not optimal_filling:
+        assert trap_num == min_traps
+        bound_below_dev = dataclasses.replace(
+            device, min_layout_traps=trap_num + 1
+        )
+        assert (
+            reg.with_automatic_layout(bound_below_dev).layout.number_of_traps
+            == bound_below_dev.min_layout_traps
+        )
+    elif trap_num < optimal_traps:
+        assert trap_num > min_traps
+        bound_above_dev = dataclasses.replace(
+            device, max_layout_traps=trap_num - 1
+        )
+        assert (
+            reg.with_automatic_layout(bound_above_dev).layout.number_of_traps
+            == bound_above_dev.max_layout_traps
+        )
+
+    with pytest.raises(TypeError, match="must be of type Device"):
+        reg.with_automatic_layout(MockDevice)
+
+    # Minimum number of traps is too high
+    with pytest.raises(RuntimeError, match="Failed to find a site"):
+        reg.with_automatic_layout(
+            dataclasses.replace(device, min_layout_traps=200)
+        )
+
+    # The Register is larger than max_traps
+    big_reg = Register.square(8, spacing=5)
+    min_traps = np.ceil(len(big_reg.qubit_ids) / max_layout_filling)
+    with pytest.raises(
+        RuntimeError, match="Failed to find a site for 2 traps"
+    ):
+        big_reg.with_automatic_layout(
+            dataclasses.replace(device, max_layout_traps=int(min_traps - 2))
+        )
+    # Without max_traps, it would still work
+    assert (
+        big_reg.with_automatic_layout(device).layout.number_of_traps
+        >= min_traps
+    )

tutorials/advanced_features/Backends for Sequence Execution.ipynb

@@ -62,7 +62,11 @@
     "Sequence execution on a QPU is done through the `QPUBackend`, which is a remote backend. Therefore, it requires a remote backend connection, which should be open from the start due to two additional QPU constraints:\n",
     "\n",
     "1. The `Device` must be chosen among the options available at the moment, which can be found through `connection.fetch_available_devices()`.\n",
-    "2. The `Register` must be defined from one of the register layouts calibrated for the chosen `Device`, which are found under `Device.calibrated_register_layouts`. Check out [this tutorial](reg_layouts.nblink) for more information on how to define a `Register` from a `RegisterLayout`.\n",
+    "2. If in the chosen device `Device.requires_layout` is `True`, the `Register` must be defined from a register layout:  \n",
+    "    -  If `Device.accepts_new_layouts` is `False`, use one of the register layouts calibrated for the chosen `Device` (found under `Device.calibrated_register_layouts`). Check out [this tutorial](reg_layouts.nblink) for more information on how to define a `Register` from a `RegisterLayout`.\n",
+    "    - Otherwise, we may choose to define our own custom layout or rely on `Register.with_automatic_layout()` to\n",
+    "    give us a register from an automatically generated register layout that fits our desired register while obeying the device constraints. \n",
+    "\n",
     "\n",
     "On the contrary, execution on emulator backends imposes no further restriction on the device and the register. We will stick to emulator backends in this tutorial, so we will forego the requirements of QPU backends in the following steps."
    ]