from argparse import ArgumentParser
from os import environ
import numpy as np
import xarray as xr
from parcels import Field
from parcels.plotting import cartopy_colorbar
from parcels.plotting import create_parcelsfig_axis
from parcels.plotting import plotfield
try:
import matplotlib.animation as animation
from matplotlib import rc
except:
anim = None
[docs]def plotTrajectoriesFile(filename, mode='2d', tracerfile=None, tracerfield='P',
tracerlon='x', tracerlat='y', recordedvar=None, movie_forward=True,
bins=20, show_plt=True, central_longitude=0):
"""Quick and simple plotting of Parcels trajectories
:param filename: Name of Parcels-generated NetCDF file with particle positions
:param mode: Type of plot to show. Supported are '2d', '3d', 'hist2d',
'movie2d' and 'movie2d_notebook'. The latter two give animations,
with 'movie2d_notebook' specifically designed for jupyter notebooks
:param tracerfile: Name of NetCDF file to show as background
:param tracerfield: Name of variable to show as background
:param tracerlon: Name of longitude dimension of variable to show as background
:param tracerlat: Name of latitude dimension of variable to show as background
:param recordedvar: Name of variable used to color particles in scatter-plot.
Only works in 'movie2d' or 'movie2d_notebook' mode.
:param movie_forward: Boolean whether to show movie in forward or backward mode (default True)
:param bins: Number of bins to use in `hist2d` mode. See also https://matplotlib.org/api/_as_gen/matplotlib.pyplot.hist2d.html
:param show_plt: Boolean whether plot should directly be show (for py.test)
:param central_longitude: Degrees East at which to center the plot
"""
environ["HDF5_USE_FILE_LOCKING"] = "FALSE"
try:
pfile = xr.open_dataset(str(filename), decode_cf=True)
except:
pfile = xr.open_dataset(str(filename), decode_cf=False)
lon = np.ma.filled(pfile.variables['lon'], np.nan)
lat = np.ma.filled(pfile.variables['lat'], np.nan)
time = np.ma.filled(pfile.variables['time'], np.nan)
z = np.ma.filled(pfile.variables['z'], np.nan)
mesh = pfile.attrs['parcels_mesh'] if 'parcels_mesh' in pfile.attrs else 'spherical'
if(recordedvar is not None):
record = np.ma.filled(pfile.variables[recordedvar], np.nan)
pfile.close()
if tracerfile is not None and mode != 'hist2d':
tracerfld = Field.from_netcdf(tracerfile, tracerfield, {'lon': tracerlon, 'lat': tracerlat})
plt, fig, ax, cartopy = plotfield(tracerfld)
if plt is None:
return # creating axes was not possible
titlestr = ' and ' + tracerfield
else:
spherical = False if mode == '3d' or mesh == 'flat' else True
plt, fig, ax, cartopy = create_parcelsfig_axis(spherical=spherical, central_longitude=central_longitude)
if plt is None:
return # creating axes was not possible
titlestr = ''
if cartopy:
for p in range(lon.shape[1]):
lon[:, p] = [ln if ln < 180 else ln - 360 for ln in lon[:, p]]
if mode == '3d':
from mpl_toolkits.mplot3d import Axes3D # noqa
plt.clf() # clear the figure
ax = fig.gca(projection='3d')
for p in range(len(lon)):
ax.plot(lon[p, :], lat[p, :], z[p, :], '.-')
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')
ax.set_zlabel('Depth')
ax.set_title('Particle trajectories')
elif mode == '2d':
if cartopy:
ax.plot(np.transpose(lon), np.transpose(lat), '.-', transform=cartopy.crs.Geodetic())
else:
ax.plot(np.transpose(lon), np.transpose(lat), '.-')
ax.set_title('Particle trajectories' + titlestr)
elif mode == 'hist2d':
_, _, _, cs = plt.hist2d(lon[~np.isnan(lon)], lat[~np.isnan(lat)], bins=bins)
cartopy_colorbar(cs, plt, fig, ax)
ax.set_title('Particle histogram')
elif mode in ('movie2d', 'movie2d_notebook'):
ax.set_xlim(np.nanmin((lon+central_longitude+180) % 360 - 180), np.nanmax((lon+central_longitude+180) % 360 - 180))
ax.set_ylim(np.nanmin(lat), np.nanmax(lat))
plottimes = np.unique(time)
if not movie_forward:
plottimes = np.flip(plottimes, 0)
if isinstance(plottimes[0], (np.datetime64, np.timedelta64)):
plottimes = plottimes[~np.isnat(plottimes)]
else:
try:
plottimes = plottimes[~np.isnan(plottimes)]
except:
pass
b = time == plottimes[0]
def timestr(plottimes, index):
if isinstance(plottimes[index], np.timedelta64):
if plottimes[-1] > np.timedelta64(1, 'h'):
return str(plottimes[index].astype('timedelta64[h]'))
elif plottimes[-1] > np.timedelta64(1, 's'):
return str(plottimes[index].astype('timedelta64[s]'))
else:
return str(plottimes[index])
if cartopy:
scat = ax.scatter(lon[b], lat[b], s=20, color='k', transform=cartopy.crs.Geodetic())
else:
scat = ax.scatter(lon[b], lat[b], s=20, color='k')
ttl = ax.set_title('Particles' + titlestr + ' at time ' + timestr(plottimes, 0))
frames = np.arange(0, len(plottimes))
def animate(t):
b = time == plottimes[t]
scat.set_offsets(np.vstack((lon[b], lat[b])).transpose())
ttl.set_text('Particle' + titlestr + ' at time ' + timestr(plottimes, t))
if recordedvar is not None:
scat.set_array(record[b])
return scat,
rc('animation', html='html5')
anim = animation.FuncAnimation(fig, animate, frames=frames, interval=100, blit=False)
else:
raise RuntimeError('mode %s not known' % mode)
if mode == 'movie2d_notebook':
plt.close()
return anim
else:
if show_plt:
plt.show()
return plt
if __name__ == "__main__":
p = ArgumentParser(description="""Quick and simple plotting of Parcels trajectories""")
p.add_argument('mode', choices=('2d', '3d', 'hist2d', 'movie2d', 'movie2d_notebook'), nargs='?',
default='movie2d', help='Type of display')
p.add_argument('-p', '--particlefile', type=str, default='MyParticle.nc',
help='Name of particle file')
p.add_argument('-f', '--tracerfile', type=str, default=None,
help='Name of tracer file to display underneath particle trajectories')
p.add_argument('-flon', '--tracerfilelon', type=str, default='x',
help='Name of longitude dimension in tracer file')
p.add_argument('-flat', '--tracerfilelat', type=str, default='y',
help='Name of latitude dimension in tracer file')
p.add_argument('-ffld', '--tracerfilefield', type=str, default='P',
help='Name of field in tracer file')
p.add_argument('-r', '--recordedvar', type=str, default=None,
help='Name of a variable recorded along trajectory')
p.add_argument('-bins', type=int, default=20,
help='Number of bins for mode=hist2d')
args = p.parse_args()
plotTrajectoriesFile(args.particlefile, mode=args.mode, tracerfile=args.tracerfile,
tracerfield=args.tracerfilefield, tracerlon=args.tracerfilelon,
tracerlat=args.tracerfilelat, recordedvar=args.recordedvar,
bins=args.bins, show_plt=True)
