mirror of
https://github.com/logos-blockchain/research.git
synced 2026-08-07 03:33:33 +00:00
115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
"""Stake-weighted slot lottery (sparse sampler).
|
|
|
|
Per node ``i`` and slot, an independent Bernoulli win with probability
|
|
``phi_f(alpha_i) = 1 - (1 - f)^alpha_i`` where ``alpha_i = w_i / D_est``. Multiple winners
|
|
in a slot are possible (a guaranteed fork).
|
|
|
|
The number of slots a node wins is exactly ``Binomial(n_slots, p_i)``, and the won slots
|
|
are a uniformly-random distinct subset — this is *distributionally identical* to drawing an
|
|
independent Bernoulli(p_i) in every slot, but avoids materialising the dense
|
|
``(n_nodes, n_slots)`` array (which was ~95% of the whole simulator's runtime). Winners are
|
|
returned as sparse ``(winner_slots, winner_nodes)`` coordinates, sorted by slot.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from joblib import Parallel, delayed
|
|
|
|
|
|
def phi(f: float, alpha: np.ndarray | float) -> np.ndarray | float:
|
|
"""Leader-lottery win probability ``1 - (1 - f)^alpha``."""
|
|
return 1.0 - (1.0 - f) ** alpha
|
|
|
|
|
|
def win_probs(stake: np.ndarray, d_est: float, f: float) -> np.ndarray:
|
|
"""Per-node win probability ``phi_f(w_i / D_est)``."""
|
|
return phi(f, stake / d_est)
|
|
|
|
|
|
def _winners_from_counts(
|
|
counts: np.ndarray, offset: int, span: int, rng: np.random.Generator
|
|
) -> tuple[list[np.ndarray], list[np.ndarray]]:
|
|
"""For each winning node, sample ``counts[i]`` distinct slots in ``[offset, offset+span)``."""
|
|
nz = np.nonzero(counts)[0]
|
|
slot_parts: list[np.ndarray] = []
|
|
node_parts: list[np.ndarray] = []
|
|
for i in nz:
|
|
c = int(counts[i])
|
|
slots_i = rng.choice(span, size=c, replace=False).astype(np.int64) + offset
|
|
slot_parts.append(slots_i)
|
|
node_parts.append(np.full(c, i, np.int64))
|
|
return slot_parts, node_parts
|
|
|
|
|
|
def _finalize(
|
|
slot_parts: list[np.ndarray], node_parts: list[np.ndarray]
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
if not slot_parts:
|
|
return np.empty(0, np.int64), np.empty(0, np.int64)
|
|
winner_slots = np.concatenate(slot_parts)
|
|
winner_nodes = np.concatenate(node_parts)
|
|
order = np.argsort(winner_slots, kind="stable")
|
|
return winner_slots[order], winner_nodes[order]
|
|
|
|
|
|
def sample_wins(
|
|
p_win: np.ndarray, n_slots: int, rng: np.random.Generator, chunk: int = 8192
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Sample lottery wins over ``n_slots`` slots (sparse; ``chunk`` kept for API compat)."""
|
|
counts = rng.binomial(n_slots, p_win)
|
|
slot_parts, node_parts = _winners_from_counts(counts, 0, n_slots, rng)
|
|
return _finalize(slot_parts, node_parts)
|
|
|
|
|
|
def sample_wins_chunked(
|
|
p_win: np.ndarray,
|
|
n_slots: int,
|
|
seedseq: np.random.SeedSequence,
|
|
n_chunks: int,
|
|
n_jobs: int = -1,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Parallel sparse lottery: partition slots into ``n_chunks`` independent ranges.
|
|
|
|
Correct because ``Binomial(n_slots, p) = sum_c Binomial(L_c, p)`` and per-chunk distinct
|
|
subsets are independent. Deterministic given ``(seedseq, n_chunks)`` — but the exact
|
|
winner identities differ from the serial sampler and *change with* ``n_chunks``, so
|
|
``n_chunks`` must be a pinned, recorded config parameter, never derived from core count.
|
|
|
|
Note: after the sparse rewrite the lottery is a small fraction of an epoch, so this has
|
|
little ROI versus across-config parallelism; it exists for the rare isolated config with
|
|
an enormous ``n_slots`` and no across-config work to fill cores.
|
|
"""
|
|
if n_chunks <= 1:
|
|
return sample_wins(p_win, n_slots, np.random.default_rng(seedseq))
|
|
bounds = np.linspace(0, n_slots, n_chunks + 1).astype(np.int64)
|
|
children = seedseq.spawn(n_chunks)
|
|
|
|
def one_chunk(c: int) -> tuple[list[np.ndarray], list[np.ndarray]]:
|
|
lo, hi = int(bounds[c]), int(bounds[c + 1])
|
|
span = hi - lo
|
|
rng = np.random.default_rng(children[c])
|
|
counts = rng.binomial(span, p_win)
|
|
return _winners_from_counts(counts, lo, span, rng)
|
|
|
|
results = Parallel(n_jobs=n_jobs, prefer="threads")(
|
|
delayed(one_chunk)(c) for c in range(n_chunks)
|
|
)
|
|
slot_parts: list[np.ndarray] = []
|
|
node_parts: list[np.ndarray] = []
|
|
for sp, npar in results:
|
|
slot_parts.extend(sp)
|
|
node_parts.extend(npar)
|
|
return _finalize(slot_parts, node_parts)
|
|
|
|
|
|
def group_by_slot(
|
|
winner_slots: np.ndarray, winner_nodes: np.ndarray
|
|
) -> tuple[np.ndarray, list[np.ndarray]]:
|
|
"""Group sorted winner coordinates into ``(active_slots, winners_per_active_slot)``."""
|
|
if winner_slots.size == 0:
|
|
return np.empty(0, np.int64), []
|
|
active_slots, starts = np.unique(winner_slots, return_index=True)
|
|
groups = np.split(winner_nodes, starts[1:])
|
|
return active_slots, groups
|