feat: clusterOrder argument for plotClusterTrajectories (#144)

R/model.R

@@ -1325,6 +1325,9 @@ setMethod('plotFittedTrajectories', 'lcModel', function(object, ...) {
 #' @inheritDotParams clusterTrajectories
 #' @param clusterLabeler A `function(clusterNames, clusterSizes)` that generates plot labels for the clusters.
 #' By default the cluster name with the proportional size is shown, see [make.clusterPropLabels].
+#' @param clusterOrder Specify which clusters to plot and the order.
+#' Can be the cluster names or index.
+#' By default, all clusters are shown.
 #' @param trajAssignments The cluster assignments for the fitted trajectories. Only used when `trajectories = TRUE` and `facet = TRUE`. See [trajectoryAssignments].
 #' @param ... Arguments passed to [clusterTrajectories()], or [ggplot2::geom_line()] for plotting the cluster trajectory lines.
 #' @return A `ggplot` object.
@@ -1337,6 +1340,15 @@ setMethod('plotFittedTrajectories', 'lcModel', function(object, ...) {
 #' if (require("ggplot2")) {
 #'   plotClusterTrajectories(model)
 #'
+#'   # show cluster sizes in labels
+#'   plotClusterTrajectories(model, clusterLabeler = make.clusterSizeLabels)
+#'
+#'   # change cluster order
+#'   plotClusterTrajectories(model, clusterOrder = c('B', 'C', 'A'))
+#'
+#'   # show only specific clusters
+#'   plotClusterTrajectories(model, clusterOrder = c('B', 'C'))
+#'
 #'   # show assigned trajectories
 #'   plotClusterTrajectories(model, trajectories = TRUE)
 #'
@@ -1358,6 +1370,7 @@ setMethod('plotClusterTrajectories', 'lcModel',
     object,
     what = 'mu',
     at = time(object),
+    clusterOrder = character(),
     clusterLabeler = make.clusterPropLabels,
     trajectories = FALSE,
     facet = !isFALSE(as.logical(trajectories[1])),
@@ -1381,6 +1394,7 @@ setMethod('plotClusterTrajectories', 'lcModel',
 
   .plotClusterTrajs(
     clusdata,
+    clusterOrder = clusterOrder,
     clusterLabels = clusterLabels,
     response = responseVariable(object, what = what),
     time = timeVariable(object),

R/trajectories.R

@@ -52,6 +52,7 @@ setMethod('plotClusterTrajectories', 'data.frame', function(
   object,
   response,
   cluster = 'Cluster',
+  clusterOrder = character(),
   clusterLabeler = make.clusterPropLabels,
   time = getOption('latrend.time'),
   center = meanNA,
@@ -77,6 +78,7 @@ setMethod('plotClusterTrajectories', 'data.frame', function(
 
   .plotClusterTrajs(
     clusTrajData,
+    clusterOrder = clusterOrder,
     clusterLabels = clusterLabels,
     response = response,
     time = time,
@@ -92,6 +94,7 @@ setMethod('plotClusterTrajectories', 'data.frame', function(
 
 .plotClusterTrajs = function(
   data,
+  clusterOrder,
   clusterLabels,
   response,
   time,
@@ -114,11 +117,15 @@ setMethod('plotClusterTrajectories', 'data.frame', function(
     is.flag(facet)
   )
 
-  data = data.table(data)
+  clusterNames = as.character(unique(data[[cluster]]))
+  clusterOrderNames = make.orderedClusterNames(clusterNames, clusterOrder, subset = TRUE)
+  clusterIndices = match(clusterOrderNames, clusterNames)
+
+  data = data.table(data)[get(cluster) %in% clusterOrderNames]
   data[, c(cluster) := factor(
     get(cluster),
-    levels = levels(get(cluster)),
-    labels = clusterLabels
+    levels = clusterNames[clusterIndices],
+    labels = clusterLabels[clusterIndices]
   )]
 
   if (isTRUE(as.logical(trajectories))) {

tests/testthat/test-plot.R

@@ -102,3 +102,10 @@ test_that('plotClusterTrajectories.lcModel with labeler', {
   expect_true(is.ggplot(p3))
 })
 
+
+test_that('plotClusterTrajectories.lcModel with ordening', {
+  p3 = plotClusterTrajectories(testModel3, clusterOrder = clusterNames(testModel3)[1:2])
+  expect_true(is.ggplot(p3))
+})
+
+