edgelist2mtx.py

#!/usr/bin/env python

"""
    edgelist2mtx.py
"""


from scipy.io import mmwrite
import numpy as np
import pandas as pd
import argparse
from scipy import sparse

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--inpath', type=str)
    parser.add_argument('--outpath', type=str)
    parser.add_argument('--symmetry', type=str, required=True)
    parser.add_argument('--field', type=str, default='pattern')
    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()
    print('reading')
    edges = pd.read_csv(args.inpath, sep='\t', header=None).values
    
    print('processing')
    rows = edges[:,0]
    cols = edges[:,1]
    vals = np.ones(edges.shape[0])
    
    unodes = np.unique(np.hstack([rows, cols]))
    if unodes.shape[0] != unodes.max() + 1:
        print('remapping nodes')
        lookup = dict(zip(unodes, range(len(unodes))))
        rows = np.array([lookup[x] for x in rows])
        cols = np.array([lookup[x] for x in cols])
        
    num_nodes = unodes.shape[0]
    
    m = sparse.csr_matrix((vals, (rows, cols)), shape=(num_nodes, num_nodes))
    m.sort_indices()
    
    print('writing (%s, %s)' % (args.field, args.symmetry))
    mmwrite(args.outpath, m, field=args.field, symmetry=args.symmetry)