Function to graph discrete character data matrix next to the tips of a plotted tree

I recently posted a couple of demos (1, 2) showing how to add a grid of discrete character data next to the tips of plotted tree.

The following is a function that automates some parts of this. I have so far not attempted to automate the space allowance for the data matrix - instead leaving that for the user to specify.

The function .check.pkg is an internal phytools function that checks to see if RColorBrewer is installed. I borrowed it from geiger.

.check.pkg<-phytools:::.check.pkg

plotTree.datamatrix<-function(tree,X,...){
    N<-Ntip(tree)
    ss<-sapply(X,function(x) levels(x))
    k<-sapply(ss,length)
    if(hasArg(fsize)) fsize<-list(...)$fsize
    else fsize<-40*par()$pin[2]/par()$pin[1]/Ntip(tree)
    if(hasArg(xexp)) xexp<-list(...)$xexp
    else xexp<-1.3
    if(hasArg(yexp)) yexp<-list(...)$yexp
    else yexp<-1.05
    if(hasArg(colors)) colors<-list(...)$colors
    else {
        chk<-.check.pkg("RColorBrewer")
        if(!chk) brewer.pal<-function(...) NULL
        else {
            palettes<-c("Accent","Dark2","Paired","Pastel1","Pastel2",
                "Set1","Set2","Set3")
            while(length(palettes)<length(k)) palettes<-c(palettes,palettes)
            BREWER.PAL<-function(k,pal) brewer.pal(max(k,3),pal)[1:k]
            colors<-mapply(setNames,mapply(BREWER.PAL,k,
                palettes[1:length(ss)]),ss,SIMPLIFY=FALSE)
            }
        }
    if(hasArg(sep)) sep<-list(...)$sep
    else sep<-0.5
    if(hasArg(srt)) srt<-list(...)$srt
    else srt<-60
    if(hasArg(space)) space<-list(...)$space
    else space<-0
    cw<-reorder(tree,"cladewise")
    X<-X[cw$tip.label,]
    plotTree(cw,plot=FALSE,fsize=fsize)
    obj<-get("last_plot.phylo",envir=.PlotPhyloEnv)
    plotTree(cw,lwd=1,ylim=c(0,obj$y.lim[2]*yexp),
        xlim=c(0,obj$x.lim[2]*xexp),fsize=fsize,
        ftype="off")
    obj<-get("last_plot.phylo",envir=.PlotPhyloEnv)
    h<-max(obj$xx)
    for(i in 1:Ntip(cw)){
        lines(c(obj$xx[i],h),rep(obj$yy[i],2),lty="dotted")
        text(h,obj$yy[i],cw$tip.label[i],cex=fsize,pos=4,font=3,offset=0.1)
    }
    s<-max(fsize*strwidth(cw$tip.label))
    start.x<-1.05*h+s
    half<-0.5*(1-space)
    for(i in 1:ncol(X)){
            text(start.x,max(obj$yy)+1,colnames(X)[i],pos=4,srt=srt,
            cex=fsize,offset=0)
        for(j in 1:nrow(X)){
            xy<-c(start.x,obj$yy[j])
            y<-c(xy[2]-half,xy[2]+half,xy[2]+half,xy[2]-half)
            asp<-(par()$usr[2]-par()$usr[1])/(par()$usr[4]-par()$usr[3])*
                par()$pin[2]/par()$pin[1]
            x<-c(xy[1]-half*asp,xy[1]-half*asp,xy[1]+half*asp,xy[1]+half*asp)
            polygon(x,y,col=colors[[i]][as.character(X[[i]][j])])
        }
        start.x<-start.x+(1+sep)*asp
    }
    obj<-list(fsize=fsize,
        colors=colors,
        end.x=start.x)
    invisible(obj)
}

I guess we can try it now:

library(phytools)
tree<-rtree(n=40)
Q1<-matrix(c(-1,1,1,-1),2,2,dimnames=list(0:1,0:1))
x1<-sim.Mk(tree,Q1)
Q2<-matrix(c(-1,1,0,1,-2,1,0,1,-1),3,3,dimnames=list(letters[1:3],
    letters[1:3]))
x2<-sim.Mk(tree,Q2)
Q3<-matrix(c(-0.5,0.5,0.5,-0.5),2,2,dimnames=list(c("rough","smooth"),
    c("rough","smooth")))
x3<-sim.Mk(tree,Q3)
Q4<-matrix(c(-3,1,1,1,1,-3,1,1,1,1,-3,1,1,1,1,-3),4,4,
    dimnames=list(LETTERS[1:4],LETTERS[1:4]))
x4<-sim.Mk(tree,Q4)
X<-data.frame(x1,x2,x3,x4)
colnames(X)<-c("Trait 1","Trait 2","Trait 3","Trait 4")
object<-plotTree.datamatrix(tree,X,sep=0,srt=90,yexp=1.1,xexp=1.1,
    fsize=0.8,space=0.2)

The object returned invisibly by the function (here object) so far contains only a few items: the font size (since that is set automatically by the function if not specified by the user), a list of the color schemes used for each column of the data matrix, and the right-most x coordinate of the plotted data. This latter element might be used, for instance, to figure out where to draw our legends. For example:

object<-plotTree.datamatrix(tree,X,sep=0,srt=70,yexp=1.05,fsize=0.8)
x<-object$end.x+diff(par()$usr[1:2])*0.01
y<-Ntip(tree)
for(i in 1:ncol(X)){
    text(x,y,colnames(X)[i],pos=4,cex=0.9,offset=0)
    add.simmap.legend(colors=object$colors[[i]],shape="square",
        prompt=FALSE,x=x,y=y-2*strheight("W")*0.9,fsize=0.9)
    y<-y-1.5*0.9*strheight("W")*(length(object$colors[[i]])-1)-6
}

That's pretty much what I was going for.

The following shows how we might show presence/absence data:

X<-sim.Mk(tree,Q1,nsim=10)
colors<-setNames(replicate(ncol(X),setNames(c("white","darkgrey"),0:1),
    simplify=FALSE),colnames(X))
object<-plotTree.datamatrix(tree,X,colors=colors,sep=0,srt=90,yexp=1.05,
    fsize=0.8)
add.simmap.legend(colors=setNames(c("white","darkgrey"),
    c("absent","present")),shape="square",prompt=FALSE,x=0,
    y=3,fsize=0.9)