HartiganWong.hpp in umappp-0.2.0

- old
+ new

@@ -104,17 +104,23 @@
     /** 
      * @brief Default parameter values for `HartiganWong`.
      */
     struct Defaults {
         /** 
-         * See `HartiganWong::set_max_iterations()`.
+         * See `set_max_iterations()` for more details.
          */
         static constexpr int max_iterations = 10;
+
+        /** 
+         * See `set_num_threads()` for more details.
+         */
+        static constexpr int num_threads = 1;
     };
 
 private:
     int maxiter = Defaults::max_iterations;
+    int nthreads = Defaults::num_threads;
 
 public:
     /**
      * @param m Maximum number of iterations.
      * More iterations increase the opportunity for convergence at the cost of more computational time.
@@ -124,10 +130,20 @@
     HartiganWong& set_max_iterations(int m = Defaults::max_iterations) {
         maxiter = m;
         return *this;
     }
 
+    /**
+     * @param n Number of threads to use.
+     *
+     * @return A reference to this `HartiganWong` object.
+     */
+    HartiganWong& set_num_threads(int n = Defaults::num_threads) {
+        nthreads = n;
+        return *this;
+    }
+
 public:
     Details<DATA_t, INDEX_t> run(int ndim, INDEX_t nobs, const DATA_t* data, CLUSTER_t ncenters, DATA_t* centers, CLUSTER_t* clusters) {
         num_dim = ndim;
         num_obs = nobs;
         data_ptr = data;
@@ -158,12 +174,17 @@
         }
 
         /* For each point I, find its two closest centres, IC1(I) and 
          * IC2(I). Assign it to IC1(I). 
          */
-        #pragma omp parallel for
+#ifndef KMEANS_CUSTOM_PARALLEL
+        #pragma omp parallel for num_threads(nthreads)
         for (INDEX_t obs = 0; obs < num_obs; ++obs) {
+#else
+        KMEANS_CUSTOM_PARALLEL(num_obs, [&](INDEX_t first, INDEX_t last) -> void {
+        for (INDEX_t obs = first; obs < last; ++obs) {
+#endif
             auto& best = ic1[obs];
             best = 0;
             DATA_t best_dist = squared_distance_from_cluster(obs, best);
 
             auto& second = ic2[obs];
@@ -184,10 +205,15 @@
                         std::swap(best_dist, second_dist);
                         std::swap(best, second);
                     }
                 }
             }
+#ifndef KMEANS_CUSTOM_PARALLEL
         }
+#else
+        }
+        }, nthreads);
+#endif
 
         /* Update cluster centres to be the average of points contained 
          * within them. 
          * NC(L) := #{units in cluster L},  L = 1..K 
          */