Skip to content

Instantly share code, notes, and snippets.

@niallrobinson
Forked from ajdawson/cube_explorer.py
Last active December 18, 2015 03:49
Show Gist options
  • Save niallrobinson/5721433 to your computer and use it in GitHub Desktop.
Save niallrobinson/5721433 to your computer and use it in GitHub Desktop.
an attempt at classifying and adding a picker. Doesn't work but you get the idea
'''
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
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 = None
self.plot_kwargs = None
self.butts = {}
self.butt_fns = {}
self.make_plot(self, *args, **kwargs)
def show(self):
plt.show()
def make_plot(self, *args, **kwargs):
"""
Makes initial plot
"""
self.fig = plt.figure(num=None)
# pop self
args = args[1:]
# Make the initial plot.
self.ax = self.fig.add_subplot(111)
self.plot_func(self.disp_data, axes=self.ax, *args, **kwargs)
self.plot_args = args
self.plot_kwargs = kwargs
if self.axes_hook is not None:
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 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
"""
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.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]])
def add_animate_buttons(self, dim, button_names_tup, slot=0, refresh_rate=0.2):
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]])
def _get_ani_fns(self, dim, refresh_rate):
def play(event):
print "Play"
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:
print self.current_slice[dim]
self.current_slice[dim] += 1
print self.current_slice[dim]
else:
self.current_slice[dim] = 0
self._refresh_plot()
def stop(event):
print "Stop"
self.playing = False
return play, stop
def _refresh_plot(self):
self.disp_data = self.cube.data[tuple(self.current_slice)]
this_plot_fn = getattr(self.ax, self.plot_func.__name__)
self.ax.clear()
this_plot_fn(self.disp_data, *self.plot_args, **self.plot_kwargs)
if self.axes_hook is not None:
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):
print "Up"
if self.current_slice[dim] < self.cube.shape[dim]-1:
print self.current_slice[dim]
self.current_slice[dim] += 1
print self.current_slice[dim]
elif circular:
self.current_slice[dim] = 0
self._refresh_plot()
elif inc_or_dec is 'dec':
def fn(event):
print "Down"
if self.current_slice[dim] > 0:
print self.current_slice[dim]
self.current_slice[dim] -= 1
print self.current_slice[dim]
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, plt.pcolormesh, [0, slice(None), slice(None)], cmap=mpl_cm.get_cmap('brewer_OrRd_09'), axes_hook=axes_hook, vmin=290, vmax=310)
# ce = CubeExplorer(cube, plt.plot, [slice(None), 50, 50], axes_hook=axes_hook)
ce.add_nav_buttons(0, ("Up", "Down"), slot=0)
ce.add_animate_buttons(0, ("Play", "Stop"), slot=1)
# ce.add_nav_buttons(1, ("Lon Up", "Lon Down"), slot=0)
# ce.add_nav_buttons(2, ("Lat Up", "Lat Down"), slot=1)
ce.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment