graphkit-learn/pygraph/kernels/treeletKernel.py

"""
@author: linlin
@references: Gaüzère B, Brun L, Villemin D. Two new graphs kernels in chemoinformatics. Pattern Recognition Letters. 2012 Nov 1;33(15):2038-47.
"""

import sys
sys.path.insert(0, "../")
import time
from collections import Counter
from itertools import chain
from functools import partial
from multiprocessing import Pool
from tqdm import tqdm

import networkx as nx
import numpy as np

from pygraph.utils.graphdataset import get_dataset_attributes
from pygraph.utils.parallel import parallel_gm

def treeletkernel(*args, 
                  sub_kernel, 
                  node_label='atom', 
                  edge_label='bond_type', 
                  n_jobs=None, 
                  verbose=True):
    """Calculate treelet graph kernels between graphs.

    Parameters
    ----------
    Gn : List of NetworkX graph
        List of graphs between which the kernels are calculated.
    /
    G1, G2 : NetworkX graphs
        Two graphs between which the kernel is calculated.
    sub_kernel : function
        The sub-kernel between 2 real number vectors. Each vector counts the
        numbers of isomorphic treelets in a graph.
    node_label : string
        Node attribute used as label. The default node label is atom.        
    edge_label : string
        Edge attribute used as label. The default edge label is bond_type.
    labeled : boolean
        Whether the graphs are labeled. The default is True.

    Return
    ------
    Kmatrix : Numpy matrix
        Kernel matrix, each element of which is the treelet kernel between 2 praphs.
    """
    # pre-process
    Gn = args[0] if len(args) == 1 else [args[0], args[1]]
    Kmatrix = np.zeros((len(Gn), len(Gn)))
    ds_attrs = get_dataset_attributes(Gn,
        attr_names=['node_labeled', 'edge_labeled', 'is_directed'],
        node_label=node_label, edge_label=edge_label)
    labeled = False
    if ds_attrs['node_labeled'] or ds_attrs['edge_labeled']:
        labeled = True
        if not ds_attrs['node_labeled']:
            for G in Gn:
                nx.set_node_attributes(G, '0', 'atom')
        if not ds_attrs['edge_labeled']:
            for G in Gn:
                nx.set_edge_attributes(G, '0', 'bond_type')
    
    start_time = time.time()
    
    # ---- use pool.imap_unordered to parallel and track progress. ----
    # get all canonical keys of all graphs before calculating kernels to save 
    # time, but this may cost a lot of memory for large dataset.
    pool = Pool(n_jobs)
    itr = zip(Gn, range(0, len(Gn)))
    if len(Gn) < 100 * n_jobs:
        chunksize = int(len(Gn) / n_jobs) + 1
    else:
        chunksize = 100
    canonkeys = [[] for _ in range(len(Gn))]
    getps_partial = partial(wrapper_get_canonkeys, node_label, edge_label, 
                            labeled, ds_attrs['is_directed'])
    if verbose:
        iterator = tqdm(pool.imap_unordered(getps_partial, itr, chunksize),
                        desc='getting canonkeys', file=sys.stdout)
    else:
        iterator = pool.imap_unordered(getps_partial, itr, chunksize)
    for i, ck in iterator:
        canonkeys[i] = ck
    pool.close()
    pool.join()
    
    # compute kernels.
    def init_worker(canonkeys_toshare):
        global G_canonkeys
        G_canonkeys = canonkeys_toshare
    do_partial = partial(wrapper_treeletkernel_do, sub_kernel)
    parallel_gm(do_partial, Kmatrix, Gn, init_worker=init_worker, 
                glbv=(canonkeys,), n_jobs=n_jobs, verbose=verbose)
    
    run_time = time.time() - start_time
    if verbose:
        print("\n --- treelet kernel matrix of size %d built in %s seconds ---" 
              % (len(Gn), run_time))
        
    return Kmatrix, run_time


def _treeletkernel_do(canonkey1, canonkey2, sub_kernel):
    """Calculate treelet graph kernel between 2 graphs.
    
    Parameters
    ----------
    canonkey1, canonkey2 : list
        List of canonical keys in 2 graphs, where each key is represented by a string.
        
    Return
    ------
    kernel : float
        Treelet Kernel between 2 graphs.
    """
    keys = set(canonkey1.keys()) & set(canonkey2.keys()) # find same canonical keys in both graphs
    vector1 = np.array([(canonkey1[key] if (key in canonkey1.keys()) else 0) for key in keys])
    vector2 = np.array([(canonkey2[key] if (key in canonkey2.keys()) else 0) for key in keys]) 
    kernel = np.sum(np.exp(-np.square(vector1 - vector2) / 2))
#    kernel = sub_kernel(vector1, vector2) 
    return kernel


def wrapper_treeletkernel_do(sub_kernel, itr):
    i = itr[0]
    j = itr[1]
    return i, j, _treeletkernel_do(G_canonkeys[i], G_canonkeys[j], sub_kernel)


def get_canonkeys(G, node_label, edge_label, labeled, is_directed):
    """Generate canonical keys of all treelets in a graph.
    
    Parameters
    ----------
    G : NetworkX graphs
        The graph in which keys are generated.
    node_label : string
        node attribute used as label. The default node label is atom.        
    edge_label : string
        edge attribute used as label. The default edge label is bond_type.
    labeled : boolean
        Whether the graphs are labeled. The default is True.
        
    Return
    ------
    canonkey/canonkey_l : dict
        For unlabeled graphs, canonkey is a dictionary which records amount of 
        every tree pattern. For labeled graphs, canonkey_l is one which keeps 
        track of amount of every treelet.
    """
    patterns = {} # a dictionary which consists of lists of patterns for all graphlet.
    canonkey = {} # canonical key, a dictionary which records amount of every tree pattern.

    ### structural analysis ###
    ### In this section, a list of patterns is generated for each graphlet, 
    ### where every pattern is represented by nodes ordered by Morgan's 
    ### extended labeling.
    # linear patterns
    patterns['0'] = G.nodes()
    canonkey['0'] = nx.number_of_nodes(G)
    for i in range(1, 6): # for i in range(1, 6):
        patterns[str(i)] = find_all_paths(G, i, is_directed)
        canonkey[str(i)] = len(patterns[str(i)])

    # n-star patterns
    patterns['3star'] = [[node] + [neighbor for neighbor in G[node]] for node in G.nodes() if G.degree(node) == 3]
    patterns['4star'] = [[node] + [neighbor for neighbor in G[node]] for node in G.nodes() if G.degree(node) == 4]
    patterns['5star'] = [[node] + [neighbor for neighbor in G[node]] for node in G.nodes() if G.degree(node) == 5]        
    # n-star patterns
    canonkey['6'] = len(patterns['3star'])
    canonkey['8'] = len(patterns['4star'])
    canonkey['d'] = len(patterns['5star'])

    # pattern 7
    patterns['7'] = [] # the 1st line of Table 1 in Ref [1]
    for pattern in patterns['3star']:
        for i in range(1, len(pattern)): # for each neighbor of node 0
            if G.degree(pattern[i]) >= 2:
                pattern_t = pattern[:]
                # set the node with degree >= 2 as the 4th node
                pattern_t[i], pattern_t[3] = pattern_t[3], pattern_t[i]
                for neighborx in G[pattern[i]]:
                    if neighborx != pattern[0]:
                        new_pattern = pattern_t + [neighborx]
                        patterns['7'].append(new_pattern)
    canonkey['7'] = len(patterns['7'])

    # pattern 11
    patterns['11'] = [] # the 4th line of Table 1 in Ref [1]
    for pattern in patterns['4star']:
        for i in range(1, len(pattern)):
            if G.degree(pattern[i]) >= 2:
                pattern_t = pattern[:]
                pattern_t[i], pattern_t[4] = pattern_t[4], pattern_t[i]
                for neighborx in G[pattern[i]]:
                    if neighborx != pattern[0]:
                        new_pattern = pattern_t + [ neighborx ]
                        patterns['11'].append(new_pattern)
    canonkey['b'] = len(patterns['11'])

    # pattern 12
    patterns['12'] = [] # the 5th line of Table 1 in Ref [1]
    rootlist = [] # a list of root nodes, whose extended labels are 3
    for pattern in patterns['3star']:
        if pattern[0] not in rootlist: # prevent to count the same pattern twice from each of the two root nodes
            rootlist.append(pattern[0])
            for i in range(1, len(pattern)):
                if G.degree(pattern[i]) >= 3:
                    rootlist.append(pattern[i])
                    pattern_t = pattern[:]
                    pattern_t[i], pattern_t[3] = pattern_t[3], pattern_t[i]
                    for neighborx1 in G[pattern[i]]:
                        if neighborx1 != pattern[0]:
                            for neighborx2 in G[pattern[i]]:
                                if neighborx1 > neighborx2 and neighborx2 != pattern[0]:
                                    new_pattern = pattern_t + [neighborx1] + [neighborx2]
#                         new_patterns = [ pattern + [neighborx1] + [neighborx2] for neighborx1 in G[pattern[i]] if neighborx1 != pattern[0] for neighborx2 in G[pattern[i]] if (neighborx1 > neighborx2 and neighborx2 != pattern[0]) ]
                                    patterns['12'].append(new_pattern)
    canonkey['c'] = int(len(patterns['12']) / 2)

    # pattern 9
    patterns['9'] = [] # the 2nd line of Table 1 in Ref [1]
    for pattern in patterns['3star']:
        for pairs in [ [neighbor1, neighbor2] for neighbor1 in G[pattern[0]] if G.degree(neighbor1) >= 2 \
            for neighbor2 in G[pattern[0]] if G.degree(neighbor2) >= 2 if neighbor1 > neighbor2 ]:
            pattern_t = pattern[:]
            # move nodes with extended labels 4 to specific position to correspond to their children
            pattern_t[pattern_t.index(pairs[0])], pattern_t[2] = pattern_t[2], pattern_t[pattern_t.index(pairs[0])]
            pattern_t[pattern_t.index(pairs[1])], pattern_t[3] = pattern_t[3], pattern_t[pattern_t.index(pairs[1])]
            for neighborx1 in G[pairs[0]]:
                if neighborx1 != pattern[0]:
                    for neighborx2 in G[pairs[1]]:
                        if neighborx2 != pattern[0]:
                            new_pattern = pattern_t + [neighborx1] + [neighborx2]
                            patterns['9'].append(new_pattern)
    canonkey['9'] = len(patterns['9'])

    # pattern 10
    patterns['10'] = [] # the 3rd line of Table 1 in Ref [1]
    for pattern in patterns['3star']:        
        for i in range(1, len(pattern)):
            if G.degree(pattern[i]) >= 2:
                for neighborx in G[pattern[i]]:
                    if neighborx != pattern[0] and G.degree(neighborx) >= 2:
                        pattern_t = pattern[:]
                        pattern_t[i], pattern_t[3] = pattern_t[3], pattern_t[i]
                        new_patterns = [ pattern_t + [neighborx] + [neighborxx] for neighborxx in G[neighborx] if neighborxx != pattern[i] ]
                        patterns['10'].extend(new_patterns)
    canonkey['a'] = len(patterns['10'])

    ### labeling information ###
    ### In this section, a list of canonical keys is generated for every 
    ### pattern obtained in the structural analysis section above, which is a 
    ### string corresponding to a unique treelet. A dictionary is built to keep
    ### track of the amount of every treelet.
    if labeled == True:
        canonkey_l = {} # canonical key, a dictionary which keeps track of amount of every treelet.

        # linear patterns
        canonkey_t = Counter(list(nx.get_node_attributes(G, node_label).values()))
        for key in canonkey_t:
            canonkey_l['0' + key] = canonkey_t[key]

        for i in range(1, 6): # for i in range(1, 6):
            treelet = []
            for pattern in patterns[str(i)]:
                canonlist = list(chain.from_iterable((G.node[node][node_label], \
                    G[node][pattern[idx+1]][edge_label]) for idx, node in enumerate(pattern[:-1])))
                canonlist.append(G.node[pattern[-1]][node_label])
                canonkey_t = ''.join(canonlist)
                canonkey_t = canonkey_t if canonkey_t < canonkey_t[::-1] else canonkey_t[::-1]
                treelet.append(str(i) + canonkey_t)
            canonkey_l.update(Counter(treelet))

        # n-star patterns
        for i in range(3, 6):
            treelet = []
            for pattern in patterns[str(i) + 'star']:
                canonlist = [ G.node[leaf][node_label] + G[leaf][pattern[0]][edge_label] for leaf in pattern[1:] ]
                canonlist.sort()
                canonkey_t = ('d' if i == 5 else str(i * 2)) + G.node[pattern[0]][node_label] + ''.join(canonlist)
                treelet.append(canonkey_t)
            canonkey_l.update(Counter(treelet))

        # pattern 7
        treelet = []
        for pattern in patterns['7']:
            canonlist = [ G.node[leaf][node_label] + G[leaf][pattern[0]][edge_label] for leaf in pattern[1:3] ]
            canonlist.sort()
            canonkey_t = '7' + G.node[pattern[0]][node_label] + ''.join(canonlist) \
                + G.node[pattern[3]][node_label] + G[pattern[3]][pattern[0]][edge_label] \
                 + G.node[pattern[4]][node_label] + G[pattern[4]][pattern[3]][edge_label]
            treelet.append(canonkey_t)
        canonkey_l.update(Counter(treelet))

        # pattern 11
        treelet = []
        for pattern in patterns['11']:
            canonlist = [ G.node[leaf][node_label] + G[leaf][pattern[0]][edge_label] for leaf in pattern[1:4] ]
            canonlist.sort()
            canonkey_t = 'b' + G.node[pattern[0]][node_label] + ''.join(canonlist) \
                + G.node[pattern[4]][node_label] + G[pattern[4]][pattern[0]][edge_label] \
                 + G.node[pattern[5]][node_label] + G[pattern[5]][pattern[4]][edge_label]
            treelet.append(canonkey_t)
        canonkey_l.update(Counter(treelet))

        # pattern 10
        treelet = []
        for pattern in patterns['10']:
            canonkey4 = G.node[pattern[5]][node_label] + G[pattern[5]][pattern[4]][edge_label]
            canonlist = [ G.node[leaf][node_label] + G[leaf][pattern[0]][edge_label] for leaf in pattern[1:3] ]
            canonlist.sort()
            canonkey0 = ''.join(canonlist)
            canonkey_t = 'a' + G.node[pattern[3]][node_label] \
                + G.node[pattern[4]][node_label] + G[pattern[4]][pattern[3]][edge_label] \
                + G.node[pattern[0]][node_label] + G[pattern[0]][pattern[3]][edge_label] \
                + canonkey4 + canonkey0
            treelet.append(canonkey_t)
        canonkey_l.update(Counter(treelet))

        # pattern 12
        treelet = []
        for pattern in patterns['12']:
            canonlist0 = [ G.node[leaf][node_label] + G[leaf][pattern[0]][edge_label] for leaf in pattern[1:3] ]
            canonlist0.sort()
            canonlist3 = [ G.node[leaf][node_label] + G[leaf][pattern[3]][edge_label] for leaf in pattern[4:6] ]
            canonlist3.sort()
            
            # 2 possible key can be generated from 2 nodes with extended label 3, select the one with lower lexicographic order.
            canonkey_t1 = 'c' + G.node[pattern[0]][node_label] \
                + ''.join(canonlist0) \
                + G.node[pattern[3]][node_label] + G[pattern[3]][pattern[0]][edge_label] \
                + ''.join(canonlist3)

            canonkey_t2 = 'c' + G.node[pattern[3]][node_label] \
                + ''.join(canonlist3) \
                + G.node[pattern[0]][node_label] + G[pattern[0]][pattern[3]][edge_label] \
                + ''.join(canonlist0)

            treelet.append(canonkey_t1 if canonkey_t1 < canonkey_t2 else canonkey_t2)
        canonkey_l.update(Counter(treelet))

        # pattern 9
        treelet = []
        for pattern in patterns['9']:
            canonkey2 = G.node[pattern[4]][node_label] + G[pattern[4]][pattern[2]][edge_label]
            canonkey3 = G.node[pattern[5]][node_label] + G[pattern[5]][pattern[3]][edge_label]
            prekey2 = G.node[pattern[2]][node_label] + G[pattern[2]][pattern[0]][edge_label]
            prekey3 = G.node[pattern[3]][node_label] + G[pattern[3]][pattern[0]][edge_label]
            if prekey2 + canonkey2 < prekey3 + canonkey3:
                canonkey_t = G.node[pattern[1]][node_label] + G[pattern[1]][pattern[0]][edge_label] \
                    + prekey2 + prekey3 + canonkey2 + canonkey3
            else:
                canonkey_t = G.node[pattern[1]][node_label] + G[pattern[1]][pattern[0]][edge_label] \
                    + prekey3 + prekey2 + canonkey3 + canonkey2
            treelet.append('9' + G.node[pattern[0]][node_label] + canonkey_t)
        canonkey_l.update(Counter(treelet))

        return canonkey_l

    return canonkey


def wrapper_get_canonkeys(node_label, edge_label, labeled, is_directed, itr_item):
    g = itr_item[0]
    i = itr_item[1]
    return i, get_canonkeys(g, node_label, edge_label, labeled, is_directed)
    

def find_paths(G, source_node, length):
    """Find all paths with a certain length those start from a source node. 
    A recursive depth first search is applied.
    
    Parameters
    ----------
    G : NetworkX graphs
        The graph in which paths are searched.
    source_node : integer
        The number of the node from where all paths start.
    length : integer
        The length of paths.
        
    Return
    ------
    path : list of list
        List of paths retrieved, where each path is represented by a list of nodes.
    """
    if length == 0:
        return [[source_node]]
    path = [[source_node] + path for neighbor in G[source_node] \
        for path in find_paths(G, neighbor, length - 1) if source_node not in path]
    return path


def find_all_paths(G, length, is_directed):
    """Find all paths with a certain length in a graph. A recursive depth first
    search is applied.
    
    Parameters
    ----------
    G : NetworkX graphs
        The graph in which paths are searched.
    length : integer
        The length of paths.
        
    Return
    ------
    path : list of list
        List of paths retrieved, where each path is represented by a list of nodes.
    """
    all_paths = []
    for node in G:
        all_paths.extend(find_paths(G, node, length))
        
    if not is_directed:
        # For each path, two presentations are retrieved from its two extremities. 
        # Remove one of them.
        all_paths_r = [path[::-1] for path in all_paths]  
        for idx, path in enumerate(all_paths[:-1]):
            for path2 in all_paths_r[idx+1::]:
                if path == path2:
                    all_paths[idx] = []
                    break
        all_paths = list(filter(lambda a: a != [], all_paths))
            
    return all_paths