Source code for psipose.ansatze.strongly_entangling

"""Strongly entangling ansatz using PennyLane's template."""

import pennylane as qml

from .base import BaseAnsatz




class StronglyEntanglingAnsatz(BaseAnsatz):
    """Strongly entangling ansatz based on PennyLane's template.

    Applies a strongly entangling layer structure where each qubit
    receives rotations around all three axes, followed by entangling
    gates between all neighboring qubits.

    This follows the circuit-centric quantum classifier approach from:

    - Schuld, Bocharov, Svore & Wiebe (2018), "Circuit-centric
      quantum classifiers" (arXiv:1804.00633). Introduces layered
      circuits with single-qubit rotations and entangling gates
      for variational quantum classification.
    - PennyLane's :func:`qml.StronglyEntanglingLayers` template,
      which implements this architecture directly.

    Mathematically, each layer :math:`l` applies:

    .. math::
        U_l(\\boldsymbol{\\theta}) = \\text{CNOT}_{\\text{chain}} \\,
        \\prod_{q=0}^{n-1} R_Z(\\theta_{l,q,2}) R_Y(\\theta_{l,q,1})
        R_X(\\theta_{l,q,0})

    where :math:`\\boldsymbol{\\theta}` has shape
    :math:`(\\text{layers}, n_{\\text{qubits}}, 3)`.

    Parameters
    ----------
    n_qubits : int
        Number of qubits.
    layers : int, default=2
        Number of repeating layers.

    Attributes
    ----------
    weight_shape_ : tuple[int, int]
        The shape of weights returned by weight_shape(): (layers, n_qubits, 3).

    Examples
    --------
    >>> ansatz = StronglyEntanglingAnsatz(n_qubits=4, layers=3)
    >>> weights = ansatz.initialize_weights(seed=42)
    >>> # weights.shape == (3, 4, 3)
    """

    def __init__(self, n_qubits, layers=2):
        super().__init__(n_qubits)
        self.layers = layers



    def circuit(self, weights, wires):
        """Build the quantum circuit using PennyLane's StronglyEntanglingLayers.

        Parameters
        ----------
        weights : array-like, shape (layers, n_qubits, 3)
            Rotation parameters for each layer and qubit. The last dimension
            contains [RX, RY, RZ] rotation angles in that order.
        wires : Iterable[int]
            Qubit wires.

        Notes
        -----
        This uses PennyLane's built-in StronglyEntanglingLayers template.
        The circuit structure follows the pattern described in Schuld et al. (2018):
        each layer consists of single-qubit rotations followed by nearest-neighbor
        CNOT entangling gates in a ring pattern.
        """
        qml.StronglyEntanglingLayers(weights, wires=wires)




    def weight_shape(self, n_qubits):
        """Return the shape of the weight array: (layers, n_qubits, 3)."""
        return (self.layers, n_qubits, 3)