From 00d32c6e6e0e6b1716c303272ac95ed2a5fd022a Mon Sep 17 00:00:00 2001
From: jgabry <jgabry@gmail.com>
Date: Fri, 17 Nov 2023 11:11:38 -0700
Subject: [PATCH] new helper function repair_variable_names

closes #317
---
 NAMESPACE                       |  1 +
 R/draws-index.R                 | 33 ++++++++++++++++++++++++++++++---
 man/set_variables.Rd            | 22 +++++++++++++++++++---
 tests/testthat/test-variables.R | 14 ++++++++++++++
 4 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/NAMESPACE b/NAMESPACE
index 2d81d65..fb03d94 100755
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -463,6 +463,7 @@ export(r_scale)
 export(rdo)
 export(rename_variables)
 export(repair_draws)
+export(repair_variable_names)
 export(resample_draws)
 export(reserved_variables)
 export(rfun)
diff --git a/R/draws-index.R b/R/draws-index.R
index 8cdfa99..0f6d796 100644
--- a/R/draws-index.R
+++ b/R/draws-index.R
@@ -145,11 +145,13 @@ remove_reserved_variable_names <- function(variables, reserved) {
 
 #' Set variable names in `draws` objects
 #'
-#' Set variable names for all variables in a [`draws`] object. Useful
-#' when using pipe operators.
+#' Set variable names for all variables in a [`draws`] object. Useful when using
+#' pipe operators. The additional helper function `repair_variable_names()` can
+#' convert variable names using periods (e.g. `"theta.1"`) to the more common
+#' square brackets (`"theta[1]"`) that are better supported by the package.
 #'
 #' @param x (draws) A [`draws`] object.
-#' @param variables (character) new variable names.
+#' @param variables (character) New variable names.
 #' @template args-methods-dots
 #'
 #' @return Returns a [`draws`] object of the same format as `x`, with
@@ -173,6 +175,31 @@ set_variables <- function(x, variables, ...) {
   return(x)
 }
 
+#' @rdname set_variables
+#' @param old_variables (character) Variable names to repair. Should be variable
+#'   names with periods to convert to square brackets (e.g., `"theta.1"` ->
+#'   `"theta[1]"`, `"theta.1.1"` -> `"theta[1,1]"`).
+#' @examples
+#' # using repair_variable_names
+#' x <- matrix(rnorm(100), ncol = 2)
+#' colnames(x) <- c("theta.1", "theta.2")
+#' repair_variable_names(colnames(x))
+#' x <- set_variables(as_draws(x), repair_variable_names(colnames(x)))
+#' variables(x)
+#'
+#' @export
+repair_variable_names <- function(old_variables) {
+  if (!all(grepl("\\.", old_variables))) {
+    stop_no_call(
+      "All names in 'old_variables' must contain at least one '.' in the name."
+    )
+  }
+  repaired_variables <- sub("\\.", "[", old_variables)
+  repaired_variables <- gsub("\\.", ",", repaired_variables)
+  repaired_variables[grep("\\[", repaired_variables)] <-
+    paste0(repaired_variables[grep("\\[", repaired_variables)], "]")
+  repaired_variables
+}
 
 #' @rdname draws-index
 #' @export
diff --git a/man/set_variables.Rd b/man/set_variables.Rd
index ba67f44..e9c04c0 100644
--- a/man/set_variables.Rd
+++ b/man/set_variables.Rd
@@ -2,24 +2,33 @@
 % Please edit documentation in R/draws-index.R
 \name{set_variables}
 \alias{set_variables}
+\alias{repair_variable_names}
 \title{Set variable names in \code{draws} objects}
 \usage{
 set_variables(x, variables, ...)
+
+repair_variable_names(old_variables)
 }
 \arguments{
 \item{x}{(draws) A \code{\link{draws}} object.}
 
-\item{variables}{(character) new variable names.}
+\item{variables}{(character) New variable names.}
 
 \item{...}{Arguments passed to individual methods (if applicable).}
+
+\item{old_variables}{(character) Variable names to repair. Should be variable
+names with periods to convert to square brackets (e.g., \code{"theta.1"} ->
+\code{"theta[1]"}, \code{"theta.1.1"} -> \code{"theta[1,1]"}).}
 }
 \value{
 Returns a \code{\link{draws}} object of the same format as \code{x}, with
 variables named as specified.
 }
 \description{
-Set variable names for all variables in a \code{\link{draws}} object. Useful
-when using pipe operators.
+Set variable names for all variables in a \code{\link{draws}} object. Useful when using
+pipe operators. The additional helper function \code{repair_variable_names()} can
+convert variable names using periods (e.g. \code{"theta.1"}) to the more common
+square brackets (\code{"theta[1]"}) that are better supported by the package.
 }
 \examples{
 x <- as_draws(matrix(rnorm(100), ncol = 2))
@@ -32,6 +41,13 @@ variables(x)
 variables(x) <- c("theta[1]", "theta[2]")
 variables(x)
 
+# using repair_variable_names
+x <- matrix(rnorm(100), ncol = 2)
+colnames(x) <- c("theta.1", "theta.2")
+repair_variable_names(colnames(x))
+x <- set_variables(as_draws(x), repair_variable_names(colnames(x)))
+variables(x)
+
 }
 \seealso{
 \code{\link{variables}}
diff --git a/tests/testthat/test-variables.R b/tests/testthat/test-variables.R
index e773d49..b058fa2 100755
--- a/tests/testthat/test-variables.R
+++ b/tests/testthat/test-variables.R
@@ -109,3 +109,17 @@ test_that("cannot set duplicate variable names", {
   expect_error(set_variables(x, c("a", "a")), "Duplicate variable names are not allowed")
 })
 
+test_that("repair_variable_names does the right conversion", {
+  expect_equal(
+    repair_variable_names(c("foo.12", "foo.12.13", "foo.12.13.14.15")),
+    c("foo[12]", "foo[12,13]", "foo[12,13,14,15]")
+  )
+})
+
+test_that("repair_variable_names errors if all names don't contain '.'", {
+  expect_error(
+    repair_variable_names(c("foo[12]", "foo.12.13")),
+    "All names in 'old_variables' must contain at least one '.' in the name."
+  )
+})
+