fix: start NUTS chains from walkers of different lineages after stuck-walker re-seeding

    removed this way carry a negligible posterior probability (see

    :func:`stuck_walker_threshold`).

    A copy continues the history of its donor, so it is not an independent

    start for a NUTS chain. ``ctx["walker_lineage"]`` records, for every

    walker, the original walker whose history it continues;

    :func:`nuts_start_walkers` takes at most one walker per lineage.

    Controlled by ``sampling.reseed_stuck_walkers`` (default true) and

    ``sampling.stuck_walker_threshold`` (``auto`` or a number of nats).

            + (f" (round {round_i + 1})" if round_i else ""),

        )

        donors = rng.choice(healthy, size=n_stuck, replace=True)

        lineage = np.array(ctx.get("walker_lineage", np.arange(num_chains)))

        lineage[stuck] = lineage[donors]

        ctx["walker_lineage"] = lineage

        new_z = {}

        for name, val in z.items():

            spread = np.std(val[healthy], axis=0)

    return total

def nuts_start_walkers(

    energies: Any, lineage: Any | None, num_chains: int

) -> np.ndarray:

    """Indices of the walkers from which the NUTS chains start.

    The lowest-energy walkers, at most one per lineage (see

    :func:`reseed_stuck_walkers`): a re-seeded walker and its donor share their

    history, and chains started from both would agree with each other for a

    while even before converging, which makes R-hat optimistic. If there are

    fewer lineages than chains, the remaining chains take the next lowest

    energies.

    """

    order = np.argsort(np.asarray(energies))

    lineage = np.arange(order.size) if lineage is None else np.asarray(lineage)

    picked: list[int] = []

    seen: set[int] = set()

    for i in order:

        if int(lineage[i]) not in seen:

            picked.append(int(i))

            seen.add(int(lineage[i]))

        if len(picked) == num_chains:

            break

    picked += [int(i) for i in order if int(i) not in picked][

        : num_chains - len(picked)

    ]

    return np.resize(np.array(picked), num_chains)

def check_convergence(

    ctx: dict[str, Any], ctx_obj: SamplingContext, ensemble_mcmc: Any, num_chains: int

) -> dict[str, Any]:

        )

        # Each chain starts from a different low-energy walker (the ensemble has

        # just passed the single-peak test), so that R-hat compares chains that

        # did not share their starting point.

        order = np.argsort(np.asarray(final_energies))

        start_idx = np.resize(order[: max(1, num_chains_if_nuts)], num_chains_if_nuts)

        # just passed the single-peak test), one per lineage, so that R-hat

        # compares chains that did not share their starting point or history.

        start_idx = nuts_start_walkers(

            final_energies, ctx.get("walker_lineage"), num_chains_if_nuts

        )

        nuts_init = {

            name: jnp.asarray(np.asarray(val)[start_idx])

            for name, val in final_positions.items()

        _log(

            plModel,

            f"  NUTS chains start from walkers {start_idx.tolist()} "

            f"(the {num_chains_if_nuts} lowest energies)",

            f"(the {num_chains_if_nuts} lowest energies, one per lineage)",

        )

        nuts_warmup_start = ttt.time()

        nuts_mcmc.warmup(

from PyAETNA.sampling.energy import batch_potential_energy

from PyAETNA.sampling.engine.phases import (

    find_stuck_walkers,

    nuts_start_walkers,

    reseed_stuck_walkers,

    stuck_walker_threshold,

)

    assert (

        "Re-seeding 3 stuck walker(s) [2, 5, 9]" in (tmp_path / "toy.log").read_text()

    )

    # each copy continues the history of a healthy donor

    lineage = np.asarray(ctx["walker_lineage"])

    healthy = [i for i in range(n_walkers) if i not in (2, 5, 9)]

    assert all(lineage[i] in healthy for i in (2, 5, 9))

    assert [lineage[i] for i in healthy] == healthy

def test_reseed_is_a_no_op_on_a_healthy_ensemble(tmp_path: Any) -> None:

    }

    assert reseed_stuck_walkers(ctx, mcmc, 16) == 0

    assert float(mcmc.last_state.z["a"][2]) == 150.0

def test_nuts_starts_take_one_walker_per_lineage() -> None:

    energies = np.array([5.0, 1.0, 2.0, 3.0, 4.0, 0.5])

    # walker 5 is a copy of walker 1, walker 2 a copy of walker 0

    lineage = np.array([0, 1, 0, 3, 4, 1])

    assert nuts_start_walkers(energies, lineage, 3).tolist() == [5, 2, 3]

    # without re-seeding every walker is its own lineage

    assert nuts_start_walkers(energies, None, 3).tolist() == [5, 1, 2]

def test_nuts_starts_with_fewer_lineages_than_chains() -> None:

    energies = np.array([3.0, 1.0, 2.0, 0.5])

    lineage = np.array([0, 0, 0, 0])

    assert nuts_start_walkers(energies, lineage, 3).tolist() == [3, 1, 2]

    # more chains than walkers: the starts repeat

    assert nuts_start_walkers(energies, None, 6).tolist() == [3, 1, 2, 0, 3, 1]