Add layout-restricting parameters to BaseDevice

pulser-core/pulser/devices/_device_datacls.py

@@ -40,14 +40,20 @@
 
 DIMENSIONS = Literal[2, 3]
 
-ALWAYS_OPTIONAL_PARAMS = ("max_sequence_duration", "max_runs")
+ALWAYS_OPTIONAL_PARAMS = (
+    "max_sequence_duration",
+    "max_runs",
+    "optimal_layout_filling",
+    "max_layout_traps",
+)
 OPTIONAL_IN_ABSTR_REPR = tuple(
     list(ALWAYS_OPTIONAL_PARAMS)
     + [
         "dmm_objects",
         "default_noise_model",
         "requires_layout",
         "accepts_new_layouts",
+        "min_layout_traps",
     ]
 )
 PARAMS_WITH_ABSTR_REPR = ("channel_objects", "channel_ids", "dmm_objects")
@@ -83,6 +89,11 @@ class BaseDevice(ABC):
         supports_slm_mask: Whether the device supports the SLM mask feature.
         max_layout_filling: The largest fraction of a layout that can be filled
             with atoms.
+        optimal_layout_filling: An optional value for the fraction of a layout
+            that should be filled with atoms.
+        min_layout_traps: The minimum number of traps a layout can have.
+        max_layout_traps: An optional value for the maximum number of traps a
+            layout can have.
         max_sequence_duration: The maximum allowed duration for a sequence
             (in ns).
         max_runs: The maximum number of runs allowed on the device. Only used
@@ -103,6 +114,9 @@ class BaseDevice(ABC):
     interaction_coeff_xy: float | None = None
     supports_slm_mask: bool = False
     max_layout_filling: float = 0.5
+    optimal_layout_filling: float | None = None
+    min_layout_traps: int = 1
+    max_layout_traps: int | None = None
     max_sequence_duration: int | None = None
     max_runs: int | None = None
     requires_layout: bool = False
@@ -141,6 +155,8 @@ def type_check(
             "max_radial_distance",
             "max_sequence_duration",
             "max_runs",
+            "min_layout_traps",
+            "max_layout_traps",
         ):
             value = getattr(self, param)
             if (
@@ -180,6 +196,40 @@ def type_check(
                 f"not {self.max_layout_filling}."
             )
 
+        if self.optimal_layout_filling is not None and not (
+            0.0 < self.optimal_layout_filling <= self.max_layout_filling
+        ):
+            raise ValueError(
+                "When defined, the optimal layout filling fraction "
+                "must be greater than 0. and less than or equal to "
+                f"`max_layout_filling` ({self.max_layout_filling}), "
+                f"not {self.optimal_layout_filling}."
+            )
+
+        if self.max_layout_traps is not None:
+            if self.max_layout_traps < self.min_layout_traps:
+                raise ValueError(
+                    "The maximum number of layout traps "
+                    f"({self.max_layout_traps}) must be greater than "
+                    "or equal to the minimum number of layout traps "
+                    f"({self.min_layout_traps})."
+                )
+            if (
+                self.max_atom_num is not None
+                and (
+                    max_atoms_ := int(
+                        self.max_layout_filling * self.max_layout_traps
+                    )
+                )
+                < self.max_atom_num
+            ):
+                raise ValueError(
+                    "With the given maximum layout filling and maximum number "
+                    f"of traps, a layout supports at most {max_atoms_} atoms, "
+                    "which is less than the maximum number of atoms allowed"
+                    f"({self.max_atom_num})."
+                )
+
         for ch_obj in self.channel_objects:
             type_check("All channels", Channel, value_override=ch_obj)
 
@@ -360,6 +410,23 @@ def validate_layout(self, layout: RegisterLayout) -> None:
                 f"{self.dimensions} dimensions."
             )
 
+        if layout.number_of_traps < self.min_layout_traps:
+            raise ValueError(
+                "The device requires register layouts to have "
+                f"at least {self.min_layout_traps} traps; "
+                f"{layout!s} has only {layout.number_of_traps}."
+            )
+
+        if (
+            self.max_layout_traps is not None
+            and layout.number_of_traps > self.max_layout_traps
+        ):
+            raise ValueError(
+                "The device requires register layouts to have "
+                f"at most {self.max_layout_traps} traps; "
+                f"{layout!s} has {layout.number_of_traps}."
+            )
+
         self._validate_coords(layout.traps_dict, kind="traps")
 
     def validate_layout_filling(
@@ -547,6 +614,11 @@ class Device(BaseDevice):
         supports_slm_mask: Whether the device supports the SLM mask feature.
         max_layout_filling: The largest fraction of a layout that can be filled
             with atoms.
+        optimal_layout_filling: An optional value for the fraction of a layout
+            that should be filled with atoms.
+        min_layout_traps: The minimum number of traps a layout can have.
+        max_layout_traps: An optional value for the maximum number of traps a
+            layout can have.
         max_sequence_duration: The maximum allowed duration for a sequence
             (in ns).
         max_runs: The maximum number of runs allowed on the device. Only used
@@ -792,6 +864,11 @@ class VirtualDevice(BaseDevice):
         supports_slm_mask: Whether the device supports the SLM mask feature.
         max_layout_filling: The largest fraction of a layout that can be filled
             with atoms.
+        optimal_layout_filling: An optional value for the fraction of a layout
+            that should be filled with atoms.
+        min_layout_traps: The minimum number of traps a layout can have.
+        max_layout_traps: An optional value for the maximum number of traps a
+            layout can have.
         max_sequence_duration: The maximum allowed duration for a sequence
             (in ns).
         max_runs: The maximum number of runs allowed on the device. Only used

pulser-core/pulser/json/abstract_repr/schemas/device-schema.json

@@ -166,6 +166,10 @@
               "description": "The largest fraction of a layout that can be filled with atoms.",
               "type": "number"
             },
+            "max_layout_traps": {
+              "description": "The maximum number of traps a layout can have.",
+              "type": "number"
+            },
             "max_radial_distance": {
               "description": "Maximum distance an atom can be from the center of the array (in µm).",
               "type": "number"
@@ -182,10 +186,18 @@
               "description": "The closest together two atoms can be (in μm).",
               "type": "number"
             },
+            "min_layout_traps": {
+              "description": "The minimum number of traps a layout can have.",
+              "type": "number"
+            },
             "name": {
               "description": "A unique name for the device.",
               "type": "string"
             },
+            "optimal_layout_filling": {
+              "description": "The optimal fraction of a layout that should be filled with atoms.",
+              "type": "number"
+            },
             "pre_calibrated_layouts": {
               "description": "Register layouts already calibrated on the device.",
               "items": {
@@ -288,6 +300,10 @@
               "description": "The largest fraction of a layout that can be filled with atoms.",
               "type": "number"
             },
+            "max_layout_traps": {
+              "description": "The maximum number of traps a layout can have.",
+              "type": "number"
+            },
             "max_radial_distance": {
               "description": "Maximum distance an atom can be from the center of the array (in µm).",
               "type": [
@@ -307,10 +323,18 @@
               "description": "The closest together two atoms can be (in μm).",
               "type": "number"
             },
+            "min_layout_traps": {
+              "description": "The minimum number of traps a layout can have.",
+              "type": "number"
+            },
             "name": {
               "description": "A unique name for the device.",
               "type": "string"
             },
+            "optimal_layout_filling": {
+              "description": "The optimal fraction of a layout that should be filled with atoms.",
+              "type": "number"
+            },
             "requires_layout": {
               "description": "Whether the register used in the sequence must be created from a register layout.  Only enforced in QPU execution.",
               "type": "boolean"

tests/test_abstract_repr.py

@@ -479,6 +479,9 @@ def check_error_raised(
         [
             (MockDevice, "max_sequence_duration", 1000),
             (MockDevice, "max_runs", 100),
+            (MockDevice, "optimal_layout_filling", 0.4),
+            (MockDevice, "min_layout_traps", 10),
+            (MockDevice, "max_layout_traps", 200),
             (MockDevice, "requires_layout", True),
             (AnalogDevice, "requires_layout", False),
             (AnalogDevice, "accepts_new_layouts", False),

tests/test_devices.py

@@ -47,6 +47,8 @@ def test_params():
         min_atom_distance=1,
         max_atom_num=None,
         max_radial_distance=None,
+        min_layout_traps=10,
+        max_layout_traps=100,
     )
 
 
@@ -122,6 +124,36 @@ def test_post_init_type_checks(test_params, param, value, msg):
             "maximum layout filling fraction must be greater than 0. and"
             " less than or equal to 1.",
         ),
+        (
+            "optimal_layout_filling",
+            0.0,
+            "When defined, the optimal layout filling fraction must be greater"
+            " than 0. and less than or equal to `max_layout_filling`",
+        ),
+        (
+            "optimal_layout_filling",
+            0.9,
+            "When defined, the optimal layout filling fraction must be greater"
+            " than 0. and less than or equal to `max_layout_filling`",
+        ),
+        (
+            "min_layout_traps",
+            0,
+            "'min_layout_traps' must be greater than zero",
+        ),
+        ("max_layout_traps", 0, None),
+        (
+            "max_atom_num",
+            100,
+            "With the given maximum layout filling and maximum number "
+            "of traps, a layout supports at most 50 atoms",
+        ),
+        (
+            "max_layout_traps",
+            9,
+            "must be greater than or equal to the minimum number of "
+            "layout traps",
+        ),
         (
             "channel_ids",
             ("rydberg_global", "rydberg_global"),
@@ -336,6 +368,22 @@ def test_validate_layout():
             )
         )
 
+    restricted_device = replace(
+        DigitalAnalogDevice, min_layout_traps=10, max_layout_traps=200
+    )
+    with pytest.raises(
+        ValueError,
+        match="The device requires register layouts to have "
+        "at least 10 traps",
+    ):
+        restricted_device.validate_layout(TriangularLatticeLayout(9, 10))
+    with pytest.raises(
+        ValueError,
+        match="The device requires register layouts to have "
+        "at most 200 traps",
+    ):
+        restricted_device.validate_layout(TriangularLatticeLayout(201, 10))
+
     valid_layout = RegisterLayout(
         Register.square(
             int(np.sqrt(DigitalAnalogDevice.max_atom_num * 2))