Source code for psipose.feature_maps.amplitude

"""Amplitude encoding: encode a vector as quantum state amplitudes.

Amplitude encoding maps a classical feature vector to the amplitudes
of a quantum state. This is a powerful encoding that can represent
exponentially many features with only n qubits (2^n amplitudes).
"""

import numpy as np
import pennylane as qml

from .base import FeatureMap




class AmplitudeFeatureMap(FeatureMap):
    """Amplitude encoding of a classical vector into quantum amplitudes.

    Maps a classical feature vector :math:`\\mathbf{x}` to the amplitudes
    of a quantum state :math:`|\\phi(\\mathbf{x})\\rangle`. Requires
    :math:`n_{\\text{qubits}}` such that
    :math:`2^{n_{\\text{qubits}}} \\ge n_{\\text{features}}`.

    This implements the state preparation method from Möttönen et al. (2004)
    "Quantum circuit for preparing arbitrary states" (arXiv:quant-ph/0407010).
    PennyLane's :func:`qml.AmplitudeEmbedding` template uses this algorithm.

    Mathematically, for a normalized feature vector
    :math:`\\mathbf{x} = (x_1, \\dots, x_d)` with :math:`\\|\\mathbf{x}\\| = 1`,
    the encoded state is:

    .. math::
        |\\phi(\\mathbf{x})\\rangle = \\sum_{i=1}^{d} x_i |i-1\\rangle

    where :math:`d = 2^{n_{\\text{qubits}}}`. If the input dimension
    is not a power of 2, it is padded with zeros to the next power of 2.

    Parameters
    ----------
    n_qubits : int, optional
        Number of qubits to use. If None, the smallest power of 2
        that can accommodate the feature dimension is used.
    normalize : bool, default=True
        Whether to normalize the input vector before encoding.

    Attributes
    ----------
    n_features_ : int
        The original number of features (set after fit).

    Examples
    --------
    >>> feature_map = AmplitudeFeatureMap()
    >>> feature_map.fit(X_train)  # determines n_qubits_ based on n_features
    >>> # X_train.shape[1] must be <= 2**n_qubits
    """

    def __init__(self, n_qubits=None, normalize=True):
        super().__init__(n_qubits)
        self.normalize = normalize
        self.n_features_ = None



    def fit(self, X):
        """Determine the number of qubits needed.

        For amplitude encoding, we need
        :math:`2^{n_{\\text{qubits}}} \\ge n_{\\text{features}}`.
        If n_qubits is not specified, choose the smallest that works.
        """
        n_features = X.shape[1]
        self.n_features_ = n_features

        if self.n_qubits is None:
            # Find smallest n such that 2^n >= n_features
            self.n_qubits_ = int(np.ceil(np.log2(max(1, n_features))))
        else:
            if 2**self.n_qubits < n_features:
                raise ValueError(
                    f"n_qubits={self.n_qubits} is too small for {n_features} features. "
                    f"Need at least {int(np.ceil(np.log2(n_features)))} qubits."
                )
            self.n_qubits_ = self.n_qubits

        return self




    def encode(self, x, wires):
        """Apply amplitude encoding to the given wires.

        Pads or truncates the input to match
        :math:`2^{\\text{len(wires)}}`.

        Note: ``normalize=False`` is passed to ``qml.AmplitudeEmbedding``
        because normalization is handled manually above (using ``self.normalize``)
        to support the zero-vector edge case.
        """
        n_wires = len(wires)
        dim = 2**n_wires

        # Prepare the amplitude vector
        x_padded = np.zeros(dim)
        n_to_use = min(len(x), dim)
        x_padded[:n_to_use] = x[:n_to_use]

        if self.normalize:
            norm = np.linalg.norm(x_padded)
            if norm > 0:
                x_padded = x_padded / norm
            else:
                x_padded[0] = 1.0  # |0> state for zero vector

        qml.AmplitudeEmbedding(x_padded, wires=wires, normalize=False, pad_with=0.0)