SFI.inference.sparse.beam module

SFI.inference.sparse.beam — Bidirectional beam search

The BeamSearchStrategy explores the support lattice by expanding every ±1 neighbour of every support in the current frontier and retaining only the beam_width best models per cardinality k.

This is the original algorithm described in the PASTIS paper.

Performance notes

• Child generation uses pure Python integer lists — no per-child JAX array allocations.

• Each child is scored individually via info_and_coeffs (the single-support JIT kernel) rather than batched through scorer.vmap_info. On CPU this is the faster path here because beam children come in many distinct cardinalities, and the vmap cache would otherwise compile O(max_k × unique_batch_sizes) shapes; per-child scoring caps the cache at O(max_k) entries.

• Frontier capping uses heapq.nlargest instead of full sort.

• Skyline update only touches cardinalities modified in the current generation.

class SFI.inference.sparse.beam.BeamSearchStrategy(*, beam_width=20, aic_patience=2, report_time=False)

Bases: SparsityStrategy

Bidirectional beam search over the support lattice.

Parameters:

• beam_width (int, default 20) – Maximum number of candidate models retained per cardinality.

• aic_patience (int, default 2) – Stop early when AIC has strictly declined for this many consecutive closed cardinality levels.

• report_time (bool, default False) – If True, log elapsed time and number of explored supports.

name: str = 'beam'

Short identifier used in SparsityResult.method.

run(scorer, *, max_k, init_supports=None, **_kwargs)

Execute the beam search.

Parameters:

• scorer (SparseScorer)

• max_k (int) – Maximum model size to consider.

• init_supports (list of tuples of int, optional) – Seed supports to inject into the initial frontier. Each entry is a tuple of basis-function indices. Useful for seeding the search with a known good model (e.g. the true support) so that the Pareto front is guaranteed to include it.

Return type:

SparsityResult