using OpenMP safely

Author: pachadotdev

R/cpp11.R

@@ -4,6 +4,18 @@ center_variables_r_ <- function(V_r, w_r, klist, tol, maxiter, interrupt_iter) {
   .Call(`_capybara_center_variables_r_`, V_r, w_r, klist, tol, maxiter, interrupt_iter)
 }
 
+felm_fit_ <- function(y_r, x_r, wt_r, control, k_list) {
+  .Call(`_capybara_felm_fit_`, y_r, x_r, wt_r, control, k_list)
+}
+
+feglm_fit_ <- function(beta_r, eta_r, y_r, x_r, wt_r, theta, family, control, k_list) {
+  .Call(`_capybara_feglm_fit_`, beta_r, eta_r, y_r, x_r, wt_r, theta, family, control, k_list)
+}
+
+feglm_offset_fit_ <- function(eta_r, y_r, offset_r, wt_r, family, control, k_list) {
+  .Call(`_capybara_feglm_offset_fit_`, eta_r, y_r, offset_r, wt_r, family, control, k_list)
+}
+
 get_alpha_ <- function(p_r, klist, control) {
   .Call(`_capybara_get_alpha_`, p_r, klist, control)
 }
@@ -23,15 +35,3 @@ group_sums_var_ <- function(M_r, jlist) {
 group_sums_cov_ <- function(M_r, N_r, jlist) {
   .Call(`_capybara_group_sums_cov_`, M_r, N_r, jlist)
 }
-
-feglm_fit_ <- function(beta_r, eta_r, y_r, x_r, wt_r, theta, family, control, k_list) {
-  .Call(`_capybara_feglm_fit_`, beta_r, eta_r, y_r, x_r, wt_r, theta, family, control, k_list)
-}
-
-feglm_offset_fit_ <- function(eta_r, y_r, offset_r, wt_r, family, control, k_list) {
-  .Call(`_capybara_feglm_offset_fit_`, eta_r, y_r, offset_r, wt_r, family, control, k_list)
-}
-
-felm_fit_ <- function(y_r, x_r, wt_r, control, k_list) {
-  .Call(`_capybara_felm_fit_`, y_r, x_r, wt_r, control, k_list)
-}

src/00_main.h

@@ -8,6 +8,10 @@ using namespace cpp11;
 
 // used across the scripts
 
+#ifdef _OPENMP
+const size_t n_threads = omp_get_max_threads();
+#endif
+
 void center_variables_(mat &V, const vec &w, const list &klist,
                        const double &tol, const size_t &maxiter,
                        const size_t &interrupt_iter);

src/01_center_variables.cpp

@@ -0,0 +1,161 @@
+#include "00_main.h"
+
+// Halperin method with aggressive precomputing and optimized memory use
+void center_variables_(mat &V, const vec &w, const list &klist,
+                       const double &tol, const size_t &maxiter,
+                       const size_t &interrupt_iter) {
+  // Auxiliary variables (fixed)
+  const size_t N = V.n_rows;
+  const size_t P = V.n_cols;
+  const size_t K = klist.size();
+  const double inv_sw = 1.0 / accu(w);
+
+  // Auxiliary variables (storage)
+  vec x(N), x0(N);
+  field<field<uvec>> group_indices(K);
+  field<vec> group_inverse_weights(K);
+
+  // Precompute indices and inverse weights once
+  for (size_t k = 0; k < K; ++k) {
+    const list &jlist = klist[k];
+    size_t J = jlist.size();
+
+    field<uvec> indices(J);
+    vec inverse_weights(J);
+
+    for (size_t j = 0; j < J; ++j) {
+      indices(j) = as_uvec(as_cpp<integers>(jlist[j]));
+      inverse_weights(j) = 1.0 / accu(w.elem(indices(j)));
+    }
+
+    group_indices(k) = std::move(indices);
+    group_inverse_weights(k) = std::move(inverse_weights);
+  }
+
+// Perform Halperin projections, parallelizing columns
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static) num_threads(n_threads) private(x, x0)
+#endif
+  for (size_t p = 0; p < P; ++p) {
+    x = V.col(p);
+    size_t interrupt = interrupt_iter;
+
+    for (size_t iter = 0; iter < maxiter; ++iter) {
+      if (iter == interrupt) {
+        // Only main thread checks for interrupts
+#ifdef _OPENMP
+        if (omp_get_thread_num() == 0) {
+          check_user_interrupt();
+        }
+#else
+        check_user_interrupt();
+#endif
+        interrupt += interrupt_iter;
+      }
+
+      x0 = x;
+
+      // Project onto group means
+      for (size_t l = 0; l < K; ++l) {
+        const size_t L = group_indices(l).size();
+        for (size_t m = 0; m < L; ++m) {
+          const uvec &coords = group_indices(l)(m);
+          const double xbar =
+              dot(w.elem(coords), x.elem(coords)) * group_inverse_weights(l)(m);
+          x.elem(coords) -= xbar;
+        }
+      }
+
+      // Check convergence (correct placement)
+      double ratio = dot(abs(x - x0) / (1.0 + abs(x0)), w) * inv_sw;
+      if (ratio < tol) {
+        break;
+      }
+    }
+
+    // Assign back at convergence
+    V.col(p) = std::move(x);
+  }
+}
+
+[[cpp11::register]] doubles_matrix<> center_variables_r_(
+    const doubles_matrix<> &V_r, const doubles &w_r, const list &klist,
+    const double &tol, const int &maxiter, const int &interrupt_iter) {
+  mat V = as_Mat(V_r);
+  vec w = as_Col(w_r);
+  center_variables_(V, w, klist, tol, maxiter, interrupt_iter);
+  return as_doubles_matrix(std::move(V));
+}
+
+// Kaczmarz demeaning
+// void center_variables_(mat &V, const vec &w, const list &klist,
+//                        const double &tol, const size_t &maxiter,
+//                        const size_t &interrupt_iter) {
+//   // Auxiliary variables (fixed)
+//   const size_t P = V.n_cols;
+//   const size_t K = klist.size();
+//   const double inv_sw = 1.0 / accu(w);
+
+//   // Auxiliary variables (storage)
+//   size_t interrupt = static_cast<size_t>(interrupt_iter);
+//   uvec coords;
+
+//   // Precompute group indices and weights parallelizing over groups
+//   field<field<uvec>> group_indices(K);
+//   field<vec> group_inverse_weights(K);
+
+// #ifdef _OPENMP
+// #pragma omp parallel for schedule(static, n_threads)
+// #endif
+//   for (size_t k = 0; k < K; ++k) {
+//     const list &jlist = klist[k];
+//     size_t J = jlist.size();
+
+//     field<uvec> indices(J);
+//     vec inverse_weights(J);
+
+//     for (size_t j = 0; j < J; ++j) {
+//       indices(j) = as_uvec(as_cpp<integers>(jlist[j]));
+//       inverse_weights(j) = 1.0 / accu(w.elem(indices(j)));
+//     }
+
+//     group_indices(k) = std::move(indices);
+//     group_inverse_weights(k) = std::move(inverse_weights);
+//   }
+
+//   // Kaczmarz iterations parallelizing over columns
+// #ifdef _OPENMP
+// #pragma omp parallel for schedule(static, n_threads)
+// #endif
+//   for (size_t p = 0; p < P; ++p) {
+//     for (size_t iter = 0; iter < maxiter; ++iter) {
+//       if (iter == interrupt) {
+//         check_user_interrupt();
+//         interrupt += 1000;
+//       }
+
+//       vec x = V.col(p);
+//       vec x0 = x;
+//       double ratio;
+
+//       for (size_t l = 0; l < K; ++l) {
+//         size_t L = group_indices(l).size();
+//         if (L == 0) continue;
+
+//         for (size_t m = 0; m < L; ++m) {
+//           const uvec &coords = group_indices(l)(m);
+//           double xbar =
+//               dot(w.elem(coords), x.elem(coords)) / accu(w.elem(coords));
+//           x.elem(coords) -= xbar;
+//         }
+//       }
+
+//       ratio = dot(abs(x - x0) / (1.0 + abs(x0)), w) * inv_sw;
+//       if (ratio < tol) {
+//         break;
+//       }
+
+//       V.col(p) = x;
+//     }
+//   }
+// }

src/04_linear_algebra.cpp renamed to src/01_linear_algebra.cpp

@@ -42,13 +42,10 @@
 
   // Generate result list
 
-  writable::list out(4);
-  out[0] = as_doubles(beta);
-  out[1] = as_doubles(fitted);
-  out[2] = as_doubles(w);
-  out[3] = as_doubles_matrix(H);
-  out.attr("names") = writable::strings(
-      {"coefficients", "fitted.values", "weights", "hessian"});
-
-  return out;
+  return writable::list(
+      {"coefficients"_nm = as_doubles(std::move(beta)),
+       "fitted.values"_nm = as_doubles(std::move(fitted)),
+       "weights"_nm = as_doubles(std::move(w)),
+       "hessian"_nm = as_doubles_matrix(std::move(H))
+  });
 }

src/02_center_variables.cpp

@@ -307,26 +307,22 @@ bool valid_mu_(const vec &mu, const FamilyType &fam) {
 
   // Generate result list
 
-  writable::list out(8);
-
-  out[0] = as_doubles(beta);
-  out[1] = as_doubles(eta);
-  out[2] = as_doubles(wt);
-  out[3] = as_doubles_matrix(H);
-  out[4] = writable::doubles({dev});
-  out[5] = writable::doubles({null_dev});
-  out[6] = writable::logicals({conv});
-  out[7] = writable::integers({static_cast<int>(iter + 1)});
-
-  out.attr("names") =
-      writable::strings({"coefficients", "eta", "weights", "hessian",
-                         "deviance", "null_deviance", "conv", "iter"});
+  writable::list out({
+      "coefficients"_nm = as_doubles(std::move(beta)),
+       "eta"_nm = as_doubles(std::move(eta)),
+       "weights"_nm = as_doubles(std::move(wt)),
+       "hessian"_nm = as_doubles_matrix(std::move(H)),
+       "deviance"_nm = writable::doubles({dev}),
+       "null_deviance"_nm = writable::doubles({null_dev}),
+       "conv"_nm = writable::logicals({conv}),
+       "iter"_nm = writable::integers({static_cast<int>(iter + 1)})
+  });
 
   if (keep_mx) {
     mat x_cpp = as_Mat(x_r);
     center_variables_(x_cpp, w, k_list, center_tol, iter_center_max,
                       iter_interrupt);
-    out.push_back({"MX"_nm = as_doubles_matrix(x_cpp)});
+    out.push_back({"MX"_nm = as_doubles_matrix(std::move(x_cpp))});
   }
 
   return out;

src/07_lm_fit.cpp renamed to src/03_lm_fit.cpp

@@ -98,5 +98,5 @@ feglm_offset_fit_(const doubles &eta_r, const doubles &y_r,
     Myadj = Myadj - yadj;
   }
 
-  return as_doubles(eta);
+  return as_doubles(std::move(eta));
 }