Skip to content

Backend Protocol

A backend starts the kernels of one device and reads its archive. The driver is written once against this protocol, so a backend that meets it runs correctly.

Two rules make the protocol work. submit_batch returns without waiting, which lets the driver queue the next chunk before it calls the monitor. A backend draws its own parent niches, because the host engine draws a fresh parent at every step and the device engine draws one per walker per pass.

The pattern follows BEAGLE, which solves the same problem in the same field. Ayres et al. (2019).

Backend

Bases: Protocol

The device interface.

A handle is opaque to the driver. A backend returns one from prepare and gets it back on every later call.

Source code in src/hifuku/backend.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
@runtime_checkable
class Backend(Protocol):
    """The device interface.

    A handle is opaque to the driver.  A backend returns one from ``prepare``
    and gets it back on every later call.
    """

    caps: Capabilities

    def prepare(self, ctx: RunContext, params) -> Any:
        """Allocate the run and return its handle.  This does no survey work."""

    def seed(self, handle) -> None:
        """Place the anchor trees and the start trees in the archive.

        A seed tree is placed without the log-likelihood gate.  Raise
        ``ValueError`` when no seed tree lands on the chart grid.
        """

    def submit_batch(self, handle, keys, n_variations: int,
                     floor: float) -> None:
        """Queue ``n_variations`` variations.  Do not synchronize.

        Draw the parents from ``keys``, which holds the niche keys the driver
        knows about.  Apply ``floor`` as the log-likelihood gate: a variation
        below the floor does not make a niche.  Place every candidate before
        this call returns control to the driver, or queue the placement on the
        same stream.
        """

    def sync(self) -> None:
        """Wait for the queued work to finish."""

    def read_stats(self, handle) -> Stats:
        """Read the cumulative survey statistics.  Call after ``sync``."""

    def read_keys(self, handle, since: int) -> np.ndarray:
        """Return the niche keys added at or after index ``since``.

        The driver holds the full list and asks only for the new entries, so a
        device backend copies a small slice and not the whole array.  The key
        values are opaque to the driver, which gives them back to
        ``submit_batch``.
        """

    def read_map(self, handle) -> np.ndarray:
        """Return the dense log-likelihood grid, ``NaN`` where unoccupied.

        This is one bounded transfer.  The array is a copy, never a live buffer.
        """

    def read_trees(self, handle) -> dict:
        """Return ``{(i, j): HifukuTree}`` for every occupied niche.

        The trees are copies.  The CPU archive changes during a run, so a
        reference to a live tree would show a later state than the checkpoint.
        """

    def finalize(self, handle) -> Any:
        """Return the completed ``EliteArchive``."""

prepare 

prepare(ctx: RunContext, params) -> Any

Allocate the run and return its handle. This does no survey work.

Source code in src/hifuku/backend.py 
128
129
def prepare(self, ctx: RunContext, params) -> Any:
    """Allocate the run and return its handle.  This does no survey work."""

seed 

seed(handle) -> None

Place the anchor trees and the start trees in the archive.

A seed tree is placed without the log-likelihood gate. Raise ValueError when no seed tree lands on the chart grid.

Source code in src/hifuku/backend.py 
131
132
133
134
135
136
def seed(self, handle) -> None:
    """Place the anchor trees and the start trees in the archive.

    A seed tree is placed without the log-likelihood gate.  Raise
    ``ValueError`` when no seed tree lands on the chart grid.
    """

submit_batch 

submit_batch(handle, keys, n_variations: int, floor: float) -> None

Queue n_variations variations. Do not synchronize.

Draw the parents from keys, which holds the niche keys the driver knows about. Apply floor as the log-likelihood gate: a variation below the floor does not make a niche. Place every candidate before this call returns control to the driver, or queue the placement on the same stream.

Source code in src/hifuku/backend.py 
138
139
140
141
142
143
144
145
146
147
def submit_batch(self, handle, keys, n_variations: int,
                 floor: float) -> None:
    """Queue ``n_variations`` variations.  Do not synchronize.

    Draw the parents from ``keys``, which holds the niche keys the driver
    knows about.  Apply ``floor`` as the log-likelihood gate: a variation
    below the floor does not make a niche.  Place every candidate before
    this call returns control to the driver, or queue the placement on the
    same stream.
    """

sync 

sync() -> None

Wait for the queued work to finish.

Source code in src/hifuku/backend.py 
149
150
def sync(self) -> None:
    """Wait for the queued work to finish."""

read_stats 

read_stats(handle) -> Stats

Read the cumulative survey statistics. Call after sync.

Source code in src/hifuku/backend.py 
152
153
def read_stats(self, handle) -> Stats:
    """Read the cumulative survey statistics.  Call after ``sync``."""

read_keys 

read_keys(handle, since: int) -> np.ndarray

Return the niche keys added at or after index since.

The driver holds the full list and asks only for the new entries, so a device backend copies a small slice and not the whole array. The key values are opaque to the driver, which gives them back to submit_batch.

Source code in src/hifuku/backend.py 
155
156
157
158
159
160
161
162
def read_keys(self, handle, since: int) -> np.ndarray:
    """Return the niche keys added at or after index ``since``.

    The driver holds the full list and asks only for the new entries, so a
    device backend copies a small slice and not the whole array.  The key
    values are opaque to the driver, which gives them back to
    ``submit_batch``.
    """

read_map 

read_map(handle) -> np.ndarray

Return the dense log-likelihood grid, NaN where unoccupied.

This is one bounded transfer. The array is a copy, never a live buffer.

Source code in src/hifuku/backend.py 
164
165
166
167
168
def read_map(self, handle) -> np.ndarray:
    """Return the dense log-likelihood grid, ``NaN`` where unoccupied.

    This is one bounded transfer.  The array is a copy, never a live buffer.
    """

read_trees 

read_trees(handle) -> dict

Return {(i, j): HifukuTree} for every occupied niche.

The trees are copies. The CPU archive changes during a run, so a reference to a live tree would show a later state than the checkpoint.

Source code in src/hifuku/backend.py 
170
171
172
173
174
175
def read_trees(self, handle) -> dict:
    """Return ``{(i, j): HifukuTree}`` for every occupied niche.

    The trees are copies.  The CPU archive changes during a run, so a
    reference to a live tree would show a later state than the checkpoint.
    """

finalize 

finalize(handle) -> Any

Return the completed EliteArchive.

Source code in src/hifuku/backend.py 
177
178
def finalize(self, handle) -> Any:
    """Return the completed ``EliteArchive``."""

Capabilities dataclass

What a backend is and what it can do.

Attributes:

Name Type Description
name str

Backend name, for example "numba-cuda" or "cpu".
unified_memory bool

True when a state read is pointer arithmetic. On a discrete device this is False and the driver gives one warning for a trees monitor. The driver never refuses the request.
defaults Mapping[str, Any]

ArchiveParams field overrides that suit this device. The mapping holds only the fields the backend wants to change, because a full parameter object cannot show which values the backend meant to set. This is the first layer of ArchiveParams.resolve.
can_read_trees bool

True when read_trees works. A backend that cannot rebuild trees sets this False, and a trees monitor then raises at setup.
Source code in src/hifuku/backend.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@dataclass(frozen=True)
class Capabilities:
    """What a backend is and what it can do.

    Attributes
    ----------
    name : str
        Backend name, for example ``"numba-cuda"`` or ``"cpu"``.
    unified_memory : bool
        True when a state read is pointer arithmetic.  On a discrete device this
        is False and the driver gives one warning for a ``trees`` monitor.  The
        driver never refuses the request.
    defaults : Mapping[str, Any]
        ``ArchiveParams`` field overrides that suit this device.  The mapping
        holds only the fields the backend wants to change, because a full
        parameter object cannot show which values the backend meant to set.
        This is the first layer of ``ArchiveParams.resolve``.
    can_read_trees : bool
        True when ``read_trees`` works.  A backend that cannot rebuild trees
        sets this False, and a ``trees`` monitor then raises at setup.
    """

    name: str
    unified_memory: bool
    defaults: Mapping[str, Any] = field(default_factory=dict)
    can_read_trees: bool = True

Stats dataclass

Survey statistics read at one checkpoint.

Every value is cumulative from the start of the run. The driver takes the difference between two checkpoints to get the coverage and precision rates, which is why a backend must not reset these counters.

Attributes:

Name Type Description
n_keys int

Occupied niches.
gain float

Total elite gain, summed over every replacement of an elite.
best float

Best log-likelihood found.
elite_range float

best minus the smallest occupied log-likelihood. This normalizes the precision rate. Only the backend can measure it cheaply, because on a device it is a reduction over the fitness array.
Source code in src/hifuku/backend.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@dataclass(frozen=True)
class Stats:
    """Survey statistics read at one checkpoint.

    Every value is cumulative from the start of the run.  The driver takes the
    difference between two checkpoints to get the coverage and precision rates,
    which is why a backend must not reset these counters.

    Attributes
    ----------
    n_keys : int
        Occupied niches.
    gain : float
        Total elite gain, summed over every replacement of an elite.
    best : float
        Best log-likelihood found.
    elite_range : float
        ``best`` minus the smallest occupied log-likelihood.  This normalizes
        the precision rate.  Only the backend can measure it cheaply, because on
        a device it is a reduction over the fitness array.
    """

    n_keys: int
    gain: float
    best: float
    elite_range: float

RunContext dataclass

The live objects that one survey runs against.

A backend gets this once, in prepare. The context holds trees, an alignment and a model, so it does not travel between hosts. A SurveyTask holds labels and indices instead, and a worker rebuilds the context from the plan.

Attributes:

Name Type Description
anchors tuple

The three anchor trees that fix the chart.
alignment, model

The alignment under survey and its substitution model.
start_trees list

Extra seed trees, normally the NJ tree of the surveyed gene.
triangle AnchorTriangle

Anchor geometry. It carries the chart metric.
taxon_table TaxonTable

The shared namespace.
Source code in src/hifuku/backend.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
@dataclass(frozen=True)
class RunContext:
    """The live objects that one survey runs against.

    A backend gets this once, in ``prepare``.  The context holds trees, an
    alignment and a model, so it does not travel between hosts.  A
    ``SurveyTask`` holds labels and indices instead, and a worker rebuilds the
    context from the plan.

    Attributes
    ----------
    anchors : tuple
        The three anchor trees that fix the chart.
    alignment, model
        The alignment under survey and its substitution model.
    start_trees : list
        Extra seed trees, normally the NJ tree of the surveyed gene.
    triangle : AnchorTriangle
        Anchor geometry.  It carries the chart metric.
    taxon_table : TaxonTable
        The shared namespace.
    """

    anchors: tuple
    alignment: Any
    model: Any
    start_trees: list
    triangle: Any
    taxon_table: Any

CpuBackend

Run the survey on the host.

A state read is a dictionary lookup, so unified_memory is True and a trees monitor costs nothing beyond building the mapping.

Source code in src/hifuku/cpu_backend.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
class CpuBackend:
    """Run the survey on the host.

    A state read is a dictionary lookup, so ``unified_memory`` is True and a
    ``trees`` monitor costs nothing beyond building the mapping.
    """

    def __init__(self):
        self.caps = Capabilities(name="cpu", unified_memory=True,
                                 defaults=dict(_CPU_DEFAULTS),
                                 can_read_trees=True)

    def prepare(self, ctx: RunContext, params) -> CpuHandle:
        """Allocate the archive and the generator.  This does no survey work."""
        archive = EliteArchive(cell=params.cell, origin=params.origin,
                               anchors=ctx.triangle,
                               taxon_table=ctx.taxon_table,
                               metric=ctx.triangle.tree_metric)
        return CpuHandle(ctx=ctx, params=params, archive=archive,
                         rng=np.random.default_rng(params.seed))

    def seed(self, handle: CpuHandle) -> None:
        """Place the anchor trees and the start trees without the gate."""
        ctx = handle.ctx
        for tree in (*ctx.anchors, *ctx.start_trees):
            key, elite = self._evaluate(handle, tree)
            handle.archive.place(key, elite)
        if handle.archive.n_filled == 0:
            raise ValueError("no seed tree landed on the chart grid")

    def submit_batch(self, handle: CpuHandle, keys, n_variations: int,
                     floor: float) -> None:
        """Run ``n_variations`` variations as whole passes.

        A host archive is always current, so this backend draws its parents
        from its own key list rather than from the ``keys`` the driver holds.  A
        parent therefore comes from the niches that exist at the moment of the
        draw, which is the sampling of the sequential survey.
        """
        params = handle.params
        n_walkers, walk_steps = params.n_walkers, params.walk_steps
        passes = max(1, n_variations // (n_walkers * walk_steps))
        for _ in range(passes):
            self._one_pass(handle, n_walkers, walk_steps, floor)

    def _one_pass(self, handle, n_walkers, walk_steps, floor) -> None:
        """Draw the parents, walk them, then place every candidate together."""
        ctx, params, archive = handle.ctx, handle.params, handle.archive
        candidates = []
        for _ in range(n_walkers):
            if not archive.keys:
                break
            parent = archive.keys[int(handle.rng.integers(len(archive.keys)))]
            tree = archive.elites[parent].tree
            for _ in range(walk_steps):
                child = _vary(tree, ctx.taxon_table, handle.rng,
                              params.sigma_slide, params.scale_width)
                if child is None:
                    continue
                tree = child
                key, elite = self._evaluate(handle, child)
                if elite.logL >= floor:
                    candidates.append((key, elite))
        if handle.on_candidates is not None:
            handle.on_candidates(candidates)
        _, gain = archive.place_batch(candidates)
        handle.gain += gain

    def _evaluate(self, handle, tree):
        """Return the ``(key, Elite)`` of one tree on this chart."""
        ctx = handle.ctx
        lam1, lam2, z = locate_tree(tree, ctx.taxon_table, ctx.triangle)
        logL = float(log_likelihood(tree, ctx.alignment, ctx.model))
        return handle.archive.niche_of(lam1, lam2), Elite(logL=logL, tree=tree,
                                                          z=z)

    def sync(self) -> None:
        """Return at once.  Host work is finished when it returns."""

    def read_stats(self, handle: CpuHandle) -> Stats:
        """Read the cumulative survey statistics.

        Non-finite elites are skipped, as they are in ``best_logL`` and in the
        device reduction.  ``max`` keeps the first item a comparison cannot
        beat, so a NaN among the values would set the best and the range by its
        position in the dictionary rather than by its magnitude.
        """
        vals = [e.logL for e in handle.archive.elites.values()
                if math.isfinite(e.logL)]
        if not vals:
            return Stats(n_keys=len(handle.archive.keys), gain=handle.gain,
                         best=float("-inf"), elite_range=0.0)
        return Stats(n_keys=len(handle.archive.keys), gain=handle.gain,
                     best=max(vals), elite_range=max(vals) - min(vals))

    def read_keys(self, handle: CpuHandle, since: int) -> np.ndarray:
        """Return the positions of the keys added at or after ``since``.

        The values index ``archive.keys``, which only grows.  They are opaque to
        the driver.
        """
        return np.arange(since, len(handle.archive.keys), dtype=np.int64)

    def read_map(self, handle: CpuHandle) -> np.ndarray:
        """Return the elite log-likelihood grid, ``NaN`` where unoccupied."""
        return np.ma.filled(handle.archive.elevation_grid().astype(np.float64),
                            np.nan)

    def read_trees(self, handle: CpuHandle) -> dict:
        """Return ``{(i, j): HifukuTree}`` for every occupied niche.

        The mapping is a new dictionary.  A move returns a new tree and never
        changes its parent, so the tree objects of a snapshot stay as they were
        at the checkpoint.
        """
        return {key: handle.archive.tree_at(key) for key in handle.archive.keys}

    def finalize(self, handle: CpuHandle) -> EliteArchive:
        """Return the completed archive."""
        return handle.archive

prepare 

prepare(ctx: RunContext, params) -> CpuHandle

Allocate the archive and the generator. This does no survey work.

Source code in src/hifuku/cpu_backend.py 
86
87
88
89
90
91
92
93
def prepare(self, ctx: RunContext, params) -> CpuHandle:
    """Allocate the archive and the generator.  This does no survey work."""
    archive = EliteArchive(cell=params.cell, origin=params.origin,
                           anchors=ctx.triangle,
                           taxon_table=ctx.taxon_table,
                           metric=ctx.triangle.tree_metric)
    return CpuHandle(ctx=ctx, params=params, archive=archive,
                     rng=np.random.default_rng(params.seed))

seed 

seed(handle: CpuHandle) -> None

Place the anchor trees and the start trees without the gate.

Source code in src/hifuku/cpu_backend.py 
 95
 96
 97
 98
 99
100
101
102
def seed(self, handle: CpuHandle) -> None:
    """Place the anchor trees and the start trees without the gate."""
    ctx = handle.ctx
    for tree in (*ctx.anchors, *ctx.start_trees):
        key, elite = self._evaluate(handle, tree)
        handle.archive.place(key, elite)
    if handle.archive.n_filled == 0:
        raise ValueError("no seed tree landed on the chart grid")

submit_batch 

submit_batch(handle: CpuHandle, keys, n_variations: int, floor: float) -> None

Run n_variations variations as whole passes.

A host archive is always current, so this backend draws its parents from its own key list rather than from the keys the driver holds. A parent therefore comes from the niches that exist at the moment of the draw, which is the sampling of the sequential survey.

Source code in src/hifuku/cpu_backend.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def submit_batch(self, handle: CpuHandle, keys, n_variations: int,
                 floor: float) -> None:
    """Run ``n_variations`` variations as whole passes.

    A host archive is always current, so this backend draws its parents
    from its own key list rather than from the ``keys`` the driver holds.  A
    parent therefore comes from the niches that exist at the moment of the
    draw, which is the sampling of the sequential survey.
    """
    params = handle.params
    n_walkers, walk_steps = params.n_walkers, params.walk_steps
    passes = max(1, n_variations // (n_walkers * walk_steps))
    for _ in range(passes):
        self._one_pass(handle, n_walkers, walk_steps, floor)

sync 

sync() -> None

Return at once. Host work is finished when it returns.

Source code in src/hifuku/cpu_backend.py 
150
151
def sync(self) -> None:
    """Return at once.  Host work is finished when it returns."""

read_stats 

read_stats(handle: CpuHandle) -> Stats

Read the cumulative survey statistics.

Non-finite elites are skipped, as they are in best_logL and in the device reduction. max keeps the first item a comparison cannot beat, so a NaN among the values would set the best and the range by its position in the dictionary rather than by its magnitude.

Source code in src/hifuku/cpu_backend.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def read_stats(self, handle: CpuHandle) -> Stats:
    """Read the cumulative survey statistics.

    Non-finite elites are skipped, as they are in ``best_logL`` and in the
    device reduction.  ``max`` keeps the first item a comparison cannot
    beat, so a NaN among the values would set the best and the range by its
    position in the dictionary rather than by its magnitude.
    """
    vals = [e.logL for e in handle.archive.elites.values()
            if math.isfinite(e.logL)]
    if not vals:
        return Stats(n_keys=len(handle.archive.keys), gain=handle.gain,
                     best=float("-inf"), elite_range=0.0)
    return Stats(n_keys=len(handle.archive.keys), gain=handle.gain,
                 best=max(vals), elite_range=max(vals) - min(vals))

read_keys 

read_keys(handle: CpuHandle, since: int) -> np.ndarray

Return the positions of the keys added at or after since.

The values index archive.keys, which only grows. They are opaque to the driver.

Source code in src/hifuku/cpu_backend.py 
169
170
171
172
173
174
175
def read_keys(self, handle: CpuHandle, since: int) -> np.ndarray:
    """Return the positions of the keys added at or after ``since``.

    The values index ``archive.keys``, which only grows.  They are opaque to
    the driver.
    """
    return np.arange(since, len(handle.archive.keys), dtype=np.int64)

read_map 

read_map(handle: CpuHandle) -> np.ndarray

Return the elite log-likelihood grid, NaN where unoccupied.

Source code in src/hifuku/cpu_backend.py 
177
178
179
180
def read_map(self, handle: CpuHandle) -> np.ndarray:
    """Return the elite log-likelihood grid, ``NaN`` where unoccupied."""
    return np.ma.filled(handle.archive.elevation_grid().astype(np.float64),
                        np.nan)

read_trees 

read_trees(handle: CpuHandle) -> dict

Return {(i, j): HifukuTree} for every occupied niche.

The mapping is a new dictionary. A move returns a new tree and never changes its parent, so the tree objects of a snapshot stay as they were at the checkpoint.

Source code in src/hifuku/cpu_backend.py 
182
183
184
185
186
187
188
189
def read_trees(self, handle: CpuHandle) -> dict:
    """Return ``{(i, j): HifukuTree}`` for every occupied niche.

    The mapping is a new dictionary.  A move returns a new tree and never
    changes its parent, so the tree objects of a snapshot stay as they were
    at the checkpoint.
    """
    return {key: handle.archive.tree_at(key) for key in handle.archive.keys}

finalize 

finalize(handle: CpuHandle) -> EliteArchive

Return the completed archive.

Source code in src/hifuku/cpu_backend.py 
191
192
193
def finalize(self, handle: CpuHandle) -> EliteArchive:
    """Return the completed archive."""
    return handle.archive

CudaBackend

Run the survey on a CUDA device through the shared driver.

Source code in src/hifuku/gpu_backend.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
class CudaBackend:
    """Run the survey on a CUDA device through the shared driver."""

    def __init__(self, gpu_params: GpuParams | None = None):
        self.gpu_params = gpu_params or GpuParams()
        self.caps = Capabilities(name="numba-cuda",
                                 unified_memory=_unified_memory(),
                                 defaults=dict(_GPU_DEFAULTS),
                                 can_read_trees=True)

    # -- setup ------------------------------------------------------------

    def prepare(self, ctx: RunContext, params) -> CudaHandle:
        """Allocate the chains, the device archive and the buffers.

        This compiles the walk kernel for the model's state count and lays out
        the niche allocation.  It runs no survey work.
        """
        _warn_if_parents_go_stale(params)

        from hifuku.gpu_walk import (
            _canon_leaf_names,
            _device_grid,
            _precompute_triangle_arrays,
            _relabel_to,
            _walk_kernel_for,
            state_dim,
        )

        r1, r2, r3 = ctx.anchors
        model, triangle = ctx.model, ctx.triangle
        M, K = params.n_walkers, params.walk_steps

        seed_trees = [r1, r2, r3, *ctx.start_trees]
        for t in seed_trees:
            _set_global_leaf_indices(t, ctx.taxon_table)

        # One shared leaf order, so the single tip-partials and leaf-index
        # arrays of the chains are correct for every walker.
        canon = _canon_leaf_names(r1, ctx.taxon_table)
        cr1, cr2, cr3 = (_relabel_to(t, canon, ctx.taxon_table)
                         for t in (r1, r2, r3))
        canon_seeds = [_relabel_to(t, canon, ctx.taxon_table) for t in seed_trees]

        n1 = (M + 2) // 3
        n2 = (M + 1) // 3
        n3 = M - n1 - n2
        chains = init_chain_trees_2d(
            [cr1] * n1 + [cr2] * n2 + [cr3] * n3,
            anchor_ref1=cr1, anchor_ref2=cr2, anchor_ref3=cr3,
            global_seed=params.seed, alignment=ctx.alignment, model=model,
        )

        sdim = state_dim(chains.n_nodes)
        grid = _device_grid(params.cell, params.origin, self.gpu_params.extent)
        n_off, n_side, grid_lo, sink = grid

        A_inv, Mbary_inv, rhs_off, P3_arr = _precompute_triangle_arrays(triangle)
        tm = triangle.tree_metric
        device_arrays = {
            "model_U": cuda.to_device(model.U.astype(np.float64)),
            "model_eigvals": cuda.to_device(model.eigvals.astype(np.float64)),
            "model_inv_sqrt_pi": cuda.to_device(model.inv_sqrt_pi.astype(np.float64)),
            "model_sqrt_pi": cuda.to_device(model.sqrt_pi.astype(np.float64)),
            "freqs": cuda.to_device(model.freqs.astype(np.float64)),
            "weights": cuda.to_device(ctx.alignment.weights.astype(np.float64)),
            "A_inv": cuda.to_device(A_inv),
            "Mbary_inv": cuda.to_device(Mbary_inv),
            "rhs_off": cuda.to_device(rhs_off),
            "P3": cuda.to_device(P3_arr),
            "metric_w": np.float64(tm.w),
            "metric_c_ref": np.float64(tm.c_ref),
            "metric_r_ref": np.float64(tm.r_ref),
        }

        n_cand = M * K
        candidates = {
            "niche": cuda.device_array(n_cand, dtype=np.int32),
            "fit": cuda.device_array(n_cand, dtype=np.float32),
            "z": cuda.device_array(n_cand, dtype=np.float64),
            "state": cuda.device_array((n_cand, sdim), dtype=np.float32),
        }

        return CudaHandle(
            ctx=ctx, params=params, chains=chains,
            archive_dev=new_device_archive(n_side * n_side + 1, sdim),
            device_arrays=device_arrays,
            walk_kernel=_walk_kernel_for(model.n_states),
            canon=list(canon), canon_seeds=canon_seeds, grid=grid,
            rng=np.random.default_rng(params.seed), candidates=candidates,
        )

    def seed(self, handle: CudaHandle) -> None:
        """Place the anchor trees and the start trees without the gate.

        The descriptor and the log-likelihood come from the host, so the seeded
        archive matches the CPU survey exactly.
        """
        from hifuku.gpu_walk import flatten_host

        ctx, params = handle.ctx, handle.params
        n_off, n_side, _grid_lo, _sink = handle.grid
        chains = init_chain_trees_2d(
            handle.canon_seeds,
            anchor_ref1=handle.canon_seeds[0], anchor_ref2=handle.canon_seeds[1],
            anchor_ref3=handle.canon_seeds[2],
            global_seed=params.seed, alignment=ctx.alignment, model=ctx.model,
        )
        host = {name: getattr(chains, attr).copy_to_host() for name, attr in (
            ("parent", "parent"), ("lc", "left_child"), ("rc", "right_child"),
            ("edge", "edge_len"), ("r1", "ref_len_by_node_1"),
            ("r2", "ref_len_by_node_2"), ("r3", "ref_len_by_node_3"),
            ("A", "A"), ("B1", "B1"), ("B2", "B2"), ("B3", "B3"))}

        scratch = EliteArchive(cell=params.cell, origin=params.origin)
        niches, fits, states = [], [], []
        for mi, tree in enumerate(handle.canon_seeds):
            lam1, lam2, _z = locate_tree(tree, ctx.taxon_table, ctx.triangle)
            i, j = scratch.niche_of(lam1, lam2)
            ci, cj = i + n_off, j + n_off
            if not (0 <= ci < n_side and 0 <= cj < n_side):
                continue          # outside the allocation; extreme, so skip
            niches.append(ci * n_side + cj)
            fits.append(float(log_likelihood(tree, ctx.alignment, ctx.model)))
            states.append(flatten_host(
                host["parent"][mi], host["lc"][mi], host["rc"][mi],
                host["edge"][mi], host["r1"][mi], host["r2"][mi], host["r3"][mi],
                host["A"][mi], host["B1"][mi], host["B2"][mi], host["B3"][mi],
                chains.n_nodes))
        if not niches:
            raise ValueError("no seed tree landed on the device allocation")
        place_candidates(
            handle.archive_dev,
            cuda.to_device(np.asarray(niches, dtype=np.int32)),
            cuda.to_device(np.asarray(fits, dtype=np.float32)),
            cuda.to_device(np.asarray(states, dtype=np.float32)),
        )

    # -- the survey loop --------------------------------------------------

    def submit_batch(self, handle: CudaHandle, keys, n_variations: int,
                     floor: float) -> None:
        """Queue whole passes of walk and place.  Do not synchronize.

        The assignments for the whole chunk upload once and the floor holds for
        the chunk, so no pass waits on the host.
        """
        params = handle.params
        M, K = params.n_walkers, params.walk_steps
        passes = max(1, n_variations // (M * K))
        if keys.size == 0:
            return

        assigns = keys[handle.rng.integers(0, keys.size, size=(passes, M))]
        d_assigns = cuda.to_device(np.ascontiguousarray(assigns, dtype=np.int32))

        ct = handle.chains
        da = handle.device_arrays
        cand = handle.candidates
        _n_off, n_side, grid_lo, sink = handle.grid
        from hifuku.gpu_archive import EMPTY_FITNESS

        for bi in range(passes):
            handle.walk_kernel[M, _CUDA_BLOCK_SIZE](
                d_assigns[bi], handle.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,
                da["model_U"], da["model_eigvals"], da["model_inv_sqrt_pi"],
                da["model_sqrt_pi"], da["freqs"], da["weights"],
                da["A_inv"], da["Mbary_inv"], da["rhs_off"], da["P3"],
                np.float64(ct.C1), np.float64(ct.C2), np.float64(ct.C3),
                da["metric_w"], da["metric_c_ref"], da["metric_r_ref"],
                cand["niche"], cand["fit"], cand["z"], cand["state"],
                np.int32(K), np.int32(ct.n_leaves), np.int32(ct.n_nodes),
                np.int32(ct.n_internal), np.int32(ct.n_patterns),
                np.int32(ct.n_ref),
                np.float64(params.sigma_slide), np.float64(params.scale_width),
                np.float64(grid_lo), np.float64(params.cell), np.int32(n_side),
                np.int32(sink), np.float32(EMPTY_FITNESS), np.float64(floor),
            )
            place_candidates(handle.archive_dev, cand["niche"], cand["fit"],
                             cand["state"])

    def sync(self) -> None:
        """Wait for the queued passes to finish."""
        cuda.synchronize()

    # -- reads ------------------------------------------------------------

    def read_stats(self, handle: CudaHandle) -> Stats:
        """Read the survey statistics as device scalars.

        The elite range is ``best`` less the smallest occupied fitness.  A rising
        relative floor can leave early elites below it, so the smallest elite is
        reduced on the device rather than taken as ``best - floor``.
        """
        dev = handle.archive_dev
        n_keys = int(dev["n_keys"].copy_to_host()[0])
        gain = float(dev["stat_gain"].copy_to_host()[0])
        best = float(dev["stat_best"].copy_to_host()[0])
        if n_keys == 0:
            return Stats(n_keys=0, gain=gain, best=float("-inf"), elite_range=0.0)
        elite_range = best - min_occupied_fitness(handle.archive_dev)
        return Stats(n_keys=n_keys, gain=gain, best=best,
                     elite_range=max(elite_range, 0.0))

    def read_keys(self, handle: CudaHandle, since: int) -> np.ndarray:
        """Copy the niche keys appended at or after ``since``, in niche order.

        Only the new slice moves, so a checkpoint does not copy the whole key
        array.  A key equal to the sink is dropped: it marks a candidate that
        fell outside the allocation.

        The slice is sorted before it is returned.  A winning thread claims its
        slot in the key array with an atomic counter, so the device stores the
        keys of one pass in the order the threads finish.  The driver hands this
        list back for parent selection, so an unsorted list would feed the
        thread schedule into the search and a survey would differ from run to
        run.  The set of keys a pass adds is fixed, so sorting each slice makes
        the whole key list a function of the run parameters.
        """
        dev = handle.archive_dev
        n_keys = int(dev["n_keys"].copy_to_host()[0])
        if n_keys <= since:
            return np.empty(0, dtype=np.int32)
        _n_off, _n_side, _grid_lo, sink = handle.grid
        fresh = dev["keys"][since:n_keys].copy_to_host()
        return np.sort(fresh[fresh != sink])

    def _occupied(self, handle: CudaHandle):
        """``(flat_indices, fitness)`` for the occupied niches."""
        _n_off, n_side, _grid_lo, _sink = handle.grid
        fitness = handle.archive_dev["fitness"].copy_to_host()[:n_side * n_side]
        return np.nonzero(fitness > _OCCUPIED)[0], fitness

    def read_map(self, handle: CudaHandle) -> np.ndarray:
        """Return the elite log-likelihood grid, ``NaN`` where unoccupied.

        The grid spans the bounding box of the filled niches, which is the same
        window the host backend reports.
        """
        n_off, n_side, _grid_lo, _sink = handle.grid
        idx, fitness = self._occupied(handle)
        if idx.size == 0:
            return np.full((1, 1), np.nan)
        rows, cols = idx // n_side, idx % n_side
        r0, r1 = int(rows.min()), int(rows.max())
        c0, c1 = int(cols.min()), int(cols.max())
        grid = np.full((r1 - r0 + 1, c1 - c0 + 1), np.nan, dtype=np.float64)
        grid[rows - r0, cols - c0] = fitness[idx]
        return grid

    def read_trees(self, handle: CudaHandle) -> dict:
        """Rebuild ``{(i, j): HifukuTree}`` from the device archive state."""
        from hifuku.gpu_walk import _tree_from_state

        n_off, n_side, _grid_lo, _sink = handle.grid
        idx, _fitness = self._occupied(handle)
        state = handle.archive_dev["state"].copy_to_host()[:n_side * n_side]
        out = {}
        for flat in idx:
            ci, cj = int(flat) // n_side, int(flat) % n_side
            out[(ci - n_off, cj - n_off)] = _tree_from_state(
                state[flat], handle.canon, handle.chains.n_nodes,
                handle.chains.n_leaves, handle.ctx.taxon_table)
        return out

    def finalize(self, handle: CudaHandle) -> EliteArchive:
        """Rebuild the host archive from the device fitness and state."""
        from hifuku.gpu_walk import _tree_from_state, _z_from_state

        ctx, params = handle.ctx, handle.params
        n_off, n_side, _grid_lo, _sink = handle.grid
        idx, fitness = self._occupied(handle)
        state = handle.archive_dev["state"].copy_to_host()[:n_side * n_side]

        archive = EliteArchive(cell=params.cell, origin=params.origin,
                               anchors=ctx.triangle, taxon_table=ctx.taxon_table,
                               metric=ctx.triangle.tree_metric)
        at_edge = False
        for flat in idx:
            ci, cj = int(flat) // n_side, int(flat) % n_side
            if ci in (0, n_side - 1) or cj in (0, n_side - 1):
                at_edge = True
            tree = _tree_from_state(state[flat], handle.canon,
                                    handle.chains.n_nodes,
                                    handle.chains.n_leaves, ctx.taxon_table)
            z = _z_from_state(state[flat], handle.chains.n_nodes,
                              handle.chains.C1, handle.chains.C2,
                              handle.chains.C3, ctx.triangle)
            archive.place((ci - n_off, cj - n_off),
                          Elite(logL=float(fitness[flat]), 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 archive

prepare 

prepare(ctx: RunContext, params) -> CudaHandle

Allocate the chains, the device archive and the buffers.

This compiles the walk kernel for the model's state count and lays out the niche allocation. It runs no survey work.

Source code in src/hifuku/gpu_backend.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
def prepare(self, ctx: RunContext, params) -> CudaHandle:
    """Allocate the chains, the device archive and the buffers.

    This compiles the walk kernel for the model's state count and lays out
    the niche allocation.  It runs no survey work.
    """
    _warn_if_parents_go_stale(params)

    from hifuku.gpu_walk import (
        _canon_leaf_names,
        _device_grid,
        _precompute_triangle_arrays,
        _relabel_to,
        _walk_kernel_for,
        state_dim,
    )

    r1, r2, r3 = ctx.anchors
    model, triangle = ctx.model, ctx.triangle
    M, K = params.n_walkers, params.walk_steps

    seed_trees = [r1, r2, r3, *ctx.start_trees]
    for t in seed_trees:
        _set_global_leaf_indices(t, ctx.taxon_table)

    # One shared leaf order, so the single tip-partials and leaf-index
    # arrays of the chains are correct for every walker.
    canon = _canon_leaf_names(r1, ctx.taxon_table)
    cr1, cr2, cr3 = (_relabel_to(t, canon, ctx.taxon_table)
                     for t in (r1, r2, r3))
    canon_seeds = [_relabel_to(t, canon, ctx.taxon_table) for t in seed_trees]

    n1 = (M + 2) // 3
    n2 = (M + 1) // 3
    n3 = M - n1 - n2
    chains = init_chain_trees_2d(
        [cr1] * n1 + [cr2] * n2 + [cr3] * n3,
        anchor_ref1=cr1, anchor_ref2=cr2, anchor_ref3=cr3,
        global_seed=params.seed, alignment=ctx.alignment, model=model,
    )

    sdim = state_dim(chains.n_nodes)
    grid = _device_grid(params.cell, params.origin, self.gpu_params.extent)
    n_off, n_side, grid_lo, sink = grid

    A_inv, Mbary_inv, rhs_off, P3_arr = _precompute_triangle_arrays(triangle)
    tm = triangle.tree_metric
    device_arrays = {
        "model_U": cuda.to_device(model.U.astype(np.float64)),
        "model_eigvals": cuda.to_device(model.eigvals.astype(np.float64)),
        "model_inv_sqrt_pi": cuda.to_device(model.inv_sqrt_pi.astype(np.float64)),
        "model_sqrt_pi": cuda.to_device(model.sqrt_pi.astype(np.float64)),
        "freqs": cuda.to_device(model.freqs.astype(np.float64)),
        "weights": cuda.to_device(ctx.alignment.weights.astype(np.float64)),
        "A_inv": cuda.to_device(A_inv),
        "Mbary_inv": cuda.to_device(Mbary_inv),
        "rhs_off": cuda.to_device(rhs_off),
        "P3": cuda.to_device(P3_arr),
        "metric_w": np.float64(tm.w),
        "metric_c_ref": np.float64(tm.c_ref),
        "metric_r_ref": np.float64(tm.r_ref),
    }

    n_cand = M * K
    candidates = {
        "niche": cuda.device_array(n_cand, dtype=np.int32),
        "fit": cuda.device_array(n_cand, dtype=np.float32),
        "z": cuda.device_array(n_cand, dtype=np.float64),
        "state": cuda.device_array((n_cand, sdim), dtype=np.float32),
    }

    return CudaHandle(
        ctx=ctx, params=params, chains=chains,
        archive_dev=new_device_archive(n_side * n_side + 1, sdim),
        device_arrays=device_arrays,
        walk_kernel=_walk_kernel_for(model.n_states),
        canon=list(canon), canon_seeds=canon_seeds, grid=grid,
        rng=np.random.default_rng(params.seed), candidates=candidates,
    )

seed 

seed(handle: CudaHandle) -> None

Place the anchor trees and the start trees without the gate.

The descriptor and the log-likelihood come from the host, so the seeded archive matches the CPU survey exactly.

Source code in src/hifuku/gpu_backend.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def seed(self, handle: CudaHandle) -> None:
    """Place the anchor trees and the start trees without the gate.

    The descriptor and the log-likelihood come from the host, so the seeded
    archive matches the CPU survey exactly.
    """
    from hifuku.gpu_walk import flatten_host

    ctx, params = handle.ctx, handle.params
    n_off, n_side, _grid_lo, _sink = handle.grid
    chains = init_chain_trees_2d(
        handle.canon_seeds,
        anchor_ref1=handle.canon_seeds[0], anchor_ref2=handle.canon_seeds[1],
        anchor_ref3=handle.canon_seeds[2],
        global_seed=params.seed, alignment=ctx.alignment, model=ctx.model,
    )
    host = {name: getattr(chains, attr).copy_to_host() for name, attr in (
        ("parent", "parent"), ("lc", "left_child"), ("rc", "right_child"),
        ("edge", "edge_len"), ("r1", "ref_len_by_node_1"),
        ("r2", "ref_len_by_node_2"), ("r3", "ref_len_by_node_3"),
        ("A", "A"), ("B1", "B1"), ("B2", "B2"), ("B3", "B3"))}

    scratch = EliteArchive(cell=params.cell, origin=params.origin)
    niches, fits, states = [], [], []
    for mi, tree in enumerate(handle.canon_seeds):
        lam1, lam2, _z = locate_tree(tree, ctx.taxon_table, ctx.triangle)
        i, j = scratch.niche_of(lam1, lam2)
        ci, cj = i + n_off, j + n_off
        if not (0 <= ci < n_side and 0 <= cj < n_side):
            continue          # outside the allocation; extreme, so skip
        niches.append(ci * n_side + cj)
        fits.append(float(log_likelihood(tree, ctx.alignment, ctx.model)))
        states.append(flatten_host(
            host["parent"][mi], host["lc"][mi], host["rc"][mi],
            host["edge"][mi], host["r1"][mi], host["r2"][mi], host["r3"][mi],
            host["A"][mi], host["B1"][mi], host["B2"][mi], host["B3"][mi],
            chains.n_nodes))
    if not niches:
        raise ValueError("no seed tree landed on the device allocation")
    place_candidates(
        handle.archive_dev,
        cuda.to_device(np.asarray(niches, dtype=np.int32)),
        cuda.to_device(np.asarray(fits, dtype=np.float32)),
        cuda.to_device(np.asarray(states, dtype=np.float32)),
    )

submit_batch 

submit_batch(handle: CudaHandle, keys, n_variations: int, floor: float) -> None

Queue whole passes of walk and place. Do not synchronize.

The assignments for the whole chunk upload once and the floor holds for the chunk, so no pass waits on the host.

Source code in src/hifuku/gpu_backend.py 
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
def submit_batch(self, handle: CudaHandle, keys, n_variations: int,
                 floor: float) -> None:
    """Queue whole passes of walk and place.  Do not synchronize.

    The assignments for the whole chunk upload once and the floor holds for
    the chunk, so no pass waits on the host.
    """
    params = handle.params
    M, K = params.n_walkers, params.walk_steps
    passes = max(1, n_variations // (M * K))
    if keys.size == 0:
        return

    assigns = keys[handle.rng.integers(0, keys.size, size=(passes, M))]
    d_assigns = cuda.to_device(np.ascontiguousarray(assigns, dtype=np.int32))

    ct = handle.chains
    da = handle.device_arrays
    cand = handle.candidates
    _n_off, n_side, grid_lo, sink = handle.grid
    from hifuku.gpu_archive import EMPTY_FITNESS

    for bi in range(passes):
        handle.walk_kernel[M, _CUDA_BLOCK_SIZE](
            d_assigns[bi], handle.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,
            da["model_U"], da["model_eigvals"], da["model_inv_sqrt_pi"],
            da["model_sqrt_pi"], da["freqs"], da["weights"],
            da["A_inv"], da["Mbary_inv"], da["rhs_off"], da["P3"],
            np.float64(ct.C1), np.float64(ct.C2), np.float64(ct.C3),
            da["metric_w"], da["metric_c_ref"], da["metric_r_ref"],
            cand["niche"], cand["fit"], cand["z"], cand["state"],
            np.int32(K), np.int32(ct.n_leaves), np.int32(ct.n_nodes),
            np.int32(ct.n_internal), np.int32(ct.n_patterns),
            np.int32(ct.n_ref),
            np.float64(params.sigma_slide), np.float64(params.scale_width),
            np.float64(grid_lo), np.float64(params.cell), np.int32(n_side),
            np.int32(sink), np.float32(EMPTY_FITNESS), np.float64(floor),
        )
        place_candidates(handle.archive_dev, cand["niche"], cand["fit"],
                         cand["state"])

sync 

sync() -> None

Wait for the queued passes to finish.

Source code in src/hifuku/gpu_backend.py 
306
307
308
def sync(self) -> None:
    """Wait for the queued passes to finish."""
    cuda.synchronize()

read_stats 

read_stats(handle: CudaHandle) -> Stats

Read the survey statistics as device scalars.

The elite range is best less the smallest occupied fitness. A rising relative floor can leave early elites below it, so the smallest elite is reduced on the device rather than taken as best - floor.

Source code in src/hifuku/gpu_backend.py 
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
def read_stats(self, handle: CudaHandle) -> Stats:
    """Read the survey statistics as device scalars.

    The elite range is ``best`` less the smallest occupied fitness.  A rising
    relative floor can leave early elites below it, so the smallest elite is
    reduced on the device rather than taken as ``best - floor``.
    """
    dev = handle.archive_dev
    n_keys = int(dev["n_keys"].copy_to_host()[0])
    gain = float(dev["stat_gain"].copy_to_host()[0])
    best = float(dev["stat_best"].copy_to_host()[0])
    if n_keys == 0:
        return Stats(n_keys=0, gain=gain, best=float("-inf"), elite_range=0.0)
    elite_range = best - min_occupied_fitness(handle.archive_dev)
    return Stats(n_keys=n_keys, gain=gain, best=best,
                 elite_range=max(elite_range, 0.0))

read_keys 

read_keys(handle: CudaHandle, since: int) -> np.ndarray

Copy the niche keys appended at or after since, in niche order.

Only the new slice moves, so a checkpoint does not copy the whole key array. A key equal to the sink is dropped: it marks a candidate that fell outside the allocation.

The slice is sorted before it is returned. A winning thread claims its slot in the key array with an atomic counter, so the device stores the keys of one pass in the order the threads finish. The driver hands this list back for parent selection, so an unsorted list would feed the thread schedule into the search and a survey would differ from run to run. The set of keys a pass adds is fixed, so sorting each slice makes the whole key list a function of the run parameters.

Source code in src/hifuku/gpu_backend.py 
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
def read_keys(self, handle: CudaHandle, since: int) -> np.ndarray:
    """Copy the niche keys appended at or after ``since``, in niche order.

    Only the new slice moves, so a checkpoint does not copy the whole key
    array.  A key equal to the sink is dropped: it marks a candidate that
    fell outside the allocation.

    The slice is sorted before it is returned.  A winning thread claims its
    slot in the key array with an atomic counter, so the device stores the
    keys of one pass in the order the threads finish.  The driver hands this
    list back for parent selection, so an unsorted list would feed the
    thread schedule into the search and a survey would differ from run to
    run.  The set of keys a pass adds is fixed, so sorting each slice makes
    the whole key list a function of the run parameters.
    """
    dev = handle.archive_dev
    n_keys = int(dev["n_keys"].copy_to_host()[0])
    if n_keys <= since:
        return np.empty(0, dtype=np.int32)
    _n_off, _n_side, _grid_lo, sink = handle.grid
    fresh = dev["keys"][since:n_keys].copy_to_host()
    return np.sort(fresh[fresh != sink])

read_map 

read_map(handle: CudaHandle) -> np.ndarray

Return the elite log-likelihood grid, NaN where unoccupied.

The grid spans the bounding box of the filled niches, which is the same window the host backend reports.

Source code in src/hifuku/gpu_backend.py 
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
def read_map(self, handle: CudaHandle) -> np.ndarray:
    """Return the elite log-likelihood grid, ``NaN`` where unoccupied.

    The grid spans the bounding box of the filled niches, which is the same
    window the host backend reports.
    """
    n_off, n_side, _grid_lo, _sink = handle.grid
    idx, fitness = self._occupied(handle)
    if idx.size == 0:
        return np.full((1, 1), np.nan)
    rows, cols = idx // n_side, idx % n_side
    r0, r1 = int(rows.min()), int(rows.max())
    c0, c1 = int(cols.min()), int(cols.max())
    grid = np.full((r1 - r0 + 1, c1 - c0 + 1), np.nan, dtype=np.float64)
    grid[rows - r0, cols - c0] = fitness[idx]
    return grid

read_trees 

read_trees(handle: CudaHandle) -> dict

Rebuild {(i, j): HifukuTree} from the device archive state.

Source code in src/hifuku/gpu_backend.py 
375
376
377
378
379
380
381
382
383
384
385
386
387
388
def read_trees(self, handle: CudaHandle) -> dict:
    """Rebuild ``{(i, j): HifukuTree}`` from the device archive state."""
    from hifuku.gpu_walk import _tree_from_state

    n_off, n_side, _grid_lo, _sink = handle.grid
    idx, _fitness = self._occupied(handle)
    state = handle.archive_dev["state"].copy_to_host()[:n_side * n_side]
    out = {}
    for flat in idx:
        ci, cj = int(flat) // n_side, int(flat) % n_side
        out[(ci - n_off, cj - n_off)] = _tree_from_state(
            state[flat], handle.canon, handle.chains.n_nodes,
            handle.chains.n_leaves, handle.ctx.taxon_table)
    return out

finalize 

finalize(handle: CudaHandle) -> EliteArchive

Rebuild the host archive from the device fitness and state.

Source code in src/hifuku/gpu_backend.py 
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
def finalize(self, handle: CudaHandle) -> EliteArchive:
    """Rebuild the host archive from the device fitness and state."""
    from hifuku.gpu_walk import _tree_from_state, _z_from_state

    ctx, params = handle.ctx, handle.params
    n_off, n_side, _grid_lo, _sink = handle.grid
    idx, fitness = self._occupied(handle)
    state = handle.archive_dev["state"].copy_to_host()[:n_side * n_side]

    archive = EliteArchive(cell=params.cell, origin=params.origin,
                           anchors=ctx.triangle, taxon_table=ctx.taxon_table,
                           metric=ctx.triangle.tree_metric)
    at_edge = False
    for flat in idx:
        ci, cj = int(flat) // n_side, int(flat) % n_side
        if ci in (0, n_side - 1) or cj in (0, n_side - 1):
            at_edge = True
        tree = _tree_from_state(state[flat], handle.canon,
                                handle.chains.n_nodes,
                                handle.chains.n_leaves, ctx.taxon_table)
        z = _z_from_state(state[flat], handle.chains.n_nodes,
                          handle.chains.C1, handle.chains.C2,
                          handle.chains.C3, ctx.triangle)
        archive.place((ci - n_off, cj - n_off),
                      Elite(logL=float(fitness[flat]), 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 archive