jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _bin_size_factor(adata):
	"""
		Organize the size factors into different groups
	"""
	size_factor = adata.obs['memento_size_factor'].values
	binned_stat = s.binned_statistic(size_factor, size_factor, bins=adata.uns['memento']['num_bins'], statistic='mean')
	bin_idx = np.clip(binned_stat[2], a_min=1, a_max=binned_stat[0].shape[0])
	approx_sf = binned_stat[0][bin_idx-1]
	max_sf = size_factor.max()
	approx_sf[size_factor == max_sf] = max_sf

	adata.uns['memento']['all_approx_size_factor'] = approx_sf
	adata.uns['memento']['approx_size_factor'] = \
		{group:approx_sf[(adata.obs['memento_group'] == group).values] for group in adata.uns['memento']['groups']}
	adata.uns['memento']['size_factor'] = \
		{group:size_factor[(adata.obs['memento_group'] == group).values] for group in adata.uns['memento']['groups']}
//  [python]
def _bin_size_factor(adata):
	"""
		Organize the size factors into different groups
	"""
	size_factor = adata.obs['memento_size_factor'].values
	binned_stat = stats.binned_statistic(size_factor, size_factor, bins=adata.uns['memento']['num_bins'], statistic='mean')
	bin_idx = np.clip(binned_stat[2], a_min=1, a_max=binned_stat[0].shape[0])
	approx_sf = binned_stat[0][bin_idx-1]
	max_sf = size_factor.max()
	approx_sf[size_factor == max_sf] = max_sf

	adata.uns['memento']['all_approx_size_factor'] = approx_sf
	adata.uns['memento']['approx_size_factor'] = \
		{group:approx_sf[(adata.obs['memento_group'] == group).values] for group in adata.uns['memento']['groups']}
	adata.uns['memento']['size_factor'] = \
		{group:size_factor[(adata.obs['memento_group'] == group).values] for group in adata.uns['memento']['groups']}
//  [python]
def plot_sfs(s, yscale='log', bins=None, n=None,
             clip_endpoints=True, label=None, plot_kwargs=None,
             ax=None):
    """Plot a site frequency spectrum.

    Parameters
    ----------
    s : array_like, int, shape (n_chromosomes,)
        Site frequency spectrum.
    yscale : string, optional
        Y axis scale.
    bins : int or array_like, int, optional
        Allele count bins.
    n : int, optional
        Number of chromosomes sampled. If provided, X axis will be plotted
        as allele frequency, otherwise as allele count.
    clip_endpoints : bool, optional
        If True, do not plot first and last values from frequency spectrum.
    label : string, optional
        Label for data series in plot.
    plot_kwargs : dict-like
        Additional keyword arguments, passed through to ax.plot().
    ax : axes, optional
        Axes on which to draw. If not provided, a new figure will be created.

    Returns
    -------
    ax : axes
        The axes on which the plot was drawn.

    """

    import matplotlib.pyplot as plt
    import numpy as np

    # check inputs
    s = asarray_ndim(s, 1)

    # setup axes
    if ax is None:
        fig, ax = plt.subplots()

    # setup data
    if bins is None:
        if clip_endpoints:
            x = np.arange(1, s.shape[0]-1)
            y = s[1:-1]
        else:
            x = np.arange(s.shape[0])
            y = s
    else:
        if clip_endpoints:
            y = np.histogram(np.arange(1, s.shape[0]-1), bins=bins, weights=s[1:-1])[0]
            x = (np.histogram(np.arange(1, s.shape[0]-1), bins=bins)[1][:-1] + np.histogram(np.arange(1, s.shape[0]-1), bins=bins)[1][1:]) / 2
        else:
            y = np.histogram(np.arange(s.shape[0]), bins=bins, weights=s)[0]
            x = (np.histogram(np.arange(s.shape[0]), bins=bins)[1][:-1] + np.histogram(np.arange(s.shape[0]), bins=bins)[1][1:]) / 2

    if n:
        # convert allele counts to allele frequencies
        x = x / n
        ax.set_xlabel('derived allele frequency')
    else:
        ax.set_xlabel('derived allele count')

    # do plotting
    if plot_kwargs is None:
        plot_kwargs = dict()
    ax.plot(x, y, label=label, **plot_kwargs)

    # tidy
    ax.set_yscale(yscale)
    ax.set_ylabel('site frequency')
    ax.autoscale(axis='x', tight=True)

    return ax
//  [python]
def plot_sfs(s, yscale='log', bins=None, n=None,
             clip_endpoints=True, label=None, plot_kwargs=None,
             ax=None):
    """Plot a site frequency spectrum.

    Parameters
    ----------
    s : array_like, int, shape (n_chromosomes,)
        Site frequency spectrum.
    yscale : string, optional
        Y axis scale.
    bins : int or array_like, int, optional
        Allele count bins.
    n : int, optional
        Number of chromosomes sampled. If provided, X axis will be plotted
        as allele frequency, otherwise as allele count.
    clip_endpoints : bool, optional
        If True, do not plot first and last values from frequency spectrum.
    label : string, optional
        Label for data series in plot.
    plot_kwargs : dict-like
        Additional keyword arguments, passed through to ax.plot().
    ax : axes, optional
        Axes on which to draw. If not provided, a new figure will be created.

    Returns
    -------
    ax : axes
        The axes on which the plot was drawn.

    """

    import matplotlib.pyplot as plt
    import scipy

    # check inputs
    s = asarray_ndim(s, 1)

    # setup axes
    if ax is None:
        fig, ax = plt.subplots()

    # setup data
    if bins is None:
        if clip_endpoints:
            x = np.arange(1, s.shape[0]-1)
            y = s[1:-1]
        else:
            x = np.arange(s.shape[0])
            y = s
    else:
        if clip_endpoints:
            y, b, _ = scipy.stats.binned_statistic(
                np.arange(1, s.shape[0]-1),
                values=s[1:-1],
                bins=bins,
                statistic='sum')
        else:
            y, b, _ = scipy.stats.binned_statistic(
                np.arange(s.shape[0]),
                values=s,
                bins=bins,
                statistic='sum')
        # use bin midpoints for plotting
        x = (b[:-1] + b[1:]) / 2

    if n:
        # convert allele counts to allele frequencies
        x = x / n
        ax.set_xlabel('derived allele frequency')
    else:
        ax.set_xlabel('derived allele count')

    # do plotting
    if plot_kwargs is None:
        plot_kwargs = dict()
    ax.plot(x, y, label=label, **plot_kwargs)

    # tidy
    ax.set_yscale(yscale)
    ax.set_ylabel('site frequency')
    ax.autoscale(axis='x', tight=True)

    return ax
//  [python]
def binned_statistic(
            x, y, bins=20, ax=None, ptype=None, pargs=None, **kwargs
        ):
    """Bin the data and plot their statistics

    Per default, this calculates the median for each bin. If you need another
    statistic (e.g. mean or std) use the keyword `statistic`.

    Args:
        x: The values which should be binned and plotted on the x-axis.
        y: The values on which the statistics should be applied. The results
            will be plotted on the y-axis.
        bins: Number of bins. Default is 20.
        ax (AxesSubplot, optional): Axes to plot in.
        ptype: Plot type. Can be *scatter* or *boxplot*.
        pargs: Plotting keyword arguments that are allowed for *ptype*.
        **kwargs: Additional key word arguments for
            `scipy.stats.binned_statistic`.

    Returns:
        The plot object.

    Examples:

    .. :code-block:: python


    """

    if ax is None:
        ax = plt.gca()

    if pargs is None:
        pargs = {}

    if ptype is None or ptype == "scatter":
        default = {
            "statistic": "median",
            "bins": bins,
            **kwargs,
        }

        statistics, bin_edges, bin_ind = s.binned_statistic(
            x, values=y, **default
        )
        bin_width = (bin_edges[1] - bin_edges[0])
        bin_centers = bin_edges[1:] - bin_width / 2

        plot = ax.plot(bin_centers, statistics, **pargs)
    elif ptype == "boxplot":
        bin_lefts = np.linspace(x.min(), x.max(), bins)
        bins_indices = np.digitize(x, bin_lefts)

        plot = ax.boxplot(
            [y[bins_indices == i] for i in range(bins)],
            **pargs
        )

        bin_width = (bin_lefts[1] - bin_lefts[0])
        bin_centers = bin_lefts[1:] + bin_width / 2
        ax.set_xticklabels([f"{center:.1f}" for center in bin_centers])
    else:
        raise ValueError(f"Unknown plot type {ptype}!")

    return plot
//  [python]
def binned_statistic(
            x, y, bins=20, ax=None, ptype=None, pargs=None, **kwargs
        ):
    """Bin the data and plot their statistics

    Per default, this calculates the median for each bin. If you need another
    statistic (e.g. mean or std) use the keyword `statistic`.

    Args:
        x: The values which should be binned and plotted on the x-axis.
        y: The values on which the statistics should be applied. The results
            will be plotted on the y-axis.
        bins: Number of bins. Default is 20.
        ax (AxesSubplot, optional): Axes to plot in.
        ptype: Plot type. Can be *scatter* or *boxplot*.
        pargs: Plotting keyword arguments that are allowed for *ptype*.
        **kwargs: Additional key word arguments for
            `scipy.stats.binned_statistic`.

    Returns:
        The plot object.

    Examples:

    .. :code-block:: python


    """

    if ax is None:
        ax = plt.gca()

    if pargs is None:
        pargs = {}

    if ptype is None or ptype == "scatter":
        default = {
            "statistic": "median",
            "bins": bins,
            **kwargs,
        }

        statistics, bin_edges, bin_ind = stats.binned_statistic(
            x, values=y, **default
        )
        bin_width = (bin_edges[1] - bin_edges[0])
        bin_centers = bin_edges[1:] - bin_width / 2

        plot = ax.plot(bin_centers, statistics, **pargs)
    elif ptype == "boxplot":
        bin_lefts = np.linspace(x.min(), x.max(), bins)
        bins_indices = np.digitize(x, bin_lefts)

        plot = ax.boxplot(
            [y[bins_indices == i] for i in range(bins)],
            **pargs
        )

        bin_width = (bin_lefts[1] - bin_lefts[0])
        bin_centers = bin_lefts[1:] + bin_width / 2
        ax.set_xticklabels([f"{center:.1f}" for center in bin_centers])
    else:
        raise ValueError(f"Unknown plot type {ptype}!")

    return plot
//  [python]
def readStatsDafsComputeStandardizationBins(
    statAndDafFileName, nBins=50, pMisPol=0.0
):
    stats = {}
    dafs = []
    pMisPolInv = 1 - pMisPol
    misPolarizedSnps, totalSnps = 0, 0
    with open(statAndDafFileName) as statAndDafFile:
        first = True
        for line in statAndDafFile:
            line = line.strip().split()
            if first:
                first = False
                header = line
                assert header[0] == "daf"
                for i in range(1, len(header)):
                    stats[header[i]] = []
            else:
                totalSnps += 1
                if random.random() >= pMisPolInv:
                    dafs.append(1 - float(line[0]))
                    misPolarizedSnps += 1
                else:
                    dafs.append(float(line[0]))
                for i in range(1, len(line)):
                    stats[header[i]].append(float(line[i]))

    statInfo = {}
    for statName in stats.keys():
        stats[statName] = np.array(stats[statName])
        nonan = ~np.isnan(stats[statName])
        score_nonan = stats[statName][nonan]
        daf_nonan = np.array(dafs)[nonan]
        bins = allel.stats.selection.make_similar_sized_bins(daf_nonan, nBins)
        mean_score, bin_edges, bin_number = scipy.stats.binned_statistic(
            daf_nonan, score_nonan, statistic=np.mean, bins=bins
        )
        std_score, _, _ = scipy.stats.binned_statistic(
            daf_nonan, score_nonan, statistic=np.std, bins=bins
        )
        statInfo[statName] = (mean_score, std_score, bins)
        sys.stderr.write(
            "mispolarized %d of %d (%f%%) "
            "SNPs when standardizing scores in %s\n"
            % (
                misPolarizedSnps,
                totalSnps,
                100 * misPolarizedSnps / float(totalSnps),
                statAndDafFileName,
            )
        )
    return statInfo
//  [python]
def readStatsDafsComputeStandardizationBins(
    statAndDafFileName, nBins=50, pMisPol=0.0
):
    stats = {}
    dafs = []
    pMisPolInv = 1 - pMisPol
    misPolarizedSnps, totalSnps = 0, 0
    with open(statAndDafFileName) as statAndDafFile:
        first = True
        for line in statAndDafFile:
            line = line.strip().split()
            if first:
                first = False
                header = line
                assert header[0] == "daf"
                for i in range(1, len(header)):
                    stats[header[i]] = []
            else:
                totalSnps += 1
                if random.random() >= pMisPolInv:
                    dafs.append(1 - float(line[0]))
                    misPolarizedSnps += 1
                else:
                    dafs.append(float(line[0]))
                for i in range(1, len(line)):
                    stats[header[i]].append(float(line[i]))

    statInfo = {}
    for statName in stats.keys():
        stats[statName] = np.array(stats[statName])
        nonan = ~np.isnan(stats[statName])
        score_nonan = stats[statName][nonan]
        daf_nonan = np.array(dafs)[nonan]
        bins = allel.stats.selection.make_similar_sized_bins(daf_nonan, nBins)
        mean_score, _, _ = scipy.stats.binned_statistic(
            daf_nonan, score_nonan, statistic=np.mean, bins=bins
        )
        std_score, _, _ = scipy.stats.binned_statistic(
            daf_nonan, score_nonan, statistic=np.std, bins=bins
        )
        statInfo[statName] = (mean_score, std_score, bins)
        sys.stderr.write(
            "mispolarized %d of %d (%f%%) "
            "SNPs when standardizing scores in %s\n"
            % (
                misPolarizedSnps,
                totalSnps,
                100 * misPolarizedSnps / float(totalSnps),
                statAndDafFileName,
            )
        )
    return statInfo
//  [python]
def _lnL(self, n_bins, x, y):
        """
                Log likelihood from Hogg 2008
                """
        N, e, _ = s.binned_statistic(x, statistic="sum", values=y, bins=n_bins)
        d = abs(e[0] - e[1])

        if any((N + self.a) < 1):
            return np.nan

        s = np.sum(N + self.a)
        L = np.sum(N * np.log((N + self.a - 1) / (d * (s - 1))))

        return L
//  [python]
def _lnL(self, n_bins, x, y):
        """
                Log likelihood from Hogg 2008
                """
        N, e, _ = stats.binned_statistic(x, statistic="sum", values=y, bins=n_bins)
        d = abs(e[0] - e[1])

        if any((N + self.a) < 1):
            return np.nan

        s = np.sum(N + self.a)
        L = np.sum(N * np.log((N + self.a - 1) / (d * (s - 1))))

        return L
//  [python]
def fit(self, x, y):
        """
                Get the optimal bin number and create those bins,
                then aggregate the labels accordingly and populate
                self.mu
                """
        self.bin_no = self._optimal_bin_no(x, y)
        self.bins = np.histogram(x, bins=self.bin_no)[1]

        pdf = lambda a: np.sum(a) / np.sum(y)
        self.mu = np.histogram(x, bins=self.bins, weights=y)[0] / np.histogram(x, bins=self.bins)[0]
        return self
//  [python]
def fit(self, x, y):
        """
                Get the optimal bin number and create those bins,
                then aggregate the labels accordingly and populate
                self.mu
                """
        self.bin_no = self._optimal_bin_no(x, y)
        self.bins = np.histogram(x, bins=self.bin_no)[1]

        pdf = lambda a: np.sum(a) / np.sum(y)
        self.mu = stats.binned_statistic(x, statistic=pdf, values=y, bins=self.bins)[0]
        return self
//  [python]
def bin_quantity(r, q, r_bin_edge):
    q_bin, _, _ = s.binned_statistic(
        r, q, statistic="mean", bins=r_bin_edge
    )
    q2_bin, _, _ = s.binned_statistic(
        r, q ** 2, statistic="mean", bins=r_bin_edge
    )
    q_sigma_bin = np.sqrt(q2_bin - q_bin ** 2)
    return q_bin, q_sigma_bin
//  [python]
def bin_quantity(r, q, r_bin_edge):
    q_bin, _, _ = stats.binned_statistic(
        r, q, statistic="mean", bins=r_bin_edge
    )
    q2_bin, _, _ = stats.binned_statistic(
        r, q ** 2, statistic="mean", bins=r_bin_edge
    )
    q_sigma_bin = np.sqrt(q2_bin - q_bin ** 2)
    return q_bin, q_sigma_bin
//  [python]
def fig_patch_rmse(patch_file:str, nbins:int=16):
    """ Binned stats for patches

    Args:
        patch_file (str): _description_
        nbins (int, optional): _description_. Defaults to 16.
    """
    lims = (-5,1)

    # Parse
    t_per, p_per = enki_utils.parse_enki_file(patch_file)

    # Outfile
    outfile = f'fig_patch_rmse_t{t_per}_p{p_per}.png'

    # Load
    patch_file = os.path.join(os.getenv("OS_OGCM"),
        'LLC', 'Enki', 'Recon', patch_file)

    f = np.load(patch_file)
    data = f['data']
    data = data.reshape((data.shape[0]*data.shape[1], 
                         data.shape[2]))

    items = f['items']
    tbl = pandas.DataFrame(data, columns=items)

    nbins = 32
    metric = 'log10_std_diff'
    stat = 'median'

    x_metric = 'log10_stdT'
    xvalues, x_lbl = anly_patches.parse_metric(tbl, x_metric)

    values, lbl = anly_patches.parse_metric(tbl, metric)

    good = np.isfinite(xvalues.values)

    # Do it
    eval_stats, x_edge, ibins = scipy.stats.binned_statistic(
        xvalues.values[good], values.values[good], statistic=stat, bins=nbins)

    # Figure
    fig = plt.figure(figsize=(8, 8))
    plt.clf()
    ax = plt.gca()

    # Patches
    plt_x = (x_edge[:-1]+x_edge[1:])/2
    ax.plot(plt_x, eval_stats, 'b', label='Patches')

    # PMC Model
    consts = (0.01, 8.)
    xval = np.linspace(lims[0], lims[1], 1000)
    yval = np.log10((10**xval + consts[0])/consts[1])
    ax.plot(xval, yval, 'r:', label=r'$\log_{10}( \, (\sigma_T + '+f'{consts[0]})/{consts[1]}'+r')$')

    # Axes
    ax.set_xlabel(x_lbl)
    ax.set_ylabel(f'{stat}({lbl})')

    ax.set_xlim(lims[0], lims[1])
    ax.set_ylim(lims[0], lims[1])

    # Labeling
    fsz = 17.
    ax.text(0.05, 0.9, f't={t_per}, p={p_per}',
            transform=ax.transAxes,
              fontsize=fsz, ha='left', color='k')

    ax.legend(loc='lower right', fontsize=fsz-2)

    plotting.set_fontsize(ax, fsz)

    # 1-1
    ax.plot(lims, lims, 'k--')

    # Grid
    ax.grid()

    plt.savefig(outfile, dpi=300)
    plt.close()
    print('Wrote {:s}'.format(outfile))
//  [python]
def convolve_spectrum(en1, spec1, en2, spec2, enbins):
    en = np.outer(en1, en2).flatten()
    conv_points = np.outer(spec1, spec2).flatten()

    spec = np.histogram(en, bins=enbins.bin_edges, weights=conv_points)[0]

    return spec
//  [python]
def convolve_spectrum(en1, spec1, en2, spec2, enbins):
    en = np.outer(en1, en2).flatten()
    conv_points = np.outer(spec1, spec2).flatten()

    spec, _, _ = scipy.stats.binned_statistic(en, conv_points, statistic='sum', bins=enbins.bin_edges)

    return spec
//  [python]
def plot_profile(x,y,name,odir='.',xtitle="",ytitle="Entries",nbins=40,lims=None,
                 stats=True, logy=False, leg=None, text=""):
    if lims==None: lims = (x.min(),x.max())
    median_result = scipy.stats.binned_statistic(x, y, bins=nbins, range=lims, statistic=lambda x: np.quantile(x,0.5))
    lo_result     = scipy.stats.binned_statistic(x, y, bins=nbins, range=lims, statistic=lambda x: np.quantile(x,0.5-0.68/2))
    hi_result     = scipy.stats.binned_statistic(x, y, bins=nbins, range=lims, statistic=lambda x: np.quantile(x,0.5+0.68/2))
    median = np.nan_to_num(median_result.statistic)
    hi = np.nan_to_num(hi_result.statistic)
    lo = np.nan_to_num(lo_result.statistic)
    hie = hi-median
    loe = median-lo
    bin_edges = median_result.bin_edges
    bin_centers = (bin_edges[:-1] + bin_edges[1:])/2.

    plt.figure(figsize=(6,4))
    plt.errorbar(x=bin_centers, y=median, yerr=[loe,hie], linestyle='none', marker='.', label=leg)

    printstr=""
    for i,b in enumerate(bin_centers):
        printstr += "{} {} {} {} \n".format(b, median[i], loe[i], hie[i])

    ax = plt.gca()
    plt.text(0.1, 0.9, name,transform=ax.transAxes)
    if text: plt.text(0.1, 0.82, text.replace('MAX','inf'), transform=ax.transAxes)
    plt.xlabel(xtitle)
    plt.ylabel(ytitle)
    if logy: plt.yscale('log')
    pname = odir+"/"+name+".pdf"
    tname = pname.replace('.pdf','.txt')
    string_to_file(tname, printstr)
    plt.savefig(pname)
    plt.close()
    return bin_centers, median, [loe,hie]
//  [python]
def get_slope(cat_param, cat_type, nTrials, SD):
    if cat_type=='uniform':
        data = sg.experiment(cat_type=cat_type, nTrials=int(nTrials), meas_SD=SD, unif_range=cat_param)
    else:
        data = sg.experiment(cat_type=cat_type, nTrials=int(nTrials), meas_SD=SD, norm_mean=1, norm_SD=cat_param)

    if data['correctness'].all():
        slope=np.nan
    else:
        false_index = ~data['correctness']
        means, edges, bin_number = s.binned_statistic(abs(data['s'][false_index]), data['conf'][false_index], statistic=np.mean, bins=nBins)[:3]

        centers = edges[0:-1]+np.diff(edges)/2

        if np.isnan(means).all():
            slope = np.nan
        else:
            idx = (means==means) & (centers==centers)
            slope = s.linregress(centers[idx], means[idx])[0] # (np.diff(means[~np.isnan(means)][[0,-1]]) # diff of first and last non-nan

    return slope
//  [python]
def get_slope(cat_param, cat_type, nTrials, SD):
    if cat_type=='uniform':
        data = sg.experiment(cat_type=cat_type, nTrials=int(nTrials), meas_SD=SD, unif_range=cat_param)
    else:
        data = sg.experiment(cat_type=cat_type, nTrials=int(nTrials), meas_SD=SD, norm_mean=1, norm_SD=cat_param)

    if data['correctness'].all():
        slope=np.nan
    else:
        false_index = ~data['correctness']
        means, edges = sps.binned_statistic(abs(data['s'][false_index]), data['conf'][false_index], statistic=np.mean, bins=nBins)[:2]

        centers = edges[0:-1]+np.diff(edges)/2

        if np.isnan(means).all():
            slope = np.nan
        else:
            idx = (means==means) & (centers==centers)
            slope = sps.linregress(centers[idx], means[idx])[0] # (np.diff(means[~np.isnan(means)][[0,-1]]) # diff of first and last non-nan

    return slope
//  [python]
def downbin_spec(specHR, lamHR, lamLR, dlam=None):
    """
    Re-bin spectum to lower resolution using :py:obj:`scipy.binned_statistic`
    with ``statistic = 'mean'``. This is a "top-hat" convolution.

    Parameters
    ----------
    specHR : array-like
        Spectrum to be degraded
    lamHR : array-like
        High-res wavelength grid
    lamLR : array-like
        Low-res wavelength grid
    dlam : array-like, optional
        Low-res wavelength width grid

    Returns
    -------
    specLR : :py:obj:`numpy.ndarray`
        Low-res spectrum
    """

    if dlam is None:
        ValueError("Please supply dlam in downbin_spec()")

    # Reverse ordering if wl vector is decreasing with index
    if len(lamLR) > 1:
        if lamHR[0] > lamHR[1]:
            lamHI = np.array(lamHR[::-1])
            spec = np.array(specHR[::-1])
        if lamLR[0] > lamLR[1]:
            lamLO = np.array(lamLR[::-1])
            dlamLO = np.array(dlam[::-1])

    # Calculate bin edges
    LRedges = np.hstack([lamLR - 0.5*dlam, lamLR[-1]+0.5*dlam[-1]])

    # Call scipy.stats.binned_statistic()
    specLR = s.binned_statistic(lamHR, specHR, statistic="mean", bins=LRedges)[0]

    return specLR
//  [python]
def downbin_spec(specHR, lamHR, lamLR, dlam=None):
    """
    Re-bin spectum to lower resolution using :py:obj:`scipy.binned_statistic`
    with ``statistic = 'mean'``. This is a "top-hat" convolution.

    Parameters
    ----------
    specHR : array-like
        Spectrum to be degraded
    lamHR : array-like
        High-res wavelength grid
    lamLR : array-like
        Low-res wavelength grid
    dlam : array-like, optional
        Low-res wavelength width grid

    Returns
    -------
    specLR : :py:obj:`numpy.ndarray`
        Low-res spectrum
    """

    if dlam is None:
        ValueError("Please supply dlam in downbin_spec()")

    # Reverse ordering if wl vector is decreasing with index
    if len(lamLR) > 1:
        if lamHR[0] > lamHR[1]:
            lamHI = np.array(lamHR[::-1])
            spec = np.array(specHR[::-1])
        if lamLR[0] > lamLR[1]:
            lamLO = np.array(lamLR[::-1])
            dlamLO = np.array(dlam[::-1])

    # Calculate bin edges
    LRedges = np.hstack([lamLR - 0.5*dlam, lamLR[-1]+0.5*dlam[-1]])

    # Call scipy.stats.binned_statistic()
    specLR = binned_statistic(lamHR, specHR, statistic="mean", bins=LRedges)[0]

    return specLR
//  [python]
def bin_stats(x, y, stat='median', nbins=128, bin_range=None):
    import scipy.stats as s
    if bin_range is None:
        bin_range = (x.min(), x.max())
    bin_width = (bin_range[1] - bin_range[0]) / nbins
    print("Computing 1-degree bin statistics: %s" % stat)
    bin_stat, bin_edges, bin_number = s.binned_statistic(x, y, statistic=stat, \
            bins=nbins, range=bin_range)
    bin_centers = bin_edges[:-1] + bin_width/2.
    bin_stat = np.ma.masked_invalid(bin_stat)
    return bin_stat, bin_edges, bin_centers
//  [python]
def bin_stats(x, y, stat='median', nbins=128, bin_range=None):
    import scipy.stats
    if bin_range is None:
        bin_range = (x.min(), x.max())
    bin_width = (bin_range[1] - bin_range[0]) / nbins
    print("Computing 1-degree bin statistics: %s" % stat)
    bin_stat, bin_edges, bin_number = scipy.stats.binned_statistic(x, y, statistic=stat, \
            bins=nbins, range=bin_range)
    bin_centers = bin_edges[:-1] + bin_width/2.
    bin_stat = np.ma.masked_invalid(bin_stat)
    return bin_stat, bin_edges, bin_centers
//  [python]
def _lnL(self, bins, x, y):
        """
            Log likelihood from Hogg 2008
            """
        N, e, _ = s.binned_statistic(x, statistic="sum", values=y, bins=bins)

        d = np.diff(e)

        if any((N + self.a) < 1):
            return np.nan

        s = np.sum(N + self.a)
        L = np.sum(N * np.log((N + self.a - 1) / (d * (s - 1))))

        return L
//  [python]
def _lnL(self, bins, x, y):
        """
            Log likelihood from Hogg 2008
            """
        N, e, _ = stats.binned_statistic(x, statistic="sum", values=y, bins=bins)

        d = np.diff(e)

        if any((N + self.a) < 1):
            return np.nan

        s = np.sum(N + self.a)
        L = np.sum(N * np.log((N + self.a - 1) / (d * (s - 1))))

        return L
//  [python]
def fit(self, x, y):
        """
            Get the optimal bin number and create those bins,
            then aggregate the labels accordingly and populate
            self.mu
            """
        self.bins = self.optimal_binning(x, y)

        pdf = lambda a: np.sum(a) / np.sum(y)
        self.mu = np.histogram(x, bins=self.bins, weights=y)[0] / np.histogram(x, bins=self.bins)[0]
        return self
//  [python]
def fit(self, x, y):
        """
            Get the optimal bin number and create those bins,
            then aggregate the labels accordingly and populate
            self.mu
            """
        self.bins = self.optimal_binning(x, y)

        pdf = lambda a: np.sum(a) / np.sum(y)
        self.mu = stats.binned_statistic(x, statistic=pdf, values=y, bins=self.bins)[0]
        return self
//  [python]
def plot_divergence_panel(ax, data, nBins=100, Bayes=True, xlabel=True, ylabel=True, absolute=True, cat=None, fade_below_prop=.3, conf_threshold=None):
    # plot confidence as a function of stimulus and correctness (if conf_threshold=None)
    # or correctness as a function of stimulus and confidence (using arbitrary conf_threshold)

    if cat is None:
        cat_idx = data['C']<100  # just get index of everything
    else:
        cat_idx = data['C']==cat
    
    if Bayes:
        conf = data['conf']
        conflabel = 'mean Bayesian\nconfidence'# $p(C=\hat{C} \mid x)$'
    else:
        conf = abs(data['x'])
        conflabel = 'measurement magnitude $|x|$'
    
    if conf_threshold is None:
        colors = correctness_colors  # red, green
        condlabels = ['incorrect', 'correct']
        indices = [data['correctness']==False, data['correctness']==True]
        y = conf
        y_label = conflabel
    else:
        colors=np.array([[150,150,150], [0,0,0]])/255.0 # gray, black
        condlabels=['high conf.', 'low conf.']
        indices = [conf>conf_threshold, conf<conf_threshold]
        y = data['correctness']
        y_label = 'correctness'

    for j, idx in enumerate(indices):
        index = idx & cat_idx

        if absolute:
            s = abs(data['s'][index])
        else:
            s = data['s'][index]

        means, edges, bin_number = stats.binned_statistic(s, y[index], statistic='mean', bins=nBins)
        
        counts = np.histogram(s, edges)[0]
        max_count_for_white_mixing = max(counts)*fade_below_prop
        white_mix_frac = np.around(np.minimum(1, np.maximum(0, (max_count_for_white_mixing - counts)/max_count_for_white_mixing)), decimals=3)
        faded_colors = mix_white(np.tile(colors[j], (len(counts), 1)), white_mix_frac[:,np.newaxis])

        centers = edges[0:-1]+np.diff(edges)/2
        xy=np.concatenate((centers[:,np.newaxis], means[:, np.newaxis]), axis=1).reshape(-1,1,2)
        segments = np.hstack([xy[:-1], xy[1:]])

        collection = LineCollection(segments, color=faded_colors)

        ax.add_collection(collection)
        
    ax.xaxis.set_major_formatter(FormatStrFormatter('%g'))
    ax.yaxis.set_major_formatter(FormatStrFormatter('%g'))
    sns.despine(ax=ax)
    
    
    if xlabel:
        ax.set_xlabel('stimulus magnitude $|s|$')
    else:
        ax.set_xticklabels('')
        
    if ylabel:
        ax.set_ylabel(y_label)
    else:
        ax.set_yticklabels('')
//  [python]
def plot_divergence_panel(ax, data, nBins=100, Bayes=True, xlabel=True, ylabel=True, absolute=True, cat=None, fade_below_prop=.3, conf_threshold=None):
    # plot confidence as a function of stimulus and correctness (if conf_threshold=None)
    # or correctness as a function of stimulus and confidence (using arbitrary conf_threshold)

    if cat is None:
        cat_idx = data['C']<100  # just get index of everything
    else:
        cat_idx = data['C']==cat
    
    if Bayes:
        conf = data['conf']
        conflabel = 'mean Bayesian\nconfidence'# $p(C=\hat{C} \mid x)$'
    else:
        conf = abs(data['x'])
        conflabel = 'measurement magnitude $|x|$'
    
    if conf_threshold is None:
        colors = correctness_colors  # red, green
        condlabels = ['incorrect', 'correct']
        indices = [data['correctness']==False, data['correctness']==True]
        y = conf
        y_label = conflabel
    else:
        colors=np.array([[150,150,150], [0,0,0]])/255.0 # gray, black
        condlabels=['high conf.', 'low conf.']
        indices = [conf>conf_threshold, conf<conf_threshold]
        y = data['correctness']
        y_label = 'correctness'

    for j, idx in enumerate(indices):
        index = idx & cat_idx

        if absolute:
            s = abs(data['s'][index])
        else:
            s = data['s'][index]

        means, edges = sps.binned_statistic(s, y[index], statistic=np.mean, bins=nBins)[0:2]
        
        counts = np.histogram(s, edges)[0]
        max_count_for_white_mixing = max(counts)*fade_below_prop
        white_mix_frac = np.around(np.minimum(1, np.maximum(0, (max_count_for_white_mixing - counts)/max_count_for_white_mixing)), decimals=3)
        faded_colors = mix_white(np.tile(colors[j], (len(counts), 1)), white_mix_frac[:,np.newaxis])

        centers = edges[0:-1]+np.diff(edges)/2
        xy=np.concatenate((centers[:,np.newaxis], means[:, np.newaxis]), axis=1).reshape(-1,1,2)
        segments = np.hstack([xy[:-1], xy[1:]])

        collection = LineCollection(segments, color=faded_colors)

        ax.add_collection(collection)
        
    ax.xaxis.set_major_formatter(FormatStrFormatter('%g'))
    ax.yaxis.set_major_formatter(FormatStrFormatter('%g'))
    sns.despine(ax=ax)
    
    
    if xlabel:
        ax.set_xlabel('stimulus magnitude $|s|$')
    else:
        ax.set_xticklabels('')
        
    if ylabel:
        ax.set_ylabel(y_label)
    else:
        ax.set_yticklabels('')
//  [python]
def downbin_spec_err(specHR, errHR, lamHR, lamLR, dlam=None):
    """
    Re-bin spectum and errors to lower resolution using :py:obj:`scipy.binned_statistic`.
    This function calculates the noise weighted mean of the points within a bin such that
    :math:`\sqrt{\sum_i \mathrm{SNR}_i}` within each :math:`i` bin is preserved.

    Parameters
    ----------
    specHR : array-like
        Spectrum to be degraded
    errHR : array-like
        One sigma errors of high-res spectrum
    lamHR : array-like
        High-res wavelength grid
    lamLR : array-like
        Low-res wavelength grid
    dlam : array-like, optional
        Low-res wavelength width grid

    Returns
    -------
    specLR : :py:obj:`numpy.ndarray`
        Low-res spectrum
    errLR : :py:obj:`numpy.ndarray`
        Low-res spectrum one sigma errors
    """

    if dlam is None:
        ValueError("Please supply dlam in downbin_spec_err()")

    # Reverse ordering if wl vector is decreasing with index
    if len(lamLR) > 1:
        if lamHR[0] > lamHR[1]:
            lamHI = np.array(lamHR[::-1])
            spec = np.array(specHR[::-1])
        if lamLR[0] > lamLR[1]:
            lamLO = np.array(lamLR[::-1])
            dlamLO = np.array(dlam[::-1])

    # Calculate bin edges
    LRedges = np.hstack([lamLR - 0.5*dlam, lamLR[-1]+0.5*dlam[-1]])

    # Calc the sum of the squares of the noise weighted spectrum in each bin
    wsum = binned_statistic(lamHR, (specHR / errHR)**2., statistic="sum", bins=LRedges)[0]

    # Calc the sum of the squares of just the weights
    w = binned_statistic(lamHR, (1. / errHR)**2., statistic="sum", bins=LRedges)[0]

    # Get weighted spectrum means: normalize noise weighted spectrum by noise weights and take square root
    specLR = np.sqrt(wsum / w)

    # Get one sigma errors: take square root of 1/weights
    errLR = np.sqrt(1. / w)

    return specLR, errLR
//  [python]
def expected_calibration_error(y, p_pred, n_bins=100, n_classes=None, p=1):
    """
    Computes the expected calibration error ECE_p.

    Computes the empirical p-expected calibration error for a vector of confidence
    estimates by binning.

    Parameters
    ----------
    y : array-like
        Ground truth labels.
    p_pred : array-like
        Array of confidence estimates.
    n_bins : int, default=15
        Number of bins of :math:`[\\frac{1}{n_{\\text{classes}},1]` for the confidence estimates.
    n_classes : int default=None
        Number of classes. Estimated from `y` and `y_pred` if not given.
    p : int, default=1
        Power of the calibration error, :math:`1 \\leq p \\leq \\infty`.

    Returns
    -------
    float
        Expected calibration error
    """
    # Check input
    y = sklearn.utils.validation.column_or_1d(y)
    y_pred = np.argmax(p_pred, axis=1)
    y_pred = sklearn.utils.validation.column_or_1d(y_pred)
    if n_classes is None:
        n_classes = np.unique(np.concatenate([y, y_pred])).shape[0]

    # Compute bin means
    bin_range = [1 / n_classes, 1]
    bins = np.linspace(bin_range[0], bin_range[1], n_bins + 1)

    # Find prediction confidence
    p_max = np.max(p_pred, axis=1)

    # Compute empirical accuracy
    empirical_acc = np.zeros(n_bins)
    for i in range(n_bins):
        indices = np.where((p_max >= bins[i]) & (p_max < bins[i+1]))[0]
        if len(indices) > 0:
            empirical_acc[i] = np.mean((y_pred[indices] == y[indices]).astype(int))

    nanindices = np.where(empirical_acc != 0)[0]

    # Perfect calibration
    calibrated_acc = np.linspace(bin_range[0] + bin_range[1] / (2 * n_bins), bin_range[1] - bin_range[1] / (2 * n_bins),
                                 n_bins)

    # Expected calibration error
    weights_ece = np.histogram(p_max, bins)[0][nanindices]
    if p < np.inf:
        ece = np.average(abs(empirical_acc[nanindices] - calibrated_acc[nanindices]) ** p,
                         weights=weights_ece)
    elif np.isinf(p):
        ece = np.max(abs(empirical_acc[nanindices] - calibrated_acc[nanindices]))

    return ece
//  [python]
def expected_calibration_error(y, p_pred, n_bins=100, n_classes=None, p=1):
    """
    Computes the expected calibration error ECE_p.

    Computes the empirical p-expected calibration error for a vector of confidence
    estimates by binning.

    Parameters
    ----------
    y : array-like
        Ground truth labels.
    p_pred : array-like
        Array of confidence estimates.
    n_bins : int, default=15
        Number of bins of :math:`[\\frac{1}{n_{\\text{classes}},1]` for the confidence estimates.
    n_classes : int default=None
        Number of classes. Estimated from `y` and `y_pred` if not given.
    p : int, default=1
        Power of the calibration error, :math:`1 \\leq p \\leq \\infty`.

    Returns
    -------
    float
        Expected calibration error
    """
    # Check input
    y = sklearn.utils.validation.column_or_1d(y)
    y_pred = np.argmax(p_pred, axis=1)
    y_pred = sklearn.utils.validation.column_or_1d(y_pred)
    if n_classes is None:
        n_classes = np.unique(np.concatenate([y, y_pred])).shape[0]

    # Compute bin means
    bin_range = [1 / n_classes, 1]
    bins = np.linspace(bin_range[0], bin_range[1], n_bins + 1)

    # Find prediction confidence
    p_max = np.max(p_pred, axis=1)

    # Compute empirical accuracy
    empirical_acc = scipy.stats.binned_statistic(p_max, (y_pred == y).astype(int),
                                                 bins=n_bins, range=bin_range)[0]
    nanindices = np.where(np.logical_not(np.isnan(empirical_acc)))[0]

    # Perfect calibration
    calibrated_acc = np.linspace(bin_range[0] + bin_range[1] / (2 * n_bins), bin_range[1] - bin_range[1] / (2 * n_bins),
                                 n_bins)

    # Expected calibration error
    weights_ece = np.histogram(p_max, bins)[0][nanindices]
    if p < np.inf:
        ece = np.average(abs(empirical_acc[nanindices] - calibrated_acc[nanindices]) ** p,
                         weights=weights_ece)
    elif np.isinf(p):
        ece = np.max(abs(empirical_acc[nanindices] - calibrated_acc[nanindices]))

    return ece
//  [python]
def get_median(x,y,nbins):
    df_train, cuts1, cuts2 = ewd(x,y, nbins)
    median, edges, binnum = s.binned_statistic(df_train[:,0], df_train[:,1], statistic='median', bins=cuts2)
    return median
//  [python]
def get_median(x,y,nbins):
    df_train, cuts1, cuts2 = ewd(x,y, nbins)
    median, edges, binnum = scipy.stats.binned_statistic(df_train[:,0], df_train[:,1], statistic='median', bins=cuts2)
    return median
//  [python]
def binned_statistics(x,y, nbins):
    """Find the mean, variance and number of counts within the bins described by ewd"""
    
    def CI_high(data, confidence=0.68):
        ## remove the lowest and highest 16% of the values
        
        a = 1.0 * np.array(data)
        n = len(a)
        b = np.sort(data)

        highest = np.int((1-(1-confidence)/2)*len(a))
        high_a = b[highest]
 
        return high_a
    
    def CI_low(data, confidence=0.68):
        ## remove the lowest and highest 16% of the values
        
        a = 1.0 * np.array(data)
        n = len(a)
        b = np.sort(data)
        lowest = np.int(((1-confidence)/2)*len(a))
        low_a = b[lowest]

        return low_a
    
    def get_mean(data):
        return np.mean(data)
    
    df_train, cuts1, cuts2 = ewd(x,y, nbins)
    mean, edges, binnum = s.binned_statistic(df_train[:,0], df_train[:,1], statistic='mean', bins=cuts2)
    var, edges, binnum = s.binned_statistic(df_train[:,0], df_train[:,1], statistic='std', bins=cuts2)
    count, edges, binnum = s.binned_statistic(df_train[:,0], df_train[:,1], statistic='count', bins=cuts2)
    low, edges, binnum = s.binned_statistic(df_train[:,0], df_train[:,1], statistic=CI_low, bins=cuts2)
    high, edges, binnum = s.binned_statistic(df_train[:,0], df_train[:,1], statistic=CI_high, bins=cuts2)
    mean2, edges, binnum = s.binned_statistic(df_train[:,0], df_train[:,1], statistic=get_mean, bins=cuts2)
    
    
    
    return mean, var, count, low, high, mean2
//  [python]
def binned_statistics(x,y, nbins):
    """Find the mean, variance and number of counts within the bins described by ewd"""
    
    def CI_high(data, confidence=0.68):
        ## remove the lowest and highest 16% of the values
        
        a = 1.0 * np.array(data)
        n = len(a)
        b = np.sort(data)

        highest = np.int((1-(1-confidence)/2)*len(a))
        high_a = b[highest]
 
        return high_a
    
    def CI_low(data, confidence=0.68):
        ## remove the lowest and highest 16% of the values
        
        a = 1.0 * np.array(data)
        n = len(a)
        b = np.sort(data)
        lowest = np.int(((1-confidence)/2)*len(a))
        low_a = b[lowest]

        return low_a
    
    def get_mean(data):
        return np.mean(data)
    
    df_train, cuts1, cuts2 = ewd(x,y, nbins)
    mean, edges, binnum = scipy.stats.binned_statistic(df_train[:,0], df_train[:,1], statistic='mean', bins=cuts2)
    var, edges, binnum = scipy.stats.binned_statistic(df_train[:,0], df_train[:,1], statistic='std', bins=cuts2)
    count, edges, binnum = scipy.stats.binned_statistic(df_train[:,0], df_train[:,1], statistic='count', bins=cuts2)
    low, edges, binnum = scipy.stats.binned_statistic(df_train[:,0], df_train[:,1], statistic=CI_low, bins=cuts2)
    high, edges, binnum = scipy.stats.binned_statistic(df_train[:,0], df_train[:,1], statistic=CI_high, bins=cuts2)
    mean2, edges, binnum = scipy.stats.binned_statistic(df_train[:,0], df_train[:,1], statistic=get_mean, bins=cuts2)
    
    
    
    return mean, var, count, low, high, mean2
//  [python]
def _new_1d_binning(self):
        if self._actual_range1[0] is None:
            # deal with the case when actual_range1 == (None, None) and
            # there is no (FOM, etc.) data
            self._actual_range1 = (0, 0)

        self._stats1, self._edges1, _ = \
            s.binned_statistic(self._slow1.data(),
                                   self._fom.data(),
                                   self._statistics(),
                                   self._n_bins1,
                                   self._actual_range1)
        np.nan_to_num(self._stats1, copy=False)

        if self._vfom is not None:
            self._vfom_heat1, _, _ = \
                s.binned_statistic(self._slow1.data(),
                                       self._vfom.data().T,
                                       self._statistics(),
                                       self._n_bins1,
                                       self._actual_range1)
            np.nan_to_num(self._vfom_heat1, copy=False)

        self._counts1, _, _ = \
            s.binned_statistic(self._slow1.data(),
                                   self._fom.data(),
                                   'count',
                                   self._n_bins1,
                                   self._actual_range1)
        np.nan_to_num(self._counts1, copy=False)
//  [python]
def _new_1d_binning(self):
        if self._actual_range1[0] is None:
            # deal with the case when actual_range1 == (None, None) and
            # there is no (FOM, etc.) data
            self._actual_range1 = (0, 0)

        self._stats1, self._edges1, _ = \
            stats.binned_statistic(self._slow1.data(),
                                   self._fom.data(),
                                   self._statistics(),
                                   self._n_bins1,
                                   self._actual_range1)
        np.nan_to_num(self._stats1, copy=False)

        if self._vfom is not None:
            self._vfom_heat1, _, _ = \
                stats.binned_statistic(self._slow1.data(),
                                       self._vfom.data().T,
                                       self._statistics(),
                                       self._n_bins1,
                                       self._actual_range1)
            np.nan_to_num(self._vfom_heat1, copy=False)

        self._counts1, _, _ = \
            stats.binned_statistic(self._slow1.data(),
                                   self._fom.data(),
                                   'count',
                                   self._n_bins1,
                                   self._actual_range1)
        np.nan_to_num(self._counts1, copy=False)
//  [python]
def binned(
    coll=None,
    key_diag=None,
    key_cam=None,
    # data to be binned
    data=None,
    statistic=None,
    # binning dimension
    bins0=None,
    bins1=None,
    bin_data0=None,
    bin_data1=None,
    # store
    store=None,
    # plotting
    plot=None,
    ax=None,
):

    # ------------
    # check

    (
        key_diag, key_cam, spectro, is2d,
        ddata, dref, units,
        dbins0, dbins1,
        statistic, store, plot,
    ) = _binned_check(**locals())

    # ----------
    # compute

    dout = {}
    for ii, kcam in enumerate(ddata.keys()):

        dati = ddata[kcam]
        bind = dbins0['ddata'][kcam]

        # binning index
        if dbins1 is None:

            for kcam in key_cam:
                dout[kcam] = coll.binning(
                    keys=data,
                    ref_key=dbins0['key'][''],
                    bins=dbins0['dedges'][kcam],
                    verb=False,
                    store=False,
                    returnas=None,
                    key_store=None,
                )

                import pdb; pdb.et_trace()      # DB

            iok = np.isfinite(dati) & np.isfinite(bind)
            out = s.binned_statistic(
                bind[iok],
                dati[iok],
                statistic=statistic,
                bins=dbins0['dedges'][kcam],
            )

        else:
            out = s.binned_statistic_2d(
                bind[iok],
                dati[iok],
                statistic=statistic,
                bins=[dbins0['dedges'][kcam], dbins1['dedges'][kcam]],
            )

        import pdb; pdb.et_trace()      # DB

        # fill dict
        # dout[kcam] = {
        #     'data': out,
        #     'units': units,
        #     'bins': 0.5*(dbin_edges[kcam][1:] + dbin_edges[kcam][:-1]),
        #     'bin_units': bin_dd[kcam]['units'],
        # }

    # ----------
    # plot

    if plot is True:
        ax = _spectro2d_binned_plot(
            key_diag=key_diag,
            # data to be binned
            data=None,
            # binning dimension
            bin_data=None,
            dout=dout,
            # plotting
            ax=ax,
        )

    # ----------
    # return

    if plot is True:
        return dout, ax
    else:
        return dout