Merge pull request #19 from soedinglab/dev

change tests

Author: benjamin-lieser

benchmark_test/felsenstein.py

@@ -46,8 +46,9 @@ def __init__(self, tree):
 
     def log_likelihood(self, S, sqrt_pi):
         if S.shape[0] == 1:
-            S = S.expand(300, -1, -1)
-            sqrt_pi = sqrt_pi.expand(300, -1)
+            # Quick hack to get single side to work for the testing, ideally it should know what L is here
+            S = S.expand(100, -1, -1)
+            sqrt_pi = sqrt_pi.expand(100, -1)
         matrices = rate_matrix_from_S(S, sqrt_pi)
         self.root.compute(matrices)
 
@@ -69,13 +70,13 @@ def gradients(tree, S, sqrt_pi):
     S.requires_grad = True
     sqrt_pi.requires_grad = True
     if S.shape[0] == 1:
-        S_expand = S.expand(300, -1, -1)
-        sqrt_pi_expand = sqrt_pi.expand(300, -1)
+        S_expand = S.expand(100, -1, -1)
+        sqrt_pi_expand = sqrt_pi.expand(100, -1)
     else:
         S_expand = S
         sqrt_pi_expand = sqrt_pi
     logP = tree.log_likelihood(S_expand, sqrt_pi_expand)
 
     logP.sum().backward()
 
-    return S.grad, sqrt_pi.grad
+    return S.grad, sqrt_pi.grad

benchmark_test/test.py

@@ -13,7 +13,7 @@
 def gen_tree(t_dtype, dim):
     torch.manual_seed(0)
     newick = read_newick('test_tree.newick')
-    L = 300
+    L = 100
     leaf_log_p = torch.randn([L, newick['num_leaf'], dim], dtype=t_dtype)
     leaf_log_p = torch.nn.functional.log_softmax(leaf_log_p, dim=2)
 
@@ -23,7 +23,7 @@ def gen_tree(t_dtype, dim):
 
 def gen_data(t_dtype, dim, seed):
     torch.manual_seed(seed)
-    L = 300
+    L = 100
     S = torch.exp(torch.randn(L,dim, dim, dtype=t_dtype))
     sqrt_pi = torch.sqrt(torch.nn.functional.softmax(torch.randn(L,dim, dtype=t_dtype), dim = 1))
     return S, sqrt_pi
@@ -108,7 +108,7 @@ def test_grads():
     helper_test("f64", 20, True)
 
 def test_grads_single_model():
+    helper_test("f64", 4, True, single_model=True)
+    helper_test("f64", 20, True, single_model=True)
     helper_test("f32", 4, True, single_model=True)
     helper_test("f32", 20, True, single_model=True)
-    helper_test("f64", 4, True, single_model=True)
-    helper_test("f64", 20, True, single_model=True)