Elite Archive (GPU)

The device elite-archive survey: an unconditional K-step walk per block, merged into the device archive by winner-take-all placement. The CPU archive is the reference for parity.

run_archive_gpu

run_archive_gpu(r1, r2, r3, alignment, model, start_trees, triangle, taxon_table, cell: float = ARCHIVE_CELL, origin: float = 0.0, extent: float = ARCHIVE_GPU_EXTENT, n_iters: int = ARCHIVE_N_ITERS, seed: int = 0, sigma_slide: float = DEFAULT_SIGMA_SLIDE, scale_width: float = DEFAULT_SCALE_WIDTH, n_walkers: int = ARCHIVE_WALK_M, walk_steps: int = ARCHIVE_WALK_K, checkpoint_every: int = ARCHIVE_CHECKPOINT_EVERY, converge: bool = True, force_full_run: bool = False, gate_window: int = ARCHIVE_GATE_WINDOW, coverage_eps: float = ARCHIVE_COVERAGE_EPS, precision_eps: float = ARCHIVE_PRECISION_EPS, logl_margin: float | None = ARCHIVE_LOGL_MARGIN, logl_floor: float | None = None, monitor=False) -> ArchiveResult

Run the GPU elite-archive survey for one alignment.

The device runs n_walkers blocks per pass. Each block walks walk_steps moves from an assigned niche and writes one candidate per move. A batched place then merges every candidate. The archive is seeded from the anchors and the start trees.

This function builds a :class:~hifuku.gpu_backend.CudaBackend and gives it to :func:~hifuku.driver.run_survey. The loop and the halt rule are in the driver, which the CPU survey also uses, so both engines halt on the trailing window of gate_window checkpoints and the CPU survey serves as the parity oracle for this one.

The returned EliteArchive and history match the CPU survey's shape.

Parameters:

Name	Type	Description	Default
r1	HifukuTree	Anchor trees that fix the chart.	required
r2	HifukuTree	Anchor trees that fix the chart.	required
r3	HifukuTree	Anchor trees that fix the chart.	required
alignment		The alignment surveyed and its substitution model.	required
model		The alignment surveyed and its substitution model.	required
start_trees	list[HifukuTree]	Extra seed trees, normally the anchor tree for the surveyed gene.	required
triangle	AnchorTriangle	Anchor geometry. It carries the chart metric.	required
taxon_table	TaxonTable	Shared namespace.	required
cell	float	Niche side and grid origin in barycentric units.	ARCHIVE_CELL
origin	float	Niche side and grid origin in barycentric units.	ARCHIVE_CELL
extent	float	Half-extent of the device niche allocation. A memory bound, not a limit on the chart domain.	ARCHIVE_GPU_EXTENT
n_iters	int	The survey budget cap in variations.	ARCHIVE_N_ITERS
seed	int	RNG seed.	0
sigma_slide	float	Proposal widths for the branch moves.	DEFAULT_SIGMA_SLIDE
scale_width	float	Proposal widths for the branch moves.	DEFAULT_SIGMA_SLIDE
n_walkers	int	Blocks per pass, and moves per block per pass.	ARCHIVE_WALK_M
walk_steps	int	Blocks per pass, and moves per block per pass.	ARCHIVE_WALK_M
converge	bool	Halt on discovery saturation when set.	True
force_full_run	bool	Run the full budget and ignore the halt.	False
gate_window	int	Trailing window in checkpoints for the halt rates.	ARCHIVE_GATE_WINDOW
coverage_eps	float	Halt thresholds on the windowed coverage and precision rates.	ARCHIVE_COVERAGE_EPS
precision_eps	float	Halt thresholds on the windowed coverage and precision rates.	ARCHIVE_COVERAGE_EPS
logl_margin	float	Relative log-likelihood gate in nats below the best tree found. The driver reads the floor once per checkpoint. Pass None for a survey that only the budget limits.	ARCHIVE_LOGL_MARGIN
logl_floor	float	Absolute log-likelihood gate. Not permitted with logl_margin.	None
monitor	(callable, None or False)	A checkpoint callback. False, the default, reports nothing. None selects the default reporter. See :mod:hifuku.monitor.	False

Returns:

Type	Description
ArchiveResult	The archive, its per-checkpoint history, and the halt state.

Source code in src/hifuku/gpu_walk.py

def run_archive_gpu(
    r1, r2, r3, alignment, model, start_trees, triangle, taxon_table,
    cell: float = ARCHIVE_CELL,
    origin: float = 0.0,
    extent: float = ARCHIVE_GPU_EXTENT,
    n_iters: int = ARCHIVE_N_ITERS,
    seed: int = 0,
    sigma_slide: float = DEFAULT_SIGMA_SLIDE,
    scale_width: float = DEFAULT_SCALE_WIDTH,
    n_walkers: int = ARCHIVE_WALK_M,
    walk_steps: int = ARCHIVE_WALK_K,
    checkpoint_every: int = ARCHIVE_CHECKPOINT_EVERY,
    converge: bool = True,
    force_full_run: bool = False,
    gate_window: int = ARCHIVE_GATE_WINDOW,
    coverage_eps: float = ARCHIVE_COVERAGE_EPS,
    precision_eps: float = ARCHIVE_PRECISION_EPS,
    logl_margin: float | None = ARCHIVE_LOGL_MARGIN,
    logl_floor: float | None = None,
    monitor=False,
) -> ArchiveResult:
    """Run the GPU elite-archive survey for one alignment.

    The device runs ``n_walkers`` blocks per pass.  Each block walks
    ``walk_steps`` moves from an assigned niche and writes one candidate per
    move.  A batched place then merges every candidate.  The archive is seeded
    from the anchors and the start trees.

    This function builds a :class:`~hifuku.gpu_backend.CudaBackend` and gives it
    to :func:`~hifuku.driver.run_survey`.  The loop and the halt rule are in the
    driver, which the CPU survey also uses, so both engines halt on the trailing
    window of ``gate_window`` checkpoints and the CPU survey serves as the parity
    oracle for this one.

    The returned ``EliteArchive`` and history match the CPU survey's shape.

    Parameters
    ----------
    r1, r2, r3 : HifukuTree
        Anchor trees that fix the chart.
    alignment, model
        The alignment surveyed and its substitution model.
    start_trees : list[HifukuTree]
        Extra seed trees, normally the anchor tree for the surveyed gene.
    triangle : AnchorTriangle
        Anchor geometry.  It carries the chart metric.
    taxon_table : TaxonTable
        Shared namespace.
    cell, origin : float
        Niche side and grid origin in barycentric units.
    extent : float
        Half-extent of the device niche allocation.  A memory bound, not a
        limit on the chart domain.
    n_iters : int
        The survey budget cap in variations.
    seed : int
        RNG seed.
    sigma_slide, scale_width : float
        Proposal widths for the branch moves.
    n_walkers, walk_steps : int
        Blocks per pass, and moves per block per pass.
    converge : bool
        Halt on discovery saturation when set.
    force_full_run : bool
        Run the full budget and ignore the halt.
    gate_window : int
        Trailing window in checkpoints for the halt rates.
    coverage_eps, precision_eps : float
        Halt thresholds on the windowed coverage and precision rates.
    logl_margin : float, optional
        Relative log-likelihood gate in nats below the best tree found.  The
        driver reads the floor once per checkpoint.  Pass None for a survey that
        only the budget limits.
    logl_floor : float, optional
        Absolute log-likelihood gate.  Not permitted with ``logl_margin``.
    monitor : callable, None or False
        A checkpoint callback.  False, the default, reports nothing.  None
        selects the default reporter.  See :mod:`hifuku.monitor`.

    Returns
    -------
    ArchiveResult
        The archive, its per-checkpoint history, and the halt state.
    """
    from hifuku.backend import RunContext
    from hifuku.driver import run_survey
    from hifuku.gpu_backend import CudaBackend
    from hifuku.params import ArchiveParams, GpuParams

    backend = CudaBackend(GpuParams(extent=extent))
    params = ArchiveParams.resolve(backend.caps.defaults, {
        "cell": cell,
        "origin": origin,
        "n_iters": n_iters,
        "seed": seed,
        "sigma_slide": sigma_slide,
        "scale_width": scale_width,
        "n_walkers": n_walkers,
        "walk_steps": walk_steps,
        "checkpoint_every": checkpoint_every,
        "converge": converge,
        "force_full_run": force_full_run,
        "gate_window": gate_window,
        "coverage_eps": coverage_eps,
        "precision_eps": precision_eps,
        "logl_margin": logl_margin,
        "logl_floor": logl_floor,
    })
    ctx = RunContext(anchors=(r1, r2, r3), alignment=alignment, model=model,
                     start_trees=list(start_trees), triangle=triangle,
                     taxon_table=taxon_table)
    return run_survey(backend, ctx, params, monitor=monitor)

evaluate_gpu

evaluate_gpu(trees, r1, r2, r3, alignment, model, triangle, taxon_table, cell: float = ARCHIVE_CELL, origin: float = 0.0, extent: float = ARCHIVE_GPU_EXTENT)

Evaluate each tree's log-likelihood and niche on the device.

Flattens the trees into archive state, then runs the seed-eval kernel, which unflattens and peels with no move. Isolates the device unflatten and Felsenstein path so it can be checked against the CPU oracle. Returns (logL, ij) where ij is an (M, 2) array of archive niche indices; a tree outside the device allocation gets the sentinel INT_MIN niche.

Source code in src/hifuku/gpu_walk.py

def evaluate_gpu(trees, r1, r2, r3, alignment, model, triangle, taxon_table,
                 cell: float = ARCHIVE_CELL, origin: float = 0.0,
                 extent: float = ARCHIVE_GPU_EXTENT):
    """Evaluate each tree's log-likelihood and niche on the device.

    Flattens the trees into archive state, then runs the seed-eval kernel, which
    unflattens and peels with no move.  Isolates the device unflatten and
    Felsenstein path so it can be checked against the CPU oracle.  Returns
    ``(logL, ij)`` where ``ij`` is an ``(M, 2)`` array of archive niche indices;
    a tree outside the device allocation gets the sentinel ``INT_MIN`` niche.
    """
    trees = list(trees)
    for t in (r1, r2, r3, *trees):
        _set_global_leaf_indices(t, taxon_table)
    M = len(trees)

    canon = _canon_leaf_names(r1, taxon_table)
    cr1, cr2, cr3 = (_relabel_to(t, canon, taxon_table) for t in (r1, r2, r3))
    ctrees = [_relabel_to(t, canon, taxon_table) for t in trees]

    ct = init_chain_trees_2d(
        ctrees, anchor_ref1=cr1, anchor_ref2=cr2, anchor_ref3=cr3,
        global_seed=0, alignment=alignment, model=model,
    )
    n_nodes = ct.n_nodes
    n_leaves = ct.n_leaves
    n_internal = ct.n_internal
    n_patterns = ct.n_patterns
    sdim = state_dim(n_nodes)
    n_off, n_side, grid_lo, sink = _device_grid(cell, origin, extent)
    tm = triangle.tree_metric
    metric_w = np.float64(tm.w)
    metric_c_ref = np.float64(tm.c_ref)
    metric_r_ref = np.float64(tm.r_ref)
    n_ref = int(ct.n_ref)

    eval_kernel = _eval_kernel_for(model.n_states)
    d_model_U = cuda.to_device(model.U.astype(np.float64))
    d_model_eigvals = cuda.to_device(model.eigvals.astype(np.float64))
    d_model_inv_sqrt_pi = cuda.to_device(model.inv_sqrt_pi.astype(np.float64))
    d_model_sqrt_pi = cuda.to_device(model.sqrt_pi.astype(np.float64))
    d_freqs = cuda.to_device(model.freqs.astype(np.float64))
    d_weights = cuda.to_device(alignment.weights.astype(np.float64))
    A_inv, Mbary_inv, rhs_off, P3_arr = _precompute_triangle_arrays(triangle)

    hp = ct.parent.copy_to_host()
    hlc = ct.left_child.copy_to_host()
    hrc = ct.right_child.copy_to_host()
    hedge = ct.edge_len.copy_to_host()
    hr1 = ct.ref_len_by_node_1.copy_to_host()
    hr2 = ct.ref_len_by_node_2.copy_to_host()
    hr3 = ct.ref_len_by_node_3.copy_to_host()
    hA = ct.A.copy_to_host()
    hB1 = ct.B1.copy_to_host()
    hB2 = ct.B2.copy_to_host()
    hB3 = ct.B3.copy_to_host()
    states = np.stack([
        flatten_host(hp[m], hlc[m], hrc[m], hedge[m], hr1[m], hr2[m], hr3[m],
                     hA[m], hB1[m], hB2[m], hB3[m], n_nodes)
        for m in range(M)
    ])
    d_state = cuda.to_device(states)
    d_assign = cuda.to_device(np.arange(M, dtype=np.int32))
    d_fit = cuda.device_array(M, dtype=np.float32)
    d_niche = cuda.device_array(M, dtype=np.int32)

    eval_kernel[M, _CUDA_BLOCK_SIZE](
        d_assign, d_state,
        ct.parent, ct.left_child, ct.right_child, ct.edge_len,
        ct.node_order, ct.node_order_pos, ct.subtree_leaf_cnt,
        ct.A, ct.B1, ct.B2, ct.B3,
        ct.ref_len_by_node_1, ct.ref_len_by_node_2, ct.ref_len_by_node_3,
        ct.d_ref_hashes_1, ct.d_ref_hashes_2, ct.d_ref_hashes_3,
        ct.leaf_global_idx,
        ct.partials, ct.P_mat, ct.scratch, ct.scratch_log_scale,
        d_model_U, d_model_eigvals, d_model_inv_sqrt_pi, d_model_sqrt_pi,
        d_freqs, d_weights,
        cuda.to_device(A_inv), cuda.to_device(Mbary_inv),
        cuda.to_device(rhs_off), cuda.to_device(P3_arr),
        np.float64(ct.C1), np.float64(ct.C2), np.float64(ct.C3),
        metric_w, metric_c_ref, metric_r_ref,
        d_fit, d_niche,
        np.int32(n_leaves), np.int32(n_nodes), np.int32(n_internal),
        np.int32(n_patterns), np.int32(n_ref),
        np.float64(grid_lo), np.float64(cell), np.int32(n_side), np.int32(sink),
    )
    cuda.synchronize()
    flat = d_niche.copy_to_host().astype(np.int64)
    ij = np.stack([flat // n_side - n_off, flat % n_side - n_off], axis=1)
    ij[flat == sink] = np.iinfo(np.int64).min
    return d_fit.copy_to_host(), ij