Likelihood (GPU)

numba-cuda implementation of Felsenstein pruning, bandwidth-optimized for NVIDIA Blackwell (sm_12x).

log_likelihood_gpu

log_likelihood_gpu(tree: HifukuTree, alignment: Alignment, model: EigenDecomp) -> float

Compute total log-likelihood via GPU Felsenstein pruning.

Dispatches to a 4-state kernel for DNA or a 20-state kernel for AA. Stateless: allocates device arrays per call.

Felsenstein (1981). Ayres et al. (2019). Ref: quarimo.

Source code in src/hifuku/gpu_pruning.py

def log_likelihood_gpu(
    tree: HifukuTree,
    alignment: Alignment,
    model: EigenDecomp,
) -> float:
    """
    Compute total log-likelihood via GPU Felsenstein pruning.

    Dispatches to a 4-state kernel for DNA or a 20-state kernel for AA.
    Stateless: allocates device arrays per call.

    Felsenstein (1981). Ayres et al. (2019). Ref: quarimo.
    """
    if alignment.seq_type == "DNA":
        return _log_likelihood_gpu_dna(tree, alignment, model)
    if alignment.seq_type == "AA":
        return _log_likelihood_gpu_aa(tree, alignment, model)
    raise NotImplementedError(
        f"GPU pruning not implemented for seq_type {alignment.seq_type!r}"
    )