Add basic class for nesmr/esmr

jean997 · Dec 16, 2024 · ecdd3b3 · ecdd3b3
1 parent 6d9ee8e
commit ecdd3b3
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 7 deletions.
diff --git a/R/esmr.R b/R/esmr.R
@@ -64,13 +64,19 @@ esmr <- function(beta_hat_X, se_X,
 
   dat <- set_data(beta_hat_Y, se_Y, beta_hat_X, se_X, R, ld_scores, RE, tau_init)
 
+  class(dat) <- c(c("esmr"), class(dat))
+  dat$is_nesmr <- ! is.null(direct_effect_template)
+  if (dat$is_nesmr) {
+    class(dat) <- c(c("nesmr"), class(dat))
+  }
+
   if(is.null(G)){
     if(dat$p == 2){
       G <- diag(dat$p)
-    }else if(!missing(direct_effect_template) & !is.null(direct_effect_template)){
+    } else if (dat$is_nesmr) {
       warning("Cannot estimate G for network problem yet.\n")
       G <- diag(dat$p)
-    }else{
+    } else{
       G <- estimate_G(beta_hat_X = dat$Y[,-1,drop =F],
                       se_X = dat$S[,-1, drop = F],
                       R = R[-1, -1, drop = FALSE],

diff --git a/R/grad_log_py.R b/R/grad_log_py.R
@@ -1,9 +1,9 @@
 ## likelihood function, very slow
-log_py <- function(fit, g_hat, fbar, max_prob = 1, nmax = Inf){
-  if(missing(g_hat)){
+log_py <- function(fit, g_hat = NULL, fbar = NULL, max_prob = 1, nmax = Inf){
+  if(is.null(g_hat)){
     g_hat <- fit$l$g_hat
   }
-  if(missing(fbar)){
+  if(is.null(fbar)){
     fbar <- fit$f$fbar
   }
   fgbar <-  fbar %*% fit$G

diff --git a/R/likelihood_funcs.R b/R/likelihood_funcs.R
@@ -42,8 +42,8 @@ calc_ell2 <- function(Y, abar, a2bar, fgbar, omega, omega_logdet, s_equal){
 
 
 #'@export
-logLik.esmr <- function(x) {
-  log_py(x)
+logLik.esmr <- function(x, ...) {
+  log_py(x, ...)
 }