Merge pull request #2 from dirtysalt/master

support to configure how many threads be used in native ALS

implicit/_implicit.pyx

@@ -9,7 +9,7 @@ cimport scipy.linalg.cython_lapack as cython_lapack
 cimport scipy.linalg.cython_blas as cython_blas
 
 @cython.boundscheck(False)
-def least_squares(Cui, double [:, :] X, double [:, :] Y, double regularization):
+def least_squares(Cui, double [:, :] X, double [:, :] Y, double regularization, int num_threads):
     cdef int [:] indptr = Cui.indptr, indices = Cui.indices
     cdef double [:] data = Cui.data
 
@@ -25,14 +25,14 @@ def least_squares(Cui, double [:, :] X, double [:, :] Y, double regularization):
     cdef double * b
     cdef int * pivot
 
-    with nogil, parallel():
+    with nogil, parallel(num_threads = num_threads):
         # allocate temp memory for each thread
         A = <double *> malloc(sizeof(double) * factors * factors)
         b = <double *> malloc(sizeof(double) * factors)
         pivot = <int *> malloc(sizeof(int) * factors)
         try:
             for u in prange(users, schedule='guided'):
-                # For each user u calculate 
+                # For each user u calculate
                 # Xu = (YtCuY + regularization*I)i^-1 * YtYCuPu
 
                 # Build up A = YtCuY + reg * I and b = YtCuPu
@@ -42,7 +42,7 @@ def least_squares(Cui, double [:, :] X, double [:, :] Y, double regularization):
                 for index in range(indptr[u], indptr[u+1]):
                     i = indices[index]
                     confidence = data[index]
-                   
+
                     # b += Yi Cui Pui
                     # Pui is implicit, its defined to be 1 for non-zero entries
                     cython_blas.daxpy(&factors, &confidence, &Y[i, 0], &one, b, &one)

implicit/implicit.py

@@ -9,7 +9,7 @@
 
 
 def alternating_least_squares(Cui, factors, regularization=0.01,
-                              iterations=15, use_native=True):
+                              iterations=15, use_native=True, num_threads=0):
     """ factorizes the matrix Cui using an implicit alternating least squares
     algorithm
 
@@ -19,6 +19,8 @@ def alternating_least_squares(Cui, factors, regularization=0.01,
         regularization (double): Regularization parameter to use
         iterations (int): Number of alternating least squares iterations to
         run
+        num_threads (int): Number of threads to run least squares iterations.
+        0 means to use all CPU cores.
 
     Returns:
         tuple: A tuple of (row, col) factors
@@ -36,14 +38,14 @@ def alternating_least_squares(Cui, factors, regularization=0.01,
 
     for iteration in range(iterations):
         s = time.time()
-        solver(Cui, X, Y, regularization)
-        solver(Ciu, Y, X, regularization)
+        solver(Cui, X, Y, regularization, num_threads)
+        solver(Ciu, Y, X, regularization, num_threads)
         log.debug("finished iteration %i in %s", iteration, time.time() - s)
 
     return X, Y
 
 
-def least_squares(Cui, X, Y, regularization):
+def least_squares(Cui, X, Y, regularization, num_threads):
     """ For each user in Cui, calculate factors Xu for them
     using least squares on Y.