mirror of
https://github.com/fatiando/moulder.git
synced 2025-12-21 02:23:46 +08:00
515 lines
18 KiB
Python
515 lines
18 KiB
Python
from __future__ import division, absolute_import
|
|
from future.builtins import super, zip
|
|
|
|
try:
|
|
import cPickle as pickle
|
|
except ImportError:
|
|
import pickle
|
|
|
|
import os
|
|
import sys
|
|
import numpy
|
|
import matplotlib
|
|
from matplotlib import pyplot, widgets, patches
|
|
from matplotlib.lines import Line2D
|
|
from matplotlib.figure import Figure
|
|
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
|
|
from PyQt5.QtGui import QIcon
|
|
from PyQt5.QtCore import Qt
|
|
from PyQt5.QtWidgets import QSizePolicy, QMainWindow, QApplication, QAction
|
|
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):
|
|
|
|
# The tolerance range for mouse clicks on vertices. In pixels.
|
|
epsilon = 5
|
|
# App instructions printed in the figure suptitle
|
|
instructions = ' | '.join([
|
|
'n: New polygon', 'd: delete', 'click: select/move', 'a: add vertex',
|
|
'r: reset view', 'esc: cancel'])
|
|
|
|
def __init__(self, parent, x, z, min_depth, max_depth,
|
|
density_range=[-2000, 2000], width=5, height=4, dpi=100):
|
|
self.fig = Figure(figsize=(width, height), dpi=dpi)
|
|
super().__init__(self.fig)
|
|
self.setParent(parent)
|
|
|
|
self.min_depth, self.max_depth = min_depth, max_depth
|
|
self._x, self._z = x, z
|
|
self.density_range = density_range
|
|
self._predicted = numpy.zeros_like(x)
|
|
self._data = None
|
|
self.cmap = pyplot.cm.RdBu_r
|
|
self.canvas = self.fig.canvas
|
|
|
|
self.polygons = []
|
|
self.lines = []
|
|
self.densities = []
|
|
|
|
# Set arbitrary density and error values (only for first implementations)
|
|
# It will be determined by sliders/entries in MoulderApp
|
|
self._density = 100
|
|
self._error = 0
|
|
|
|
# Data min and max (only for first implementations)
|
|
# They will be determined when data is imported
|
|
self.dmin, self.dmax = 0, 0
|
|
|
|
self._figure_setup()
|
|
self._init_markers()
|
|
self._connect()
|
|
|
|
@property
|
|
def x(self):
|
|
return self._x
|
|
|
|
@x.setter
|
|
def x(self, new_value):
|
|
self._x = numpy.asarray(new_value)
|
|
|
|
@property
|
|
def z(self):
|
|
return self._z
|
|
|
|
@z.setter
|
|
def z(self, new_value):
|
|
self._z = numpy.asarray(new_value)
|
|
|
|
@property
|
|
def data(self):
|
|
return self._data
|
|
|
|
@data.setter
|
|
def data(self, new_data):
|
|
self._data = new_data
|
|
|
|
@property
|
|
def density(self):
|
|
return self._density
|
|
|
|
@density.setter
|
|
def density(self, new_value):
|
|
self._density = new_value
|
|
|
|
@property
|
|
def error(self):
|
|
return self._error
|
|
|
|
@error.setter
|
|
def error(self, new_value):
|
|
self._error = new_value
|
|
|
|
@property
|
|
def predicted(self):
|
|
self._predicted = talwani.gz(self.x, self.z, self.model)
|
|
if self.error > 0:
|
|
self._predicted = utils.contaminate(self.predicted, self.error)
|
|
return self._predicted
|
|
|
|
@property
|
|
def model(self):
|
|
"""
|
|
The polygon model drawn as :class:`fatiando.mesher.Polygon` objects.
|
|
"""
|
|
m = [Polygon(p.xy, {'density': d}, force_clockwise=True)
|
|
for p, d in zip(self.polygons, self.densities)]
|
|
return m
|
|
|
|
def set_meassurement_points(self, x, z):
|
|
self.x = x
|
|
self.z = z
|
|
self.modelax.set_xlim(self.x.min(), self.x.max())
|
|
self.modelax.set_ylim(self.max_depth, self.min_depth) # y inverted axe
|
|
self.predicted_line.remove()
|
|
self.predicted_line, = self.dataax.plot(self.x, self.predicted, '-r')
|
|
self._update_data_plot()
|
|
|
|
def _figure_setup(self):
|
|
self.dataax, self.modelax = self.fig.subplots(2, 1, sharex=True)
|
|
self.dataax.set_ylabel("Gravity Anomaly [mGal]")
|
|
self.dataax.set_ylim((-200, 200))
|
|
self.dataax.grid(True)
|
|
self.modelax.set_xlabel("x [m]")
|
|
self.modelax.set_ylabel("z [m]")
|
|
self.modelax.set_xlim(self.x.min(), self.x.max())
|
|
self.modelax.set_ylim(self.min_depth, self.max_depth)
|
|
self.modelax.grid(True)
|
|
self.modelax.invert_yaxis()
|
|
self.predicted_line, = self.dataax.plot(self.x, self.predicted, '-r')
|
|
self.canvas.draw()
|
|
|
|
def _init_markers(self):
|
|
self._ivert = None
|
|
self._ipoly = None
|
|
self._lastevent = None
|
|
self._drawing = False
|
|
self._add_vertex = False
|
|
self._xy = []
|
|
self._drawing_plot = None
|
|
self.background = None
|
|
|
|
def _connect(self):
|
|
"""
|
|
Connect the matplotlib events to their callback methods.
|
|
"""
|
|
# Make the proper callback connections
|
|
self.canvas.mpl_connect('button_press_event',
|
|
self._button_press_callback)
|
|
self.canvas.mpl_connect('button_release_event',
|
|
self._button_release_callback)
|
|
self.canvas.mpl_connect('motion_notify_event',
|
|
self._mouse_move_callback)
|
|
|
|
def _density2color(self, density):
|
|
"""
|
|
Map density values to colors using the given *cmap* attribute.
|
|
|
|
Parameters:
|
|
|
|
* density : 1d-array
|
|
The density values of the model polygons
|
|
|
|
Returns
|
|
|
|
* colors : 1d-array
|
|
The colors mapped to each density value (returned by a matplotlib
|
|
colormap object.
|
|
|
|
"""
|
|
dmin, dmax = self.density_range
|
|
return self.cmap((density - dmin)/(dmax - dmin))
|
|
|
|
def _make_polygon(self, vertices, density):
|
|
"""
|
|
Create a polygon for drawing.
|
|
|
|
Polygons are matplitlib.patches.Polygon objects for the fill and
|
|
matplotlib.lines.Line2D for the contour.
|
|
|
|
Parameters:
|
|
|
|
* vertices : list of [x, z]
|
|
List of the [x, z] coordinate pairs of each vertex of the polygon
|
|
* density : float
|
|
The density of the polygon (used to set the color)
|
|
|
|
Returns:
|
|
|
|
* polygon, line
|
|
The matplotlib Polygon and Line2D objects
|
|
|
|
"""
|
|
poly = patches.Polygon(vertices, animated=False, alpha=0.9,
|
|
color=self._density2color(density))
|
|
x, y = list(zip(*poly.xy))
|
|
line = Line2D(x, y, **LINE_ARGS)
|
|
return poly, line
|
|
|
|
def _update_data_plot(self):
|
|
"""
|
|
Update the predicted data plot in the *dataax*.
|
|
|
|
Adjusts the xlim of the axes to fit the data.
|
|
"""
|
|
predicted = self.predicted
|
|
self.predicted_line.set_ydata(predicted)
|
|
vmin = 1.2*min(predicted.min(), self.dmin)
|
|
vmax = 1.2*max(predicted.max(), self.dmax)
|
|
self.dataax.set_ylim(vmin, vmax)
|
|
self.dataax.grid(True)
|
|
self.canvas.draw()
|
|
|
|
def _set_error_callback(self, value):
|
|
"""
|
|
Callback when error slider is edited
|
|
"""
|
|
self.error = value
|
|
self._update_data_plot()
|
|
|
|
def _set_density_callback(self, value):
|
|
"""
|
|
Callback when density slider is edited
|
|
"""
|
|
if self._ipoly is not None:
|
|
self.densities[self._ipoly] = value
|
|
self.polygons[self._ipoly].set_color(self._density2color(value))
|
|
# self._update_data()
|
|
self._update_data_plot()
|
|
self.canvas.draw()
|
|
|
|
def _get_polygon_vertice_id(self, event):
|
|
"""
|
|
Find out which vertex of which polygon the event happened in.
|
|
|
|
If the click was inside a polygon (not on a vertex), identify that
|
|
polygon.
|
|
|
|
Returns:
|
|
|
|
* p, v : int, int
|
|
p: the index of the polygon the event happened in or None if
|
|
outside all polygons.
|
|
v: the index of the polygon vertex that was clicked or None if the
|
|
click was not on a vertex.
|
|
|
|
"""
|
|
distances = []
|
|
indices = []
|
|
for poly in self.polygons:
|
|
x, y = poly.get_transform().transform(poly.xy).T
|
|
d = numpy.sqrt((x - event.x)**2 + (y - event.y)**2)
|
|
distances.append(d.min())
|
|
indices.append(numpy.argmin(d))
|
|
p = numpy.argmin(distances)
|
|
if distances[p] >= self.epsilon:
|
|
# Check if the event was inside a polygon
|
|
x, y = event.x, event.y
|
|
p, v = None, None
|
|
for i, poly in enumerate(self.polygons):
|
|
if poly.contains_point([x, y]):
|
|
p = i
|
|
break
|
|
else:
|
|
v = indices[p]
|
|
last = len(self.polygons[p].xy) - 1
|
|
if v == 0 or v == last:
|
|
v = [0, last]
|
|
return p, v
|
|
|
|
def _add_new_vertex(self, event):
|
|
"""
|
|
Add new vertex to polygon
|
|
"""
|
|
vertices = self.polygons[self._ipoly].get_xy()
|
|
x, y = vertices[:, 0], vertices[:, 1]
|
|
# Compute the angle between the vectors to each pair of
|
|
# vertices corresponding to each line segment of the polygon
|
|
x1, y1 = x[:-1], y[:-1]
|
|
x2, y2 = numpy.roll(x1, -1), numpy.roll(y1, -1)
|
|
u = numpy.vstack((x1 - event.xdata, y1 - event.ydata)).T
|
|
v = numpy.vstack((x2 - event.xdata, y2 - event.ydata)).T
|
|
angle = numpy.arccos(numpy.sum(u*v, 1) /
|
|
numpy.sqrt(numpy.sum(u**2, 1)) /
|
|
numpy.sqrt(numpy.sum(v**2, 1)))
|
|
position = angle.argmax() + 1
|
|
x = numpy.hstack((x[:position], event.xdata, x[position:]))
|
|
y = numpy.hstack((y[:position], event.ydata, y[position:]))
|
|
new_vertices = numpy.vstack((x, y)).T
|
|
return new_vertices
|
|
|
|
def _button_press_callback(self, event):
|
|
"""
|
|
What actions to perform when a mouse button is clicked
|
|
"""
|
|
if event.inaxes != self.modelax:
|
|
return
|
|
if event.button == 1 and not self._drawing and self.polygons:
|
|
self._lastevent = event
|
|
if not self._add_vertex:
|
|
for line, poly in zip(self.lines, self.polygons):
|
|
poly.set_animated(False)
|
|
line.set_animated(False)
|
|
line.set_color([0, 0, 0, 0])
|
|
self.canvas.draw()
|
|
# Find out if a click happened on a vertice
|
|
# and which vertice of which polygon
|
|
self._ipoly, self._ivert = self._get_polygon_vertice_id(event)
|
|
if self._ipoly is not None:
|
|
# self.density_slider.set_val(self.densities[self._ipoly])
|
|
self.polygons[self._ipoly].set_animated(True)
|
|
self.lines[self._ipoly].set_animated(True)
|
|
self.lines[self._ipoly].set_color([0, 1, 0, 0])
|
|
self.canvas.draw()
|
|
self.background = self.canvas.copy_from_bbox(
|
|
self.modelax.bbox)
|
|
self.modelax.draw_artist(self.polygons[self._ipoly])
|
|
self.modelax.draw_artist(self.lines[self._ipoly])
|
|
self.canvas.blit(self.modelax.bbox)
|
|
else:
|
|
# If a polygon is selected, we will add a new vertex by
|
|
# removing the polygon and inserting a new one with the extra
|
|
# vertex.
|
|
if self._ipoly is not None:
|
|
vertices = self._add_new_vertex(event)
|
|
density = self.densities[self._ipoly]
|
|
polygon, line = self._make_polygon(vertices, density)
|
|
self.polygons[self._ipoly].remove()
|
|
self.lines[self._ipoly].remove()
|
|
self.polygons.pop(self._ipoly)
|
|
self.lines.pop(self._ipoly)
|
|
self.polygons.insert(self._ipoly, polygon)
|
|
self.lines.insert(self._ipoly, line)
|
|
self.modelax.add_patch(polygon)
|
|
self.modelax.add_line(line)
|
|
self.lines[self._ipoly].set_color([0, 1, 0, 0])
|
|
self.canvas.draw()
|
|
# self._update_data()
|
|
self._update_data_plot()
|
|
elif self._drawing:
|
|
if event.button == 1:
|
|
self._xy.append([event.xdata, event.ydata])
|
|
self._drawing_plot.set_data(list(zip(*self._xy)))
|
|
self.canvas.restore_region(self.background)
|
|
self.modelax.draw_artist(self._drawing_plot)
|
|
self.canvas.blit(self.modelax.bbox)
|
|
elif event.button == 3:
|
|
if len(self._xy) >= 3:
|
|
poly, line = self._make_polygon(self._xy, self.density)
|
|
self.polygons.append(poly)
|
|
self.lines.append(line)
|
|
self.densities.append(self.density)
|
|
self.modelax.add_patch(poly)
|
|
self.modelax.add_line(line)
|
|
self._drawing_plot.remove()
|
|
self._drawing_plot = None
|
|
self._xy = None
|
|
self._drawing = False
|
|
self._ipoly = len(self.polygons) - 1
|
|
self.lines[self._ipoly].set_color([0, 1, 0, 0])
|
|
self.dataax.set_title(self.instructions)
|
|
self.canvas.draw()
|
|
# self._update_data()
|
|
self._update_data_plot()
|
|
|
|
def _button_release_callback(self, event):
|
|
"""
|
|
Reset place markers on mouse button release
|
|
"""
|
|
if event.inaxes != self.modelax:
|
|
return
|
|
if event.button != 1:
|
|
return
|
|
if self._add_vertex:
|
|
self._add_vertex = False
|
|
if self._ivert is None and self._ipoly is None:
|
|
return
|
|
self.background = None
|
|
for line, poly in zip(self.lines, self.polygons):
|
|
poly.set_animated(False)
|
|
line.set_animated(False)
|
|
self.canvas.draw()
|
|
self._ivert = None
|
|
# self._ipoly is only released when clicking outside
|
|
# the polygons
|
|
self._lastevent = None
|
|
# self._update_data()
|
|
self._update_data_plot()
|
|
|
|
def _key_press_callback(self, event_key):
|
|
"""
|
|
What to do when a key is pressed on the keyboard.
|
|
"""
|
|
if event_key == 'd':
|
|
if self._drawing and self._xy:
|
|
self._xy.pop()
|
|
if self._xy:
|
|
self._drawing_plot.set_data(list(zip(*self._xy)))
|
|
else:
|
|
self._drawing_plot.set_data([], [])
|
|
self.canvas.restore_region(self.background)
|
|
self.modelax.draw_artist(self._drawing_plot)
|
|
self.canvas.blit(self.modelax.bbox)
|
|
elif self._ivert is not None:
|
|
poly = self.polygons[self._ipoly]
|
|
line = self.lines[self._ipoly]
|
|
if len(poly.xy) > 4:
|
|
verts = numpy.atleast_1d(self._ivert)
|
|
poly.xy = numpy.array([xy for i, xy in enumerate(poly.xy)
|
|
if i not in verts])
|
|
line.set_data(list(zip(*poly.xy)))
|
|
# self._update_data()
|
|
self._update_data_plot()
|
|
self.canvas.restore_region(self.background)
|
|
self.modelax.draw_artist(poly)
|
|
self.modelax.draw_artist(line)
|
|
self.canvas.blit(self.modelax.bbox)
|
|
self._ivert = None
|
|
elif self._ipoly is not None:
|
|
self.polygons[self._ipoly].remove()
|
|
self.lines[self._ipoly].remove()
|
|
self.polygons.pop(self._ipoly)
|
|
self.lines.pop(self._ipoly)
|
|
self.densities.pop(self._ipoly)
|
|
self._ipoly = None
|
|
self.canvas.draw()
|
|
# self._update_data()
|
|
self._update_data_plot()
|
|
elif event_key == 'n':
|
|
self._ivert = None
|
|
self._ipoly = None
|
|
for line, poly in zip(self.lines, self.polygons):
|
|
poly.set_animated(False)
|
|
line.set_animated(False)
|
|
line.set_color([0, 0, 0, 0])
|
|
self.canvas.draw()
|
|
self.background = self.canvas.copy_from_bbox(self.modelax.bbox)
|
|
self._drawing = True
|
|
self._xy = []
|
|
self._drawing_plot = Line2D([], [], **LINE_ARGS)
|
|
self._drawing_plot.set_animated(True)
|
|
self.modelax.add_line(self._drawing_plot)
|
|
self.dataax.set_title(' | '.join([
|
|
'left click: set vertice', 'right click: finish',
|
|
'esc: cancel']))
|
|
self.canvas.draw()
|
|
elif event_key == 'escape':
|
|
if self._add_vertex:
|
|
self._add_vertex = False
|
|
else:
|
|
self._drawing = False
|
|
self._xy = []
|
|
if self._drawing_plot is not None:
|
|
self._drawing_plot.remove()
|
|
self._drawing_plot = None
|
|
for line, poly in zip(self.lines, self.polygons):
|
|
poly.set_animated(False)
|
|
line.set_animated(False)
|
|
line.set_color([0, 0, 0, 0])
|
|
self.canvas.draw()
|
|
elif event_key == 'r':
|
|
self.modelax.set_xlim(self.x.min(), self.x.max())
|
|
self.modelax.set_ylim(self.min_depth, self.max_depth)
|
|
self._update_data_plot()
|
|
elif event_key == 'a':
|
|
self._add_vertex = not self._add_vertex
|
|
|
|
def _mouse_move_callback(self, event):
|
|
"""
|
|
Handle things when the mouse move.
|
|
"""
|
|
if event.inaxes != self.modelax:
|
|
return
|
|
if event.button != 1:
|
|
return
|
|
if self._ivert is None and self._ipoly is None:
|
|
return
|
|
if self._add_vertex:
|
|
return
|
|
x, y = event.xdata, event.ydata
|
|
p = self._ipoly
|
|
v = self._ivert
|
|
if self._ivert is not None:
|
|
self.polygons[p].xy[v] = x, y
|
|
else:
|
|
dx = x - self._lastevent.xdata
|
|
dy = y - self._lastevent.ydata
|
|
self.polygons[p].xy[:, 0] += dx
|
|
self.polygons[p].xy[:, 1] += dy
|
|
self.lines[p].set_data(list(zip(*self.polygons[p].xy)))
|
|
self._lastevent = event
|
|
self.canvas.restore_region(self.background)
|
|
self.modelax.draw_artist(self.polygons[p])
|
|
self.modelax.draw_artist(self.lines[p])
|
|
self.canvas.blit(self.modelax.bbox)
|