ls

# the name of the file that you wish to open
specfilename = '20151111'
# the name of the x column
x = 'MCMY'
# the name of the detector (y column)
y = 'PD21'
# the name of the monitor column
monitor = 'SRcur'
# the scans that you wish to process
scans = [108, 110, 112, 114]
# the name of the output file that you are going to write
output_file_name = '-'.join([specfilename, x, y, monitor]) + '-' + '_'.join([str(scan) for scan in scans])
print('output file name is going to be: %s' % output_file_name)

"""
Explanation: Alignment and interpolation of inelastic scattering data
Some required parameters (This section requires editing!)
End of explanation
"""


interpolation_mode = 'linear'
# The number to divide the step size by 
# use a value < 1 for more interpolated points
# use a value > 1 for less interpolated points
densify_interpolated_axis_factor = 1

"""
Explanation: Specify the kind of interpolation you want to use as a string 
'linear'
'nearest'
'zero'
'slinear'
'quadratic
'cubic'

where 'slinear', 'quadratic' and 'cubic' refer to a spline interpolation of first, second or third order) or as an integer specifying the order of the spline interpolator to use.
End of explanation
"""


# Some imports that are required for this notebook
import matplotlib.pyplot as plt
from lmfit.models import LorentzianModel, LinearModel
import numpy as np
import os
import pandas as pd
%matplotlib inline

"""
Explanation: The boring stuff
Imports
End of explanation
"""


class Specfile:
    def __init__(self, filename):
        self.filename = os.path.abspath(filename)
        with open(self.filename, 'r') as f:
            scan_data = f.read().split('#S')
        scan_data = [section.split('\n') for section in scan_data]
        self.header = scan_data.pop(0)
        self.scans = {}
        for scan in scan_data:
            sid = int(scan[0].split()[0])
            self.scans[sid] = Specscan(self, scan)
    
    def __getitem__(self, key):
        return self.scans[key]
    
    def __len__(self):
        return len(self.scans)-1
    
    def __iter__(self):
        return (self.scans[sid] for sid in sorted(self.scans.keys()))

class Specscan:
    def __init__(self, specfile, raw_scan_data):
        self.specfile = specfile
        self.raw_scan_data = raw_scan_data
        header_row = self.raw_scan_data.pop(0).split()
        self.scan_id = int(header_row.pop(0))
        self.scan_command = header_row.pop(0)
        self.scan_args = header_row
        for row in self.raw_scan_data:
            if row.startswith('#L'):
                self.col_names = row.split()[1:]
        scan_data = [row.split() for row in self.raw_scan_data 
                     if not row.startswith('#') if row]
        self.scan_data = pd.DataFrame(data=scan_data, columns=self.col_names, dtype=float)
    def __repr__(self):
        return 'Specfile("%s")[%s]' % (self.specfile.filename, self.scan_id)

    def __str__(self):
        return str(self.scan_data)
    
    def __len__(self):
        return len(self.scan_data)

    def plot(self, column_names=None, x=None):
        if x is None:
            x = self.scan_data.columns[0]
        if column_names is None:
            column_names = self.scan_data.columns
        ncols = 2
        nrows = int(np.ceil(len(column_names)/ncols))
        try:
            self.ncols
            self.nrows
        except AttributeError:
            self.ncols = 0
            self.nrows = 0
        if self.ncols != ncols or self.nrows != nrows:
            self.ncols, self.nrows = ncols, nrows
            self.fig, self.axes = plt.subplots(nrows=nrows,
                                               ncols=ncols,
                                               figsize=(5*ncols, 2*nrows))
        self.arts = {}
        for data, ax in zip(column_names, self.axes.ravel()):
            ax.cla()
            self.arts[data] = ax.plot(self.scan_data[x], self.scan_data[data], label=data)
            ax.legend(loc=0)
            

def fit(x, y, bounds=None):
    """Fit a lorentzian + linear background to `field` in `scan`
    
    Parameters
    ----------
    scan : Specscan object
    field : The field to fit
    bounds : The +/- range to fit the data to
    
    Returns
    -------
    fit : lmfit.model.ModelFit
        The results of fitting the data to a linear + lorentzian peak
    
    Examples
    --------
    >>> fit = fit_lorentzian(scan.scan_data)
    >>> fit.plot()
    """
    lorentzian = LorentzianModel()   
    linear = LinearModel()
    center = x[np.argmax(y)]
    if bounds is None:
        lower, upper = 0, len(x)
    else:
        lower = center - bounds
        upper = center + bounds
        if lower < 0:
            lower = 0
        if upper > len(x):
            upper = len(x)
    bounds = slice(lower, upper)
#     print("Using bounds = {}".format(bounds))
    y = y[bounds]
    x = x[bounds]
#     print("Using x = {}".format(x))
#     print("Using y = {}".format(y))
    lorentzian_params = lorentzian.guess(y, x=x, center=center)
    linear_params = linear.guess(y, x=x)
    lorentzian_params.update(linear_params)
    model = lorentzian + linear
    return model.fit(y, x=x, params=lorentzian_params)

def plotter(xy):
    fig, ax = plt.subplots()
    arts = {}
    for sid, (xdata, ydata) in zip(scans, xy):
        arts[sid] = ax.plot(xdata, ydata, '-o', label=sid)
    ax.legend(loc=0)
    

"""
Explanation: Defining required objects and functions
End of explanation
"""


f = Specfile(specfilename)

"""
Explanation: The exciting stuff!
Open the spec file with the Specfile object defined above
End of explanation
"""


raw = [(
        f[scan_id].scan_data[x].values,
        f[scan_id].scan_data[y].values
    ) for scan_id in scans]

# Use the plotter helper function defined above
plotter(raw)

"""
Explanation: Plot the raw data to make sure it looks ok
End of explanation
"""


normalized = [(
        x, 
        y / f[scan_id].scan_data[monitor].values
    ) for scan_id, (x, y) in zip(scans, raw)]

"""
Explanation: Normalize the data by the monitor value defined in the first cell
End of explanation
"""


# Use the plotter helper function defined above
plotter(normalized)

"""
Explanation: Plot the normalized data to make sure it looks ok
End of explanation
"""


fits = [fit(xdata, ydata) for (xdata, ydata) in normalized]

"""
Explanation: Fit to a linear model plus a lorentzian model
End of explanation
"""


fig_kws = {'figsize': [15, 15]}
for sid, f in zip(scans, fits):
    title_dict = {'title': 'scan_id: %s' % sid}
    f.plot(numpoints=len(f.data)*10, ax_res_kws=title_dict,
           ax_fit_kws=title_dict, fig_kws=fig_kws)

"""
Explanation: Plot the fits to make sure they look ok!
End of explanation
"""


zeroed = [(np.array(f.userkws['x']-f.params['center'], dtype=float), f.data) for f in fits]

"""
Explanation: Shift the peaks to a zero energy transfer based on the center of the Lorentzian model
End of explanation
"""


plotter(zeroed)

"""
Explanation: Plot the shifted peaks to make sure they look ok!
End of explanation
"""


diff = np.average([np.average(np.diff(x)) for x, y in zeroed])
minval = np.min([np.min(x) for x, y in zeroed])
maxval = np.max([np.max(x) for x, y in zeroed])

diff, minval, maxval

new_axis = np.arange(minval, maxval, diff / densify_interpolated_axis_factor)

"""
Explanation: Compute the new axis onto which the data should be interpolated
This new axis will be a linear space from the minimum value of all scans to the maximum value of all scans with a step size equal to the average step size between all data points
End of explanation
"""


from scipy.interpolate import interp1d

interpolaters = [interp1d(x, y, kind=interpolation_mode, 
                          bounds_error=False, 
                          fill_value=np.nan) 
                 for x, y in zeroed]

# Create a dict of the interpolated values so it can easily be passed to pandas
interpolated = {sid: interpolator(new_axis) 
          for sid, interpolator in zip(scans, interpolaters)}

df = pd.DataFrame(interpolated, index=new_axis)

"""
Explanation: Set up the interpolator
End of explanation
"""


df.plot(style='-o')

"""
Explanation: Plot the new interpolated values
End of explanation
"""


df.sum(axis=1).plot(style='-o')

"""
Explanation: Just take the straight sum of all data in the data frame
End of explanation
"""


(df.sum(axis=1) / df.count(axis=1)).plot(style='-o')

"""
Explanation: Ok, that looks a little weird, let's divide the sum by the number of channels at each x value
End of explanation
"""


df.dropna().sum(axis=1).plot(style='-o')

"""
Explanation: Hmm, that still looks a little strange, lets just sum the channels where they all have values
End of explanation
"""


df.dropna().sum(axis=1).to_csv(output_file_name, encoding='ascii', sep=',')

"""
Explanation: Output the data
End of explanation
"""


!cat 20151111-MCMY-PD21-SRcur-108_110_112_114

"""
Explanation: Show that we actually output the data...
End of explanation
"""


!pwd

print(output_file_name)

"""
Explanation: If you are working in the notebook and want to download this data file to your local computer, follow these steps:

Figure out where that file lives with !pwd
Remember the file name given by: print(output_file_name)
Now that you know where the file lives, go back to the home page of jupyter hub by right clicking on the Jupyter icon in the top-left of this page and selecting "Open in new tab"
Then go to that new tab, navigate to the folder that was shown by !pwd and open the file named output_file_name, shown above
Finally, open the file by left-clicking on it.  Once it opens, select File -> Download
End of explanation
"""


