Created
July 11, 2013 12:25
-
-
Save niallrobinson/5975003 to your computer and use it in GitHub Desktop.
Cube explorer
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 iris | |
| import iris.plot as iplt | |
| import matplotlib.pyplot as plt | |
| from matplotlib.widgets import Button | |
| import numpy as np | |
| import matplotlib.cm as mpl_cm | |
| import iris.quickplot as qplt | |
| import time | |
| from iris.exceptions import CoordinateNotFoundError | |
| import sys | |
| class CubeExplorer(object): | |
| def __init__(self, cube, plot_func, current_slice, *args, **kwargs): | |
| """ | |
| Args: | |
| * cube: cube of data to plot | |
| * plot_func: pointer to plotting function | |
| * current_slice: index tuple which gives a slice of cube | |
| which is compatible with cube | |
| """ | |
| self.cube = cube | |
| self.plot_func = plot_func | |
| self.current_slice = current_slice | |
| self.disp_data = cube.data[tuple(current_slice)] | |
| self.ax = None | |
| self.axes_hook = kwargs.pop('axes_hook', None) | |
| self.plot_args = args | |
| self.plot_kwargs = kwargs | |
| self.butts = {} | |
| self.butt_fns = {} | |
| self.make_plot() | |
| def show(self): | |
| plt.show() | |
| 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 add_nav_buttons(self, dim, button_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 | |
| * button_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=0.85) | |
| if type(self.current_slice[dim]) is slice: | |
| raise TypeError("Cannot iterate over a displayed dimension") | |
| self.butts[button_names_tup[0]] = Button(plt.axes([0.875, 0.85-(slot*0.15), 0.11, 0.05]), button_names_tup[0]) | |
| self.butt_fns[button_names_tup[0]] = self._get_nav_fn(dim, 'inc', circular) | |
| self.butts[button_names_tup[0]].on_clicked(self.butt_fns[button_names_tup[0]]) | |
| self.butts[button_names_tup[1]] = Button(plt.axes([0.875, 0.79-(slot*0.15), 0.11, 0.05]), button_names_tup[1]) | |
| self.butt_fns[button_names_tup[1]] = self._get_nav_fn(dim, 'dec', circular) | |
| self.butts[button_names_tup[1]].on_clicked(self.butt_fns[button_names_tup[1]]) | |
| # set axis back to plot | |
| plt.sca(self.ax) | |
| def add_animate_buttons(self, dim, button_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 | |
| * button_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[button_names_tup[0]] = Button(plt.axes([0.875, 0.85-(slot*0.15), 0.11, 0.05]), button_names_tup[0]) | |
| self.butts[button_names_tup[1]] = Button(plt.axes([0.875, 0.79-(slot*0.15), 0.11, 0.05]), button_names_tup[1]) | |
| play_fn, stop_fn = self._get_ani_fns(dim, refresh_rate) | |
| self.butt_fns[button_names_tup[0]] = play_fn | |
| self.butt_fns[button_names_tup[1]] = stop_fn | |
| self.butts[button_names_tup[0]].on_clicked(self.butt_fns[button_names_tup[0]]) | |
| self.butts[button_names_tup[1]].on_clicked(self.butt_fns[button_names_tup[1]]) | |
| # set axis back to plot | |
| plt.sca(self.ax) | |
| 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 | |
| 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 | |
| if __name__ == '__main__': | |
| cube = iris.load_cube(iris.sample_data_path('GloSea4', 'ensemble_001.pp')) | |
| def axes_hook(ax): | |
| ax.coastlines() | |
| ax.set_title('Depth slices') | |
| ce = CubeExplorer(cube, qplt.plot, [0, 50, slice(None)])#, cmap=mpl_cm.get_cmap('brewer_OrRd_09'), axes_hook=axes_hook) | |
| ce.add_nav_buttons('time', ("Up", "Down"), slot=0) | |
| ce.add_animate_buttons(0, ("Play", "Stop"), slot=1) | |
| ce.add_na |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
sorry - last line is wrong due to copy/paste incompetency. Gist is being flakey and I can't update it, however. Just replace line 191 with ce.show()