feat: added tomography result class

src/bloqade/analysis/tomography/__init__.py

@@ -0,0 +1,3 @@
+"""Helper class for computing tomography."""
+
+from .tomography import TomographyResult as TomographyResult

src/bloqade/analysis/tomography/tomography.py

@@ -0,0 +1,122 @@
+"""Minimal single-qubit tomography helpers."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+
+BASES = ("X", "Y", "Z")
+
+
+def _density_matrix_from_bloch(bloch: Mapping[str, float]) -> np.ndarray:
+    required_keys = set(BASES)
+    if set(bloch) != required_keys:
+        raise ValueError("Single-qubit tomography requires X, Y, and Z keys.")
+
+    x = float(bloch["X"])
+    y = float(bloch["Y"])
+    z = float(bloch["Z"])
+    return 0.5 * np.array(
+        [[1.0 + z, x - 1j * y], [x + 1j * y, 1.0 - z]],
+        dtype=np.complex128,
+    )
+
+
+def _bloch_mapping_from_sequence(
+    bloch: np.ndarray | Sequence[float],
+) -> dict[str, float]:
+    bloch_arr = np.asarray(bloch, dtype=np.float64)
+    return {
+        "X": float(bloch_arr[0]),
+        "Y": float(bloch_arr[1]),
+        "Z": float(bloch_arr[2]),
+    }
+
+
+def _validate_single_qubit_bloch_vector(
+    bloch: np.ndarray | Sequence[float],
+    *,
+    tol: float,
+) -> dict[str, float]:
+    if tol < 0:
+        raise ValueError("tol must be non-negative.")
+    bloch_arr = np.asarray(bloch, dtype=np.float64)
+    if bloch_arr.shape != (3,):
+        raise ValueError("bloch must be a length-3 vector.")
+    if not np.all(np.isfinite(bloch_arr)):
+        raise ValueError("bloch components must be finite.")
+    bloch_norm_squared = float(np.dot(bloch_arr, bloch_arr))
+    if bloch_norm_squared > 1.0 + tol:
+        raise ValueError("Single-qubit Bloch vector must have squared norm <= 1.")
+    return _bloch_mapping_from_sequence(bloch_arr)
+
+
+def _single_qubit_fidelity(
+    density_matrix: np.ndarray,
+    target_density_matrix: np.ndarray,
+) -> float:
+    overlap = float(np.real(np.trace(density_matrix @ target_density_matrix)))
+    det_product = float(
+        np.real(np.linalg.det(density_matrix))
+        * np.real(np.linalg.det(target_density_matrix))
+    )
+    return overlap + 2.0 * math.sqrt(max(det_product, 0.0))
+
+
+@dataclass(frozen=True, init=False)
+class TomographyResult:
+    """Point-estimate single-qubit tomography result."""
+
+    density_matrix: np.ndarray
+
+    def __init__(
+        self,
+        shots_by_basis: Mapping[str, np.ndarray],
+    ) -> None:
+        """
+        Create a tomography result by computing the density matrix from the shots per basis.
+
+        Args:
+            shots_by_basis (Mapping[str, np.ndarray]): A mapping of each basis to an array of shots (0/1's) in each basis.
+        """
+        if set(shots_by_basis) != set(BASES):
+            raise ValueError("Single-qubit tomography requires X, Y, and Z keys.")
+
+        bloch: dict[str, float] = {}
+        for basis in BASES:
+            shots = np.asarray(shots_by_basis[basis])
+            if shots.ndim != 1:
+                raise ValueError(
+                    "TomographyResult expects each basis to have shape (shots,)."
+                )
+            if shots.size == 0:
+                raise ValueError(f"{basis}-basis shots cannot be empty.")
+            if not np.all((shots == 0) | (shots == 1)):
+                raise ValueError("Tomography shots must contain only zero or one.")
+
+            shots = shots.astype(np.uint8, copy=False)
+            prob_meas_one = float(np.mean(shots))
+            bloch[basis] = 1.0 - 2.0 * prob_meas_one
+
+        object.__setattr__(self, "density_matrix", _density_matrix_from_bloch(bloch))
+
+    # NOTE: if you want to add more generic methods for fidelity, to density matrices, just define a new method "fidelity_to_density_mat".
+    def fidelity_bloch(
+        self,
+        target_bloch: np.ndarray | Sequence[float],
+        tol: float = 1e-10,
+    ) -> float:
+        """Return the fidelity to a target state from its Bloch vector."""
+
+        target_density_matrix = _density_matrix_from_bloch(
+            _validate_single_qubit_bloch_vector(target_bloch, tol=tol)
+        )
+        return _single_qubit_fidelity(self.density_matrix, target_density_matrix)
+
+
+__all__ = [
+    "TomographyResult",
+]

test/analysis/tomography/test_tomography.py

@@ -0,0 +1,98 @@
+import numpy as np
+import pytest
+
+from bloqade.analysis.tomography import TomographyResult
+
+
+def test_reconstructs_single_qubit_density_matrix():
+    result = TomographyResult(
+        {
+            "X": np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 1]),
+            "Y": np.array([0, 0, 0, 1, 1, 1, 1, 1, 1, 1]),
+            "Z": np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 1]),
+        }
+    )
+
+    expected_bloch = {"X": -0.2, "Y": -0.4, "Z": 0.6}
+    expected_density_matrix = 0.5 * np.array(
+        [[1.6, -0.2 + 0.4j], [-0.2 - 0.4j, 0.4]],
+        dtype=np.complex128,
+    )
+
+    np.testing.assert_allclose(result.density_matrix, expected_density_matrix)
+
+    reconstructed_bloch = np.array(
+        [
+            2.0 * result.density_matrix[0, 1].real,
+            -2.0 * result.density_matrix[0, 1].imag,
+            (result.density_matrix[0, 0] - result.density_matrix[1, 1]).real,
+        ]
+    )
+    np.testing.assert_allclose(reconstructed_bloch, list(expected_bloch.values()))
+
+
+def test_fidelity_for_pure_target():
+    result = TomographyResult(
+        {
+            "X": np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 1]),
+            "Y": np.array([0, 0, 0, 1, 1, 1, 1, 1, 1, 1]),
+            "Z": np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 1]),
+        }
+    )
+    target = np.ones(3) / np.sqrt(3.0)
+
+    fidelity = result.fidelity_bloch(target)
+
+    measured_bloch = np.array([-0.2, -0.4, 0.6])
+    expected_fidelity = 0.5 * (1.0 + measured_bloch @ target)
+
+    assert fidelity == pytest.approx(expected_fidelity)
+
+
+@pytest.mark.parametrize("missing_basis", ["X", "Y", "Z"])
+def test_requires_every_basis(missing_basis):
+    shots = {
+        "X": np.array([0, 1]),
+        "Y": np.array([0, 1]),
+        "Z": np.array([0, 1]),
+    }
+    del shots[missing_basis]
+
+    with pytest.raises(ValueError, match="requires X, Y, and Z keys"):
+        TomographyResult(shots)
+
+
+@pytest.mark.parametrize(
+    ("bad_shots", "message"),
+    [
+        (np.array([]), "cannot be empty"),
+        (np.array([[0], [1]]), r"shape \(shots,\)"),
+        (np.array([0.0, 0.5, 1.0]), "only zero or one"),
+    ],
+)
+def test_rejects_invalid_shots(bad_shots, message):
+    with pytest.raises(ValueError, match=message):
+        TomographyResult(
+            {
+                "X": bad_shots,
+                "Y": np.array([0, 1]),
+                "Z": np.array([0, 1]),
+            }
+        )
+
+
+@pytest.mark.parametrize(
+    "target",
+    [np.array([1.0, 0.0]), np.array([2.0, 0.0, 0.0]), np.array([np.nan, 0, 0])],
+)
+def test_rejects_invalid_target_bloch_vectors(target):
+    result = TomographyResult(
+        {
+            "X": np.array([0, 1]),
+            "Y": np.array([0, 1]),
+            "Z": np.array([0, 1]),
+        }
+    )
+
+    with pytest.raises(ValueError):
+        result.fidelity_bloch(target)