Fast row-wise, parallel calculations on large, sparse matrix

Question

I have a big (50k*150k) matrix with co-occurrences of nouns (rows) and adjectives (columns). As most nouns don't co-occurr with most adjectives, this matrix is very sparse,>99.9%, so I'm using the CSR format. For every pair of nouns I need to calculate the Jason-Shannon Divergence to asses how similar the co-occurrence distributions with all adjectives are between each pair. I have sixteen cores I can use, so I run the calculations in parallel, such that each core processes every sixteenth noun.

It's not going very fast and I'm wondering whether there's a faster way to do this.

from __future__ import division
from multiprocessing import Pool
from scipy.stats import entropy
import codecs
import os
import scipy
import scipy.sparse
# --- data & parameters--- #
frequenciesAdjectives = codecs.open('frequencies.txt', 'r', 'utf-8')
results = codecs.open('results.txt', 'w', 'utf-8')
temporaryFilesPath = "/temp/"
cores = 16
# --- functions --- #
# calculate jensen shannon divergence
def JSD(p, q):
 p = p
 q = q
 m = 0.5 * (p+q)
 jsd = 0.5 * (entropy(p, m) + entropy(q, m))
 return jsd
# calculate JSD for every i-th noun with every other noun and write to temp file, where i is the number of cores used, and return filepath
def getJSDs(n):
 # open temporary file
 filename = "temp"+str(n).zfill(2)
 fullpath = temporaryFilesPath+filename
 temp = codecs.open(fullpath, "w", "utf-8")
 # shortcut write function
 tempwrite = temp.write
 # calculate JSD for each noun pair and write to temporary file
 for index, noun1 in enumerate(nouns[n::cores]):
 index = index*cores+n
 first = sparseMatrix.getrow(index).toarray()[0]
 tempwrite("here")
 for index2, noun2 in enumerate(nouns[index:]):
 index2 += index
 second = sparseMatrix.getrow(index2).toarray()[0]
 divergence = JSD(first, second)
 tempwrite(u"{noun1}\t{noun2}\t{divergence}\n".format(**locals()))
 temp.close()
 return fullpath
# --- processing --- #
# set up parameter variables for sparse matrix
nouns = []
adjectiveIDs = {} # column id of each adjective in matrix
frequencies = [] # non-zero entries in matrix
positions = [] # column id of frequency in corresponding position in "frequencies"
indices = [0] # frequencies[indices[i]:indices[i+]] = non-zero entries of line i of matrix
# ignore file header
frequenciesAdjectives.readline()
# incrementally get sparse matrix parameters (turn frequencies into probabilites)
for line in frequenciesAdjectives:
 line = line.strip().lower().split("\t")
 noun = line[0]
 nouns.append(noun)
 adjectiveList = [pair.split(" ") for pair in line[2:]]
 total = sum([int(frequency) for _,frequency in adjectiveList])
 for pair in adjectiveList:
 adjective, frequency = pair
 probability = int(frequency)/total
 position = adjectiveIDs.setdefault(adjective, len(adjectiveIDs))
 frequencies.append(probability)
 positions.append(position)
 indices.append(len(frequencies))
# turn lists into arrays
frequencies = scipy.array(frequencies)
positions = scipy.array(positions)
indices = scipy.array(indices)
# create sparse matrix from parameter arrays and delete arrays
sparseMatrix = scipy.sparse.csr_matrix((frequencies, positions, indices), shape=(len(nouns), len(adjectiveIDs)))
del frequencies, positions, indices, adjectiveIDs
# calculate JSDs in parallel and get list of temporary files
pool = Pool()
tempFiles = pool.map(getJSDs, range(cores))
pool.close()
pool.join()
# shortcut results.write and write header
resultswrite = results.write
resultswrite(u"noun1\tnoun2\tjensenShannonDivergence\n")
# combine temporary files into results file and delete
for path in tempFiles:
 tempfile = codecs.open(path, "r", "utf-8")
 for line in tempfile:
 resultswrite(line)
 tempfile.close()
 os.remove(path)

The format of the file from which I read in the frequencies is this:

"noun\tnounFrequency\tadjective1\tcooccurrenceFrequency1\tadjective2\tcooccurrenceFrequency2\n"

I can't use the dense matrix format, because that uses up all of the RAM. I doubt I can get the entropy function to run faster than SciPy - I tried using a for loop and numba, but it didn't come out faster and I don't know any C. I don't really know parallel computing, so I bet my way of using pool.map isn't optimal. Would it maybe be faster if I copied the global matrix variable into a local variable for each function?

Answer 1

I rewrote the entire thing, now it's two orders of magnitude (500x or so!) faster.

Key changes include:

• Using a list of tuples (noun, adjectiveDictionary) instead of the sparse matrix and dropping the 0 elements completely
• Using math.log instead of scipy.log (huge difference!)
• Switching from Python 2 to Python 3
• Not dividing the tasks up manually, instead using a generator and the queue that comes with imap_unordered, for which it was crucial to find the right chunksize.

#!/usr/bin/env/python3
from multiprocessing import Pool
from collections import Counter
from math import log2
# from numba import jit
# --- data & parameters--- #
frequenciesAdjectives = open('/home/christian/results/gender/frequenciesAdjectivesGivenNouns_UK.txt', 'r')
results = open('/home/christian/results/gender/JensenShannonDivergences_ukWaC.txt', 'w')
cores = 16
# --- functions --- #
# calculates Jason-Shannon Divergence from tuple of two nouns and their associated adjective probabilities in dictionaries p and q
def JSD(nounTuple):
 noun1, noun2, p, q = nounTuple
 jsd = 0.0
 m = p + q
 for key in m:
 m_key = 0.5 * m[key]
 if key in p:
 p_key = p[key]
 jsd += 0.5 * p_key * log2(p_key/m_key)
 if key in q:
 q_key = q[key]
 jsd += 0.5 * q_key * log2(q_key/m_key)
 return noun1, noun2, jsd
def jobGenerator(tuples):
 for index, (noun, adjectives) in enumerate(tuples):
 for noun2, adjectives2 in tuples[index:]:
 yield noun, noun2, adjectives, adjectives2
# --- processing --- #
# ignore header
frequenciesAdjectives.readline()
# make list of tuples of nouns and dictionaries containing their preceding adjective frequencies
nounAdjectives = []
for line in frequenciesAdjectives:
 adjectives = Counter()
 line = line.strip().lower().split("\t")
 noun = line[0]
 adjectiveList = [pair.split(" ") for pair in line[2:]]
 frequencySum = sum(int(frequency) for _, frequency in adjectiveList)
 for adjective, frequency in adjectiveList:
 probability = int(frequency)/frequencySum
 adjectives[adjective] = probability
 nounAdjectives.append((noun, adjectives))
# make generator of (noun, noun2, adjectives, adjectives2)-tuples
jobs = jobGenerator(nounAdjectives)
# shortcut results.write and write header
resultswrite = results.write
resultswrite(u"noun1\tnoun2\tjensenShannonDivergence")
# calculate JSDs in parallel and write to file
pool = Pool(cores)
for noun1, noun2, jsd in pool.imap_unordered(JSD, jobs, chunksize=500000):
 resultswrite(u"\n{noun1}\t{noun2}\t{jsd}".format_map(locals()))
pool.close()