Improves Pseudolikelihood computation performance

.github/workflows/test-coverage.yaml

@@ -41,7 +41,7 @@ jobs:
 
       - name: Upload test results
         if: failure()
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: coverage-test-failures
           path: ${{ runner.temp }}/package

DESCRIPTION

@@ -1,7 +1,7 @@
 Package: mrf2d
 Type: Package
 Title: Markov Random Field Models for Image Analysis
-Version: 1.0.1
+Version: 1.0.2
 Authors@R:
     person("Victor", "Freguglia", email = "victorfreguglia@gmail.com",
              role = c("aut", "cre"),
@@ -29,7 +29,7 @@ LinkingTo: Rcpp,
            RcppArmadillo
 Roxygen: list(markdown = TRUE)
 RdMacros: Rdpack
-RoxygenNote: 7.2.3
+RoxygenNote: 7.3.3
 Suggests:
     testthat (>= 2.1.0),
     covr,

src/Makevars

@@ -0,0 +1,2 @@
+PKG_CXXFLAGS = $(SHLIB_OPENMP_CXXFLAGS)
+PKG_LIBS = $(SHLIB_OPENMP_CXXFLAGS)

src/Makevars.win

@@ -0,0 +1,2 @@
+PKG_CXXFLAGS = $(SHLIB_OPENMP_CXXFLAGS)
+PKG_LIBS = $(SHLIB_OPENMP_CXXFLAGS)

src/conditional_probability.cpp

@@ -15,22 +15,26 @@ NumericVector conditional_probabilities_mrf(const IntegerMatrix &Z,
                                             const int N, const int M,
                                             const int n_R, const int C){
 
-  NumericVector probs(C+1);
-  double this_prob;
+  NumericVector log_pots(C+1, 0.0);
   int dx, dy;
   int x = position[0] -1; int y = position[1] -1;
 
-  for(int value = 0; value <= C; value++){
-    this_prob = 0.0;
-    for(int i = 0; i < n_R; i++){
-      dx = R(i,0); dy = R(i,1);
-      if(0 <= x+dx && x+dx < N && 0 <= y+dy && y+dy < M){
-        this_prob = this_prob + theta(value, Z(x+dx, y+dy), i);}
-      if(0 <= x-dx && x-dx < N && 0 <= y-dy && y-dy < M){
-        this_prob = this_prob + theta(Z(x-dx, y-dy), value, i);}
+  for(int i = 0; i < n_R; i++){
+    dx = R(i,0); dy = R(i,1);
+    if(0 <= x+dx && x+dx < N && 0 <= y+dy && y+dy < M){
+      int nb = Z(x+dx, y+dy);
+      for(int value = 0; value <= C; value++)
+        log_pots[value] += theta(value, nb, i);
+    }
+    if(0 <= x-dx && x-dx < N && 0 <= y-dy && y-dy < M){
+      int nb = Z(x-dx, y-dy);
+      for(int value = 0; value <= C; value++)
+        log_pots[value] += theta(nb, value, i);
     }
-    probs[value] = exp(this_prob);
   }
+  // log-sum-exp trick: subtract max before exponentiation to prevent overflow
+  double max_lp = max(log_pots);
+  NumericVector probs = exp(log_pots - max_lp);
   return(probs/sum(probs));
 }
 
@@ -44,24 +48,28 @@ NumericVector conditional_probabilities_mrf_sub(const IntegerMatrix &Z,
                                                 const int N, const int M,
                                                 const int n_R, const int C){
 
-  NumericVector probs(C+1);
-  double this_prob;
+  NumericVector log_pots(C+1, 0.0);
   int dx, dy;
   int x = position[0] -1; int y = position[1] -1;
 
-  for(int value = 0; value <= C; value++){
-    this_prob = 0.0;
-    for(int i = 0; i < n_R; i++){
-      dx = R(i,0); dy = R(i,1);
-      if(0 <= x+dx && x+dx < N && 0 <= y+dy && y+dy < M){
-        if(sub_mat(x+dx, y+dy)){
-          this_prob = this_prob + theta(value, Z(x+dx, y+dy), i);}}
-      if(0 <= x-dx && x-dx < N && 0 <= y-dy && y-dy < M){
-        if(sub_mat(x-dx, y-dy)){
-          this_prob = this_prob + theta(Z(x-dx, y-dy), value, i);}}
-    }
-    probs[value] = exp(this_prob);
+  for(int i = 0; i < n_R; i++){
+    dx = R(i,0); dy = R(i,1);
+    if(0 <= x+dx && x+dx < N && 0 <= y+dy && y+dy < M){
+      if(sub_mat(x+dx, y+dy)){
+        int nb = Z(x+dx, y+dy);
+        for(int value = 0; value <= C; value++)
+          log_pots[value] += theta(value, nb, i);
+      }}
+    if(0 <= x-dx && x-dx < N && 0 <= y-dy && y-dy < M){
+      if(sub_mat(x-dx, y-dy)){
+        int nb = Z(x-dx, y-dy);
+        for(int value = 0; value <= C; value++)
+          log_pots[value] += theta(nb, value, i);
+      }}
   }
+  // log-sum-exp trick: subtract max before exponentiation to prevent overflow
+  double max_lp = max(log_pots);
+  NumericVector probs = exp(log_pots - max_lp);
   return(probs/sum(probs));
 }
 

src/pseudo_likelihood.cpp

@@ -1,9 +1,89 @@
 #include <RcppArmadillo.h>
 #include "mrf2d.h"
+#ifdef _OPENMP
+#include <omp.h>
+#endif
 
 using namespace Rcpp;
 
+// Computes log P(Z_ij = zij | Z_rest) directly in log-space:
+//   log P = s_zij - log_sum_exp(s_0, ..., s_C)
+// where s_v = sum of theta potentials for value v with all neighbors.
+// Avoids: per-pixel NumericVector allocation, redundant exp/log roundtrip,
+// and numerical overflow via the log-sum-exp trick.
+static inline double log_cond_pixel(const IntegerMatrix &Z,
+                                    const IntegerMatrix &R,
+                                    const arma::fcube &theta,
+                                    const int N, const int M,
+                                    const int n_R, const int C,
+                                    const int i, const int j,
+                                    std::vector<double> &log_probs){
+  int dx, dy;
+
+  for(int v = 0; v <= C; v++) log_probs[v] = 0.0;
+
+  // Neighbors outer, values inner:
+  // - Z(neighbor) read once per direction instead of C+1 times
+  // - bounds check done once per direction
+  // - theta(v, nb, r) with v sequential = cache-friendly (Armadillo col-major)
+  for(int r = 0; r < n_R; r++){
+    dx = R(r, 0); dy = R(r, 1);
+    if(i+dx >= 0 && i+dx < N && j+dy >= 0 && j+dy < M){
+      int nb = Z(i+dx, j+dy);
+      for(int v = 0; v <= C; v++)
+        log_probs[v] += theta(v, nb, r);
+    }
+    if(i-dx >= 0 && i-dx < N && j-dy >= 0 && j-dy < M){
+      int nb = Z(i-dx, j-dy);
+      for(int v = 0; v <= C; v++)
+        log_probs[v] += theta(nb, v, r);
+    }
+  }
+
+  double max_lp = *std::max_element(log_probs.begin(), log_probs.end());
+  double sum_exp = 0.0;
+  for(int v = 0; v <= C; v++)
+    sum_exp += std::exp(log_probs[v] - max_lp);
+
+  return log_probs[Z(i, j)] - max_lp - std::log(sum_exp);
+}
+
+static inline double log_cond_pixel_sub(const IntegerMatrix &Z,
+                                        const LogicalMatrix &sub_mat,
+                                        const IntegerMatrix &R,
+                                        const arma::fcube &theta,
+                                        const int N, const int M,
+                                        const int n_R, const int C,
+                                        const int i, const int j,
+                                        std::vector<double> &log_probs){
+  int dx, dy;
+
+  for(int v = 0; v <= C; v++) log_probs[v] = 0.0;
+
+  for(int r = 0; r < n_R; r++){
+    dx = R(r, 0); dy = R(r, 1);
+    if(i+dx >= 0 && i+dx < N && j+dy >= 0 && j+dy < M && sub_mat(i+dx, j+dy)){
+      int nb = Z(i+dx, j+dy);
+      for(int v = 0; v <= C; v++)
+        log_probs[v] += theta(v, nb, r);
+    }
+    if(i-dx >= 0 && i-dx < N && j-dy >= 0 && j-dy < M && sub_mat(i-dx, j-dy)){
+      int nb = Z(i-dx, j-dy);
+      for(int v = 0; v <= C; v++)
+        log_probs[v] += theta(nb, v, r);
+    }
+  }
+
+  double max_lp = *std::max_element(log_probs.begin(), log_probs.end());
+  double sum_exp = 0.0;
+  for(int v = 0; v <= C; v++)
+    sum_exp += std::exp(log_probs[v] - max_lp);
+
+  return log_probs[Z(i, j)] - max_lp - std::log(sum_exp);
+}
+
 // [[Rcpp::depends(RcppArmadillo)]]
+// [[Rcpp::plugins(openmp)]]
 // [[Rcpp::export]]
 double log_pl_mrf(const IntegerMatrix Z,
                   const IntegerMatrix R,
@@ -13,22 +93,31 @@ double log_pl_mrf(const IntegerMatrix Z,
   int n_R = R.nrow();
   int C = theta.n_rows - 1;
   double log_pl = 0.0;
-  double this_cond_prob = 0.0;
-  IntegerVector position(2);
-  int zij = -1;
 
+#ifdef _OPENMP
+  #pragma omp parallel reduction(+:log_pl)
+  {
+    std::vector<double> log_probs(C + 1); // thread-local buffer
+    #pragma omp for collapse(2) schedule(static)
+    for(int i = 0; i < N; i++){
+      for(int j = 0; j < M; j++){
+        log_pl += log_cond_pixel(Z, R, theta, N, M, n_R, C, i, j, log_probs);
+      }
+    }
+  }
+#else
+  std::vector<double> log_probs(C + 1);
   for(int i = 0; i < N; i++){
     for(int j = 0; j < M; j++){
-      zij = Z(i,j);
-      position[0] = i+1; position[1] = j+1;
-      this_cond_prob = conditional_probabilities_mrf(Z, position, R, theta, N, M, n_R, C)(zij);
-      log_pl += log(this_cond_prob);
+      log_pl += log_cond_pixel(Z, R, theta, N, M, n_R, C, i, j, log_probs);
     }
   }
+#endif
   return(log_pl);
 }
 
 // [[Rcpp::depends(RcppArmadillo)]]
+// [[Rcpp::plugins(openmp)]]
 // [[Rcpp::export]]
 double log_pl_mrf_sub(const IntegerMatrix Z,
                       const LogicalMatrix sub_mat,
@@ -39,19 +128,29 @@ double log_pl_mrf_sub(const IntegerMatrix Z,
   int n_R = R.nrow();
   int C = theta.n_rows - 1;
   double log_pl = 0.0;
-  double this_cond_prob = 0.0;
-  IntegerVector position(2);
-  int zij = -1;
 
+#ifdef _OPENMP
+  #pragma omp parallel reduction(+:log_pl)
+  {
+    std::vector<double> log_probs(C + 1); // thread-local buffer
+    #pragma omp for collapse(2) schedule(static)
+    for(int i = 0; i < N; i++){
+      for(int j = 0; j < M; j++){
+        if(sub_mat(i,j)){
+          log_pl += log_cond_pixel_sub(Z, sub_mat, R, theta, N, M, n_R, C, i, j, log_probs);
+        }
+      }
+    }
+  }
+#else
+  std::vector<double> log_probs(C + 1);
   for(int i = 0; i < N; i++){
     for(int j = 0; j < M; j++){
       if(sub_mat(i,j)){
-        zij = Z(i,j);
-        position[0] = i+1; position[1] = j+1;
-        this_cond_prob = conditional_probabilities_mrf_sub(Z, sub_mat, position, R, theta, N, M, n_R, C)(zij);
-        log_pl += log(this_cond_prob);
+        log_pl += log_cond_pixel_sub(Z, sub_mat, R, theta, N, M, n_R, C, i, j, log_probs);
       }
     }
   }
+#endif
   return(log_pl);
 }