def generate(self, py_ast, funcvars):
# Replace occurences of intrinsic objects in Python AST
transformer = IntrinsicTransformer(self.fieldset, self.ptype)
py_ast = transformer.visit(py_ast)
# Untangle Pythonic tuple-assignment statements
py_ast = TupleSplitter().visit(py_ast)
# Generate C-code for all nodes in the Python AST
self.visit(py_ast)
self.ccode = py_ast.ccode
# Insert variable declarations for non-instrinsics
# Make sure that repeated variables are not declared more than
# once. If variables occur in multiple Kernels, give a warning
used_vars = []
funcvars_copy = copy(funcvars) # editing a list while looping over it is dangerous
for kvar in funcvars:
if kvar in used_vars:
if kvar not in ['particle', 'fieldset', 'time']:
logger.warning(kvar+" declared in multiple Kernels")
funcvars_copy.remove(kvar)
else:
used_vars.append(kvar)
funcvars = funcvars_copy
for kvar in self.kernel_vars + self.array_vars:
if kvar in funcvars:
funcvars.remove(kvar)
self.ccode.body.insert(0, c.Value('ErrorCode', 'err'))
if len(funcvars) > 0:
self.ccode.body.insert(0, c.Value("type_coord", ", ".join(funcvars)))
if len(transformer.tmp_vars) > 0:
self.ccode.body.insert(0, c.Value("float", ", ".join(transformer.tmp_vars)))
return self.ccode
def to_dict(self, pfile, time, deleted_only=False):
"""Convert all Particle data from one time step to a python dictionary.
:param pfile: ParticleFile object requesting the conversion
:param time: Time at which to write ParticleSet
:param deleted_only: Flag to write only the deleted Particles
returns two dictionaries: one for all variables to be written each outputdt,
and one for all variables to be written once
"""
data_dict = {}
data_dict_once = {}
time = time.total_seconds() if isinstance(time, delta) else time
pd = self.particle_data
if pfile.lasttime_written != time and \
(pfile.write_ondelete is False or deleted_only is not False):
if pd['id'].size == 0:
logger.warning(
"ParticleSet is empty on writing as array at time %g" %
time)
else:
if deleted_only is not False:
to_write = deleted_only
else:
to_write = _to_write_particles(pd, time)
if np.any(to_write) > 0:
for var in pfile.var_names:
data_dict[var] = pd[var][to_write]
pfile.maxid_written = np.maximum(pfile.maxid_written,
np.max(data_dict['id']))
pset_errs = (to_write & (pd['state'] != OperationCode.Delete)
& np.less(1e-3,
np.abs(time - pd['time']),
where=np.isfinite(pd['time'])))
if np.count_nonzero(pset_errs) > 0:
logger.warning_once(
'time argument in pfile.write() is {}, but particles have time {}'
.format(time, pd['time'][pset_errs]))
if time not in pfile.time_written:
pfile.time_written.append(time)
if len(pfile.var_names_once) > 0:
first_write = (_to_write_particles(pd, time) & np.isin(
pd['id'], pfile.written_once, invert=True))
if np.any(first_write):
data_dict_once['id'] = np.array(
pd['id'][first_write]).astype(dtype=np.int64)
for var in pfile.var_names_once:
data_dict_once[var] = pd[var][first_write]
pfile.written_once.extend(pd['id'][first_write])
if deleted_only is False:
pfile.lasttime_written = time
return data_dict, data_dict_once
def write(self, pset, time, sync=True, deleted_only=False):
"""Write :class:`parcels.particleset.ParticleSet` data to file
:param pset: ParticleSet object to write
:param time: Time at which to write ParticleSet
:param sync: Optional argument whether to write data to disk immediately. Default is True
"""
if self.dataset is None:
self.open_dataset()
if isinstance(time, delta):
time = time.total_seconds()
if self.lasttime_written != time and \
(self.write_ondelete is False or deleted_only is True):
if pset.size > 0:
first_write = [
p for p in pset if (p.fileid < 0 or len(self.idx) == 0)
and _is_particle_started_yet(p, time)
] # len(self.idx)==0 in case pset is written to new ParticleFile
for p in first_write:
p.fileid = self.lasttraj # particle id in current file
self.lasttraj += 1
self.idx = np.append(self.idx, np.zeros(len(first_write)))
for p in pset:
if _is_particle_started_yet(p, time):
i = p.fileid
self.id[i, self.idx[i]] = p.id
self.time[i, self.idx[i]] = p.time
self.lat[i, self.idx[i]] = p.lat
self.lon[i, self.idx[i]] = p.lon
self.z[i, self.idx[i]] = p.depth
for var in self.user_vars:
getattr(self, var)[i,
self.idx[i]] = getattr(p, var)
if p.state != ErrorCode.Delete and not np.allclose(
p.time, time):
logger.warning_once(
'time argument in pfile.write() is %g, but a particle has time %g.'
% (time, p.time))
for p in first_write:
for var in self.user_vars_once:
getattr(self, var)[p.fileid] = getattr(p, var)
else:
logger.warning(
"ParticleSet is empty on writing as array at time %g" %
time)
if not deleted_only:
self.idx += 1
self.lasttime_written = time
if sync:
self.sync()
def toDictionary(self, pfile, time, deleted_only=False):
"""
Convert all Particle data from one time step to a python dictionary.
:param pfile: ParticleFile object requesting the conversion
:param time: Time at which to write ParticleSet
:param deleted_only: Flag to write only the deleted Particles
returns two dictionaries: one for all variables to be written each outputdt,
and one for all variables to be written once
This function depends on the specific collection in question and thus needs to be specified in specific
derivative classes.
"""
data_dict = {}
data_dict_once = {}
time = time.total_seconds() if isinstance(time, delta) else time
indices_to_write = []
if pfile.lasttime_written != time and \
(pfile.write_ondelete is False or deleted_only is not False):
if self._data['id'].size == 0:
logger.warning("ParticleSet is empty on writing as array at time %g" % time)
else:
if deleted_only is not False:
if type(deleted_only) not in [list, np.ndarray] and deleted_only in [True, 1]:
indices_to_write = np.where(np.isin(self._data['state'],
[OperationCode.Delete]))[0]
elif type(deleted_only) in [list, np.ndarray]:
indices_to_write = deleted_only
else:
indices_to_write = _to_write_particles(self._data, time)
if np.any(indices_to_write) > 0:
for var in pfile.var_names:
data_dict[var] = self._data[var][indices_to_write]
pfile.maxid_written = np.maximum(pfile.maxid_written, np.max(data_dict['id']))
pset_errs = ((self._data['state'][indices_to_write] != OperationCode.Delete) & np.greater(np.abs(time - self._data['time'][indices_to_write]), 1e-3, where=np.isfinite(self._data['time'][indices_to_write])))
if np.count_nonzero(pset_errs) > 0:
logger.warning_once('time argument in pfile.write() is {}, but particles have time {}'.format(time, self._data['time'][pset_errs]))
# ==== this function should probably move back somewhere into the particle-file instead of the to_dict ==== #
if time not in pfile.time_written:
pfile.time_written.append(time)
if len(pfile.var_names_once) > 0:
first_write = (_to_write_particles(self._data, time) & _is_particle_started_yet(self._data, time) & np.isin(self._data['id'], pfile.written_once, invert=True))
if np.any(first_write):
data_dict_once['id'] = np.array(self._data['id'][first_write]).astype(dtype=np.int64)
for var in pfile.var_names_once:
data_dict_once[var] = self._data[var][first_write]
pfile.written_once.extend(np.array(self._data['id'][first_write]).astype(dtype=np.int64).tolist())
if deleted_only is False:
pfile.lasttime_written = time
return data_dict, data_dict_once
def convert_pset_to_dict(self, pset, time, deleted_only=False):
"""Convert all Particle data from one time step to a python dictionary.
:param pset: ParticleSet object to write
:param time: Time at which to write ParticleSet
:param deleted_only: Flag to write only the deleted Particles
returns two dictionaries: one for all variables to be written each outputdt,
and one for all variables to be written once
"""
data_dict = {}
data_dict_once = {}
time = time.total_seconds() if isinstance(time, delta) else time
if self.lasttime_written != time and \
(self.write_ondelete is False or deleted_only is True):
if pset.size == 0:
logger.warning("ParticleSet is empty on writing as array at time %g" % time)
else:
for var in self.var_names:
data_dict[var] = np.nan * np.zeros(len(pset))
i = 0
for p in pset:
if p.dt*p.time <= p.dt*time:
for var in self.var_names:
data_dict[var][i] = getattr(p, var)
if p.state != ErrorCode.Delete and not np.allclose(p.time, time):
logger.warning_once('time argument in pfile.write() is %g, but a particle has time %g.' % (time, p.time))
self.maxid_written = np.max([self.maxid_written, p.id])
i += 1
save_ind = np.isfinite(data_dict["id"])
for key in self.var_names:
data_dict[key] = data_dict[key][save_ind]
if time not in self.time_written:
self.time_written.append(time)
if len(self.var_names_once) > 0:
first_write = [p for p in pset if (p.id not in self.written_once) and _is_particle_started_yet(p, time)]
data_dict_once['id'] = np.nan * np.zeros(len(first_write))
for var in self.var_names_once:
data_dict_once[var] = np.nan * np.zeros(len(first_write))
i = 0
for p in first_write:
self.written_once.append(p.id)
data_dict_once['id'][i] = p.id
for var in self.var_names_once:
data_dict_once[var][i] = getattr(p, var)
i += 1
if not deleted_only:
self.lasttime_written = time
return data_dict, data_dict_once
def __init__(self, fieldset, ptype, pyfunc=None, funcname=None,
funccode=None, py_ast=None, funcvars=None, c_include="", delete_cfiles=True):
self.fieldset = fieldset
self.ptype = ptype
self._lib = None
self.delete_cfiles = delete_cfiles
self._cleanup_files = None
self._cleanup_lib = None
# Derive meta information from pyfunc, if not given
self.funcname = funcname or pyfunc.__name__
if pyfunc is AdvectionRK4_3D:
warning = False
if isinstance(fieldset.W, Field) and fieldset.W.creation_log != 'from_nemo' and \
fieldset.W._scaling_factor is not None and fieldset.W._scaling_factor > 0:
warning = True
if type(fieldset.W) in [SummedField, NestedField]:
for f in fieldset.W:
if f.creation_log != 'from_nemo' and f._scaling_factor is not None and f._scaling_factor > 0:
warning = True
if warning:
logger.warning_once('Note that in AdvectionRK4_3D, vertical velocity is assumed positive towards increasing z.\n'
' If z increases downward and w is positive upward you can re-orient it downwards by setting fieldset.W.set_scaling_factor(-1.)')
elif pyfunc is AdvectionAnalytical:
if ptype.uses_jit:
raise NotImplementedError('Analytical Advection only works in Scipy mode')
if fieldset.U.interp_method != 'cgrid_velocity':
raise NotImplementedError('Analytical Advection only works with C-grids')
if fieldset.U.grid.gtype not in [GridCode.CurvilinearZGrid, GridCode.RectilinearZGrid]:
raise NotImplementedError('Analytical Advection only works with Z-grids in the vertical')
if funcvars is not None:
self.funcvars = funcvars
elif hasattr(pyfunc, '__code__'):
self.funcvars = list(pyfunc.__code__.co_varnames)
else:
self.funcvars = None
self.funccode = funccode or inspect.getsource(pyfunc.__code__)
# Parse AST if it is not provided explicitly
self.py_ast = py_ast or parse(fix_indentation(self.funccode)).body[0]
if pyfunc is None:
# Extract user context by inspecting the call stack
stack = inspect.stack()
try:
user_ctx = stack[-1][0].f_globals
user_ctx['math'] = globals()['math']
user_ctx['ParcelsRandom'] = globals()['ParcelsRandom']
user_ctx['random'] = globals()['random']
user_ctx['StateCode'] = globals()['StateCode']
user_ctx['OperationCode'] = globals()['OperationCode']
user_ctx['ErrorCode'] = globals()['ErrorCode']
except:
logger.warning("Could not access user context when merging kernels")
user_ctx = globals()
finally:
del stack # Remove cyclic references
# Compile and generate Python function from AST
py_mod = parse("")
py_mod.body = [self.py_ast]
exec(compile(py_mod, "<ast>", "exec"), user_ctx)
self.pyfunc = user_ctx[self.funcname]
else:
self.pyfunc = pyfunc
if version_info[0] < 3:
numkernelargs = len(inspect.getargspec(self.pyfunc).args)
else:
numkernelargs = len(inspect.getfullargspec(self.pyfunc).args)
assert numkernelargs == 3, \
'Since Parcels v2.0, kernels do only take 3 arguments: particle, fieldset, time !! AND !! Argument order in field interpolation is time, depth, lat, lon.'
self.name = "%s%s" % (ptype.name, self.funcname)
# Generate the kernel function and add the outer loop
if self.ptype.uses_jit:
kernelgen = KernelGenerator(fieldset, ptype)
kernel_ccode = kernelgen.generate(deepcopy(self.py_ast),
self.funcvars)
self.field_args = kernelgen.field_args
self.vector_field_args = kernelgen.vector_field_args
fieldset = self.fieldset
for f in self.vector_field_args.values():
Wname = f.W.ccode_name if f.W else 'not_defined'
for sF_name, sF_component in zip([f.U.ccode_name, f.V.ccode_name, Wname], ['U', 'V', 'W']):
if sF_name not in self.field_args:
if sF_name != 'not_defined':
self.field_args[sF_name] = getattr(f, sF_component)
self.const_args = kernelgen.const_args
loopgen = LoopGenerator(fieldset, ptype)
if path.isfile(c_include):
with open(c_include, 'r') as f:
c_include_str = f.read()
else:
c_include_str = c_include
self.ccode = loopgen.generate(self.funcname, self.field_args, self.const_args,
kernel_ccode, c_include_str)
if MPI:
mpi_comm = MPI.COMM_WORLD
mpi_rank = mpi_comm.Get_rank()
basename = path.join(get_cache_dir(), self._cache_key) if mpi_rank == 0 else None
basename = mpi_comm.bcast(basename, root=0)
basename = basename + "_%d" % mpi_rank
else:
basename = path.join(get_cache_dir(), "%s_0" % self._cache_key)
self.src_file = "%s.c" % basename
self.lib_file = "%s.%s" % (basename, 'dll' if platform == 'win32' else 'so')
self.log_file = "%s.log" % basename
def from_particlefile(cls,
fieldset,
pclass,
filename,
restart=True,
restarttime=None,
repeatdt=None,
lonlatdepth_dtype=None,
**kwargs):
"""Initialise the ParticleSet from a netcdf ParticleFile.
This creates a new ParticleSet based on locations of all particles written
in a netcdf ParticleFile at a certain time. Particle IDs are preserved if restart=True
:param fieldset: :mod:`parcels.fieldset.FieldSet` object from which to sample velocity
:param pclass: mod:`parcels.particle.JITParticle` or :mod:`parcels.particle.ScipyParticle`
object that defines custom particle
:param filename: Name of the particlefile from which to read initial conditions
:param restart: Boolean to signal if pset is used for a restart (default is True).
In that case, Particle IDs are preserved.
:param restarttime: time at which the Particles will be restarted. Default is the last time written.
Alternatively, restarttime could be a time value (including np.datetime64) or
a callable function such as np.nanmin. The last is useful when running with dt < 0.
:param repeatdt: Optional interval (in seconds) on which to repeat the release of the ParticleSet
:param lonlatdepth_dtype: Floating precision for lon, lat, depth particle coordinates.
It is either np.float32 or np.float64. Default is np.float32 if fieldset.U.interp_method is 'linear'
and np.float64 if the interpolation method is 'cgrid_velocity'
"""
if repeatdt is not None:
logger.warning(
'Note that the `repeatdt` argument is not retained from %s, and that '
'setting a new repeatdt will start particles from the _new_ particle '
'locations.' % filename)
pfile = xr.open_dataset(str(filename), decode_cf=True)
pfile_vars = [v for v in pfile.data_vars]
vars = {}
to_write = {}
for v in pclass.getPType().variables:
if v.name in pfile_vars:
vars[v.name] = np.ma.filled(pfile.variables[v.name], np.nan)
elif v.name not in ['xi', 'yi', 'zi', 'ti', 'dt', '_next_dt', 'depth', 'id', 'fileid', 'state'] \
and v.to_write:
raise RuntimeError(
'Variable %s is in pclass but not in the particlefile' %
v.name)
to_write[v.name] = v.to_write
vars['depth'] = np.ma.filled(pfile.variables['z'], np.nan)
vars['id'] = np.ma.filled(pfile.variables['trajectory'], np.nan)
if isinstance(vars['time'][0, 0], np.timedelta64):
vars['time'] = np.array(
[t / np.timedelta64(1, 's') for t in vars['time']])
if restarttime is None:
restarttime = np.nanmax(vars['time'])
elif callable(restarttime):
restarttime = restarttime(vars['time'])
else:
restarttime = restarttime
inds = np.where(vars['time'] == restarttime)
for v in vars:
if to_write[v] is True:
vars[v] = vars[v][inds]
elif to_write[v] == 'once':
vars[v] = vars[v][inds[0]]
if v not in ['lon', 'lat', 'depth', 'time', 'id']:
kwargs[v] = vars[v]
if restart:
pclass.setLastID(0) # reset to zero offset
else:
vars['id'] = None
return cls(fieldset=fieldset,
pclass=pclass,
lon=vars['lon'],
lat=vars['lat'],
depth=vars['depth'],
time=vars['time'],
pid_orig=vars['id'],
lonlatdepth_dtype=lonlatdepth_dtype,
repeatdt=repeatdt,
**kwargs)
def convert_pset_to_dict(self, pset, time, deleted_only=False):
"""Convert all Particle data from one time step to a python dictionary.
:param pset: ParticleSet object to write
:param time: Time at which to write ParticleSet
:param deleted_only: Flag to write only the deleted Particles
returns two dictionaries: one for all variables to be written each outputdt,
and one for all variables to be written once
"""
data_dict = {}
data_dict_once = {}
time = time.total_seconds() if isinstance(time, delta) else time
if self.lasttime_written != time and \
(self.write_ondelete is False or deleted_only is True):
if pset.size == 0:
logger.warning(
"ParticleSet is empty on writing as array at time %g" %
time)
else:
if deleted_only:
pset_towrite = pset
elif pset[0].dt > 0:
pset_towrite = [
p for p in pset
if time <= p.time < time + p.dt and np.isfinite(p.id)
]
else:
pset_towrite = [
p for p in pset
if time + p.dt < p.time <= time and np.isfinite(p.id)
]
if len(pset_towrite) > 0:
for var in self.var_names:
data_dict[var] = np.array(
[getattr(p, var) for p in pset_towrite])
self.maxid_written = np.max(
[self.maxid_written,
np.max(data_dict['id'])])
pset_errs = [
p for p in pset_towrite
if p.state != ErrorCode.Delete and abs(time -
p.time) > 1e-3
]
for p in pset_errs:
logger.warning_once(
'time argument in pfile.write() is %g, but a particle has time % g.'
% (time, p.time))
if time not in self.time_written:
self.time_written.append(time)
if len(self.var_names_once) > 0:
first_write = [
p for p in pset if (p.id not in self.written_once)
and _is_particle_started_yet(p, time)
]
data_dict_once['id'] = np.array(
[p.id for p in first_write])
for var in self.var_names_once:
data_dict_once[var] = np.array(
[getattr(p, var) for p in first_write])
self.written_once += [p.id for p in first_write]
if not deleted_only:
self.lasttime_written = time
return data_dict, data_dict_once
def plotfield(field,
show_time=None,
domain=None,
depth_level=0,
projection=None,
land=True,
vmin=None,
vmax=None,
savefile=None,
**kwargs):
"""Function to plot a Parcels Field
:param show_time: Time at which to show the Field
:param domain: dictionary (with keys 'N', 'S', 'E', 'W') defining domain to show
:param depth_level: depth level to be plotted (default 0)
:param projection: type of cartopy projection to use (default PlateCarree)
:param land: Boolean whether to show land. This is ignored for flat meshes
:param vmin: minimum colour scale (only in single-plot mode)
:param vmax: maximum colour scale (only in single-plot mode)
:param savefile: Name of a file to save the plot to
:param animation: Boolean whether result is a single plot, or an animation
"""
if type(field) is VectorField:
spherical = True if field.U.grid.mesh == 'spherical' else False
field = [field.U, field.V]
plottype = 'vector'
elif type(field) is Field:
spherical = True if field.grid.mesh == 'spherical' else False
field = [field]
plottype = 'scalar'
else:
raise RuntimeError('field needs to be a Field or VectorField object')
if field[0].grid.gtype in [
GridCode.CurvilinearZGrid, GridCode.CurvilinearSGrid
]:
logger.warning(
'Field.show() does not always correctly determine the domain for curvilinear grids. '
'Use plotting with caution and perhaps use domain argument as in the NEMO 3D tutorial'
)
plt, fig, ax, cartopy = create_parcelsfig_axis(spherical,
land,
projection=projection)
if plt is None:
return None, None, None, None # creating axes was not possible
data = {}
plotlon = {}
plotlat = {}
for i, fld in enumerate(field):
show_time = fld.grid.time[0] if show_time is None else show_time
if fld.grid.defer_load:
fld.fieldset.computeTimeChunk(show_time, 1)
(idx, periods) = fld.time_index(show_time)
show_time -= periods * (fld.grid.time_full[-1] - fld.grid.time_full[0])
if show_time > fld.grid.time[-1] or show_time < fld.grid.time[0]:
raise TimeExtrapolationError(show_time, field=fld, msg='show_time')
latN, latS, lonE, lonW = parsedomain(domain, fld)
if isinstance(fld.grid, CurvilinearGrid):
plotlon[i] = fld.grid.lon[latS:latN, lonW:lonE]
plotlat[i] = fld.grid.lat[latS:latN, lonW:lonE]
else:
plotlon[i] = fld.grid.lon[lonW:lonE]
plotlat[i] = fld.grid.lat[latS:latN]
if i > 0 and not np.allclose(plotlon[i], plotlon[0]):
raise RuntimeError(
'VectorField needs to be on an A-grid for plotting')
if fld.grid.time.size > 1:
if fld.grid.zdim > 1:
data[i] = np.squeeze(
fld.temporal_interpolate_fullfield(idx,
show_time))[depth_level,
latS:latN,
lonW:lonE]
else:
data[i] = np.squeeze(
fld.temporal_interpolate_fullfield(idx,
show_time))[latS:latN,
lonW:lonE]
else:
if fld.grid.zdim > 1:
data[i] = np.squeeze(fld.data)[depth_level, latS:latN,
lonW:lonE]
else:
data[i] = np.squeeze(fld.data)[latS:latN, lonW:lonE]
if plottype == 'vector':
if field[0].interp_method == 'cgrid_velocity':
logger.warning_once(
'Plotting a C-grid velocity field is achieved via an A-grid projection, reducing the plot accuracy'
)
d = np.empty_like(data[0])
d[:-1, :] = (data[0][:-1, :] + data[0][1:, :]) / 2.
d[-1, :] = data[0][-1, :]
data[0] = d
d = np.empty_like(data[0])
d[:, :-1] = (data[0][:, :-1] + data[0][:, 1:]) / 2.
d[:, -1] = data[0][:, -1]
data[1] = d
spd = data[0]**2 + data[1]**2
speed = np.where(spd > 0, np.sqrt(spd), 0)
vmin = speed.min() if vmin is None else vmin
vmax = speed.max() if vmax is None else vmax
if isinstance(field[0].grid, CurvilinearGrid):
x, y = plotlon[0], plotlat[0]
else:
x, y = np.meshgrid(plotlon[0], plotlat[0])
u = np.where(speed > 0., data[0] / speed, 0)
v = np.where(speed > 0., data[1] / speed, 0)
if cartopy:
cs = ax.quiver(np.asarray(x),
np.asarray(y),
np.asarray(u),
np.asarray(v),
speed,
cmap=plt.cm.gist_ncar,
clim=[vmin, vmax],
scale=50,
transform=cartopy.crs.PlateCarree())
else:
cs = ax.quiver(x,
y,
u,
v,
speed,
cmap=plt.cm.gist_ncar,
clim=[vmin, vmax],
scale=50)
else:
vmin = data[0].min() if vmin is None else vmin
vmax = data[0].max() if vmax is None else vmax
assert len(data[0].shape) == 2
if field[0].interp_method == 'cgrid_tracer':
d = data[0][1:, 1:]
elif field[0].interp_method == 'cgrid_velocity':
if field[0].fieldtype == 'U':
d = np.empty_like(data[0])
d[:-1, :-1] = (data[0][1:, :-1] + data[0][1:, 1:]) / 2.
elif field[0].fieldtype == 'V':
d = np.empty_like(data[0])
d[:-1, :-1] = (data[0][:-1, 1:] + data[0][1:, 1:]) / 2.
else: # W
d = data[0][1:, 1:]
else: # if A-grid
d = (data[0][:-1, :-1] + data[0][1:, :-1] + data[0][:-1, 1:] +
data[0][1:, 1:]) / 4.
d = np.where(data[0][:-1, :-1] == 0, 0, d)
d = np.where(data[0][1:, :-1] == 0, 0, d)
d = np.where(data[0][1:, 1:] == 0, 0, d)
d = np.where(data[0][:-1, 1:] == 0, 0, d)
if cartopy:
cs = ax.pcolormesh(plotlon[0],
plotlat[0],
d,
transform=cartopy.crs.PlateCarree())
else:
cs = ax.pcolormesh(plotlon[0], plotlat[0], d)
if cartopy is None:
ax.set_xlim(np.nanmin(plotlon[0]), np.nanmax(plotlon[0]))
ax.set_ylim(np.nanmin(plotlat[0]), np.nanmax(plotlat[0]))
elif domain is not None:
ax.set_extent([
np.nanmin(plotlon[0]),
np.nanmax(plotlon[0]),
np.nanmin(plotlat[0]),
np.nanmax(plotlat[0])
],
crs=cartopy.crs.PlateCarree())
cs.cmap.set_over('k')
cs.cmap.set_under('w')
cs.set_clim(vmin, vmax)
cartopy_colorbar(cs, plt, fig, ax)
timestr = parsetimestr(field[0].grid.time_origin, show_time)
titlestr = kwargs.pop('titlestr', '')
if field[0].grid.zdim > 1:
if field[0].grid.gtype in [
GridCode.CurvilinearZGrid, GridCode.RectilinearZGrid
]:
gphrase = 'depth'
depth_or_level = field[0].grid.depth[depth_level]
else:
gphrase = 'level'
depth_or_level = depth_level
depthstr = ' at %s %g ' % (gphrase, depth_or_level)
else:
depthstr = ''
if plottype == 'vector':
ax.set_title(titlestr + 'Velocity field' + depthstr + timestr)
else:
ax.set_title(titlestr + field[0].name + depthstr + timestr)
if not spherical:
ax.set_xlabel('Zonal distance [m]')
ax.set_ylabel('Meridional distance [m]')
plt.draw()
if savefile:
plt.savefig(savefile)
logger.info('Plot saved to ' + savefile + '.png')
plt.close()
return plt, fig, ax, cartopy
def __init__(self, fieldset, ptype, pyfunc=None, funcname=None,
funccode=None, py_ast=None, funcvars=None, c_include="", delete_cfiles=True):
super(KernelSOA, self).__init__(fieldset=fieldset, ptype=ptype, pyfunc=pyfunc, funcname=funcname, funccode=funccode, py_ast=py_ast, funcvars=funcvars, c_include=c_include, delete_cfiles=delete_cfiles)
# Derive meta information from pyfunc, if not given
self.check_fieldsets_in_kernels(pyfunc)
if funcvars is not None:
self.funcvars = funcvars
elif hasattr(pyfunc, '__code__'):
self.funcvars = list(pyfunc.__code__.co_varnames)
else:
self.funcvars = None
self.funccode = funccode or inspect.getsource(pyfunc.__code__)
# Parse AST if it is not provided explicitly
self.py_ast = py_ast or parse(BaseKernel.fix_indentation(self.funccode)).body[0]
if pyfunc is None:
# Extract user context by inspecting the call stack
stack = inspect.stack()
try:
user_ctx = stack[-1][0].f_globals
user_ctx['math'] = globals()['math']
user_ctx['ParcelsRandom'] = globals()['ParcelsRandom']
user_ctx['random'] = globals()['random']
user_ctx['StateCode'] = globals()['StateCode']
user_ctx['OperationCode'] = globals()['OperationCode']
user_ctx['ErrorCode'] = globals()['ErrorCode']
except:
logger.warning("Could not access user context when merging kernels")
user_ctx = globals()
finally:
del stack # Remove cyclic references
# Compile and generate Python function from AST
py_mod = parse("")
py_mod.body = [self.py_ast]
exec(compile(py_mod, "<ast>", "exec"), user_ctx)
self._pyfunc = user_ctx[self.funcname]
else:
self._pyfunc = pyfunc
numkernelargs = self.check_kernel_signature_on_version()
assert numkernelargs == 3, \
'Since Parcels v2.0, kernels do only take 3 arguments: particle, fieldset, time !! AND !! Argument order in field interpolation is time, depth, lat, lon.'
self.name = "%s%s" % (ptype.name, self.funcname)
# Generate the kernel function and add the outer loop
if self.ptype.uses_jit:
kernelgen = KernelGenerator(fieldset, ptype)
kernel_ccode = kernelgen.generate(deepcopy(self.py_ast),
self.funcvars)
self.field_args = kernelgen.field_args
self.vector_field_args = kernelgen.vector_field_args
fieldset = self.fieldset
for f in self.vector_field_args.values():
Wname = f.W.ccode_name if f.W else 'not_defined'
for sF_name, sF_component in zip([f.U.ccode_name, f.V.ccode_name, Wname], ['U', 'V', 'W']):
if sF_name not in self.field_args:
if sF_name != 'not_defined':
self.field_args[sF_name] = getattr(f, sF_component)
self.const_args = kernelgen.const_args
loopgen = LoopGenerator(fieldset, ptype)
if path.isfile(self._c_include):
with open(self._c_include, 'r') as f:
c_include_str = f.read()
else:
c_include_str = self._c_include
self.ccode = loopgen.generate(self.funcname, self.field_args, self.const_args,
kernel_ccode, c_include_str)
src_file_or_files, self.lib_file, self.log_file = self.get_kernel_compile_files()
if type(src_file_or_files) in (list, dict, tuple, np.ndarray):
self.dyn_srcs = src_file_or_files
else:
self.src_file = src_file_or_files
def _get_initial_chunk_dictionary(self):
"""
[private function - not to be called from outside the class]
Super-function that maps and correlates the requested chunksize with the requested parcels dimensions, variables
and the NetCDF-available dimensions. Thus, it takes care to remove chunksize arguments that are not in the
Parcels- or NetCDF dimensions, or whose chunking would be omitted due to an empty chunk dimension.
The function retuns the corrected chunksize dictionary. The function also initializes the chunk_map.
The chunk map maps the array index dimension to the requested chunksize.
Apart from resolving the different requested version of the chunksize, the function also test-executes the
chunk request. If this initial test fails, as a last resort, we execute a heuristic to map the requested
parcels dimensions to the dimension signature of the most-parameterized NetCDF variable, and heuristically
try to map its parameters to the parcels dimensions with the class-wide name-map.
"""
# ==== check-opening requested dataset to access metadata ==== #
# ==== file-opening and dimension-reading does not require a decode or lock ==== #
self.dataset = xr.open_dataset(str(self.filename), decode_cf=False, engine=self.netcdf_engine, chunks={}, lock=False)
self.dataset['decoded'] = False
# ==== self.dataset temporarily available ==== #
init_chunk_dict = {}
init_chunk_map = {}
if isinstance(self.chunksize, dict):
init_chunk_dict, init_chunk_map = self._get_initial_chunk_dictionary_by_dict_()
elif self.chunksize == 'auto':
av_mem = psutil.virtual_memory().available
chunk_cap = av_mem * (1/8) * (1/3)
if 'array.chunk-size' in da_conf.config.keys():
chunk_cap = da_utils.parse_bytes(da_conf.config.get('array.chunk-size'))
else:
predefined_cap = da_conf.get('array.chunk-size')
if predefined_cap is not None:
chunk_cap = da_utils.parse_bytes(predefined_cap)
else:
logger.info_once("Unable to locate chunking hints from dask, thus estimating the max. chunk size heuristically. Please consider defining the 'chunk-size' for 'array' in your local dask configuration file (see http://oceanparcels.org/faq.html#field_chunking_config and https://docs.dask.org).")
loni, lonname, lonvalue = self._is_dimension_in_dataset('lon')
lati, latname, latvalue = self._is_dimension_in_dataset('lat')
if lati is not None and loni is not None and lati >= 0 and loni >= 0:
pDim = int(math.floor(math.sqrt(chunk_cap/np.dtype(np.float64).itemsize)))
init_chunk_dict[latname] = min(latvalue, pDim)
init_chunk_map[lati] = min(latvalue, pDim)
init_chunk_dict[lonname] = min(lonvalue, pDim)
init_chunk_map[loni] = min(lonvalue, pDim)
timei, timename, timevalue = self._is_dimension_in_dataset('time')
if timei is not None and timei >= 0:
init_chunk_dict[timename] = min(1, timevalue)
init_chunk_map[timei] = min(1, timevalue)
depthi, depthname, depthvalue = self._is_dimension_in_dataset('depth')
if depthi is not None and depthi >= 0:
init_chunk_dict[depthname] = max(1, depthvalue)
init_chunk_map[depthi] = max(1, depthvalue)
# ==== closing check-opened requested dataset ==== #
self.dataset.close()
# ==== check if the chunksize reading is successful. if not, load the file ONCE really into memory and ==== #
# ==== deduce the chunking from the array dims. ==== #
if len(init_chunk_dict) == 0 and self.chunksize not in [False, None, 'auto']:
self.autochunkingfailed = True
raise DaskChunkingError(self.__class__.__name__, "No correct mapping found between Parcels- and NetCDF dimensions! Please correct the 'FieldSet(..., chunksize={...})' parameter and try again.")
else:
self.autochunkingfailed = False
try:
self.dataset = xr.open_dataset(str(self.filename), decode_cf=True, engine=self.netcdf_engine, chunks=init_chunk_dict, lock=False)
if isinstance(self.chunksize, dict):
self.chunksize = init_chunk_dict
except:
logger.warning("Chunking with init_chunk_dict = {} failed - Executing Dask chunking 'failsafe'...".format(init_chunk_dict))
self.autochunkingfailed = True
if not self.autochunkingfailed:
init_chunk_dict = self._failsafe_parse_()
if isinstance(self.chunksize, dict):
self.chunksize = init_chunk_dict
finally:
self.dataset.close()
self.chunk_mapping = init_chunk_map
self.dataset = None
# ==== self.dataset not available ==== #
return init_chunk_dict
def write(self, pset, time, sync=True, deleted_only=False):
"""Write :class:`parcels.particleset.ParticleSet`
All data from one time step is saved to one NPY-file using a python
dictionary. The data is saved in the folder 'out'.
:param pset: ParticleSet object to write
:param time: Time at which to write ParticleSet
:param sync: Optional argument whether to write data to disk immediately. Default is True
"""
if self.dataset is None:
self.open_dataset()
if isinstance(time, delta):
time = time.total_seconds()
if self.lasttime_written != time and \
(self.write_ondelete is False or deleted_only is True):
if pset.size > 0:
first_write = [p for p in pset if (p.fileid < 0 or len(self.idx) == 0) and (p.dt*p.time <= p.dt*time or np.isnan(p.dt))] # len(self.idx)==0 in case pset is written to new ParticleFile
for p in first_write:
p.fileid = self.lasttraj
self.lasttraj += 1
self.idx = np.append(self.idx, np.zeros(len(first_write)))
size = len(pset)
# dictionary for temporary hold data
tmp = {}
tmp["ids"], tmp["time"], tmp["lat"], tmp["lon"], tmp["z"] =\
map(lambda x: np.zeros(x), [size,size,size,size,size])
for var in self.user_vars:
tmp[var] = np.zeros(size)
for key in tmp.keys():
tmp[key][:] = np.nan
i = 0
for p in pset:
if p.dt*p.time <= p.dt*time:
tmp["ids"][i] = p.id
tmp["time"][i] = time
tmp["lat"][i] = p.lat
tmp["lon"][i] = p.lon
tmp["z"][i] = p.depth
for var in self.user_vars:
tmp[var][i] = getattr(p, var)
if p.state != ErrorCode.Delete and not np.allclose(p.time, time):
logger.warning_once('time argument in pfile.write() is %g, but a particle has time %g.' % (time, p.time))
i += 1
if not os.path.exists(self.npy_path):
os.mkdir(self.npy_path)
save_ind = np.isfinite(tmp["ids"])
for key in tmp.keys():
tmp[key] = tmp[key][save_ind]
np.save(os.path.join(self.npy_path,str(time)),tmp)
for p in first_write:
for var in self.user_vars_once:
getattr(self, var)[p.fileid] = getattr(p, var)
else:
logger.warning("ParticleSet is empty on writing as array at time %g" % time)
if not deleted_only:
self.idx += 1
self.lasttime_written = time
if sync:
self.sync()
def __init__(self,
fieldset,
ptype,
pyfunc=None,
funcname=None,
funccode=None,
py_ast=None,
funcvars=None,
c_include=""):
self.fieldset = fieldset
self.ptype = ptype
self._lib = None
# Derive meta information from pyfunc, if not given
self.funcname = funcname or pyfunc.__name__
if pyfunc is AdvectionRK4_3D:
logger.warning_once(
'Note that positive vertical velocity is assumed DOWNWARD by AdvectionRK4_3D'
)
if funcvars is not None:
self.funcvars = funcvars
elif hasattr(pyfunc, '__code__'):
self.funcvars = list(pyfunc.__code__.co_varnames)
else:
self.funcvars = None
self.funccode = funccode or inspect.getsource(pyfunc.__code__)
# Parse AST if it is not provided explicitly
self.py_ast = py_ast or parse(fix_indentation(self.funccode)).body[0]
if pyfunc is None:
# Extract user context by inspecting the call stack
stack = inspect.stack()
try:
user_ctx = stack[-1][0].f_globals
user_ctx['math'] = globals()['math']
user_ctx['random'] = globals()['random']
user_ctx['ErrorCode'] = globals()['ErrorCode']
except:
logger.warning(
"Could not access user context when merging kernels")
user_ctx = globals()
finally:
del stack # Remove cyclic references
# Compile and generate Python function from AST
py_mod = Module(body=[self.py_ast])
exec(compile(py_mod, "<ast>", "exec"), user_ctx)
self.pyfunc = user_ctx[self.funcname]
else:
self.pyfunc = pyfunc
self.name = "%s%s" % (ptype.name, self.funcname)
# Generate the kernel function and add the outer loop
if self.ptype.uses_jit:
kernelgen = KernelGenerator(fieldset, ptype)
kernel_ccode = kernelgen.generate(deepcopy(self.py_ast),
self.funcvars)
self.field_args = kernelgen.field_args
self.vector_field_args = kernelgen.vector_field_args
fieldset = self.fieldset
for fname in self.vector_field_args:
f = getattr(fieldset, fname)
Wname = f.W.name if f.W else 'not_defined'
for sF in [f.U.name, f.V.name, Wname]:
if sF not in self.field_args:
try:
self.field_args[sF] = getattr(fieldset, sF)
except:
continue
self.const_args = kernelgen.const_args
loopgen = LoopGenerator(fieldset, ptype)
if path.isfile(c_include):
with open(c_include, 'r') as f:
c_include_str = f.read()
else:
c_include_str = c_include
self.ccode = loopgen.generate(self.funcname, self.field_args,
self.const_args, kernel_ccode,
c_include_str)
basename = path.join(get_cache_dir(), self._cache_key)
self.src_file = "%s.c" % basename
self.lib_file = "%s.%s" % (basename,
'dll' if platform == 'win32' else 'so')
self.log_file = "%s.log" % basename
def __init__(self,
fieldset,
ptype,
pyfunc=None,
funcname=None,
funccode=None,
py_ast=None,
funcvars=None,
c_include=""):
self.fieldset = fieldset
self.ptype = ptype
self._lib = None
# Derive meta information from pyfunc, if not given
self.funcname = funcname or pyfunc.__name__
if pyfunc is AdvectionRK4_3D:
warning = False
if isinstance(fieldset.W, Field) and fieldset.W.creation_log != 'from_nemo' and \
fieldset.W._scaling_factor is not None and fieldset.W._scaling_factor > 0:
warning = True
if type(fieldset.W) in [SummedField, NestedField]:
for f in fieldset.W:
if f.creation_log != 'from_nemo' and f._scaling_factor is not None and f._scaling_factor > 0:
warning = True
if warning:
logger.warning_once(
'Note that in AdvectionRK4_3D, vertical velocity is assumed positive towards increasing z.\n'
' If z increases downward and w is positive upward you can re-orient it downwards by setting fieldset.W.set_scaling_factor(-1.)'
)
if funcvars is not None:
self.funcvars = funcvars
elif hasattr(pyfunc, '__code__'):
self.funcvars = list(pyfunc.__code__.co_varnames)
else:
self.funcvars = None
self.funccode = funccode or inspect.getsource(pyfunc.__code__)
# Parse AST if it is not provided explicitly
self.py_ast = py_ast or parse(fix_indentation(self.funccode)).body[0]
if pyfunc is None:
# Extract user context by inspecting the call stack
stack = inspect.stack()
try:
user_ctx = stack[-1][0].f_globals
user_ctx['math'] = globals()['math']
user_ctx['random'] = globals()['random']
user_ctx['ErrorCode'] = globals()['ErrorCode']
except:
logger.warning(
"Could not access user context when merging kernels")
user_ctx = globals()
finally:
del stack # Remove cyclic references
# Compile and generate Python function from AST
py_mod = Module(body=[self.py_ast])
exec(compile(py_mod, "<ast>", "exec"), user_ctx)
self.pyfunc = user_ctx[self.funcname]
else:
self.pyfunc = pyfunc
assert len(inspect.getargspec(self.pyfunc).args) == 3, \
'Since Parcels v2.0, kernels do only take 3 arguments: particle, fieldset, time !! AND !! Argument order in field interpolation is time, depth, lat, lon.'
self.name = "%s%s" % (ptype.name, self.funcname)
# Generate the kernel function and add the outer loop
if self.ptype.uses_jit:
kernelgen = KernelGenerator(fieldset, ptype)
kernel_ccode = kernelgen.generate(deepcopy(self.py_ast),
self.funcvars)
self.field_args = kernelgen.field_args
self.vector_field_args = kernelgen.vector_field_args
fieldset = self.fieldset
for fname in self.vector_field_args:
f = getattr(fieldset, fname)
Wname = f.W.name if f.W else 'not_defined'
for sF in [f.U.name, f.V.name, Wname]:
if sF not in self.field_args:
try:
self.field_args[sF] = getattr(fieldset, sF)
except:
continue
self.const_args = kernelgen.const_args
loopgen = LoopGenerator(fieldset, ptype)
if path.isfile(c_include):
with open(c_include, 'r') as f:
c_include_str = f.read()
else:
c_include_str = c_include
self.ccode = loopgen.generate(self.funcname, self.field_args,
self.const_args, kernel_ccode,
c_include_str)
basename = path.join(get_cache_dir(), self._cache_key)
self.src_file = "%s.c" % basename
self.lib_file = "%s.%s" % (basename,
'dll' if platform == 'win32' else 'so')
self.log_file = "%s.log" % basename
