R/plottest.R

In alemiani/explora: Explore and Visualize Data

plottest <- function(x, y, group = NULL, subplots = TRUE, sameXYaxes = TRUE, onlyResults = FALSE) {

  if(!missing(group)) {
    if(nlevels(as.factor(group)) > 3) stop('Hey, no more than three groups in here!!')
  }



  mycol = rgb(245, 245, 255, max = 255, alpha = 220) ### set legend color

  ############################################################################################################ make Numbers
  y.name <- paste(substitute(y))[[3]]
  x.name <- paste(substitute(x))[[3]]

  formula <- paste(y.name,x.name, sep = ' ~ ')
  main <- formula

  ###################################################### 1 group

  if(is.null(group)) {
    result.all <- cor.test(x, y)
    N <- paste('(N=',result.all$parameter+2,')', sep = '')
    a.S <- summary(lm(y ~ x))
    text.1G <- paste(paste('r =', round(result.all$estimate, 2)),
                     paste('p =', round(result.all$p.value, 4)),N,
                     paste('[R2=', round(a.S$r.squared,2),']', sep=''),
                     sep = '; ')
    results <- paste(formula, text.1G, sep = '\n')
  }

  ###################################################### 2 groups
  if(nlevels(as.factor(group)) >= 2) {
    group       <- droplevels(as.factor(group))
    gg          <- levels(as.factor(group))
    lev         <- (as.numeric(as.factor(group)))
    levels(lev) <- c("black", "red")

    result.all <- cor.test(x, y)
    N <- paste('(N=',result.all$parameter+2,')', sep = '')
    a.S <- summary(lm(y ~ x))
    text.1G <- paste(paste('r =', round(result.all$estimate, 2)),
                     paste('p =', round(result.all$p.value, 4)),N,
                     paste('[R2=', round(a.S$r.squared,2),']', sep=''),
                     sep = '; ')
    results <- paste(formula, text.1G, sep = '\n')


    a.S <- summary(lm(y ~ x))
    black.S <- summary(lm(y[lev==1] ~ x[lev==1]))
    red.S <- summary(lm(y[lev==2] ~ x[lev==2]))

    result.black <- cor.test(x[lev==1], y[lev==1])
    N <- paste('(N=',result.black$parameter+2,')', sep = '')
    text.black <- paste(paste(paste(substitute(gg))[[1]],'(black) :',
                              'r =', round(result.black$estimate, 2)),
                        paste('p =', round(result.black$p.value, 4)),N,
                        paste('[R2=', round(black.S$r.squared,2),']', sep=''),
                        sep = '; ')

    result.red <- cor.test(x[lev==2], y[lev==2])
    N <- paste('(N=',result.red$parameter+2,')', sep = '')
    text.red <- paste(paste(paste(substitute(gg))[[2]],'(red) :',
                            'r =', round(result.red$estimate, 2)),
                      paste('p =', round(result.red$p.value, 4)),N,
                      paste('[R2=', round(red.S$r.squared,2),']', sep=''),
                      sep = '; ')

    text.2G <- paste(text.1G, text.black, text.red, sep = '\n')
    results <- paste(formula, text.2G, sep = '\n')

    ###################################################### 3 groups
    if(nlevels(as.factor(group)) == 3) {
      group       <- droplevels(as.factor(group))
      gg          <- levels(as.factor(group))
      lev         <- (as.numeric(as.factor(group)))
      levels(lev) <- c("black", "red", "green")

      result.green <- cor.test(x[lev==3], y[lev==3])
      N <- paste('(N=',result.green$parameter+2,')', sep = '')
      green.S <- summary(lm(y[lev==3] ~ x[lev==3]))
      text.green <- paste(paste(paste(substitute(gg))[[3]],'(green) :',
                                'r =', round(result.green$estimate, 2)),
                          paste('p =', round(result.green$p.value, 4)),N,
                          paste('[R2=', round(green.S$r.squared,2),']', sep=''),
                          sep = '; ')
      text.3G <- paste(text.1G, text.black, text.red, text.green, sep = '\n')
      results <- paste(formula, text.3G, sep = '\n')
    }

  }

  if(onlyResults == TRUE) {
    cat(paste(results,sep = '\n'))
  } else {




  ############################################################################################################ make Plots
  if(sameXYaxes == TRUE) { ########################## Set XY limits
    xlim <- range(x, na.rm = T)
    ylim <- range(y, na.rm = T)
    } else {
      xlim <- NULL
      ylim <- NULL
    }

  ###################################################### plot 1 group
  if(missing(group)) {

    plot(x, y,
         main = main,
         xlab = x.name,
         ylab = y.name,
         xlim = xlim,
         ylim = ylim
         )
    abline(lm(y ~ x), col = 'blue', lwd = 2, lty = 6)
    legend('top', text.1G, xjust = 1, yjust = 0.5, adj = 0, box.lwd = 0, box.col = mycol, bg = mycol)

    cat(paste(formula, text.1G, sep = '\n'))

  }


  else {
    ##############################################################################. Plot 2 groups

    if(nlevels(as.factor(group)) >= 2) {

      if(subplots == TRUE) {
        layout(mat=matrix(c(1,1,2,3), ncol=2, byrow=TRUE))

        plot(x, y, pch = lev, col = lev,
             main = main,
             xlab = x.name,
             ylab = y.name,
             xlim = xlim,
             ylim = ylim
        )
        abline(lm(y ~ x), col = 'blue', lwd = 2, lty = 6)
        abline(lm(y[lev==1] ~ x[lev==1]), col = 'black')
        abline(lm(y[lev==2] ~ x[lev==2]), col = 'red')
        legend('top', text.2G, xjust = 1, yjust = 0.5, adj = 0, box.lwd = 0, box.col = mycol, bg = mycol)

        # make subplots
        plot(x[lev==1], y[lev==1],
             xlab = x.name,
             ylab = y.name,
             main = paste(substitute(gg))[[1]],
             xlim = xlim,
             ylim = ylim
        )
        abline(lm(y[lev==1] ~ x[lev==1]), col = 'blue', lwd = 2, lty = 6)

        plot(x[lev==2], y[lev==2],
             xlab = x.name,
             ylab = y.name,
             main = paste(substitute(gg))[[2]],
             xlim = xlim,
             ylim = ylim
        )
        abline(lm(y[lev==2] ~ x[lev==2]), col = 'blue', lwd = 2, lty = 6)

        par(mfrow=c(1,1))
        } else {
          plot(x, y, pch = lev, col = lev,
               main = main,
               xlab = x.name,
               ylab = y.name,
               xlim = xlim,
               ylim = ylim
          )
          abline(lm(y ~ x), col = 'blue', lwd = 2, lty = 6)
          abline(lm(y[lev==1] ~ x[lev==1]), col = 'black')
          abline(lm(y[lev==2] ~ x[lev==2]), col = 'red')
          legend('top', text.2G, xjust = 1, yjust = 0.5, adj = 0, box.lwd = 0, box.col = mycol, bg = mycol)
        }
      if(nlevels(as.factor(group)) == 2) {
        cat(paste(formula, text.2G,sep = '\n'))

      }




    }




    ############################################################################## 3 groups

    if(nlevels(as.factor(group)) == 3) {

      if(subplots == TRUE) {

        layout(mat=matrix(c(1,1,1, 2,3,4), ncol=3, byrow=TRUE))

        plot(x, y, pch = lev, col = lev,
             main = main,
             xlab = x.name,
             ylab = y.name,
             xlim = xlim,
             ylim = ylim
        )
        abline(lm(y ~ x), col = 'blue', lwd = 2, lty = 6)
        abline(lm(y[lev==1] ~ x[lev==1]), col = 'black')
        abline(lm(y[lev==2] ~ x[lev==2]), col = 'red')
        abline(lm(y[lev==3] ~ x[lev==3]), col = 'green')
        legend('top', text.3G, xjust = 1, yjust = 0.5, adj = 0, box.lwd = 0, box.col = mycol, bg = mycol)


        # make subplots
        plot(x[lev==1], y[lev==1],
             xlab = x.name,
             ylab = y.name,
             main = paste(substitute(gg))[[1]],
             xlim = xlim,
             ylim = ylim
        )
        abline(lm(y[lev==1] ~ x[lev==1]), col = 'blue', lwd = 2, lty = 6)

        plot(x[lev==2], y[lev==2],
             xlab = x.name,
             ylab = y.name,
             main = paste(substitute(gg))[[2]],
             xlim = xlim,
             ylim = ylim
        )
        abline(lm(y[lev==2] ~ x[lev==2]), col = 'blue', lwd = 2, lty = 6)

        plot(x[lev==3], y[lev==3],
             xlab = x.name,
             ylab = y.name,
             main = paste(substitute(gg))[[3]],
             xlim = xlim,
             ylim = ylim
        )
        abline(lm(y[lev==3] ~ x[lev==3]), col = 'blue', lwd = 2, lty = 6)

        par(mfrow=c(1,1))
      } else {
        plot(x, y, pch = lev, col = lev,
             main = main,
             xlab = x.name,
             ylab = y.name,
             xlim = xlim,
             ylim = ylim
        )
        abline(lm(y ~ x), col = 'blue', lwd = 2, lty = 6)
        abline(lm(y[lev==1] ~ x[lev==1]), col = 'black')
        abline(lm(y[lev==2] ~ x[lev==2]), col = 'red')
        abline(lm(y[lev==3] ~ x[lev==3]), col = 'green')
        legend('top', text.3G, xjust = 1, yjust = 0.5, adj = 0, box.lwd = 0, box.col = mycol, bg = mycol)

        cat(paste(formula, text.3G,sep = '\n'))
      }

      cat(paste(formula, text.3G,sep = '\n'))
    }

  }

  }
}



# set.seed(1235)
# x1 <- rnorm(50,10,5)
# x2 <- c(rep(0, 25), rep(1,25))
# y <- x1*2 + x2*5 + x1*x2*2 + rnorm(50,0,10)
# plottest(df$x1, df$y, df$group)