Telkamp7/issue3

NAMESPACE

@@ -1,9 +1,9 @@
 # Generated by roxygen2: do not edit by hand
 
+S3method(predict,aedseo)
 export("%>%")
 export(aedseo)
 export(fit_growth_rate)
-export(predict_growth_rate)
 export(tsd)
 importFrom(lifecycle,deprecated)
 importFrom(magrittr,"%>%")

R/aedseo.R

@@ -55,7 +55,6 @@
 #'
 #' # Print the AEDSEO results
 #' print(aedseo_results)
-#'
 aedseo <- function(
     tsd,
     k = 5,
@@ -106,5 +105,11 @@ aedseo <- function(
     )
   }
 
-  return(res)
+  # Turn the results into an `aedseo` class
+  ans <- tibble::new_tibble(
+    x = res,
+    class = "aedseo"
+  )
+
+  return(ans)
 }

R/predict_growth_rate.R → R/predict.R

@@ -12,10 +12,11 @@
 #' the result of the `aedseo()` function.
 #' @param n_step An integer specifying the number of future time steps for
 #' which you want to predict growth rates.
+#' @param ... Additional arguement affecting the predictions produced.
 #'
-#' @return  A tibble containing the predicted growth rates, including time,
-#' estimated growth rate, lower confidence interval, and upper confidence
-#' interval for the specified number of future time steps.
+#' @return  A tibble S3 object called `aedseo` containing the predicted growth
+#' rates, including time, estimated growth rate, lower confidence interval,
+#' and upper confidence interval for the specified number of future time steps.
 #' @export
 #'
 #' @importFrom rlang .data
@@ -43,14 +44,14 @@
 #' )
 #'
 #' # Predict growth rates for the next 5 time steps
-#' prediction <- predict_growth_rate(object = aedseo_results, n_step = 5)
+#' prediction <- predict(object = aedseo_results, n_step = 5)
 #'
 #' # Print the prediction
 #' print(prediction)
 #'
-predict_growth_rate <- function(object, n_step) {
+predict.aedseo <- function(object, n_step, ...) {
   # Calculate the prediction
-  prediction <- dplyr::last(object) %>%
+  ans <- dplyr::last(object) %>%
     dplyr::reframe(
       t = 0:n_step,
       time = .data$reference_time + t,
@@ -60,5 +61,5 @@ predict_growth_rate <- function(object, n_step) {
     )
 
   # Return
-  return(prediction)
+  return(ans)
 }

tests/testthat/test-predict_growth_rate.R → tests/testthat/test-predict.R

@@ -25,7 +25,7 @@ test_that("Returns the desired length", {
   n_step <- 5
 
   # Predict growth rates for the next 5 time steps
-  prediction <- predict_growth_rate(object = aedseo_results, n_step = n_step)
+  prediction <- predict(object = aedseo_results, n_step = n_step)
 
   # Return the number of prediction + the initial observation
   expect_length(prediction$estimate, n_step + 1)