Improve performance

ndunnewind · Oct 18, 2024 · 667deab · 667deab
1 parent fda33db
commit 667deab
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 44 deletions.
diff --git a/R/nma.R b/R/nma.R
@@ -728,10 +728,7 @@ nma <- function(network,
 
       idat_agd_arm <- idat_agd_arm %>%
         dplyr::group_by(.data$.study) %>%
-        dplyr::mutate(.mu = dplyr::if_else(
-          .data$.trt != ref_trt & ref_trt %in% .data$.trt,
-          as.numeric(.data$.study), 0
-        )) %>%
+        dplyr::mutate(.mu = as.integer(.data$.trt != ref_trt & ref_trt %in% .data$.trt)) %>%
         dplyr::ungroup()
     }
 
@@ -876,7 +873,11 @@ nma <- function(network,
     xbar <- NULL
   }
 
-  xbar_mu <- if (".mu" %in% all.vars(regression)) calculate_baseline_risk(network, link)
+  if (".mu" %in% all.vars(regression)) {
+    xbar_mu <- calculate_baseline_risk(network, link)
+  } else {
+    xbar_mu <- 0
+  }
 
   # Make NMA formula
   nma_formula <- make_nma_formula(regression,
@@ -1490,7 +1491,7 @@ nma.fit <- function(ipd_x, ipd_y,
   col_trt <- grepl("^(\\.trt|\\.contr)[^:]+$", x_names)
   col_omega <- x_names == ".omegaTRUE"
   col_reg <- !col_study & !col_trt & !col_omega
-  col_br <- col_reg & grepl("\\.mu", x_names)
+  col_brmr <- col_reg & grepl("\\.mu", x_names)
 
   n_trt <- sum(col_trt) + 1
 
@@ -1661,10 +1662,9 @@ nma.fit <- function(ipd_x, ipd_y,
     # Offsets
     has_offset = has_offsets,
     offsets = if (has_offsets) as.array(c(ipd_offset, agd_arm_offset, agd_contrast_offset)) else numeric(),
-    br_n = 0L,
-    br_index = matrix(0L, 0L, 2L),
-    br_mu = integer(),
-    mu_center = 0
+    brmr_n_col = sum(col_brmr),
+    brmr_col = which(col_brmr),
+    xbar_mu = xbar_mu
   )
 
   # Add priors
@@ -1719,24 +1719,6 @@ nma.fit <- function(ipd_x, ipd_y,
   # Set chain_id to make CHAIN_ID available in data block
   stanargs$chain_id <- 1L
 
-  # Baseline risk meta-regression
-  if (!is.null(xbar_mu)) {
-    stopifnot(length(col_br) == ncol(X_all))
-
-    br_index <- which(
-      t(t(X_all) != 0 & col_br),
-      arr.ind = TRUE
-    )
-
-    standat <- purrr::list_modify(
-      standat,
-      br_n = nrow(br_index),
-      br_index = br_index,
-      br_mu = as.array(as.integer(X_all[br_index])),
-      mu_center = xbar_mu
-    )
-  }
-
   # Call Stan model for given likelihood
 
   # -- Normal likelihood

diff --git a/inst/stan/binomial_1par.stan b/inst/stan/binomial_1par.stan
@@ -34,6 +34,10 @@ transformed parameters {
       X_agd_arm * beta_tilde + offset_agd_arm :
       X_agd_arm * beta_tilde;
 
+    if (brmr_n_col > 0) {
+      eta_agd_arm_noRE += (X_agd_arm[,1:totns] * mu - xbar_mu - 1) .* (X_agd_arm[,brmr_col] * beta_tilde[brmr_col]);
+    }
+
     if (nint_max > 1) { // -- If integration points are used --
       if (RE) {
         if (link == 1) { // logit link

diff --git a/inst/stan/include/data_common.stan b/inst/stan/include/data_common.stan
@@ -55,6 +55,11 @@ corr_matrix[RE ? max(which_RE) : 1] RE_cor; // RE correlation matrix
 // -- Node-splitting --
 int<lower=0, upper=1> nodesplit; // Node-splitting flag (yes = 1)
 
+// -- Baseline risk meta-regression --
+int<lower=0> brmr_n_col;
+array[brmr_n_col] int<lower=1> brmr_col;
+real xbar_mu;
+
 // -- Priors --
 int<lower=0,upper=3> prior_intercept_dist;
 real prior_intercept_location;
@@ -76,8 +81,3 @@ int<lower=0,upper=3> prior_reg_dist;
 real prior_reg_location;
 real<lower=0> prior_reg_scale;
 real<lower=0> prior_reg_df;
-
-int<lower=0> br_n;
-array[br_n, 2] int<lower=1> br_index;
-array[br_n] int<lower=1> br_mu;
-real mu_center;
diff --git a/inst/stan/include/transformed_data_common.stan b/inst/stan/include/transformed_data_common.stan
@@ -33,7 +33,11 @@ int totns = ns_ipd + ns_agd_arm; // + ns_agd_contrast;
 // All treatments vector
 array[narm_ipd + narm_agd_arm + ni_agd_contrast] int<lower=1> trt = append_array(append_array(ipd_trt, agd_arm_trt), agd_contrast_trt);
 
+// Split Q matrix or X matrix into IPD and AgD rows
 matrix[0, nX] Xdummy;
+matrix[ni_ipd, nX] X_ipd = ni_ipd ? X[1:ni_ipd] : Xdummy;
+matrix[nint_max * ni_agd_arm, nX] X_agd_arm = ni_agd_arm ? X[(ni_ipd + 1):(ni_ipd + nint_max * ni_agd_arm)] : Xdummy;
+matrix[nint_max * ni_agd_contrast, nX] X_agd_contrast = ni_agd_contrast ? X[(ni_ipd + nint_max * ni_agd_arm + 1):(ni_ipd + nint_max * (ni_agd_arm + ni_agd_contrast))] : Xdummy;
 
 // Split offsets into IPD and AgD rows
 vector[0] odummy;

diff --git a/inst/stan/include/transformed_parameters_common.stan b/inst/stan/include/transformed_parameters_common.stan
@@ -48,17 +48,6 @@ if (nodesplit) {
   omega[1] = allbeta[totns + nt];
 }
 
-// -- Baseline risk meta-regression --
-matrix[ni_ipd + nint_max * (ni_agd_arm + ni_agd_contrast), nX] X2 = X;
-for (i in 1:br_n) {
-  X2[br_index[i, 1], br_index[i, 2]] = mu[br_mu[i]] - mu_center;
-}
-
-// Split Q matrix or X matrix into IPD and AgD rows
-matrix[ni_ipd, nX] X_ipd = ni_ipd ? X2[1:ni_ipd] : Xdummy;
-matrix[nint_max * ni_agd_arm, nX] X_agd_arm = ni_agd_arm ? X2[(ni_ipd + 1):(ni_ipd + nint_max * ni_agd_arm)] : Xdummy;
-matrix[nint_max * ni_agd_contrast, nX] X_agd_contrast = ni_agd_contrast ? X2[(ni_ipd + nint_max * ni_agd_arm + 1):(ni_ipd + nint_max * (ni_agd_arm + ni_agd_contrast))] : Xdummy;
-
 // -- IPD model --
 // We define the IPD and AgD models here in the transformed parameters block,
 // as the linear predictors are required to calculate the log likelihood