def test_from_netcdf_chunking(mode, time_periodic, chunksize, deferLoad):
fnameU = path.join(path.dirname(__file__), 'test_data', 'perlinfieldsU.nc')
fnameV = path.join(path.dirname(__file__), 'test_data', 'perlinfieldsV.nc')
ufiles = [
fnameU,
] * 4
vfiles = [
fnameV,
] * 4
timestamps = np.arange(0, 4, 1) * 86400.0
timestamps = np.expand_dims(timestamps, 1)
files = {'U': ufiles, 'V': vfiles}
variables = {'U': 'vozocrtx', 'V': 'vomecrty'}
dimensions = {'lon': 'nav_lon', 'lat': 'nav_lat'}
fieldset = FieldSet.from_netcdf(files,
variables,
dimensions,
timestamps=timestamps,
time_periodic=time_periodic,
deferred_load=deferLoad,
allow_time_extrapolation=True if
time_periodic in [False, None] else False,
chunksize=chunksize)
pset = ParticleSet.from_line(fieldset,
size=1,
pclass=ptype[mode],
start=(0.5, 0.5),
finish=(0.5, 0.5))
pset.execute(AdvectionRK4, dt=1, runtime=1)
def test_moving_eddies_fwdbwd(mode, npart=2):
method = AdvectionRK4
fieldset = moving_eddies_fieldset()
# Determine particle class according to mode
pset = ParticleSet.from_line(fieldset=fieldset,
size=npart,
pclass=ptype[mode],
start=(3.3, 46.),
finish=(3.3, 47.8))
# Execte for 14 days, with 30sec timesteps and hourly output
runtime = delta(days=1)
dt = delta(minutes=5)
outputdt = delta(hours=1)
print("MovingEddies: Advecting %d particles for %s" %
(npart, str(runtime)))
pset.execute(method,
runtime=runtime,
dt=dt,
output_file=pset.ParticleFile(name="EddyParticlefwd",
outputdt=outputdt))
print("Now running in backward time mode")
pset.execute(method,
endtime=0,
dt=-dt,
output_file=pset.ParticleFile(name="EddyParticlebwd",
outputdt=outputdt))
assert (pset[0].lon > 3.2 and 45.9 < pset[0].lat < 46.1)
assert (pset[1].lon > 3.2 and 47.7 < pset[1].lat < 47.9)
return pset
def test_brownian_example(mode, npart=3000):
fieldset = brownian_fieldset()
# Set diffusion constants.
fieldset.Kh_meridional = 100.
fieldset.Kh_zonal = 100.
# Set random seed
random.seed(123456)
ptcls_start = 300000. # Start all particles at same location in middle of grid.
pset = ParticleSet.from_line(fieldset=fieldset, size=npart, pclass=ptype[mode],
start=(ptcls_start, ptcls_start),
finish=(ptcls_start, ptcls_start))
endtime = delta(days=1)
dt = delta(hours=1)
interval = delta(hours=1)
k_brownian = pset.Kernel(two_dim_brownian_flat)
pset.execute(k_brownian, endtime=endtime, dt=dt, interval=interval,
output_file=pset.ParticleFile(name="BrownianParticle"),
show_movie=False)
lats = np.array([particle.lat for particle in pset.particles])
lons = np.array([particle.lon for particle in pset.particles])
expected_std_lat = np.sqrt(2*fieldset.Kh_meridional*endtime.total_seconds())
expected_std_lon = np.sqrt(2*fieldset.Kh_zonal*endtime.total_seconds())
assert np.allclose(np.std(lats), expected_std_lat, rtol=.1)
assert np.allclose(np.std(lons), expected_std_lon, rtol=.1)
assert np.allclose(np.mean(lons), ptcls_start, rtol=.1)
assert np.allclose(np.mean(lats), ptcls_start, rtol=.1)
def moving_eddies_example(fieldset,
npart=2,
mode='jit',
verbose=False,
method=AdvectionRK4):
"""Configuration of a particle set that follows two moving eddies
:arg fieldset: :class FieldSet: that defines the flow field
:arg npart: Number of particles to initialise"""
# Determine particle class according to mode
pset = ParticleSet.from_line(fieldset=fieldset,
size=npart,
pclass=ptype[mode],
start=(3.3, 46.),
finish=(3.3, 47.8))
if verbose:
print("Initial particle positions:\n%s" % pset)
# Execute for 1 week, with 1 hour timesteps and hourly output
runtime = delta(days=7)
print("MovingEddies: Advecting %d particles for %s" %
(npart, str(runtime)))
pset.execute(method,
runtime=runtime,
dt=delta(hours=1),
output_file=pset.ParticleFile(name="EddyParticle",
outputdt=delta(hours=1)),
moviedt=None)
if verbose:
print("Final particle positions:\n%s" % pset)
return pset
def moving_eddies_example(grid,
npart=2,
mode='jit',
verbose=False,
method=AdvectionRK4):
"""Configuration of a particle set that follows two moving eddies
:arg grid: :class Grid: that defines the flow field
:arg npart: Number of particles to intialise"""
# Determine particle class according to mode
pset = ParticleSet.from_line(grid=grid,
size=npart,
pclass=ptype[mode],
start=(3.3, 46.),
finish=(3.3, 47.8))
if verbose:
print("Initial particle positions:\n%s" % pset)
# Execte for 21 days, with 5min timesteps and hourly output
endtime = delta(days=21)
print("MovingEddies: Advecting %d particles for %s" %
(npart, str(endtime)))
pset.execute(method,
endtime=endtime,
dt=delta(minutes=5),
output_file=pset.ParticleFile(name="EddyParticle"),
interval=delta(hours=1),
show_movie=False)
if verbose:
print("Final particle positions:\n%s" % pset)
return pset
def stommel_example(npart=1,
mode='jit',
verbose=False,
method=AdvectionRK4,
grid_type='A',
outfile="StommelParticle.nc",
repeatdt=None,
maxage=None,
write_fields=True):
timer.fieldset = timer.Timer('FieldSet', parent=timer.stommel)
fieldset = stommel_fieldset(grid_type=grid_type)
if write_fields:
filename = 'stommel'
fieldset.write(filename)
timer.fieldset.stop()
# Determine particle class according to mode
timer.pset = timer.Timer('Pset', parent=timer.stommel)
timer.psetinit = timer.Timer('Pset_init', parent=timer.pset)
ParticleClass = JITParticle if mode == 'jit' else ScipyParticle
class MyParticle(ParticleClass):
p = Variable('p', dtype=np.float32, initial=0.)
p_start = Variable('p_start', dtype=np.float32, initial=fieldset.P)
age = Variable('age', dtype=np.float32, initial=0.)
pset = ParticleSet.from_line(fieldset,
size=npart,
pclass=MyParticle,
repeatdt=repeatdt,
start=(10e3, 5000e3),
finish=(100e3, 5000e3),
time=0)
if verbose:
print("Initial particle positions:\n%s" % pset)
# Execute for 30 days, with 1hour timesteps and 12-hourly output
runtime = delta(days=600)
dt = delta(hours=1)
outputdt = delta(days=5)
maxage = runtime.total_seconds() if maxage is None else maxage
fieldset.add_constant('maxage', maxage)
print("Stommel: Advecting %d particles for %s" % (npart, runtime))
timer.psetinit.stop()
timer.psetrun = timer.Timer('Pset_run', parent=timer.pset)
pset.execute(method + pset.Kernel(UpdateP) + pset.Kernel(AgeP),
runtime=runtime,
dt=dt,
moviedt=None,
output_file=pset.ParticleFile(name=outfile,
outputdt=outputdt))
if verbose:
print("Final particle positions:\n%s" % pset)
timer.psetrun.stop()
timer.pset.stop()
return pset
def test_pset_create_line(fieldset, mode, lonlatdepth_dtype, npart=100):
lon = np.linspace(0, 1, npart, dtype=lonlatdepth_dtype)
lat = np.linspace(1, 0, npart, dtype=lonlatdepth_dtype)
pset = ParticleSet.from_line(fieldset, size=npart, start=(0, 1), finish=(1, 0),
pclass=ptype[mode], lonlatdepth_dtype=lonlatdepth_dtype)
assert np.allclose([p.lon for p in pset], lon, rtol=1e-12)
assert np.allclose([p.lat for p in pset], lat, rtol=1e-12)
assert isinstance(pset[0].lat, lonlatdepth_dtype)
def run_longitudinalshear(fieldset, npart, outfilename):
pset = ParticleSet.from_line(fieldset, size=npart, pclass=JITParticle,
start=(0, -30), finish=(0, 60))
outfile = pset.ParticleFile(name=outfilename, outputdt=delta(days=1))
pset.execute(AdvectionRK4, runtime=delta(days=57), dt=delta(minutes=5),
output_file=outfile)
def test_recursive_errorhandling(mode, xdim=2, ydim=2):
"""Example script to show how recursaive error handling can work.
In this example, a set of Particles is started at Longitude 0.5.
These are run through a Kernel that throws an error if the
Longitude is smaller than 0.7.
The error Kernel then draws a new random number between 0 and 1
Importantly, the 'normal' Kernel and Error Kernel keep iterating
until a particle does have a longitude larger than 0.7.
This behaviour can be useful if particles need to be 'pushed out'
from e.g. land. Note however that current under-the-hood
implementation is not extremely efficient, so code could be slow."""
dimensions = {
'lon': np.linspace(0., 1., xdim, dtype=np.float32),
'lat': np.linspace(0., 1., ydim, dtype=np.float32)
}
data = {
'U': np.zeros((ydim, xdim), dtype=np.float32),
'V': np.zeros((ydim, xdim), dtype=np.float32)
}
fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
# Set minimum value for valid longitudes (i.e. all longitudes < minlon are 'land')
fieldset.add_constant('minlon', 0.7)
# create a ParticleSet with all particles starting at centre of Field
pset = ParticleSet.from_line(fieldset=fieldset,
pclass=ptype[mode],
start=(0.5, 0.5),
finish=(0.5, 0.5),
size=10)
def TestLon(particle, fieldset, time):
"""Kernel to check whether a longitude is larger than fieldset.minlon.
If not, the Kernel throws an error"""
if particle.lon <= fieldset.minlon:
return ErrorCode.Error
def Error_RandomiseLon(particle, fieldset, time):
"""Error handling kernel that draws a new longitude.
Note that this new longitude can be smaller than fieldset.minlon"""
particle.lon = random.uniform(0., 1.)
random.seed(123456)
# The .execute below is only run for one timestep. Yet the
# recovery={ErrorCode.Error: Error_RandomiseLon} assures Parcels keeps
# attempting to move all particles beyond 0.7 longitude
pset.execute(pset.Kernel(TestLon),
runtime=1,
dt=1,
recovery={ErrorCode.Error: Error_RandomiseLon})
assert (np.array([p.lon for p in pset]) > fieldset.minlon).all()
def pensinsula_example(fieldset,
npart,
mode='jit',
degree=1,
verbose=False,
output=True,
method=AdvectionRK4):
"""Example configuration of particle flow around an idealised Peninsula
:arg filename: Basename of the input fieldset
:arg npart: Number of particles to intialise"""
# First, we define a custom Particle class to which we add a
# custom variable, the initial stream function value p.
# We determine the particle base class according to mode.
class MyParticle(ptype[mode]):
# JIT compilation requires a-priori knowledge of the particle
# data structure, so we define additional variables here.
p = Variable('p', dtype=np.float32, initial=0.)
p_start = Variable('p_start', dtype=np.float32, initial=fieldset.P)
def __repr__(self):
"""Custom print function which overrides the built-in"""
return "P(%.4f, %.4f)[p=%.5f, p_start=%f]" % (self.lon, self.lat,
self.p, self.p_start)
# Initialise particles
x = 3. * (1. / 1.852 / 60) # 3 km offset from boundary
y = (fieldset.U.lat[0] + x, fieldset.U.lat[-1] - x
) # latitude range, including offsets
pset = ParticleSet.from_line(fieldset,
size=npart,
pclass=MyParticle,
start=(x, y[0]),
finish=(x, y[1]))
if verbose:
print("Initial particle positions:\n%s" % pset)
# Advect the particles for 24h
time = delta(hours=24)
dt = delta(minutes=5)
k_adv = pset.Kernel(method)
k_p = pset.Kernel(UpdateP)
out = pset.ParticleFile(name="MyParticle") if output else None
interval = delta(hours=1) if output else -1
print("Peninsula: Advecting %d particles for %s" % (npart, str(time)))
pset.execute(k_adv + k_p,
endtime=time,
dt=dt,
output_file=out,
interval=interval)
if verbose:
print("Final particle positions:\n%s" % pset)
return pset
def run_stommelgyre(fieldset, outfilename):
class MyParticle(JITParticle):
psi = Variable('psi', dtype=np.float32, initial=fieldset.Psi)
pset = ParticleSet.from_line(fieldset, size=4, pclass=MyParticle, time=0.,
start=(100, 5000), finish=(1000, 5000))
outfile = pset.ParticleFile(name=outfilename, outputdt=delta(hours=24))
pset.execute(AdvectionRK4 + pset.Kernel(UpdatePsi), runtime=delta(days=50),
dt=delta(minutes=5), output_file=outfile)
def test_pset_create_line(grid, mode, npart=100):
lon = np.linspace(0, 1, npart, dtype=np.float32)
lat = np.linspace(1, 0, npart, dtype=np.float32)
pset = ParticleSet.from_line(grid,
size=npart,
start=(0, 1),
finish=(1, 0),
pclass=ptype[mode])
assert np.allclose([p.lon for p in pset], lon, rtol=1e-12)
assert np.allclose([p.lat for p in pset], lat, rtol=1e-12)
def stommel_pset(fieldset, npart=1):
class MyParticle(JITParticle):
p = Variable('p', dtype=np.float32, initial=0.)
p_start = Variable('p_start', dtype=np.float32, initial=fieldset.P)
return ParticleSet.from_line(fieldset,
size=npart,
pclass=MyParticle,
start=(10e3, 5000e3),
finish=(1000e3, 5000e3),
time=0)
def test_pset_create_with_time(fieldset, mode, npart=100):
lon = np.linspace(0, 1, npart, dtype=np.float32)
lat = np.linspace(1, 0, npart, dtype=np.float32)
time = 5.
pset = ParticleSet(fieldset, lon=lon, lat=lat, pclass=ptype[mode], time=time)
assert np.allclose([p.time for p in pset], time, rtol=1e-12)
pset = ParticleSet.from_list(fieldset, lon=lon, lat=lat, pclass=ptype[mode],
time=[time]*npart)
assert np.allclose([p.time for p in pset], time, rtol=1e-12)
pset = ParticleSet.from_line(fieldset, size=npart, start=(0, 1), finish=(1, 0),
pclass=ptype[mode], time=time)
assert np.allclose([p.time for p in pset], time, rtol=1e-12)
def run_mitgcm_zonally_reentrant(mode):
"""Function that shows how to load MITgcm data in a zonally periodic domain"""
data_path = path.join(path.dirname(__file__), "MITgcm_example_data/")
filenames = {
"U": data_path + "mitgcm_UV_surface_zonally_reentrant.nc",
"V": data_path + "mitgcm_UV_surface_zonally_reentrant.nc",
}
variables = {"U": "UVEL", "V": "VVEL"}
dimensions = {
"U": {
"lon": "XG",
"lat": "YG",
"time": "time"
},
"V": {
"lon": "XG",
"lat": "YG",
"time": "time"
},
}
fieldset = FieldSet.from_mitgcm(filenames,
variables,
dimensions,
mesh="flat")
fieldset.add_periodic_halo(zonal=True)
fieldset.add_constant('domain_width', 1000000)
def periodicBC(particle, fieldset, time):
if particle.lon < 0:
particle.lon += fieldset.domain_width
elif particle.lon > fieldset.domain_width:
particle.lon -= fieldset.domain_width
# Release particles 5 cells away from the Eastern boundary
pset = ParticleSet.from_line(
fieldset,
pclass=ptype[mode],
start=(fieldset.U.grid.lon[-5], fieldset.U.grid.lat[5]),
finish=(fieldset.U.grid.lon[-5], fieldset.U.grid.lat[-5]),
size=10,
)
pfile = ParticleFile("MIT_particles_" + str(mode) + ".nc",
pset,
outputdt=delta(days=1))
kernels = AdvectionRK4 + pset.Kernel(periodicBC)
pset.execute(kernels,
runtime=delta(days=5),
dt=delta(minutes=30),
output_file=pfile)
pfile.close()
def run_timeoscillation(fieldset, outfilename):
pset = ParticleSet.from_line(fieldset,
pclass=JITParticle,
size=20,
start=(-10000, 0),
finish=(10000, 0))
outfile = pset.ParticleFile(name=outfilename, outputdt=delta(hours=3))
pset.execute(AdvectionRK4,
runtime=delta(days=4),
dt=delta(minutes=5),
output_file=outfile)
def test_fieldset_constant(mode):
data, dimensions = generate_fieldset(100, 100)
fieldset = FieldSet.from_data(data, dimensions)
westval = -0.2
eastval = 0.3
fieldset.add_constant('movewest', westval)
fieldset.add_constant('moveeast', eastval)
assert fieldset.movewest == westval
pset = ParticleSet.from_line(fieldset, size=1, pclass=ptype[mode],
start=(0.5, 0.5), finish=(0.5, 0.5))
pset.execute(pset.Kernel(addConst), dt=1, runtime=1)
assert abs(pset[0].lon - (0.5 + westval + eastval)) < 1e-4
def run_brownian(fieldset, npart, outfilename):
# Set diffusion constants.
fieldset.Kh_meridional = 100.
fieldset.Kh_zonal = 100.
# Set random seed
random.seed(123456)
pset = ParticleSet.from_line(fieldset=fieldset, size=npart, pclass=JITParticle,
start=(0., 0.), finish=(0., 0.))
pset.execute(two_dim_brownian_flat, runtime=delta(days=1), dt=delta(minutes=5))
pset.ParticleFile(name=outfilename).write(pset, pset[0].time)
def run_radialrotation(fieldset, outfilename):
# Define a ParticleSet
pset = ParticleSet.from_line(fieldset,
size=4,
pclass=JITParticle,
start=(0, 1000),
finish=(0, 4000))
# Advect the particles for 24h
outfile = pset.ParticleFile(name=outfilename, outputdt=delta(hours=1))
pset.execute(AdvectionRK4,
runtime=delta(days=1),
dt=delta(minutes=5),
output_file=outfile)
def test_sampling_multigrids_non_vectorfield(mode, npart):
xdim, ydim = 100, 200
U = Field('U',
np.zeros((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32))
V = Field('V',
np.zeros((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32))
B = Field('B',
np.ones((3 * ydim, 4 * xdim), dtype=np.float32),
lon=np.linspace(0., 1., 4 * xdim, dtype=np.float32),
lat=np.linspace(0., 1., 3 * ydim, dtype=np.float32))
fieldset = FieldSet(U, V)
fieldset.add_field(B, 'B')
fieldset.add_constant('sample_depth', 2.5)
assert fieldset.U.grid is fieldset.V.grid
assert fieldset.U.grid is not fieldset.B.grid
class TestParticle(ptype[mode]):
sample_var = Variable('sample_var', initial=0.)
pset = ParticleSet.from_line(fieldset,
pclass=TestParticle,
start=[0.3, 0.3],
finish=[0.7, 0.7],
size=npart)
def test_sample(particle, fieldset, time):
particle.sample_var += fieldset.B[time, fieldset.sample_depth,
particle.lat, particle.lon]
kernels = pset.Kernel(AdvectionRK4) + pset.Kernel(test_sample)
pset.execute(kernels, runtime=10, dt=1)
assert np.allclose(pset.sample_var, 10.0)
if mode == 'jit':
assert len(pset.xi.shape) == 2
assert pset.xi.shape[0] == len(pset.lon)
assert pset.xi.shape[1] == fieldset.gridset.size
assert np.all(pset.xi >= 0)
assert np.all(pset.xi[:, fieldset.B.igrid] < xdim * 4)
assert np.all(pset.xi[:, 0] < xdim)
assert pset.yi.shape[0] == len(pset.lon)
assert pset.yi.shape[1] == fieldset.gridset.size
assert np.all(pset.yi >= 0)
assert np.all(pset.yi[:, fieldset.B.igrid] < ydim * 3)
assert np.all(pset.yi[:, 0] < ydim)
def test_grid_constant(mode):
u, v, lon, lat, depth, time = generate_grid(100, 100)
grid = Grid.from_data(u, lon, lat, v, lon, lat, depth, time)
westval = -0.2
eastval = 0.3
grid.add_constant('movewest', westval)
grid.add_constant('moveeast', eastval)
assert grid.movewest == westval
pset = ParticleSet.from_line(grid,
size=1,
pclass=ptype[mode],
start=(0.5, 0.5),
finish=(0.5, 0.5))
pset.execute(pset.Kernel(addConst), dt=1, runtime=1)
assert abs(pset[0].lon - (0.5 + westval + eastval)) < 1e-4
def stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4):
timer.fieldset = timer.Timer('FieldSet', parent=timer.stommel)
fieldset = stommel_fieldset()
filename = 'stommel'
fieldset.write(filename)
timer.fieldset.stop()
# Determine particle class according to mode
timer.pset = timer.Timer('Pset', parent=timer.stommel)
timer.psetinit = timer.Timer('Pset_init', parent=timer.pset)
ParticleClass = JITParticle if mode == 'jit' else ScipyParticle
class MyParticle(ParticleClass):
p = Variable('p', dtype=np.float32, initial=0.)
p_start = Variable('p_start', dtype=np.float32, initial=fieldset.P)
pset = ParticleSet.from_line(fieldset,
size=npart,
pclass=MyParticle,
start=(100, 5000),
finish=(200, 5000),
time=0)
if verbose:
print("Initial particle positions:\n%s" % pset)
# Execute for 30 days, with 1hour timesteps and 12-hourly output
runtime = delta(days=30)
dt = delta(hours=1)
outputdt = delta(hours=12)
print("Stommel: Advecting %d particles for %s" % (npart, runtime))
timer.psetinit.stop()
timer.psetrun = timer.Timer('Pset_run', parent=timer.pset)
pset.execute(method + pset.Kernel(UpdateP),
runtime=runtime,
dt=dt,
moviedt=None,
output_file=pset.ParticleFile(name="StommelParticle",
outputdt=outputdt))
if verbose:
print("Final particle positions:\n%s" % pset)
timer.psetrun.stop()
timer.pset.stop()
return pset
def test_sampling_out_of_bounds_time(mode,
allow_time_extrapolation,
k_sample_p,
xdim=10,
ydim=10,
tdim=10):
lon = np.linspace(0., 1., xdim, dtype=np.float32)
lat = np.linspace(0., 1., ydim, dtype=np.float32)
time = np.linspace(0., 1., tdim, dtype=np.float64)
U = np.zeros((xdim, ydim, tdim), dtype=np.float32)
V = np.zeros((xdim, ydim, tdim), dtype=np.float32)
P = np.ones((xdim, ydim, 1), dtype=np.float32) * time
grid = Grid.from_data(U,
lon,
lat,
V,
lon,
lat,
time=time,
mesh='flat',
field_data={'P': np.asarray(P, dtype=np.float32)},
allow_time_extrapolation=allow_time_extrapolation)
pset = ParticleSet.from_line(grid,
size=1,
pclass=pclass(mode),
start=(0.5, 0.5),
finish=(0.5, 0.5))
if allow_time_extrapolation:
pset.execute(k_sample_p, starttime=-1.0, endtime=-0.9, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 0.0, rtol=1e-5)
else:
with pytest.raises(RuntimeError):
pset.execute(k_sample_p, starttime=-1.0, endtime=-0.9, dt=0.1)
pset.execute(k_sample_p, starttime=0.0, endtime=0.1, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 0.0, rtol=1e-5)
pset.execute(k_sample_p, starttime=0.5, endtime=0.6, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 0.5, rtol=1e-5)
pset.execute(k_sample_p, starttime=1.0, endtime=1.1, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 1.0, rtol=1e-5)
if allow_time_extrapolation:
pset.execute(k_sample_p, starttime=2.0, endtime=2.1, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 1.0, rtol=1e-5)
else:
with pytest.raises(RuntimeError):
pset.execute(k_sample_p, starttime=2.0, endtime=2.1, dt=0.1)
def test_sampling_out_of_bounds_time(mode,
allow_time_extrapolation,
k_sample_p,
xdim=10,
ydim=10,
tdim=10):
dimensions = {
'lon': np.linspace(0., 1., xdim, dtype=np.float32),
'lat': np.linspace(0., 1., ydim, dtype=np.float32),
'time': np.linspace(0., 1., tdim, dtype=np.float64)
}
data = {
'U': np.zeros((xdim, ydim, tdim), dtype=np.float32),
'V': np.zeros((xdim, ydim, tdim), dtype=np.float32),
'P': np.ones((xdim, ydim, 1), dtype=np.float32) * dimensions['time']
}
fieldset = FieldSet.from_data(
data,
dimensions,
mesh='flat',
allow_time_extrapolation=allow_time_extrapolation)
pset = ParticleSet.from_line(fieldset,
size=1,
pclass=pclass(mode),
start=(0.5, 0.5),
finish=(0.5, 0.5))
if allow_time_extrapolation:
pset.execute(k_sample_p, starttime=-1.0, endtime=-0.9, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 0.0, rtol=1e-5)
else:
with pytest.raises(RuntimeError):
pset.execute(k_sample_p, starttime=-1.0, endtime=-0.9, dt=0.1)
pset.execute(k_sample_p, starttime=0.0, endtime=0.1, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 0.0, rtol=1e-5)
pset.execute(k_sample_p, starttime=0.5, endtime=0.6, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 0.5, rtol=1e-5)
pset.execute(k_sample_p, starttime=1.0, endtime=1.1, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 1.0, rtol=1e-5)
if allow_time_extrapolation:
pset.execute(k_sample_p, starttime=2.0, endtime=2.1, dt=0.1)
assert np.allclose(np.array([p.p for p in pset]), 1.0, rtol=1e-5)
else:
with pytest.raises(RuntimeError):
pset.execute(k_sample_p, starttime=2.0, endtime=2.1, dt=0.1)
def run_pensinsula(fieldset, npart, outfilename):
class MyParticle(JITParticle):
psi = Variable('psi', dtype=np.float32, initial=fieldset.Psi)
# Initialise particles
x = 3. * (1. / 1.852 / 60)
y = (fieldset.U.lat[0] + x, fieldset.U.lat[-1] - x)
pset = ParticleSet.from_line(fieldset,
size=npart,
pclass=MyParticle,
start=(x, y[0]),
finish=(x, y[1]),
time=0.)
# Advect the particles for 24h
outfile = pset.ParticleFile(name=outfilename, outputdt=delta(hours=1))
pset.execute(AdvectionRK4 + pset.Kernel(UpdatePsi),
runtime=delta(hours=24),
dt=delta(minutes=5),
output_file=outfile)
def stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4):
grid = stommel_grid()
filename = 'stommel'
grid.write(filename)
# Determine particle class according to mode
ParticleClass = JITParticle if mode == 'jit' else ScipyParticle
class MyParticle(ParticleClass):
p = Variable('p', dtype=np.float32, initial=0.)
p_start = Variable('p_start', dtype=np.float32, initial=0.)
pset = ParticleSet.from_line(grid,
size=npart,
pclass=MyParticle,
start=(100, 5000),
finish=(200, 5000))
for particle in pset:
particle.p_start = grid.P[0., particle.lon, particle.lat]
if verbose:
print("Initial particle positions:\n%s" % pset)
# Execute for 50 days, with 5min timesteps and hourly output
runtime = delta(days=50)
dt = delta(minutes=5)
interval = delta(hours=12)
print("Stommel: Advecting %d particles for %s" % (npart, runtime))
pset.execute(method + pset.Kernel(UpdateP),
runtime=runtime,
dt=dt,
interval=interval,
output_file=pset.ParticleFile(name="StommelParticle"),
show_movie=False)
if verbose:
print("Final particle positions:\n%s" % pset)
return pset
def test_vector_fields(mode, swapUV):
lon = np.linspace(0., 10., 12, dtype=np.float32)
lat = np.linspace(0., 10., 10, dtype=np.float32)
U = np.ones((10, 12), dtype=np.float32)
V = np.zeros((10, 12), dtype=np.float32)
data = {'U': U, 'V': V}
dimensions = {'U': {'lat': lat, 'lon': lon},
'V': {'lat': lat, 'lon': lon}}
fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
if swapUV: # we test that we can freely edit whatever UV field
UV = VectorField('UV', fieldset.V, fieldset.U)
fieldset.add_vector_field(UV)
pset = ParticleSet.from_line(fieldset, size=1, pclass=ptype[mode],
start=(0.5, 0.5), finish=(0.5, 0.5))
pset.execute(AdvectionRK4, dt=1, runtime=1)
if swapUV:
assert abs(pset[0].lon - .5) < 1e-9
assert abs(pset[0].lat - 1.5) < 1e-9
else:
assert abs(pset[0].lon - 1.5) < 1e-9
assert abs(pset[0].lat - .5) < 1e-9
def rotation_example(grid, mode='jit', method=AdvectionRK4):
npart = 2 # Test two particles on the rotating grid.
pset = ParticleSet.from_line(
grid,
size=npart,
pclass=ptype[mode],
start=(30., 30.),
finish=(30., 50.)) # One particle in centre, one on periphery of grid.
endtime = delta(hours=17)
dt = delta(minutes=5)
interval = delta(hours=1)
pset.execute(method,
endtime=endtime,
dt=dt,
interval=interval,
output_file=pset.ParticleFile(name="RadialParticle"),
show_movie=False)
return pset
def rotation_example(fieldset, mode='jit', method=AdvectionRK4):
npart = 2 # Test two particles on the rotating fieldset.
pset = ParticleSet.from_line(
fieldset,
size=npart,
pclass=ptype[mode],
start=(30., 30.),
finish=(30.,
50.)) # One particle in centre, one on periphery of Field.
runtime = delta(hours=17)
dt = delta(minutes=5)
outputdt = delta(hours=1)
pset.execute(method,
runtime=runtime,
dt=dt,
moviedt=None,
output_file=pset.ParticleFile(name="RadialParticle",
outputdt=outputdt))
return pset
def test_sampling_multigrids_non_vectorfield_from_file(mode,
npart,
tmpdir,
chs,
filename='test_subsets'
):
xdim, ydim = 100, 200
filepath = tmpdir.join(filename)
U = Field('U',
np.zeros((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32))
V = Field('V',
np.zeros((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32))
B = Field('B',
np.ones((3 * ydim, 4 * xdim), dtype=np.float32),
lon=np.linspace(0., 1., 4 * xdim, dtype=np.float32),
lat=np.linspace(0., 1., 3 * ydim, dtype=np.float32))
fieldset = FieldSet(U, V)
fieldset.add_field(B, 'B')
fieldset.write(filepath)
fieldset = None
ufiles = [
filepath + 'U.nc',
] * 4
vfiles = [
filepath + 'V.nc',
] * 4
bfiles = [
filepath + 'B.nc',
] * 4
timestamps = np.arange(0, 4, 1) * 86400.0
timestamps = np.expand_dims(timestamps, 1)
files = {'U': ufiles, 'V': vfiles, 'B': bfiles}
variables = {'U': 'vozocrtx', 'V': 'vomecrty', 'B': 'B'}
dimensions = {'lon': 'nav_lon', 'lat': 'nav_lat'}
fieldset = FieldSet.from_netcdf(files,
variables,
dimensions,
timestamps=timestamps,
allow_time_extrapolation=True,
field_chunksize=chs)
fieldset.add_constant('sample_depth', 2.5)
assert fieldset.U.grid is fieldset.V.grid
assert fieldset.U.grid is not fieldset.B.grid
class TestParticle(ptype[mode]):
sample_var = Variable('sample_var', initial=0.)
pset = ParticleSet.from_line(fieldset,
pclass=TestParticle,
start=[0.3, 0.3],
finish=[0.7, 0.7],
size=npart)
def test_sample(particle, fieldset, time):
particle.sample_var += fieldset.B[time, fieldset.sample_depth,
particle.lat, particle.lon]
kernels = pset.Kernel(AdvectionRK4) + pset.Kernel(test_sample)
pset.execute(kernels, runtime=10, dt=1)
assert np.allclose(pset.sample_var, 10.0)
if mode == 'jit':
assert len(pset.xi.shape) == 2
assert pset.xi.shape[0] == len(pset.lon)
assert pset.xi.shape[1] == fieldset.gridset.size
assert np.all(pset.xi >= 0)
assert np.all(pset.xi[:, fieldset.B.igrid] < xdim * 4)
assert np.all(pset.xi[:, 0] < xdim)
assert pset.yi.shape[0] == len(pset.lon)
assert pset.yi.shape[1] == fieldset.gridset.size
assert np.all(pset.yi >= 0)
assert np.all(pset.yi[:, fieldset.B.igrid] < ydim * 3)
assert np.all(pset.yi[:, 0] < ydim)
