fix: AIC and BIC from the maximum likelihood, ESS per second in the log,...

- The log warning about walkers still stuck after the re-seeding rounds now says that it is harmless when NUTS follows.

- These were found by the paper's NUTS benchmarks: until then the default route had run only on models with GPs,

  which are always in float64.

- **AIC and BIC (2026-10-04)** were computed from the log-posterior of the MAP (likelihood × prior × Jacobian)

  instead of the maximum likelihood, so they moved with the prior widths. They now use the maximum of the

  likelihood within the prior support (BFGS from the best of the MAP and the best posterior samples), printed as

  `ln L (maximum)` and stored with χ² in the new `goodness_of_fit` section of `<config>_results.json`. AICc uses

  the standard correction 2k(k+1)/(N−k−1) (it divided by N−k). The `map` section's score is now called

  `map_logposterior` (was `map_loglike`).

- The log's summary tables print the effective samples per second (the column read 0.00); with nested sampling

  the table is titled POSTERIOR SUMMARY and has no ESS and R-hat columns, which say nothing about independent

  draws.

### Upgrading from 3.x

    )

    # ── Log summary table ─────────────────────────────────────────────────────

    sampling_time = (

        (load_results(files.results).get("run") or {}).get("sampling_elapsed_s")

        if os.path.exists(files.results)

        else None

    )

    print_derived_quantities_table(

        derived_result,

        num_samples_kept,

        num_chains,

        sampling_time=None,

        sampling_time=sampling_time,

        log_func=_log,

        nested=data.get("sampler_type") == "ns",

    )

    # ── Write PDF ─────────────────────────────────────────────────────────────

    num_chains: int,

    sampling_time: float | None = None,

    log_func: Any = print,

    nested: bool = False,

) -> None:

    """Print a DERIVED QUANTITIES summary table to the log.

    With ``nested=True`` (nested-sampling posterior) the ESS, ESS/s and R-hat

    columns are left out: they carry no information on independent draws.

    Mirrors the MCMC summary table format for visual consistency.

    Three sections:

    1. Sampled derived quantities (ecc, omega, u1, u2, orbital geometry,

        log_func("")

    if derived_samples:

        log_func(

        header = (

            f"{'parameter':<30} {'mean':>12} {'std':>12} {'q16':>12} "

            f"{'median':>12} {'q84':>12} {'MAP':>12} {'ess':>10} "

            f"{'ess/s':>10} {'rhat':>8}"

            f"{'median':>12} {'q84':>12} {'MAP':>12}"

        )

        if not nested:

            header += f" {'ess':>10} {'ess/s':>10} {'rhat':>8}"

        log_func(header)

        log_func("-" * 180)

        for name in derived_order:

            # Skip if this parameter was already listed as a user prior

                "" if (np.isnan(rhat) or rhat < 1.01) else " *" if rhat < 1.1 else " **"

            )

            map_val = derived_map.get(name, np.nan)

            log_func(

            row = (

                f"{name:<30} {_fmt(np.mean(vals))} {_fmt(np.std(vals))} "

                f"{_fmt(np.percentile(vals, 16))} {_fmt(np.median(vals))} "

                f"{_fmt(np.percentile(vals, 84))} {_fmt(map_val)}"

                f"{ess:>10.1f} {ess_per_s:>10.2f} {rhat:>7.4f}{rhat_flag}"

            )

            if not nested:

                row += f" {ess:>10.1f} {ess_per_s:>10.2f} {rhat:>7.4f}{rhat_flag}"

            log_func(row)

        log_func("-" * 180)

    if derived_fixed:

    residuals: np.ndarray,

    dy: np.ndarray,

    n_params: int,

    map_loglike: float | None = None,

    max_loglike: float | None = None,

) -> dict[str, float]:

    """

    Compute goodness-of-fit statistics for a set of residuals and uncertainties.

        n_params: Global number of free parameters (used as a conservative DOF

            estimate; exact per-visit DOF bookeeping is non-trivial when some

            parameters are visit-scoped).

        map_loglike: Log-likelihood at MAP for information criteria (optional;

            approximated from chi2 when ``None``).

        max_loglike: Maximum log-likelihood for the information criteria

            (optional; approximated from chi2 when ``None``).  Not the

            log-posterior of the MAP, which also contains the prior.

    Returns:

        dict with keys: chi2, chi2_reduced, dof, bic, aic, rms, normalized_rms,

    chi2_red = chi2 / dof if dof > 0 else np.nan

    # ── information criteria ─────────────────────────────────────────────────

    if map_loglike is not None:

        log_likelihood = map_loglike

    if max_loglike is not None:

        log_likelihood = max_loglike

    else:

        # Gaussian log-likelihood approximation (adequate for ICs even without GP)

        log_likelihood = -0.5 * (

            'map_idx': Index of MAP sample (post-discard indexing, or None)

            'map_params': Dict of MAP parameter values from argmin(potential_energy)

            'map_energy': MAP potential energy

            'map_loglike': -potential_energy at MAP (log-posterior score, not pure log-likelihood)

            'map_logposterior': -potential_energy at MAP (log-posterior score, not a log-likelihood)

            'map_source': How MAP was determined ('potential_energy' or 'median_fallback')

            'median_params': Dict of median parameter values

            'stuck_params': List of stuck parameter names

    if potential_energy_kept is not None and len(potential_energy_kept) > 0:

        map_idx = int(np.argmin(potential_energy_kept))

        map_energy = float(potential_energy_kept[map_idx])

        map_loglike = -map_energy

        map_logposterior = -map_energy

        # For scalar params: take the flat index.  For vector params: locate the

        # (chain, step) position and extract the full vector at that step.

        map_params = {}

        )

        map_idx = None

        map_energy = None

        map_loglike = None

        map_logposterior = None

        map_params = median_params.copy()

        map_source = "median_fallback"

        "map_idx": map_idx,

        "map_params": map_params,

        "map_energy": map_energy,

        "map_loglike": map_loglike,

        "map_logposterior": map_logposterior,

        "map_source": map_source,

        "median_params": median_params,

        "stuck_params": stuck_params,

    log_func: Callable[[str], None] = print,

    show_ptail: bool = False,

    discard_extra_warmup: int = 0,

    nested: bool = False,

) -> None:

    """

    Print formatted MCMC summary table.

    With ``nested=True`` (a nested-sampling posterior) the table is titled

    POSTERIOR SUMMARY and has no ESS, ESS/s and R-hat columns: the draws are

    one sequence of independent resamples, for which ESS is about the number

    of draws and R-hat about 1 whatever the run, so neither says anything about

    convergence.  The log reports the posterior ESS of the weighted dead

    points instead.

    Args:

        result: Dict from compute_mcmc_diagnostics()

        num_chains: Number of chains

        log_func: Function to call for logging (default: print)

        show_ptail: If True, compute and show P(tail) column after MAP (for final summary)

        discard_extra_warmup: Number of extra warmup samples discarded (for title only, default: 0)

        nested: The samples come from nested sampling (see above).

    """

    diagnostics = result["diagnostics"]

    quantiles = result["quantiles"]

    log_func("")

    log_func("=" * 180)

    if discard_extra_warmup > 0:

    if nested:

        log_func("POSTERIOR SUMMARY (nested sampling, independent draws)")

    elif discard_extra_warmup > 0:

        log_func(f"MCMC SUMMARY (after rejecting {discard_extra_warmup} extra warmup)")

    else:

        log_func("MCMC SUMMARY")

    # Choose column header based on whether we're showing P(tail)

    if show_ptail:

        log_func(

            f"{'parameter':<30} {'mean':>12} {'std':>12} {'q16':>12} {'median':>12} {'q84':>12} {'MAP':>12} {'P(tail)':>12} {'ess':>10} {'ess/s':>10} {'rhat':>8}"

        )

    else:

        log_func(

            f"{'parameter':<30} {'mean':>12} {'std':>12} {'q16':>12} {'median':>12} {'q84':>12} {'MAP':>12} {'ess':>10} {'ess/s':>10} {'rhat':>8}"

    header = (

        f"{'parameter':<30} {'mean':>12} {'std':>12} {'q16':>12} {'median':>12} "

        f"{'q84':>12} {'MAP':>12}"

    )

    if show_ptail:

        header += f" {'P(tail)':>12}"

    if not nested:

        header += f" {'ess':>10} {'ess/s':>10} {'rhat':>8}"

    log_func(header)

    log_func("-" * 180)

    for name in param_order:

                ptail_str = f"P(>)={100 * p_tail:.1f}%"

            ptail_s = f"{ptail_str:>12}"

        row = f"{name:<30} {mean_s} {std_s} {q16_s} {med_s} {q84_s} {map_s}"

        if show_ptail:

            log_func(

                f"{name:<30} {mean_s} {std_s} {q16_s} {med_s} {q84_s} {map_s} {ptail_s} "

                f"{diag['ess']:>10.1f} {diag['ess_per_sec']:>10.2f} {diag['rhat']:>7.4f}{rhat_flag}"

            )

        else:

            log_func(

                f"{name:<30} {mean_s} {std_s} {q16_s} {med_s} {q84_s} {map_s} "

                f"{diag['ess']:>10.1f} {diag['ess_per_sec']:>10.2f} {diag['rhat']:>7.4f}{rhat_flag}"

            row += f" {ptail_s}"

        if not nested:

            row += (

                f" {diag['ess']:>10.1f} {diag['ess_per_sec']:>10.2f} "

                f"{diag['rhat']:>7.4f}{rhat_flag}"

            )

        log_func(row)

    log_func("-" * 180)

    if nested:

        log_func(

            "# Independent draws: convergence is judged from dlogz and the "

            "posterior ESS in the sampling section"

        )

    else:

        log_func("# * = rhat > 1.01 (warning), ** = rhat > 1.1 (poor convergence)")

    if show_ptail:

        log_func(

        )

    log_func("=" * 180)

    # Warnings

    if result["stuck_params"] or result["low_ess_params"]:

    # Warnings (low ESS is meaningless for independent nested-sampling draws)

    if result["stuck_params"] or (result["low_ess_params"] and not nested):

        log_func("")

        log_func("WARNING: Chain convergence issues detected!")

        if result["stuck_params"]:

                f"  Parameters with no variation (stuck chain): {result['stuck_params']}"

            )

            log_func("  -> Chain likely got stuck at initialization or boundary")

        if result["low_ess_params"]:

        if result["low_ess_params"] and not nested:

            log_func(

                f"  Parameters with very low ESS (<10): {result['low_ess_params']}"

            )

import numpy as np

from PyAETNA.sampling.eccentricity import add_derived_eccentricity

from PyAETNA.sampling.energy import batch_potential_energy, optimize_potential_energy

from PyAETNA.sampling.energy import (

    batch_log_likelihood,

    batch_potential_energy,

    optimize_log_likelihood,

    optimize_potential_energy,

)

def find_map_parameters(

        dict with keys:

            'map_params'         : Dict of MAP parameter values (free parameters only).

            'map_idx'            : Index of best sample (or None if from optimization).

            'map_loglike'        : Log-posterior at MAP (loglike + logprior, constrained).

            'map_logposterior'   : Log-posterior score at the MAP: minus the

                                   potential energy (log-likelihood + log-prior

                                   + Jacobian).  Not a log-likelihood: the

                                   information criteria use

                                   :func:`find_max_log_likelihood`.

            'map_source'         : One of 'potential_energy', 'last_samples',

                                   'optimization', 'optimization_partial',

                                   'median_fallback'.

    result: dict[str, Any] = {

        "map_params": None,

        "map_idx": None,

        "map_loglike": None,

        "map_logposterior": None,

        "map_source": None,

        "optimization_result": None,

        "median_params": median_params,

        # constrained space without the Jacobian, so the number printed beside

        # the MAP was on a different scale from the one used to rank it.)

        map_energy = float(potential_energy[map_idx])

        map_logposterior = -map_energy

        samples_per_chain = len(potential_energy) // num_chains

        chain_idx = map_idx // samples_per_chain

        sample_in_chain = map_idx % samples_per_chain

        _log(f"Best sample: chain {chain_idx}, sample {sample_in_chain}")

        result["map_params"] = map_params

        result["map_idx"] = map_idx

        result["map_loglike"] = map_logposterior

        result["map_logposterior"] = -map_energy

        result["map_source"] = "potential_energy"

    else:

        _log("No potential_energy available - evaluating chain endpoints...")

        if ep_params is not None and ep_energy is not None:

            result["map_params"] = ep_params

            result["map_idx"] = None

            result["map_loglike"] = -ep_energy  # Convert energy back to log-posterior

            result["map_logposterior"] = (

                -ep_energy

            )  # Convert energy back to log-posterior

            result["map_source"] = "last_samples"

            _log(f"Best chain endpoint: chain {ep_chain_idx} (energy: {ep_energy:.4f})")

        else:

                        refined, param_config.get("planet_configs", [])

                    )

                    result["map_params"] = refined

                    result["map_loglike"] = -opt["energy"]

                    result["map_logposterior"] = -opt["energy"]

                    result["map_idx"] = None

                    # BFGS reports converged=False when it exhausts its budget or

                    # the line search stalls, yet still returns a strictly better

    if result["map_params"] is None:

        _log("FALLBACK: Using median parameters (no potential_energy available)")

        result["map_params"] = median_params.copy()

        result["map_loglike"] = None

        result["map_logposterior"] = None

        result["map_source"] = "median_fallback"

    _log(f"MAP source: {result['map_source']}")

    # if result["map_loglike"] is not None:

    #     _log(f"MAP energy: {-result['map_loglike']:.4f}")

    _log("=" * 90)

    return result

    except Exception as e:

        _log(f"    ERROR evaluating chain endpoints: {e}")

        return None, None, None

def find_max_log_likelihood(

    samples: dict[str, Any],

    potential_energy: np.ndarray | None,

    map_params: dict[str, Any],

    param_config: dict[str, Any],

    plModel: Any,

    plData: Any,

    model_fn: Any,

    log_func: Any | None = None,

    maxiter: int = 200,

    n_candidates: int = 256,

    batch_size: int = 64,

) -> dict[str, Any] | None:

    """Maximum of the data-term log-likelihood, for the information criteria.

    AIC and BIC are defined on the maximum likelihood, not on the MAP: the MAP

    also weighs the prior, so it moves with the prior widths.  The search

    starts from the best of the MAP and the ``n_candidates`` lowest-energy

    posterior samples, then climbs the likelihood with BFGS within the prior

    support (:func:`~PyAETNA.sampling.energy.optimize_log_likelihood`).

    Args:

        samples: Posterior samples kept for the report (after any discard).

        potential_energy: Their potential energies, or ``None``.

        map_params: The MAP of :func:`find_map_parameters`.

        param_config: Parameter configuration (``free_params``).

        plModel, plData: Model and data objects.

        model_fn: The NumPyro model.

        log_func: Logging function.

        maxiter: Maximum BFGS iterations.

        n_candidates: Posterior samples evaluated to pick the starting point.

        batch_size: Samples evaluated at once (lower it for large GPs).

    Returns:

        ``{"max_loglike", "start_loglike", "start", "status"}``, or ``None``

        if the likelihood cannot be evaluated.

    """

    _log = log_func if log_func is not None else print

    latent = [name for name in param_config.get("free_params", {}) if name in samples]

    if model_fn is None or not latent:

        return None

    n = len(np.asarray(samples[latent[0]]))

    if potential_energy is not None and len(potential_energy) == n:

        energy = np.where(np.isfinite(potential_energy), potential_energy, np.inf)

        order = np.argsort(energy)[: min(n_candidates, n)]

    else:

        order = np.linspace(0, n - 1, min(n_candidates, n)).astype(int)

    candidates = {

        name: np.concatenate(

            [

                [float(map_params.get(name, samples[name][order[0]]))],

                samples[name][order],

            ]

        )

        for name in latent

    }

    _prev_x64 = jax.config.jax_enable_x64

    jax.config.update("jax_enable_x64", True)

    try:

        loglikes = []

        for i in range(0, len(order) + 1, batch_size):

            batch = {

                k: jnp.asarray(v[i : i + batch_size], dtype=jnp.float64)

                for k, v in candidates.items()

            }

            loglikes.append(batch_log_likelihood(model_fn, plData, plModel, batch))

        loglike = np.concatenate(loglikes)

        loglike = np.where(np.isfinite(loglike), loglike, -np.inf)

        best = int(np.argmax(loglike))

        if not np.isfinite(loglike[best]):

            _log("WARNING: the log-likelihood is not finite at any posterior sample")

            return None

        start = {k: float(v[best]) for k, v in candidates.items()}

        opt = optimize_log_likelihood(model_fn, plData, plModel, start, maxiter=maxiter)

    except Exception as exc:

        _log(f"WARNING: maximum-likelihood search failed: {exc}")

        return None

    finally:

        jax.config.update("jax_enable_x64", _prev_x64)

    start_loglike = float(loglike[best])

    max_loglike = -float(opt["energy"]) if opt["improved"] else start_loglike

    _log(

        f"Maximum likelihood (for AIC and BIC): ln L = {max_loglike:.4f}, "

        f"from {'the MAP' if best == 0 else 'a posterior sample'} "

        f"(ln L = {start_loglike:.4f}) refined with BFGS "

        f"({opt['nit']} iterations)"

    )

    return {

        "max_loglike": max_loglike,

        "start_loglike": start_loglike,

        "start": "map" if best == 0 else "posterior_sample",

        "status": int(opt["status"]),

    }