WIP

Author: awohns

tsdate/core.py

@@ -258,15 +258,28 @@ def get_mut_lik_fixed_node(self, edge):
 
         mutations_on_edge = self.mut_edges[edge.id]
         child_time = self.ts.node(edge.child).time
-        #assert child_time == 0
-        # Temporary hack - we should really take a more precise likelihood
-        return self._lik(
-            mutations_on_edge,
-            edge.span,
-            self.timediff,
-            self.mut_rate,
-            normalize=self.normalize,
-        )
+        if child_time == 0:
+            return self._lik(
+                mutations_on_edge,
+                edge.span,
+                self.timediff,
+                self.mut_rate,
+                normalize=self.normalize,
+            )
+        else:
+            timediff = self.timepoints - child_time + 1e-8
+            # Temporary hack - we should really take a more precise likelihood
+            likelihood = self._lik(
+                mutations_on_edge,
+                edge.span,
+                timediff,
+                self.mut_rate,
+                normalize=False,
+            )
+            likelihood[timediff < 0] = 0
+            
+            return likelihood
+
 
     def get_mut_lik_lower_tri(self, edge):
         """
@@ -677,10 +690,13 @@ def inside_pass(self, *, normalize=True, cache_inside=False, progress=None):
                     )
                     edge_lik = self.lik.get_inside(daughter_val, edge)
                 val = self.lik.combine(val, edge_lik)
+                if np.all(val == 0):
+                    raise ValueError
                 if cache_inside:
                     g_i[edge.id] = edge_lik
             norm[parent] = np.max(val) if normalize else 1
             inside[parent] = self.lik.reduce(val, norm[parent])
+            
         if cache_inside:
             self.g_i = self.lik.reduce(g_i, norm[self.ts.tables.edges.child, None])
         # Keep the results in this object
@@ -732,10 +748,10 @@ def outside_pass(
                 if ignore_oldest_root:
                     if edge.parent == self.ts.num_nodes - 1:
                         continue
-                #if edge.parent in self.fixednodes:
-                #    raise RuntimeError(
-                #        "Fixed nodes cannot currently be parents in the TS"
-                #    )
+                if edge.parent in self.fixednodes:
+                    raise RuntimeError(
+                        "Fixed nodes cannot currently be parents in the TS"
+                    )
                 # Geometric scaling works exactly for all nodes fixed in graph
                 # but is an approximation when times are unknown.
                 spanfrac = edge.span / self.spans[child]
@@ -1065,10 +1081,6 @@ def get_dates(
 
     :return: tuple(mn_post, posterior, timepoints, eps, nodes_to_date)
     """
-    # Stuff yet to be implemented. These can be deleted once fixed
-    #for sample in tree_sequence.samples():
-    #    if tree_sequence.node(sample).time != 0:
-    #        raise NotImplementedError("Samples must all be at time 0")
     fixed_nodes = set(tree_sequence.samples())
 
     # Default to not creating approximate priors unless ts has > 1000 samples

tsdate/prior.py

@@ -1031,6 +1031,8 @@ def build_grid(
         inference and a discretised time grid
     :rtype:  base.NodeGridValues Object
     """
+    #tree_sequence = tree_sequence.simplify(tree_sequence.samples())
+
     if Ne <= 0:
         raise ValueError("Parameter 'Ne' must be greater than 0")
     if approximate_priors:
@@ -1042,12 +1044,6 @@ def build_grid(
                 "Can't set approx_prior_size if approximate_prior is False"
             )
 
-    #contmpr_ts, node_map = util.reduce_to_contemporaneous(tree_sequence)
-    #if contmpr_ts.num_nodes != tree_sequence.num_nodes:
-    #    raise ValueError(
-    #        "Passed tree sequence is not simplified and/or contains "
-    #        "noncontemporaneous samples"
-    #    )
     span_data = SpansBySamples(tree_sequence, progress=progress)
 
     base_priors = ConditionalCoalescentTimes(
@@ -1079,8 +1075,6 @@ def build_grid(
         raise ValueError("time_slices must be an integer or a numpy array of floats")
 
     prior_params = base_priors.get_mixture_prior_params(span_data)
-    # Map the nodes in the prior params back to the node ids in the original ts
-    #prior_params = prior_params_contmpr[node_map, :]
     # Set all fixed nodes (i.e. samples) to have 0 variance
     priors = fill_priors(
         prior_params,