From ebd9d80c34b5f5d1866476de2d736e78ffd725e0 Mon Sep 17 00:00:00 2001
From: Klaus Schliep <klaus.schliep@gmail.com>
Date: Tue, 16 Jan 2024 17:48:20 +0100
Subject: [PATCH] add_edge_length assign now also edge length to an splits
 graph (networx)

---
 R/draw_CI.R            | 18 ++++++++++++++----
 man/add_edge_length.Rd |  3 ++-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/R/draw_CI.R b/R/draw_CI.R
index c4d78dff..a92dea8d 100644
--- a/R/draw_CI.R
+++ b/R/draw_CI.R
@@ -1,5 +1,6 @@
 edge_length_matrix <- function(tree, trees, rooted=TRUE){
   if(!inherits(trees, "multiPhylo")) stop("Trees must be of class multiPhylo!")
+  if(inherits(tree, "networx")) rooted <- FALSE
   trees <- .uncompressTipLabel(trees) |> .compressTipLabel(ref=tree$tip.label)
   if(!rooted){
     trees <- unroot(trees)
@@ -13,7 +14,9 @@ edge_length_matrix <- function(tree, trees, rooted=TRUE){
     el[x$edge[,2]] <- x$edge.length
     el
   }
-  bp <- bip(tree)
+  if(inherits(tree, "networx")) bp <- tree$splits
+  else bp <- bip(tree)
+
   if(!rooted) bp <- SHORTwise(bp)
   m <- length(bp)
   res <- matrix(NA_real_, length(trees), m)
@@ -35,7 +38,8 @@ edge_length_matrix <- function(tree, trees, rooted=TRUE){
 ##' @title Assign and compute edge lengths from a sample of trees
 ##' @description This command can infer some average edge lengths and assign
 ##' them from a (bootstrap/MCMC) sample.
-##' @param tree tree where edge lengths are assigned to.
+##' @param tree a phylogenetic tree or splitnetwork where edge lengths are
+##' assigned to.
 ##' @param trees an object of class multiPhylo, where the average for the edges
 ##' is computed from.
 ##' @param fun a function to compute the average (default is median).
@@ -63,8 +67,14 @@ add_edge_length <- function(tree, trees, fun=\(x)median(na.omit(x)),
   if(!rooted) tree <- unroot(tree)
   X <- edge_length_matrix(tree, trees, rooted)
   nh <- apply(X, 2, fun)
-  if(rooted) tree$edge.length <- nh[tree$edge[,1]] - nh[tree$edge[,2]]
-  else tree$edge.length <- nh[tree$edge[,2]]
+  if(inherits(tree, "networx")){
+    tree$edge.length <- nh[tree$splitIndex]
+    attr(tree$splits, "weights") <- nh
+  }
+  else{
+    if(rooted) tree$edge.length <- nh[tree$edge[,1]] - nh[tree$edge[,2]]
+    else tree$edge.length <- nh[tree$edge[,2]]
+  }
   tree
 }
 
diff --git a/man/add_edge_length.Rd b/man/add_edge_length.Rd
index 734c8874..337e6a90 100644
--- a/man/add_edge_length.Rd
+++ b/man/add_edge_length.Rd
@@ -8,7 +8,8 @@ add_edge_length(tree, trees, fun = function(x) median(na.omit(x)),
   rooted = all(is.rooted(trees)))
 }
 \arguments{
-\item{tree}{tree where edge lengths are assigned to.}
+\item{tree}{a phylogenetic tree or splitnetwork where edge lengths are
+assigned to.}
 
 \item{trees}{an object of class multiPhylo, where the average for the edges
 is computed from.}