R/maketikz.R

In causalDisco: Tools for Causal Discovery on Observational Data

#' Generate Latex tikz code for plotting a temporal DAG or PDAG. 
#'
#' @details Note that it is necessary to read in relevant tikz libraries in the 
#' Latex preamble. The relevant lines of code are (depending a bit on parameter settings): \cr
#' \code{\\usepackage{tikz}} \cr
#' \code{\\usetikzlibrary{arrows,shapes,snakes,automata,backgrounds,petri}} \cr
#' \code{\\usepackage{pgfplots}}
#'
#' @param model \code{tpdag} or \code{tamat} object to plot. 
#' @param xjit How much should nodes within a period be jittered horizontally.
#' @param yjit Vertical distance between nodes within a period.
#' @param markperiods If \code{TRUE}, gray boxes are drawn behind each period. 
#' @param xpgap Horizontal gap between different periods. 
#' @param annotateEdges If \code{TRUE}, add a text annotation to edges. If \code{annotationlabels}
#' are supplied, these labels will be used. Otherwise, the value in the inputted adjacency matrix corresponding
#' to the edge will be used. 
#' @param addAxis If \code{TRUE}, a horizontal axis with period labels are added. 
#' @param varLabels Optional labels for nodes (variables). Should be given as a named list, where 
#' the name is the variable name, and the entry is the label, e.g. \code{list(vname = "Label for vname")}. 
#' @param periodLabels Optional labels for periods. Should be given as a named list, where 
#' the name is the period name (as stored in the \code{tamat}), and the entry is the label, 
#' e.g. \code{list(periodname = "Label for period")}. 
#' @param annotationLabels Optional labels for edge annotations. Only used if \code{annotateEdges = TRUE}. 
#' Should be given as a named list, where the name is the edge annotation (as stored in the \code{tamat}), 
#' and the entry is the label, e.g. \code{list(h = "High")}. 
#' @param clipboard If \code{TRUE}, the tikz code is not printed, but instead copied to the clipboard,
#' so it can easily be pasted into a Latex document. 
#' @param colorAnnotate Named list of colors to use to mark edge annotations instead of labels. This 
#' overrules \code{annotateEdges} and both are not available at the same time. The list should be given with
#' annotations as names and colors as entries, e.g. \code{list(h = "blue")}. 
#' 
#' @return Silently returns a character vector with lines of tikz code. The function 
#' furthermore has a side-effect. If \code{clipboard = TRUE}, the side-effect is that the tikz 
#' code is also copied to the clipboard. If \code{clipboard = FALSE}, the tikz code is instead printed  
#' in the console.
#'
#' @importFrom clipr write_clip
#' @export
maketikz <- function(model, xjit = 2, yjit = 2,
                       markperiods = TRUE, xpgap = 4,
                       annotateEdges = NULL,
                       addAxis = TRUE,
                       varLabels = NULL,
                       periodLabels = NULL,
                       annotationLabels = NULL,
                       clipboard = TRUE,
                       colorAnnotate = NULL) {
  if ("tpdag" %in% class(model)) {
    tamat <- model$amat
    order <- model$order
  } else if ("tamat" %in% class(model)) {
    tamat <- model
    order <- attr(model, "order")
  } else {
    stop("Input model must be of class tpdag or tamat")
  }
  
  amat <- tamat$amat
  
  if (!is.null(colorAnnotate) & is.null(annotateEdges)) {
    annotateEdges <- FALSE
  }
  
  if (is.null(annotateEdges)) {
    annotateEdges <- !all(amat %in% c(0,1))
  } 
  
  vars <- colnames(amat)
  nvar <- length(vars)
  nperiod <- length(order)
  
  vparts <- strsplit(vars, split = "_")
  periods <- sapply(vparts, function(x) x[1])
  vnames <- sapply(vparts, function(x) paste(x[-1], collapse = "\\_"))
  neachperiod <- as.numeric(table(periods)[order])
  maxpvar <- max(neachperiod)[1]
  maxypos <- maxpvar * yjit
  
  if (is.null(varLabels)) {
    varLabels <- vnames
  } else {
    varLabels <- unlist(varLabels[rownames(amat)])
  }
  
  if (is.null(periodLabels)) {
    periodLabels <- order 
  }
  
  out <- c("%TIKZ FIG MADE BY CAUSALDISCO", "\\begin{tikzpicture}")
  
  out <- c(out,  paste("[every node/.style={font=\\small, align = center},",
                       "every edge/.append style={nodes={font=\\itshape\\scriptsize}}]"))
  if (annotateEdges & !is.null(annotationLabels)) {
    for (i in 1:length(annotationLabels)) {
      thisAnnotation <- names(annotationLabels)[i]
      thisLab <- annotationLabels[[i]]
      amat[amat == thisAnnotation] <- thisLab
    }
  }
  
  pcounter <- 1
  
  for (i in 1:nvar) {
    thisperiod <- periods[i]
    thisvname <- varLabels[i]
    thispno <- which(order == thisperiod)
    thisptotal <- sum(periods == thisperiod)
    
    xpos <- (thispno-1)*xpgap + (thispno-1)*xjit + (i %% 2) * xjit
    ydist <- floor(maxypos / neachperiod[thispno])
    ypos <- (maxypos - ((neachperiod[thispno] - 1) * ydist))/2 + ydist * (pcounter-1) 
    
    out <- c(out, paste("\\node (", i, ") at (", xpos, ",", ypos, ") {", 
                        thisvname, "};", sep = ""))
    
    if (addAxis & pcounter == 1) {
      out <- c(out, paste("\\node at (", (thispno-1)*xpgap + (thispno-1)*xjit  + xjit/2, 
                          ",", "-0.5", ") {", 
                          periodLabels[thispno], "};", sep = ""))
    }
    
    if (pcounter == thisptotal) {
      pcounter <- 1 
    } else { 
      pcounter <- pcounter + 1
    }
  }
  
  allundir <- list()
  
  for (i in 1:nvar) {
    thesechildren <- which(amat[, i] != 0)
    theseparents <- which(amat[i, ] != 0)
    theseundir <- intersect(thesechildren, theseparents)
    
    thesetruechildren <- setdiff(thesechildren, theseundir) 
    
    theseundir <- theseundir[theseundir < i] #only store when smaller, avoids duplicates
    if (length(theseundir) > 0 ) {
      allundir[[i]] <- theseundir
    }
    if (length(thesetruechildren) > 0) {
      if (!annotateEdges & is.null(colorAnnotate)) {
        out <- c(out, paste("\\draw [->] (", i, ") edge (", thesetruechildren, ");", sep = ""))
      } 
      if (annotateEdges) {
        #  browser()
        out <- c(out, paste("\\draw [->] (", i, ") edge node [above,sloped] {", amat[thesetruechildren, i],
                            "} (", thesetruechildren, ");", sep = ""))
      }
      if (!is.null(colorAnnotate)) {
        out <- c(out, paste("\\draw [->, ", unlist(colorAnnotate[amat[thesetruechildren, i]]), 
                            "] (", i, ") edge (", thesetruechildren, ");", sep = ""))
      }
    }
  }
  
  n_undir <- length(allundir)
  if (n_undir > 0) {
    for (i in 1:length(allundir)) {
      theseneigh <- allundir[[i]]
      if (length(theseneigh) > 0) {
        if (!annotateEdges) {
          out <- c(out, paste("\\draw [-] (", i, ") edge (", theseneigh, ");", sep = ""))
        }
        if (annotateEdges) {
          out <- c(out, paste("\\draw [-] (", i, ") edge node [above,sloped] {", amat[theseneigh, i],
                              "} (", theseneigh, ");", sep = ""))
        }
      }
    }
  }
  
  if (addAxis) {
    max_xpos <- (nperiod-1)*xpgap + xjit*nperiod
    out <- c(out, paste("\\draw [-] (-1,0) edge (", max_xpos + 1, ",0);", sep = ""))
  }
  
  if (markperiods) { 
    out <- c(out, "\\begin{pgfonlayer}{background}", 
             "\\filldraw [join=round,black!10]")
    
    for (j in 1:nperiod) {
      xpos_1 <- (j - 1) * xpgap + (j - 1) * xjit
      ypos_1 <- 0
      xpos_2 <- xpos_1 + xjit
      ypos_2 <- (maxpvar) * yjit 
      
      out <- c(out, paste("(", xpos_1, ",", ypos_1, ") rectangle (",
                          xpos_2, ",", ypos_2, ")", sep = ""))
    }
    out <- c(out, "; \\end{pgfonlayer}")
  }
  
  out <- c(out, "\\end{tikzpicture}")
  
  if (clipboard) {
    write_clip(out)
  } else {
    cat(paste(out, collapse = "\n"))
  }
  invisible(out)
}