## Bruce Swihart
## March 03, 2009

## bruce.swihart@gmail.com
## Johns Hopkins School of Public Health.

## Supplemental Materials for Lasagna Plots

## some functions and examples
## please feel free to improve upon the code.

library(colorRamps)
library(RColorBrewer)
library(fields)




#######################
##
### Functions
##
#######################

## we require these packages

library(colorRamps)
library(RColorBrewer)
library(fields)


## lasagna() uses image(), but manipulates the matrix so the image
## rendered is that of just painting the elements of the matrix.

lasagna<- function(X,...){ image(t(X)[,(nrow(X):1)], col =
rgb.palette(256),...  )}


## lasagna.leg() uses image.plot() to get the legend, manipulates the
## matrix the same way lasagna() does.

lasagna.leg <- function(X,...){
image.plot(t(X)[,(nrow(X):1)], col = rgb.palette(256),...  )}


## must specify colors, for gradients, see:
## http://www.ficml.org/jemimap/style/color/wheel.html
rgb.palette <- colorRampPalette(c( c("#ddff77","#77ff66","#00cc33","#007722","#003300"),"black"))
## OR
## below defines Matlab's "jet" palette
rgb.palette <- colorRampPalette(c("#00007F", "blue", "#007FFF",
"cyan", "#7FFF7F","yellow", "#FF7F00", "red", "#7F0000")) 



## within row sort for continuous outcomes
wr.cont <- function(X, naLast=F){
  sortedWR <-  apply( X, 1, function(W) sort(W,na.last = naLast )) 
  sortedWR <- t(sortedWR) # transposed b/c every sorting function should return an intuit matrix
  sortedWR
}


## within column sort for continuous outcomes
wc.cont <- function(X, naLast=F ){
  sorted <- apply( X, 2, function(W) sort(W,na.last = naLast ) )
  sorted
}

## entire column sort for continuous outcomes
ec <- function(X,orderVar=c(), naLast=F){
    if (length(orderVar) == 1){
        perc <- apply( X, 2, function(W){ sum(W==orderVar)  })
        return(X[ ,order(perc)])
    }
    X[ ,order(orderVar)]
}


## within column for discrete outcomes
wc.disc <- function(X, orderVar=c(), colorSeq, naLast=F ){
  ##make priority mask
  for(i in 1:length(colorSeq) ){        X[ X==colorSeq[i] ] <- i  }
  sorted <- apply( X, 2, function(W) sort(W,na.last = naLast ) )
  ## undo mask
  for(i in 1:length(colorSeq)  ){        sorted[ sorted==i ] <- colorSeq[i] }
  sorted
}

## within row for discrete
wr.disc <- function(X, orderVar=c(), colorSeq, naLast=F){
  ##make priority mask
  for(i in 1:length(colorSeq) ){        X[ X==colorSeq[i] ] <- i  }
  sortedWR <-  apply( X, 1, function(W) sort(W,na.last = naLast )) 
  sortedWR <- t(sortedWR) # transposed b/c every sorting function should return an intuit matrix
  ## undo mask
  for(i in 1:length(colorSeq)  ){        sortedWR[ sortedWR==i ] <- colorSeq[i] }
  sortedWR
}



## a continuous example, simulated data

## time points, say months.  24 mos = 2 yrs
t <- 1:24
## make a random matrix for demographic group 1
e1 <- matrix(rnorm(240, 3, 2), nrow=10, ncol=24)
## make a random matrix for demographic group 2
e2 <- matrix(rnorm(240, 0, 1), nrow=10, ncol=24)
## generate oscillating repeated measures...
x <- matrix( sin(t), nrow=10, ncol=24, byrow=T)
## put into the matrices
g1 <- x+e1
g2 <- x+e2
## make group id vector
id <- c( rep(1,10), rep(2,10))
## make random vector, to scramble up the groups
scramble <- runif(20)


## g.s, a simulated dataset, with no particular order to subjects:
id.s <- id[order(scramble)]
g <- rbind(g1,g2)
g.s <- g[ order(scramble), ]

## lasagna plot of random
lasagna(g.s)
## lasagna plot with legend
lasagna.leg(g.s)

## group subjects according to group status
## do an entire row sort on group id.
g <- g.s[ order(id.s),]
lasagna(g)

## a within column sort, within groups.
## break groups apart
g1 <- g[ 1:10,]
g2 <- g[11:20,]
## do within column sorting on each group
g1.wc <- wc.cont(g1)
g2.wc <- wc.cont(g2)
## recombine
g.wc <- rbind(g1.wc, g2.wc)
## plot
lasagna(g.wc)

## an entire column sort
## applied to g.wc, within groups,
## based on mean outcome within time point
## across subjects.
g1.wc <- g.wc[ 1:10,]
g2.wc <- g.wc[11:20,]
## get the means on which to order the columns of the plot
mean1 <- apply(g1.wc,2,mean)
mean2 <- apply(g2.wc,2,mean)
## do within column sorting on each group
g1.ec <- g1.wc[,order(mean1)]
g2.ec <- g2.wc[,order(mean2)]
## recombine
g.ec <- rbind(g1.ec, g2.ec)
## plot
lasagna(g.ec)


## a within row sort, no need to break groups
g.wr <- wr.cont(g)
## plot
lasagna(g.wr)


## see them all on the same 4-plot
win.graph()
par(mfrow=c(2,3))
lasagna(g.s)
title("scrambled")
lasagna(g)
title("entire-row sort to organize groups")
lasagna(g.wr)
title("within-row sort")
lasagna(g.wc)
title("within-column, within-group sort ")
lasagna(g.ec)
title("entire-column, within-group sort")



## a discrete example
## in discrete outcomes (say, sleep state)
## one knows all the states and thus has more control over how to
## organize the sorts, and thus the need for different functions for
## continuous outcome and discrete outcomes.

## time points, say months.  24 mos = 2 yrs
t <- 1:24
## make a random matrix for demographic group 1
e1 <- matrix(rbinom(240, 2, .9), nrow=10, ncol=24)
## make a random matrix for demographic group 2
e2 <- matrix(rbinom(240, 2, .1), nrow=10, ncol=24)
## generate oscillating repeated measures...
x <- matrix( sin(t), nrow=10, ncol=24, byrow=T)
## put into the matrices
g1 <- ceiling(x+e1)
g2 <- floor(x+e2)
## make group id vector
id <- c( rep(1,10), rep(2,10))
## make random vector, to scramble up the groups
scramble <- runif(20)

## put groups together and check to make sure the
## outcomes are c(-1,0,1,2,3) only.  If not, rerun
## the data generating code above before continuing with
## the discrete example.
g <- rbind(g1,g2)
table(g)

## g.s, a simulated dataset, with no particular order to subjects:
id.s <- id[order(scramble)]
g.s <- g[ order(scramble), ]

## only have 5 states, only need 5 colors
rgb.palette <- colorRampPalette(c("#00007F", "#7FFF7F","yellow","red", "#7F0000")) 



## lasagna plot of random
lasagna(g.s)
## lasagna plot with legend,
## need to specify states to make legend only
## present colors that are in the plot.
lasagna.leg(g.s,
            nlevel = 5 ,
            axis.args=list( at=c(-1+.25,0,1,2,3-.25),
            labels=c("-1","0","1","2","3") ) )



## group subjects according to group status
## do an entire row sort on group id.
g <- g.s[ order(id.s),]
lasagna(g)

## a within column sort, within groups.
## break groups apart
g1 <- g[ 1:10,]
g2 <- g[11:20,]
## do within column sorting on each group, specifiy the order of ranking
## we add an arbitrarily large number (1000) to transform the matrix
## of states to one containing no element that could be its rank.
## for instance, we have 5 states, which means we have ranks
## 1,2,3,4,5.  We must have an input that does not contain
## those numbers, or the sorting will not work properly.
## Thus, we are left with this "hack" of transforming
## the matrix and colorSeq by 1000, and then subtracting 1000
## from the matrix returned by the function.
g1.wc <- wc.disc(g1+1000, colorSeq=c(-1,0,1,2,3)+1000 ) - 1000
g2.wc <- wc.disc(g2+1000, colorSeq=c(-1,0,1,2,3)+1000 ) - 1000
## recombine
g.wc <- rbind(g1.wc, g2.wc)
## plot
lasagna(g.wc)
## change the order, so you can better compare different states
g1.wc <- wc.disc(g1+1000, colorSeq=c(3,2,1,-1,0)+1000)  - 1000
g2.wc <- wc.disc(g2+1000, colorSeq=c(3,2,1,-1,0)+1000)  - 1000
## recombine
g.wc <- rbind(g1.wc, g2.wc)
## plot
win.graph()
lasagna(g.wc)



## an entire column sort
## applied to g.wc, within groups,
## based on mean outcome within time point
## across subjects.
g1.wc <- g.wc[ 1:10,]
g2.wc <- g.wc[11:20,]
## get the means on which to order the columns of the plot
mean1 <- apply(g1.wc,2,mean)
mean2 <- apply(g2.wc,2,mean)
## do within column sorting on each group
g1.ec <- g1.wc[,order(mean1)]
g2.ec <- g2.wc[,order(mean2)]
## recombine
g.ec <- rbind(g1.ec, g2.ec)
## plot
lasagna(g.ec)


## a within row sort, no need to break groups
g.wr <- wr.disc(g+1000,colorSeq=c(-1,0,1,2,3)+1000)-1000
## plot
lasagna(g.wr)


## see them all on the same 4-plot
win.graph()
par(mfrow=c(2,3))
lasagna(g.s)
title("discrete scrambled")
lasagna(g)
title("discrete entire-row sort to organize groups")
lasagna(g.wr)
title("discrete within-row sort")
lasagna(g.wc)
title("discrete within-column, within-group sort ")
lasagna(g.ec)
title("discrete entire-column, within-group sort")