qubo_lab/util/queries.py

from .kSAT import kSAT
import bson
import dimod

class Instance_scope_query:
    
    def __init__(self, database):
        self.__database = database
        self.__query = None
        self.__instance_ids = []
        self.__instance_id_iterator = 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"
            },
            {
                "$project": {"instance_id": "$instance._id"}
            }
        ])
        
        self.__instance_ids = []
        for doc in self.__query:
            self.__instance_ids.append(doc["instance_id"])
        
        self.__instance_id_iterator = iter(self.__instance_ids)
    
    def __iter__(self):
        return self
    
    def __next__(self):
        instance_id = self.__instance_id_iterator.__next__()
        
        document = self.__database["instances"].find_one({"_id": instance_id})
        
        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
        self.__qubo_ids = []
        self.__qubo_id_iterator = 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"
            },
            {
                "$project": {"qubo_id": "$qubo._id"}
            }
        ])
        
        self.__qubo_ids = []
        for doc in self.__query:
            self.__qubo_ids.append(doc["qubo_id"])
        
        self.__qubo_id_iterator = iter(self.__qubo_ids)
            
    def __len__(self):
        return self.query.count_documents({})
            
    def __iter__(self):
        return self
    
    def __next__(self):
        qubo_filter = {"_id": self.__qubo_id_iterator.__next__()}
        
        return self.__database["wmis_qubos"].find_one(qubo_filter)
    
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
        self.__ids = []
        self.__id_iterator = 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"
            },
            {
                "$project":
                {
                    "instance_id": True,
                    "wmis_qubo_id": "$wmis_qubo._id"
                }
            },
            {
                "$lookup":
                {
                    "from": "embeddings",
                    "let":
                    {
                        "qubo_id": "$wmis_qubo_id",
                        "solver_graph_id": bson.ObjectId(solver_graph_id)
                    },
                    "pipeline":
                    [
                    {
                        "$project":
                        {
                            "qubo": True,
                            "solver_graph": True,
                        }
                    },
                    {
                        "$match":
                        {
                            "$expr":
                            {
                                "$and":
                                [
                                {"$eq": ["$qubo", "$$qubo_id"]},
                                {"$eq": ["$solver_graph", "$$solver_graph_id"]}
                                ]
                            }
                        }
                    },
                    {
                        "$project":
                        {
                            "_id": True
                        }
                    }       
                    ],
                    "as": "embeddings_id"
                }
            },
            {
                "$unwind": "$embeddings_id"
            },
            {
                "$project":
                {
                    "instance_id": True,
                    "wmis_qubo_id": True,
                    "embeddings_id": "$embeddings_id._id"
                }
            }
        ])
                
        self.__ids = []
        for doc in self.__query:
            self.__ids.append({
                "instance_id": doc["instance_id"],
                "wmis_qubo_id": doc["wmis_qubo_id"],
                "embeddings_id": doc["embeddings_id"]
            })
        
        self.__id_iterator = iter(self.__ids)
            
    def __iter__(self):
        return self
    
    def __next__(self):
        ids = self.__id_iterator.__next__()
        
        doc = {}
        
        doc["instance_id"] = ids["instance_id"]
        
        qubo_filter = {"_id": ids["wmis_qubo_id"]}
        doc["wmis_qubo"] = self.__database["wmis_qubos"].find_one(qubo_filter)
        
        embeddings_filter = {"_id": ids["embeddings_id"]}
        doc["embeddings"] = self.__database["embeddings"].find_one(embeddings_filter)
        
        return doc
    
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"]["embeddings"]:
            data["embeddings"].append(read_raw_embedding(raw_emb))
            
        return data
    
class WMIS_result_scope_query_raw:
    def __init__(self, database):
        self.__database = database
        self.__query = None
        
    def query(self, scope, collection):
        self.__query = self.__database["experiment_scopes"].aggregate([
            {
                "$match": {"_id": scope}
            },
            {
                "$unwind": "$instances"
            },
            {
                "$project": {"instance_id": "$instances"}
            },
            {
                "$lookup":
                {
                    "from": collection,
                    "localField": "instance_id",  
                    "foreignField": "instance",
                    "as": "result"
                }
            },
            {
                "$unwind": "$result"
            },
            {
                "$replaceRoot": {"newRoot": "$result"}
            }
        ])
    
    def __iter__(self):
        return self
    
    def __next__(self):
        return self.__query.next()
    
class Minisat_run_scope_query_raw:
    def __init__(self, database):
        self.__database = database
        self.__query = None
        
    def query(self, scope, collection):
        self.__query = self.__database["experiment_scopes"].aggregate([
            {
                "$match": {"_id": scope}
            },
            {
                "$unwind": "$instances"
            },
            {
                "$project": {"instance_id": "$instances"}
            },
            {
                "$lookup":
                {
                    "from": collection,
                    "localField": "instance_id",  
                    "foreignField": "instance",
                    "as": "run"
                }
            },
            {
                "$unwind": "$run"
            },
            {
                "$replaceRoot": {"newRoot": "$run"}
            }
        ])
    
    def __iter__(self):
        return self
    
    def __next__(self):
        return self.__query.next()
    

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 = {}
    
    if "embedding" in raw_embedding:
        raw_embedding = raw_embedding["embedding"]
        
    for entry in raw_embedding:
        emb[tuple(entry[0])] = entry[1]
        
    return emb

def read_raw_wmis_sample_set(raw_sample_set):
    sample_set_data = raw_sample_set.copy()
    
    sample_set_data["variable_labels"] = []
    
    for label in raw_sample_set["variable_labels"]:
        sample_set_data["variable_labels"].append(tuple(label))
        
    return dimod.SampleSet.from_serializable(sample_set_data)

def get_instance_by_id(collection, id):
    doc = collection.find_one({"_id": bson.ObjectId(id)})
    
    sat = kSAT()
    
    for clause in doc["clauses"]:
            sat.addClause(clause);
        
    return sat