# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:desy6]
#     language: python
#     name: conda-env-desy6-py
# ---

# +
import fitsio
import proplot as pplt
import numpy as np
import glob
import tqdm
import sys

from mattspy.stats import mad

# +
fnames = glob.glob("./results/meas*.fits")

d = []
for fname in tqdm.tqdm(fnames, ncols=79):
    d.append(fitsio.read(fname, ext="data"))
d = np.concatenate(d)

aps = fitsio.read(fname, ext="aps")

sys.stderr.flush()
sys.stdout.flush()
print("# of samples:", d.shape[0])

ap_ind = np.where(aps == 1.5)[0][0]

# +
s2n_bins = np.linspace(5, 50, 15)

def _resample_stat(x, func, n=500, rng=None):
    if rng is None:
        rng = np.random.RandomState()
        
    vals = []
    for _ in range(n):
        vals.append(func(rng.choice(x, size=len(x), replace=True)))
    return func(x), np.std(vals)


def _resample_stat_two(x, w, func, n=500, rng=None):
    if rng is None:
        rng = np.random.RandomState()
        
    vals = []
    for _ in range(n):
        inds = rng.choice(len(x), size=len(x), replace=True)
        vals.append(func(x[inds], w[inds]))
    return func(x, w), np.std(vals)


def _fvar(x, xerr):
    ex = np.std(x)
    ex_err = np.sqrt(np.sum(xerr**2)/len(x))
    return ex/ex_err-1

nb = 100
chis = []
chi2s = []
out_5 = []
ms2n = []
ffluxes = []
ffluxes_ap = []

for i in range(len(s2n_bins)-1):
    min_s2n = s2n_bins[i]
    max_s2n = s2n_bins[i+1]
    
    msk = (
        (d["flux_flags"][:, ap_ind] == 0) 
        & (d["true_s2n"][:, ap_ind] > min_s2n)
        & (d["true_s2n"][:, ap_ind] < max_s2n)
    )
    ms2n.append(np.mean(d["s2n"][msk, ap_ind]))
    chi_num = (d["flux"][msk, ap_ind] - d["true_ap_flux"][msk, ap_ind])
    chi_den = d["flux_err"][msk, ap_ind]
    chi = chi_num/chi_den
    fflux = (d["true_ap_flux"][msk, ap_ind] - d["true_flux"][msk, ap_ind])/d["true_flux"][msk, ap_ind] + 1
    fflux_ap = (d["flux"][msk, ap_ind] - d["true_ap_flux"][msk, ap_ind])/d["true_ap_flux"][msk, ap_ind]
    chis.append(_resample_stat(chi, lambda x: np.mean(x), n=nb))
    chi2s.append(_resample_stat_two(chi_num, chi_den, _fvar, n=nb))
    out_5.append(_resample_stat(chi, lambda x: np.sum(np.abs(x) > 5)/np.sum(msk), n=nb))
    ffluxes.append(_resample_stat(fflux, np.mean, n=nb))
    ffluxes_ap.append(_resample_stat(fflux_ap, np.mean, n=nb))
    print(
        "%10.3f %10.3f %10.3f %10.3f %10.3f %10.3e %10.3e" % (
            min_s2n, max_s2n,
            np.sum(msk), 
            np.mean(chi),
            _fvar(chi_num, chi_den),
            np.sum(np.abs(chi) > 5)/np.sum(msk),
            np.mean(fflux)
        )
    )
    
chis = np.array(chis)
chi2s = np.array(chi2s)
out_5 = np.array(out_5)
ms2n = np.array(ms2n)
ffluxes = np.array(ffluxes)
ffluxes_ap = np.array(ffluxes_ap)


# +
def _errplot(ax, x, ar):
    ax.plot(
        x,
        ar[:, 0],
    )
    ax.fill_between(
        x,
        ar[:, 0] - ar[:, 1]*3,
        ar[:, 0] + ar[:, 1]*3,
        alpha=0.5
    )
    ax.grid(False)


with pplt.rc.context(textusetex=False):
    fig, axs = pplt.subplots(nrows=3, ncols=1, sharex=3, sharey=0, figsize=(6, 5))

    axs[0, 0].axhline(1, color='k', linestyle="dashed")
    _errplot(axs[0, 0], ms2n, ffluxes)
    axs[0, 0].set_ylabel(r"$F_{\rm ap}/F_{\rm true}$")

    axs[1, 0].axhline(0, color='k', linestyle="dashed")
    _errplot(axs[1, 0], ms2n, ffluxes_ap)
    axs[1, 0].set_ylabel(r"$\Delta F_{\rm ap}/F_{\rm ap}$")

    axs[2, 0].axhline(0, color='k', linestyle="dashed")
    _errplot(axs[2, 0], ms2n, chi2s)
    axs[2, 0].set_xlabel("input signal-to-noise")
    axs[2, 0].set_ylabel(r"$\Delta \sigma_{F_{\rm ap}}/\sigma_{F_{\rm ap}}$")
    
    fig.savefig("test-noround-nodetect.pdf")
# -



StarcoderdataJupyterScriptsDedupFiltered