feat: Add Iceberg architecture builder and annotation

src/guppyft/code/iceberg/encode.py

@@ -0,0 +1,206 @@
+"""Builder for public configuration of the Iceberg QEC architecture."""
+
+from copy import deepcopy
+from dataclasses import asdict, dataclass, field, replace
+from enum import StrEnum, auto
+from typing import Any, Self
+
+from hugr import Hugr
+from hugr.package import Package
+from hugr.utils import JsonType
+
+from guppyft.encode import AnnotatingArchitectureInstance, annotate_encoding
+
+
+class DynamicQubitReuse(StrEnum):
+    """Configuration for how dynamic qubits are allocated or reused. Fresh qubits are
+    from block allocations and have not yet been allocated. When dynamic qubits are
+    freed, depending on the configuration, future dynamic qubit allocations can reset
+    and reuse these qubits.
+
+    When all dynamic qubits from a logical block have been freed, the block itself is
+    freed and becomes available to be reallocated."""
+
+    disabled = auto()
+    """Only allocate fresh qubits from logical block allocations. Do not reset and reuse
+    dynamic qubits."""
+    if_necessary = auto()
+    """Prefer using fresh logical qubits. Only reuse logical qubits if no more blocks
+    are available to be allocated."""
+    enabled = auto()
+    """Use available fresh qubits, then freed qubits, before preparing another block."""
+
+
+class InitStyle(StrEnum):
+    """The initialisation style of an Iceberg block."""
+
+    HALF_DEPTH = auto()
+    LOG_DEPTH = auto()
+
+
+@dataclass
+class IcebergQEDPolicy:
+    """Policy to determine when QED cycles are injected.
+
+    Each logical block accumulates a cost based on `costs`. Once a block's accumulated
+    cost reaches `threshold`, a QED cycle is performed on that block and its counter is
+    reset.
+
+    Attributes:
+        threshold: Threshold at which a QED cycle is triggered.
+        costs: See `IcebergQEDPolicy.OperationCosts`.
+    """
+
+    @dataclass
+    class OperationCosts:
+        """Configuration for operation costs that are used to determined when QED cycles
+        should be performed during runtime."""
+
+        idle: float = 0.0
+        alloc_dynq: float = 0.0
+        x: float = 0.0
+        y: float = 0.0
+        z: float = 0.0
+        x_all: float = 0.0
+        y_all: float = 0.0
+        z_all: float = 0.0
+        x_all_except_at: float = 0.0
+        z_all_except_at: float = 0.0
+        x_all_except_z_at: float = 0.0
+        z_all_except_x_at: float = 0.0
+        rx: float = 0.0
+        rz: float = 0.0
+        pauli_exp_x_all: float = 0.0
+        pauli_exp_y_all: float = 0.0
+        pauli_exp_z_all: float = 0.0
+        pauli_exp_x_all_except_at: float = 0.0
+        pauli_exp_z_all_except_at: float = 0.0
+        pauli_exp_x_all_except_z_at: float = 0.0
+        pauli_exp_z_all_except_x_at: float = 0.0
+        h_all: float = 0.0
+        cx_all: float = 0.0
+        cx_intrablock: float = 0.0
+        cx_interblock: float = 0.0
+        xx_intrablock: float = 0.0
+        yy_intrablock: float = 0.0
+        zz_intrablock: float = 0.0
+        xx_phase_intrablock: float = 0.0
+        xx_phase_interblock: float = 0.0
+        yy_phase_intrablock: float = 0.0
+        yy_phase_interblock: float = 0.0
+        zz_phase_intrablock: float = 0.0
+        zz_phase_interblock: float = 0.0
+        fan_in: float = 0.0
+        fan_out: float = 0.0
+        swap_intrablock: float = 0.0
+
+        def __setattr__(self, key: str, value: Any) -> None:
+            """Set a non-negative cost for a known logical operation."""
+            if not hasattr(self, key):
+                raise KeyError(f"Unknown cost key: {key}")
+            if value < 0:
+                raise ValueError(f"Op cost cannot be negative: received {value}")
+            super().__setattr__(key, value)
+
+    threshold: int = 1
+    costs: OperationCosts = field(default_factory=OperationCosts)
+
+    def to_json(self) -> JsonType:
+        """Serialise the QED policy to JSON."""
+        return asdict(self)
+
+
+@dataclass(frozen=True, kw_only=True)
+class _IcebergOptions:
+    init_style: InitStyle
+    qed_policy: IcebergQEDPolicy
+    syndrome_during_measurement: bool
+    state_prep_rus_attempts: int
+    dynamic_qubit_reuse: DynamicQubitReuse
+
+
+@dataclass(frozen=True)
+class IcebergInstance(AnnotatingArchitectureInstance):
+    """An Iceberg architecture instance built by `IcebergBuilder.build`."""
+
+    _n_blocks: int
+    _block_size: int
+    _options: _IcebergOptions
+
+    def annotate(self, hugr: Package | Hugr[Any]) -> None:
+        """Annotate the current parameters of the Iceberg code architecture on the given
+        HUGR's metadata. See `guppyft.encode.annotate_encoding` for more information."""
+        annotate_encoding(
+            hugr,
+            "iceberg",
+            {
+                "schema_version": 1,
+                "n_blocks": self._n_blocks,
+                "block_size": self._block_size,
+                "init_style": self._options.init_style,
+                "qed_policy": self._options.qed_policy.to_json(),
+                "syndrome_during_measurement": self._options.syndrome_during_measurement,  # noqa: E501 ruff format keeps changing this line back
+                "state_prep_rus_attempts": self._options.state_prep_rus_attempts,
+                "dynamic_qubit_reuse": self._options.dynamic_qubit_reuse,
+            },
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class IcebergBuilder:
+    """Iceberg architecture builder class for creating `IcebergInstance` objects."""
+
+    _options: _IcebergOptions = field(
+        default=_IcebergOptions(
+            init_style=InitStyle.LOG_DEPTH,
+            qed_policy=IcebergQEDPolicy(),
+            syndrome_during_measurement=False,
+            state_prep_rus_attempts=3,
+            dynamic_qubit_reuse=DynamicQubitReuse.enabled,
+        )
+    )
+
+    def with_qed_policy(self, qed_policy: IcebergQEDPolicy) -> Self:
+        """Set the QED policy."""
+        return replace(self, _options=replace(self._options, qed_policy=qed_policy))
+
+    def with_init_style(self, style: InitStyle) -> Self:
+        """Set the Iceberg block initialization style."""
+        return replace(self, _options=replace(self._options, init_style=style))
+
+    def with_state_prep_rus_attempts(self, max_attempts: int) -> Self:
+        """Set the number of attempts state prep is allowed before it fails."""
+        if max_attempts < 1:
+            raise ValueError(
+                f"Cannot set the maximum state prep RUS attempts lower than 1,"
+                f" {max_attempts} given"
+            )
+        return replace(
+            self, _options=replace(self._options, state_prep_rus_attempts=max_attempts)
+        )
+
+    def with_syndrome_extraction_during_measurement(self, enabled: bool) -> Self:
+        """Set whether syndromes are extracted and checked when measuring a block."""
+        return replace(
+            self, _options=replace(self._options, syndrome_during_measurement=enabled)
+        )
+
+    def with_dynamic_qubit_reuse(self, conf: DynamicQubitReuse) -> Self:
+        """Set whether dynamic qubits should be reset and reused."""
+        return replace(self, _options=replace(self._options, dynamic_qubit_reuse=conf))
+
+    def build(self, *, n_blocks: int, block_size: int) -> IcebergInstance:
+        """Build an `IcebergInstance` configured for `n_blocks` logical blocks each
+        with `block_size` logical qubits."""
+        if n_blocks < 0:
+            raise ValueError(
+                f"Cannot set the number of blocks lower than 0, {n_blocks} given"
+            )
+        if block_size < 2 or block_size % 2 != 0:
+            raise ValueError(
+                f"Block size has to be even and at least 2, {block_size} given"
+            )
+
+        return IcebergInstance(
+            _n_blocks=n_blocks, _block_size=block_size, _options=deepcopy(self._options)
+        )

src/guppyft/code/steane/encode.py

@@ -100,7 +100,7 @@ class QECPolicy:
     Attributes:
         style: The style of syndrome extraction to use (see `QECStyle`).
         threshold: Threshold at which a QEC cycle is triggered.
-        costs: See `OperationCosts`.
+        costs: See `QECPolicy.OperationCosts`.
     """
 
     @dataclass

tests/code/iceberg/test_iceberg_encode.py

@@ -0,0 +1,170 @@
+from dataclasses import fields
+from typing import Any
+
+import pytest
+from hugr import Hugr
+from hugr.package import Package
+
+from guppyft.code.iceberg.encode import (
+    DynamicQubitReuse,
+    IcebergBuilder,
+    IcebergQEDPolicy,
+    InitStyle,
+)
+from guppyft.extensions import iceberg_ops
+
+
+def test_iceberg_invalid_max_attempts() -> None:
+    arch = IcebergBuilder()
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot set the maximum state prep RUS attempts lower than 1, 0 given",
+    ):
+        arch.with_state_prep_rus_attempts(0)
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot set the maximum state prep RUS attempts lower than 1, -5 given",
+    ):
+        arch.with_state_prep_rus_attempts(-5)
+
+
+def test_iceberg_invalid_n_blocks() -> None:
+    arch = IcebergBuilder()
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot set the number of blocks lower than 0, -1 given",
+    ):
+        arch.build(n_blocks=-1, block_size=4)
+
+
+def test_iceberg_invalid_block_size() -> None:
+    arch = IcebergBuilder()
+    with pytest.raises(
+        ValueError,
+        match=r"Block size has to be even and at least 2, 0 given",
+    ):
+        arch.build(n_blocks=2, block_size=0)
+    with pytest.raises(
+        ValueError,
+        match=r"Block size has to be even and at least 2, -2 given",
+    ):
+        arch.build(n_blocks=2, block_size=-2)
+    with pytest.raises(
+        ValueError,
+        match=r"Block size has to be even and at least 2, 3 given",
+    ):
+        arch.build(n_blocks=2, block_size=3)
+
+
+def test_annotate_iceberg_encoding() -> None:
+    hugr = Hugr[Any]()
+    pkg = Package([hugr])
+
+    policy = IcebergQEDPolicy()
+    policy.threshold = 42
+    policy.costs.x = 1000.0
+    arch = (
+        IcebergBuilder()
+        .with_init_style(InitStyle.HALF_DEPTH)
+        .with_qed_policy(policy)
+        .with_syndrome_extraction_during_measurement(True)
+        .with_state_prep_rus_attempts(5)
+        .with_dynamic_qubit_reuse(DynamicQubitReuse.disabled)
+        .build(n_blocks=4, block_size=6)
+    )
+    arch.annotate(pkg)
+
+    assert hugr[hugr.module_root].metadata["qnt.nexus.compilation"] == {
+        "schema_version": 1,
+        "ft_encoding": {
+            "name": "iceberg",
+            "params": {
+                "schema_version": 1,
+                "n_blocks": 4,
+                "block_size": 6,
+                "init_style": "half_depth",
+                "qed_policy": {
+                    "threshold": 42,
+                    "costs": {
+                        "idle": 0.0,
+                        "alloc_dynq": 0.0,
+                        "x": 1000.0,
+                        "y": 0.0,
+                        "z": 0.0,
+                        "x_all": 0.0,
+                        "y_all": 0.0,
+                        "z_all": 0.0,
+                        "x_all_except_at": 0.0,
+                        "z_all_except_at": 0.0,
+                        "x_all_except_z_at": 0.0,
+                        "z_all_except_x_at": 0.0,
+                        "rx": 0.0,
+                        "rz": 0.0,
+                        "pauli_exp_x_all": 0.0,
+                        "pauli_exp_y_all": 0.0,
+                        "pauli_exp_z_all": 0.0,
+                        "pauli_exp_x_all_except_at": 0.0,
+                        "pauli_exp_z_all_except_at": 0.0,
+                        "pauli_exp_x_all_except_z_at": 0.0,
+                        "pauli_exp_z_all_except_x_at": 0.0,
+                        "h_all": 0.0,
+                        "cx_all": 0.0,
+                        "cx_intrablock": 0.0,
+                        "cx_interblock": 0.0,
+                        "xx_intrablock": 0.0,
+                        "yy_intrablock": 0.0,
+                        "zz_intrablock": 0.0,
+                        "xx_phase_intrablock": 0.0,
+                        "xx_phase_interblock": 0.0,
+                        "yy_phase_intrablock": 0.0,
+                        "yy_phase_interblock": 0.0,
+                        "zz_phase_intrablock": 0.0,
+                        "zz_phase_interblock": 0.0,
+                        "fan_in": 0.0,
+                        "fan_out": 0.0,
+                        "swap_intrablock": 0.0,
+                    },
+                },
+                "syndrome_during_measurement": True,
+                "state_prep_rus_attempts": 5,
+                "dynamic_qubit_reuse": "disabled",
+            },
+        },
+    }
+
+
+def test_annotated_policy_immutable() -> None:
+    hugr = Hugr[Any]()
+    pkg = Package([hugr])
+
+    policy = IcebergQEDPolicy(threshold=2)
+    policy.costs.x = 4.0
+    instance = IcebergBuilder().with_qed_policy(policy).build(n_blocks=4, block_size=6)
+
+    policy.threshold = 99
+    policy.costs.x = 20.0
+    instance.annotate(pkg)
+
+    compilation = hugr[hugr.module_root].metadata["qnt.nexus.compilation"]
+    policy_val = compilation["ft_encoding"]["params"]["qed_policy"]  # type: ignore[index,call-overload]
+    assert policy_val["threshold"] == 2  # type: ignore[index,call-overload]
+    assert policy_val["costs"]["x"] == 4.0  # type: ignore[index,call-overload]
+
+
+def test_costs_have_existing_ops() -> None:
+    costs = {cost.name for cost in fields(IcebergQEDPolicy.OperationCosts)}
+    assert "idle" in costs
+    costs.remove("idle")
+
+    ops = {
+        op
+        for op in iceberg_ops().operations
+        if (
+            (("dynq" not in op and "alloc" not in op) or op == "alloc_dynq")
+            and "measure" not in op
+            and op
+            not in ("free", "check_pre_block", "borrow", "restore", "take", "put")
+        )
+    }
+
+    assert costs == ops