docs: branches figure with its script; fix the guide's claims on separate modes and efficiency

"""Draw the flowchart of the two sampling branches shown in ``sampling.md``.

Run after changing the sampling pipeline, so that the figure matches it:

    python docs/config/make_branches_figure.py

It writes ``docs/config/img/branches.png``.

"""

from __future__ import annotations

from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt

HERE = Path(__file__).resolve().parent

INK, INK2, GRID = "#0b0b0b", "#52514e", "#dcdad4"

COLUMN_WIDTH = 3.46  # inches

plt.rcParams.update({"font.size": 8, "pdf.fonttype": 42})

def branches() -> None:

    """Flowchart of the two sampling branches."""

    from matplotlib.patches import FancyBboxPatch

    fig, ax = plt.subplots(figsize=(COLUMN_WIDTH, 3.6))

    ax.set_xlim(0, 10)

    ax.set_ylim(0, 12.4)

    ax.axis("off")

    fs = 5.9

    def box(

        x: float,

        y: float,

        text: str,

        w: float = 4.5,

        h: float = 0.95,

        color: str = GRID,

        bold: bool = False,

        key: str = "",

    ) -> None:

        ax.add_patch(

            FancyBboxPatch(

                (x - w / 2, y - h / 2),

                w,

                h,

                boxstyle="round,pad=0.02,rounding_size=0.18",

                fc=color,

                ec="none",

            )

        )

        if key:

            ax.text(

                x,

                y + 0.2,

                text,

                ha="center",

                va="center",

                fontsize=fs,

                color=INK,

                fontweight="bold",

            )

            ax.text(

                x,

                y - 0.22,

                key,

                ha="center",

                va="center",

                fontsize=fs - 0.4,

                color=INK,

                family="monospace",

            )

        else:

            ax.text(

                x,

                y,

                text,

                ha="center",

                va="center",

                fontsize=fs,

                color=INK,

                fontweight="bold" if bold else "normal",

                linespacing=1.15,

            )

    def arrow(x0: float, y0: float, x1: float, y1: float) -> None:

        ax.annotate(

            "",

            xy=(x1, y1),

            xytext=(x0, y0),

            arrowprops={

                "arrowstyle": "-|>",

                "color": INK2,

                "lw": 0.9,

                "shrinkA": 0,

                "shrinkB": 0,

                "mutation_scale": 7,

            },

        )

    blue, aqua = "#d6e6f9", "#d1efe3"

    box(5.0, 11.8, "YAML configuration + data table", w=6.4, bold=True)

    box(2.4, 10.3, "Branch 1", color=blue, key="production_sampler: ns")

    box(7.6, 10.3, "Branch 2", color=aqua, key="production_sampler: auto")

    arrow(4.0, 11.3, 2.6, 10.8)

    arrow(6.0, 11.3, 7.4, 10.8)

    box(

        2.4,

        8.2,

        "nested slice sampling\nfrom the prior,\nvectorised over\nlive points",

        h=2.1,

        color=blue,

    )

    arrow(2.4, 9.8, 2.4, 9.25)

    box(7.6, 8.75, "Phase 1: ensemble warmup\n(ESS or AIES)", color=aqua)

    box(7.6, 7.45, "re-seed stuck walkers;\nPhase 2: energy-spread test", color=aqua)

    arrow(7.6, 9.8, 7.6, 9.25)

    arrow(7.6, 8.25, 7.6, 7.95)

    box(6.45, 5.65, "Phase 3a\nNUTS", w=2.2, color=aqua)

    box(8.75, 5.65, "Phase 3b\nensemble", w=2.2, color=aqua)

    arrow(7.3, 6.95, 6.65, 6.15)

    arrow(7.9, 6.95, 8.55, 6.15)

    ax.text(

        6.9, 6.5, "unimodal", fontsize=fs - 0.6, color=INK2, ha="right", va="center"

    )

    ax.text(

        8.2, 6.5, "multimodal", fontsize=fs - 0.6, color=INK2, ha="left", va="center"

    )

    box(2.4, 5.65, "posterior + $\\ln\\mathcal{Z}$", color=blue)

    arrow(2.4, 7.15, 2.4, 6.15)

    box(7.6, 4.05, "posterior; AIC, BIC", color=aqua)

    arrow(6.45, 5.15, 7.2, 4.55)

    arrow(8.75, 5.15, 8.0, 4.55)

    box(5.0, 2.4, "report: fit, residuals, noise budget,\ndiagnostics, MAP", w=7.4)

    arrow(2.4, 5.15, 4.0, 2.9)

    arrow(7.6, 3.55, 6.0, 2.9)

    box(

        5.0,

        0.75,

        "optional physics stage: derived quantities,\ntemperatures, albedos",

        w=7.4,

    )

    arrow(5.0, 1.9, 5.0, 1.25)

    out = HERE / "img" / "branches.png"

    out.parent.mkdir(parents=True, exist_ok=True)

    fig.savefig(out, dpi=300, bbox_inches="tight", pad_inches=0.02)

    plt.close(fig)

    print(f"wrote {out}")

if __name__ == "__main__":

    branches()

| `production_sampler` | `ns` | `auto` (default), or `nuts`, `ess`, `aies` |

| Bayesian evidence ($\ln Z$) | **Yes**, from the same run | No |

| Model comparison | Bayes factors from $\ln Z$ | AIC, AICc and BIC (in the report) |

| Separate solutions (multimodality) | Each solution gets its correct posterior weight | Ensemble walkers may populate several solutions, but with weights set by where they started |

| Separate solutions (multimodality) | All but one are lost: run one solution at a time, with priors around each | Ensemble walkers may populate several solutions, but with weights set by where they started |

| Warmup | None: it starts from the prior | Ensemble warmup (`auto`), NUTS adaptation, or none (`ess`/`aies`) |

**Rule of thumb.** If you need Bayesian evidence to compare models (circular

vs eccentric, one vs two planets, with or without a GP), or you suspect

several separate solutions, use **Branch 1**. Otherwise use **Branch 2**, whose

`auto` mode is the fastest route to a well-sampled posterior with a single peak.

vs eccentric, one vs two planets, with or without a GP), use **Branch 1**. It was

also the more efficient route in every benchmark of the software paper (5 to 400 times more effective samples per

second than `auto` → NUTS, up to 106 free parameters). If you suspect several separate solutions (e.g. period

aliases), run one solution at a time with priors around each: PyAETNA's nested sampler keeps only one of them.

**Branch 2** is the alternative when no evidence is needed or as an independent check.

![The two sampling branches](img/branches.png)

| Route | Pros | Cons | Use it when |

|---|---|---|---|

| **`ns`** (Branch 1) | $\ln Z$ for model comparison; correct weights for separate solutions; independent posterior draws, no burn-in to judge; starts from the prior, so it cannot get stuck near a bad starting point; clear stopping rule (`dlogz_if_ns`); vectorised on GPU, parallel on CPU cores; checkpointed | More likelihood evaluations than NUTS on a single-peaked posterior; cost grows with the information gained over the prior (so very wide priors cost more) and roughly with the square of the number of free parameters; needs proper priors (all PyAETNA priors are) | You need Bayes factors; the posterior may have several solutions; you want an independent check of an MCMC result |

| **`auto` → NUTS** (Branch 2, Phase 3a) | Most efficient sampler for a single-peaked posterior, including many parameters (tens to hundreds, e.g. TTV fits); nearly independent samples; ensemble warmup finds the peak and the parameter scales first; checkpointed | No $\ln Z$; NUTS cannot move between separate solutions; needs gradients (always available in PyAETNA) | Default choice for parameter estimation |

| **`ns`** (Branch 1) | $\ln Z$ for model comparison; the most efficient route in the paper's benchmarks; independent posterior draws, no burn-in to judge; starts from the prior, so it cannot get stuck near a bad starting point; clear stopping rule (`dlogz_if_ns`); vectorised on GPU, parallel on CPU cores; checkpointed | More likelihood evaluations than NUTS on a single-peaked posterior; cost grows with the information gained over the prior (so very wide priors cost more) and roughly with the square of the number of free parameters; needs proper priors (all PyAETNA priors are); of separate solutions all but one are lost, and $\ln Z$ is overestimated by several nats on strongly curved posteriors | Recommended choice; you need Bayes factors; you want an independent check of an MCMC result |

| **`auto` → NUTS** (Branch 2, Phase 3a) | Gradient-based moves; ensemble warmup finds the peak and the parameter scales first; checkpointed | No $\ln Z$; NUTS cannot move between separate solutions; 5 to 400 times less efficient than nested sampling in the paper's benchmarks (up to 106 parameters; `adapt_mass_matrix_if_nuts` helps at many parameters) | Parameter estimation when no evidence is needed; a check of a nested-sampling result |

| **`auto` → ensemble** (Branch 2, Phase 3b) | Gradient-free; robust to rough likelihood surfaces; checkpointed | No $\ln Z$; walkers cannot move between separate solutions, so their relative weights are not the posterior's; mixing degrades above ~20–30 parameters; convergence must be judged from autocorrelation and R-hat | The warmup finds walkers at very different energies and you want to keep MCMC; prefer `ns` if the solutions matter |

| **`nuts`** (direct) | Skips the warmup; full mass-matrix and step-size adaptation | Starts near the prior median, so it can fail when that is far from the posterior; no $\ln Z$; single-peaked posteriors only | You already know the posterior is single-peaked and well inside the priors |

| **`ess`** / **`aies`** (direct) | Skips the warmup; gradient-free | Burn-in must be discarded by hand (`reports.discard_extra_warmup`); stuck walkers are not re-seeded; same limits as ensemble production above | Quick exploration, or reproducing an earlier ensemble run |

5. **Phase 3b, multimodal → ensemble.** If the spread is > threshold (always,

   with `chain_energy_threshold: 0.0`), the ensemble continues for

   `num_samples_if_ensemble` iterations per walker. If the separate solutions

   matter scientifically, switch to Branch 1, which weights them correctly.

   matter scientifically, run Branch 1 once per solution, with priors around each.

### Direct variants

| Requires gradients | **Yes** | No | No | No |

| Cost per step | High (≈ 8–16× model evaluation) | Low–medium (≈ 1–2×) | Low (≈ 1×) | Medium (slice steps) |

| Convergence ($D$ params) | Fast (50–500 steps) | Moderate–slow (500–5000+) | Moderate–slow (500–5000+) | Fixed by `dlogz_if_ns` |

| Separate solutions | No | Walkers can sit in several, but weights are wrong | As ESS | **Yes, correct weights** |

| Separate solutions | No | Walkers can sit in several, but weights are wrong | As ESS | No: all but one lost (one run per solution) |

| Bayesian evidence | No | No | No | **Yes** |

| Parallelism | Independent chains | Ensemble of walkers | Ensemble of walkers | Batch of replacement chains |