In [1]:
import glob
filenames=glob.glob('/data/20130606/grids/LMA_20130606_0[2-4]*_flash_extent.nc')
print(filenames)
In [2]:
save_files = False
if save_files:
import matplotlib
matplotlib.use('Agg')
else:
%matplotlib qt
In [3]:
import numpy as np
from datetime import datetime, timedelta
from lmatools.vis.multiples_nc import centers_to_edges
import pupynere as nc
import itertools
from lmatools.grid.grid_collection import LMAgridFileCollection
In [4]:
"""
Here we will provide a class to accept a file collection object
and produce a series of plots, one per time frame.
As a matplotlib animation object?
"""
def update_pcolor(pcolor, x, y, C):
"""updates coordinates and scalar data for a pcolormesh plot
basically a hybrid of matplotlib.axes.Axes.pcolormesh and matplotlib.collections.QuadMesh
# Tests that updated pcolormesh is equivalent to creating a new pcolormesh
import matplotlib.pyplot as plt
from acuity.MPLutils.managerhelpers import update_pcolor
import numpy as np
x = np.arange(10)
y = np.arange(10)
x, y = np.meshgrid(x,y)
z=x*y
pc = plt.pcolormesh(x,y,z, shading='flat')
plt.show()
# x += 5
# y += 5
# z = z**2.0
x2 = x[:-5,:]
y2 = y[:-5,:]
z2 = z[:-5,:]
update_pcolor(pc,x2,y2,z2)
if True:
pc2 = plt.pcolormesh(x2,y2,z2, shading='flat')
print 'Comparing attrs in pc2'
for key in pc2.__dict__:
if np.asarray(pc2.__dict__[key] != pc.__dict__[key]).all(): print key
print '------------------------'
print 'Comparing attrs in pc'
for key in pc.__dict__:
if np.asarray(pc2.__dict__[key] != pc.__dict__[key]).all(): print key
plt.draw()
"""
import matplotlib.transforms as transforms
from matplotlib.collections import QuadMesh
assert isinstance(pcolor, QuadMesh)
# Axes.pcolormesh.__init__()
Ny, Nx = x.shape
coords = np.zeros(((Nx * Ny), 2), dtype=float)
# no effort made to handle masked arrays here
C = np.ravel(C[0:Ny-1, 0:Nx-1])
coords[:, 0] = x.ravel()
coords[:, 1] = y.ravel()
# QuadMesh.__init__()
pcolor._meshWidth = Nx-1
pcolor._meshHeight = Ny-1
pcolor._bbox = transforms.Bbox.unit()
pcolor._bbox.update_from_data_xy(coords.reshape(
(Nx * Ny, 2)))
coords=coords.reshape((Ny, Nx, 2))
pcolor._coordinates = coords
# end quadmesh init
if pcolor.get_array() != None:
pcolor.set_array(C)
class GridAnimation(object):
def __init__(self, gridcollection, vmin, vmax, ax=None, titler=None, grid_label='Data', log=False):
""" gridcollection is an iterable that yields
t, xedge, yedge, density
where t is a datetime object, and xedge, yedge, and density
are ready for use in pcolormesh.
The optional titler function accepts a datetime object for the current frame,
and should return the title to be set on the axes for each frame.
If ax is not passed in, a new figure with a single set of axes
will be created along with a colorbar.
The class can used as follows, for an instance of this class a:
matplotlib.animation.FuncAnimation(a.ax.figure, a.animate, frames=a.framer,
init_func=a.setup, interval=20, blit=True)
"""
self.gridcollection = gridcollection
if titler is None:
self.titler = self.title_default
else:
self.titler = titler
self.grid_label = grid_label
if ax is not None:
self.ax = axes
else:
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from matplotlib.colors import LogNorm, Normalize
from matplotlib.colorbar import ColorbarBase, make_axes
fig = plt.figure()
self.ax = fig.add_subplot(1,1,1)
cbar_ax, cbar_kw = make_axes(self.ax)
self.cbar_ax = cbar_ax
if log:
self.norm = LogNorm(vmin=vmin, vmax=vmax)
else:
self.norm = Normalize(vmin=vmin, vmax=vmax)
cbar_kw['norm'] = self.norm
self.cmap = get_cmap('gist_earth')
cbar_kw['cmap'] = self.cmap
self.pc = None
self.cbar = ColorbarBase(self.cbar_ax, **cbar_kw)
self.cbar.set_label(self.grid_label)
def title_default(self, t):
return "{0}".format(t)
def setup(self):
title_art = self.ax.set_title('')
if self.pc is None:
# self.ax.artists.remove(self.pc)
x, y = np.meshgrid(np.array([0,1]),
np.array([0,1]))
c = np.ones((1,1))
self.pc = self.ax.pcolormesh(x,y,c, cmap=self.cmap, norm=self.norm)
return self.pc, title_art
def animate(self, payload):
t, xedge, yedge, data = payload
x,y = np.meshgrid(xedge, yedge)
update_pcolor(self.pc, x,y,data)
title_art = self.ax.set_title(self.titler(t))
return self.pc, title_art
def framer(self):
for a in self.gridcollection:
yield a
# self.area_range = (0.0, 100000.0)
# self.rate_range = (0, 10000)
# self.source_range = (0, 100000)
# self.init_range = (0,100)
In [5]:
# names are 'x', 'y' for data that have been map projected
# or 'longtiude', 'latitude' otherwise
NCs = LMAgridFileCollection(filenames, 'flash_extent', x_name='x', y_name='y')
In [6]:
GA = GridAnimation(NCs, 1, 1000, ax=None, titler=None, grid_label='Flashes per pixel', log=True)
In [7]:
t = datetime(2013,6,6,2,30,0)
xedge, yedge, data = NCs.data_for_time(t)
print xedge.min(), yedge.min()
In [8]:
#GA.ax.axis((xedge.min(), xedge.max(), yedge.min(), yedge.max()))
GA.ax.axis((-200.0, 200, -200.0, 200.0))
from matplotlib.animation import FuncAnimation
FA = FuncAnimation(GA.ax.figure, GA.animate, frames=GA.framer,
init_func=GA.setup, interval=100.0, blit=False)
In [9]:
if save_files:
# It's not clear that the anim's save function knows to save the right number of frames.
# Sometimes, for short animations, it seems to start resaving frames.
import matplotlib.animation as anim
anim.writers.list()
FA.save('0529-0530.png', writer='imagemagick_file')#, clear_temp=False)
else:
# FIXME: Blitting of the title artist isn't working right.
# Is the text object being recreated each time?
# If so, try changing the text artist instead of recreating new text with set_title?
from matplotlib.pyplot import show
show()
In [ ]:
In [ ]:
In [ ]:
In [ ]: