From a531a500f693838751934e5ec0e26b37caaf7f4f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Henrique=20Silv=C3=A9rio?= Date: Mon, 9 Oct 2023 13:48:37 +0200 Subject: [PATCH] Abstract representation fixes (#595) * Fix `phase_shift()` deserialization * Support multi-target serialization * Bump version to v0.15.2 * Improve comments in the serializer * Serialize target to single int whenever possible --- VERSION.txt | 2 +- .../pulser/json/abstract_repr/deserializer.py | 12 +++- .../schemas/sequence-schema.json | 11 ++- .../pulser/json/abstract_repr/serializer.py | 39 ++++++++--- pulser-core/pulser/sequence/sequence.py | 8 ++- tests/test_abstract_repr.py | 68 +++++++++++++++++-- 6 files changed, 119 insertions(+), 21 deletions(-) diff --git a/VERSION.txt b/VERSION.txt index e815b861f..4312e0d0c 100644 --- a/VERSION.txt +++ b/VERSION.txt @@ -1 +1 @@ -0.15.1 +0.15.2 diff --git a/pulser-core/pulser/json/abstract_repr/deserializer.py b/pulser-core/pulser/json/abstract_repr/deserializer.py index e34111a56..fa304259e 100644 --- a/pulser-core/pulser/json/abstract_repr/deserializer.py +++ b/pulser-core/pulser/json/abstract_repr/deserializer.py @@ -63,7 +63,7 @@ VARIABLE_TYPE_MAP = {"int": int, "float": float} -ExpReturnType = Union[int, float, ParamObj] +ExpReturnType = Union[int, float, list, ParamObj] @overload @@ -76,6 +76,13 @@ def _deserialize_parameter(param: float, vars: dict[str, Variable]) -> float: pass +@overload +def _deserialize_parameter( + param: list[int], vars: dict[str, Variable] +) -> list[int]: + pass + + @overload def _deserialize_parameter( param: dict[str, str], vars: dict[str, Variable] @@ -84,7 +91,7 @@ def _deserialize_parameter( def _deserialize_parameter( - param: Union[int, float, dict[str, Any]], + param: Union[int, float, list[int], dict[str, Any]], vars: dict[str, Variable], ) -> Union[ExpReturnType, Variable]: """Deserialize a parameterized object. @@ -213,6 +220,7 @@ def _deserialize_operation(seq: Sequence, op: dict, vars: dict) -> None: seq.phase_shift_index( _deserialize_parameter(op["phi"], vars), *[_deserialize_parameter(t, vars) for t in op["targets"]], + basis=op["basis"], ) elif op["op"] == "pulse": phase = _deserialize_parameter(op["phase"], vars) diff --git a/pulser-core/pulser/json/abstract_repr/schemas/sequence-schema.json b/pulser-core/pulser/json/abstract_repr/schemas/sequence-schema.json index 96c42236b..734709b99 100644 --- a/pulser-core/pulser/json/abstract_repr/schemas/sequence-schema.json +++ b/pulser-core/pulser/json/abstract_repr/schemas/sequence-schema.json @@ -724,8 +724,15 @@ "type": "string" }, "target": { - "$ref": "#/definitions/ParametrizedNum", - "description": "New target atom index" + "anyOf": [ + { + "$ref": "#/definitions/ParametrizedNum" + }, + { + "$ref": "#/definitions/ParametrizedNumArray" + } + ], + "description": "New target atom index (or indices)" } }, "required": [ diff --git a/pulser-core/pulser/json/abstract_repr/serializer.py b/pulser-core/pulser/json/abstract_repr/serializer.py index b5e2e07c8..31354f1a0 100644 --- a/pulser-core/pulser/json/abstract_repr/serializer.py +++ b/pulser-core/pulser/json/abstract_repr/serializer.py @@ -16,19 +16,20 @@ import inspect import json +from collections.abc import Iterable from itertools import chain -from typing import TYPE_CHECKING, Any -from typing import Sequence as abcSequence -from typing import Union, cast +from typing import TYPE_CHECKING, Any, Union, cast import numpy as np +import pulser from pulser.json.abstract_repr.signatures import SIGNATURES from pulser.json.abstract_repr.validation import validate_abstract_repr from pulser.json.exceptions import AbstractReprError from pulser.json.utils import stringify_qubit_ids if TYPE_CHECKING: + from pulser.parametrized import Parametrized from pulser.register.base_register import QubitId from pulser.sequence import Sequence from pulser.sequence._call import _Call @@ -154,17 +155,29 @@ def serialize_abstract_sequence( for var in seq._variables.values(): res["variables"][var.name]["value"] = [var.dtype()] * var.size + def unfold_targets( + target_ids: QubitId | Iterable[QubitId], + ) -> QubitId | list[QubitId]: + if isinstance(target_ids, (int, str)): + return target_ids + + targets = list(cast(Iterable, target_ids)) + return targets if len(targets) > 1 else targets[0] + def convert_targets( - target_ids: Union[QubitId, abcSequence[QubitId]] + target_ids: Union[QubitId, Iterable[QubitId]], + force_list_out: bool = False, ) -> Union[int, list[int]]: - target_array = np.array(target_ids) + target_array = np.array(unfold_targets(target_ids)) og_dim = target_array.ndim if og_dim == 0: target_array = target_array[np.newaxis] indices = seq.get_register(include_mappable=True).find_indices( target_array.tolist() ) - return indices[0] if og_dim == 0 else indices + if force_list_out or og_dim > 0: + return indices + return indices[0] def get_kwarg_default(call_name: str, kwarg_name: str) -> Any: sig = inspect.signature(getattr(seq, call_name)) @@ -230,10 +243,20 @@ def remove_kwarg_if_default( ) elif "target" in call.name: data = get_all_args(("qubits", "channel"), call) + target: Parametrized | int | list[int] if call.name == "target": target = convert_targets(data["qubits"]) elif call.name == "target_index": - target = data["qubits"] + if isinstance( + data["qubits"], pulser.parametrized.Parametrized + ): + # The qubit indices are given through a variable + target = data["qubits"] + else: + # Either a single index or a sequence of indices + target = cast( + Union[int, list], unfold_targets(data["qubits"]) + ) else: raise AbstractReprError(f"Unknown call '{call.name}'.") operations.append( @@ -269,7 +292,7 @@ def remove_kwarg_if_default( elif "phase_shift" in call.name: targets = call.args[1:] if call.name == "phase_shift": - targets = convert_targets(targets) + targets = convert_targets(targets, force_list_out=True) elif call.name != "phase_shift_index": raise AbstractReprError(f"Unknown call '{call.name}'.") operations.append( diff --git a/pulser-core/pulser/sequence/sequence.py b/pulser-core/pulser/sequence/sequence.py index 504b77b81..aba14b16a 100644 --- a/pulser-core/pulser/sequence/sequence.py +++ b/pulser-core/pulser/sequence/sequence.py @@ -447,7 +447,9 @@ def current_phase_ref( ) if basis not in self._basis_ref: - raise ValueError("No declared channel targets the given 'basis'.") + raise ValueError( + f"No declared channel targets the given 'basis' ('{basis}')." + ) return self._basis_ref[basis][qubit].phase.last_phase @@ -2079,7 +2081,9 @@ def _phase_shift( _index: bool = False, ) -> None: if basis not in self._basis_ref: - raise ValueError("No declared channel targets the given 'basis'.") + raise ValueError( + f"No declared channel targets the given 'basis' ('{basis}')." + ) target_ids = self._check_qubits_give_ids(*targets, _index=_index) if not self.is_parametrized(): diff --git a/tests/test_abstract_repr.py b/tests/test_abstract_repr.py index f02605525..d85d85dd5 100644 --- a/tests/test_abstract_repr.py +++ b/tests/test_abstract_repr.py @@ -41,7 +41,7 @@ from pulser.json.exceptions import AbstractReprError, DeserializeDeviceError from pulser.parametrized.decorators import parametrize from pulser.parametrized.paramobj import ParamObj -from pulser.parametrized.variable import VariableItem +from pulser.parametrized.variable import Variable, VariableItem from pulser.register.register_layout import RegisterLayout from pulser.register.special_layouts import TriangularLatticeLayout from pulser.sequence._call import _Call @@ -260,7 +260,10 @@ def sequence(self, request): reg = Register(qubits) device = request.param seq = Sequence(reg, device) - seq.declare_channel("digital", "raman_local", initial_target="control") + + seq.declare_channel( + "digital", "raman_local", initial_target=("control",) + ) seq.declare_channel( "rydberg", "rydberg_local", initial_target="control" ) @@ -291,7 +294,7 @@ def sequence(self, request): seq.align("digital", "rydberg") seq.add(pi_pulse, "rydberg") seq.phase_shift(1.0, "control", "target", basis="ground-rydberg") - seq.target("target", "rydberg") + seq.target({"target"}, "rydberg") seq.add(two_pi_pulse, "rydberg") seq.delay(100, "digital") @@ -348,6 +351,12 @@ def test_values(self, abstract): assert abstract["operations"][0] == { "op": "target", "channel": "digital", + "target": 0, # tuple[int] is still serialized as int + } + + assert abstract["operations"][1] == { + "op": "target", + "channel": "rydberg", "target": 0, } @@ -414,6 +423,12 @@ def test_values(self, abstract): "post_phase_shift": 0.0, } + assert abstract["operations"][8] == { + "op": "target", + "channel": "rydberg", + "target": 1, + } + assert abstract["operations"][10] == { "op": "delay", "channel": "digital", @@ -897,6 +912,33 @@ def test_dmm_slm_mask(self, triangular_lattice, is_empty): assert abstract["operations"][3]["op"] == "pulse" assert abstract["operations"][3]["channel"] == "rydberg_global" + def test_multi_qubit_target(self): + seq_ = Sequence(Register.square(2, prefix="q"), MockDevice) + var_targets = seq_.declare_variable("var_targets", dtype=int, size=4) + + seq_.declare_channel( + "rydberg_local", "rydberg_local", initial_target=("q0", "q1") + ) + seq_.target(["q3", "q2"], "rydberg_local") + seq_.target_index(var_targets, "rydberg_local") + seq_.target(["q0"], "rydberg_local") + seq_.target_index(var_targets[2], "rydberg_local") + + abstract = json.loads(seq_.to_abstract_repr()) + + assert all(op["op"] == "target" for op in abstract["operations"]) + assert abstract["operations"][0]["target"] == [0, 1] + assert abstract["operations"][1]["target"] == [3, 2] + assert abstract["operations"][2]["target"] == { + "variable": "var_targets" + } + assert abstract["operations"][3]["target"] == 0 + assert abstract["operations"][4]["target"] == { + "expression": "index", + "lhs": {"variable": "var_targets"}, + "rhs": 2, + } + def _get_serialized_seq( operations: list[dict] = [], @@ -1185,6 +1227,7 @@ def test_deserialize_variables(self, without_default): "op", [ {"op": "target", "target": 2, "channel": "digital"}, + {"op": "target", "target": [1, 2], "channel": "digital"}, {"op": "delay", "time": 500, "channel": "global"}, {"op": "align", "channels": ["digital", "global"]}, { @@ -1215,7 +1258,9 @@ def test_deserialize_variables(self, without_default): ids=_get_op, ) def test_deserialize_non_parametrized_op(self, op): - s = _get_serialized_seq(operations=[op]) + s = _get_serialized_seq( + operations=[op], device=json.loads(MockDevice.to_abstract_repr()) + ) _check_roundtrip(s) seq = Sequence.from_abstract_repr(json.dumps(s)) @@ -1240,6 +1285,7 @@ def test_deserialize_non_parametrized_op(self, op): elif op["op"] == "phase_shift": assert c.name == "phase_shift_index" assert c.args == tuple([op["phi"], *op["targets"]]) + assert c.kwargs["basis"] == "digital" elif op["op"] == "pulse": assert c.name == "add" assert c.kwargs["channel"] == op["channel"] @@ -1390,12 +1436,17 @@ def test_deserialize_measurement(self): "op", [ {"op": "target", "target": var1, "channel": "digital"}, + { + "op": "target", + "target": {"variable": "var1"}, + "channel": "digital", + }, {"op": "delay", "time": var2, "channel": "global"}, { "op": "phase_shift", "phi": var1, "targets": [2, var1], - "basis": "digital", + "basis": "ground-rydberg", }, { "op": "pulse", @@ -1438,7 +1489,10 @@ def test_deserialize_parametrized_op(self, op): c = seq._to_build_calls[0] if op["op"] == "target": assert c.name == "target_index" - assert isinstance(c.kwargs["qubits"], VariableItem) + target_type = ( + VariableItem if "expression" in op["target"] else Variable + ) + assert isinstance(c.kwargs["qubits"], target_type) assert c.kwargs["channel"] == op["channel"] elif op["op"] == "delay": assert c.name == "delay" @@ -1452,6 +1506,8 @@ def test_deserialize_parametrized_op(self, op): assert c.args[1] == 2 # qubit 2 is variable assert isinstance(c.args[2], VariableItem) + # basis is fixed + assert c.kwargs["basis"] == "ground-rydberg" elif op["op"] == "pulse": assert c.name == "add" assert c.kwargs["channel"] == op["channel"]