diff --git a/DESCRIPTION b/DESCRIPTION
index f98d6fd..6934346 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -14,7 +14,7 @@ Description: A progression model for repeated measures (PMRM)
   This package implements frequentist PMRMs by
   Raket (2022) <doi:10.1002/sim.9581> using
   'RTMB' by Kristensen (2016) <doi:10.18637/jss.v070.i05>.
-Version: 0.0.2.9000
+Version: 0.0.2.9001
 License: MIT + file LICENSE
 URL: https://github.com/openpharma/pmrm,
   https://openpharma.github.io/pmrm/
diff --git a/NEWS.md b/NEWS.md
index 041dae1..354ac6c 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,6 +1,6 @@
-# pmrm 0.0.2.9000 (development)
-
+# pmrm development version
 
+* Use all of `pmrm_data()` and rebuild all ordered factors in `predict()` (#5, @lbenz-lilly).
 
 # pmrm 0.0.2
 
diff --git a/R/pmrm_data.R b/R/pmrm_data.R
index 04fbf91..608b26a 100644
--- a/R/pmrm_data.R
+++ b/R/pmrm_data.R
@@ -49,6 +49,9 @@
 #'   Set `covariates` to `~ 0` (default) to opt out of covariate adjustment.
 #'   The intercept term is removed from the model matrix `W`
 #'   whether or not the formula begins with `~ 0.
+#' @param subset `TRUE` if `data` is a subset of data without
+#'   all levels of `visit` or `arm`, `FALSE` to expect a full dataset
+#'   when validating.
 pmrm_data <- function(
   data,
   outcome,
@@ -56,7 +59,8 @@ pmrm_data <- function(
   patient,
   visit,
   arm,
-  covariates = ~0
+  covariates = ~0,
+  subset = FALSE
 ) {
   assert(is.data.frame(data), message = "data must be a data frame.")
   labels <- list(
@@ -68,7 +72,7 @@ pmrm_data <- function(
     covariates = covariates
   )
   data <- pmrm_data_new(data = data, labels = labels)
-  pmrm_data_validate(data)
+  pmrm_data_validate(data, subset = subset)
   pmrm_data_clean(data)
 }
 
@@ -82,7 +86,7 @@ pmrm_data_labels <- function(data) {
   attr(data, "pmrm_data_labels")
 }
 
-pmrm_data_validate <- function(data) {
+pmrm_data_validate <- function(data, subset = FALSE) {
   pmrm_predictors_validate(data)
   labels <- pmrm_data_labels(data)
   assert(
@@ -97,14 +101,16 @@ pmrm_data_validate <- function(data) {
       "In addition, all non-missing values must be finite."
     )
   )
-  assert(
-    length(unique(data[[labels$arm]])) > 1L,
-    message = "data[[arm]] must have more than one unique element."
-  )
-  assert(
-    length(unique(data[[labels$visit]])) > 1L,
-    message = "data[[visit]] must have more than one unique element."
-  )
+  if (!subset) {
+    assert(
+      length(unique(data[[labels$arm]])) > 1L,
+      message = "data[[arm]] must have more than one unique element."
+    )
+    assert(
+      length(unique(data[[labels$visit]])) > 1L,
+      message = "data[[visit]] must have more than one unique element."
+    )
+  }
 }
 
 pmrm_predictors_validate <- function(data) {
@@ -204,6 +210,12 @@ pmrm_predictors_validate <- function(data) {
 pmrm_data_clean <- function(data) {
   labels <- pmrm_data_labels(data)
   data <- data[!is.na(data[[labels$outcome]]), , drop = FALSE] # nolint
+  data <- pmrm_data_clean_factors(data)
+  data[order(data[[labels$patient]], data[[labels$visit]]), , drop = FALSE] # nolint
+}
+
+pmrm_data_clean_factors <- function(data) {
+  labels <- pmrm_data_labels(data)
   data[[labels$patient]] <- factor(data[[labels$patient]], ordered = FALSE)
   for (name in c(labels$visit, labels$arm)) {
     value <- data[[name]]
@@ -218,5 +230,5 @@ pmrm_data_clean <- function(data) {
       data[[name]] <- ordered(value, levels = sort(levels(value)))
     }
   }
-  data[order(data[[labels$patient]], data[[labels$visit]]), , drop = FALSE] # nolint
+  data
 }
diff --git a/R/predict.R b/R/predict.R
index d43ae83..ba98c48 100644
--- a/R/predict.R
+++ b/R/predict.R
@@ -70,19 +70,9 @@ predict.pmrm_fit <- function(
     . <= 1,
     message = "confidence must have length 1 and be between 0 and 1."
   )
-  labels <- pmrm_data_labels(object$data)
-  data_new <- pmrm_data_new(data = data, labels = labels)
-  pmrm_predictors_validate(data_new)
-  data_new[[labels$outcome]] <- NA_real_
+  data_new <- pmrm_predict_data_new(object, data)
   data_old <- object$data
-  patient <- labels$patient
-  data_new[[patient]] <- paste0("new_", data_new[[patient]])
-  data_old[[patient]] <- paste0("old_", data_old[[patient]])
-  data_combined <- dplyr::bind_rows(data_old, data_new)
-  data_combined[[patient]] <- factor(
-    data_combined[[patient]],
-    levels = unique(data_combined[[patient]])
-  )
+  data_combined <- pmrm_predict_data_combined(object, data_new, data_old)
   constants <- pmrm_constants(
     data = data_combined,
     visit_times = object$constants$visit_times,
@@ -114,7 +104,9 @@ predict.pmrm_fit <- function(
   summary <- tibble::as_tibble(summary[rownames(summary) == name, ])
   summary <- summary[-seq_len(nrow(data_old)), , drop = FALSE] # nolint
   z <- stats::qnorm(p = (1 - confidence) / 2, lower.tail = FALSE)
+  labels <- pmrm_data_labels(object$data)
   tibble::tibble(
+    patient = data_new[[labels$patient]],
     arm = data[[labels$arm]],
     visit = data[[labels$visit]],
     time = data[[labels$time]],
@@ -125,5 +117,51 @@ predict.pmrm_fit <- function(
   )
 }
 
+pmrm_predict_data_combined <- function(object, data_new, data_old) {
+  labels <- pmrm_data_labels(object$data)
+  patient <- labels$patient
+  data_new[[patient]] <- paste0("new_", data_new[[patient]])
+  data_old[[patient]] <- paste0("old_", data_old[[patient]])
+  data_combined <- dplyr::bind_rows(data_old, data_new)
+  data_combined[[patient]] <- factor(
+    data_combined[[patient]],
+    levels = unique(data_combined[[patient]])
+  )
+  data_combined
+}
+
+pmrm_predict_data_new <- function(object, data) {
+  labels <- pmrm_data_labels(object$data)
+  data[[labels$outcome]] <- 0
+  data <- pmrm_data(
+    data = data,
+    outcome = labels$outcome,
+    time = labels$time,
+    patient = labels$patient,
+    visit = labels$visit,
+    arm = labels$arm,
+    covariates = labels$covariates,
+    subset = TRUE
+  )
+  data[[labels$outcome]] <- NA_real_
+  assert(
+    all(levels(data[[labels$visit]]) %in% levels(object$data[[labels$visit]])),
+    message = paste(
+      "visit levels in new data must be a subset",
+      "of visit levels in original data."
+    )
+  )
+  assert(
+    all(levels(data[[labels$arm]]) %in% levels(object$data[[labels$arm]])),
+    message = paste(
+      "arm levels in new data must be a subset",
+      "of arm levels in original data."
+    )
+  )
+  levels(data[[labels$visit]]) <- levels(object$data[[labels$visit]])
+  levels(data[[labels$arm]]) <- levels(object$data[[labels$arm]])
+  data
+}
+
 #' @export
 stats::predict
diff --git a/man/pmrm_data.Rd b/man/pmrm_data.Rd
index 0ac0c00..9b253bf 100644
--- a/man/pmrm_data.Rd
+++ b/man/pmrm_data.Rd
@@ -4,7 +4,16 @@
 \alias{pmrm_data}
 \title{Internal function to prepare data}
 \usage{
-pmrm_data(data, outcome, time, patient, visit, arm, covariates = ~0)
+pmrm_data(
+  data,
+  outcome,
+  time,
+  patient,
+  visit,
+  arm,
+  covariates = ~0,
+  subset = FALSE
+)
 }
 \arguments{
 \item{data}{A \code{tibble} or data frame with one row per patient visit.
@@ -53,6 +62,10 @@ any missing values.
 Set \code{covariates} to \code{~ 0} (default) to opt out of covariate adjustment.
 The intercept term is removed from the model matrix \code{W}
 whether or not the formula begins with `~ 0.}
+
+\item{subset}{\code{TRUE} if \code{data} is a subset of data without
+all levels of \code{visit} or \code{arm}, \code{FALSE} to expect a full dataset
+when validating.}
 }
 \value{
 A \code{tibble} of class \code{"pmrm_data"}
diff --git a/tests/testthat/test-predict.R b/tests/testthat/test-predict.R
index f6a3144..76dfa4c 100644
--- a/tests/testthat/test-predict.R
+++ b/tests/testthat/test-predict.R
@@ -37,6 +37,7 @@ test_that("predict() on small data with proportional decline model", {
       sort(colnames(out)),
       sort(
         c(
+          "patient",
           "estimate",
           "standard_error",
           "lower",
@@ -89,6 +90,7 @@ test_that("predict() on small data with non-proportional slowing model", {
       sort(colnames(out)),
       sort(
         c(
+          "patient",
           "estimate",
           "standard_error",
           "lower",