From 13c3b3b15b8e9f0fca2ef583f3e8449e6dc9eda2 Mon Sep 17 00:00:00 2001
From: alexmccreight <57416850+alexmccreight@users.noreply.github.com>
Date: Tue, 21 Oct 2025 11:14:32 -0500
Subject: [PATCH] ash fix. unit testing fix

---
 R/sufficient_stats_methods.R | 23 ++++++++++++-----------
 R/susie_utils.R              |  7 ++++---
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/R/sufficient_stats_methods.R b/R/sufficient_stats_methods.R
index b074d105..66bc0e82 100644
--- a/R/sufficient_stats_methods.R
+++ b/R/sufficient_stats_methods.R
@@ -335,8 +335,8 @@ update_variance_components.ss <- function(data, params, model, ...) {
     # Update the sparse effect variance
     sparse_var <- mean(colSums(model$alpha * model$V))
 
-    # Update sigma2 conditional on sparse effect variance
-    mom_result <- mom_unmappable(data, params, model, omega, sparse_var, est_tau2 = FALSE, est_sigma2 = TRUE)
+    # Update sigma2
+    mom_result <- mom_unmappable(data, params, model, omega, sparse_var, est_tau2 = TRUE, est_sigma2 = TRUE)
 
     # Compute diagXtOmegaX and XtOmega for mr.ash using sparse effect variance and MoM residual variance
     omega_res <- compute_omega_quantities(data, sparse_var, mom_result$sigma2)
@@ -344,13 +344,13 @@ update_variance_components.ss <- function(data, params, model, ...) {
 
     # Create ash variance grid
     est_sa2 <- create_ash_grid(
-      PIP     = model$alpha,
-      mu      = model$mu,
-      omega   = omega,
-      tausq   = sparse_var,
-      sigmasq = mom_result$sigma2,
-      n       = data$n
-    )
+       PIP     = model$alpha,
+       mu      = model$mu,
+       omega   = omega,
+       tausq   = sparse_var,
+       sigmasq = mom_result$sigma2,
+       n       = data$n
+     )
 
     # Call mr.ash directly with pre-computed quantities
     mrash_output <- mr.ash.alpha.mccreight::mr.ash(
@@ -372,9 +372,10 @@ update_variance_components.ss <- function(data, params, model, ...) {
 
     return(list(
       sigma2 = mrash_output$sigma2,
-      tau2   = sum(est_sa2 * mrash_output$pi),
+      tau2   = sum(mrash_output$data$sa2 * mrash_output$pi),
       theta  = mrash_output$beta,
-      ash_pi = mrash_output$pi
+      ash_pi = mrash_output$pi,
+      sa2    = mrash_output$data$sa2
     ))
   } else {
     # Use default method for standard SuSiE
diff --git a/R/susie_utils.R b/R/susie_utils.R
index e9880d81..25d2eaa6 100644
--- a/R/susie_utils.R
+++ b/R/susie_utils.R
@@ -46,7 +46,8 @@ muffled_cov2cor <- function(x) {
 # Check for symmetric matrix.
 #' @keywords internal
 is_symmetric_matrix <- function(x) {
-  if (requireNamespace("Rfast", quietly = TRUE)) {
+  if (is.matrix(x) && is.numeric(x) && !isS4(x) &&
+      requireNamespace("Rfast", quietly = TRUE)) {
     return(Rfast::is.symmetric(x))
   } else {
     return(Matrix::isSymmetric(x))
@@ -976,12 +977,12 @@ create_ash_grid <- function(PIP, mu, omega, tausq, sigmasq, n,
   L_eff <- max(L_eff, 1)
 
   # Set lower bound
-  s2_min <- sigmasq / (n * max(h2_snp, 0.1))
+  s2_min <- sigmasq / n
   s2_min <- max(s2_min, 1e-8)
 
   # Set upper bound
   var_y <- sigmasq / (1 - h2_snp)
-  s2_max <- (h2_snp / sqrt(L_eff)) * var_y
+  s2_max <- h2_snp * var_y
   s2_max <- max(s2_max, 10 * s2_min)
 
   # Quantile-based adaptive grid