Make CoalescentModelInterface expect a TimeTreeModelInterface

4ment · Mar 10, 2024 · f34a571 · f34a571
1 parent 2216fc0
commit f34a571
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/benchmarks/benchmark.py b/benchmarks/benchmark.py
@@ -204,18 +204,18 @@ def ratio_transform_jacobian(args):
     tree_model()
 
     @benchmark
-    def fn(internal_heights):
+    def fn():
         return torch.tensor([tree_model.inst.transform_jacobian()])
 
     @benchmark
-    def fn_grad(internal_heights):
+    def fn_grad():
         return torch.tensor(tree_model.inst.gradient_transform_jacobian())
 
-    total_time, log_det_jac = fn(args.replicates, internal_heights)
+    total_time, log_det_jac = fn(args.replicates)
     print(f"  {args.replicates} evaluations: {total_time} ({log_det_jac})")
 
     internal_heights.requires_grad = True
-    grad_total_time, grad_log_det_jac = fn_grad(args.replicates, internal_heights)
+    grad_total_time, grad_log_det_jac = fn_grad(args.replicates)
     print(f"  {args.replicates} gradient evaluations: {grad_total_time}")
 
     if args.output:

diff --git a/torchtree_physher/_version.py b/torchtree_physher/_version.py
@@ -2,4 +2,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "1.0.0-dev3"
+__version__ = "1.0.0-dev4"
diff --git a/torchtree_physher/physher.cpp b/torchtree_physher/physher.cpp
@@ -229,20 +229,20 @@ PYBIND11_MODULE(physher, m) {
 
   py::class_<ConstantCoalescentModelInterface, CoalescentModelInterface>(
       m, "ConstantCoalescentModel")
-      .def(py::init<double, TreeModelInterface *>());
+      .def(py::init<double, TimeTreeModelInterface *>());
 
   py::class_<PiecewiseConstantCoalescentInterface, CoalescentModelInterface>(
       m, "PiecewiseConstantCoalescentModel")
-      .def(py::init<const std::vector<double>, TreeModelInterface *>());
+      .def(py::init<const std::vector<double>, TimeTreeModelInterface *>());
 
   py::class_<PiecewiseConstantCoalescentGridInterface,
              CoalescentModelInterface>(m,
                                        "PiecewiseConstantCoalescentGridModel")
-      .def(py::init<const std::vector<double>, TreeModelInterface *, double>());
+      .def(py::init<const std::vector<double>, TimeTreeModelInterface *, double>());
 
   py::class_<PiecewiseLinearCoalescentGridInterface, CoalescentModelInterface>(
       m, "PiecewiseLinearCoalescentGridModel")
-      .def(py::init<const std::vector<double>, TreeModelInterface *, double>());
+      .def(py::init<const std::vector<double>, TimeTreeModelInterface *, double>());
 
   py::class_<CTMCScaleModelInterface, CallableModelInterface>(m,
                                                               "CTMCScaleModel")