diff --git a/interactive.py b/interactive.py index 70ab198..20b0407 100644 --- a/interactive.py +++ b/interactive.py @@ -21,12 +21,16 @@ from PyQt5.QtWidgets import QMenu, QWidget, QVBoxLayout, QMessageBox from PyQt5.QtWidgets import QSlider, QHBoxLayout, QLabel, QDialog from PyQt5.QtWidgets import QDialogButtonBox - from fatiando import utils from fatiando.gravmag import talwani from fatiando.mesher import Polygon +LINE_ARGS = dict( + linewidth=2, linestyle='-', color='k', marker='o', + markerfacecolor='k', markersize=5, animated=False, alpha=0.6 + ) + class Moulder(FigureCanvasQTAgg): """ Interactive 2D forward modeling using polygons. @@ -108,104 +112,63 @@ class Moulder(FigureCanvasQTAgg): instructions = ' | '.join([ 'n: New polygon', 'd: delete', 'click: select/move', 'esc: cancel']) - - - def __init__(self, parent, area, x, z, data=None, - density_range=[-2000, 2000], width=5, height=4, dpi=100, - **kwargs): + def __init__(self, parent, width=5, height=4, dpi=100): self.fig = Figure(figsize=(width, height), dpi=dpi) super().__init__(self.fig) self.setParent(parent) + self._plotted = False - self.area = area - self.x, self.z = numpy.asarray(x), numpy.asarray(z) - self.density_range = density_range - self.data = data - # Used to set the ylims for the data axes. - if data is None: - self.dmin, self.dmax = 0, 0 - else: - self.dmin, self.dmax = data.min(), data.max() - self.predicted = kwargs.get('predicted', numpy.zeros_like(x)) - self.error = kwargs.get('error', 0) - self.cmap = kwargs.get('cmap', pyplot.cm.RdBu_r) - self.line_args = dict( - linewidth=2, linestyle='-', color='k', marker='o', - markerfacecolor='k', markersize=5, animated=False, alpha=0.6) + self.cmap = pyplot.cm.RdBu_r + + self._x = None + self._z = None + self._min_depth = 0 + self._max_depth = 10000 + + self.predicted = None + self.error = None + + self.predicted_line = [] + self.data = None self.polygons = [] self.lines = [] - self.densities = kwargs.get('densities', []) - vertices = kwargs.get('vertices', []) - for xy, dens in zip(vertices, self.densities): - poly, line = self._make_polygon(xy, dens) - self.polygons.append(poly) - self.lines.append(line) + self.densities = [] - def save_predicted(self, fname): - """ - Save the predicted data to a text file. + @property + def plotted(self): + return self._plotted - Data will be saved in 3 columns separated by spaces: x z data + @property + def x(self): + return self._x - Parameters: + @x.setter + def x(self, new_value): + self._x = numpy.asarray(new_value) - * fname : string or file-like object - The name of the output file or an open file-like object. + @property + def z(self): + return self._z - """ - numpy.savetxt(fname, numpy.transpose([self.x, self.z, self.predicted])) + @z.setter + def z(self, new_value): + self._z = numpy.asarray(new_value) - def save(self, fname): - """ - Save the application state into a pickle file. + @property + def min_depth(self): + return self._min_depth - Use this to persist the application. You can later reload the entire - object, with the drawn model and data, using the - :meth:`~fatiando.gravmag.interactive.Moulder.load` method. + @min_depth.setter + def min_depth(self, new_value): + self._min_depth = new_value - Parameters: + @property + def max_depth(self): + return self._max_depth - * fname : string - The name of the file to save the application. The extension doesn't - matter (use ``.pkl`` if in doubt). - - """ - with open(fname, 'w') as f: - vertices = [numpy.asarray(p.xy) for p in self.polygons] - state = dict(area=self.area, x=self.x, - z=self.z, data=self.data, - density_range=self.density_range, - cmap=self.cmap, - predicted=self.predicted, - vertices=vertices, - densities=self.densities, - error=self.error) - pickle.dump(state, f) - - @classmethod - def load(cls, fname): - """ - Restore an application from a pickle file. - - The pickle file should have been generated by the - :meth:`~fatiando.gravmag.interactive.Moulder.save` method. - - Parameters: - - * fname : string - The name of the file. - - Returns: - - * app : Moulder object - The restored application. You can continue using it as if nothing - had happened. - - """ - with open(fname) as f: - state = pickle.load(f) - app = cls(**state) - return app + @max_depth.setter + def max_depth(self, new_value): + self._max_depth = new_value @property def model(self): @@ -217,27 +180,7 @@ class Moulder(FigureCanvasQTAgg): return m def run(self): - """ - Start the application for drawing. - - Will pop-up a window with a place for drawing the model (below) and a - place with the predicted (and, optionally, observed) data (top). - - Follow the instruction on the figure title. - - When done, close the window to resume program execution. - """ self._figure_setup() - # Sliders to control the density and the error in the data - self.density_slider = widgets.Slider( - self.fig.add_axes([0.10, 0.01, 0.30, 0.02]), 'Density', - self.density_range[0], self.density_range[1], valinit=0., - valfmt='%6.0f kg/m3') - self.error_slider = widgets.Slider( - self.fig.add_axes([0.60, 0.01, 0.30, 0.02]), 'Error', - 0, 5, valinit=self.error, valfmt='%1.2f mGal') - # Put instructions on figure title - self.dataax.set_title(self.instructions) # Markers for mouse click events self._ivert = None self._ipoly = None @@ -263,34 +206,12 @@ class Moulder(FigureCanvasQTAgg): # Make the proper callback connections self.canvas.mpl_connect('button_press_event', self._button_press_callback) - self.canvas.mpl_connect('key_press_event', - self._key_press_callback) +# self.canvas.mpl_connect('key_press_event', +# self._key_press_callback) self.canvas.mpl_connect('button_release_event', self._button_release_callback) self.canvas.mpl_connect('motion_notify_event', self._mouse_move_callback) - self.density_slider.on_changed(self._set_density_callback) - self.error_slider.on_changed(self._set_error_callback) - - def plot(self, figsize=(10, 8), dpi=70): - """ - Make a plot of the data and model for embedding in IPython notebooks - - Doesn't require ``%matplotlib inline`` to embed the plot (as that would - not allow the app to run). - - Parameters: - - * figsize : list = (width, height) - The figure size in inches. - * dpi : float - The number of dots-per-inch for the figure resolution. - - """ - self._update_data_plot() - pyplot.close(self.fig) - data = print_figure(self.fig, dpi=dpi) - return Image(data=data) def _figure_setup(self, **kwargs): """ @@ -310,32 +231,12 @@ class Moulder(FigureCanvasQTAgg): kwargs['sharex'] = True axes = self.fig.subplots(2, 1, **kwargs) ax1, ax2 = axes - self.predicted_line, = ax1.plot(self.x, self.predicted, '-r') - if self.data is not None: - self.data_line, = ax1.plot(self.x, self.data, '.k') ax1.set_ylabel('Gravity anomaly (mGal)') ax2.set_xlabel('x (m)', labelpad=-10) - ax1.set_xlim(self.area[:2]) + ax1.set_xlim(self.x.min(), self.x.max()) ax1.set_ylim((-200, 200)) ax1.grid(True) - tmp = ax2.pcolor(numpy.array([self.density_range]), cmap=self.cmap) - tmp.set_visible(False) - pyplot.colorbar(tmp, orientation='horizontal', - pad=0.08, aspect=80).set_label(r'Density (kg/cm3)') - # Remake the polygons and lines to make sure they belong to the right - # axis coordinates - vertices = [p.xy for p in self.polygons] - newpolygons, newlines = [], [] - for xy, dens in zip(vertices, self.densities): - poly, line = self._make_polygon(xy, dens) - newpolygons.append(poly) - newlines.append(line) - ax2.add_patch(poly) - ax2.add_line(line) - self.polygons = newpolygons - self.lines = newlines - ax2.set_xlim(self.area[:2]) - ax2.set_ylim(self.area[2:]) + ax2.set_ylim(self.min_depth, self.max_depth) ax2.grid(True) ax2.invert_yaxis() ax2.set_ylabel('z (m)') @@ -590,11 +491,11 @@ class Moulder(FigureCanvasQTAgg): self._update_data() self._update_data_plot() - def _key_press_callback(self, event): + def key_press_callback(self, event_key): """ What to do when a key is pressed on the keyboard. """ - if event.key == 'd': + if event_key == 'd': if self._drawing and self._xy: self._xy.pop() if self._xy: @@ -629,7 +530,7 @@ class Moulder(FigureCanvasQTAgg): self.canvas.draw() self._update_data() self._update_data_plot() - elif event.key == 'n': + elif event_key == 'n': self._ivert = None self._ipoly = None for line, poly in zip(self.lines, self.polygons): @@ -647,7 +548,7 @@ class Moulder(FigureCanvasQTAgg): 'left click: set vertice', 'right click: finish', 'esc: cancel'])) self.canvas.draw() - elif event.key == 'escape': + elif event_key == 'escape': if self._add_vertex: self._add_vertex = False else: @@ -661,11 +562,11 @@ class Moulder(FigureCanvasQTAgg): line.set_animated(False) line.set_color([0, 0, 0, 0]) self.canvas.draw() - elif event.key == 'r': + elif event_key == 'r': self.modelax.set_xlim(self.area[:2]) self.modelax.set_ylim(self.area[2:]) self._update_data_plot() - elif event.key == 'a': + elif event_key == 'a': self._add_vertex = not self._add_vertex def _mouse_move_callback(self, event): diff --git a/moulder.py b/moulder.py index c6d0099..42a4cfa 100644 --- a/moulder.py +++ b/moulder.py @@ -29,10 +29,19 @@ class MoulderApp(QMainWindow): self.init_ui() self.set_callbacks() + self.canvas = Moulder(self, width=5, height=4, dpi=100) #self.canvas = GravityModelCanvas(self, # width=5, height=4, dpi=100) - #self.canvas.setFocus() - #self.setCentralWidget(self.canvas) + # self.canvas.setFocus() + self.setCentralWidget(self.canvas) + + def keyPressEvent(self, event): + print(event) + keys_dict = {Qt.Key_N: "n", Qt.Key_R: "r", + Qt.Key_A: "a", Qt.Key_D: "d", + Qt.Key_Escape: "escape"} + if self.canvas.plotted and event.key in keys_dict.keys(): + self.canvas.key_press_callback(keys_dict[event.key]) def closeEvent(self, event): event.ignore() @@ -91,10 +100,8 @@ class MoulderApp(QMainWindow): new_model_dialog = NewModelDialog(parent=self) new_model_dialog.exec_() if new_model_dialog.is_completed(): - x = new_model_dialog.x - z = new_model_dialog.z - area = (x.min(), x.max(), z.min(), 10000) - self.canvas = Moulder(self, area, x, z) + self.canvas.x = new_model_dialog.x + self.canvas.z = new_model_dialog.z self.canvas.run() self.setCentralWidget(self.canvas)