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,
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,
) -> ArchiveResult:
"""Run the GPU elite-archive survey for one alignment.
The device runs ``n_walkers`` blocks per batch, each an unconditional walk of
``walk_steps`` moves from a uniformly selected occupied niche. The archive is
seeded from the anchors and the start trees; the batch loop walks then places,
halting on the same discovery-saturation signal as ``archive.run_archive``.
The returned ``EliteArchive`` and history match the CPU survey's shape, so the
CPU archive is the reference for parity.
"""
if logl_margin is not None and logl_floor is not None:
raise ValueError("pass logl_margin or logl_floor, not both")
M = int(n_walkers)
K = int(walk_steps)
seed_trees = [r1, r2, r3, *start_trees]
for t in seed_trees:
_set_global_leaf_indices(t, taxon_table)
# Relabel every tree to one shared leaf order so the single tip-partials and
# leaf-global-index arrays of ChainTrees2D are correct for every chain.
canon = _canon_leaf_names(r1, taxon_table)
cr1, cr2, cr3 = (_relabel_to(t, canon, taxon_table) for t in (r1, r2, r3))
canon_seeds = [_relabel_to(t, canon, taxon_table) for t in seed_trees]
# Chains that host the walk. Their contents are overwritten each batch by
# the unflatten, but init allocates the Felsenstein arrays and the
# topology-invariant tip partials.
n1 = (M + 2) // 3
n2 = (M + 1) // 3
n3 = M - n1 - n2
ct = init_chain_trees_2d(
[cr1] * n1 + [cr2] * n2 + [cr3] * n3,
anchor_ref1=cr1, anchor_ref2=cr2, anchor_ref3=cr3,
global_seed=seed, alignment=alignment, model=model,
)
n_nodes = ct.n_nodes
n_leaves = ct.n_leaves
n_internal = ct.n_internal
n_patterns = ct.n_patterns
n_ref = ct.n_ref
sdim = state_dim(n_nodes)
n_off, n_side, grid_lo, sink = _device_grid(cell, origin, extent)
n_niche = n_side * n_side
walk_kernel = _walk_kernel_for(model.n_states)
# Device model and anchor-geometry arrays.
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)
d_A_inv = cuda.to_device(A_inv)
d_Mbary_inv = cuda.to_device(Mbary_inv)
d_rhs_off = cuda.to_device(rhs_off)
d_P3 = cuda.to_device(P3_arr)
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)
archive_dev = new_device_archive(n_niche + 1, sdim)
# Seed the archive from the anchors and start trees (host descriptor and logL,
# to match the CPU seeding exactly).
seed_ct = init_chain_trees_2d(
canon_seeds, anchor_ref1=cr1, anchor_ref2=cr2, anchor_ref3=cr3,
global_seed=seed, alignment=alignment, model=model,
)
h = {
"parent": seed_ct.parent.copy_to_host(),
"lc": seed_ct.left_child.copy_to_host(),
"rc": seed_ct.right_child.copy_to_host(),
"edge": seed_ct.edge_len.copy_to_host(),
"r1": seed_ct.ref_len_by_node_1.copy_to_host(),
"r2": seed_ct.ref_len_by_node_2.copy_to_host(),
"r3": seed_ct.ref_len_by_node_3.copy_to_host(),
"A": seed_ct.A.copy_to_host(),
"B1": seed_ct.B1.copy_to_host(),
"B2": seed_ct.B2.copy_to_host(),
"B3": seed_ct.B3.copy_to_host(),
}
tmp = EliteArchive(cell=cell, origin=origin)
seed_cn, seed_cf, seed_cs = [], [], []
for mi, tree in enumerate(canon_seeds):
lam1, lam2, _z = locate_tree(tree, taxon_table, triangle)
i, j = tmp.niche_of(lam1, lam2)
ci, cj = i + n_off, j + n_off
if not (0 <= ci < n_side and 0 <= cj < n_side):
continue # seed outside the device allocation (extreme; skip)
seed_cn.append(ci * n_side + cj)
seed_cf.append(float(log_likelihood(tree, alignment, model)))
seed_cs.append(flatten_host(
h["parent"][mi], h["lc"][mi], h["rc"][mi], h["edge"][mi],
h["r1"][mi], h["r2"][mi], h["r3"][mi],
h["A"][mi], h["B1"][mi], h["B2"][mi], h["B3"][mi], n_nodes,
))
if not seed_cn:
raise ValueError("no seed tree landed on the device allocation")
place_candidates(
archive_dev,
cuda.to_device(np.asarray(seed_cn, dtype=np.int32)),
cuda.to_device(np.asarray(seed_cf, dtype=np.float32)),
cuda.to_device(np.asarray(seed_cs, dtype=np.float32)),
)
# Candidate buffers, reused across batches.
n_cand = M * K
d_out_niche = cuda.device_array(n_cand, dtype=np.int32)
d_out_fit = cuda.device_array(n_cand, dtype=np.float32)
d_out_z = cuda.device_array(n_cand, dtype=np.float64)
d_out_state = cuda.device_array((n_cand, sdim), dtype=np.float32)
rng = np.random.default_rng(seed)
n_per_batch = M * K
n_batches = max(1, n_iters // n_per_batch)
fit0 = archive_dev["fitness"].copy_to_host()[:n_niche]
filled0 = fit0 > -1.0e29
history = [(0, int(filled0.sum()),
float(fit0[filled0].max()) if filled0.any() else float("-inf"),
float("nan"), float("nan"))]
converged = False
halt_iter = -1
total = 0
best_run = float(fit0[filled0].max()) if filled0.any() else -1.0e30
# Host-side survey state, refreshed at each checkpoint from device scalars and
# the newly appended keys, with no per-interval fitness copy.
_n_keys = int(archive_dev["n_keys"].copy_to_host()[0])
host_keys = archive_dev["keys"][:_n_keys].copy_to_host()
host_keys = host_keys[host_keys != sink]
n_keys_copied = _n_keys
n_filled_prev = _n_keys
gain_prev = float(archive_dev["stat_gain"].copy_to_host()[0])
consec_below = 0
def _floor_from(best):
if logl_floor is not None:
return float(logl_floor)
if logl_margin is not None:
return best - float(logl_margin)
return -1.0e30
batches_done = 0
while batches_done < n_batches:
if host_keys.size == 0:
break
# Launch a checkpoint interval of walk-and-place pairs back to back on the
# default stream, with the assigns for the whole interval uploaded once
# and the floor held fixed. No host sync inside the interval, so the GPU
# runs the pairs without waiting on the host.
chunk = min(ARCHIVE_STREAM_CHECK, n_batches - batches_done)
assigns = host_keys[rng.integers(0, host_keys.size, size=(chunk, M))].astype(np.int32)
d_assigns = cuda.to_device(assigns)
floor = _floor_from(best_run)
for bi in range(chunk):
walk_kernel[M, _CUDA_BLOCK_SIZE](
d_assigns[bi], archive_dev["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.rng_state,
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_lengths_1,
ct.d_ref_hashes_2, ct.d_ref_lengths_2,
ct.d_ref_hashes_3, ct.d_ref_lengths_3,
ct.partials, ct.P_mat, ct.scratch, ct.scratch_log_scale,
ct.leaf_global_idx,
d_model_U, d_model_eigvals, d_model_inv_sqrt_pi, d_model_sqrt_pi,
d_freqs, d_weights,
d_A_inv, d_Mbary_inv, d_rhs_off, d_P3,
np.float64(ct.C1), np.float64(ct.C2), np.float64(ct.C3),
metric_w, metric_c_ref, metric_r_ref,
d_out_niche, d_out_fit, d_out_z, d_out_state,
np.int32(K), np.int32(n_leaves), np.int32(n_nodes),
np.int32(n_internal), np.int32(n_patterns), np.int32(n_ref),
np.float64(sigma_slide), np.float64(scale_width),
np.float64(grid_lo), np.float64(cell), np.int32(n_side),
np.int32(sink), np.float32(EMPTY_FITNESS), np.float64(floor),
)
place_candidates(archive_dev, d_out_niche, d_out_fit, d_out_state)
batches_done += chunk
total = batches_done * n_per_batch
# Checkpoint: one sync, then the interval stats read as device scalars
# (filled count, accumulated gain, running best), normalized per batch so
# the halt thresholds match the CPU survey. The elite range is the spread
# of the live elites, best minus the smallest occupied fitness, matching
# the CPU ``max(vals) - min(vals)``. A rising relative floor can leave
# early elites below it, so the smallest elite is reduced on the device
# rather than proxied by ``best - floor``.
cuda.synchronize()
n_keys_now = int(archive_dev["n_keys"].copy_to_host()[0])
gain_now = float(archive_dev["stat_gain"].copy_to_host()[0])
best_now = float(archive_dev["stat_best"].copy_to_host()[0])
n_filled = n_keys_now
n_new = n_filled - n_filled_prev
gain = gain_now - gain_prev
best_run = max(best_run, best_now)
elite_range = (best_run - min_occupied_fitness(archive_dev)) or 1.0
denom = chunk * max(n_filled, 1)
cov = n_new / denom
prec = gain / (denom * elite_range)
history.append((total, n_filled, best_run, cov, prec))
# Refresh the selectable keys by copying only the newly appended entries.
if n_keys_now > n_keys_copied:
new_keys = archive_dev["keys"][n_keys_copied:n_keys_now].copy_to_host()
new_keys = new_keys[new_keys != sink]
host_keys = (np.concatenate([host_keys, new_keys])
if host_keys.size else new_keys)
n_keys_copied = n_keys_now
gain_prev = gain_now
n_filled_prev = n_filled
consec_below = (consec_below + 1
if _survey_converged(cov, prec, coverage_eps, precision_eps)
else 0)
if converge and not force_full_run and consec_below >= ARCHIVE_STREAM_HALT:
converged = True
halt_iter = total
break
# Reconstruct the host EliteArchive from the device fitness and state.
final_fit = archive_dev["fitness"].copy_to_host()[:n_niche]
final_state = archive_dev["state"].copy_to_host()[:n_niche]
leaf_names = list(canon)
archive = EliteArchive(cell=cell, origin=origin, anchors=triangle,
taxon_table=taxon_table, metric=triangle.tree_metric)
at_edge = False
for idx in np.nonzero(final_fit > -1.0e29)[0]:
ci, cj = int(idx) // n_side, int(idx) % n_side
if ci == 0 or ci == n_side - 1 or cj == 0 or cj == n_side - 1:
at_edge = True
tree = _tree_from_state(
final_state[idx], leaf_names, n_nodes, n_leaves, taxon_table)
z = _z_from_state(final_state[idx], n_nodes, ct.C1, ct.C2, ct.C3, triangle)
archive.place((ci - n_off, cj - n_off),
Elite(logL=float(final_fit[idx]), tree=tree, z=z))
if at_edge:
import warnings
warnings.warn(
"elite archive reached the device allocation edge; the log-likelihood "
"contour exceeds `extent` and was clipped. Increase `extent`.",
RuntimeWarning, stacklevel=2,
)
return ArchiveResult(
archive=archive,
history=np.array(history, dtype=np.float64),
converged=converged,
halt_iter=halt_iter,
)