def log_likelihood(
tree: HifukuTree,
alignment: Alignment,
model: EigenDecomp,
) -> float:
"""
Compute total log-likelihood via the Felsenstein pruning algorithm.
Partial likelihoods are float32 with per-node per-pattern rescaling.
Log-scale accumulators and the final weighted sum are float64.
Parameters
----------
tree : HifukuTree (unrooted binary, node_order set)
alignment : Alignment - patterns uint8 (n_local, n_patterns),
weights int32 (n_patterns,)
model : EigenDecomp - DNA (4-state) or protein (20-state)
Returns
-------
float : total log-likelihood (sum over all sites)
Notes
-----
- Felsenstein (1981) : peeling recursion.
- Ayres et al. (2019) : rescaling, 4-bit tip partials.
- Fourment et al. (2025) : rescaled pruning recursion.
"""
if alignment.seq_type == "DNA":
if model.n_states != 4:
raise ValueError(
f"DNA alignment requires a 4-state model, got {model.n_states}"
)
elif alignment.seq_type == "AA":
if model.n_states != 20:
raise ValueError(
f"AA alignment requires a 20-state model, got {model.n_states}"
)
else:
raise ValueError(f"Unknown seq_type: {alignment.seq_type!r}")
n_patterns = alignment.patterns.shape[1]
# Map tree leaf index → alignment row index
aln_label_to_local: dict[str, int] = {
lbl: j for j, lbl in enumerate(alignment.index.labels)
}
for i in range(tree.n_leaves):
lbl = tree.names[i]
if lbl not in aln_label_to_local:
raise ValueError(f"Tree leaf '{lbl}' not found in alignment")
# --- Tip partials ---
if alignment.seq_type == "DNA":
# Ayres et al. (2019) : 4-bit encoding → partial vector
tip_buf = _tip_partials_dna(alignment.patterns)
else:
tip_buf = _tip_partials_aa(alignment.patterns)
# partials[k] is a (4, n_patterns) float32 array for node k
partials: list[np.ndarray] = [None] * tree.n_nodes # type: ignore[list-item]
for i in range(tree.n_leaves):
local_idx = aln_label_to_local[tree.names[i]]
partials[i] = tip_buf[local_idx].copy() # (4, n_patterns) float32
# log-scale accumulator per pattern, float64
log_scale = np.zeros(n_patterns, dtype=np.float64)
# --- Internal nodes: post-order from tree.node_order ---
# Felsenstein 1981: partial[k][s] = prod_{c ∈ children(k)} sum_x P(t_c)[s,x] * partial[c][x]
for k in tree.node_order:
lc = int(tree.left_child[k])
rc = int(tree.right_child[k])
t_l = float(tree.edge_len[lc])
t_r = float(tree.edge_len[rc])
# Transition matrices cast to float32 for partial multiplication
P_l = model.p_t(t_l).astype(np.float32) # (4, 4)
P_r = model.p_t(t_r).astype(np.float32)
contrib_l = P_l @ partials[lc] # (4, n_patterns)
contrib_r = P_r @ partials[rc]
pk = contrib_l * contrib_r # element-wise product (4, n_patterns)
# Per-pattern rescaling: divide by max state value per pattern.
# Ayres et al. (2019) : rescalePartials strategy.
scale = pk.max(axis=0) # (n_patterns,)
nonzero = scale > _RESCALE_THRESHOLD
if nonzero.any():
safe = np.where(nonzero, scale, np.float32(1.0))
pk /= safe[None, :]
log_scale += np.where(nonzero, np.log(safe.astype(np.float64)), 0.0)
partials[k] = pk
# --- Root reduction ---
root = tree.root
lc = int(tree.left_child[root])
rc = int(tree.right_child[root])
t_l = float(tree.edge_len[lc])
t_r = float(tree.edge_len[rc])
P_l = model.p_t(t_l).astype(np.float32)
P_r = model.p_t(t_r).astype(np.float32)
root_partial = (P_l @ partials[lc]) * (P_r @ partials[rc]) # (4, n_patterns)
# Felsenstein 1981: L_k = sum_s π_s * partial_root[s,k]
freqs_f32 = model.freqs.astype(np.float32)
site_L = freqs_f32 @ root_partial # (n_patterns,) float32
site_logL = np.log(site_L.astype(np.float64)) + log_scale # float64
return float(np.dot(alignment.weights.astype(np.float64), site_logL))