Add pure-Python result post-processing helpers

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Added
 
+- Pure-Python result post-processing helpers for converting probabilities to counts, relabeling state keys to bitstrings, computing marginals, and calculating Z-parity expectation values.
 - `QctrlQaoaJobCreationPayload` and `QctrlQaoaJobInput` for submitting Q-CTRL QAOA maxcut combinatorial-optimization jobs via `create_job`. The `create_job` body union now also accepts `QctrlQaoaJobCreationPayload`.
 - `cost_model` optional field on `BaseJob`, `GetCircuitJobResponse`, and `GetJobResponse`, typed as `CostModel` (`"quantum_compute_time"` or `"execution_time"`).
 

custom-templates/package_init.py.jinja

@@ -1,5 +1,5 @@
 {% from "helpers.jinja" import safe_docstring %}
-{% set modules = ["exceptions", "extensions", "gates", "ionq_client", "pagination", "polling", "session"] %}
+{% set modules = ["exceptions", "extensions", "gates", "ionq_client", "pagination", "polling", "results", "session"] %}
 {{ safe_docstring(package_description) }}
 from . import {{ modules | join(", ") }}
 from .client import AuthenticatedClient, Client  # noqa: F401

ionq_core/results.py

@@ -0,0 +1,163 @@
+# SPDX-FileCopyrightText: 2026 IonQ, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pure-Python helpers for post-processing IonQ probability results.
+
+IonQ result payloads encode measured basis states as integer strings. These
+helpers treat qubit 0 as the least significant bit of that integer key.
+"""
+
+__all__ = ["expectation_z", "marginal", "probabilities_to_counts", "relabel_to_bitstrings"]
+
+import math
+from collections.abc import Iterable, Mapping
+
+
+def probabilities_to_counts(probabilities: Mapping[str, float], shots: int) -> dict[str, int]:
+    """Convert a probability distribution to integer shot counts.
+
+    Counts are rounded with the largest-remainder method and always sum exactly
+    to ``shots``. Input probabilities must be finite, non-negative, and sum to 1
+    within floating-point tolerance.
+
+    Args:
+        probabilities: Mapping from integer-encoded state keys to probabilities.
+        shots: Number of shots to distribute.
+
+    Returns:
+        A mapping with the same keys and integer counts summing to ``shots``.
+    """
+    if shots < 0:
+        raise ValueError("shots must be non-negative")
+
+    scaled: list[tuple[str, int, float, int]] = []
+    total_probability = 0.0
+    for order, (key, probability) in enumerate(probabilities.items()):
+        probability = _validate_probability(probability, f"probabilities[{key!r}]")
+        total_probability += probability
+        expected_count = probability * shots
+        floor_count = math.floor(expected_count)
+        scaled.append((key, floor_count, expected_count - floor_count, order))
+
+    if not scaled:
+        if shots == 0:
+            return {}
+        raise ValueError("probabilities must not be empty when shots is positive")
+
+    if not math.isclose(total_probability, 1.0, rel_tol=1e-12, abs_tol=1e-12):
+        raise ValueError("probabilities must sum to 1")
+
+    counts = {key: floor_count for key, floor_count, _, _ in scaled}
+    remaining = shots - sum(counts.values())
+    if remaining:
+        for key, _, _, _ in sorted(scaled, key=lambda item: (-item[2], item[3]))[:remaining]:
+            counts[key] += 1
+    return counts
+
+
+def relabel_to_bitstrings(probabilities: Mapping[str, float], num_qubits: int) -> dict[str, float]:
+    """Relabel integer-encoded state keys to zero-padded bitstrings.
+
+    Args:
+        probabilities: Mapping from integer-encoded state keys to probabilities.
+        num_qubits: Width of the measured register.
+
+    Returns:
+        A mapping from bitstring labels to probabilities.
+    """
+    _validate_num_qubits(num_qubits)
+    return {
+        format(_parse_state_key(key, num_qubits), f"0{num_qubits}b"): _validate_probability(
+            probability, f"probabilities[{key!r}]"
+        )
+        for key, probability in probabilities.items()
+    }
+
+
+def marginal(probabilities: Mapping[str, float], qubits: Iterable[int], num_qubits: int) -> dict[str, float]:
+    """Return the probability marginal over a subset of qubits.
+
+    Qubit indices are interpreted little-endian: qubit 0 is the least
+    significant bit of the integer state key. The returned marginal key uses the
+    order supplied in ``qubits``; the first selected qubit becomes bit 0 in the
+    reduced integer key.
+
+    Args:
+        probabilities: Mapping from integer-encoded state keys to probabilities.
+        qubits: Qubits to keep in the returned marginal.
+        num_qubits: Width of the measured register.
+
+    Returns:
+        A probability mapping over the selected qubits.
+    """
+    selected_qubits = _normalize_qubits(qubits, num_qubits)
+    reduced: dict[str, float] = {}
+    for key, probability in probabilities.items():
+        state = _parse_state_key(key, num_qubits)
+        probability = _validate_probability(probability, f"probabilities[{key!r}]")
+        reduced_key = str(_project_state(state, selected_qubits))
+        reduced[reduced_key] = reduced.get(reduced_key, 0.0) + probability
+    return reduced
+
+
+def expectation_z(probabilities: Mapping[str, float], num_qubits: int) -> float:
+    """Compute the parity expectation value for ``Z`` on every measured qubit.
+
+    Args:
+        probabilities: Mapping from integer-encoded state keys to probabilities.
+        num_qubits: Width of the measured register.
+
+    Returns:
+        The expectation value ``sum(p(x) * (-1) ** popcount(x))``.
+    """
+    _validate_num_qubits(num_qubits)
+    expectation = 0.0
+    for key, probability in probabilities.items():
+        state = _parse_state_key(key, num_qubits)
+        probability = _validate_probability(probability, f"probabilities[{key!r}]")
+        expectation += probability if state.bit_count() % 2 == 0 else -probability
+    return expectation
+
+
+def _normalize_qubits(qubits: Iterable[int], num_qubits: int) -> tuple[int, ...]:
+    _validate_num_qubits(num_qubits)
+    selected_qubits = tuple(qubits)
+    seen: set[int] = set()
+    for qubit in selected_qubits:
+        if qubit < 0 or qubit >= num_qubits:
+            raise ValueError(f"qubit index {qubit} is outside the {num_qubits}-qubit register")
+        if qubit in seen:
+            raise ValueError(f"qubit index {qubit} is repeated")
+        seen.add(qubit)
+    return selected_qubits
+
+
+def _parse_state_key(key: str, num_qubits: int) -> int:
+    try:
+        state = int(key)
+    except ValueError as exc:
+        raise ValueError(f"state key {key!r} is not an integer") from exc
+
+    if state < 0 or state >= 1 << num_qubits:
+        raise ValueError(f"state key {key!r} is outside the {num_qubits}-qubit register")
+    return state
+
+
+def _project_state(state: int, qubits: tuple[int, ...]) -> int:
+    reduced_state = 0
+    for output_bit, qubit in enumerate(qubits):
+        reduced_state |= ((state >> qubit) & 1) << output_bit
+    return reduced_state
+
+
+def _validate_num_qubits(num_qubits: int) -> None:
+    if num_qubits < 1:
+        raise ValueError("num_qubits must be at least 1")
+
+
+def _validate_probability(probability: float, label: str) -> float:
+    if not math.isfinite(probability):
+        raise ValueError(f"{label} must be finite")
+    if probability < 0:
+        raise ValueError(f"{label} must be non-negative")
+    return probability

tests/test_results.py

@@ -0,0 +1,87 @@
+import math
+
+import pytest
+
+import ionq_core
+from ionq_core import expectation_z, marginal, probabilities_to_counts, relabel_to_bitstrings
+
+BELL_PROBABILITIES = {"0": 0.5, "3": 0.5}
+
+
+def test_results_helpers_are_reexported():
+    assert "results" in ionq_core.__all__
+    assert "probabilities_to_counts" in ionq_core.__all__
+
+
+def test_relabel_to_bitstrings_uses_zero_padded_integer_keys():
+    assert relabel_to_bitstrings(BELL_PROBABILITIES, 2) == {"00": 0.5, "11": 0.5}
+
+
+def test_probabilities_to_counts_uses_largest_remainder_rounding():
+    counts = probabilities_to_counts({"0": 1 / 3, "1": 1 / 3, "2": 1 / 3}, 10)
+
+    assert counts == {"0": 4, "1": 3, "2": 3}
+    assert sum(counts.values()) == 10
+
+
+def test_probabilities_to_counts_keeps_exact_counts_when_no_remainder():
+    assert probabilities_to_counts({"0": 0.2, "1": 0.3, "2": 0.5}, 10) == {"0": 2, "1": 3, "2": 5}
+
+
+def test_probabilities_to_counts_allows_empty_zero_shots_distribution():
+    assert probabilities_to_counts({}, 0) == {}
+
+
+def test_probabilities_to_counts_rejects_invalid_inputs():
+    with pytest.raises(ValueError, match="shots"):
+        probabilities_to_counts({"0": 1.0}, -1)
+
+    with pytest.raises(ValueError, match="empty"):
+        probabilities_to_counts({}, 1)
+
+    with pytest.raises(ValueError, match="finite"):
+        probabilities_to_counts({"0": math.inf}, 1)
+
+    with pytest.raises(ValueError, match="non-negative"):
+        probabilities_to_counts({"0": -0.1, "1": 1.1}, 1)
+
+    with pytest.raises(ValueError, match="sum to 1"):
+        probabilities_to_counts({"0": 0.4, "1": 0.5}, 10)
+
+
+def test_marginal_keeps_requested_qubits_in_requested_order():
+    probabilities = {"1": 0.25, "2": 0.75}
+
+    assert marginal(BELL_PROBABILITIES, [0], 2) == {"0": 0.5, "1": 0.5}
+    assert marginal(probabilities, [1, 0], 2) == {"2": 0.25, "1": 0.75}
+    assert marginal(BELL_PROBABILITIES, [], 2) == {"0": 1.0}
+
+
+def test_marginal_rejects_invalid_qubits():
+    with pytest.raises(ValueError, match="at least 1"):
+        marginal({"0": 1.0}, [0], 0)
+
+    with pytest.raises(ValueError, match="outside"):
+        marginal({"0": 1.0}, [2], 2)
+
+    with pytest.raises(ValueError, match="repeated"):
+        marginal({"0": 1.0}, [0, 0], 2)
+
+
+def test_expectation_z_computes_parity_expectation():
+    assert expectation_z(BELL_PROBABILITIES, 2) == 1.0
+    assert expectation_z({"1": 0.25, "2": 0.25, "3": 0.5}, 2) == 0.0
+
+
+def test_state_key_validation():
+    with pytest.raises(ValueError, match="not an integer"):
+        relabel_to_bitstrings({"zero": 1.0}, 2)
+
+    with pytest.raises(ValueError, match="outside"):
+        relabel_to_bitstrings({"4": 1.0}, 2)
+
+    with pytest.raises(ValueError, match="outside"):
+        expectation_z({"-1": 1.0}, 2)
+
+    with pytest.raises(ValueError, match="at least 1"):
+        relabel_to_bitstrings({"0": 1.0}, 0)