defragmented temporary memory allocations in sparse_dot_topn routines

Author: ParticularMiner

setup.py

@@ -35,6 +35,7 @@ def finalize_options(self):
                                     './sparse_dot_topn/sparse_dot_topn_source.cpp'
                                 ],
                          extra_compile_args=extra_compile_args,
+                         define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')],
                          language='c++')
 
 original_ext = Extension('sparse_dot_topn.sparse_dot_topn',
@@ -43,6 +44,7 @@ def finalize_options(self):
                                     './sparse_dot_topn/sparse_dot_topn_source.cpp'
                                 ],
                          extra_compile_args=extra_compile_args,
+                         define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')],
                          language='c++')
 
 threaded_ext = Extension('sparse_dot_topn.sparse_dot_topn_threaded',
@@ -52,6 +54,7 @@ def finalize_options(self):
                              './sparse_dot_topn/sparse_dot_topn_parallel.cpp'
                             ],
                          extra_compile_args=extra_compile_args,
+                         define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')],
                          language='c++')
 
 setup(

sparse_dot_topn/awesome_cossim_topn.py

@@ -25,23 +25,25 @@ def awesome_cossim_topn(
     """
     This function will return a matrix C in CSR format, where
     C = [sorted top n results > lower_bound for each row of A * B].
-    If return_best_topn=True it will also return best_topn (the 
-    true maximum number of elements > lower_bound per row of A * B).
+    If return_best_topn=True then best_topn
+    (the true maximum number of elements > lower_bound per row of A * B)
+    will also be returned in a tuple together with C as (C, best_topn).
 
     Input:
         A and B: two CSR matrices
-        ntop: n top results
-        lower_bound: a threshold that the element of A*B must greater than
-                     use_threads: use multi-thread or not
+        ntop: top n results
+        lower_bound: a threshold that the element of A*B must be greater than
+        use_threads: use multi-thread or not
         n_jobs: number of thread, must be >= 1
-        ntop_is_flexible: if True, memory management will be handed over to C/C++ if 
-                          python's attempt at allocating memory fails.
-        mem_manager_is_C: (this is mainly for testing purposes) if True, will force
-                          memory management to be handed over to C/C++. Should be
-                          used only when ntop >= number of columns of B or 
-                          ntop_is_flexible=True.  Defaults to False.
-        return_best_topn: if True, will return best_topn together with C as a tuple:
-                          (C, best_topn)
+        ntop_is_flexible: (default: False) if True, memory management will be handed 
+                          over to C/C++ whenever python's attempt at allocating
+                          memory fails.
+        mem_manager_is_C: (default: False) this is mainly for testing purposes. if 
+                          True, will force memory management to be handed over to
+                          C/C++. Should be used only when ntop >= number of columns 
+                          of B or ntop_is_flexible=True.
+        return_best_topn: (default: False) if True, will return best_topn together 
+                          with C as a tuple: (C, best_topn)
 
     Output:
         C: result matrix (returned alone, if return_best_topn=False)
@@ -80,11 +82,13 @@ def awesome_cossim_topn(
             return output
 
     # filled matrices from here on
-    indptr = np.empty(M+1, dtype=idx_dtype)
+    indptr = np.empty(M + 1, dtype=idx_dtype)
     try:
         indices = np.empty(nnz_max, dtype=idx_dtype)
         data = np.empty(nnz_max, dtype=A.dtype)
+        
         if mem_manager_is_C: raise MemoryError    # This is mainly for testing purposes
+        
     except MemoryError:
         # if mem_manager_is_C: print('Exception raised! Continuing ...', flush=True)
         if ntop_is_flexible or ntop >= N:
@@ -107,7 +111,6 @@ def awesome_cossim_topn(
                     lower_bound,
                     indptr
                 )
-                
             else:
 
                 indices, data, best_topn = ct_thread.sparse_dot_free_threaded(
@@ -120,14 +123,19 @@ def awesome_cossim_topn(
                     lower_bound,
                     indptr, n_jobs
                 )
-
         else:
+
             if mem_manager_is_C:
-                raise Exception('When mem_manager_is_C=True, set ntop >= N, or set ntop_is_flexible=True')
+                raise Exception(
+                    'When mem_manager_is_C=True, set ntop >= B.shape[1], or set ntop_is_flexible=True'
+                )
             else:
-                raise Exception('Not enough memory!  Data array is too large. Try reducing the value of ntop.')
-            
+                raise Exception(
+                    'Not enough memory!  Data array is too large. Try reducing the value of ntop.'
+                    'or set ntop_is_flexible=True'
+                )
     else:
+        # no exception was raised; then use old function (as it is expected to be the fastest)
 
         best_topn_arr = np.full(1, 0, dtype=idx_dtype)
 
@@ -144,7 +152,6 @@ def awesome_cossim_topn(
                 lower_bound,
                 indptr, indices, data, best_topn_arr
             )
-    
         else:
             if n_jobs < 1:
                 err_str = 'Whenever you select the multi-thread mode, n_job must be greater than or equal to 1!'
@@ -161,9 +168,9 @@ def awesome_cossim_topn(
                 lower_bound,
                 indptr, indices, data, best_topn_arr, n_jobs
             )
-        
         best_topn = best_topn_arr[0]
 
+    # prepare and return the output:
     output = csr_matrix((data, indices, indptr), shape=(M, N))
     if return_best_topn:
         return output, best_topn