Example-02-step-3-Solution.py

# SPDX-License-Identifier: Apache-2.0 AND CC-BY-NC-4.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Defining functions to generate the Hamiltonian and Kernel for a given graph
# Necessary packages
import networkx as nx
from networkx import algorithms
from networkx.algorithms import community
import cudaq
from cudaq import spin
from cudaq.qis import *
import numpy as np
from typing import List, Tuple
from mpi4py import MPI


# Getting information about platform
cudaq.set_target("nvidia")
target = cudaq.get_target()

# Setting up MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
num_qpus = comm.Get_size()


#######################################################
# Step 1
####################################################### 
# Function to return a dictionary of subgraphs of the input graph 
# using the greedy modularity maximization algorithm

def subgraphpartition(G,n):
    """Divide the graph up into at most n subgraphs
    
    Parameters
    ----------
    G: networkX.Graph 
        Graph that we want to subdivdie
    n : int
        n is the maximum number of subgraphs in the partition
    Returns
    -------
    dict of str : networkX.Graph
        Dictionary of networkX graphs with a string as the key
    """
    greedy_partition = community.greedy_modularity_communities(G, weight=None, resolution=1.1, cutoff=1, best_n=n)
    number_of_subgraphs = len(greedy_partition)

    graph_dictionary = {}
    graph_names=[]
    for i in range(number_of_subgraphs):
        name='G'+str(i)
        graph_names.append(name)

    for i in range(number_of_subgraphs):
        nodelist = sorted(list(greedy_partition[i]))
        graph_dictionary[graph_names[i]] = nx.subgraph(G, nodelist)
     
    return(graph_dictionary) 


if rank ==0:
    # Defining the example graph
    # Random graph parameters
    n = 35  # numnber of nodes
    m = 80  # number of edges
    seed = 20160  # seed random number generators for reproducibility

    # Use seed for reproducibility
    sampleGraph2 = nx.gnm_random_graph(n, m, seed=seed)

    # Subdividing the graph
    num_subgraphs_limit = min(12, len(sampleGraph2.nodes())) # maximum number of subgraphs for the partition
    subgraph_dictionary = subgraphpartition(sampleGraph2,num_subgraphs_limit)

    # Assign the subgraphs to the QPUs
    number_of_subgraphs = len(sorted(subgraph_dictionary))
    number_of_subgraphs_per_qpu = int(np.ceil(number_of_subgraphs/num_qpus))

    keys_on_qpu ={}
    
    for q in range(num_qpus):
        keys_on_qpu[q]=[]
        for k in range(number_of_subgraphs_per_qpu):
            if (k*num_qpus+q < number_of_subgraphs):
                key = sorted(subgraph_dictionary)[k*num_qpus+q]
                keys_on_qpu[q].append(key)        
    print('Subgraph problems to be computed on each processor have been assigned')
        

    # Distribute the subgraph data to the QPUs
    for i in range(num_qpus):
        subgraph_to_qpu ={}
        for k in keys_on_qpu[i]:
            subgraph_to_qpu[k]= subgraph_dictionary[k]
        if i != 0:
            comm.send(subgraph_to_qpu, dest=i, tag=rank)
        else:
            assigned_subgraph_dictionary = subgraph_to_qpu
else:
    # Receive the subgraph data
    assigned_subgraph_dictionary= comm.recv(source=0, tag=0)
    print("Processor {} received {} from processor {}".format(rank,assigned_subgraph_dictionary, 0))



#######################################################
# Step 2
####################################################### 

# Define a function to generate the Hamiltonian for a max cut problem using the graph G

def hamiltonian_max_cut(sources : List[int], targets : List[int]): 
    """Hamiltonian for finding the max cut for the graph  with edges defined by the pairs generated by source and target edges
        
    Parameters
    ----------
    sources: List[int] 
        list of the source vertices for edges in the graph
    targets: List[int]
        list of the target vertices for the edges in the graph

    Returns
    -------
    cudaq.SpinOperator
        Hamiltonian for finding the max cut of the graph defined by the given edges
    """
    hamiltonian = 0
    # Since our vertices may not be a list from 0 to n, or may not even be integers,
    
    for i in range(len(sources)):
        # Add a term to the Hamiltonian for the edge (u,v)
        qubitu = sources[i]
        qubitv = targets[i]
        hamiltonian += 0.5*(spin.z(qubitu)*spin.z(qubitv)-spin.i(qubitu)*spin.i(qubitv))
    
    return hamiltonian

# Problem Kernel

@cudaq.kernel
def qaoaProblem(qubit_0 : cudaq.qubit, qubit_1 : cudaq.qubit, alpha : float):
    """Build the QAOA gate sequence between two qubits that represent an edge of the graph
    Parameters
    ----------
    qubit_0: cudaq.qubit 
        Qubit representing the first vertex of an edge
    qubit_1: cudaq.qubit 
        Qubit representing the second vertex of an edge
    alpha: float
        Free variable   
        

    """
    x.ctrl(qubit_0, qubit_1)
    rz(2.0*alpha, qubit_1)
    x.ctrl(qubit_0, qubit_1)

# Mixer Kernel
@cudaq.kernel
def qaoaMixer(qubit_0 : cudaq.qubit, beta : float):
    """Build the QAOA gate sequence that is applied to each qubit in the mixer portion of the circuit
    Parameters
    ----------
    qubit_0: cudaq.qubit 
        Qubit 
    beta: float
        Free variable   
        

    """
    rx(2.0*beta, qubit_0)


# We now define the kernel_qaoa function which will be the QAOA circuit for our graph
# Since the QAOA circuit for max cut depends on the structure of the graph, 
# we'll feed in global concrete variable values into the kernel_qaoa function for the qubit_count, layer_count, edges_src, edges_tgt.
# The types for these variables are restricted to Quake Values (e.g. qubit, int, List[int], ...)
# The thetas plaeholder will be our free parameters (the alphas and betas in the circuit diagrams depicted above)
@cudaq.kernel
def kernel_qaoa(qubit_count :int, layer_count: int, edges_src: List[int], edges_tgt: List[int], thetas : List[float]):
    """Build the QAOA circuit for max cut of the graph with given edges and nodes
    Parameters
    ----------
    qubit_count: int 
        Number of qubits in the circuit, which is the same as the number of nodes in our graph
    layer_count : int 
        Number of layers in the QAOA kernel
    edges_src: List[int]
        List of the first (source) node listed in each edge of the graph, when the edges of the graph are listed as pairs of nodes
    edges_tgt: List[int]
        List of the second (target) node listed in each edge of the graph, when the edges of the graph are listed as pairs of nodes
    thetas: List[float]
        Free variables to be optimized   
        

    """
    # Let's allocate the qubits
    qreg = cudaq.qvector(qubit_count)
    
    # And then place the qubits in superposition
    h(qreg)
    
    # Each layer has two components: the problem kernel and the mixer
    for i in range(layer_count):
        # Add the problem kernel to each layer
        for edge in range(len(edges_src)):
            qubitu = edges_src[edge]
            qubitv = edges_tgt[edge]
            qaoaProblem(qreg[qubitu], qreg[qubitv], thetas[i])
        # Add the mixer kernel to each layer
        for j in range(qubit_count):
            qaoaMixer(qreg[j],thetas[i+layer_count])
            
def find_optimal_parameters(G, layer_count, seed):
    """Function for finding the optimal parameters of QAOA for the max cut of a graph
    Parameters
    ----------
    G: networkX graph 
        Problem graph whose max cut we aim to find
    layer_count : int 
        Number of layers in the QAOA circuit
    seed : int
        Random seed for reproducibility of results
        
    Returns
    -------
    list[float]
        Optimal parameters for the QAOA applied to the given graph G
    """
    parameter_count: int = 2 * layer_count

    # Problem parameters
    nodes = sorted(list(nx.nodes(G)))
    qubit_src = []
    qubit_tgt = []
    for u, v in nx.edges(G):
        # We can use the index() command to read out the qubits associated with the vertex u and v.
        qubit_src.append(nodes.index(u))
        qubit_tgt.append(nodes.index(v))
    # The number of qubits we'll need is the same as the number of vertices in our graph
    qubit_count : int = len(nodes)
    # Each layer of the QAOA kernel contains 2 parameters
    parameter_count : int = 2*layer_count
    
    # Specify the optimizer and its initial parameters. 
    optimizer = cudaq.optimizers.COBYLA()
    np.random.seed(seed)
    optimizer.initial_parameters = np.random.uniform(-np.pi, np.pi,
                                                     parameter_count)   

    # Pass the kernel, spin operator, and optimizer to `cudaq.vqe`.
    optimal_expectation, optimal_parameters = cudaq.vqe(
        kernel=kernel_qaoa,
        spin_operator=hamiltonian_max_cut(qubit_src, qubit_tgt),
        argument_mapper=lambda parameter_vector: (qubit_count, layer_count, qubit_src, qubit_tgt, parameter_vector),
        optimizer=optimizer,
        parameter_count=parameter_count)

    return optimal_parameters
def qaoa_for_graph(G, layer_count, shots, seed):
    """Function for finding the max cut of a graph using QAOA
    
    Parameters
    ----------
    G: networkX graph 
        Problem graph whose max cut we aim to find
    layer_count : int 
        Number of layers in the QAOA circuit
    shots : int
        Number of shots in the sampling subroutine
    seed : int
        Random seed for reproducibility of results

    Returns
    -------
    str
        Binary string representing the max cut coloring of the vertinces of the graph
    """

    
    parameter_count: int = 2 * layer_count

    # Problem parameters
    nodes = sorted(list(nx.nodes(G)))
    qubit_src = []
    qubit_tgt = []
    for u, v in nx.edges(G):
        # We can use the index() command to read out the qubits associated with the vertex u and v.
        qubit_src.append(nodes.index(u))
        qubit_tgt.append(nodes.index(v))
    # The number of qubits we'll need is the same as the number of vertices in our graph
    qubit_count : int = len(nodes)
    # Each layer of the QAOA kernel contains 2 parameters
    parameter_count : int = 2*layer_count
    
    optimal_parameters = find_optimal_parameters(G, layer_count, seed)

    # Print the optimized parameters
    print("Optimal parameters = ", optimal_parameters)
    
    # Sample the circuit
    counts = cudaq.sample(kernel_qaoa, qubit_count, layer_count, qubit_src, qubit_tgt, optimal_parameters, shots_count=shots)
    print('most_probable outcome = ',counts.most_probable())
    results = str(counts.most_probable())
    return results


############################################################################
# On GPU with rank r, compute the subgraph solutions for the 
# subgraphs in assigned_subgraph_dictionary that live on GPU r
############################################################################
layer_count =1
results = {}
new_seed_for_each_graph = rank # to give each subgraph solution different initial parameters
for key in assigned_subgraph_dictionary:
    G = assigned_subgraph_dictionary[key]
    results[key] = qaoa_for_graph(G, layer_count, shots = 10000, seed=6543+new_seed_for_each_graph)
    new_seed_for_each_graph+=1
    print('The max cut QAOA coloring for the subgraph',key,'is',results[key])
print('The results dictionary variable on GPU',rank,'is',results)

############################################################################
# Exercise to copy over the subgraph solutions from the individual GPUs
# back to GPU 0.
#############################################################################
# Let's introduce another MPI function that will be useful to
# iterate over all the GPUs. 
size = comm.Get_size()

# Copy the results over to QPU 0 for consolidation 
if rank!=0:
    comm.send(results, dest=0, tag=0)
    #print("{} sent by processor {}".format(results, rank))
else:
    for j in range(1,size,1):
        colors = comm.recv(source=j, tag=0)
        print("Received {} from processor {}".format(colors, j))
        for key in colors:
            results[key]=colors[key]
    #print("The results dictionary on GPU 0 =", results)
 
 
    #######################################################
    # Step 3
    ####################################################### 
    ############################################################################
    # Merge results on QPU 0
    ############################################################################  
    
    # Add color attribute to subgraphs and sampleGraph2 to record the subgraph solutions
    # Plot sampleGraph2 with node colors inherited from the subgraph solutions

    subgraphColors={}

    for key in subgraph_dictionary:
        subgraphColors[key]=[int(i) for i in results[key]]
    

    for key in subgraph_dictionary:
        G = subgraph_dictionary[key]
        for v in sorted(list(nx.nodes(G))):
            G.nodes[v]['color']=subgraphColors[key][sorted(list(nx.nodes(G))).index(v)]
            sampleGraph2.nodes[v]['color']=G.nodes[v]['color']
    
    
    # A function that takes as input a subgraph partition (in the form of a graph dictionary) and a vertex.
    # The function should return the key associated with the subgraph that contains the given vertex.

    def subgraph_of_vertex(graph_dictionary, vertex):
        """
        A function that takes as input a subgraph partition (in the form of a graph dictionary) and a vertex.
        The function should return the key associated with the subgraph that contains the given vertex.
    
        Parameters
        ----------
        graph_dictionary: dict of networkX.Graph with str as keys 
        v : int
            v is a name for a vertex

        Returns
        -------
        str
            the key associated with the subgraph that contains the given vertex.
        """
        # in case a vertex does not appear in the graph_dictionary, return the empty string
        location = 'Vertex is not in the subgraph_dictionary'

        for key in graph_dictionary:
            if vertex in graph_dictionary[key].nodes():
                location = key
        return location

    
    # First let's define a function that constructs the border graph

    def border(G, subgraph_dictionary):
        """Build a graph made up of border vertices from the subgraph partition
    
        Parameters
        ----------
        G: networkX.Graph 
            Graph whose max cut we want to find
        subgraph_dictionary: dict of networkX graph with str as keys 
            Each graph in the dictionary should be a subgraph of G
        
        Returns
        -------
        networkX.Graph
            Subgraph of G made up of only the edges connecting subgraphs in the subgraph dictionary
        """   
        borderG = nx.Graph()
        for u,v in G.edges():
            border = True
            for key in subgraph_dictionary:
                SubG = subgraph_dictionary[key]
                edges = list(nx.edges(SubG))
                if (u,v) in edges:
                    border = False
            if border==True:
                borderG.add_edge(u,v)
        
        return borderG
    
    
    # Create the borderGraph
    borderGraph = border(sampleGraph2, subgraph_dictionary)
    
    
    # Define the Hamiltonian for applying QAOA to the variables
    # s_i where s_i = 1 means we will not flip the subgraph Gi's colors
    # and s_i = -1 means we will flip the colors of subgraph G_i

    def mHamiltonian(merger):
        """Hamiltonian for finding the optimal swap schedule for the subgraph partitioning encoded in the merger graph
        
        Parameters
        ----------
        merger: networkX.Graph 
            Weighted graph

        Returns
        -------
        cudaq.SpinOperator
            Hamiltonian for finding the optimal swap schedule for the subgraph partitioning encoded in the merger graph
        """  

        mergerHamiltonian = 0
        mergerNodes = sorted(list(merger.nodes()))
    # Add Hamiltonian terms for edges within a subgraph that contain a border element
        for u, v in merger.edges():
            qubitu = mergerNodes.index(u)
            qubitv = mergerNodes.index(v)
            mergerHamiltonian+= -1*merger[u][v]['penalty']*(spin.z(qubitu))*(spin.z(qubitv))
        return mergerHamiltonian
    
    
    # Define the mergerGraph and color code the vertices
    # according to the subgraph that the vertex represents
    def createMergerGraph(border, subgraphs):
        """Build a graph containing a vertex for each subgraph 
        and edges between vertices are added if there is an edge between
        the corresponding subgraphs
    
        Parameters
        ----------
        border: networkX.Graph 
            Graph of connections between vertices in distinct subgraphs
        subgraphs: dict of networkX graph with str as keys 
            The nodes of border should be a subset of the the graphs in the subgraphs dictionary
        
        Returns
        -------
        networkX.Graph
            Merger graph containing a vertex for each subgraph 
            and edges between vertices are added if there is an edge between
            the corresponding subgraphs
        """   

        H = nx.Graph()
    
        for u, v in border.edges():
            subgraph_id_for_u = subgraph_of_vertex(subgraphs, u)
            subgraph_id_for_v = subgraph_of_vertex(subgraphs, v)
            if subgraph_id_for_u != subgraph_id_for_v:
                H.add_edge(subgraph_id_for_u, subgraph_id_for_v)   
        return H


    mergerGraph = createMergerGraph(borderGraph, subgraph_dictionary)

    # Add attribute to capture the penalties of changing subgraph colors

    # Initialize all the penalties to 0
    nx.set_edge_attributes(mergerGraph, int(0), 'penalty')

    # Compute penalties for each edge
    for i, j in mergerGraph.edges():
        penalty_ij = 0
        for u in subgraph_dictionary[i]:
            for neighbor_u in nx.all_neighbors(sampleGraph2, u):
                if neighbor_u in subgraph_dictionary[j]:
                    if str(sampleGraph2.nodes[u]['color']) != str(sampleGraph2.nodes[neighbor_u]['color']):
                        penalty_ij += 1
                    else:
                        penalty_ij += -1
    mergerGraph[i][j]['penalty'] = penalty_ij


    # Graph the penalties of each edge
    edge_labels = nx.get_edge_attributes(mergerGraph, 'penalty')

    
    # Run QAOA on the merger subgraph to identify which subgraphs
    # if any should change colors

    layer_count_merger = 1 # set arbitrarily
    parameter_count_merger: int = 2 * layer_count_merger

    # Specify the optimizer and its initial parameters. Make it repeatable.
    cudaq.set_random_seed(101)
    optimizer_merger = cudaq.optimizers.COBYLA()
    np.random.seed(101)
    optimizer_merger.initial_parameters = np.random.uniform(-np.pi, np.pi,
                                                         parameter_count_merger)  
    optimizer_merger.max_iterations=150

    merger_nodes = list(mergerGraph.nodes())
    qubit_count = len(merger_nodes)
    merger_edge_src = []
    merger_edge_tgt = []
    for u, v in nx.edges(mergerGraph):
        # We can use the index() command to read out the qubits associated with the vertex u and v.
        merger_edge_src.append(merger_nodes.index(u))
        merger_edge_tgt.append(merger_nodes.index(v))

    # Pass the kernel, spin operator, and optimizer to `cudaq.vqe`.
    optimal_expectation, optimal_parameters = cudaq.vqe(
        kernel=kernel_qaoa,
        spin_operator=mHamiltonian(mergerGraph),
        argument_mapper=lambda parameter_vector: (qubit_count, layer_count, merger_edge_src, merger_edge_tgt, parameter_vector),
        optimizer=optimizer_merger,
        parameter_count=parameter_count_merger, 
        shots = 10000)

    # Print the optimized value and its parameters
    print("Optimal value = ", optimal_expectation)
    print("Optimal parameters = ", optimal_parameters)

    # Sample the circuit using the optimized parameters
    sample_number=15000
    counts = cudaq.sample(kernel_qaoa, qubit_count, layer_count, merger_edge_src, merger_edge_tgt, optimal_parameters, shots_count=5000)
    print(f"most_probable = {counts.most_probable()}")

    # Merger results
    mergerResultsString=counts.most_probable()

    ## Record a new coloring of the sampleGraph2 according to the merger results
    # with a node vertex attribute 'new_color'
    flipGraphColors={}

    mergerNodes = sorted(list(nx.nodes(mergerGraph)))
    for u in mergerNodes:
        indexu = mergerNodes.index(u)
        flipGraphColors[u]=int(mergerResultsString[indexu])
   
    for key in subgraph_dictionary:
        if flipGraphColors[key]==1:
            for u in subgraph_dictionary[key].nodes():
                sampleGraph2.nodes[u]['new_color'] = 1 - sampleGraph2.nodes[u]['color']
        else:
            for u in subgraph_dictionary[key].nodes():
                sampleGraph2.nodes[u]['new_color'] = sampleGraph2.nodes[u]['color']

    # Compute the new cut of the larger graph based on the new colors 
    max_cut = 0
    max_cut_edges = []
    for u, v in sampleGraph2.edges():
        if str(sampleGraph2.nodes[u]['new_color']) != str(sampleGraph2.nodes[v]['new_color']): 
                max_cut+=1
                max_cut_edges.append((u,v))
    print('The max cut value approximated from the Divide and Conquer QAOA is',max_cut)