diff --git a/DESCRIPTION b/DESCRIPTION
index ee7681e..7c1f5aa 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,7 +1,7 @@
 Package: miaTime
 Type: Package
 Title: Microbiome Time Series Analysis
-Version: 0.99.1
+Version: 0.99.2
 Authors@R:
      c(person(given = "Leo", family = "Lahti", role = c("aut"),
               email = "leo.lahti@iki.fi",
diff --git a/R/getBaselineDivergence.R b/R/getBaselineDivergence.R
index 8cb7881..f0b225c 100644
--- a/R/getBaselineDivergence.R
+++ b/R/getBaselineDivergence.R
@@ -54,11 +54,9 @@
 #' \code{colData} that identifies the baseline samples to be used.
 #' (Default: \code{NULL})
 #'
-#' @param name \code{Character scalar}. Specifies a column name for storing
-#' divergence results. (Default: \code{"divergence"})
-#'
-#' @param name.time \code{Character scalar}. Specifies a column name for storing
-#' time differences. (Default: \code{"time_diff"})
+#' @param name \code{Character vector}. Specifies a column name for storing
+#' divergence results.
+#' (Default: \code{c("divergence", "time_diff", "ref_samples")})
 #'
 #' @param ... Optional arguments passed into
 #' \code{\link[mia:addDivergence]{mia::addDivergence()}}.
@@ -86,8 +84,8 @@
 #'     reference = "reference",
 #'     group = "subject",
 #'     time.col = "time",
-#'     name = "divergence_from_baseline",
-#'     name.time = "time_from_baseline",
+#'     name = c("divergence_from_baseline",
+#'         "time_from_baseline", "reference_samples"),
 #'     assay.type = "relabundance",
 #'     method = "bray")
 #'
@@ -153,7 +151,7 @@ setMethod("getBaselineDivergence", signature = c(x = "SummarizedExperiment"),
         reference <- args[["reference"]]
         time_res <- .get_time_difference(x, time.col, reference)
         # Create a DF to return
-        res <- .convert_divergence_to_df(x, res, time_res, ...)
+        res <- .convert_divergence_to_df(x, res, time_res, reference, ...)
         return(res)
     }
 )
@@ -161,12 +159,14 @@ setMethod("getBaselineDivergence", signature = c(x = "SummarizedExperiment"),
 #' @rdname addBaselineDivergence
 #' @export
 setMethod("addBaselineDivergence", signature = c(x = "SummarizedExperiment"),
-    function(x, name = "divergence", name.time = "time_diff", ...){
+    function(
+        x, name = c("divergence", "time_diff", "ref_samples"), ...){
+        .check_input(name, c("character vector"), length = 3L)
         # Calculate divergence
         res <- getBaselineDivergence(x, ...)
         # Add to colData
         res <- as.list(res) |> unname()
-        x <- .add_values_to_colData(x, res, list(name, name.time), ...)
+        x <- .add_values_to_colData(x, res, name, ...)
         return(x)
     }
 )
@@ -321,7 +321,7 @@ setMethod("addBaselineDivergence", signature = c(x = "SummarizedExperiment"),
     return(res)
 }
 
-# This function get time difference between a sample and its referene sample
+# This function get time difference between a sample and its reference sample
 .get_time_difference <- function(x, time.col, reference){
     # Get timepoints
     time_point <- x[[time.col]]
@@ -334,13 +334,12 @@ setMethod("addBaselineDivergence", signature = c(x = "SummarizedExperiment"),
 
 # This function converts time divergence results to DF object
 .convert_divergence_to_df <- function(
-        x, res, time_res, name = "divergence", name.time = "time_diff", ...){
-    #
-    temp <- .check_input(name, list("character scalar"))
-    #
-    temp <- .check_input(name.time, list("character scalar"))
+        x, res, time_res, reference,
+        name = c("divergence", "time_diff", "ref_samples"), ...){
+    # Validate 'name' param
+    temp <- .check_input(name, list("character vector"), length = 3L)
     #
-    df <- DataFrame(res, time_res, row.names = colnames(x))
-    colnames(df) <- c(name, name.time)
+    df <- DataFrame(res, time_res, x[[reference]], row.names = colnames(x))
+    colnames(df) <- name
     return(df)
 }
diff --git a/R/getStepwiseDivergence.R b/R/getStepwiseDivergence.R
index 6d4bbad..6237fa9 100644
--- a/R/getStepwiseDivergence.R
+++ b/R/getStepwiseDivergence.R
@@ -1,13 +1,13 @@
 #' @name addStepwiseDivergence
 #' @export
-#' 
+#'
 #' @title
 #' Beta diversity between consecutive time steps
-#' 
+#'
 #' @description
 #' Calculates sample dissimilarity between consecutive time points along with
 #' time difference.
-#' 
+#'
 #' @details
 #' These functions calculate time-wise divergence, meaning each sample is
 #' compared to the previous i-th sample, where i is the specified time
@@ -15,17 +15,17 @@
 #' divergence by comparing all samples with each other. However, it is often
 #' more meaningful to calculate divergence within a specific patient or group
 #' (see the \code{group} parameter).
-#' 
+#'
 #' @return
 #' \code{getStepwiseDivergence} returns \code{DataFrame} object
 #' containing the sample dissimilarity and corresponding time difference between
 #' samples. \code{addStepwiseDivergence}, on the other hand, returns a
 #' \code{\link[SummarizedExperiment:SummarizedExperiment-class]{SummarizedExperiment}}
 #' object with these results in its \code{colData}.
-#' 
+#'
 #' @inheritParams addBaselineDivergence
-#' 
-#' @param time.interval \code{Integer scalar}. Indicates the increment between 
+#'
+#' @param time.interval \code{Integer scalar}. Indicates the increment between
 #' time steps. By default, the function compares each sample to the
 #' previous one. If you need to take every second, every third, or so, time
 #' step, then increase this accordingly. (Default: \code{1L})
@@ -35,7 +35,7 @@
 #'
 #' data(hitchip1006)
 #' tse <- transformAssay(hitchip1006, method = "relabundance")
-#' 
+#'
 #' # Calculate divergence
 #' tse <- addStepwiseDivergence(
 #'     tse,
@@ -44,10 +44,10 @@
 #'     time.col = "time",
 #'     assay.type = "relabundance"
 #'     )
-#' 
+#'
 #' @seealso
 #' \code{\link[mia:addDivergence]{mia::addDivergence()}}
-#' 
+#'
 NULL
 
 #' @rdname addStepwiseDivergence
@@ -102,7 +102,7 @@ setMethod("getStepwiseDivergence", signature = c(x = "ANY"),
         reference <- args[["reference"]]
         time_res <- .get_time_difference(x, time.col, reference)
         # Create a DF to return
-        res <- .convert_divergence_to_df(x, res, time_res, ...)
+        res <- .convert_divergence_to_df(x, res, time_res, reference, ...)
         return(res)
     }
 )
@@ -110,12 +110,14 @@ setMethod("getStepwiseDivergence", signature = c(x = "ANY"),
 #' @rdname addStepwiseDivergence
 #' @export
 setMethod("addStepwiseDivergence", signature = c(x = "SummarizedExperiment"),
-    function(x, name = "divergence", name.time = "time_diff", ...){
+    function(
+        x, name = c("divergence", "time_diff", "ref_samples"), ...){
+        .check_input(name, c("character vector"), length = 3L)
         # Calculate divergence
         res <- getStepwiseDivergence(x,  ...)
         # Add to colData
         res <- as.list(res) |> unname()
-        x <- .add_values_to_colData(x, res, list(name, name.time), ...)
+        x <- .add_values_to_colData(x, res, name, ...)
         return(x)
     }
 )
diff --git a/R/utils.R b/R/utils.R
index c096ef7..c9ebfc3 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -6,7 +6,7 @@
 # Borrowed from HoloFoodR.
 .check_input <- function(
         variable, supported_class, supported_values = NULL, limits = NULL,
-        variable_name = .get_name_in_parent(variable)){
+        length = NULL, variable_name = .get_name_in_parent(variable)){
     # Convert supported classes to character
     classes_char <- lapply(supported_class, function(class){
         if( is.null(class) ){
@@ -54,6 +54,11 @@
         msg <- paste0(msg, ")")
     }
 
+    # If length was provided
+    if( !is.null(length) ){
+        msg <- paste0(msg, " The length must be ", length, ".")
+    }
+
     # List all the input types. Run the check if the variable must be that type.
     # If correct type was found, change the result to TRUE.
     input_correct <- FALSE
@@ -128,6 +133,10 @@
             input_correct <- FALSE
         }
     }
+    # Check length if provided
+    if( !is.null(length) && length(variable) != length ){
+        input_correct <- FALSE
+    }
     # Give error if variable was not correct type
     if( !input_correct ){
         stop(msg, call. = FALSE)
diff --git a/man/addBaselineDivergence.Rd b/man/addBaselineDivergence.Rd
index 48a03fd..ed798f9 100644
--- a/man/addBaselineDivergence.Rd
+++ b/man/addBaselineDivergence.Rd
@@ -21,7 +21,11 @@ addBaselineDivergence(x, ...)
   ...
 )
 
-\S4method{addBaselineDivergence}{SummarizedExperiment}(x, name = "divergence", name.time = "time_diff", ...)
+\S4method{addBaselineDivergence}{SummarizedExperiment}(
+  x,
+  name = c("divergence", "time_diff", "ref_samples"),
+  ...
+)
 }
 \arguments{
 \item{x}{A
@@ -49,11 +53,9 @@ used in the dissimilarity estimation. (Default: \code{"counts"})}
 Method is passed to the function that is specified by \code{dis.fun}.
 (Default: \code{"bray"})}
 
-\item{name}{\code{Character scalar}. Specifies a column name for storing
-divergence results. (Default: \code{"divergence"})}
-
-\item{name.time}{\code{Character scalar}. Specifies a column name for storing
-time differences. (Default: \code{"time_diff"})}
+\item{name}{\code{Character vector}. Specifies a column name for storing
+divergence results.
+(Default: \code{c("divergence", "time_diff", "ref_samples")})}
 }
 \value{
 \code{getBaselineDivergence} returns \code{DataFrame} object
@@ -106,8 +108,8 @@ tse <- addBaselineDivergence(
     reference = "reference",
     group = "subject",
     time.col = "time",
-    name = "divergence_from_baseline",
-    name.time = "time_from_baseline",
+    name = c("divergence_from_baseline", 
+        "time_from_baseline", "reference_samples"),
     assay.type = "relabundance",
     method = "bray")
 
diff --git a/man/addStepwiseDivergence.Rd b/man/addStepwiseDivergence.Rd
index d8ddd9c..20f948f 100644
--- a/man/addStepwiseDivergence.Rd
+++ b/man/addStepwiseDivergence.Rd
@@ -21,7 +21,11 @@ addStepwiseDivergence(x, ...)
   ...
 )
 
-\S4method{addStepwiseDivergence}{SummarizedExperiment}(x, name = "divergence", name.time = "time_diff", ...)
+\S4method{addStepwiseDivergence}{SummarizedExperiment}(
+  x,
+  name = c("divergence", "time_diff", "ref_samples"),
+  ...
+)
 }
 \arguments{
 \item{x}{A
@@ -50,11 +54,9 @@ step, then increase this accordingly. (Default: \code{1L})}
 Method is passed to the function that is specified by \code{dis.fun}.
 (Default: \code{"bray"})}
 
-\item{name}{\code{Character scalar}. Specifies a column name for storing
-divergence results. (Default: \code{"divergence"})}
-
-\item{name.time}{\code{Character scalar}. Specifies a column name for storing
-time differences. (Default: \code{"time_diff"})}
+\item{name}{\code{Character vector}. Specifies a column name for storing
+divergence results.
+(Default: \code{c("divergence", "time_diff", "ref_samples")})}
 }
 \value{
 \code{getStepwiseDivergence} returns \code{DataFrame} object
diff --git a/tests/testthat/test-getBaselineDivergence.R b/tests/testthat/test-getBaselineDivergence.R
index 057b9da..1cdeede 100644
--- a/tests/testthat/test-getBaselineDivergence.R
+++ b/tests/testthat/test-getBaselineDivergence.R
@@ -4,7 +4,8 @@ test_that("addBaselineDivergence output", {
     tse <- hitchip1006
     tse2 <- addBaselineDivergence(
         tse, group = "subject", time.col = "time",
-        name = "divergence_from_baseline", name.time = "time_from_baseline")
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     # Input and output classes should match
     expect_equal(class(tse), class(tse2))
     # A subject to check time difference calculation
@@ -13,7 +14,7 @@ test_that("addBaselineDivergence output", {
         which(tse2[["subject"]] == "843"), "time_from_baseline"]
     expect_true( all(time2 == time_diff_2) )
     # Test divergences
-    inds0 <- which(tse2[["subject"]] == "843")  
+    inds0 <- which(tse2[["subject"]] == "843")
     inds <- which(tse2[["subject"]] == "843")
     original.divergence <- as.matrix(
         vegan::vegdist(t(assay(tse[, inds0], "counts"))))[, 1]
@@ -25,11 +26,12 @@ test_that("addBaselineDivergence output", {
 test_that("Divergence in baseline other than 0", {
     data(hitchip1006)
     tse <- hitchip1006
-    # Should also work when baseline is not 0  
+    # Should also work when baseline is not 0
     inds <- which(tse[["subject"]] == "843")[2:5]
     tse2 <- addBaselineDivergence(
         tse[, inds], group = "subject", time.col = "time",
-        name = "divergence_from_baseline", name.time = "time_from_baseline")
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     time2 <- tse[, inds][["time"]] - min(tse[, inds][["time"]])
     time_diff_2 <- tse2[["time_from_baseline"]]
     expect_true( all(time2 == time_diff_2) )
@@ -41,39 +43,44 @@ test_that("addBaselineDivergence reference", {
     tse <- hitchip1006
     # Just pick 1 subject with many time points
     # The baseline time point 0 is Sample-843
-    tse <- tse[, tse[["subject"]] == "843"] 
+    tse <- tse[, tse[["subject"]] == "843"]
     tse2 <- addBaselineDivergence(tse, group = "subject", time.col = "time")
     # Define the baseline sample manually
     tse3 <- addBaselineDivergence(
         tse, time.col = "time", group = "subject", reference = "Sample-843",
-        name.time = "time_from_baseline", name = "divergence_from_baseline")
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     tse4 <- addBaselineDivergence(
         tse, time.col = "time", group = "subject", reference = "Sample-1075",
-        name.time = "time_from_baseline", name = "divergence_from_baseline")
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     # Now the times from baseline should be shifted and dissimilarities differ
-    # Sample baseline when the zero time baseline is automatically checked or 
+    # Sample baseline when the zero time baseline is automatically checked or
     # manually set
     expect_true(all(tse2$time_from_baseline==tse3$time_from_baseline))
     # The shifted case (different, middle sample as baseline)
     expect_true(all(tse3$time_from_baseline == tse4$time_from_baseline + 0.7))
-    
+
     tse5 <- addBaselineDivergence(
         tse[, tse[["subject"]] == "843"], group = "subject",
-        time.col = "time", name.time = "time_from_baseline",
-        name = "divergence_from_baseline")      
+        time.col = "time",
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     tse6 <- addBaselineDivergence(
         tse, group = "subject", time.col = "time",
-        name.time = "time_from_baseline", name = "divergence_from_baseline")
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     tse7 <- addBaselineDivergence(
-        tse, group = "subject", time.col = "time", reference = "Sample-1075", 
-        name.time = "time_from_baseline", name = "divergence_from_baseline")  
+        tse, group = "subject", time.col = "time", reference = "Sample-1075",
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     expect_identical(
-        colData(tse5)["Sample-843", "time_from_baseline"], 
+        colData(tse5)["Sample-843", "time_from_baseline"],
         colData(tse6)["Sample-843", "time_from_baseline"])
     expect_identical(
-        colData(tse5)["Sample-843", "time_from_baseline"] - 0.7, 
+        colData(tse5)["Sample-843", "time_from_baseline"] - 0.7,
         colData(tse7)["Sample-843", "time_from_baseline"])
-    
+
     tse <- hitchip1006
     subjects <- unique(tse$subject)
     # Test with full baseline list
@@ -82,15 +89,17 @@ test_that("addBaselineDivergence reference", {
     baselines[names(baselines) == tse[, "Sample-843"][["subject"]]] <-
         "Sample-1075"
     tse8 <- addBaselineDivergence(
-        tse, group = "subject", time.col = "time", reference = baselines, 
-        name.time = "time_from_baseline", name = "divergence_from_baseline")
+        tse, group = "subject", time.col = "time", reference = baselines,
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     expect_identical(
-        colData(tse7)["Sample-843", "time_from_baseline"], 
+        colData(tse7)["Sample-843", "time_from_baseline"],
         colData(tse8)["Sample-843", "time_from_baseline"])
     tse[["reference_sam"]] <- baselines[ match(tse$subject, names(baselines)) ]
     res <- addBaselineDivergence(
-        tse, group = "subject", time.col = "time", reference = "reference_sam", 
-        name.time = "time_from_baseline", name = "divergence_from_baseline")
+        tse, group = "subject", time.col = "time", reference = "reference_sam",
+        name = c("divergence_from_baseline", "time_from_baseline",
+            "reference_samples"))
     ref <- getDivergence(tse, reference = "reference_sam")
     expect_equal(res[["divergence_from_baseline"]], ref)
 })
@@ -104,7 +113,7 @@ test_that("Test altExp", {
         tse, group = "subject", time.col = "time", altexp = "Family")
     altExp(tse, "Family_test") <-  addBaselineDivergence(
         altExp(tse, "Family"), group = "subject", time.col = "time",
-        name = "val", name.time = "time_val")
+        name = c("val", "time_val", "ref"))
     # Time differences should still match
     expect_equal(
         altExp(tse, "Family")$divergence, altExp(tse, "Family_test")$val)
@@ -120,21 +129,25 @@ test_that(".get_reference_samples with different time intervals", {
     res <- getBaselineDivergence(
         tse, group = "subject", time.col = "time",
         assay.type = "counts", method = "euclidean")
-    expect_equal(colData(tse)[, c("divergence", "time_diff")], res)
+    expect_equal(
+        colData(tse)[, c("divergence", "time_diff", "ref_samples")], res)
 })
 
 # Basic SummarizedExperiment for testing
 col_data <- DataFrame(
     time = c(0, 1, 2, 1, 2, 0),
     group = c("A", "A", "A", "B", "B", "B"),
-    row.names = c("Sample1", "Sample2", "Sample3", "Sample4", "Sample5", "Sample6"))
+    row.names = c("Sample1", "Sample2",
+        "Sample3", "Sample4", "Sample5", "Sample6"))
 count_data <- matrix(c(10, 20, 30, 40, 50, 60), ncol = 6, byrow = TRUE)
-se <- SummarizedExperiment(assays = list(counts = count_data), colData = col_data)
+se <- SummarizedExperiment(assays = list(counts = count_data),
+                           colData = col_data)
 
 # Input validation for getBaselineDivergence
 test_that("getBaselineDivergence input validations", {
     expect_error(getBaselineDivergence(se, time.col = "nonexistent"))
-    expect_error(getBaselineDivergence(se, time.col = "time", assay.type = "unknown"))
+    expect_error(getBaselineDivergence(se, time.col = "time",
+        assay.type = "unknown"))
     expect_error(getBaselineDivergence(se, group = "nonexistent"))
     expect_error(getBaselineDivergence(se, reference = "nonexistent"))
     expect_error(getBaselineDivergence(se, name = "nonexistent"))
@@ -164,8 +177,9 @@ test_that("addBaselineDivergence adds columns to colData", {
 # Custom column naming test for addBaselineDivergence
 test_that("addBaselineDivergence handles custom column names", {
     se_result <- addBaselineDivergence(
-        se, time.col = "time", name = "custom_div",
-        name.time = "custom_time_diff")
+        se, time.col = "time",
+        name = c("custom_div", "custom_time_diff",
+            "custom_ref_samples"))
     expect_true("custom_div" %in% colnames(colData(se_result)))
     expect_true("custom_time_diff" %in% colnames(colData(se_result)))
 })
@@ -174,7 +188,9 @@ test_that("addBaselineDivergence handles custom column names", {
 test_that(".add_reference_samples_to_coldata assigns correct baselines", {
     res <- .add_reference_samples_to_coldata(
         se, time.col = "time", group = "group")
-    expect_true("temporal_reference_for_divergence" %in% colnames(colData(res[[1]])))
+    expect_true(
+        "temporal_reference_for_divergence" %in%
+        colnames(colData(res[[1]])))
 })
 
 # Reference sample assignments
@@ -198,13 +214,24 @@ test_that(".get_time_difference calculates correct time diff", {
 
 # Convert divergence to DataFrame
 test_that(".convert_divergence_to_df formats correctly", {
+    col_data <- DataFrame(
+        reference <- c("sample1", "sample2", "sample3", "sample4",
+            "sample5", "sample6")
+    )
+    colnames(col_data) <- c("reference")
     divergence <- c(0.1, 0.2, 0.3, 0, NA, 2)
     time_diff <- c(0, 1, 2, 1, 0, NA)
+    se <- SummarizedExperiment(
+        assays = list(),
+        colData = col_data
+    )
+    reference <- "reference"
     df <- .convert_divergence_to_df(
-        se, divergence, time_diff, name = "test_div",
-        name.time = "test_time_diff")
+        se, divergence, time_diff, reference,
+        name = c("test_div", "test_time_diff", "test_reference_samples"))
     expect_s4_class(df, "DataFrame")
-    expect_equal(colnames(df), c("test_div", "test_time_diff"))
+    expect_equal(colnames(df),
+        c("test_div", "test_time_diff", "test_reference_samples"))
     expect_equal(df$test_div, divergence)
     expect_equal(df$test_time_diff, time_diff)
 })
@@ -233,11 +260,21 @@ test_that("addBaselineDivergence with custom reference sample", {
 
 # Test that postprocessing works with NA values
 test_that(".convert_divergence_to_df with NA divergence values", {
+    col_data <- DataFrame(
+        reference <- c("sample1", "sample2", "sample3", "sample4",
+            "sample5", "sample6")
+    )
     divergence <- c(0.1, NA, 0.3, NA, 0.5, 0.6)
     time_diff <- c(0, 1, 2, 1, 0, NA)
+    colnames(col_data) <- c("reference")
+    se <- SummarizedExperiment(
+        assays = list(),
+        colData = col_data
+    )
+    reference <- "reference"
     df <- .convert_divergence_to_df(
-        se, divergence, time_diff, name = "test_div",
-        name.time = "test_time_diff")
+        se, divergence, time_diff, reference,
+        name = c("test_div", "test_time_diff", "test_reference_samples"))
     expect_s4_class(df, "DataFrame")
     expect_true(all(is.na(df$test_div[is.na(divergence)])))
 })
diff --git a/tests/testthat/test-getStepwiseDivergence.R b/tests/testthat/test-getStepwiseDivergence.R
index 85c3e1e..c61a19f 100644
--- a/tests/testthat/test-getStepwiseDivergence.R
+++ b/tests/testthat/test-getStepwiseDivergence.R
@@ -5,7 +5,7 @@ test_that("Basic functionality of addStepwiseDivergence", {
     tse2 <- addStepwiseDivergence(
         tse, group = "subject", time.interval = 1, time.col = "time",
         assay.type="counts", dis.fun = vegan::vegdist, method = "bray",
-        name.time = "time_difference")
+        name = c("divergence", "time_diff", "ref_samples"))
     expect_equal(class(tse), class(tse2))
 })
 
@@ -16,7 +16,7 @@ test_that("Adding new colData field with existing name generates warning", {
   tse[["time_difference"]] <- NA
   expect_warning(addStepwiseDivergence(
       tse, group = "subject", time.interval = 1, time.col = "time",
-      name.time = "time_difference"))
+      name = c("divergence", "time_difference", "ref_samples")))
 })
 
 # Test: Time difference calculation for a specific subject
@@ -24,9 +24,9 @@ test_that("Time difference calculation is correct for a specific subject", {
   data(hitchip1006)
   tse <- hitchip1006
   tse2 <- addStepwiseDivergence(
-    tse, group = "subject", time.interval = 1, time.col = "time",
-    assay.type="counts", dis.fun = vegan::vegdist, method = "bray",
-    name.time = "time_difference")
+      tse, group = "subject", time.interval = 1, time.col = "time",
+      assay.type="counts", dis.fun = vegan::vegdist, method = "bray",
+      name = c("divergence", "time_difference", "ref_samples"))
   obs_diff <- colData(tse2)[
       which(tse2[["subject"]] == "843"), "time_difference"]
   exp_diff <- c(NA, diff(colData(tse)[
@@ -41,10 +41,11 @@ test_that("addStepwiseDivergence with n > 1 calculates divergences correctly", {
   tse2 <- addStepwiseDivergence(
       tse, group = "subject", time.interval = 2, time.col = "time",
       assay.type = "counts", dis.fun = vegan::vegdist, method = "bray",
-      name.time = "time_difference")
+      name = c("divergence", "time_difference", "ref_samples"))
   time_interval <- 2
   time <- colData(tse2)[which(tse2[["subject"]] == "843"), "time"]
-  time_diff <- colData(tse2)[which(tse2[["subject"]] == "843"), "time_difference"]
+  time_diff <- colData(tse2)[which(tse2[["subject"]] == "843"),
+                             "time_difference"]
   divergence_number <- length(time) - time_interval
   divergence_calculated <- length(which(!is.na(time_diff) == TRUE))
   expect_equal(divergence_number, divergence_calculated)
@@ -55,9 +56,9 @@ test_that("Interval check for divergence calculation", {
   data(hitchip1006)
   tse <- hitchip1006
   tse2 <- addStepwiseDivergence(
-    tse, group = "subject", time.interval = 2, time.col = "time",
-    assay.type = "counts", dis.fun = vegan::vegdist, method = "bray",
-    name.time = "time_difference")
+      tse, group = "subject", time.interval = 2, time.col = "time",
+      assay.type = "counts", dis.fun = vegan::vegdist, method = "bray",
+      name = c("divergence", "time_difference", "ref_samples"))
   time <- colData(tse2)[which(tse2[["subject"]] == "843"), "time"]
   calculated_diff <- time[(1 + 2):length(time)] -
       time[seq_len(length(time) - 2)]
@@ -75,8 +76,7 @@ test_that("Single time point results in NA divergence values", {
   tse2 <- addStepwiseDivergence(
       tse2, group = "subject", time.interval = 1, time.col = "time",
       assay.type = "counts", dis.fun = vegan::vegdist, method = "bray",
-      name = "time_divergence",
-      name.time = "time_difference")
+      name = c("time_divergence", "time_diff", "ref_samples"))
   expect_true(all(is.na(colData(tse2)[
       which(duplicated(tse2[["subject"]]) == FALSE),
       "time_divergence"])))
@@ -89,11 +89,11 @@ test_that("Comparing vegan distances (bray vs euclidean)", {
   tse2 <- addStepwiseDivergence(
       tse, group = "subject", time.interval = 1, time.col = "time",
       assay.type = "counts", dis.fun = vegan::vegdist, method = "bray",
-      name.time = "timedifference", name = "timedivergence")
+      name = c("timedivergence", "timedifference", "ref_samples"))
   tse2 <- addStepwiseDivergence(
       tse2, group = "subject", time.interval = 1, time.col = "time",
       assay.type = "counts", dis.fun = vegan::vegdist, method = "euclidean",
-      name.time = "timedifference2", name = "timedivergence2")
+      name = c("timedivergence2", "timedifference2", "ref_samples2"))
   expect_true(identical(tse2$timedifference, tse2$timedifference2))
   expect_true(!identical(tse2$timedivergence, tse2$timedivergence2))
 })
@@ -106,15 +106,17 @@ test_that("AltExp functionality in addStepwiseDivergence", {
   tse <- addStepwiseDivergence(
       tse, group = "subject", time.interval = 1, time.col = "time",
       altexp = "Family")
-  altExp(tse, "Family_test") <- addStepwiseDivergence(
-    altExp(tse, "Family"), group = "subject", time.interval = 1,
-    time.col = "time", name.time = "timedifference", name = "timedivergence")
+
+  altExp(tse, "Family") <- addStepwiseDivergence(
+      altExp(tse, "Family"), group = "subject", time.interval = 1,
+      time.col = "time",
+      name = c("timedivergence", "timedifference", "ref_samples2"))
   expect_equal(
       altExp(tse, "Family")$time_diff,
-      altExp(tse, "Family_test")$timedifference)
+      altExp(tse, "Family")$timedifference)
   expect_equal(
       altExp(tse, "Family")$divergence,
-      altExp(tse, "Family_test")$timedivergence)
+      altExp(tse, "Family")$timedivergence)
 })
 
 # Test: getStepwiseDivergence output type
@@ -176,21 +178,25 @@ test_that(".get_reference_samples with different time intervals", {
     res <- getStepwiseDivergence(
         tse, group = "subject", time.interval = 2, time.col = "time",
         assay.type = "counts", method = "euclidean")
-    expect_equal(colData(tse)[, c("divergence", "time_diff")], res)
+    expect_equal(
+        colData(tse)[, c("divergence", "time_diff", "ref_samples")], res)
 })
 
 # Basic SummarizedExperiment for testing
 col_data <- DataFrame(
     time = c(0, 1, 2, 1, 2, 0),
     group = c("A", "A", "A", "B", "B", "B"),
-    row.names = c("Sample1", "Sample2", "Sample3", "Sample4", "Sample5", "Sample6"))
+    row.names = c("Sample1", "Sample2", "Sample3", "Sample4",
+        "Sample5", "Sample6"))
 count_data <- matrix(c(10, 20, 30, 40, 50, 60), ncol = 6, byrow = TRUE)
-se <- SummarizedExperiment(assays = list(counts = count_data), colData = col_data)
+se <- SummarizedExperiment(assays = list(counts = count_data),
+                           colData = col_data)
 
 # Input validation for getStepwiseDivergence
 test_that("getStepwiseDivergence input validations", {
     expect_error(getStepwiseDivergence(se, time.col = "nonexistent"))
-    expect_error(getStepwiseDivergence(se, time.col = "time", assay.type = "unknown"))
+    expect_error(getStepwiseDivergence(se, time.col = "time",
+        assay.type = "unknown"))
     expect_error(getStepwiseDivergence(se, group = "nonexistent"))
     expect_error(getStepwiseDivergence(se, reference = "nonexistent"))
     expect_error(getStepwiseDivergence(se, name = "nonexistent"))
@@ -201,7 +207,9 @@ test_that("getStepwiseDivergence input validations", {
 test_that("getStepwiseDivergence dissimilarity calculation", {
     result <- getStepwiseDivergence(se, time.col = "time", method = "bray")
     expect_s4_class(result, "DataFrame")
-    expect_true(all(c("divergence", "time_diff") %in% colnames(result)))
+    expect_true(
+        all(c("divergence", "time_diff", "ref_samples") %in%
+        colnames(result)))
 })
 
 # Correct time difference calculation test
@@ -220,8 +228,9 @@ test_that("addStepwiseDivergence adds columns to colData", {
 # Custom column naming test for addStepwiseDivergence
 test_that("addStepwiseDivergence handles custom column names", {
     se_result <- addStepwiseDivergence(
-        se, time.col = "time", name = "custom_div",
-        name.time = "custom_time_diff")
+        se, time.col = "time",
+        name = c("custom_div", "custom_time_diff", "ref_samples"))
+    # Test
     expect_true("custom_div" %in% colnames(colData(se_result)))
     expect_true("custom_time_diff" %in% colnames(colData(se_result)))
 })
@@ -230,7 +239,9 @@ test_that("addStepwiseDivergence handles custom column names", {
 test_that(".add_reference_samples_to_coldata assigns correct baselines", {
     res <- .add_reference_samples_to_coldata(
         se, time.col = "time", group = "group")
-    expect_true("temporal_reference_for_divergence" %in% colnames(colData(res[[1]])))
+    expect_true(
+        "temporal_reference_for_divergence" %in%
+        colnames(colData(res[[1]])))
 })
 
 # Reference sample assignments
@@ -262,6 +273,7 @@ test_that("getStepwiseDivergence unsupported method", {
 test_that("addStepwiseDivergence with custom reference sample", {
     res <- getStepwiseDivergence(
         se, time.col = "time", group = "group")
+    #
     se[["reference"]] <- c(NA, "Sample1", "Sample2", "Sample6", "Sample4", NA)
     time_diff <- c(NA, 1, 1, 1, 1, NA)
     ref <- getDivergence(se, reference = "reference")
diff --git a/vignettes/articles/minimalgut.Rmd b/vignettes/articles/minimalgut.Rmd
index 4c27792..3835499 100644
--- a/vignettes/articles/minimalgut.Rmd
+++ b/vignettes/articles/minimalgut.Rmd
@@ -135,8 +135,8 @@ tse <- addStepwiseDivergence(
     group = "StudyIdentifier",
     time.interval = 1,
     time.col = "Time.hr",
-    name = "divergence_from_previous_step",
-    name.time = "time_from_previous_step"
+    name = c("divergence_from_previous_step", 
+        "time_from_previous_step", "reference_samples")
     )
 ```
 
diff --git a/vignettes/miaTime.Rmd b/vignettes/miaTime.Rmd
index 9524e70..6f88a61 100755
--- a/vignettes/miaTime.Rmd
+++ b/vignettes/miaTime.Rmd
@@ -74,7 +74,8 @@ data(hitchip1006)
 tse <- hitchip1006
 
 res <- getBaselineDivergence(
-    tse, time.col = "time", group = "sample", name = "baseline")
+    tse, time.col = "time", group = "sample", 
+    name = c("baseline", "time_diff", "ref_samples"))
 res |> head()
 ```