From 117bfad6eadd91769db91052f93aa6f98a67dab7 Mon Sep 17 00:00:00 2001
From: EvaLiyt <liyutong0617@gmail.com>
Date: Tue, 4 Jun 2024 14:59:25 +1200
Subject: [PATCH] add getBranchRates method in TimeTree.java

---
 .../src/main/java/lphy/base/evolution/tree/TimeTree.java | 9 +++++++++
 .../test/java/lphy/base/evolution/tree/TimeTreeTest.java | 9 +++++++++
 2 files changed, 18 insertions(+)

diff --git a/lphy-base/src/main/java/lphy/base/evolution/tree/TimeTree.java b/lphy-base/src/main/java/lphy/base/evolution/tree/TimeTree.java
index e71fd7ad..fbc434c4 100644
--- a/lphy-base/src/main/java/lphy/base/evolution/tree/TimeTree.java
+++ b/lphy-base/src/main/java/lphy/base/evolution/tree/TimeTree.java
@@ -332,4 +332,13 @@ public TimeTreeNode getLabeledNode(String label) {
         }
         return null;
     }
+
+    @MethodInfo(description = "get the array of all branch rates in the tree.")
+    public Double[] getBranchRates(){
+        List<Double> branchRates = new ArrayList<>();
+        for (TimeTreeNode node: nodes){
+            branchRates.add(node.getBranchRate());
+        }
+        return branchRates.toArray(Double[]::new);
+    }
 }
diff --git a/lphy-base/src/test/java/lphy/base/evolution/tree/TimeTreeTest.java b/lphy-base/src/test/java/lphy/base/evolution/tree/TimeTreeTest.java
index 9879d8c6..b7404f31 100644
--- a/lphy-base/src/test/java/lphy/base/evolution/tree/TimeTreeTest.java
+++ b/lphy-base/src/test/java/lphy/base/evolution/tree/TimeTreeTest.java
@@ -104,4 +104,13 @@ void getOldestNode() {
         );
     }
 
+    @Test
+    void getBranchRatesTest() {
+        Double[] branchRates = tree.getBranchRates();
+        List<TimeTreeNode> nodes = tree.getNodes();
+        for (int i = 0; i<branchRates.length; i++){
+            // ensure the index matching
+            assertEquals(nodes.get(i).getBranchRate(), branchRates[i]);
+        }
+    }
 }