load_db

load_db.py

from gremlin_python import statics

from opencog.atomspace import AtomSpace, types
from opencog.scheme_wrapper import scheme_eval

from gremlin_python.structure.graph import Graph
from gremlin_python.process.graph_traversal import __
from gremlin_python.process.strategies import *
from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
from gremlin_python.driver import client


def measure(func):
    def wrapped(*args, **kwargs):
        start = time.monotonic()
        result = func(*args, **kwargs)
        end = time.monotonic()

        if end - start > 0.8:
            import pdb;pdb.set_trace()
        return result
    return wrapped

@measure
def addVertex(label, _id, name, wrap="wrap"):
    vertex = g.V().has("atomhash", _id)\
            .has("value", name)\
            .has("type", label).fold()\
            .coalesce(__.unfold(),
                      __.addV(label)\
                              .property("atomhash", _id)\
                              .property("type", label)\
                              .property("value", name))\
                              .aggregate(wrap)
    return vertex


def putNode(g, node):
    return addVertex(str(node.type), str(hash(node)), node.name)


def getArgType(idx):
    return 'edge_' + str(idx)

@measure
def getEdgeV(type_, id_):
    t = g.V()\
            .has("type", type_)\
            .has("atomhash", id_)

    # todo: if vertex exists need full traversal due to
    # possible has collision, so upsert here not possible

    for _ in t:
        return _


@measure
def putLink(g, link):
    vertex = getEdgeV(str(link.type), str(hash(link)))
    if vertex is not None:
        return vertex
    vertex = addVertex(str(link.type), str(hash(link)), "", wrap='ve')

    expression = vertex
    for i, atom in enumerate(link.out):
        expression = expression.select("ve").unfold().addE(getArgType(i)).to(putAtom(g, atom, 1))

    return expression.select("ve").unfold()


@measure
def putAtom(g, atom, depth=0):
    if atom.is_node():
        result = putNode(g, atom)
    elif atom.is_link():
        result = putLink(g, atom)
    else:
        raise RuntimeError("Unknown atom type!");
    if depth == 0 and hasattr(result, 'next'):
        res = result.next()
        return res, result
    return result


def load_kb(path):
    scheme_eval(atomspace, '(use-modules (opencog bioscience))')
    scheme_eval(atomspace, '(add-to-load-path ".")')
    scheme_eval(atomspace, '(load-from-path "' + path + '")')

atomspace = AtomSpace()

load_kb("kbs/ChEBI2Reactome_PE_Pathway.txt.scm")
#load_kb("test.scm")
graph = Graph()

g = graph.traversal().withRemote(DriverRemoteConnection('ws://localhost:8182/gremlin','g'))

import time
import pdb;pdb.set_trace()
print(len(atomspace))
start = time.monotonic()
for i,atom in enumerate(atomspace):
    start = time.monotonic()
    if atom.is_node():
        continue
    v = putAtom(g, atom)
    if i and i % 100 == 0:
        print(i)
    end = time.monotonic()
    if end - start > 1.0:
        print(end - start)



def exprToGremlin(expr):
    if len(expr) == 1:
        if isinstance(expr[0], list):
            return exprToGremlin(expr[0])
        return '.' + expr[0] + '()'
    if isinstance(expr[0], str):
        return '.' + expr[0] + '(' + ', '.join(expr[1:]) + ')'
    return ''.join(exprToGremlin(e) for e in expr)