fix: predict.bassetfit now accepts newdata=NULL argument

jsilve24 · May 3, 2024 · 006426c · 006426c
1 parent 36adb3d
commit 006426c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/R/bassetfit_s3.R b/R/bassetfit_s3.R
@@ -62,10 +62,14 @@ verify.bassetfit <- function(m, ...){
 #' 
 #' @export
 #' @importFrom stats median predict runif
-predict.bassetfit <- function(object, newdata, response="Lambda", size=NULL, 
+predict.bassetfit <- function(object, newdata=NULL, response="Lambda", size=NULL, 
                               use_names=TRUE, summary=FALSE, iter=NULL,
                               from_scratch=FALSE, ...){
   req(object, c("Lambda", "Sigma"))
+  if (is.null(newdata)) {
+    req(object, c("X"))
+    newdata <- object$X
+  } 
   newdata <- vec_to_mat(newdata)
 
   l <- store_coord(object)
@@ -198,4 +202,4 @@ predict.bassetfit <- function(object, newdata, response="Lambda", size=NULL,
   }
   if (response=="Y") return(Ypred)
   stop("response parameter not recognized")
-}
+}