Source code for quchip.approximations

"""Explicit strategies for reducing an authored Hamiltonian."""

from __future__ import annotations

from abc import ABC
from dataclasses import dataclass
from typing import Any, Iterable


def _freeze_bands(value: Iterable[tuple[int, ...]] | None) -> frozenset[tuple[int, ...]] | None:
    """Validate and freeze an optional replacement band selection."""
    if value is None:
        return None
    try:
        bands = frozenset(value)
    except TypeError as error:
        raise TypeError("keep_bands must be an iterable of integer tuples.") from error
    if not bands:
        raise ValueError("keep_bands must contain at least one integer tuple.")
    valid = all(
        isinstance(band, tuple)
        and band
        and all(not isinstance(number, bool) and isinstance(number, int) for number in band)
        for band in bands
    )
    if not valid:
        raise TypeError("keep_bands must be an iterable of non-empty integer tuples.")
    if any(tuple(-weight for weight in band) not in bands for band in bands):
        raise ValueError("keep_bands must include the adjoint of every band (all weights negated).")
    return bands




class Approximation(ABC):
    """Immutable engine strategy applied after the engine assembles the authored physics."""

    filters_terms: bool = False



    def keeps_operator_band(self, weights: tuple[int, ...]) -> bool:
        """Return whether a structural operator band survives reduction.

        Parameters
        ----------
        weights : tuple[int, ...]
            Energy-basis weight ``column - row`` per endpoint. Annihilation has
            weight ``+1`` and creation has weight ``-1``.
        """
        del weights
        return True




    def conserves_excitation_number(self) -> bool:
        """Return whether every retained band conserves the total energy-level index.

        Elimination uses this structural guarantee to keep its exact reduction
        and captured coordinate maps in total-excitation sectors.
        """
        return False




    def to_dict(self) -> dict[str, Any]:
        """Return the stable serialized strategy tag."""
        if type(self) is Exact:
            return {"type": "Exact"}
        if type(self) is RWA:
            data: dict[str, Any] = {"type": "RWA"}
            if self.keep_bands is not None:
                data["keep_bands"] = [list(band) for band in sorted(self.keep_bands)]
            return data
        raise TypeError(f"{type(self).__name__} must define its own stable approximation serialization.")




    @staticmethod
    def from_dict(data: Any) -> "Approximation":
        """Restore one explicit strategy and reject retired schemas.

        Parameters
        ----------
        data : mapping
            Serialized ``Exact`` or ``RWA`` record.
        """
        if not isinstance(data, dict):
            raise TypeError(
                "approximation must be a serialized Exact or RWA strategy, "
                f"got {type(data).__name__}."
            )
        strategy_type = data.get("type")
        if strategy_type == "Exact" and set(data) == {"type"}:
            return Exact()
        if strategy_type == "RWA" and set(data) <= {"type", "keep_bands"}:
            bands = data.get("keep_bands")
            return RWA(None if bands is None else {tuple(band) for band in bands})
        if strategy_type in {"Exact", "RWA"}:
            fields = sorted(set(data) - {"type", "keep_bands"})
            raise TypeError(f"Unsupported serialized approximation fields: {fields}")
        raise ValueError(f"Unknown approximation strategy {strategy_type!r}.")






@dataclass(frozen=True)
class Exact(Approximation):
    """Keep every term in the authored finite-dimensional Hamiltonian."""





@dataclass(frozen=True, init=False)
class RWA(Approximation):
    """First-order structural rotating-wave selection.

    The default keeps bands whose energy-level index weights sum to zero.
    This counts levels, not energy. Uneven spectra can have near-resonant
    transitions in discarded bands. Compare with ``Exact()`` or supply
    ``keep_bands``, which replaces the default selection.

    Parameters
    ----------
    keep_bands : iterable of tuple[int, ...] or None, default None
        Operator-band weights to keep. ``None`` keeps only bands whose weights
        sum to zero. Tuples follow the operator support order. Each tuple must
        be non-empty and integral, with its negated tuple also present.
        This adjoint closure preserves Hermiticity.

    References
    ----------
    Cohen-Tannoudji, Dupont-Roc & Grynberg, *Atom-Photon Interactions*
    (Wiley, 1992), Ch. IV.
    """

    keep_bands: frozenset[tuple[int, ...]] | None
    filters_terms = True

    def __init__(self, keep_bands: Iterable[tuple[int, ...]] | None = None) -> None:
        object.__setattr__(self, "keep_bands", _freeze_bands(keep_bands))



    def keeps_operator_band(self, weights: tuple[int, ...]) -> bool:
        """Return whether this RWA policy keeps ``weights``.

        Parameters
        ----------
        weights : tuple[int, ...]
            Energy-basis weight ``column - row`` per endpoint. Annihilation has
            weight ``+1`` and creation has weight ``-1``.
        """
        if self.keep_bands is not None:
            return weights in self.keep_bands
        return sum(weights) == 0




    def conserves_excitation_number(self) -> bool:
        """Return whether every retained band has zero total weight."""
        return self.keep_bands is None or all(sum(band) == 0 for band in self.keep_bands)




def require_approximation(value: Any) -> Approximation:
    """Validate a public approximation value without Boolean coercion."""
    if not isinstance(value, Approximation):
        raise TypeError(f"approximation must be an Exact or RWA strategy, got {type(value).__name__}.")
    return value


__all__ = ["Approximation", "Exact", "RWA"]