feat: Add Iceberg architecture builder and annotation

src/guppyft/code/iceberg/encode.py

@@ -0,0 +1,194 @@
+"""Builder for public configuration of the Iceberg QEC architecture."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import asdict, dataclass, field, replace
+from enum import StrEnum, auto
+from typing import TYPE_CHECKING, Any, Self
+
+from guppyft.encode import annotate_encoding
+
+if TYPE_CHECKING:
+    from hugr import Hugr
+    from hugr.package import Package
+    from hugr.utils import JsonType
+
+
+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 IcebergQECPolicy:
+    """Policy to determine when QEC cycles are injected.
+
+    Each logical block accumulates a cost based on `costs`. Once a
+    block's accumulated cost reaches `threshold`, a QEC cycle of the given
+    `style` is performed on that block and its counter is reset.
+
+    Attributes:
+        threshold: Threshold at which a QEC cycle is triggered.
+        costs: See `OperationCosts`.
+    """
+
+    @dataclass
+    class OperationCosts:
+        """Configuration for operation costs. Used e.g. for applying QEC cycles."""
+
+        idle: float = 0.0
+        all_x: float = 0.0
+        all_y: float = 0.0
+        all_z: float = 0.0
+        all_but_one_x: float = 0.0
+        all_but_one_z: float = 0.0
+        all_h: float = 0.0
+        x_with_all_but_one_z: float = 0.0
+        z_with_all_but_one_x: float = 0.0
+        cx_transversal: float = 0.0
+        yy_phase_intrablock: float = 0.0
+        xx_phase_interblock: float = 0.0
+        yy_phase_interblock: float = 0.0
+        alloc_dynq: float = 0.0
+        x: float = 0.0
+        y: float = 0.0
+        z: float = 0.0
+        xx: float = 0.0
+        yy: float = 0.0
+        zz: float = 0.0
+        cx_intrablock: float = 0.0
+        cx_interblock: float = 0.0
+        fan_in: float = 0.0
+        fan_out: float = 0.0
+        swap: float = 0.0
+        rx: float = 0.0
+        ry: float = 0.0
+        rz: float = 0.0
+        all_rx: float = 0.0
+        all_ry: float = 0.0
+        all_rz: float = 0.0
+        all_but_one_rx: float = 0.0
+        all_but_one_rz: float = 0.0
+        xx_phase_intrablock: float = 0.0
+        zz_phase_intrablock: float = 0.0
+        zz_phase_interblock: 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 QEC policy to JSON."""
+        return asdict(self)
+
+
+@dataclass(frozen=True, kw_only=True)
+class _IcebergOptions:
+    init_style: InitStyle
+    qec_policy: IcebergQECPolicy
+    syndrome_during_measurement: bool
+    state_prep_rus_attempts: int
+    dynamic_qubit_reuse: DynamicQubitReuse
+
+
+@dataclass(frozen=True)
+class IcebergInstance:
+    """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,
+                "qec_policy": self._options.qec_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,
+            qec_policy=IcebergQECPolicy(),
+            syndrome_during_measurement=False,
+            state_prep_rus_attempts=3,
+            dynamic_qubit_reuse=DynamicQubitReuse.enabled,
+        )
+    )
+
+    def with_qec_policy(self, qec_policy: IcebergQECPolicy) -> Self:
+        """Set the QEC policy."""
+        return replace(self, _options=replace(self._options, qec_policy=qec_policy))
+
+    def with_init_style(self, style: InitStyle) -> Self:
+        """Set the state 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 log depth state prep is allowed before it
+        fails."""
+        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."""
+        return IcebergInstance(
+            _n_blocks=n_blocks, _block_size=block_size, _options=deepcopy(self._options)
+        )

tests/code/iceberg/test_iceberg_encode.py

@@ -0,0 +1,105 @@
+from typing import Any
+
+from hugr import Hugr
+from hugr.package import Package
+
+from guppyft.code.iceberg.encode import (
+    DynamicQubitReuse,
+    IcebergBuilder,
+    IcebergQECPolicy,
+    InitStyle,
+)
+
+
+def test_annotate_steane_encoding() -> None:
+    hugr = Hugr[Any]()
+    pkg = Package([hugr])
+
+    policy = IcebergQECPolicy()
+    policy.threshold = 42
+    policy.costs.x = 1000.0
+    arch = (
+        IcebergBuilder()
+        .with_init_style(InitStyle.HALF_DEPTH)
+        .with_qec_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",
+                "qec_policy": {
+                    "threshold": 42,
+                    "costs": {
+                        "idle": 0.0,
+                        "all_x": 0.0,
+                        "all_y": 0.0,
+                        "all_z": 0.0,
+                        "all_but_one_x": 0.0,
+                        "all_but_one_z": 0.0,
+                        "all_h": 0.0,
+                        "x_with_all_but_one_z": 0.0,
+                        "z_with_all_but_one_x": 0.0,
+                        "cx_transversal": 0.0,
+                        "yy_phase_intrablock": 0.0,
+                        "xx_phase_interblock": 0.0,
+                        "yy_phase_interblock": 0.0,
+                        "alloc_dynq": 0.0,
+                        "x": 1000.0,
+                        "y": 0.0,
+                        "z": 0.0,
+                        "xx": 0.0,
+                        "yy": 0.0,
+                        "zz": 0.0,
+                        "cx_intrablock": 0.0,
+                        "cx_interblock": 0.0,
+                        "fan_in": 0.0,
+                        "fan_out": 0.0,
+                        "swap": 0.0,
+                        "rx": 0.0,
+                        "ry": 0.0,
+                        "rz": 0.0,
+                        "all_rx": 0.0,
+                        "all_ry": 0.0,
+                        "all_rz": 0.0,
+                        "all_but_one_rx": 0.0,
+                        "all_but_one_rz": 0.0,
+                        "xx_phase_intrablock": 0.0,
+                        "zz_phase_intrablock": 0.0,
+                        "zz_phase_interblock": 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 = IcebergQECPolicy(threshold=2)
+    policy.costs.x = 4.0
+    instance = IcebergBuilder().with_qec_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"]["qec_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]