diff --git a/src/bloqade/analysis/tomography/__init__.py b/src/bloqade/analysis/tomography/__init__.py new file mode 100644 index 000000000..46de40470 --- /dev/null +++ b/src/bloqade/analysis/tomography/__init__.py @@ -0,0 +1,3 @@ +"""Helper class for computing tomography.""" + +from .tomography import TomographyResult as TomographyResult diff --git a/src/bloqade/analysis/tomography/tomography.py b/src/bloqade/analysis/tomography/tomography.py new file mode 100644 index 000000000..900002fd7 --- /dev/null +++ b/src/bloqade/analysis/tomography/tomography.py @@ -0,0 +1,123 @@ +"""Minimal single-qubit tomography helpers.""" + +from __future__ import annotations + +import math +from dataclasses import field, dataclass +from collections.abc import Mapping, Sequence + +import numpy as np + +BASES = set(("X", "Y", "Z")) + + +def _density_matrix_from_bloch(bloch: Mapping[str, float]) -> np.ndarray: + required_keys = 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 +class TomographyResult: + """Point-estimate single-qubit tomography result.""" + + shots_by_basis: Mapping[str, np.ndarray] + + density_matrix: np.ndarray = field(init=False) + + def __post_init__( + self, + ) -> 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(self.shots_by_basis) != BASES: + raise ValueError("Single-qubit tomography requires X, Y, and Z keys.") + + bloch: dict[str, float] = {} + for basis in BASES: + shots = np.asarray(self.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 + + 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", +] diff --git a/test/analysis/tomography/test_tomography.py b/test/analysis/tomography/test_tomography.py new file mode 100644 index 000000000..20053bbbe --- /dev/null +++ b/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)