Source code for atlas_q.ir_enhanced.state_tomography

"""
IR-Enhanced Quantum State Tomography
======================================

Applies IR coherence-guided adaptive sampling for efficient quantum
state reconstruction.

Key Insight:
-----------
Full tomography requires 4^n measurements (exponential!)
Compressed sensing helps but still measurement-heavy

IR approach:
1. Measure subset to estimate coherence matrix
2. Identify high-mutual-information measurement pairs
3. Prioritize measurements with high coherence
4. Reconstruct state from grouped measurements

Target: 10-1000× measurement reduction

Author: ATLAS-Q + IR Integration
Date: November 2025
"""

from dataclasses import dataclass
from typing import List, Optional

import numpy as np




@dataclass
class TomographyStrategy:
    """Result of IR tomography planning"""

    measurement_basis: List[str]  # Pauli basis to measure
    measurement_order: List[int]  # Priority order
    grouping: List[List[int]]  # Commuting measurement groups
    n_measurements: int  # Total measurements needed
    compression_factor: float  # vs full tomography
    method: str





def generate_pauli_basis(n_qubits: int, max_weight: Optional[int] = None) -> List[str]:
    """
    Generate Pauli basis for n qubits.

    Parameters
    ----------
    n_qubits : int
        Number of qubits
    max_weight : int, optional
        Maximum Pauli weight (Hamming weight)
        If None, generates full basis (4^n terms)

    Returns
    -------
    List[str]
        Pauli strings
    """
    if max_weight is None:
        max_weight = n_qubits

    paulis = []
    basis_ops = ['I', 'X', 'Y', 'Z']

    def generate_recursive(current, depth):
        if depth == n_qubits:
            weight = sum(1 for p in current if p != 'I')
            if weight <= max_weight:
                paulis.append(''.join(current))
            return

        for op in basis_ops:
            generate_recursive(current + [op], depth + 1)

    generate_recursive([], 0)
    return paulis





def ir_state_tomography(
    n_qubits: int,
    max_weight: int = 2,
    target_measurements: Optional[int] = None,
    adaptive: bool = True
) -> TomographyStrategy:
    """
    IR-enhanced state tomography measurement planning.

    Generates optimal measurement strategy for quantum state reconstruction.

    Parameters
    ----------
    n_qubits : int
        Number of qubits to reconstruct
    max_weight : int
        Maximum Pauli weight (reduces measurement count)
    target_measurements : int, optional
        Target number of measurements (if None, uses all up to max_weight)
    adaptive : bool
        Use adaptive measurement selection

    Returns
    -------
    TomographyStrategy
        Measurement plan with grouping and prioritization

    Examples
    --------
    >>> # 4-qubit tomography with weight-2 Paulis
    >>> strategy = ir_state_tomography(n_qubits=4, max_weight=2)
    >>> print(f"Measurements: {strategy.n_measurements}")
    >>> print(f"Compression: {strategy.compression_factor:.1f}×")
    """
    from .vqe_grouping import (
        check_group_commutativity,
        estimate_pauli_coherence_matrix,
        group_by_variance_minimization,
    )

    # Generate Pauli basis
    pauli_basis = generate_pauli_basis(n_qubits, max_weight=max_weight)

    n_paulis = len(pauli_basis)
    full_basis_size = 4**n_qubits

    # Estimate coherence (use uniform weights initially)
    weights = np.ones(n_paulis)
    Sigma = estimate_pauli_coherence_matrix(weights, pauli_basis)

    # Prioritize measurements by coherence structure
    # High-coherence measurements provide more information
    coherence_scores = np.sum(np.abs(Sigma), axis=1)
    measurement_order = np.argsort(-coherence_scores).tolist()

    # Group commuting measurements
    grouping = group_by_variance_minimization(
        Sigma, weights,
        max_group_size=min(10, n_paulis // 4),
        pauli_strings=pauli_basis,
        check_commutativity=True
    )

    # Determine final measurement count
    if target_measurements is not None:
        n_measurements = min(target_measurements, n_paulis)
    else:
        n_measurements = n_paulis

    compression_factor = full_basis_size / n_measurements

    return TomographyStrategy(
        measurement_basis=pauli_basis,
        measurement_order=measurement_order,
        grouping=grouping,
        n_measurements=n_measurements,
        compression_factor=compression_factor,
        method=f"ir_adaptive" if adaptive else "ir_static"
    )





def tomography_measurement_groups(
    strategy: TomographyStrategy
) -> List[List[str]]:
    """
    Get measurement groups from tomography strategy.

    Parameters
    ----------
    strategy : TomographyStrategy
        Tomography plan

    Returns
    -------
    List[List[str]]
        Grouped Pauli measurements

    Examples
    --------
    >>> strategy = ir_state_tomography(n_qubits=3, max_weight=2)
    >>> groups = tomography_measurement_groups(strategy)
    >>> for i, group in enumerate(groups[:3]):
    ...     print(f"Group {i}: {group}")
    """
    groups_paulis = []
    for group_indices in strategy.grouping:
        group_paulis = [strategy.measurement_basis[i] for i in group_indices]
        groups_paulis.append(group_paulis)

    return groups_paulis