moulder/interactive.py

from __future__ import division, absolute_import
from future.builtins import super, zip

import numpy
from matplotlib import pyplot, patches
from matplotlib.lines import Line2D
from matplotlib.figure import Figure
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
from PyQt5.QtCore import Qt, pyqtSignal

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'])

    # Signal when selected polygon changes
    polygon_selected = pyqtSignal(float)

    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)
        self.dataax, self.modelax = self.fig.subplots(2, 1, sharex=True)
        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.predicted_line = None
        self.cmap = pyplot.cm.RdBu_r
        self.canvas = self.fig.canvas

        self.polygons = []
        self.lines = []
        self.densities = []

        # Initialize density and error values
        self._density = 0
        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, value):
        """
        Callback when density slider is edited
        """
        self._density = value
        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()

    @property
    def error(self):
        return self._error

    @error.setter
    def error(self, value):
        """
        Callback when error slider is edited
        """
        self._error = value
        self._update_data_plot()

    @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 add_vertex(self):
        self._add_vertex = not self._add_vertex

    def new_polygon(self):
        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()

    def delete_polygon(self):
        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()

    def cancel_drawing(self):
        if self._add_vertex:
            self._add_vertex = False
        else:
            self.dataax.set_title(self.instructions)
            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()

    def set_meassurement_points(self, x, z):
        self.x = x
        self.z = z
        self._figure_setup()
        self._update_data_plot()

    def _figure_setup(self):
        self.dataax.set_title(self.instructions)
        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()
        if self.predicted_line is not None:
            self.predicted_line.remove()
        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 _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:
                    # Emit signal: selected polygon changed (sends density)
                    self.polygon_selected.emit(self.densities[self._ipoly])
                    # 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 _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)

    def keyPressEvent(self, event):
        """
        What to do when a key is pressed on the keyboard.
        """
        if event.key() == Qt.Key_D:
            self.delete_polygon()
        elif event.key() == Qt.Key_N:
            self.new_polygon()
        elif event.key() == Qt.Key_Escape:
            self.cancel_drawing()
        elif event.key() == Qt.Key_R:
            self.modelax.set_xlim(self.x.min(), self.x.max())
            self.modelax.set_ylim(self.max_depth, self.min_depth)
            self._update_data_plot()
        elif event.key() == Qt.Key_A:
            self.add_vertex()