[add] plot arrivals option added

[change] code reformated by PyCharm

pylot/core/util/widgets.py

@@ -12,6 +12,7 @@ import warnings
 import copy
 import datetime
 import numpy as np
+import matplotlib.pyplot as plt
 
 try:
     import pyqtgraph as pg
@@ -590,8 +591,7 @@ class WaveformWidgetPG(QtGui.QWidget):
 class WaveformWidget (FigureCanvas):
     def __init__(self, parent=None, xlabel='x', ylabel='y', title='Title'):
 
-        self._parent = None
-        self.setParent(parent)
+        self._parent = parent
         self.figure = Figure ()
         self.figure.set_facecolor ((.92, .92, .92))
         # attribute plotdict is a dictionary connecting position and a name
@@ -867,6 +867,8 @@ class PickDlg(QDialog):
 
         # initialize plotting widget
         self.multicompfig = WaveformWidget (self)
+        self.phaseplot = PhasePlotWidget(self)
+        self.phaseplot.hide()
 
         # setup ui
         self.setupUi ()
@@ -901,8 +903,12 @@ class PickDlg(QDialog):
                 self.model = TauPyModel (model)
                 self.get_arrivals ()
                 self.drawArrivals ()
+                self.activateArrivalsButton (True)
+            else:
+                self.activateArrivalsButton (False)
         except Exception as e:
             print ('Warning: Could not init expected picks from taup: {}'.format (e))
+            self.activateArrivalsButton (False)
             
         # init pick delete (with right click)
         self.connect_pick_delete ()
@@ -960,6 +966,9 @@ class PickDlg(QDialog):
         # self.p_button.setToolTip('Hotkey: "1"')
         # self.s_button.setToolTip('Hotkey: "2"')
 
+        self.plot_arrivals_button = QPushButton ('Plot arrivals')
+        self.plot_arrivals_button.setCheckable(True)
+
         # create accept/reject button
         self.accept_button = QPushButton ('&Accept Picks')
         self.reject_button = QPushButton ('&Reject Picks')
@@ -990,10 +999,14 @@ class PickDlg(QDialog):
             _dialtoolbar.addWidget (self.reject_button)
         else:
             _dialtoolbar.addWidget (self.nextStation)
+        _dialtoolbar.addWidget (self.plot_arrivals_button)
 
         # layout the innermost widget
         _innerlayout = QVBoxLayout ()
-        _innerlayout.addWidget(self.multicompfig)
+        _innerinnerlayout = QtGui.QHBoxLayout()
+        _innerinnerlayout.addWidget (self.multicompfig)
+        _innerinnerlayout.addWidget (self.phaseplot)
+        _innerlayout.addLayout(_innerinnerlayout)
 
         # add button box to the dialog
         _buttonbox = QDialogButtonBox (QDialogButtonBox.Ok |
@@ -1017,6 +1030,7 @@ class PickDlg(QDialog):
         self.reject_button.clicked.connect (self.reject)
         self.accept_button.clicked.connect (self.disable_ar_buttons)
         self.reject_button.clicked.connect (self.disable_ar_buttons)
+        self.plot_arrivals_button.clicked.connect (self.toggle_arrivals_plot)
         _buttonbox.accepted.connect (self.accept)
         _buttonbox.rejected.connect (self.reject)
 
@@ -1024,10 +1038,34 @@ class PickDlg(QDialog):
         self.setLayout (_outerlayout)
         self.resize (1280, 720)
 
+    def activateArrivalsButton(self, val=True):
+        self.plot_arrivals_button.setEnabled (val)
+
+    def toggle_arrivals_plot(self):
+        if self.plot_arrivals_button.isChecked():
+            self.plot_arrivals()
+        else:
+            self.hide_arrivals_plot()
+
+    def hide_arrivals_plot(self):
+        self.phaseplot.hide()
+
+    def plot_arrivals(self):
+        if self.phaseplot.new:
+            self.get_arrivals(True)
+            ax = self.phaseplot.ax
+            self.arrivals.plot(ax=ax, show=False)
+            ax.legend()
+            self.phaseplot.new = False
+        self.phaseplot.show()
+        self.phaseplot.draw()
+
     def setDirty(self, bool):
         self._dirty = bool
 
-    def get_arrivals(self):
+    def get_arrivals(self, plot=False):
+        func = {True: self.model.get_ray_paths_geo,
+                False: self.model.get_travel_times_geo}
         phases = self.prepare_phases ()
         station_id = self.data.traces[0].get_id ()
         parser = self.parent ().metadata[1]
@@ -1037,7 +1075,7 @@ class PickDlg(QDialog):
             source_origin = origins[0]
         else:
             raise ValueError ('No source origin given.')
-        arrivals = self.model.get_travel_times_geo(source_origin.depth,
+        arrivals = func[plot] (source_origin.depth,
                                source_origin.latitude,
                                source_origin.longitude,
                                station_coords['latitude'],
@@ -1049,7 +1087,7 @@ class PickDlg(QDialog):
         settings = QtCore.QSettings ()
         p_phases = settings.value ('p_phases')
         s_phases = settings.value ('s_phases')
-        phases = p_phases + s_phases
+        phases = p_phases + ',' + s_phases
         phases = phases.split (',')
         phases = [phase.strip () for phase in phases]
         return phases
@@ -1924,6 +1962,15 @@ class PickDlg(QDialog):
             self.qmb.show ()
 
 
+class PhasePlotWidget(FigureCanvas):
+    def __init__(self, parent=None):
+        self._parent = parent
+        self.fig = Figure()
+        self.ax = self.fig.add_subplot(111, projection='polar')
+        self.new = True
+        super (PhasePlotWidget, self).__init__ (self.fig)
+
+
 class TuneAutopicker (QWidget):
     update = QtCore.Signal (str)
     '''