Support for 3D registers and layouts in the abstract representation

pulser-core/pulser/json/abstract_repr/deserializer.py

@@ -16,7 +16,7 @@
 
 import dataclasses
 import json
-from typing import TYPE_CHECKING, Any, Type, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Literal, Type, Union, cast, overload
 
 import jsonschema
 import jsonschema.exceptions
@@ -57,7 +57,7 @@
 
 if TYPE_CHECKING:
     from pulser.noise_model import NoiseModel
-    from pulser.register.base_register import BaseRegister
+    from pulser.register import Register, Register3D
     from pulser.sequence import Sequence
 
 
@@ -384,15 +384,28 @@ def _deserialize_layout(layout_obj: dict[str, Any]) -> RegisterLayout:
 
 def _deserialize_register(
     qubits: list[dict[str, Any]], layout: RegisterLayout | None
-) -> BaseRegister:
+) -> Register:
     coords = [(q["x"], q["y"]) for q in qubits]
     qubit_ids = [q["name"] for q in qubits]
     if layout:
         trap_ids = layout.get_traps_from_coordinates(*coords)
         reg = layout.define_register(*trap_ids, qubit_ids=qubit_ids)
     else:
         reg = pulser.Register(dict(zip(qubit_ids, coords)))
-    return reg
+    return cast(pulser.Register, reg)
+
+
+def _deserialize_register3d(
+    qubits: list[dict[str, Any]], layout: RegisterLayout | None
+) -> Register3D:
+    coords = [(q["x"], q["y"], q["z"]) for q in qubits]
+    qubit_ids = [q["name"] for q in qubits]
+    if layout:
+        trap_ids = layout.get_traps_from_coordinates(*coords)
+        reg = layout.define_register(*trap_ids, qubit_ids=qubit_ids)
+    else:
+        reg = pulser.Register3D(dict(zip(qubit_ids, coords)))
+    return cast(pulser.Register3D, reg)
 
 
 def _deserialize_noise_model(noise_model_obj: dict[str, Any]) -> NoiseModel:
@@ -495,11 +508,14 @@ def deserialize_abstract_sequence(obj_str: str) -> Sequence:
     layout = _deserialize_layout(obj["layout"]) if "layout" in obj else None
 
     # Register
-    reg: Union[BaseRegister, MappableRegister]
+    reg: Register | Register3D | MappableRegister
     qubits = obj["register"]
     if {"name", "x", "y"} == qubits[0].keys():
-        # Regular register
+        # Regular 2D register
         reg = _deserialize_register(qubits, layout)
+    elif {"name", "x", "y", "z"} == qubits[0].keys():
+        # Regular 3D register
+        reg = _deserialize_register3d(qubits, layout)
     else:
         # Mappable register
         assert (
@@ -589,20 +605,59 @@ def deserialize_abstract_layout(obj_str: str) -> RegisterLayout:
     return _deserialize_layout(json.loads(obj_str))
 
 
-def deserialize_abstract_register(obj_str: str) -> BaseRegister:
+@overload
+def deserialize_abstract_register(
+    obj_str: str, required_dim: Literal[2]
+) -> Register:
+    pass
+
+
+@overload
+def deserialize_abstract_register(
+    obj_str: str, required_dim: Literal[3]
+) -> Register3D:
+    pass
+
+
+@overload
+def deserialize_abstract_register(obj_str: str) -> Register | Register3D:
+    pass
+
+
+def deserialize_abstract_register(
+    obj_str: str, required_dim: Literal[None, 2, 3] = None
+) -> Register | Register3D:
     """Deserialize a register from an abstract JSON object.
 
     Args:
-        obj_str: the JSON string representing the register encoded
+        obj_str: The JSON string representing the register encoded
             in the abstract JSON format.
+        required_dim: If defined, ensures the register is of the
+            specified dimensionality.
 
     Returns:
         The Register instance.
     """
+    if required_dim not in (None, 2, 3):
+        raise ValueError(
+            "When specified, 'required_dim' must be 2 or 3, "
+            f"not {required_dim!s}."
+        )
     validate_abstract_repr(obj_str, "register")
     obj = json.loads(obj_str)
     layout = _deserialize_layout(obj["layout"]) if "layout" in obj else None
-    return _deserialize_register(qubits=obj["register"], layout=layout)
+    qubits = obj["register"]
+    dim_ = len(set(qubits[0]) - {"name"})
+    # These conditions should be enforced by the schema
+    assert dim_ == 2 or dim_ == 3
+    assert layout is None or layout.dimensionality == dim_
+    if required_dim is not None and required_dim != dim_:
+        raise ValueError(
+            f"The provided register must be in {required_dim}D, not {dim_}D."
+        )
+    if dim_ == 3:
+        return _deserialize_register3d(qubits=qubits, layout=layout)
+    return _deserialize_register(qubits=qubits, layout=layout)
 
 
 def deserialize_abstract_noise_model(obj_str: str) -> NoiseModel:

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

@@ -421,7 +421,7 @@
                       "type": "boolean"
                     },
                     "red_shift_coeff": {
-                      "description": "The weight coefficient of the blue beam's contribution to the lightshift.",
+                      "description": "The weight coefficient of the red beam's contribution to the lightshift.",
                       "type": "number"
                     }
                   },
@@ -755,7 +755,7 @@
                       "type": "boolean"
                     },
                     "red_shift_coeff": {
-                      "description": "The weight coefficient of the blue beam's contribution to the lightshift.",
+                      "description": "The weight coefficient of the red beam's contribution to the lightshift.",
                       "type": "number"
                     }
                   },
@@ -1103,7 +1103,7 @@
                       "type": "boolean"
                     },
                     "red_shift_coeff": {
-                      "description": "The weight coefficient of the blue beam's contribution to the lightshift.",
+                      "description": "The weight coefficient of the red beam's contribution to the lightshift.",
                       "type": "number"
                     }
                   },
@@ -1410,7 +1410,7 @@
                       "type": "boolean"
                     },
                     "red_shift_coeff": {
-                      "description": "The weight coefficient of the blue beam's contribution to the lightshift.",
+                      "description": "The weight coefficient of the red beam's contribution to the lightshift.",
                       "type": "number"
                     }
                   },

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

@@ -4,8 +4,18 @@
   "$schema": "http://json-schema.org/draft-07/schema#",
   "definitions": {
     "Layout": {
+      "anyOf": [
+        {
+          "$ref": "#/definitions/Layout2D"
+        },
+        {
+          "$ref": "#/definitions/Layout3D"
+        }
+      ],
+      "description": "Layout with the positions of the traps. A selection of up to 50% of these traps makes up the Register."
+    },
+    "Layout2D": {
       "additionalProperties": false,
-      "description": "Layout with the positions of the traps. A selection of up to 50% of these traps makes up the Register.",
       "properties": {
         "coordinates": {
           "description": "The trap coordinates in µm.",
@@ -28,6 +38,31 @@
         "coordinates"
       ],
       "type": "object"
+    },
+    "Layout3D": {
+      "additionalProperties": false,
+      "properties": {
+        "coordinates": {
+          "description": "The trap coordinates in µm.",
+          "items": {
+            "items": {
+              "type": "number"
+            },
+            "maxItems": 3,
+            "minItems": 3,
+            "type": "array"
+          },
+          "type": "array"
+        },
+        "slug": {
+          "description": "An optional name for the layout.",
+          "type": "string"
+        }
+      },
+      "required": [
+        "coordinates"
+      ],
+      "type": "object"
     }
   }
 }

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

@@ -26,15 +26,103 @@
       ],
       "type": "object"
     },
+    "Atom3D": {
+      "additionalProperties": false,
+      "properties": {
+        "name": {
+          "$ref": "#/definitions/QubitId",
+          "description": "Name of the atom."
+        },
+        "x": {
+          "description": "x-position in µm",
+          "type": "number"
+        },
+        "y": {
+          "description": "y-position in µm",
+          "type": "number"
+        },
+        "z": {
+          "description": "z-position in µm",
+          "type": "number"
+        }
+      },
+      "required": [
+        "name",
+        "x",
+        "y",
+        "z"
+      ],
+      "type": "object"
+    },
+    "Layout2D": {
+      "additionalProperties": false,
+      "properties": {
+        "coordinates": {
+          "description": "The trap coordinates in µm.",
+          "items": {
+            "items": {
+              "type": "number"
+            },
+            "maxItems": 2,
+            "minItems": 2,
+            "type": "array"
+          },
+          "type": "array"
+        },
+        "slug": {
+          "description": "An optional name for the layout.",
+          "type": "string"
+        }
+      },
+      "required": [
+        "coordinates"
+      ],
+      "type": "object"
+    },
+    "Layout3D": {
+      "additionalProperties": false,
+      "properties": {
+        "coordinates": {
+          "description": "The trap coordinates in µm.",
+          "items": {
+            "items": {
+              "type": "number"
+            },
+            "maxItems": 3,
+            "minItems": 3,
+            "type": "array"
+          },
+          "type": "array"
+        },
+        "slug": {
+          "description": "An optional name for the layout.",
+          "type": "string"
+        }
+      },
+      "required": [
+        "coordinates"
+      ],
+      "type": "object"
+    },
     "QubitId": {
       "description": "Name for a qubit.",
       "type": "string"
     },
     "Register": {
+      "anyOf": [
+        {
+          "$ref": "#/definitions/Register2D"
+        },
+        {
+          "$ref": "#/definitions/Register3D"
+        }
+      ]
+    },
+    "Register2D": {
       "additionalProperties": false,
       "properties": {
         "layout": {
-          "$ref": "layout-schema.json",
+          "$ref": "#/definitions/Layout2D",
           "description": "The trap layout underlying the register."
         },
         "register": {
@@ -49,6 +137,26 @@
         "register"
       ],
       "type": "object"
+    },
+    "Register3D": {
+      "additionalProperties": false,
+      "properties": {
+        "layout": {
+          "$ref": "#/definitions/Layout3D",
+          "description": "The trap layout underlying the register."
+        },
+        "register": {
+          "description": "A 3D register containing a set of atoms.",
+          "items": {
+            "$ref": "#/definitions/Atom3D"
+          },
+          "type": "array"
+        }
+      },
+      "required": [
+        "register"
+      ],
+      "type": "object"
     }
   }
 }