finished the plotting function

R/interfaceFunctions.R

@@ -94,10 +94,13 @@ fitModel <- function(data, trait.colums, factor.column, n.pop, n.burnin = 5000,
 
 #' Plot population assignments
 #'
-#' Plots population assignment probabilities per line. Uses \code{data.table} and \code{ggplot2} if available.
+#' Plot method for MuGaMix objects (generated by \code{fitModel()}). Plots population assignment probabilities for each line. Uses \code{data.table} and \code{ggplot2} if available.
+#' If the data set contains many lines, their names may not show up well. In such cases, the user can further modify plot parameters by using the plot objects returned by the function.
 #'
 #' @param obj a \code{mugamix} object generated by the \code{fitModel} function.
 #'
+#' @return a plot object (either \code{ggplot} or \code{barplot})
+#'
 #' @export
 plot.mugamix <- function(obj){
 	# Calculate medians across samples and chains
@@ -109,14 +112,24 @@ plot.mugamix <- function(obj){
 	popAss   <- NULL
 	if (requireNamespace("data.table", quietly=TRUE)) { # if we have data.table installed
 		popP   <- data.table::as.data.table(popP)
-		popP   <- popP[, line := obj$lineIDs]
+		data.table::set(popP, NULL, "line", obj$lineIDs)
 		popP   <- data.table::setorderv(popP, as.character(1:obj$n.pops), rep(-1, obj$n.pops))
 		popAss <- data.table::melt(popP, measure=as.character(1:obj$n.pops), variable.name="population", value.name="p")
-		popAss <- popAss[, line := factor(line, levels=unique(line))]
+		data.table::set(popAss, NULL, "line", factor(popAss[,"line"], levels=unique(popAss[,"line"])))
 	} else {
 		popP      <- as.data.frame(popP)
 		popP$line <- obj$lineIDs
 		popP      <- popP[do.call(order, -popP[,1:obj$n.pops]),]
-		popAss    <- data.frame(p=unlist(popP[,1:obj$n.pops]), line=rep(as.character(popP[,"line"]), times=obj$n.pops))
+		popAss    <- data.frame(p=unlist(popP[,1:obj$n.pops]),
+						line=factor(rep(as.character(popP[,"line"]), times=obj$n.pops), levels=unique(as.character(popP[,"line"]))),
+						population=rep(as.character(1:obj$n.pops), each=Nln))
+	}
+	if (requireNamespace("ggplot2", quietly=TRUE)) {
+		return(ggplot2::ggplot(data=popAss, ggplot2::aes(x=line, y=p, fill=population)) + ggplot2::geom_col())
+	} else {
+		popMat <- t(as.matrix(popP[,1:obj$n.pops]))
+		colnames(popMat) <- unlist(popP[,"line"])
+		return(barplot(popMat, border=NA, col=rainbow(obj$n.pops), ylab="p", las=2))
 	}
 }
+

README.md

@@ -9,3 +9,5 @@ The package is still in development. Basic model fitting with one level of repli
 
 To install, make sure you have the `devtools` package on your system, and then run `install_github("tonymugen/MuGaMix")`. It should work under any Unix-like system, but has not been tested under Windows.
 
+To fit the model, run the `fitModel` function (`?fitModel` for help). There is a plot method for the object returned by this function, which shows population assignment probabilities for each line, similar to the plots that are used with `STRUCTURE`.
+