embedding.py

'''
Author: Edoardo Pasetto
'''

import numpy as np
from dwave.system import DWaveSampler
from dwave_networkx import pegasus_graph
from dwave.embedding.pegasus import find_clique_embedding
from dwave.embedding import is_valid_embedding
import networkx as nx
from minorminer import find_embedding

#%%
embedding_path=''    # insert the path where the embeddings are going to be stored 
embedding_ind_path=embedding_path+'independent_graph\\'

#%%

def get_clique_emebedding(dim):
    '''
    searches for an embedding of a clique of dimension specified by dim
    '''

    var_list=list(range(dim))


    sampler=DWaveSampler()
    topology=sampler.properties['topology']
    active_qubits=sampler.properties['qubits']
    active_couplers=sampler.properties['couplers']

    new_pegasus=pegasus_graph(topology['shape'][0], node_list=active_qubits, edge_list=active_couplers)
    embedding=find_clique_embedding(k=var_list, target_graph= new_pegasus)

    return embedding


def load_embedding(path):
    '''
    loads an embedding from a .txt file saved using the function save_embedding
    returns the embedding as dict that can then be passed to an EmbeddingComposite
    '''
    r=open(path)
    lines=r.readlines()

    r.close()

    retrieved_embedding={}
    for l in range(len(lines)):
        data=np.fromstring(lines[l], sep=',').astype('int32')
        data=list(data)
        key=data[0]
        value=data[1:]

        retrieved_embedding[key]=value



    return retrieved_embedding





def check_clique_embedding(emb):
    '''
    checks if a given clique embedding is valid
    works only for cliques, for a more general case see check_embedding
    '''

    sampler=DWaveSampler()
    topology=sampler.properties['topology']
    active_qubits=sampler.properties['qubits']
    active_couplers=sampler.properties['couplers']

    new_pegasus=pegasus_graph(topology['shape'][0], node_list=active_qubits, edge_list=active_couplers)

    var_list=list(emb.keys())
    complete_graph= nx.complete_graph(var_list)

    res=is_valid_embedding(emb, complete_graph, new_pegasus)


    return res


def check_embedding(embedding,bqm):
    '''
    checks if an embedding is valid for a given problem, which is defined by a bqm
    it generates a graph from bqm and then checks if the graph can be minor-embedded in the QPU using the given embedding
    '''
    sampler=DWaveSampler()
    topology=sampler.properties['topology']
    active_qubits=sampler.properties['qubits']
    active_couplers=sampler.properties['couplers']

    new_pegasus=pegasus_graph(topology['shape'][0], node_list=active_qubits, edge_list=active_couplers)

    var_list=list(bqm.variables)

    G=nx.Graph()
    G.add_nodes_from(var_list)
    G.add_edges_from(list(bqm.quadratic.keys()))

    res=is_valid_embedding(embedding, G, new_pegasus)


    return res




def save_embedding(emb, path):
    '''
    saves the embedding as .txt file in the given path:
        for each row the first number is the logical qubit, while the other numbers in the row are the physical qubits used to create the chain
    '''


    f=open(path+'embedding.txt','w+')
    for k in emb.keys():
        f.write(str(k))

        for val in emb[k]:
            f.write(',')
            f.write(str(val))

        f.write('\n')
    f.close()





#%%


def find_independents_cliques_embedding(N,q):
    '''
    Parameters
    ----------
    N : int
        number of independents cliques with no
        shared variables (in the QKS code this corresponds to the number of episodes E)
    q : int
        dimension of the single clique

    Returns
    embedding on the pegasus graph

    '''
    sampler=DWaveSampler()
    topology=sampler.properties['topology']
    active_qubits=sampler.properties['qubits']
    active_couplers=sampler.properties['couplers']

    new_pegasus=pegasus_graph(topology['shape'][0], node_list=active_qubits, edge_list=active_couplers)

    nodes_table=np.zeros((N,q))
    for i in range(len(nodes_table)):
        for j in range(nodes_table.shape[1]):
            nodes_table[i,j]=q*i + j

    nodes_table.astype('int32')


    nodes_dict={i:list(nodes_table[i,:]) for i in range(nodes_table.shape[0])}

    G=nx.Graph()

    for k in nodes_dict.keys():
        #iterate through the independent sets of nodes/variables
        c_nodes=nodes_dict[k]
        G.add_nodes_from(c_nodes)

        edges_list=[]

        for i in range(len(c_nodes)):
            for j in range(i+1,len(c_nodes)):
                edges_list.append((c_nodes[i],c_nodes[j]))

        G.add_edges_from(edges_list)

    emb=find_embedding(S=G, T=new_pegasus)

    return emb, nodes_dict


#%%