diff python/utils.py @ 840:15a82ebc62c4

utils for sparse matrix
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Fri, 08 Jul 2016 11:41:29 -0400
parents e01cabca4c55
children 90b7d6e19c55
line wrap: on
line diff
--- a/python/utils.py	Wed Jul 06 18:05:03 2016 -0400
+++ b/python/utils.py	Fri Jul 08 11:41:29 2016 -0400
@@ -7,7 +7,8 @@
 from math import sqrt, ceil, floor
 from scipy.stats import kruskal, shapiro
 from scipy.spatial import distance
-from numpy import zeros, array, exp, sum as npsum, int as npint, arange, cumsum, median, isnan, ones, convolve,  dtype, isnan, NaN, mean, ma, isinf
+from scipy.sparse import dok_matrix
+from numpy import zeros, array, exp, sum as npsum, int as npint, arange, cumsum, median, isnan, ones, convolve,  dtype, isnan, NaN, mean, ma, isinf, savez, load as npload
 
 
 datetimeFormat = "%Y-%m-%d %H:%M:%S"
@@ -483,6 +484,17 @@
         out.close()
     return result
 
+def saveDokMatrix(filename, m):
+    'Saves a dok_matrix using savez'
+    savez(filename, shape = m.shape, keys = m.keys(), values = m.values())
+
+def loadDokMatrix(filename):
+    'Loads a dok_matrix saved using the above saveDokMatrix'
+    data = npload(filename)
+    m = dok_matrix(tuple(data['shape']))
+    for k, v in zip(data['keys'], data['values']):
+        m[tuple(k)] = v
+    return m
 
 #########################
 # regression analysis using statsmodels (and pandas)