moulder/interactive.py

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


class Moulder(FigureCanvasQTAgg):
    """
    Interactive 2D forward modeling using polygons.

    A matplotlib GUI application. Allows drawing and manipulating polygons and
    computes their predicted data automatically. Also permits contaminating the
    data with gaussian pseudo-random error for producing synthetic data sets.

    Uses :mod:`fatiando.gravmag.talwani` for computations.

    *Moulder* objects can be persisted to Python pickle files using the
    :meth:`~fatiando.gravmag.interactive.Moulder.save` method and later
    restored using :meth:`~fatiando.gravmag.interactive.Moulder.load`.

    .. warning::

        Cannot be used with ``%matplotlib inline`` on IPython notebooks because
        the app uses the matplotlib plot window. You can still embed the
        generated model and data figure on notebooks using the
        :meth:`~fatiando.gravmag.interactive.Moulder.plot` method.

    Parameters:

    * area : list = (x1, x2, z1, z2)
        The limits of the model drawing area, in meters.
    * x, z : 1d-arrays
        The x- and z-coordinates of the computation points (places where
        predicted data will be computed). In meters.
    * data : None or 1d-array
        Observed data measured at *x* and *z*. Will plot this with black dots
        along the predicted data.
    * density_range : list = [min, max]
        The minimum and maximum values allowed for the density. Determines the
        limits of the density slider of the application. In kg.m^-3. Defaults
        to [-2000, 2000].
    * kwargs : dict
        Other keyword arguments used to restore the state of the application.
        Used by the :meth:`~fatiando.gravmag.interactive.Moulder.load` method.
        Not intended for general use.

    Examples:

    Make the Moulder object and start the app::

        import numpy as np
        area = (0, 10e3, 0, 5e3)
        # Calculate on 100 points
        x = np.linspace(area[0], area[1], 100)
        z = np.zeros_like(x)
        app = Moulder(area, x, z)
        app.run()
        # This will pop-up a window with the application (like the screenshot
        # below). Start drawing (follow the instruction in the figure title).
        # When satisfied, close the window to resume execution.

    .. image:: ../_static/images/Moulder-screenshot.png
        :alt: Screenshot of the Moulder GUI


    After closing the plot window, you can access the model and data from the
    *Moulder* object::

        app.model  # The drawn model as fatiando.mesher.Polygon
        app.predicted  # 1d-array with the data predicted by the model
        # You can save the predicted data to use later
        app.save_predicted('data.txt')
        # You can also save the application and resume it later
        app.save('application.pkl')
        # Close this session/IPython notebook/etc.
        # To resume drawing later:
        app = Moulder.load('application.pkl')
        app.run()

    """

    # 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', 'esc: cancel'])



    def __init__(self, parent, area, x, z, data=None,
                 density_range=[-2000, 2000], width=5, height=4, dpi=100,
                 **kwargs):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)

        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.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)

    def save_predicted(self, fname):
        """
        Save the predicted data to a text file.

        Data will be saved in 3 columns separated by spaces: x  z  data

        Parameters:

        * fname : string or file-like object
            The name of the output file or an open file-like object.

        """
        numpy.savetxt(fname, numpy.transpose([self.x, self.z, self.predicted]))

    def save(self, fname):
        """
        Save the application state into a pickle file.

        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.

        Parameters:

        * 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

    @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 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
        self._lastevent = None
        self._drawing = False
        self._add_vertex = False
        self._xy = []
        self._drawing_plot = None
        # Used to blit the model plot and make
        # rendering faster
        self.background = None
        # Connect event callbacks
        self._connect()
        self._update_data()
        self._update_data_plot()
        self.canvas.draw()
        #pyplot.show()

    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('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):
        """
        Setup the plot figure with labels, titles, ticks, etc.

        Sets the *canvas*, *dataax*, *modelax*, *polygons* and *lines*
        attributes.

        Parameters:

        * kwargs : dict
            Keyword arguments passed to ``pyplot.subplots``.

        """
        sharex = kwargs.get('sharex')
        if not sharex:
            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_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.grid(True)
        ax2.invert_yaxis()
        ax2.set_ylabel('z (m)')
        #self.fig.subplots_adjust(top=0.95, left=0.1, right=0.95, bottom=0.06,
        #                         hspace=0.1)
        self.canvas = self.fig.canvas
        self.dataax = axes[0]
        self.modelax = axes[1]
        self.fig.canvas.draw()

    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, **self.line_args)
        return poly, line

    def _update_data(self):
        """
        Recalculate the predicted data (optionally with random error)
        """
        self.predicted = talwani.gz(self.x, self.z, self.model)
        if self.error > 0:
            self.predicted = utils.contaminate(self.predicted, self.error)

    def _update_data_plot(self):
        """
        Update the predicted data plot in the *dataax*.

        Adjusts the xlim of the axes to fit the data.
        """
        self.predicted_line.set_ydata(self.predicted)
        vmin = 1.2*min(self.predicted.min(), self.dmin)
        vmax = 1.2*max(self.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()
        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:
                    density = self.density_slider.val
                    poly, line = self._make_polygon(self._xy, density)
                    self.polygons.append(poly)
                    self.lines.append(line)
                    self.densities.append(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):
        """
        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([], [], **self.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.area[:2])
            self.modelax.set_ylim(self.area[2:])
            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)