from __future__ import division, print_function
import numpy as np
from lmfit.models import VoigtModel
from scipy.signal import argrelmax
import matplotlib.pyplot as plt


def lamda_from_bragg(th, d, n):
    return 2 * d * np.sin(th / 2.) / n


def find_peaks(chi, sides=6, intensity_threshold=0):
    # Find all potential peaks
    preliminary_peaks = argrelmax(chi, order=20)[0]

    # peaks must have at least sides pixels of data to work with
    preliminary_peaks2 = preliminary_peaks[
        np.where(preliminary_peaks < len(chi) - sides)]

    # make certain that a peak has a drop off which causes the peak height to
    # be more than twice the height at sides pixels away
    criteria = chi[preliminary_peaks2] >= 2 * chi[preliminary_peaks2 + sides]
    criteria *= chi[preliminary_peaks2] >= 2 * chi[preliminary_peaks2 - sides]
    criteria *= chi[preliminary_peaks2] >= intensity_threshold

    peaks = preliminary_peaks[np.where(criteria)]

    left_idxs = peaks - sides
    right_idxs = peaks + sides
    peak_centers = peaks
    left_idxs[left_idxs < 0] = 0
    right_idxs[right_idxs > len(chi)] = len(chi)
    return left_idxs, right_idxs, peak_centers


def get_wavelength_from_std_tth(x, y, d_spacings, ns, plot=False):
    """
    Return the wavelength from a two theta scan of a standard

    Parameters
    ----------
    x: ndarray
        the two theta coordinates
    y: ndarray
        the detector intensity
    d_spacings: ndarray
        the dspacings of the standard
    ns: ndarray
        the multiplicity of the reflection
    plot: bool
        If true plot some of the intermediate data
    Returns
    -------
    float:
        The average wavelength
    float:
        The standard deviation of the wavelength
    """
    l, r, c = find_peaks(y)
    lmfit_centers = []
    for lidx, ridx, peak_center in zip(l, r, c):
        mod = VoigtModel()
        pars = mod.guess(y[lidx: ridx],
                         x=x[lidx: ridx])
        out = mod.fit(y[lidx: ridx], pars,
                      x=x[lidx: ridx])
        lmfit_centers.append(out.values['center'])
    lmfit_centers = np.asarray(lmfit_centers)
    if plot:
        plt.plot(x, y)
        plt.plot(x[c], y[c], 'ro')
        plt.show()

    wavelengths = []
    l_peaks = lmfit_centers[lmfit_centers < 0.]
    r_peaks = lmfit_centers[lmfit_centers > 0.]
    for peak_set in [r_peaks, l_peaks[::-1]]:
        for peak_center, d, n in zip(peak_set, d_spacings, ns):
            tth = np.deg2rad(np.abs(peak_center))
            wavelengths.append(lamda_from_bragg(tth, d, n))
    return np.average(wavelengths), np.std(wavelengths)


if __name__ == '__main__':
    import os

    calibration_file = os.path.join('../../data/LaB6_d.txt')

    # step 0 load data
    d_spacings = np.loadtxt(calibration_file)
    for data_file in ['../../data/Lab6_67p8.chi', '../../data/Lab6_67p6.chi']:
        a = np.loadtxt(data_file)
        wavechange = []
        b = np.linspace(.1, 3, 100)
        for dx in b:
            x = a[:, 0]
            x = np.hstack((np.zeros(1), x))
            x = np.hstack((-x[::-1], x))
            y = a[:, 1]
            y = np.hstack((np.zeros(1), y))
            y = np.hstack((y[::-1], y))

            x = x[:] + dx
            y = y[:]
            wavechange.append(get_wavelength_from_std_tth(x, y, d_spacings,
                                                      np.ones(d_spacings.shape),
                                                      )[0])
        plt.plot(b, wavechange)
    plt.show()

@@ -2623,16 +2623,1223 @@
 gths)%0A%0A%0A
+from bluesky.callbacks import CollectThenCompute%0A%0A%0Aclass ComputeWavelength(CollectThenCompute):%0A    %22%22%22%0A    Example%0A    -------%0A    %3E%3E%3E cw = ComputeWavelgnth('tth_cal', 'some_detector', d_spacings, ns) %0A    %3E%3E%3E RE(scan(...), cw)%0A    %22%22%22%0A    CONVERSION_FACTOR = 12.3984  # keV-Angstroms%0A    def __init__(self, x_name, y_name, d_spacings, ns=None):%0A        self.x_name = x_name%0A        self.y_name = y_name%0A        self.d_spacings = d_spacings%0A        self.wavelength = None%0A        self.wavelength_std = None%0A        if ns is None:%0A            self.ns = np.ones(self.d_spacings.shape)%0A        else:%0A            self.ns = ns%0A%0A    @property%0A    def energy(self):%0A        if self.wavelength is None:%0A            return None%0A        else:%0A            return self.CONVERSION_FACTOR / self.wavelength%0A%0A    def compute(self):%0A        x = %5B%5D%0A        y = %5B%5D%0A        for event in self._events:%0A            x.append(event%5B'data'%5D%5Bself.x_name%5D)%0A            y.append(event%5B'data'%5D%5Bself.y_name%5D)%0A%0A        self.wavelength, self.wavelength_std = get_wavelength_from_std_tth(x, y, self.d_spacings, self.ns)%0A        print('wavelength', self.wavelength, '+-', self.wavelength_std)%0A        print('energy', self.energy)%0A    %0A%22%22%22%0A
 if __nam
@@ -4714,20 +4714,24 @@
 nge)%0A    plt.show()%0A
+%22%22%22%0A

Add bluesky callback wrapped around @CJ-Wright code