qpe_toolbox.circuit.qaoa
========================

.. py:module:: qpe_toolbox.circuit.qaoa


Functions
---------

.. autoapisummary::

   qpe_toolbox.circuit.qaoa.brute_force_maxcut
   qpe_toolbox.circuit.qaoa.generate_community_graph
   qpe_toolbox.circuit.qaoa.qaoa_energy
   qpe_toolbox.circuit.qaoa.study_optimization_time_costs
   qpe_toolbox.circuit.qaoa.compute_qaoa_contraction_costs


Module Contents
---------------

.. py:function:: brute_force_maxcut(graph_matrix, terms)

   Compute the Max-Cut value of a graph by brute-force enumeration.

   This function enumerates all possible non-trivial bipartitions of the
   vertex set and computes the corresponding cut value. It is intended
   for small graphs due to its exponential complexity.

   :param graph_matrix: Adjacency matrix of the graph. An edge between vertices ``i`` and
                        ``j`` is assumed to exist if ``G[i, j] == 1`` or ``G[j, i] == 1``.
   :type graph_matrix: array_like, shape (N, N)
   :param terms: Dictionary mapping edges to their weights. Keys must be tuples
                 ``(i, j)`` with ``i < j``, and values are the corresponding edge
                 weights.
   :type terms: dict

   :returns: * **maxcut** (*float*) -- Maximum cut value over all bipartitions.
             * **winners** (*list of tuple*) -- List of bipartitions (each given as a tuple of vertex indices)
               achieving the maximum cut value. Each tuple represents one side
               of the bipartition.

   .. rubric:: Notes

   - Bipartitions differing only by exchange of the two sets are treated
     as equivalent; only subsets of size ``1`` to ``N-1`` are enumerated.
   - The computational complexity scales as :math:`O(2^N)` and is only
     practical for small graphs.


.. py:function:: generate_community_graph(N, *, N_comm=4, rng=None)

   Generate a random graph with community (block) structure.

   The graph is generated using a stochastic block model, where vertices
   are divided into communities of random sizes and intra-community edges
   are more likely than inter-community edges.

   :param N: Total number of vertices in the graph.
   :type N: int
   :param N_comm: Target number of communities. The actual number of non-trivial
                  communities may be smaller if some generated sizes are <= 1.
                  Default is 4.
   :type N_comm: int, optional
   :param rng: Random number generator used to sample community sizes.
               If ``None``, a default generator is created.
   :type rng: :numpy-random:`numpy.random.Generator <generator>`, optional

   :returns: **G** -- A graph generated according to the stochastic block model.
   :rtype: :class:`networkx.Graph`

   .. rubric:: Notes

   - Community sizes are randomly generated and add up to ``N``.
   - The probability matrix is chosen such that:
       - Intra-community edge probability is 0.5.
       - Inter-community edge probability is ``0.5 / (N_comm - 1)``.


.. py:function:: qaoa_energy(x, terms, opt)

   Evaluate the QAOA energy for a given parameter vector.

   This function constructs a QAOA circuit for the Max-Cut Hamiltonian
   and computes the expectation value of the cost Hamiltonian using
   local expectation values.

   :param x: QAOA parameter vector of length ``2p``, where ``p`` is the number
             of QAOA layers. The first ``p`` entries correspond to the angles
             ``gamma``, and the last ``p`` entries correspond to the angles
             ``beta``.
   :type x: array_like
   :param terms: Dictionary mapping edges ``(i, j)`` to their weights in the
                 cost Hamiltonian.
   :type terms: dict
   :param opt: Optimizer object for tensor network contraction ordering. Should be
               an instance of ``cotengra``'s ``ReusableHyperOptimizer`` (or compatible
               optimizer). Reusing this object across multiple calls improves
               performance by reusing previously computed contraction strategies.
   :type opt: :cotengra-api:`ReusableHyperOptimizer`

   :returns: **energy** -- Negative expectation value of the cost Hamiltonian evaluated
             for the given QAOA parameters.
   :rtype: float

   .. rubric:: Notes

   - Expectation values are computed using
     ``circ.local_expectation`` with the JAX backend.
   - The returned value is real; any small imaginary part arising from
     numerical errors is discarded.


.. py:function:: study_optimization_time_costs(hamilt_terms, hyperopt, bounds, *, batch_size=5, num_iter=20, verbosity=0, seed=None)

   Measure time costs for batched parameter suggestions and evaluations in an ``optuna`` loop.

   This function runs a ``CMA-ES`` optimization study using ``optuna`` and records the
   time taken for three steps in each iteration:
   1. Asking the study for new parameter suggestions (``ask_time``)
   2. Evaluating the cost function (here, the QAOA energy) (``cost_time``)
   3. Updating the study with the results (``tell_time``)

   :param hamilt_terms: Dictionary representing the Hamiltonian to optimize, typically mapping
                        edges or terms to weights, compatible with the ``energy`` function.
   :type hamilt_terms: dict
   :param hyperopt: Optimizer object used in the ``energy`` function to control tensor network
                    contraction ordering. Reusing this object across multiple calls improves
                    performance.
   :type hyperopt: :cotengra-api:`ReusableHyperOptimizer`
   :param bounds: Bounds for each parameter in the optimization. Each tuple should be
                  ``(lower_bound, upper_bound)``.
   :type bounds: sequence of tuple of float
   :param batch_size: Number of parameter sets evaluated per iteration. Default is 5.
   :type batch_size: int, optional
   :param num_iter: Number of optimization iterations. Default is 20.
   :type num_iter: int, optional
   :param verbosity: Controls printing of intermediate results. If ``>= 1``, prints the
                     lowest energy in each batch. Default is 0.
   :type verbosity: int, optional
   :param seed: Random seed passed to the ``optuna`` sampler for reproducibility.
                Default is None.
   :type seed: int or None, optional

   :returns: * **ask_time** (*list of float*) -- List of times (in seconds) spent asking the study for new parameter suggestions
               in each iteration.
             * **tell_time** (*list of float*) -- List of times (in seconds) spent updating the study with evaluated energies
               in each iteration.
             * **cost_time** (*list of float*) -- List of times (in seconds per Hamiltonian term per parameter set) spent
               computing the energies in each iteration.
             * **study** (:optuna-api:`study.Study`) -- The ``optuna`` study object after all iterations. Can be used to query
               best parameters, best value, or continue optimization.

   .. rubric:: Notes

   - ``cost_time`` is normalized by the number of Hamiltonian terms and batch size
     to give an average per term per parameter set.


.. py:function:: compute_qaoa_contraction_costs(graph_dict, hyperopt, *, circuit_depths=(2, 3, 4), verbosity=0, description='None')

   Compute contraction widths and costs for QAOA circuits on multiple graphs.

   Generates QAOA circuits of varying depth for each graph and estimates the
   tensor network contraction width (W) and log-scaled contraction cost (C)
   using rehearsal contractions.

   :param graph_dict: Dictionary mapping graph identifiers to entries. Each entry must have at
                      least a key ``"terms"`` listing the edges as pairs of vertices.
   :type graph_dict: dict
   :param hyperopt: Optimizer object used for tensor network contraction rehearsal. Reusing this
                    object across multiple graphs improves performance.
   :type hyperopt: :cotengra-api:`ReusableHyperOptimizer`
   :param circuit_depths: List of QAOA circuit depths (number of layers) to analyze. Default is (2, 3, 4).
   :type circuit_depths: sequence of int, optional
   :param verbosity: Level of output. If >= 1, prints progress for each graph and each depth.
                     Default is 0.
   :type verbosity: int, optional
   :param description: Text description of this hyperoptimizer run (e.g., "greedy", "random").
                       Default is "None".
   :type description: str, optional

   :returns: Updated copy of ``graph_dict`` with QAOA contraction metrics. For each
             graph, a new numbered entry is added under ``"hyperoptimizers"``,
             containing contraction width ``W``, cost ``C`` for each depth ``p``,
             and the given description. For example::

                 "0": {
                     "terms": [[0, 1], [0, 2]],
                     "N": 3,
                     "E": 2,
                     "hyperoptimizers": {
                         "1": {
                             "p=2": {"W": 2.3509, "C": 2.5964},
                             "p=3": {"W": 2.3772, "C": 2.9702},
                             "description": "greedy"
                         }
                     }
                 }
   :rtype: dict

   .. rubric:: Notes

   - For each graph, a random QAOA parameter initialization is used (``gammas`` and ``betas``).
   - Contraction rehearsal is performed using :quimb-api:`Circuit.local_expectation_rehearse` for each
     term in the Hamiltonian.
   - The average width ``W`` is computed over all local contraction trees.
   - The total contraction cost ``C`` is computed using a numerically stable log-sum-exp
     over all local contraction costs.