traceplot and densityplot get layout argument

7 files changed, 28 insertions, 18 deletions

diff --git a/DESCRIPTION b/DESCRIPTION
index 2f88f39..ac508e6 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,6 +1,6 @@
 Package: jagsUI
-Version: 1.5.1.9102
-Date: 2021-05-23
+Version: 1.5.1.9103
+Date: 2021-06-16
 Title: A Wrapper Around 'rjags' to Streamline 'JAGS' Analyses
 Author: Ken Kellner <contact@kenkellner.com>
 Maintainer: Ken Kellner <contact@kenkellner.com>
diff --git a/R/densityplot.R b/R/densityplot.R
index 7d346e1..0c6f987 100644
--- a/R/densityplot.R
+++ b/R/densityplot.R
@@ -1,11 +1,11 @@
 #Get density plots for series of parameters
-densityplot <- function(x, parameters=NULL, per_plot=9, ask=NULL){
+densityplot <- function(x, parameters=NULL, layout=NULL, ask=NULL){
   
   #Check input class and get basic plot settings
   check_class(x)  
   if(is.null(ask))
     ask <- grDevices::dev.interactive(orNone = TRUE)
-  plot_info <- get_plot_info(x, parameters, per_plot, ask)
+  plot_info <- get_plot_info(x, parameters, layout, ask)
   
   #Handle par()
   old_par <- graphics::par(plot_info$new_par)
diff --git a/R/get_plot_info.R b/R/get_plot_info.R
index 86d4389..e3220ff 100644
--- a/R/get_plot_info.R
+++ b/R/get_plot_info.R
@@ -1,6 +1,8 @@
 
 #General function for setting up plots
-get_plot_info <- function(x, parameters, per_plot, ask, Rhat_min=NULL){
+# Called by densityplot, traceplot, and plot.jagsUI
+# plot.jagsUI only uses the 'params' component in the output, ignores the rest
+get_plot_info <- function(x, parameters, layout, ask, Rhat_min=NULL){
 
   #Expand non-scalar parameters and check they exist
   all_params <- param_names(x$samples)
@@ -23,16 +25,24 @@ get_plot_info <- function(x, parameters, per_plot, ask, Rhat_min=NULL){
     if(length(parameters)==0) stop("No parameters > Rhat_min")
   }
 
-  #Reduce max panels per plot if larger than number of parameters
-  if(length(parameters) <= per_plot){
-    per_plot <- length(parameters)
-    ask=FALSE
+  # Fix 'layout'
+  if(is.null(layout)) {
+    if(length(parameters) <= 9){
+      per_plot0 <- length(parameters)
+      ask <- FALSE
+      layout <- c(ceiling(sqrt(per_plot0)), round(sqrt(per_plot0)))
+    } else {
+      layout <- c(3,3)
+    }
+  } else if(length(layout) != 2) {
+    layout <- rep(layout[1], 2)
   }
-
+  per_plot <- prod(layout)
+  
   #Set up new par settings
   new_par <- list(mar=c(1.5,1.5,2.5,1), oma=c(3,3,0,0), ask=ask)
   if(per_plot > 1)
-    new_par$mfrow <- c(ceiling(sqrt(per_plot)), round(sqrt(per_plot)))
+    new_par$mfrow <- layout
 
   list(params=parameters, new_par=new_par, per_plot=per_plot)
 }
diff --git a/R/plot.R b/R/plot.R
index 79a256e..17108cf 100644
--- a/R/plot.R
+++ b/R/plot.R
@@ -4,7 +4,7 @@ plot.jagsUI <- function(x, parameters=NULL, per_plot=4, ask=NULL, ...){
 
   if(is.null(ask))
     ask <- grDevices::dev.interactive(orNone = TRUE)
-  plot_info <- get_plot_info(x, parameters, per_plot, ask)
+  plot_info <- get_plot_info(x, parameters, NULL, ask)
   dims <- c(min(length(plot_info$params), per_plot), 2)
   if(length(plot_info$params) <= per_plot)
     ask <- FALSE
diff --git a/R/traceplot.R b/R/traceplot.R
index e029d18..212b038 100644
--- a/R/traceplot.R
+++ b/R/traceplot.R
@@ -1,12 +1,12 @@
 #Get traceplots for series of parameters
 traceplot <- function(x, parameters=NULL, Rhat_min=NULL,
-                      per_plot=9, ask=NULL){
+                      layout=NULL, ask=NULL){
 
   #Check input class and get basic plot settings
   check_class(x)
   if(is.null(ask))
     ask <- grDevices::dev.interactive(orNone = TRUE)
-  plot_info <- get_plot_info(x, parameters, per_plot, ask, Rhat_min)
+  plot_info <- get_plot_info(x, parameters, layout, ask, Rhat_min)
 
   #Handle par()
   old_par <- graphics::par(plot_info$new_par)
diff --git a/man/densityplot.Rd b/man/densityplot.Rd
index efe5631..bf72c4c 100644
--- a/man/densityplot.Rd
+++ b/man/densityplot.Rd
@@ -4,13 +4,13 @@
 \title{Density plots of JAGS output}
 
 \usage{
-  densityplot(x, parameters=NULL, per_plot=9, ask=NULL)
+  densityplot(x, parameters=NULL, layout=NULL, ask=NULL)
 }
 
 \arguments{
   \item{x}{A jagsUI object}
   \item{parameters}{A vector of names (as characters) of parameters to plot. Parameter names must match parameters included in the model. Calling non-scalar parameters without subsetting (e.g. \code{alpha}) will plot all values of \code{alpha}. If \code{parameters=NULL}, all parameters will be plotted.}
-  \item{per_plot}{Maximum number of parameters to include on each plot.}
+  \item{layout}{A length 2 vector with the number of rows and columns to display in the plot. The default is 3 x 3, or smaller if there are fewer parameters to plot.}
   \item{ask}{If \code{TRUE}, ask user for confirmation before generating each new plot; the default is to ask when output is going to the screen, not when it is going to a file.}
 }
 
diff --git a/man/traceplot.Rd b/man/traceplot.Rd
index 444940e..a1e2f75 100644
--- a/man/traceplot.Rd
+++ b/man/traceplot.Rd
@@ -4,14 +4,14 @@
 \title{Traceplots of JAGS output}
 
 \usage{
-  traceplot(x, parameters=NULL, Rhat_min=NULL, per_plot=9, ask=NULL)
+  traceplot(x, parameters=NULL, Rhat_min=NULL, layout=NULL, ask=NULL)
 }
 
 \arguments{
   \item{x}{A jagsUI object}
   \item{parameters}{A vector of names (as characters) of parameters to plot. Parameter names must match parameters included in the model. Calling non-scalar parameters without subsetting (e.g. \code{alpha}) will plot all values of \code{alpha}. If \code{parameters=NULL}, all parameters will be plotted.}
   \item{Rhat_min}{If provided, only plot parameters with Rhat values that exceed the provided value. A good min value to start with is 1.05.}
-  \item{per_plot}{Maximum number of parameters to include on each plot.}
+  \item{layout}{A length 2 vector with the number of rows and columns to display in the plot. The default is 3 x 3, or smaller if there are fewer parameters to plot.}
   \item{ask}{If \code{TRUE}, ask user for confirmation before generating each new plot; the default is to ask when output is going to the screen, not when it is going to a file.}
 }