From 974038dd54e08b037d71167c563f7e2d32182298 Mon Sep 17 00:00:00 2001 From: Yan Wong Date: Mon, 22 Aug 2022 16:37:33 +0100 Subject: [PATCH] Add unity normalization constant to unfixed leaf nodes --- tsdate/core.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tsdate/core.py b/tsdate/core.py index 72b571b2..f895058b 100644 --- a/tsdate/core.py +++ b/tsdate/core.py @@ -649,9 +649,7 @@ def inside_pass(self, *, normalize=True, cache_inside=False, progress=None): ) # It is possible that a simple node is non-fixed, in which case we want to # provide an inside array that reflects the prior distribution - nonfixed_samples = np.intersect1d( - self.priors.nonfixed_node_ids(), self.ts.samples() - ) + nonfixed_samples = np.intersect1d(inside.nonfixed_node_ids(), self.ts.samples()) for u in nonfixed_samples: # this is in the same probability space as the prior, so we should be # OK just to copy the prior values straight in (but we should check they @@ -663,6 +661,8 @@ def inside_pass(self, *, normalize=True, cache_inside=False, progress=None): (self.ts.num_edges, self.lik.grid_size), self.lik.identity_constant ) norm = np.full(self.ts.num_nodes, np.nan) + to_visit = np.zeros(self.ts.num_nodes, dtype=bool) + to_visit[inside.nonfixed_node_ids()] = True # Iterate through the nodes via groupby on parent node for parent, edges in tqdm( self.edges_by_parent_asc(), @@ -707,6 +707,12 @@ def inside_pass(self, *, normalize=True, cache_inside=False, progress=None): g_i[edge.id] = edge_lik norm[parent] = np.max(val) if normalize else 1 inside[parent] = self.lik.reduce(val, norm[parent]) + to_visit[parent] = False + + # There may be nodes that are not parents but are also not fixed (e.g. + # undated sample nodes). These need an identity normalization constant + for unfixed_unvisited in np.where(to_visit)[0]: + norm[unfixed_unvisited] = 1 if cache_inside: self.g_i = self.lik.reduce(g_i, norm[self.ts.tables.edges.child, None])