qaoa_circuit_portfolio.py

###############################################################################
# // SPDX-License-Identifier: Apache-2.0
# // Copyright : JP Morgan Chase & Co
###############################################################################
import qiskit
import numpy as np
from .portfolio_optimization import yield_all_indices_cosntrained, get_configuration_cost
from qiskit import QuantumCircuit, execute, Aer, QuantumRegister
from qiskit.circuit import ParameterVector
from .utils import reverse_array_index_bit_order, state_to_ampl_counts


def generate_dicke_state_fast(N, K):
    """
    Generate the dicke state with yield function
    """
    index = yield_all_indices_cosntrained(N, K)
    s = np.zeros(2**N)
    for i in index:
        s[i] = 1
    s = 1 / np.sqrt(np.sum(s)) * s
    return s


def get_cost_circuit(po_problem, qc, gamma):
    """
    Construct the problem Hamiltonian layer for QAOA circuit
    H = 0.5*q\sum_{i=1}^{n-1} \sum_{j=i+1}^n \sigma_{ij}Z_i Z_j + 0.5 \sum_i (-q\sum_{j=1}^n{\sigma_ij} + \mu_i) Z_i + Constant
    """
    q = po_problem["q"]
    means = po_problem["means"]
    cov = po_problem["cov"]
    N = po_problem["N"]
    for i in range(N):
        qc.rz((means[i] - q * np.sum(cov[i, :])) * gamma, i)  # there is a 0.5 inside rz and rzz
    for i in range(N - 1):
        for j in range(i + 1, N):
            qc.rzz(q * cov[i, j] * gamma, i, j)
    return qc


def get_dicke_init(N, K):
    """
    Generate dicke state in gates
    """
    from qokit.dicke_state_utils import dicke_simple

    # can be other dicke state implementaitons
    return dicke_simple(N, K)


def get_mixer_Txy(qc, beta, minus=False, T=None):
    """
    H_{even} = \sum_{i is even} (X_i X_{i+1} + Y_i Y_{i+1})
    H_{odd} = \sum_{i is odd} (X_i X_{i+1} + Y_i Y_{i+1})
    H_{last} = X_{end} X_0 + Y_{end} Y_0
    U = {exp[-i*angle*H_{even}] exp[-i*angle*H_{odd}] exp[-i*angle*H_{last}]} ^ T  #repeat T times
    """
    if minus == True:
        beta = -beta
    N = len(qc._qubits)
    if T is None:
        T = 1
    beta = beta / T
    for _ in range(int(T)):
        for i in range(0, N - 1, 2):
            # even Exp[-j*angle*(XX+YY)]
            # qc.append(qiskit.circuit.library.XXPlusYYGate(4 * beta), [i, i + 1])
            qc.append(qiskit.circuit.library.XXPlusYYGate(4 * beta), [N - 2 - i, N - 1 - i])
        for i in range(1, N - 1, 2):
            # odd Exp[-j*angle*(XX+YY)]
            # qc.append(qiskit.circuit.library.XXPlusYYGate(4 * beta), [i, i + 1])
            qc.append(qiskit.circuit.library.XXPlusYYGate(4 * beta), [N - 2 - i, N - 1 - i])
        # last uniary
        qc.append(qiskit.circuit.library.XXPlusYYGate(4 * beta), [N - 1, 0])
    return qc


def get_mixer_RX(qc, beta):
    """A layer of RX gates"""
    N = len(qc._qubits)
    for i in range(N):
        qc.rx(2 * beta, i)
    return qc


def get_qaoa_circuit(
    po_problem,
    gammas,
    betas,
    depth,
    ini="dicke",
    mixer="trotter_ring",
    T=1,
    ini_state=None,
    save_state=True,
    minus=False,
):
    """
    Put all ingredients together to build up a qaoa circuit
    Minus is for define mixer with a minus sign, for checking phase diagram

    po_problem is generated by qokit.portfolio_optimization.get_problem
    """
    N = po_problem["N"]
    K = po_problem["K"]
    if ini_state is not None:
        q = QuantumRegister(N)
        circuit = QuantumCircuit(q)
        circuit.initialize(ini_state, [q[i] for i in range(N)])
    else:
        if ini.lower() == "dicke":
            circuit = get_dicke_init(N, K)
        elif ini.lower() == "uniform":
            circuit = get_uniform_init(N)
        else:
            raise ValueError("Undefined initial circuit")
    for i in range(depth):
        circuit = get_cost_circuit(po_problem, circuit, gammas[i])
        if mixer.lower() == "trotter_ring":
            circuit = get_mixer_Txy(circuit, betas[i], minus=minus, T=T)  # minus should be false
        elif mixer.lower() == "rx":
            circuit = get_mixer_RX(circuit, betas[i])
        else:
            raise ValueError("Undefined mixer circuit")
    if save_state is False:
        circuit.measure_all()
    return circuit


def get_parameterized_qaoa_circuit(
    po_problem, depth, ini="dicke", mixer="trotter_ring", T=1, ini_state=None, save_state=True, minus=False, return_parameter_vectors: bool = False
):
    """
    Put all ingredients together to build up a qaoa circuit parameterized by gamma & beta
    Minus is for define mixer with a minus sign, for checking phase diagram
    """
    N = po_problem["N"]
    K = po_problem["K"]
    if ini_state is not None:
        q = QuantumRegister(N)
        circuit = QuantumCircuit(q)
        circuit.initialize(ini_state, [q[i] for i in range(N)])
    else:
        if ini.lower() == "dicke":
            circuit = get_dicke_init(N, K)
        elif ini.lower() == "uniform":
            circuit = get_uniform_init(N)
        else:
            raise ValueError("Undefined initial circuit")

    betas = ParameterVector("beta", depth)
    gammas = ParameterVector("gamma", depth)

    for i in range(depth):
        circuit = get_cost_circuit(po_problem, circuit, gammas[i])
        if mixer.lower() == "trotter_ring":
            circuit = get_mixer_Txy(circuit, betas[i], minus=minus, T=T)  # minus should be false
        elif mixer.lower() == "rx":
            circuit = get_mixer_RX(circuit, betas[i])
        else:
            raise ValueError("Undefined mixer circuit")
    if save_state is False:
        circuit.measure_all()
    if return_parameter_vectors:
        return circuit, betas, gammas
    else:
        return circuit


def get_energy_expectation(po_problem, samples):
    """Compute energy expectation from measurement samples"""
    expectation_value = 0
    N_total = 0
    for config, count in samples.items():
        expectation_value += count * get_configuration_cost(po_problem, config)
        N_total += count
    expectation_value = expectation_value / N_total

    return expectation_value


def get_energy_expectation_sv(po_problem, samples):
    """Compute energy expectation from full state vector"""
    expectation_value = 0
    # convert state vector to dictionary
    samples = state_to_ampl_counts(samples)
    for config, wf in samples.items():
        expectation_value += (np.abs(wf) ** 2) * get_configuration_cost(po_problem, config)

    return expectation_value


def invert_counts(counts):
    """convert qubit order for measurement samples"""
    return {k[::-1]: v for k, v in counts.items()}


def measure_circuit(circuit, n_trials=1024, save_state=True):
    """Get the output from circuit, either measured samples or full state vector"""
    if save_state is False:
        backend = Aer.get_backend("qasm_simulator")
        job = execute(circuit, backend, shots=n_trials)
        result = job.result()
        bitstrings = invert_counts(result.get_counts())
        return bitstrings
    else:
        backend = Aer.get_backend("statevector_simulator")
        result = execute(circuit, backend).result()
        state = result.get_statevector()
        return reverse_array_index_bit_order(state)


def circuit_measurement_function(
    po_problem,
    p,
    ini="dicke",
    mixer="trotter_ring",
    T=None,
    ini_state=None,
    n_trials=1024,
    save_state=True,
    minus=False,
):
    """Helper function to define the objective function to optimize"""

    def f(x):
        gammas = x[0:p]
        betas = x[p:]
        circuit = get_qaoa_circuit(
            po_problem,
            ini=ini,
            mixer=mixer,
            T=T,
            ini_state=ini_state,
            gammas=gammas,
            betas=betas,
            depth=p,
            save_state=save_state,
            minus=minus,
        )
        samples = measure_circuit(circuit, n_trials=n_trials, save_state=save_state)
        if save_state is False:
            energy_expectation_value = get_energy_expectation(po_problem, samples)
        else:
            energy_expectation_value = get_energy_expectation_sv(po_problem, samples)
        return energy_expectation_value

    return f