Last active
December 19, 2015 15:28
-
-
Save niallrobinson/5976287 to your computer and use it in GitHub Desktop.
cube explorer with picker
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| ''' | |
| Proof of concept class for making interactive plot object which allow the addition | |
| of navigation buttons. | |
| Created on Jul 3, 2013 | |
| @author: nrobin | |
| ''' | |
| import matplotlib.cm as mpl_cm | |
| import matplotlib.pyplot as plt | |
| from matplotlib.widgets import Button | |
| import numpy as np | |
| import iris | |
| from iris.exceptions import CoordinateNotFoundError | |
| import iris.plot as iplt | |
| import iris.quickplot as qplt | |
| class ButtonSetup(object): | |
| """ | |
| Convenience class to store setup for Cube | |
| Explorer buttons. | |
| Args: | |
| * col_margin: | |
| Value for setting the amount of space | |
| on the right of the plot for displaying | |
| buttons. | |
| * initial_but_a_pos: | |
| Positions of button a axis | |
| in slot one. | |
| * initial_but_b_pos: | |
| Positions of button b axis | |
| in slot one. | |
| * slot_diff: | |
| The vertical offset between sets | |
| of buttons. | |
| """ | |
| def __init__(self, | |
| col_margin = 0.85, | |
| initial_but_a_pos = [0.875, 0.85, 0.11, 0.05], | |
| initial_but_b_pos = [0.875, 0.79, 0.11, 0.05], | |
| slot_diff = 0.15): | |
| self.col_margin = col_margin | |
| self.initial_but_a_pos = initial_but_a_pos | |
| self.initial_but_b_pos = initial_but_b_pos | |
| self.slot_diff = slot_diff | |
| class CubeExplorer(object): | |
| def __init__(self, cube, plot_func, current_slice, *args, **kwargs): | |
| """ | |
| Args: | |
| * cube: | |
| Cube of data to plot. | |
| * plot_func: | |
| Pointer to a plotting function | |
| which is compatible with the slice defined | |
| by current_slice. | |
| * current_slice: | |
| Index tuple which gives a slice of cube | |
| which is compatible with plot_func. | |
| """ | |
| self.cube = cube | |
| self.plot_func = plot_func | |
| self.current_slice = current_slice | |
| self.axes_hook = kwargs.pop('axes_hook', None) | |
| self.button_setup = kwargs.pop('button_setup', | |
| ButtonSetup()) | |
| # self._ax is set when the plot is made | |
| self._ax = None | |
| self._plot_args = args | |
| self._plot_kwargs = kwargs | |
| # _butts and _butt_fns are added to with | |
| # various class methods | |
| self._butts = {} | |
| self._butt_fns = {} | |
| self._make_plot() | |
| def show(self): | |
| plt.show() | |
| def add_nav_buttons(self, dim, names_tup, slot=0, circular=False): | |
| """ | |
| Adds a set of two buttons to the plot window which allow incrementing | |
| or decrementing over the specified dimension | |
| Args: | |
| * dim: | |
| Dimension number or name to traverse. | |
| * names_tup: | |
| Tuple of two strings to be the names of the | |
| increment and decrement buttons respectively. | |
| * slot: | |
| Level of the plot to display this button set. | |
| * circular: | |
| Boolean - to loop round when limit is reached or not. | |
| """ | |
| if type(dim) is str: | |
| dim, = self.cube.coord_dims(self.cube.coord(dim)) | |
| plt.subplots_adjust(right=self.button_setup.col_margin) | |
| if type(self.current_slice[dim]) is slice: | |
| raise TypeError("Cannot iterate over a displayed dimension") | |
| but_pos_a = self.button_setup.initial_but_a_pos | |
| but_pos_a[1] -= slot*self.button_setup.slot_diff | |
| self._butts[names_tup[0]] = Button(plt.axes(but_pos_a), names_tup[0]) | |
| self._butt_fns[names_tup[0]] = self._get_nav_fn(dim, 'inc', circular) | |
| self._butts[names_tup[0]].on_clicked(self._butt_fns[names_tup[0]]) | |
| but_pos_b = self.button_setup.initial_but_b_pos | |
| but_pos_b[1] -= slot*self.button_setup.slot_diff | |
| self._butts[names_tup[1]] = Button(plt.axes(but_pos_b), names_tup[1]) | |
| self._butt_fns[names_tup[1]] = self._get_nav_fn(dim, 'dec', circular) | |
| self._butts[names_tup[1]].on_clicked(self._butt_fns[names_tup[1]]) | |
| # set axis back to plot | |
| plt.sca(self._ax) | |
| def add_animate_buttons(self, dim, names_tup, slot=0, refresh_rate=0.2): | |
| """ | |
| Adds a set of two buttons to start/stop cycling through a dimension | |
| of a cube. | |
| Args: | |
| * dim: | |
| Dimension number/name to traverse. | |
| * names_tup: | |
| Tuple of two strings to be the names of the | |
| play and stop buttons respectively. | |
| * slot: | |
| Level of the plot to display this button set. | |
| * refresh_rate: | |
| Number of seconds to wait between each frame. | |
| """ | |
| if type(dim) is str: | |
| dim, = self.cube.coord_dims(self.cube.coord(dim)) | |
| plt.subplots_adjust(right=0.85) | |
| self._butts[names_tup[0]] = Button(plt.axes([0.875, 0.85-(slot*0.15), 0.11, 0.05]), names_tup[0]) | |
| self._butts[names_tup[1]] = Button(plt.axes([0.875, 0.79-(slot*0.15), 0.11, 0.05]), names_tup[1]) | |
| play_fn, stop_fn = self._get_ani_fns(dim, refresh_rate) | |
| self._butt_fns[names_tup[0]] = play_fn | |
| self._butt_fns[names_tup[1]] = stop_fn | |
| self._butts[names_tup[0]].on_clicked(self._butt_fns[names_tup[0]]) | |
| self._butts[names_tup[1]].on_clicked(self._butt_fns[names_tup[1]]) | |
| # set axis back to plot | |
| plt.sca(self._ax) | |
| def add_picker(self, plot_func, *args, **kwargs): | |
| """ | |
| Adds picker functionality to a Cube Explorer object. | |
| The user can then select points from which to display | |
| data in one of the non-displayed dimension in a pop | |
| up plot by clicking on the cube explorer plot. | |
| Args: | |
| * plot_func: | |
| A plotting function that accepts a cube | |
| consisting of all the non-displayed dimensions. | |
| """ | |
| ax_hook = kwargs.pop('ax_hook', None) | |
| # add picker option to plot | |
| self._plot_kwargs['picker'] = kwargs.pop('picker', True) | |
| self._refresh_plot() | |
| def _on_pick(event): | |
| if event.artist.get_axes() != self._ax: | |
| return | |
| s = [slice(None)] * len(self.cube.shape) | |
| picker_plot_dims = [i for i, v in enumerate(self.current_slice) if type(v) is slice] | |
| xpt = event.mouseevent.xdata | |
| s[picker_plot_dims[0]] = int(xpt) | |
| if len(picker_plot_dims) == 2: | |
| ypt = event.mouseevent.ydata | |
| s[picker_plot_dims[1]] = int(ypt) | |
| plt.figure(num=None) | |
| plot_func(self.cube[tuple(s)], *args, **kwargs) | |
| if ax_hook != None: | |
| ax = plt.gca() | |
| ax_hook(ax) | |
| plt.show() | |
| self._ax.figure.canvas.mpl_connect('pick_event', _on_pick) | |
| def _make_plot(self): | |
| """ | |
| Makes initial plot | |
| """ | |
| self.fig = plt.figure(num=None) | |
| # Make the initial plot. | |
| self.pl = self.plot_func(self.cube[tuple(self.current_slice)], *self._plot_args, **self._plot_kwargs) | |
| self._ax = plt.gca() | |
| if self.axes_hook is not None: | |
| self.axes_hook(self._ax) | |
| def _refresh_plot(self): | |
| """ | |
| Refreshes the displayed plot to display the slice defined | |
| by current_slice. | |
| """ | |
| self._ax.clear() | |
| self.plot_func(self.cube[tuple(self.current_slice)], *self._plot_args, **self._plot_kwargs) | |
| if self.axes_hook is not None: | |
| self.axes_hook(self._ax) | |
| self.fig.canvas.draw() | |
| def _get_nav_fn(self, dim, inc_or_dec, circular): | |
| """ | |
| Returns increment and decrement button functions for a dimension. | |
| """ | |
| if inc_or_dec is 'inc': | |
| def fn(event): | |
| if self.current_slice[dim] < self.cube.shape[dim]-1: | |
| self.current_slice[dim] += 1 | |
| elif circular: | |
| self.current_slice[dim] = 0 | |
| self._refresh_plot() | |
| elif inc_or_dec is 'dec': | |
| def fn(event): | |
| if self.current_slice[dim] > 0: | |
| self.current_slice[dim] -= 1 | |
| elif circular: | |
| self.current_slice[dim] = self.cube.shape[dim] | |
| self._refresh_plot() | |
| return fn | |
| def _get_ani_fns(self, dim, refresh_rate): | |
| def play(event): | |
| self.playing = True | |
| while True: | |
| self.fig.canvas.start_event_loop(timeout=refresh_rate) | |
| if self.playing: | |
| if self.current_slice[dim] < self.cube.shape[dim]-1: | |
| self.current_slice[dim] += 1 | |
| else: | |
| self.current_slice[dim] = 0 | |
| self._refresh_plot() | |
| def stop(event): | |
| self.playing = False | |
| return play, stop | |
| if __name__ == '__main__': | |
| cube = iris.load_cube(iris.sample_data_path('GloSea4', 'ensemble_001.pp')) | |
| def axes_hook_main(ax): | |
| ax.coastlines() | |
| ax.set_title('Depth slices') | |
| def axes_hook_picker(ax): | |
| ax.set_title('Time series') | |
| ce = CubeExplorer(cube, iplt.pcolormesh, [0, slice(None), slice(None)], cmap=mpl_cm.get_cmap('brewer_OrRd_09'), axes_hook=axes_hook_main) | |
| ce.add_nav_buttons('time', ("Up", "Down"), slot=0) | |
| ce.add_animate_buttons(0, ("Play", "Stop"), slot=1) | |
| ce.add_picker(qplt.pcolormesh, ax_hook=axes_hook_picker) | |
| ce.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment