jiahuimbzuai/CoQuIR huggingface dataset:
//  [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 = scpstats.binned_statistic(
                bind[iok],
                dati[iok],
                statistic=statistic,
                bins=dbins0['dedges'][kcam],
            )

        else:
            out = scpstats.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
//  [python]
def calc_corr_over_elec_distance(acq_data, acq_ch, elec_pos, bins=20, method='spearman', exclude_zero_dist=True):
    '''
    Calculates mean absolute correlation between acq_data across channels with the same distance between them.
    
    Args:
        acq_data (nt, nch): acquisition data indexed by acq_ch
        acq_ch (nelec): 1-indexed list of acquisition channels that are connected to electrodes
        elec_pos (nelec, 2): x, y position of each electrode
        bins (int or array): input into scipy.stats.binned_statistic, can be a number or a set of bins
        method (str, optional): correlation method to use ('pearson' or 'spearman')
        exclude_zero_dist (bool, optional): whether to exclude distances that are equal to zero. default True
        
    Returns:
        tuple: tuple containing:
            | **dist (nbins):** electrode distance at each bin
            | **corr (nbins):** correlation at each bin

    '''
    dist = utils.calc_euclid_dist_mat(elec_pos)
    if method == 'spearman':
        c, _ = stats.spearmanr(acq_data, axis=0)
    elif method == 'pearson':
        c = np.corrcoef(acq_data.T)
    else:
        raise ValueError(f"Unknown correlation method {method}")
    
    c_ = c[np.ix_(acq_ch-1, acq_ch-1)] # note use of open mesh to get the right logical index
    
    if exclude_zero_dist:
        zero_dist = dist == 0
        dist = dist[~zero_dist]
        c_ = c_[~zero_dist]
        
    corr, edges, _ = stats.binned_statistic(dist.flatten(), np.abs(c_.flatten()), statistic='mean', bins=bins)
    dist = (edges[:-1] + edges[1:]) / 2

    return dist, corr
//  [python]
def get_normalized_dispersion(
    mat_mean: np.ndarray, mat_var: np.ndarray, nbins: int = 20
) -> np.ndarray:
    """Calculates the normalized dispersion.

    The dispersion is calculated for each feature
    and then normalized to see how its dispersion compares to samples that had a
    similar mean value.
    """
    # See equation in https://academic.oup.com/nar/article/40/10/4288/2411520
    # If a negative binomial is parameterized with mean m, and variance = m + d * m^2
    # then this d = dispersion as calculated below
    mat_disp: np.ndarray = (mat_var - mat_mean) / np.square(mat_mean)

    quantiles = np.percentile(mat_mean, np.arange(0, 100, 100 // nbins))
    quantiles = np.append(quantiles, mat_mean.max())

    # merge bins with no difference in value
    quantiles = np.unique(quantiles)

    if len(quantiles) <= 1:
        # pathological case: the means are all identical. just return raw dispersion.
        return mat_disp

    # calc median dispersion per bin
    disp_meds = np.zeros(len(quantiles) -1)
    for i in range(len(quantiles)-1):
        indices = np.where((mat_mean >= quantiles[i]) & (mat_mean < quantiles[i+1]))[0]
        if len(indices)>0:
            disp_meds[i] = np.median(mat_disp[indices])

    disp_bins = np.digitize(mat_mean, quantiles)

    # calc median absolute deviation of dispersion per bin
    disp_meds_arr = disp_meds[disp_bins - 1]  # 0th bin is empty since our quantiles start from 0
    disp_abs_dev = abs(mat_disp - disp_meds_arr)
    disp_mads = np.zeros(len(quantiles) -1)
    for i in range(len(quantiles)-1):
        indices = np.where((mat_mean >= quantiles[i]) & (mat_mean < quantiles[i+1]))[0]
        if len(indices)>0:
            disp_mads[i] = np.median(disp_abs_dev[indices])

    # calculate normalized dispersion
    disp_mads_arr = disp_mads[disp_bins - 1]
    disp_norm = (mat_disp - disp_meds_arr) / disp_mads_arr
    return disp_norm
//  [python]
def get_normalized_dispersion(
    mat_mean: np.ndarray, mat_var: np.ndarray, nbins: int = 20
) -> np.ndarray:
    """Calculates the normalized dispersion.

    The dispersion is calculated for each feature
    and then normalized to see how its dispersion compares to samples that had a
    similar mean value.
    """
    # See equation in https://academic.oup.com/nar/article/40/10/4288/2411520
    # If a negative binomial is parameterized with mean m, and variance = m + d * m^2
    # then this d = dispersion as calculated below
    mat_disp: np.ndarray = (mat_var - mat_mean) / np.square(mat_mean)

    quantiles = np.percentile(mat_mean, np.arange(0, 100, 100 // nbins))
    quantiles = np.append(quantiles, mat_mean.max())

    # merge bins with no difference in value
    quantiles = np.unique(quantiles)

    if len(quantiles) <= 1:
        # pathological case: the means are all identical. just return raw dispersion.
        return mat_disp

    # calc median dispersion per bin
    (disp_meds, _, disp_bins) = scipy.stats.binned_statistic(
        mat_mean, mat_disp, statistic="median", bins=quantiles
    )

    # calc median absolute deviation of dispersion per bin
    disp_meds_arr = disp_meds[disp_bins - 1]  # 0th bin is empty since our quantiles start from 0
    disp_abs_dev = abs(mat_disp - disp_meds_arr)
    (disp_mads, _, disp_bins) = scipy.stats.binned_statistic(
        mat_mean, disp_abs_dev, statistic="median", bins=quantiles
    )

    # calculate normalized dispersion
    disp_mads_arr = disp_mads[disp_bins - 1]
    disp_norm = (mat_disp - disp_meds_arr) / disp_mads_arr
    return disp_norm
//  [python]
def sph_histogram(self, distance, maxbin, minbin, step):

        """Function to create a histogram of the distances
 

       Parameters
       -----------

        distance: np.array
            distances to be binned

        maxbin : int
            largest bin value

        minbin : int
            smalles bin value

        step : int
            bin size



        

       Returns
       --------
 
       all_histograms: np.array
            array with the distances values binned

        """

        import numpy as np
        import scipy.stats as stats 
    
        x=len(np.arange(minbin,maxbin,step)[:-1])
    

        all_histograms=np.zeros(x)



        binned_space=np.arange(minbin,maxbin,step)
    
        binned_space=np.array(binned_space)


        l = [item for item in distance]

        flat_list = [item for sublist in l for item in sublist]

    
        all_histograms+=np.histogram(flat_list, bins=binned_space)[0]
   
    
        all_histograms=np.array(all_histograms)


        return all_histograms
//  [python]
def sph_histogram(self, distance, maxbin, minbin, step):

        """Function to create a histogram of the distances
 

       Parameters
       -----------

        distance: np.array
            distances to be binned

        maxbin : int
            largest bin value

        minbin : int
            smalles bin value

        step : int
            bin size



        

       Returns
       --------
 
       all_histograms: np.array
            array with the distances values binned

        """

        import numpy as np
        import scipy.stats as stats 
    
        x=len(np.arange(minbin,maxbin,step)[:-1])
    

        all_histograms=np.zeros(x)



        binned_space=np.arange(minbin,maxbin,step)
    
        binned_space=np.array(binned_space)


        l = [item for item in distance]

        flat_list = [item for sublist in l for item in sublist]

    
        all_histograms+=stats.binned_statistic(flat_list, flat_list, 'count', bins=binned_space).statistic
   
    
        all_histograms=np.array(all_histograms)


        return all_histograms
//  [python]
def compute_weighted_bin_sum(xvals, yvals, xbins, weights):
    """Add `yvals * weights` in `xbins`.

    Parameters
    ----------
    xvals : array_like
        Values to be binned
    yvals : array_like
        Values to compute statistics on
    xbins: array_like
        Bin edges to sort into
    weights: array_like, None, optional
        Weights for sum.

    Returns
    -------
    numpy.ndarray
        Sum of `yvals * weights` in `xbins`.
    """
    return binned_statistic(xvals, yvals * weights, statistic="sum", bins=xbins)[0]
//  [python]
def fetch_bw(f, chrom_idx_lst, window, bins=-1):
    print(f, window)
    res = []
    for row in chrom_idx_lst:
        # print(row)
        if (row[0] == 'chr8') & (row[1] > 146000000):
            # print(row)
            y = bbi.fetch(str(f), row[0], row[1], row[1]+window)
            z, _, _ = s.binned_statistic(y, y, bins=bins)
            res.append(z)

        else:
            res.append(bbi.fetch(str(f), row[0], row[1], row[1]+window, bins=bins))

    x = np.array(res)
    # x = np.array([bbi.fetch(str(f), row[0], row[1], row[1]+window, bins=bins) for row in chrom_idx_lst])
    # return x
    return (np.around(x, decimals=2)*100)
//  [python]
def fetch_bw(f, chrom_idx_lst, window, bins=-1):
    print(f, window)
    res = []
    for row in chrom_idx_lst:
        # print(row)
        if (row[0] == 'chr8') & (row[1] > 146000000):
            # print(row)
            y = bbi.fetch(str(f), row[0], row[1], row[1]+window)
            z, _ , _ = scipy.stats.binned_statistic(y, y, bins=bins)
            res.append(z)

        else:
            res.append(bbi.fetch(str(f), row[0], row[1], row[1]+window, bins=bins))

    x = np.array(res)
    # x = np.array([bbi.fetch(str(f), row[0], row[1], row[1]+window, bins=bins) for row in chrom_idx_lst])
    # return x
    return (np.around(x, decimals=2)*100)
//  [python]
def fill(self, data, weights=None):
        """Fills the histogram with given data and weights.

        Parameters
        ----------
        data : numpy.ndarray, pandas.Series
            Sequence which is filled into the histogram.
        weights : numpy.ndarray, pandas.Series, optional
            Sequence which is used as event weights. If no weights are
            given, each event is weighted with 1.0. Default is None.
        """
        if weights is None:
            weights = np.ones_like(data)
        if isinstance(weights, list):
            weights = np.array(weights)

        self._bin_counts += np.histogram(data, bins=self._bin_edges, range=self._limits, weights=weights)[0]

        self._bin_errors_sq += np.histogram(data, bins=self._bin_edges, range=self._limits, weights=weights**2)[0]
//  [python]
def fill(self, data, weights=None):
        """Fills the histogram with given data and weights.

        Parameters
        ----------
        data : numpy.ndarray, pandas.Series
            Sequence which is filled into the histogram.
        weights : numpy.ndarray, pandas.Series, optional
            Sequence which is used as event weights. If no weights are
            given, each event is weighted with 1.0. Default is None.
        """
        if weights is None:
            weights = np.ones_like(data)
        if isinstance(weights, list):
            weights = np.array(weights)

        self._bin_counts += scipy.stats.binned_statistic(
            data,
            weights,
            statistic='sum',
            bins=self._bin_edges,
            range=self._limits
        )[0]

        self._bin_errors_sq += scipy.stats.binned_statistic(
            data,
            weights ** 2,
            statistic='sum',
            bins=self._bin_edges,
            range=self._limits
        )[0]
//  [python]
def fill(self, data, weights=None):
        """Fills the histogram with given data. If no weights are
        given, each data point is weighted with 1.0.

        Parameters
        ----------
        data : np.array
            Data to be filled in the histogram.
        weights: np.array, optional
            Weights for each entry in the histogram.
        """
        if isinstance(data, list):
            data = np.array(data)
        if weights is not None and isinstance(weights, list):
            weights = np.array(weights)

        if weights is None:
            weights = np.ones_like(data)

        if len(data) != len(weights):
            raise ValueError(
                "Shape of data array does not match weight array.")
        self._bin_counts += np.histogram(data, bins=self._bin_edges, weights=weights)[0]

        self._bin_entries += np.histogram(data, bins=self._bin_edges, weights=weights)[0]

        self._bin_errors_sq += np.histogram(data, bins=self._bin_edges, weights=weights**2)[0]
//  [python]
def fill(self, data, weights=None):
        """Fills the histogram with given data. If no weights are
        given, each data point is weighted with 1.0.

        Parameters
        ----------
        data : np.array
            Data to be filled in the histogram.
        weights: np.array, optional
            Weights for each entry in the histogram.
        """
        if isinstance(data, list):
            data = np.array(data)
        if weights is not None and isinstance(weights, list):
            weights = np.array(weights)

        if weights is None:
            weights = np.ones_like(data)

        if len(data) != len(weights):
            raise ValueError(
                "Shape of data array does not match weight array.")
        self._bin_counts += scipy.stats.binned_statistic(
            x=data, values=weights, statistic="sum", bins=self._bin_edges
        )[0]

        self._bin_entries += scipy.stats.binned_statistic(
            x=data, values=weights, statistic="count", bins=self._bin_edges
        )[0]

        self._bin_errors_sq += scipy.stats.binned_statistic(
            x=data, values=weights ** 2, statistic="sum", bins=self._bin_edges
        )[0]
//  [python]
def anly_patches(patch_file:str, nbins:int=32, model:str='std'):
    """ Analyze the patches

    Args:
        patch_file (str): _description_
        nbins (int, optional): _description_. Defaults to 32.
        nfit (int, optional): Number of parameters for the fit. Defaults to 2.

    Returns:
        _type_: _description_
    """

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

    print(f'Loading: {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)

    metric = 'log10_std_diff'
    stat = 'median'

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

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

    good1 = np.isfinite(xvalues.values)
    good2 = np.isfinite(values.values)
    good = good1 & good2

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

    # Fit
    x = (x_edge[:-1]+x_edge[1:])/2
    gd_eval = np.isfinite(eval_stats)

    gd_x = 10**x[gd_eval]
    gd_y = 10**eval_stats[gd_eval]

    if model == 'std':
        fit_model = two_param_model 
        p0=[0.01, 10.]
    elif model == 'denom':
        fit_model = denom_model
        p0 = None
    else:
        raise ValueError(f'Unknown model: {model}')
    popt, pcov = scipy.optimize.curve_fit(
        fit_model, gd_x, gd_y, p0=p0, sigma=0.1*gd_y)
    return x_edge, eval_stats, stat, x_lbl, lbl, popt, tbl
//  [python]
def dm_density(
    data: SWIFTDataset, ax: plt.Axes, radial_bins: np.array, center: np.array
) -> None:

    r_dm = np.sqrt(np.sum((data.dark_matter.coordinates - center) ** 2, axis=1))

    mass, _, _ = s.binned_statistic(
        x=r_dm,
        values=data.dark_matter.masses,
        statistic="sum",
        bins=radial_bins.to(r_dm.units),
    )

    # binned_statistic strips unit info :/
    mass *= data.dark_matter.masses.units

    density = (mass / bin_volumes(radial_bins)).to(dm_density_units)

    centers = bin_centers(radial_bins).to(radial_bin_units)

    ax.plot(centers, density)

    ax.loglog()

    ax.set_ylabel(f"Dark Matter Density $\\left[{density.units.latex_repr}\\right]$")
    ax.set_xlabel(f"Radius [${centers.units.latex_repr}$]")

    ax.set_xlim(centers[0], centers[-1])

    return
//  [python]
def dm_density(
    data: SWIFTDataset, ax: plt.Axes, radial_bins: np.array, center: np.array
) -> None:

    r_dm = np.sqrt(np.sum((data.dark_matter.coordinates - center) ** 2, axis=1))

    mass, _, _ = stat.binned_statistic(
        x=r_dm,
        values=data.dark_matter.masses,
        statistic="sum",
        bins=radial_bins.to(r_dm.units),
    )

    # binned_statistic strips unit info :/
    mass *= data.dark_matter.masses.units

    density = (mass / bin_volumes(radial_bins)).to(dm_density_units)

    centers = bin_centers(radial_bins).to(radial_bin_units)

    ax.plot(centers, density)

    ax.loglog()

    ax.set_ylabel(f"Dark Matter Density $\\left[{density.units.latex_repr}\\right]$")
    ax.set_xlabel(f"Radius [${centers.units.latex_repr}$]")

    ax.set_xlim(centers[0], centers[-1])

    return
//  [python]
def gas_density(
    data: SWIFTDataset, ax: plt.Axes, radial_bins: np.array, center: np.array
) -> None:

    r_gas = np.sqrt(np.sum((data.gas.coordinates - center) ** 2, axis=1))

    mass, _, _ = s.binned_statistic(
        x=r_gas,
        values=data.gas.masses,
        statistic="sum",
        bins=radial_bins.to(r_gas.units),
    )

    # binned_statistic strips unit info :/
    mass *= data.gas.masses.units

    density = (mass / bin_volumes(radial_bins)).to(gas_density_units)

    centers = bin_centers(radial_bins).to(radial_bin_units)

    ax.plot(centers, density, label="SWIFT")

    ax.loglog()

    ax.set_ylabel(f"Gas Density $\\left[{density.units.latex_repr}\\right]$")
    ax.set_xlabel(f"Radius [${centers.units.latex_repr}$]")

    ax.set_xlim(centers[0], centers[-1])

    for scheme_name in extra_data_schemes:
        path = f"{extra_data_directory}/rho_gas/{scheme_name}.csv"
        R, rho = np.loadtxt(path, delimiter=",").T

        # Correct h-factors
        R /= h
        rho *= h * h
        ax.plot(R, rho, linestyle="dashed", lw=1, label=scheme_name)

    ax.legend()

    return
//  [python]
def gas_density(
    data: SWIFTDataset, ax: plt.Axes, radial_bins: np.array, center: np.array
) -> None:

    r_gas = np.sqrt(np.sum((data.gas.coordinates - center) ** 2, axis=1))

    mass, _, _ = stat.binned_statistic(
        x=r_gas,
        values=data.gas.masses,
        statistic="sum",
        bins=radial_bins.to(r_gas.units),
    )

    # binned_statistic strips unit info :/
    mass *= data.gas.masses.units

    density = (mass / bin_volumes(radial_bins)).to(gas_density_units)

    centers = bin_centers(radial_bins).to(radial_bin_units)

    ax.plot(centers, density, label="SWIFT")

    ax.loglog()

    ax.set_ylabel(f"Gas Density $\\left[{density.units.latex_repr}\\right]$")
    ax.set_xlabel(f"Radius [${centers.units.latex_repr}$]")

    ax.set_xlim(centers[0], centers[-1])

    for scheme_name in extra_data_schemes:
        path = f"{extra_data_directory}/rho_gas/{scheme_name}.csv"
        R, rho = np.loadtxt(path, delimiter=",").T

        # Correct h-factors
        R /= h
        rho *= h * h
        ax.plot(R, rho, linestyle="dashed", lw=1, label=scheme_name)

    ax.legend()

    return
//  [python]
def bin_plot(x, y):
    x = np.array(x)
    bins = np.percentile(x, np.arange(0,100))
    s, edges, _ = stats.binned_statistic(x, y, statistic='mean', bins=bins)
    return sbn.scatterplot(edges[:-1] + np.diff(edges) / 2, s)
//  [python]
def bin_plot(x, y):
    x = np.array(x)
    bins = np.percentile(x, np.arange(0,100))
    s, edges, _ = scipy.stats.binned_statistic(x, y, statistic='mean', bins=bins)
    return sbn.scatterplot(edges[:-1] + np.diff(edges) / 2, s)
//  [python]
def temperature(
    data: SWIFTDataset, ax: plt.Axes, radial_bins: np.array, center: np.array
) -> None:

    r_gas = np.sqrt(np.sum((data.gas.coordinates - center) ** 2, axis=1))
    temperatures = get_gas_temperatures(data)

    temp, _, _ = s.binned_statistic(
        x=r_gas, values=temperatures, statistic="mean", bins=radial_bins.to(r_gas.units)
    )

    # binned_statistic strips unit info :/
    temp *= temperatures.units

    kT = (temp * unyt.kb).to(temperature_units)

    centers = bin_centers(radial_bins).to(radial_bin_units)

    ax.plot(centers, kT)

    ax.semilogx()

    ax.set_ylabel(f"Gas Temperature (kT) [${kT.units.latex_repr}$]")
    ax.set_xlabel(f"Radius [${centers.units.latex_repr}$]")

    for scheme_name in extra_data_schemes:
        path = f"{extra_data_directory}/kT/{scheme_name}.csv"
        R, kT = np.loadtxt(path, delimiter=",").T

        # Correct h-factors
        R /= h
        ax.plot(R, kT, linestyle="dashed", lw=1, label=scheme_name)

    ax.set_xlim(centers[0], centers[-1])

    return
//  [python]
def temperature(
    data: SWIFTDataset, ax: plt.Axes, radial_bins: np.array, center: np.array
) -> None:

    r_gas = np.sqrt(np.sum((data.gas.coordinates - center) ** 2, axis=1))
    temperatures = get_gas_temperatures(data)

    temp, _, _ = stat.binned_statistic(
        x=r_gas, values=temperatures, statistic="mean", bins=radial_bins.to(r_gas.units)
    )

    # binned_statistic strips unit info :/
    temp *= temperatures.units

    kT = (temp * unyt.kb).to(temperature_units)

    centers = bin_centers(radial_bins).to(radial_bin_units)

    ax.plot(centers, kT)

    ax.semilogx()

    ax.set_ylabel(f"Gas Temperature (kT) [${kT.units.latex_repr}$]")
    ax.set_xlabel(f"Radius [${centers.units.latex_repr}$]")

    for scheme_name in extra_data_schemes:
        path = f"{extra_data_directory}/kT/{scheme_name}.csv"
        R, kT = np.loadtxt(path, delimiter=",").T

        # Correct h-factors
        R /= h
        ax.plot(R, kT, linestyle="dashed", lw=1, label=scheme_name)

    ax.set_xlim(centers[0], centers[-1])

    return
//  [python]
def entropy(
    data: SWIFTDataset, ax: plt.Axes, radial_bins: np.array, center: np.array
) -> None:

    r_gas = np.sqrt(np.sum((data.gas.coordinates - center) ** 2, axis=1))

    # Need temperatures _and_ densities here.
    temperatures = get_gas_temperatures(data)

    temp, edges, _ = s.binned_statistic(
        x=r_gas, values=temperatures, statistic="mean", bins=radial_bins.to(r_gas.units)
    )

    # binned_statistic strips unit info :/
    temp *= temperatures.units

    mass, _, _ = s.binned_statistic(
        x=r_gas,
        values=data.gas.masses,
        statistic="sum",
        bins=radial_bins.to(r_gas.units),
    )

    # binned_statistic strips unit info :/
    mass *= data.gas.masses.units

    density = mass / bin_volumes(radial_bins)

    # Really need the number density
    density = density / unyt.mh

    entropy = unyt.kb * temp / (density ** (2.0 / 3.0))
    entropy.convert_to_units(entropy_units)

    centers = bin_centers(radial_bins).to(radial_bin_units)

    ax.plot(centers, entropy)

    ax.loglog()

    ax.set_ylabel(f"Gas Entropy (T $n_e^{{2/3}}$) [${entropy.units.latex_repr}$]")
    ax.set_xlabel(f"Radius [${centers.units.latex_repr}$]")

    for scheme_name in extra_data_schemes:
        path = f"{extra_data_directory}/S/{scheme_name}.csv"
        R, S = np.loadtxt(path, delimiter=",").T

        # Correct h-factors
        R /= h
        S /= h * h
        ax.plot(R, S, linestyle="dashed", lw=1, label=scheme_name)

    ax.set_xlim(centers[0], centers[-1])

    return
//  [python]
def entropy(
    data: SWIFTDataset, ax: plt.Axes, radial_bins: np.array, center: np.array
) -> None:

    r_gas = np.sqrt(np.sum((data.gas.coordinates - center) ** 2, axis=1))

    # Need temperatures _and_ densities here.
    temperatures = get_gas_temperatures(data)

    temp, _, _ = stat.binned_statistic(
        x=r_gas, values=temperatures, statistic="mean", bins=radial_bins.to(r_gas.units)
    )

    # binned_statistic strips unit info :/
    temp *= temperatures.units

    mass, _, _ = stat.binned_statistic(
        x=r_gas,
        values=data.gas.masses,
        statistic="sum",
        bins=radial_bins.to(r_gas.units),
    )

    # binned_statistic strips unit info :/
    mass *= data.gas.masses.units

    density = mass / bin_volumes(radial_bins)

    # Really need the number density
    density = density / unyt.mh

    entropy = unyt.kb * temp / (density ** (2.0 / 3.0))
    entropy.convert_to_units(entropy_units)

    centers = bin_centers(radial_bins).to(radial_bin_units)

    ax.plot(centers, entropy)

    ax.loglog()

    ax.set_ylabel(f"Gas Entropy (T $n_e^{{2/3}}$) [${entropy.units.latex_repr}$]")
    ax.set_xlabel(f"Radius [${centers.units.latex_repr}$]")

    for scheme_name in extra_data_schemes:
        path = f"{extra_data_directory}/S/{scheme_name}.csv"
        R, S = np.loadtxt(path, delimiter=",").T

        # Correct h-factors
        R /= h
        S /= h * h
        ax.plot(R, S, linestyle="dashed", lw=1, label=scheme_name)

    ax.set_xlim(centers[0], centers[-1])

    return
//  [python]
def bin_quantity(x, q, nbin):
    # we let binned_statistic figure out the bins and then reuse them for the
    # standard deviation mean
    qbin, bins, _ = s.binned_statistic(x, q, statistic="mean", bins=nbin)
    q2bin, _, _ = s.binned_statistic(x, q ** 2, statistic="mean", bins=bins)
    qsigma = np.sqrt(q2bin - qbin ** 2)
    return 0.5 * (bins[1:] + bins[:-1]), qbin, qsigma
//  [python]
def bin_quantity(x, q, nbin):
    # we let binned_statistic figure out the bins and then reuse them for the
    # standard deviation mean
    qbin, bins, _ = stats.binned_statistic(x, q, statistic="mean", bins=nbin)
    q2bin, _, _ = stats.binned_statistic(x, q ** 2, statistic="mean", bins=bins)
    qsigma = np.sqrt(q2bin - qbin ** 2)
    return 0.5 * (bins[1:] + bins[:-1]), qbin, qsigma
//  [python]
def get_normalized_dispersion(mat_mean, mat_var, nbins=20):
    """ Calculates the normalized dispersion.  The dispersion is calculated for each feature
        and then normalized to see how its dispersion compares to samples that had a
        similar mean value.
    """
    # See equation in https://academic.oup.com/nar/article/40/10/4288/2411520
    # If a negative binomial is parameterized with mean m, and variance = m + d * m^2
    # then this d = dispersion as calculated below
    mat_disp = (mat_var - mat_mean) / np.square(mat_mean)

    quantiles = np.percentile(mat_mean, np.arange(0, 100, 100 / nbins))
    quantiles = np.append(quantiles, mat_mean.max())

    # merge bins with no difference in value
    quantiles = np.unique(quantiles)

    if len(quantiles) <= 1:
        # pathological case: the means are all identical. just return raw dispersion.
        return mat_disp

    # calc median dispersion per bin
    disp_meds = np.zeros(len(quantiles)-1)
    for i in range(len(quantiles)-1):
        mask = (mat_mean >= quantiles[i]) & (mat_mean < quantiles[i+1])
        disp_meds[i] = np.median(mat_disp[mask])

    disp_bins = np.digitize(mat_mean, quantiles)

    # calc median absolute deviation of dispersion per bin
    disp_meds_arr = disp_meds[disp_bins-1] # 0th bin is empty since our quantiles start from 0
    disp_abs_dev = abs(mat_disp - disp_meds_arr)
    disp_mads = np.zeros(len(quantiles)-1)
    for i in range(len(quantiles)-1):
        mask = (mat_mean >= quantiles[i]) & (mat_mean < quantiles[i+1])
        disp_mads[i] = np.median(disp_abs_dev[mask])

    # calculate normalized dispersion
    disp_mads_arr = disp_mads[disp_bins-1]
    disp_norm = (mat_disp - disp_meds_arr) / disp_mads_arr
    return disp_norm
//  [python]
def get_normalized_dispersion(mat_mean, mat_var, nbins=20):
    """ Calculates the normalized dispersion.  The dispersion is calculated for each feature
        and then normalized to see how its dispersion compares to samples that had a
        similar mean value.
    """
    # See equation in https://academic.oup.com/nar/article/40/10/4288/2411520
    # If a negative binomial is parameterized with mean m, and variance = m + d * m^2
    # then this d = dispersion as calculated below
    mat_disp = (mat_var - mat_mean) / np.square(mat_mean)

    quantiles = np.percentile(mat_mean, np.arange(0, 100, 100 / nbins))
    quantiles = np.append(quantiles, mat_mean.max())

    # merge bins with no difference in value
    quantiles = np.unique(quantiles)

    if len(quantiles) <= 1:
        # pathological case: the means are all identical. just return raw dispersion.
        return mat_disp

    # calc median dispersion per bin
    (disp_meds, _, disp_bins) = scipy.stats.binned_statistic(mat_mean, mat_disp, statistic='median', bins=quantiles)

    # calc median absolute deviation of dispersion per bin
    disp_meds_arr = disp_meds[disp_bins-1] # 0th bin is empty since our quantiles start from 0
    disp_abs_dev = abs(mat_disp - disp_meds_arr)
    (disp_mads, _, disp_bins) = scipy.stats.binned_statistic(mat_mean, disp_abs_dev, statistic='median', bins=quantiles)

    # calculate normalized dispersion
    disp_mads_arr = disp_mads[disp_bins-1]
    disp_norm = (mat_disp - disp_meds_arr) / disp_mads_arr
    return disp_norm
//  [python]
def get_psd_of_depth_list(depth_list):
    avg_psd = []
    for depth in depth_list:
        psd = np.fft.fft2((depth - np.mean(depth)) / np.std(depth))  # [480, 640]
        psd = np.abs(psd) ** 2
        psd = psd.flatten()

        avg_psd.append(psd)

    psd = np.array(avg_psd)
    psd = np.mean(psd, 0)

    kfreq_x = np.fft.fftfreq(depth.shape[0]) * depth.shape[0]
    kfreq_y = np.fft.fftfreq(depth.shape[1]) * depth.shape[1]

    kfreq2D = np.meshgrid(kfreq_y, kfreq_x)
    knrm = np.sqrt(kfreq2D[0] ** 2 + kfreq2D[1] ** 2)

    knrm = knrm.flatten()

    kbins = np.arange(0.5, depth.shape[0] // 2 + 1, 1.)
    kvals = 0.5 * (kbins[1:] + kbins[:-1])

    Abins, _, _ = s.binned_statistic(knrm, psd, statistic="mean", bins=kbins)
    Abins *= np.pi * (kbins[1:] ** 2 - kbins[:-1] ** 2)
    return kvals, Abins
//  [python]
def get_psd_of_depth_list(depth_list):
    avg_psd = []
    for depth in depth_list:
        psd = np.fft.fft2((depth - np.mean(depth)) / np.std(depth))  # [480, 640]
        psd = np.abs(psd) ** 2
        psd = psd.flatten()

        avg_psd.append(psd)

    psd = np.array(avg_psd)
    psd = np.mean(psd, 0)

    kfreq_x = np.fft.fftfreq(depth.shape[0]) * depth.shape[0]
    kfreq_y = np.fft.fftfreq(depth.shape[1]) * depth.shape[1]

    kfreq2D = np.meshgrid(kfreq_y, kfreq_x)
    knrm = np.sqrt(kfreq2D[0] ** 2 + kfreq2D[1] ** 2)

    knrm = knrm.flatten()

    kbins = np.arange(0.5, depth.shape[0] // 2 + 1, 1.)
    kvals = 0.5 * (kbins[1:] + kbins[:-1])

    Abins, _, _ = stats.binned_statistic(knrm, psd, statistic="mean", bins=kbins)
    Abins *= np.pi * (kbins[1:] ** 2 - kbins[:-1] ** 2)
    return kvals, Abins
//  [python]
def estimateBehav(self, binsize=0.25):
        ratemap = self.ratemap
        gridcntr = self.gridcenter
        spks = self._spks
        speed = self._obj.position.speed
        t = self._obj.position.t
        x = self._obj.position.x
        y = self._obj.position.y

        maze = self._obj.epochs.maze
        tmz = np.arange(maze[0], maze[1], binsize)
        actualposx = np.histogram(t, bins=tmz, weights=x)[0] / np.histogram(t, bins=tmz)[0]
        actualposy = np.histogram(t, bins=tmz, weights=y)[0] / np.histogram(t, bins=tmz)[0]
        meanspeed = np.histogram(t[1:], bins=tmz, weights=speed)[0] / np.histogram(t[1:], bins=tmz)[0]
        actualpos = np.vstack((actualposx, actualposy))

        spkcount = np.asarray([np.histogram(cell, bins=tmz)[0] for cell in spks])

        """ 
        ===========================
        Probability is calculated using this formula
        prob = (1 / nspike!)* ((0.1 * frate)^nspike) * exp(-0.1 * frate)
        =========================== 
        """

        Ncells = len(spks)
        cell_prob = np.zeros((ratemap.shape[1], spkcount.shape[1], Ncells))
        for cell in range(Ncells):
            cell_spkcnt = spkcount[cell, :][np.newaxis, :]
            cell_ratemap = ratemap[cell, :][:, np.newaxis]

            coeff = 1 / (factorial(cell_spkcnt))
            # broadcasting
            cell_prob[:, :, cell] = (((0.1 * cell_ratemap) ** cell_spkcnt) * coeff) * (
                np.exp(-0.1 * cell_ratemap)
            )

        posterior = np.prod(cell_prob, axis=2)
        posterior /= np.sum(posterior, axis=0)
        self.posterior = posterior
        self.decodedPos = gridcntr[:, np.argmax(self.posterior, axis=0)]
        self.decodingtime = tmz
        self.actualPos = actualpos
        self.speed = meanspeed
//  [python]
def estimateBehav(self, binsize=0.25):
        ratemap = self.ratemap
        gridcntr = self.gridcenter
        spks = self._spks
        speed = self._obj.position.speed
        t = self._obj.position.t
        x = self._obj.position.x
        y = self._obj.position.y

        maze = self._obj.epochs.maze
        tmz = np.arange(maze[0], maze[1], binsize)
        actualposx = binned_statistic(t, values=x, bins=tmz)[0]
        actualposy = binned_statistic(t, values=y, bins=tmz)[0]
        meanspeed = binned_statistic(t[1:], speed, bins=tmz)[0]
        actualpos = np.vstack((actualposx, actualposy))

        spkcount = np.asarray([np.histogram(cell, bins=tmz)[0] for cell in spks])

        """ 
        ===========================
        Probability is calculated using this formula
        prob = (1 / nspike!)* ((0.1 * frate)^nspike) * exp(-0.1 * frate)
        =========================== 
        """

        Ncells = len(spks)
        cell_prob = np.zeros((ratemap.shape[1], spkcount.shape[1], Ncells))
        for cell in range(Ncells):
            cell_spkcnt = spkcount[cell, :][np.newaxis, :]
            cell_ratemap = ratemap[cell, :][:, np.newaxis]

            coeff = 1 / (factorial(cell_spkcnt))
            # broadcasting
            cell_prob[:, :, cell] = (((0.1 * cell_ratemap) ** cell_spkcnt) * coeff) * (
                np.exp(-0.1 * cell_ratemap)
            )

        posterior = np.prod(cell_prob, axis=2)
        posterior /= np.sum(posterior, axis=0)
        self.posterior = posterior
        self.decodedPos = gridcntr[:, np.argmax(self.posterior, axis=0)]
        self.decodingtime = tmz
        self.actualPos = actualpos
        self.speed = meanspeed
//  [python]
def plot_eta(current_dir,models,phys):
    
    fig,ax=plt.subplots()
    plt.figure(figsize=(6,4))
    
    for model in models.split(','):
        input_dir=os.path.join(current_dir,model,'verify_input_calQ.csv')
        output_dir=os.path.join(current_dir,model,'verify_decoded_calQ.csv')
        
        input_Q=load_data(input_dir)
        output_Q=load_data(output_dir)
        
        indices = range(0,(len(input_Q)))
        
        emd_values = np.array([emd(input_Q[i],output_Q[j]) for i, j in zip(indices,indices)])
    
        eta=[]
        for i in indices:
            eta=np.append(eta,phys[i][0])
            
        x=eta
        y=emd_values
        
        nbins=10
        lims=None
        stats=True
        if lims==None: lims = (x.min(),x.max())
        median_result = s.binned_statistic(x, y, bins=nbins, range=lims, statistic=lambda x: np.quantile(x,0.5))
        lo_result     = s.binned_statistic(x, y, bins=nbins, range=lims, statistic=lambda x: np.quantile(x,0.5-0.68/2))
        hi_result     = s.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.errorbar(x=bin_centers, y=median, yerr=[loe,hie], label=model)
        fig=plt.legend()
    
    plt.legend()
    plt.xlabel(r'$\eta$')
    plt.ylabel('EMD')
    plt.savefig(current_dir+'emd_v_eta.pdf',dpi=600)
//  [python]
def plot_eta(current_dir,models,phys):
    
    fig,ax=plt.subplots()
    plt.figure(figsize=(6,4))
    
    for model in models.split(','):
        input_dir=os.path.join(current_dir,model,'verify_input_calQ.csv')
        output_dir=os.path.join(current_dir,model,'verify_decoded_calQ.csv')
        
        input_Q=load_data(input_dir)
        output_Q=load_data(output_dir)
        
        indices = range(0,(len(input_Q)))
        
        emd_values = np.array([emd(input_Q[i],output_Q[j]) for i, j in zip(indices,indices)])
    
        eta=[]
        for i in indices:
            eta=np.append(eta,phys[i][0])
            
        x=eta
        y=emd_values
        
        nbins=10
        lims=None
        stats=True
        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.errorbar(x=bin_centers, y=median, yerr=[loe,hie], label=model)
        fig=plt.legend()
    
    plt.legend()
    plt.xlabel(r'$\eta$')
    plt.ylabel('EMD')
    plt.savefig(current_dir+'emd_v_eta.pdf',dpi=600)
//  [python]
def reference_cv_alpha_div_2_loss(cv_hist_edges, n_nonempty_bins, potentials1, potentials2, cvs, metad_rbias):
    """Reference implementation of the CVMetaDKLDivLoss loss function.

    Assumes kT is 1. Return the loss and its gradient w.r.t. potentials2.
    """
    metad_rbias = metad_rbias.detach().numpy()
    reduced_work = (potentials2 - potentials1).detach().numpy() + metad_rbias

    cv_logsumexp_w = s.binned_statistic(cvs, reduced_work, statistic=lambda x: scipy.special.logsumexp(x), bins=cv_hist_edges)[0]
    cv_normalization_factors = s.binned_statistic(cvs, metad_rbias, statistic=lambda x: scipy.special.logsumexp(x), bins=cv_hist_edges)[0]

    loss = np.nansum(cv_logsumexp_w - cv_normalization_factors) / n_nonempty_bins
    cv_indices = np.digitize(cvs, cv_hist_edges) - 1
    loss_prime = np.exp(reduced_work - cv_logsumexp_w[cv_indices]) / n_nonempty_bins

    return torch.tensor(loss, dtype=potentials2.dtype), torch.tensor(loss_prime, dtype=potentials2.dtype)
//  [python]
def reference_cv_alpha_div_2_loss(cv_hist_edges, n_nonempty_bins, potentials1, potentials2, cvs, metad_rbias):
    """Reference implementation of the CVMetaDKLDivLoss loss function.

    Assumes kT is 1. Return the loss and its gradient w.r.t. potentials2.
    """
    metad_rbias = metad_rbias.detach().numpy()
    reduced_work = (potentials2 - potentials1).detach().numpy() + metad_rbias

    cv_logsumexp_w, _, _ = scipy.stats.binned_statistic(cvs, reduced_work, scipy.special.logsumexp, bins=cv_hist_edges)
    cv_normalization_factors, _, _ = scipy.stats.binned_statistic(cvs, metad_rbias, scipy.special.logsumexp, bins=cv_hist_edges)

    loss = np.nansum(cv_logsumexp_w - cv_normalization_factors) / n_nonempty_bins
    cv_indices = np.digitize(cvs, cv_hist_edges) - 1
    loss_prime = np.exp(reduced_work - cv_logsumexp_w[cv_indices]) / n_nonempty_bins

    return torch.tensor(loss, dtype=potentials2.dtype), torch.tensor(loss_prime, dtype=potentials2.dtype)
//  [python]
def pbc_per_dimension(self, cluster_sel_positions, box, dimension, bins_min, bins_max, bins_step):
      """Function to find the broken parts of the micelle across the pbc and make it whole

         Parameters
         ------------
         
         spatial_clustering_results : np.array
            Numpy array with the positions of the resids 
            of interest. It needs to have the shape of n_molec x 3
         
         box : np.array
            Numpy array with the MDanalysis strat u.dimensions. Array with the x y and z dimensions
            follow by the angles of the box.

         dimension : class.int
            Dimension number that wants to be evaluated

         bins_min : class.int
            Smallest distance to evaluate if the pbc are broken (smallest x-y distance position)

         bins_max : class.int
            Biggest distance to evaluate if the pbc are broken (largest x-y distance position)

         bins_step : class.int
            Step between bins to evaluate broken pbc
            
         Returns
         --------
         
         cluster_sel_positions_return : np.array
            Array with the atom positions of the micelle made whole across the pbc
      
      """

      import MDAnalysis as mda 
      import MDAnalysis.analysis.distances
      import numpy as np
      import scipy.stats as stats

      cluster_sel_positions_return=cluster_sel_positions[:,dimension]

      y = np.histogram(cluster_sel_positions_return,
                                     bins=np.arange(bins_min, box[0]+bins_max, bins_step))[0]
        
       
      a = np.histogram(cluster_sel_positions_return,
                                   bins=np.arange(bins_min, box[0]+bins_max, bins_step))[1]
        

        
      move_above=False
      filled=np.where(y!=0)[0]

      move_above = [a[filled[i+1]] for i in range(len(filled)-1) if filled[i]!=filled[i+1]-1 ]

          
      if move_above!= False and len(move_above)>0:
         array = cluster_sel_positions_return

         for item in move_above:
            cluster_sel_positions_return[np.where(array > item)] = array[np.where(array > item)]-box[dimension] 

       
      return cluster_sel_positions_return
//  [python]
def pbc_per_dimension(self, cluster_sel_positions, box, dimension, bins_min, bins_max, bins_step):
      """Function to find the broken parts of the micelle across the pbc and make it whole

         Parameters
         ------------
         
         spatial_clustering_results : np.array
            Numpy array with the positions of the resids 
            of interest. It needs to have the shape of n_molec x 3
         
         box : np.array
            Numpy array with the MDanalysis strat u.dimensions. Array with the x y and z dimensions
            follow by the angles of the box.

         dimension : class.int
            Dimension number that wants to be evaluated

         bins_min : class.int
            Smallest distance to evaluate if the pbc are broken (smallest x-y distance position)

         bins_max : class.int
            Biggest distance to evaluate if the pbc are broken (largest x-y distance position)

         bins_step : class.int
            Step between bins to evaluate broken pbc
            
         Returns
         --------
         
         cluster_sel_positions_return : np.array
            Array with the atom positions of the micelle made whole across the pbc
      
      """

      import MDAnalysis as mda 
      import MDAnalysis.analysis.distances
      import numpy as np
      import scipy.stats as stats

      cluster_sel_positions_return=cluster_sel_positions[:,dimension]

      y = stats.binned_statistic(cluster_sel_positions_return,
                                     cluster_sel_positions_return,
                                     bins=np.arange(bins_min, box[0]+bins_max, bins_step),
                                     statistic='count').statistic
        
       
      a = stats.binned_statistic(cluster_sel_positions_return,
                                   cluster_sel_positions_return,
                                   bins=np.arange(bins_min, box[0]+bins_max, bins_step),
                                   statistic='count').bin_edges
        

        
      move_above=False
      filled=np.where(y!=0)[0]

      move_above = [a[filled[i+1]] for i in range(len(filled)-1) if filled[i]!=filled[i+1]-1 ]

          
      if move_above!= False and len(move_above)>0:
         array = cluster_sel_positions_return

         for item in move_above:
            cluster_sel_positions_return[np.where(array > item)] = array[np.where(array > item)]-box[dimension] 

       
      return cluster_sel_positions_return
//  [python]
def get_continuum(lam0, spec0, espec0, ccfconf=None):
    """Determine the continuum of the spectrum by fitting a spline

    Parameters
    ----------

    lam0: numpy array
        The wavelength vector
    spec0: numpy array
        The spectral vector
    espec0: numpy array
        The vector of spectral uncertainties
    ccfconf: CCFConfig object
        The CCF configuration object

    Returns
    -------
    cont: numpy array
        The continuum vector

"""

    lammin = lam0.min()
    N = np.log(lam0.max() / lammin) / np.log(1 + ccfconf.splinestep / 3e5)
    N = int(np.ceil(N))
    # Determine the nodes of the spline used for the continuum fit
    nodes = lammin * np.exp(
        np.arange(N) * np.log(1 + ccfconf.splinestep / 3e5))
    nodesedges = lammin * np.exp(
        (-0.5 + np.arange(N + 1)) * np.log(1 + ccfconf.splinestep / 3e5))
    medspec = np.median(spec0)
    if medspec <= 0:
        medspec = np.abs(medspec)
        if medspec == 0:
            medspec = 1
        logging.warning('The spectrum has a median that is non-positive...')

    BS = s.binned_statistic(lam0, spec0, 'median', bins=nodesedges)
    p0 = np.log(np.maximum(BS.statistic, 1e-3 * medspec))
    p0[~np.isfinite(p0)] = np.log(medspec)

    ret = scipy.optimize.least_squares(fit_resid,
                                       p0,
                                       loss='soft_l1',
                                       args=((nodes, lam0, spec0, espec0),
                                             False))
    cont = fit_resid(ret['x'], (nodes, lam0, spec0, espec0), True)
    return cont