X-Git-Url: https://git.donarmstrong.com/?a=blobdiff_plain;f=R%2FrTrait.R;h=2c18d0940af08ada1fcaf1a67c8cba9b5402b1ed;hb=2e7fdb027f189e7c8eb771dadaf195c43bd41d5f;hp=642591e6e2523d2ac0da057b861f1831546f6621;hpb=b41db3e6053bfb0e4b9000d5de678d65a1af9670;p=ape.git

diff --git a/R/rTrait.R b/R/rTrait.R
index 642591e..2c18d09 100644
--- a/R/rTrait.R
+++ b/R/rTrait.R
@@ -1,8 +1,8 @@
-## rTrait.R (2010-02-03)
+## rTrait.R (2011-04-02)
 
 ##   Trait Evolution
 
-## Copyright 2010 Emmanuel Paradis
+## Copyright 2010-2011 Emmanuel Paradis
 
 ## This file is part of the R-package `ape'.
 ## See the file ../COPYING for licensing issues.
@@ -10,7 +10,7 @@
 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)
+             ancestor = FALSE, root.value = 1, ...)
 {
     if (is.null(phy$edge.length))
         stop("tree has no branch length")
@@ -18,7 +18,7 @@ rTraitDisc <-
         stop("at least one branch length negative")
 
     if (is.character(model)) {
-        switch(model, "ER" = {
+        switch(toupper(model), "ER" = {
                    if (length(rate) != 1)
                        stop("`rate' must have one element")
                    Q <- matrix(rate, k, k)
@@ -56,12 +56,14 @@ rTraitDisc <-
 
     if (is.function(model)) {
         environment(model) <- environment() # to find 'k'
-        for (i in N:1) x[des[i]] <- model(x[anc[i]], el[i])
+        for (i in N:1) x[des[i]] <- model(x[anc[i]], el[i], ...)
     } else {
-        diag(Q) <- -rowSums(Q)
         freq <- rep(freq, each = k)
+        Q <- Q * freq
+        diag(Q) <- 0
+        diag(Q) <- -rowSums(Q)
         for (i in N:1) {
-            p <- matexpo(Q * freq * el[i])[x[anc[i]], ]
+            p <- matexpo(Q * el[i])[x[anc[i]], ]
             x[des[i]] <- .Internal(sample(k, size = 1, FALSE, prob = p))
         }
     }
@@ -79,8 +81,8 @@ rTraitDisc <-
 }
 
 rTraitCont <-
-    function(phy, model = "BM", sigma = 0.1, alpha = 1,
-             theta = 0, ancestor = FALSE, root.value = 0)
+    function(phy, model = "BM", sigma = 0.1, alpha = 1, theta = 0,
+             ancestor = FALSE, root.value = 0, linear = TRUE, ...)
 {
     if (is.null(phy$edge.length))
         stop("tree has no branch length")
@@ -100,9 +102,9 @@ rTraitCont <-
 
     if (is.function(model)) {
         environment(model) <- environment()
-        for (i in N:1) x[des[i]] <- model(x[anc[i]], el[i])
+        for (i in N:1) x[des[i]] <- model(x[anc[i]], el[i], ...)
     } else {
-        model <- pmatch(model, c("BM", "OU"))
+        model <- pmatch(toupper(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")
@@ -113,8 +115,12 @@ rTraitCont <-
             if (length(theta) == 1) theta <- rep(theta, N)
             else if (length(theta) != N)
                 stop("'theta' must have one or Nedge(phy) elements")
+            if (!linear) model <- model + 1L
         }
-        .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")
+        .C("rTraitCont", as.integer(model), as.integer(N),
+           as.integer(anc - 1L), as.integer(des - 1L), el,
+           as.double(sigma), as.double(alpha), as.double(theta), x,
+           DUP = FALSE, NAOK = TRUE, PACKAGE = "ape")
     }
 
     if (ancestor) {
@@ -126,3 +132,41 @@ rTraitCont <-
     }
     x
 }
+
+rTraitMult <-
+    function(phy, model, p = 1, root.value = rep(0, p), ancestor = FALSE,
+             as.factor = NULL, ...)
+{
+    phy <- reorder(phy, "pruningwise")
+    n <- length(phy$tip.label)
+    m <- phy$Nnode
+    N <- dim(phy$edge)[1]
+    ROOT <- n + 1L
+
+    x <- matrix(0, n + m, p)
+    x[ROOT, ] <- root.value
+
+    anc <- phy$edge[, 1]
+    des <- phy$edge[, 2]
+
+    el <- phy$edge.length
+    if (is.null(el)) el <- numeric(N)
+
+    for (i in N:1) x[des[i], ] <- model(x[anc[i], ], el[i], ...)
+
+    if (ancestor) {
+        if (is.null(phy$node.label)) phy <- makeNodeLabel(phy)
+        rownames(x) <- c(phy$tip.label, phy$node.label)
+    } else {
+        x <- x[1:n, , drop = FALSE]
+        rownames(x) <- phy$tip.label
+    }
+    x <- as.data.frame(x)
+    if (!is.null(as.factor)) {
+        for (i in as.factor) {
+            y <- x[, i]
+            x[, i] <- factor(y, labels = LETTERS[1:length(unique(y))])
+        }
+    }
+    x
+}