Source code for psipose.ansatze._base

"""Base ansatz class for parameterized quantum circuits."""

import abc

import numpy as np




class BaseAnsatz(abc.ABC):
    """Abstract base class for quantum ansatze.

    An ansatz is a parameterized quantum circuit that can be trained
    by adjusting its parameters (weights).

    Parameters
    ----------
    n_qubits : int
        Number of qubits in the circuit.
    """

    def __init__(self, n_qubits):
        self.n_qubits = n_qubits
        self.n_params = None



    @abc.abstractmethod
    def circuit(self, weights, wires):
        """Build the parameterized quantum circuit.

        Parameters
        ----------
        weights : array-like
            Trainable parameters for the ansatz. Shape depends on the ansatz.
        wires : Iterable[int]
            Qubit wires to apply the circuit to.
        """
        pass




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

        Parameters
        ----------
        n_qubits : int
            Number of qubits.

        Returns
        -------
        tuple[int, ...] or int
            Shape of the weight array. Can be a single integer (total params)
            or a tuple (e.g., (n_layers, n_qubits, 3) for layered ansatze).
        """
        raise NotImplementedError("Subclasses must implement weight_shape")




    def initialize_weights(self, seed=None):
        """Initialize random weights.

        Parameters
        ----------
        seed : int, optional
            Random seed for reproducibility.

        Returns
        -------
        ndarray
            Randomly initialized weights with the correct shape.
        """
        rng = np.random.default_rng(seed)
        shape = self.weight_shape(self.n_qubits)
        return rng.uniform(0, 2 * np.pi, shape)