qubo_lab/util/queries.py

from .kSAT import kSAT
import bson

class Instance_scope_query:
    
    def __init__(self, database):
        self.__database = database
        self.__query = None
        
    def query(self, scope):
        self.__query = self.__database["experiment_scopes"].aggregate([
            {
                "$match": {"_id": scope}
            },
            {
                "$unwind": "$instances"
            },
            {
                "$lookup":
                {
                    "from": "instances",
                    "localField": "instances",
                    "foreignField": "_id",
                    "as": "instance"
                }
            },
            {
                "$unwind": "$instance"
            },
            {
                "$replaceRoot": {"newRoot": "$instance"}
            }
        ])
    
    def __iter__(self):
        return self
    
    def __next__(self):
        document = self.__query.next()
        
        return self.__document_to_sat(document)
    
    def __document_to_sat(self, document):
        sat = kSAT()
        
        for clause in document["clauses"]:
            sat.addClause(clause);
        
        return sat, document["_id"]

class WMIS_scope_query_raw:
    
    def __init__(self, database):
        self.__database = database
        self.__query = None
    
    def query(self, scope):
        self.__query = self.__database["experiment_scopes"].aggregate([
            {
                "$match": {"_id": scope}
            },
            {
                "$unwind": "$instances"
            },
            {
                "$lookup":
                {
                    "from": "wmis_qubos",
                    "localField": "instances",
                    "foreignField": "instance",
                    "as": "qubo"
                }
            },
            {
                "$unwind": "$qubo"
            },
            {
                "$replaceRoot": {"newRoot": "$qubo"}
            }
        ])
            
    def __iter__(self):
        return self
    
    def __next__(self):
        return self.__query.next()
    
class WMIS_scope_query (WMIS_scope_query_raw):
    
    def __next__(self):
        doc = super(WMIS_scope_query, self).__next__()
        
        
        return read_raw_qubo(doc["qubo"]), doc["_id"]
    
class WMIS_solver_input_scope_query_raw:
    
    def __init__(self, database):
        self.__database = database
        self.__query = None
        
    def query(self, scope, solver_graph_id):
        self.__query = self.__database["experiment_scopes"].aggregate([
            {
                "$match": {"_id": scope}
            },
            {
                "$unwind": "$instances"
            },
            {
                "$project": {"instance_id": "$instances"}
            },
            {
                "$lookup":
                {
                    "from": "wmis_qubos",
                    "localField": "instance_id",  
                    "foreignField": "instance",
                    "as": "wmis_qubo"
                }
            },
            {
                "$unwind": "$wmis_qubo"
            },
            {
                "$lookup":
                {
                    "from": "embeddings",
                    "let":
                    {
                        "qubo_id": "$wmis_qubo._id",
                        "solver_graph_id": bson.ObjectId(solver_graph_id)
                    },
                    "pipeline":
                    [
                    {
                        "$match":
                        {
                            "$expr":
                            {
                                "$and":
                                [
                                {"$eq": ["$qubo", "$$qubo_id"]},
                                {"$eq": ["$solver_graph", "$$solver_graph_id"]}
                                ]
                            }
                        }
                    },
                    {
                        "$project": { "list": "$embeddings" }
                    }
                    ],
                    "as": "embeddings"
                }
            },
            {
                "$match": {"embeddings": {"$exists": True, "$not": {"$size": 0}}}
            },
            {
                "$unwind": "$embeddings"
            },
            {
                "$project":
                {
                    "_id": False,
                    "instance_id": True,
                    "wmis_qubo": True,
                    "embeddings": True
                }
            }
        ])
            
    def __iter__(self):
        return self
    
    def __next__(self):
        return self.__query.next()
    
class WMIS_solver_input_scope_query (WMIS_solver_input_scope_query_raw):
    
    def __next__(self):
        doc = super(WMIS_solver_input_scope_query, self).__next__()
        
        data = {}
        
        data["instance_id"] = doc["instance_id"]
        
        data["qubo_id"] = doc["wmis_qubo"]["_id"]
        data["qubo"] = read_raw_qubo(doc["wmis_qubo"]["qubo"])
        
        data["embeddings_id"] = doc["embeddings"]["_id"]
        
        data["embeddings"] = []
        for raw_emb in doc["embeddings"]["list"]:
            data["embeddings"].append(read_raw_embedding(raw_emb))
            
        return data

def load_embedding(collection, qubo_id, solver_graph_id):
    doc = collection.find_one(
        {
            "qubo": bson.ObjectId(qubo_id),
            "solver_graph": bson.ObjectId(solver_graph_id)
        }
    )
    
    if doc == None:
        return None
    
    if doc["embeddings"] == None:
        return None
    
    if len(doc["embeddings"]) == 0:
        return None
    
    return read_raw_embedding(doc["embeddings"][0])

def get_id_of_solver_graph(collection, nx_graph_data):
    doc = collection.find_one({"data": nx_graph_data})
    
    if doc == None:
        return None
    
    return doc["_id"]

def load_WMIS_qubo_of_instance(collection, instance_id):
    doc = load_WMIS_qubo_of_instance_raw(collection, instance_id)
    
    if doc == None:
        return None
    
    return read_raw_qubo(doc["qubo"]), doc["_id"]

def get_WMIS_qubo_id_of_instance(collection, instance_id):
    doc = load_WMIS_qubo_of_instance_raw(collection, instance_id)
    
    if doc == None:
        return None
    
    return doc["_id"]

def load_WMIS_qubo_of_instance_raw(collection, instance_id):
    return collection.find_one({"instance": bson.ObjectId(instance_id)})

def read_raw_qubo(raw_qubo):
    qubo = {}
    
    for entry in raw_qubo:
        energy = entry[1]
        raw_coupler = entry[0]
        node1 = tuple(raw_coupler[0])
        node2 = tuple(raw_coupler[1])
        
        qubo[(node1, node2)] = energy
        
    return qubo

def read_raw_embedding(raw_embedding):
    emb = {}
    
    for entry in raw_embedding:
        emb[tuple(entry[0])] = entry[1]
        
    return emb