concept-graph

#!/usr/bin/env python3

# TechKnAcq: Concept Graph
# Jonathan Gordon

import sys
import os
import tempfile
import glob
import random
import subprocess
import click
import re

from collections import defaultdict

from mallet import Mallet
from techknacq.corpus import Corpus
from techknacq.conceptgraph import ConceptGraph

# Parameters

MALLET_PATH = 'ext/mallet/bin/mallet'

LDA_TOPICS = 300
LDA_ITERATIONS = 1000


class TopicDependency:
    """Python wrapper class for TechKnAcq-Core Java code for predicting
    dependencies among LDA topics."""

    jar = 'ext/techknacq-core/target/techknacq-core.jar'

    def __init__(self, corpus, model, method='ce', threshold=0.05):
        self.corpus = corpus
        self.model = model

        rand_prefix = hex(random.randint(0, 0xffffff))[2:] + '-'
        self.prefix = os.path.join(tempfile.gettempdir(), rand_prefix)

        if os.path.exists('data/alledge.tsv'):
            print('Using existing alledge.tsv file.')
        else:
            doc_comp_file = self.convert_comp_file()
            citation_network_file = self.build_citation_network()
            infomap_tree_file, infomap_flow_file = self.get_infomap()

            print('Computing concept dependencies.')

            deps = glob.glob('ext/techknacq-core/target/dependency/*.jar')

            cmd = ['java', '-cp', ':'.join([self.jar] + deps),
                   'edu.isi.techknacq.graph.ComparisonOnAlledges',
                   model.wtkfile, infomap_tree_file, doc_comp_file,
                   str(len(model.topics)), citation_network_file,
                   infomap_flow_file]

            if subprocess.call(cmd) != 0:
                sys.stderr.write('Running topic-dependency code failed.\n')
                sys.stderr.write(' '.join(cmd) + '\n')
                sys.exit(1)

            os.rename('alledge.tsv', 'data/alledge.tsv')

        self.read_alledges(method, threshold)

        # self.generate_topicgraph()


    def read_alledges(self, method, threshold):
        """Read the cross-entropy edges from the alledges.tsv output file."""

        self.edges = defaultdict(dict)
        for line in open('data/alledge.tsv'):
            try:
                tokens = line.strip().split()
                if tokens[0] == 'sid':
                    continue

                t1 = tokens[0]
                t2 = tokens[2]

                if method == 'simword':
                    weight = float(tokens[5])
                elif method == 'ce':
                    weight = float(tokens[7])
                elif method == 'avg':
                    weight = (float(tokens[5]) + float(tokens[7]))/2.0
                elif method == 'sum':
                    weight = float(tokens[5]) + float(tokens[7])

                # Flip negative weights.
                if weight < 0:
                    weight = -1.0 * weight
                    t1, t2 = t2, t1
                if weight > threshold:
                    self.edges[t1][t2] = max(self.edges[t1].get(t2, 0.0),
                                             weight)
                    # Only keep the strongest edge direction; don't allow
                    # bidirectional edges.
                    # if t1 in self.edges[t2]:
                    #     if self.edges[t2][t1] > self.edges[t1][t2]:
                    #         self.edges[t1][t2] = 0.0
                    #     else:
                    #         self.edges[t2][t1] = 0.0
            except:
                sys.stderr.write('Error processing alledges line:\n')
                sys.stderr.write(line)


    def generate_topicgraph(self):
        """Produce the topicgraph.txt file used by techknacq-server."""

        with open('topicgraph.txt', 'w') as fout:
            fout.write('%d\n' % (len(self.edges)))
            for t1 in self.edges:
                good_edges = sorted([(y, self.edges[t1][y])
                                     for y in self.edges[t1]],
                                    key=lambda x: self.edges[t1].get(x, 0),
                                    reverse=True)
                fout.write('%s,%d' % (t1, len(good_edges)))
                for t2, w in good_edges:
                    fout.write(':%s,%s' % (t2, w))
                fout.write('\n')


    def convert_comp_file(self):
        """Convert Mallet composition file to the format needed by the
        topic dependency code."""

        print('Converting Mallet document-topic composition file.')

        fname = self.prefix + 'concept2doc.txt'

        def pairwise(iterable):
            """s -> (s0,s1), (s2,s3), (s4, s5), ..."""
            a = iter(iterable)
            return zip(a, a)

        with open(fname, 'w') as fout:
            for line in open(self.model.dtfile):
                if line[0] == '#':
                    continue
                elts = line.strip().split()[1:]
                fout.write(elts[0].replace('file:', ''))
                for tid, pct in sorted(pairwise(elts[1:]),
                                       key=lambda x: x[0]):
                    fout.write('\ttopic' + tid + ':' + pct)
                fout.write('\n')

        return fname


    def get_infomap(self):
        # If there's saved Infomap output in the current directory,
        # use it.
        if os.path.exists('data/infomap-tree.txt') and \
           os.path.exists('data/infomap-flow.txt'):
            print('Using saved infomap output.')
            return 'data/infomap-tree.txt', 'data/infomap-flow.txt'

        print('Building Infomap graphs.')
        treefname = self.prefix + 'infomap-tree.txt'
        flowfname = self.prefix + 'infomap-flow.txt'

        # Generate Linhong-style names for each topic.
        keynames = []
        for t in range(len(self.model.topics)):
            keynames.append('-'.join([x[0] for x in
                                      self.model.topic_pairs(t)[:5]]) + '-')

        # Generate input file for Infomap.
        with open(self.prefix + 'infomap.net', 'w') as fout:
            fout.write('*Vertices %d\n' % (len(self.model.topics)))
            for v in range(len(self.model.topics)):
                fout.write('%d "%s"\n' % (v, keynames[v]))
            out_text = ''
            edge_count = 0
            for i, line in enumerate(open(self.model.cofile)):
                for j, count in enumerate(re.split('[^0-9]+', line.strip())):
                    if int(count) > 0 and j > i:
                        out_text += '%d %d %s\n' % (i, j, count)
                        edge_count += 1
            fout.write('*Edges %d\n' % (edge_count))
            fout.write(out_text)

        cmd = ['ext/infomap/Infomap', '-z',
               '--flow-network', self.prefix + 'infomap.net',
               tempfile.gettempdir()]
        if subprocess.call(cmd) != 0:
            sys.stderr.write('Running Infomap failed.\n')
            sys.exit(1)

        os.remove(self.prefix + 'infomap.net')

        os.rename(self.prefix + 'infomap.tree', treefname)
        os.rename(self.prefix + 'infomap.flow', flowfname)

        return treefname, flowfname


    def build_citation_network(self):
        # Use a pre-existing citation network.
        if os.path.exists('cite.txt'):
            print('Using saved citation network.')
            return 'cite.txt'

        print('Building citation network.')
        fname = self.prefix + 'cite.txt'

        with open(fname, 'w') as fout:
            for doc in self.corpus:
                try:
                    for ref in doc.references:
                        fout.write('%s ==> %s\n' % (doc.id, ref))
                except:
                    print('Error with doc', doc)

        return fname

###

@click.command()
@click.option('--method', default='avg',
              type=click.Choice(['ce', 'simword', 'avg', 'sum']),
              help='Method for computing concept dependencies.')
@click.option('--threshold', default=0.0005)
@click.option('--num-topics', default=LDA_TOPICS)
@click.argument('corpusdir', type=click.Path(exists=True))
@click.argument('topic_prefix', required=False)
def main(corpusdir, topic_prefix, method, threshold, num_topics):
    rand_prefix = hex(random.randint(0, 0xffffff))[2:] + '-'
    prefix = os.path.join(tempfile.gettempdir(), rand_prefix)

    cg = ConceptGraph()

    corpus = Corpus(corpusdir)
    # corpus.fix_text()

    cg.add_docs(corpus)

    if not topic_prefix:
        print('Generating topic model.')
        mallet_corpus = prefix + 'corpus'
        os.makedirs(mallet_corpus)
        corpus.export(mallet_corpus, abstract=False, form='text')
        model = Mallet(MALLET_PATH, mallet_corpus, num_topics=num_topics,
                       iters=LDA_ITERATIONS, bigrams=True)
    else:
        print('Loading topic model.')
        model = Mallet(MALLET_PATH, prefix=topic_prefix)
        if len(model.topic_doc[0]) < len(corpus.docs):
            print('Found more documents in the corpus (%d)' %
                  (len(corpus.docs)), end=' ')
            print('than in the topic model (%d).' % (len(model.topic_doc[0])))
            print('Inferring topics for corpus documents.')
            mallet_corpus = prefix + 'corpus'
            os.makedirs(mallet_corpus)
            corpus.export(mallet_corpus, abstract=False, form='text')
            model.infer_topics(mallet_corpus)

    if os.path.exists('data/alt-dt.txt'):
        print('Loading alternative document-topic composition.')
        new_dt = [{} for n in range(len(model.topics))]
        for line in open('data/alt-dt.txt'):
            elts = line.strip().split('\t')
            tnum = int(elts[0])
            for e in elts[1:]:
                doc, weight = e.split(':')
                new_dt[tnum][doc] = float(weight)

        for i in range(len(model.topics)):
            vals = dict(model.topic_doc[i])
            for v in set(vals.keys()) | set(new_dt[i].keys()):
                if v in vals and v in new_dt[i]:
                    vals[v] = (vals[v] + new_dt[i][v])/2.0
            model.topic_doc[i] = sorted(vals.items(), key=lambda x: x[1],
                                        reverse=True)

    cg.add_concepts(model)

    dep = TopicDependency(corpus, model, method=method, threshold=threshold)
    cg.add_dependencies(dep.edges)

    cg.export(prefix + 'cg.json',
              provenance=method + ' ' + str(threshold))
    print('Concept graph:', prefix + 'cg.json')



if __name__ == '__main__':
    main()