def read_plot_ts_slice (file_path, grid, lon0=None, lat0=None, time_index=None, t_start=None, t_end=None, time_average=False, hmin=None, hmax=None, zmin=None, zmax=None, tmin=None, tmax=None, smin=None, smax=None, date_string=None, fig_name=None, second_file_path=None):
# Make sure we'll end up with a single record in time
if time_index is None and not time_average:
print 'Error (read_plot_ts_slice): either specify time_index or set time_average=True.'
sys.exit()
if date_string is None and time_index is not None:
# Determine what to write about the date
date_string = parse_date(file_path=file_path, time_index=time_index)
if not isinstance(grid, Grid):
# This is the path to the NetCDF grid file, not a Grid object
# Make a grid object from it
grid = Grid(grid)
# Read temperature
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path, 'THETA')
else:
file_path_use = file_path
temp = mask_3d(read_netcdf(file_path_use, 'THETA', time_index=time_index, t_start=t_start, t_end=t_end, time_average=time_average), grid)
# Read salinity
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path, 'SALT')
else:
file_path_use = file_path
salt = mask_3d(read_netcdf(file_path_use, 'SALT', time_index=time_index, t_start=t_start, t_end=t_end, time_average=time_average), grid)
# Plot
ts_slice_plot(temp, salt, grid, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, tmin=tmin, tmax=tmax, smin=smin, smax=smax, date_string=date_string, fig_name=fig_name)
def read_plot_slice (var, file_path, grid, lon0=None, lat0=None, time_index=None, t_start=None, t_end=None, time_average=False, hmin=None, hmax=None, zmin=None, zmax=None, vmin=None, vmax=None, date_string=None, fig_name=None, second_file_path=None):
# Make sure we'll end up with a single record in time
if time_index is None and not time_average:
print 'Error (read_plot_slice): either specify time_index or set time_average=True.'
sys.exit()
if date_string is None and time_index is not None:
# Determine what to write about the date
date_string = parse_date(file_path=file_path, time_index=time_index)
if not isinstance(grid, Grid):
# This is the path to the NetCDF grid file, not a Grid object
# Make a grid object from it
grid = Grid(grid)
# Read necessary variables from NetCDF file and mask appropriately
if var in ['temp', 'tminustf']:
# Read temperature. Some of these variables need more than temperature and so second_file_path might be set.
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path, 'THETA')
else:
file_path_use = file_path
temp = mask_3d(read_netcdf(file_path_use, 'THETA', time_index=time_index, t_start=t_start, t_end=t_end, time_average=time_average), grid)
if var in ['salt', 'tminustf']:
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path, 'SALT')
else:
file_path_use = file_path
salt = mask_3d(read_netcdf(file_path_use, 'SALT', time_index=time_index, t_start=t_start, t_end=t_end, time_average=time_average), grid)
if var == 'u':
u = mask_3d(read_netcdf(file_path, 'UVEL', time_index=time_index, t_start=t_start, t_end=t_end, time_average=time_average), grid, gtype='u')
if var == 'v':
v = mask_3d(read_netcdf(file_path, 'VVEL', time_index=time_index, t_start=t_start, t_end=t_end, time_average=time_average), grid, gtype='v')
# Plot
if var == 'temp':
slice_plot(temp, grid, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, vmin=vmin, vmax=vmax, title=r'Temperature ($^{\circ}$C)', date_string=date_string, fig_name=fig_name)
elif var == 'salt':
slice_plot(salt, grid, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, vmin=vmin, vmax=vmax, title='Salinity (psu)', date_string=date_string, fig_name=fig_name)
elif var == 'tminustf':
slice_plot(t_minus_tf(temp, salt, grid), grid, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, vmin=vmin, vmax=vmax, ctype='plusminus', title=r'Difference from in-situ freezing point ($^{\circ}$C)', date_string=date_string, fig_name=fig_name)
elif var == 'u':
slice_plot(u, grid, gtype='u', lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, vmin=vmin, vmax=vmax, ctype='plusminus', title='Zonal velocity (m/s)', date_string=date_string, fig_name=fig_name)
elif var == 'v':
slice_plot(v, grid, gtype='v', lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, vmin=vmin, vmax=vmax, ctype='plusminus', title='Zonal velocity (m/s)', date_string=date_string, fig_name=fig_name)
else:
print 'Error (read_plot_slice): variable key ' + str(var) + ' does not exist'
sys.exit()
def plot_aice_minmax(file_path, grid, year, fig_name=None, monthly=True):
if not isinstance(grid, Grid):
# This is the path to the NetCDF grid file, not a Grid object
# Make a grid object from it
grid = Grid(grid)
# Read sea ice area and the corresponding dates
aice = mask_land_zice(read_netcdf(file_path, 'SIarea'),
grid,
time_dependent=True)
time = netcdf_time(file_path, monthly=monthly)
# Find the range of dates we care about
t_start, t_end = select_year(time, year)
# Trim the arrays to these dates
aice = aice[t_start:t_end, :]
time = time[t_start:t_end]
# Find the indices of min and max sea ice area
t_min, t_max = find_aice_min_max(aice, grid)
# Wrap up in lists for easy iteration
aice_minmax = [aice[t_min, :], aice[t_max, :]]
time_minmax = [time[t_min], time[t_max]]
# Plot
fig, gs, cax = set_panels('1x2C1')
for t in range(2):
lon, lat, aice_plot = cell_boundaries(aice_minmax[t], grid)
ax = plt.subplot(gs[0, t])
shade_land_zice(ax, grid)
img = ax.pcolormesh(lon, lat, aice_plot, vmin=0, vmax=1)
latlon_axes(ax, lon, lat)
if t == 1:
# Don't need latitude labels a second time
ax.set_yticklabels([])
plt.title(parse_date(date=time_minmax[t]), fontsize=18)
# Colourbar
plt.colorbar(img, cax=cax, orientation='horizontal')
# Main title above
plt.suptitle('Min and max sea ice area', fontsize=22)
finished_plot(fig, fig_name=fig_name)
def plot_everything_diff (output_dir='./', baseline_dir=None, timeseries_file='timeseries.nc', grid_path='../grid/', fig_dir='.', option='last_year', unravelled=False, monthly=True, key='WSS'):
# Check that baseline_dir is set
# It's a keyword argument on purpose so that the user can't mix up which simulation is which.
if baseline_dir is None:
print 'Error (plot_everything_diff): must set baseline_dir'
sys.exit()
# Make sure proper directories, and rename so 1=baseline, 2=current
output_dir_1 = real_dir(baseline_dir)
output_dir_2 = real_dir(output_dir)
fig_dir = real_dir(fig_dir)
# Build lists of output files in each directory
output_files_1 = build_file_list(output_dir_1, unravelled=unravelled)
output_files_2 = build_file_list(output_dir_2, unravelled=unravelled)
# Build the grid
grid = Grid(grid_path)
# Timeseries through the entire simulation
var_names = ['fris_mass_balance', 'eta_avg', 'seaice_area', 'fris_temp', 'fris_salt', 'fris_age']
for var in var_names:
read_plot_timeseries_diff(var, output_dir_1+timeseries_file, output_dir_2+timeseries_file, precomputed=True, fig_name=fig_dir+'timeseries_'+var+'_diff.png', monthly=monthly)
# Now figure out which time indices to use for plots with no time dependence
file_path_1, file_path_2, time_index_1, time_index_2, t_start_1, t_start_2, t_end_1, t_end_2, time_average = select_common_time(output_files_1, output_files_2, option=option, monthly=monthly)
# Set date string
if option == 'last_year':
date_string = 'year beginning ' + parse_date(file_path=file_path_1, time_index=t_start_1)
elif option == 'last_month':
date_string = parse_date(file_path=file_path_1, time_index=time_index_1)
# Now make lat-lon plots
var_names = ['ismr', 'bwtemp', 'bwsalt', 'sst', 'sss', 'aice', 'hice', 'hsnow', 'mld', 'eta', 'vel', 'velice', 'bwage', 'iceprod']
if key == 'WSK':
figsize = (10,6)
else:
figsize = (8,6)
for var in var_names:
if var == 'iceprod':
vmin = -2
vmax = 2
else:
vmin = None
vmax = None
read_plot_latlon_diff(var, file_path_1, file_path_2, grid=grid, time_index=time_index_1, t_start=t_start_1, t_end=t_end_1, time_average=time_average, time_index_2=time_index_2, t_start_2=t_start_2, t_end_2=t_end_2, date_string=date_string, fig_name=fig_dir+var+'_diff.png', figsize=figsize, vmin=vmin, vmax=vmax)
# Zoom into some variables
if key=='WSK' and var in ['ismr', 'bwtemp', 'bwsalt', 'vel', 'bwage']:
if var == 'bwage':
vmin = -5
vmax = None
else:
vmin = None
vmax = None
read_plot_latlon_diff(var, file_path_1, file_path_2, grid=grid, time_index=time_index_1, t_start=t_start_1, t_end=t_end_1, time_average=time_average, time_index_2=time_index_2, t_start_2=t_start_2, t_end_2=t_end_2, zoom_fris=True, date_string=date_string, fig_name=fig_dir+var+'_zoom_diff.png', vmin=vmin, vmax=vmax)
if var == 'vel':
# Call the other options for vertical transformations
for vel_option in ['sfc', 'bottom']:
read_plot_latlon_diff(var, file_path_1, file_path_2, grid=grid, time_index=time_index_1, t_start=t_start_1, t_end=t_end_1, time_average=time_average, time_index_2=time_index_2, t_start_2=t_start_2, t_end_2=t_end_2, vel_option=vel_option, date_string=date_string, fig_name=fig_dir+var+'_'+vel_option+'_diff.png')
# Slice plots
read_plot_ts_slice_diff(file_path_1, file_path_2, grid=grid, lon0=-40, hmax=-75, zmin=-1450, time_index=time_index_1, t_start=t_start_1, t_end=t_end_1, time_average=time_average, time_index_2=time_index_2, t_start_2=t_start_2, t_end_2=t_end_2, date_string=date_string, fig_name=fig_dir+'ts_slice_filchner_diff.png')
read_plot_ts_slice_diff(file_path_1, file_path_2, grid=grid, lon0=-55, hmax=-72, time_index=time_index_1, t_start=t_start_1, t_end=t_end_1, time_average=time_average, time_index_2=time_index_2, t_start_2=t_start_2, t_end_2=t_end_2, date_string=date_string, fig_name=fig_dir+'ts_slice_ronne_diff.png')
if key == 'WSK':
read_plot_ts_slice_diff(file_path_1, file_path_2, grid=grid, lon0=-25, zmin=-2000, time_index=time_index_1, t_start=t_start_1, t_end=t_end_1, time_average=time_average, time_index_2=time_index_2, t_start_2=t_start_2, t_end_2=t_end_2, date_string=date_string, fig_name=fig_dir+'ts_slice_eweddell_diff.png')
def read_plot_latlon(var,
file_path,
grid,
time_index=None,
t_start=None,
t_end=None,
time_average=False,
vmin=None,
vmax=None,
zoom_fris=False,
xmin=None,
xmax=None,
ymin=None,
ymax=None,
date_string=None,
fig_name=None,
second_file_path=None,
change_points=None,
tf_option='min',
vel_option='avg'):
# Make sure we'll end up with a single record in time
if time_index is None and not time_average:
print 'Error (read_plot_latlon): either specify time_index or set time_average=True.'
sys.exit()
if date_string is None and time_index is not None:
# Determine what to write about the date
date_string = parse_date(file_path=file_path, time_index=time_index)
if not isinstance(grid, Grid):
# This is the path to the NetCDF grid file, not a Grid object
# Make a grid object from it
grid = Grid(grid)
# Read necessary variables from NetCDF file(s), and mask appropriately
if var == 'ismr':
shifwflx = mask_except_zice(
read_netcdf(file_path,
'SHIfwFlx',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average), grid)
if var in ['bwtemp', 'sst', 'tminustf']:
# Read temperature. Some of these variables need more than temperature and so second_file_path might be set.
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path, 'THETA')
else:
file_path_use = file_path
temp = mask_3d(
read_netcdf(file_path_use,
'THETA',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average), grid)
if var in ['bwsalt', 'sss', 'tminustf']:
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path, 'SALT')
else:
file_path_use = file_path
salt = mask_3d(
read_netcdf(file_path_use,
'SALT',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average), grid)
if var == 'aice':
aice = mask_land_zice(
read_netcdf(file_path,
'SIarea',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average), grid)
if var == 'hice':
hice = mask_land_zice(
read_netcdf(file_path,
'SIheff',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average), grid)
if var == 'mld':
mld = mask_land_zice(
read_netcdf(file_path,
'MXLDEPTH',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average), grid)
if var == 'eta':
eta = mask_land_zice(
read_netcdf(file_path,
'ETAN',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average), grid)
if var == 'saltflx':
saltflx = mask_land_zice(
read_netcdf(file_path,
'SIempmr',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average), grid)
if var == 'vel':
# First read u
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path, 'UVEL')
else:
file_path_use = file_path
u = mask_3d(read_netcdf(file_path_use,
'UVEL',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average),
grid,
gtype='u')
# Now read v
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path, 'VVEL')
else:
file_path_use = file_path
v = mask_3d(read_netcdf(file_path_use,
'VVEL',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average),
grid,
gtype='v')
if var == 'velice':
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path,
'SIuice')
else:
file_path_use = file_path
uice = mask_land_zice(read_netcdf(file_path_use,
'SIuice',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average),
grid,
gtype='u')
if second_file_path is not None:
file_path_use = find_variable(file_path, second_file_path,
'SIvice')
else:
file_path_use = file_path
vice = mask_land_zice(read_netcdf(file_path_use,
'SIvice',
time_index=time_index,
t_start=t_start,
t_end=t_end,
time_average=time_average),
grid,
gtype='v')
# Plot
if var == 'ismr':
plot_ismr(shifwflx,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
change_points=change_points,
date_string=date_string,
fig_name=fig_name)
elif var == 'bwtemp':
plot_bw('temp',
temp,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'bwsalt':
plot_bw('salt',
salt,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'sst':
plot_ss('temp',
temp,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'sss':
plot_ss('salt',
salt,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'aice':
plot_2d_noshelf('aice',
aice,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'hice':
plot_2d_noshelf('hice',
hice,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'mld':
plot_2d_noshelf('mld',
mld,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'eta':
plot_2d_noshelf('eta',
eta,
grid,
ctype='plusminus',
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'saltflx':
plot_2d_noshelf('saltflx',
saltflx,
grid,
ctype='plusminus',
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'tminustf':
plot_tminustf(temp,
salt,
grid,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
tf_option=tf_option,
date_string=date_string,
fig_name=fig_name)
elif var == 'vel':
plot_vel(u,
v,
grid,
vel_option=vel_option,
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
elif var == 'velice':
plot_vel(uice,
vice,
grid,
vel_option='ice',
vmin=vmin,
vmax=vmax,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
date_string=date_string,
fig_name=fig_name)
else:
print 'Error (read_plot_latlon): variable key ' + str(
var) + ' does not exist'
sys.exit()
def animate_latlon_coupled (var, output_dir='./', file_name='output.nc', segment_dir=None, vmin=None, vmax=None, change_points=None, mov_name=None, fig_name_beg=None, fig_name_end=None, figsize=(8,6)):
import matplotlib.animation as animation
output_dir = real_dir(output_dir)
segment_dir = check_segment_dir(output_dir, segment_dir)
file_paths = segment_file_paths(output_dir, segment_dir, file_name)
# Inner function to read and process data from a single file
def read_process_data (file_path, var_name, grid, mask_option='3d', gtype='t', lev_option=None, ismr=False, psi=False):
data = read_netcdf(file_path, var_name)
if mask_option == '3d':
data = mask_3d(data, grid, gtype=gtype, time_dependent=True)
elif mask_option == 'except_ice':
data = mask_except_ice(data, grid, gtype=gtype, time_dependent=True)
elif mask_option == 'land':
data = mask_land(data, grid, gtype=gtype, time_dependent=True)
elif mask_option == 'land_ice':
data = mask_land_ice(data, grid, gtype=gtype, time_dependent=True)
else:
print 'Error (read_process_data): invalid mask_option ' + mask_option
sys.exit()
if lev_option is not None:
if lev_option == 'top':
data = select_top(data)
elif lev_option == 'bottom':
data = select_bottom(data)
else:
print 'Error (read_process_data): invalid lev_option ' + lev_option
sys.exit()
if ismr:
data = convert_ismr(data)
if psi:
data = np.sum(data, axis=-3)*1e-6
return data
all_data = []
all_grids = []
all_dates = []
# Loop over segments
for file_path in file_paths:
print 'Processing ' + file_path
# Build the grid
grid = Grid(file_path)
# Read and process the variable we need
ctype = 'basic'
gtype = 't'
include_shelf = var not in ['aice', 'hice', 'mld']
if var == 'ismr':
data = read_process_data(file_path, 'SHIfwFlx', grid, mask_option='except_ice', ismr=True)
title = 'Ice shelf melt rate (m/y)'
ctype = 'ismr'
elif var == 'bwtemp':
data = read_process_data(file_path, 'THETA', grid, lev_option='bottom')
title = 'Bottom water temperature ('+deg_string+'C)'
elif var == 'bwsalt':
data = read_process_data(file_path, 'SALT', grid, lev_option='bottom')
title = 'Bottom water salinity (psu)'
elif var == 'draft':
data = mask_except_ice(grid.draft, grid)
title = 'Ice shelf draft (m)'
elif var == 'aice':
data = read_process_data(file_path, 'SIarea', grid, mask_option='land_ice')
title = 'Sea ice concentration'
elif var == 'hice':
data = read_process_data(file_path, 'SIheff', grid, mask_option='land_ice')
title = 'Sea ice thickness (m)'
elif var == 'mld':
data = read_process_data(file_path, 'MXLDEPTH', grid, mask_option='land_ice')
title = 'Mixed layer depth (m)'
elif var == 'eta':
data = read_process_data(file_path, 'ETAN', grid, mask_option='land')
title = 'Free surface (m)'
elif var == 'psi':
data = read_process_data(file_path, 'PsiVEL', grid, psi=True)
title = 'Vertically integrated streamfunction (Sv)'
ctype = 'plusminus'
else:
print 'Error (animate_latlon): invalid var ' + var
sys.exit()
# Loop over timesteps
if var == 'draft':
# Just one timestep
all_data.append(data)
all_grids.append(grid)
all_dates.append(parse_date(file_path=file_path, time_index=0))
else:
for t in range(data.shape[0]):
# Extract the data from this timestep
# Save it and the grid to the long lists
all_data.append(data[t,:])
all_grids.append(grid)
all_dates.append(parse_date(file_path=file_path, time_index=t))
extend = get_extend(vmin=vmin, vmax=vmax)
if vmin is None:
vmin = np.amax(data)
for elm in all_data:
vmin = min(vmin, np.amin(elm))
if vmax is None:
vmax = np.amin(data)
for elm in all_data:
vmax = max(vmax, np.amax(elm))
num_frames = len(all_data)
# Make the first and last frames as stills
tsteps = [0, -1]
fig_names = [fig_name_beg, fig_name_end]
for t in range(2):
latlon_plot(all_data[tsteps[t]], all_grids[tsteps[t]], gtype=gtype, ctype=ctype, vmin=vmin, vmax=vmax, change_points=change_points, title=title, date_string=all_dates[tsteps[t]], figsize=figsize, fig_name=fig_names[t])
# Now make the animation
fig, ax = plt.subplots(figsize=figsize)
# Inner function to plot a frame
def plot_one_frame (t):
img = latlon_plot(all_data[t], all_grids[t], ax=ax, gtype=gtype, ctype=ctype, vmin=vmin, vmax=vmax, change_points=change_points, title=title+'\n'+all_dates[t], make_cbar=False)
if t==0:
return img
# First frame
img = plot_one_frame(0)
plt.colorbar(img, extend=extend)
# Function to update figure with the given frame
def animate(t):
ax.cla()
plot_one_frame(t)
# Call this for each frame
anim = animation.FuncAnimation(fig, func=animate, frames=range(num_frames))
writer = animation.FFMpegWriter(bitrate=500, fps=10)
anim.save(mov_name, writer=writer)
if mov_name is not None:
print 'Saving ' + mov_name
anim.save(mov_name, writer=writer)
else:
plt.show()