new files for trait simulations

[ape.git] / R / rTrait.R

diff --git a/R/rTrait.R b/R/rTrait.R
new file mode 100644 (file)
index 0000000..3359df6
--- /dev/null
+++ b/R/rTrait.R
@@ -0,0 +1,127 @@
+## rTrait.R (2010-01-11)
+
+##   Trait Evolution
+
+## Copyright 2010 Emmanuel Paradis
+
+## This file is part of the R-package `ape'.
+## See the file ../COPYING for licensing issues.
+
+rTraitDisc <-
+    function(phy, model = "ER", k = if (is.matrix(model)) ncol(model) else 2,
+             rate = 0.1, states = LETTERS[1:k], freq = rep(1/k, k),
+             ancestor = FALSE, root.value = 1)
+{
+    if (is.null(phy$edge.length))
+        stop("tree has no branch length")
+    if (any(phy$edge.length < 0))
+        stop("at least one branch length negative")
+
+    if (is.character(model)) {
+        switch(model, "ER" = {
+                   if (length(rate) != 1)
+                       stop("`rate' must have one element")
+                   Q <- matrix(rate, k, k)
+               }, "ARD" = {
+                   if (length(rate) != k*(k - 1))
+                       stop("`rate' must have k(k - 1) elements")
+                   Q <- matrix(0, k, k)
+                   Q[col(Q) != row(Q)] <- rate
+               }, "SYM" = {
+                   if (length(rate) != k*(k - 1)/2)
+                       stop("`rate' must have k(k - 1)/2 elements")
+                   Q <- matrix(0, k, k)
+                   sel <- col(Q) < row(Q)
+                   Q[sel] <- rate
+                   Q <- t(Q)
+                   Q[sel] <- rate
+               })
+    }
+    if (is.matrix(model)) {
+        Q <- model
+        if (ncol(Q) != nrow(Q))
+            stop("the matrix given as `model' must be square")
+    }
+
+    phy <- reorder(phy, "pruningwise")
+    n <- length(phy$tip.label)
+    N <- dim(phy$edge)[1]
+    ROOT <- n + 1L
+    x <- integer(n + phy$Nnode)
+    x[ROOT] <- as.integer(root.value)
+
+    anc <- phy$edge[, 1]
+    des <- phy$edge[, 2]
+    el <- phy$edge.length
+
+    if (is.function(model)) {
+        environment(model) <- environment() # to find 'k'
+        for (i in N:1) x[des[i]] <- model(x[anc[i]], el[i])
+    } else {
+        diag(Q) <- -rowSums(Q)
+        for (i in N:1) {
+            p <- matexpo(Q * freq * el[i])[x[anc[i]], ]
+            x[des[i]] <- .Internal(sample(k, size = 1, FALSE, prob = p))
+        }
+    }
+
+    if (ancestor) {
+        if (is.null(phy$node.label)) phy <- makeNodeLabel(phy)
+        names(x) <- c(phy$tip.label, phy$node.label)
+    } else {
+        x <- x[1:n]
+        names(x) <- phy$tip.label
+    }
+    class(x) <- "factor"
+    levels(x) <- states
+    x
+}
+
+rTraitCont <-
+    function(phy, model = "BM", sigma = 0.1, alpha = 1,
+             theta = 0, ancestor = FALSE, root.value = 0)
+{
+    if (is.null(phy$edge.length))
+        stop("tree has no branch length")
+    if (any(phy$edge.length < 0))
+        stop("at least one branch length negative")
+
+    phy <- reorder(phy, "pruningwise")
+    n <- length(phy$tip.label)
+    N <- dim(phy$edge)[1]
+    ROOT <- n + 1L
+    x <- numeric(n + phy$Nnode)
+    x[ROOT] <- root.value
+
+    anc <- phy$edge[, 1]
+    des <- phy$edge[, 2]
+    el <- phy$edge.length
+
+    if (is.function(model)) {
+        environment(model) <- environment()
+        for (i in N:1) x[des[i]] <- model(x[anc[i]], el[i])
+    } else {
+        model <- pmatch(model, c("BM", "OU"))
+        if (length(sigma) == 1) sigma <- rep(sigma, N)
+        else if (length(sigma) != N)
+            stop("'sigma' must have one or Nedge(phy) elements")
+        if (model == 2) { # "OU"
+            if (length(alpha) == 1) alpha <- rep(alpha, N)
+            else if (length(alpha) != N)
+                stop("'alpha' must have one or Nedge(phy) elements")
+            if (length(theta) == 1) theta <- rep(theta, N)
+            else if (length(theta) != N)
+                stop("'theta' must have one or Nedge(phy) elements")
+        }
+        .C("rTraitCont", as.integer(model), as.integer(N), as.integer(anc - 1L), as.integer(des - 1L), el, sigma, alpha, theta, x, DUP = FALSE, NAOK = TRUE, PACKAGE = "ape")
+    }
+
+    if (ancestor) {
+        if (is.null(phy$node.label)) phy <- makeNodeLabel(phy)
+        names(x) <- c(phy$tip.label, phy$node.label)
+    } else {
+        x <- x[1:n]
+        names(x) <- phy$tip.label
+    }
+    x
+}