Extract sample stats from PyStan fit.

def get_sample_stats(fit, log_likelihood=None):
    """Extract sample stats from PyStan fit."""
    dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64}

    ndraws = [s - w for s, w in zip(fit.sim["n_save"], fit.sim["warmup2"])]

    extraction = OrderedDict()
    for chain, (pyholder, ndraws) in enumerate(zip(fit.sim["samples"], ndraws)):
        if chain == 0:
            for key in pyholder["sampler_param_names"]:
                extraction[key] = []
        for key, values in zip(pyholder["sampler_param_names"], pyholder["sampler_params"]):
            extraction[key].append(values[-ndraws:])

    data = OrderedDict()
    for key, values in extraction.items():
        values = np.stack(values, axis=0)
        dtype = dtypes.get(key)
        values = values.astype(dtype)
        name = re.sub("__$", "", key)
        name = "diverging" if name == "divergent" else name
        data[name] = values

    # log_likelihood
    if log_likelihood is not None:
        log_likelihood_data = get_draws(fit, variables=log_likelihood)
        data["log_likelihood"] = log_likelihood_data[log_likelihood]

    # lp__
    stat_lp = get_draws(fit, variables="lp__")
    data["lp"] = stat_lp["lp__"]

    return data