prasopes.py

#!/usr/bin/env python3

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.widgets import SpanSelector
from PyQt5 import QtCore
from PyQt5 import QtWidgets
#from rawFin2 import load_raw
from rawParse import load_raw
import sys
import matplotlib
import numpy as np
import logging
matplotlib.use("Qt5Agg")


def zoom_factory(ax, base_scale, plot=None, an=None):
    """returns zooming functionality to axis"""
    def zoom_fun(event, chart, ann):
        """zoom when scrolling"""
        if event.inaxes == ax:
            if event.button == 'up':
                #zoom in
                scale_factor = 1/base_scale
            elif event.button == 'down':
                #zoom out
                scale_factor = base_scale
            else:
                #should not happen
                scale_factor = 1
                print(event.button)

            if event.key == 'shift':
                data = event.ydata
                new_top = data + (ax.get_ylim()[1] - data)*scale_factor
                ax.set_ylim([new_top*-0.01,new_top])
            else:
                data = event.xdata
                x_left = data - ax.get_xlim()[0]
                x_right = ax.get_xlim()[1] - data
                ax.set_xlim([data - x_left*scale_factor,
                            data + x_right*scale_factor])
            if ann is not None:
                if type(chart['x']) != type(None):
                    ann_spec(event.inaxes, chart['x'], chart['y'], ann)
            ax.figure.canvas.draw()

    fig = ax.get_figure()  # get the figure of interest
    # attach the call back
    fig.canvas.mpl_connect('scroll_event',
                           lambda event: zoom_fun(event, plot, an))

    return zoom_fun


def pan_factory(axis, plot=None, annotations=None):
    """pan spectrum when you press a button"""
    def pan_fun(event, ax):
        #rescale to origin if doubleclicked
        if event.dblclick and event.inaxes == ax:
            ax.get_figure()
            ax.autoscale(True)
            ax.set_ylim(ax.get_ylim()[1]*-0.01,ax.get_ylim()[1]*1.1)
            if annotations is not None:
                ann_spec(ax, plot['x'], plot['y'], annotations)
            ax.figure.canvas.draw()
        #otherwise pan
        elif event.button == 1 and event.inaxes == ax:
            ax.start_pan(event.x, event.y, event.button)
            id_drag = fig.canvas.mpl_connect(
                'motion_notify_event',
                lambda action: drag_fun(action, ax))
            id_release = fig.canvas.mpl_connect(
                'button_release_event',
                lambda action: drag_end(
                    action, id_drag, id_release, plot, annotations, ax))

    def drag_fun(event, ax):
        ax.drag_pan(1, 'x', event.x, event.y)
        ax.figure.canvas.draw()

    def drag_end(event, id_drag, id_release, chart, ann, ax):
        if event.button == 1:
            fig.canvas.mpl_disconnect(id_drag)
            fig.canvas.mpl_disconnect(id_release)
            if ann is not None:
                if type(chart['x']) != type(None):
                    ann_spec(ax, chart['x'], chart['y'], ann)
            ax.figure.canvas.draw()
    fig = axis.get_figure()
    fig.canvas.mpl_connect('button_press_event',
                           lambda action: pan_fun(action, axis))

def pick_masses(x_min, x_max, ax, annotations, plot):
    """zoom the spectrum in x axis by mass range"""
    ax.set_xlim(x_min, x_max)
    ann_spec(spectrum, plot['x'], plot['y'], annotations)
    ax.figure.canvas.draw()

def pick_times(x_min, x_max, mpl_spectrum, data_set, mpl_chromatogram,
               annotations, mass_spect):
    """plot averaged spectrum of subselected part of the chromatogram"""
    start_scan = None
    end_scan = None 
    for i, times in enumerate(data_set['chrom_dat'][0, :]):
        if times > x_min and start_scan is None:
            start_scan = i
        if times > x_max and end_scan is None:
            end_scan = i
    mpl_spectrum.clear()
    annotations.clear()
    mass_spect['x'] = data_set['masses']
    mass_spect['y'] = np.mean(data_set['matrix'][start_scan:end_scan]
                              , axis=0)
    spectrum_plot(spectrum, mass_spect['x'], mass_spect['y'], annotations)
    mpl_chromatogram.clear()
    chrom_plot(mpl_chromatogram, data_set['chrom_dat'][0, :],
               data_set['chrom_dat'][1, :])
    mpl_chromatogram.plot(data_set['chrom_dat'][0, start_scan:end_scan],
                          data_set['chrom_dat'][1, start_scan:end_scan],
                          'b.')

def populate(mpl_chromatogram, mpl_spectrum, data_set, time_sel, an,
             mass_spect):
    """populate the GUI plots with desired dataset"""
    an.clear()
    def update_span_selector(chrom, spec, d_set, ms_spec):
        time_selector = SpanSelector(
            chrom, lambda x_min, x_max: pick_times(
                       x_min, x_max,spec, d_set, chrom, an, ms_spec),
            'horizontal',
            useblit=True,
            rectprops=dict(alpha=0.15, facecolor='purple'), button=3)
        return time_selector
    mpl_spectrum.clear()
    mpl_chromatogram.clear()
    mass_spect['x'] = data_set['masses']
    mass_spect['y'] = np.mean(data_set['matrix'], axis=0)
    spectrum_plot(mpl_spectrum, mass_spect['x'], mass_spect['y'], an)
    chrom_plot(mpl_chromatogram, data_set['chrom_dat'][0, :],
               data_set['chrom_dat'][1, :])
    time_sel[0] = update_span_selector(mpl_chromatogram, mpl_spectrum,
                                       data_set, mass_spect)
    mpl_chromatogram.figure.canvas.draw()
    mpl_spectrum.figure.canvas.draw()

def ann_spec(ms_spec, mass, intensities, annotations, ann_limit=0.01):
    """annotate spectrum
    
    First define borders where the annotation should occur. Then the
    function finds local maxima of the data set by forcycling through
    the dataset and then substracts found local minimas so it contains
    only subset of the most intense one which are in reasonable distance
    from each other. Then the function remove old annotation and
    annotate the spectrum by abovementioned subsection of local maximas.
    """
    def sub_peaks(peakz, x_value, y_value, xrange, yrange, coef_x=10,
                  coef_y=10):
        """Returns reasonable subselection of local maximas"""
        sort_peaks = sorted(peakz,
                key=lambda h: y_value[h],
                reverse=True)
        subtracted_peaks=[]
        red_x = xrange / coef_x
        red_y = yrange / coef_y
        for peak in sort_peaks:
            add = True
            yv = y_value[peak]
            xv = x_value[peak]
            for u in subtracted_peaks:
                if abs(yv - y_value[u]) < red_y \
                        and abs(xv - x_value[u]) < red_x:
                    add = False
            if add:
                subtracted_peaks.append(peak)
        return subtracted_peaks

    borders=ms_spec.get_xlim()
    start = np.argmax(mass>borders[0])
    end = np.argmax(mass>borders[1])
    if end == 0:
        end = np.argmax(mass)
    peaks = []
    l = ms_spec.get_ylim()[1] * ann_limit
    for i,intensity in enumerate(intensities[start:end], start=start):
        #end selects the first element which not to include.
        #Thus q[i+1] always exists as it points to that last element.
        if intensity > l and intensity > q[i-1] and intensity > q[i+1]:
            peaks.append(i)
    s_peaks = sub_peaks(peaks, mass, intensities,
                        ms_spec.get_xlim()[1] - ms_spec.get_xlim()[0],
                        ms_spec.get_ylim()[1] - ms_spec.get_ylim()[0])
    
    for intensity in annotations:
        intensity.remove()
    annotations.clear()

    for i in s_peaks:
        annotations.append(ms_spec.annotate(
            '{0:.2f}'.format(mass[i]), xy=(mass[i], intensities[i]),
            textcoords='data'))

def spectrum_plot(spc, mass, intensity,an):
    """Define and populate spectrum"""
    spc.plot(mass, intensity)
    spc.set_title("Spectrum:", loc ="right")
    spc.set_xlabel("m/z")
    spc.set_ylabel("ion count")
    spc.set_ylim(spc.get_ylim()[1] * -0.01, )
    spc.ticklabel_format(scilimits=(0, 0), axis='y')
    ann_spec(spc, mass, intensity, an)

def chrom_plot(chrom, times, tic):
    """Define and populate chromatogram"""
    chrom.plot(times, tic)
    chrom.set_ylim(chrom.get_ylim()[1] * -0.011,
                   chrom.get_ylim()[1] * 1.1)
    chrom.set_title("Chromatogram:", loc ="right")
    chrom.set_xlabel("time (min)")
    chrom.set_ylabel("total ion count")
    chrom.ticklabel_format(scilimits=(0, 0), axis='y')

def open_file(chrom, spc, d_set, time_sel, an, ms_s):
    """Returns chromatogram, masses and ion intensities"""
    filename=QtWidgets.QFileDialog.getOpenFileName(
            caption = "Open spectrum",
            filter="Finnigan RAW files (*.raw, *.RAW)")[0]
    if filename is '':
        return
    d_set['chrom_dat'], d_set['masses'], d_set['matrix'] \
        = load_raw(filename)
    populate(chrom, spc, d_set, time_sel, an, ms_s)


if __name__=="__main__":
    #ds for data_set
    ds = dict(chrom_dat=None,masses=None,matrix=None)
    #mass spectrometry set
    ms = dict(x=None,y=None)

    p_logger = logging.getLogger('parseLogger')
    logging.basicConfig()
    p_logger.setLevel("WARN")

    graph = Figure(figsize=(5,4),dpi=100)
    graph.patch.set_facecolor("None")

    chromatogram = graph.add_subplot(211,facecolor=(1,1,1,0.8))
    spectrum = graph.add_subplot(212,facecolor=(1,1,1,0.8))
    graph.tight_layout()

    mpl_canvas = FigureCanvas(graph)
    mpl_canvas.setStyleSheet("background-color:transparent;")
    mpl_canvas.setAutoFillBackground(False)
    mpl_canvas.setFocusPolicy( QtCore.Qt.ClickFocus )
    
    timeSelector=[None]
    annotation=[]

    pan_factory(chromatogram)
    zoom_factory(chromatogram, 1.15)
    pan_factory(spectrum, ms, annotation)
    zoom_factory(spectrum, 1.15, ms, annotation)
    mass_selector = SpanSelector(
        spectrum, lambda x_min, x_max: pick_masses(
            x_min, x_max, spectrum, annotation, ms),
        'horizontal', useblit=True, rectprops=dict(alpha=0.15,
                                                   facecolor='purple'),
        button=3)

    app = QtWidgets.QApplication(sys.argv)
    main_window = QtWidgets.QMainWindow()

    file_menu = QtWidgets.QMenu('&File', main_window)
    main_window.menuBar().addMenu(file_menu)
    file_menu.addAction('&Open..', 
        lambda: open_file(chromatogram, spectrum, ds, timeSelector,
                          annotation, ms),
        QtCore.Qt.CTRL + QtCore.Qt.Key_O)
    file_menu.addAction('&Quit', main_window.close,
                        QtCore.Qt.CTRL + QtCore.Qt.Key_Q)

    main_widget = QtWidgets.QWidget(main_window)
    main_window.setCentralWidget(main_widget)

    layout = QtWidgets.QVBoxLayout(main_widget)
    layout.addWidget(mpl_canvas)
    mpl_canvas.setFocus()


    if len(sys.argv) == 2:
        raw_file=sys.argv[1]
        ds['chrom_dat'],ds['masses'],ds['matrix'] = load_raw(raw_file)
        populate(chromatogram, spectrum, ds, timeSelector, annotation,
                 ms)
    else:
        spectrum_plot(spectrum, [0], [0], annotation)
        chrom_plot(chromatogram, [0], [0])

    main_window.show()
    sys.exit(app.exec_())