def test_random_kernel_concat(fieldset, mode, concat):
class TestParticle(ptype[mode]):
p = Variable('p', dtype=np.float32)
pset = ParticleSet(fieldset, pclass=TestParticle, lon=0, lat=0)
def RandomKernel(particle, fieldset, time):
particle.p += ParcelsRandom.uniform(0, 1)
def AddOne(particle, fieldset, time):
particle.p += 1.
kernels = pset.Kernel(RandomKernel) + pset.Kernel(
AddOne) if concat else RandomKernel
pset.execute(kernels, dt=0)
assert pset.p > 1 if concat else pset.p < 1
def test_pset_from_field(mode, xdim=10, ydim=20, npart=10000):
np.random.seed(123456)
dimensions = {
'lon': np.linspace(0., 1., xdim, dtype=np.float32),
'lat': np.linspace(0., 1., ydim, dtype=np.float32)
}
startfield = np.ones((xdim, ydim), dtype=np.float32)
for x in range(xdim):
startfield[x, :] = x
data = {
'U': np.zeros((xdim, ydim), dtype=np.float32),
'V': np.zeros((xdim, ydim), dtype=np.float32),
'start': startfield
}
fieldset = FieldSet.from_data(data,
dimensions,
mesh='flat',
transpose=True)
densfield = Field(name='densfield',
data=np.zeros((xdim + 1, ydim + 1), dtype=np.float32),
lon=np.linspace(-1. / (xdim * 2),
1. + 1. / (xdim * 2),
xdim + 1,
dtype=np.float32),
lat=np.linspace(-1. / (ydim * 2),
1. + 1. / (ydim * 2),
ydim + 1,
dtype=np.float32),
transpose=True)
fieldset.add_field(densfield)
pset = ParticleSet.from_field(fieldset,
size=npart,
pclass=pclass(mode),
start_field=fieldset.start)
pdens = pset.density(field_name='densfield', relative=True)[:-1, :-1]
assert np.allclose(np.transpose(pdens),
startfield / np.sum(startfield),
atol=1e-2)
def test_advection_zonal(lon, lat, depth, mode, npart=10):
""" Particles at high latitude move geographically faster due to
the pole correction in `GeographicPolar`.
"""
data2D = {
'U': np.ones((lon.size, lat.size), dtype=np.float32),
'V': np.zeros((lon.size, lat.size), dtype=np.float32)
}
data3D = {
'U': np.ones((lon.size, lat.size, depth.size), dtype=np.float32),
'V': np.zeros((lon.size, lat.size, depth.size), dtype=np.float32)
}
dimensions = {'lon': lon, 'lat': lat}
fieldset2D = FieldSet.from_data(data2D,
dimensions,
mesh='spherical',
transpose=True)
pset2D = ParticleSet(fieldset2D,
pclass=ptype[mode],
lon=np.zeros(npart, dtype=np.float32) + 20.,
lat=np.linspace(0, 80, npart, dtype=np.float32))
pset2D.execute(AdvectionRK4, runtime=delta(hours=2), dt=delta(seconds=30))
assert (np.diff(np.array([p.lon for p in pset2D])) > 1.e-4).all()
dimensions['depth'] = depth
fieldset3D = FieldSet.from_data(data3D,
dimensions,
mesh='spherical',
transpose=True)
pset3D = ParticleSet(fieldset3D,
pclass=ptype[mode],
lon=np.zeros(npart, dtype=np.float32) + 20.,
lat=np.linspace(0, 80, npart, dtype=np.float32),
depth=np.zeros(npart, dtype=np.float32) + 10.)
pset3D.execute(AdvectionRK4, runtime=delta(hours=2), dt=delta(seconds=30))
assert (np.diff(np.array([p.lon for p in pset3D])) > 1.e-4).all()
def test_nearest_neighbour_interpolation3D(mode, k_sample_p, npart=81):
dims = (2, 2, 2)
dimensions = {
'lon': np.linspace(0., 1., dims[0], dtype=np.float32),
'lat': np.linspace(0., 1., dims[1], dtype=np.float32),
'depth': np.linspace(0., 1., dims[2], dtype=np.float32)
}
data = {
'U': np.zeros(dims, dtype=np.float32),
'V': np.zeros(dims, dtype=np.float32),
'P': np.zeros(dims, dtype=np.float32)
}
data['P'][0, 1, 1] = 1.
fieldset = FieldSet.from_data(data,
dimensions,
mesh='flat',
transpose=True)
fieldset.P.interp_method = 'nearest'
xv, yv = np.meshgrid(np.linspace(0, 1.0, int(np.sqrt(npart))),
np.linspace(0, 1.0, int(np.sqrt(npart))))
# combine a pset at 0m with pset at 1m, as meshgrid does not do 3D
pset = ParticleSet(fieldset,
pclass=pclass(mode),
lon=xv.flatten(),
lat=yv.flatten(),
depth=np.zeros(npart))
pset2 = ParticleSet(fieldset,
pclass=pclass(mode),
lon=xv.flatten(),
lat=yv.flatten(),
depth=np.ones(npart))
pset.add(pset2)
pset.execute(k_sample_p, endtime=1, dt=1)
assert np.allclose(pset.p[(pset.lon < 0.5) & (pset.lat > 0.5) &
(pset.depth > 0.5)],
1.0,
rtol=1e-5)
assert np.allclose(pset.p[(pset.lon > 0.5) | (pset.lat < 0.5) &
(pset.depth < 0.5)],
0.0,
rtol=1e-5)
def test_moving_eddies_fwdbwd(mode, mesh, tmpdir, npart=2):
method = AdvectionRK4
fieldset = moving_eddies_fieldset(mesh=mesh)
# Determine particle class according to mode
lons = [3.3, 3.3
] if fieldset.U.grid.mesh == 'spherical' else [3.3e5, 3.3e5]
lats = [46., 47.8] if fieldset.U.grid.mesh == 'spherical' else [1e5, 2.8e5]
pset = ParticleSet(fieldset=fieldset,
pclass=ptype[mode],
lon=lons,
lat=lats)
# 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)))
outfile = tmpdir.join("EddyParticlefwd")
pset.execute(method,
runtime=runtime,
dt=dt,
output_file=pset.ParticleFile(name=outfile,
outputdt=outputdt))
print("Now running in backward time mode")
outfile = tmpdir.join("EddyParticlebwd")
pset.execute(method,
endtime=0,
dt=-dt,
output_file=pset.ParticleFile(name=outfile,
outputdt=outputdt))
assert np.allclose([p.lon for p in pset], lons)
assert np.allclose([p.lat for p in pset], lats)
return pset
def run_corefootprintparticles(dirwrite, outfile, lonss, latss, time):
snapshots = snapshot_function(date(1980, 1, 3), date(2010, 12, 30),
delta(days=3))
fieldset = set_ofes_fieldset(snapshots)
fieldset.add_periodic_halo(zonal=True)
fieldset.B.allow_time_extrapolation = True
fieldset.add_constant('dwellingdepth', np.float(dd))
fieldset.add_constant('sinkspeed', sp / 86400.)
fieldset.add_constant('maxage', 300000. * 86400)
fieldset.add_constant('surface', 2.6)
class DinoOfesParticle(JITParticle):
temp = Variable('temp', dtype=np.float32, initial=np.nan)
age = Variable('age', dtype=np.float32, initial=0.)
salin = Variable('salin', dtype=np.float32, initial=np.nan)
lon0 = Variable('lon0', dtype=np.float32, initial=0., to_write=False)
lat0 = Variable('lat0', dtype=np.float32, initial=0., to_write=False)
depth0 = Variable('depth0',
dtype=np.float32,
initial=0.,
to_write=False)
time0 = Variable('time0', dtype=np.float32, initial=0., to_write=False)
pset = ParticleSet.from_list(fieldset=fieldset,
pclass=DinoOfesParticle,
lon=lonss.tolist(),
lat=latss.tolist(),
time=time) #depth=depths,
pfile = ParticleFile(dirwrite + outfile, pset, outputdt=delta(days=3))
kernels = pset.Kernel(initials) + Sink + pset.Kernel(
AdvectionRK4_3D) + Age + periodicBC
pset.execute(kernels,
runtime=delta(days=2555),
dt=delta(minutes=-5),
output_file=pfile,
recovery={ErrorCode.ErrorOutOfBounds: DeleteParticle})
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 run_nemo_curvilinear(mode, outfile):
"""Function that shows how to read in curvilinear grids, in this case from NEMO"""
data_path = path.join(path.dirname(__file__), 'NemoCurvilinear_data/')
filenames = {
'U': {
'lon': data_path + 'mesh_mask.nc4',
'lat': data_path + 'mesh_mask.nc4',
'data': data_path + 'U_purely_zonal-ORCA025_grid_U.nc4'
},
'V': {
'lon': data_path + 'mesh_mask.nc4',
'lat': data_path + 'mesh_mask.nc4',
'data': data_path + 'V_purely_zonal-ORCA025_grid_V.nc4'
}
}
variables = {'U': 'U', 'V': 'V'}
dimensions = {'lon': 'glamf', 'lat': 'gphif'}
field_set = FieldSet.from_nemo(filenames, variables, dimensions)
# Now run particles as normal
npart = 20
lonp = 30 * np.ones(npart)
latp = [i for i in np.linspace(-70, 88, npart)]
def periodicBC(particle, pieldSet, time, dt):
if particle.lon > 180:
particle.lon -= 360
pset = ParticleSet.from_list(field_set, ptype[mode], lon=lonp, lat=latp)
pfile = ParticleFile(outfile, pset, outputdt=delta(days=1))
kernels = pset.Kernel(AdvectionRK4) + periodicBC
pset.execute(kernels,
runtime=delta(days=1) * 160,
dt=delta(hours=6),
output_file=pfile)
assert np.allclose([pset[i].lat - latp[i] for i in range(len(pset))],
0,
atol=2e-2)
def test_pset_from_field(xdim=10, ydim=10, npart=10000):
np.random.seed(123456)
dimensions = {
'lon': np.linspace(0., 1., xdim, dtype=np.float32),
'lat': np.linspace(0., 1., ydim, dtype=np.float32)
}
startfield = np.ones((xdim, ydim), dtype=np.float32)
for x in range(xdim):
startfield[x, :] = x
data = {
'U': np.zeros((xdim, ydim), dtype=np.float32),
'V': np.zeros((xdim, ydim), dtype=np.float32),
'start': startfield
}
fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
pset = ParticleSet.from_field(fieldset,
size=npart,
pclass=JITParticle,
start_field=fieldset.start)
pdens = pset.density(area_scale=False, relative=True)
assert np.allclose(pdens, startfield / np.sum(startfield), atol=5e-3)
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 test_multigrids_pointer(mode):
lon_g0 = np.linspace(0, 1e4, 21, dtype=np.float32)
lat_g0 = np.linspace(0, 1000, 2, dtype=np.float32)
depth_g0 = np.zeros((5, lat_g0.size, lon_g0.size), dtype=np.float32)
def bath_func(lon):
return lon / 1000. + 10
bath = bath_func(lon_g0)
zdim = depth_g0.shape[0]
for i in range(lon_g0.size):
for k in range(zdim):
depth_g0[k, :, i] = bath[i] * k / (zdim - 1)
grid_0 = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0)
grid_1 = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0)
u_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32)
v_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32)
w_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32)
u_field = Field('U', u_data, grid=grid_0)
v_field = Field('V', v_data, grid=grid_0)
w_field = Field('W', w_data, grid=grid_1)
field_set = FieldSet(u_field, v_field, fields={'W': w_field})
assert (u_field.grid == v_field.grid)
assert (u_field.grid == w_field.grid
) # w_field.grid is now supposed to be grid_1
pset = ParticleSet.from_list(field_set,
ptype[mode],
lon=[0],
lat=[0],
depth=[1])
for i in range(10):
pset.execute(AdvectionRK4_3D, runtime=1000, dt=500)
def moving_eddies_example(fieldset,
outfile,
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
start = (3.3, 46.) if fieldset.U.grid.mesh == 'spherical' else (3.3e5, 1e5)
finish = (3.3, 47.8) if fieldset.U.grid.mesh == 'spherical' else (3.3e5,
2.8e5)
pset = ParticleSet.from_line(fieldset=fieldset,
size=npart,
pclass=ptype[mode],
start=start,
finish=finish)
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=outfile,
outputdt=delta(hours=1)),
moviedt=None)
if verbose:
print("Final particle positions:\n%s" % pset)
return pset
def test_write_timebackward(fieldset, mode, tmpdir):
outfilepath = tmpdir.join("pfile_write_timebackward.nc")
def Update_lon(particle, fieldset, time):
particle.lon -= 0.1 * particle.dt
pset = ParticleSet(fieldset,
pclass=JITParticle,
lat=np.linspace(0, 1, 3),
lon=[0, 0, 0],
time=[1, 2, 3])
pfile = pset.ParticleFile(name=outfilepath, outputdt=1.)
pset.execute(pset.Kernel(Update_lon), runtime=4, dt=-1., output_file=pfile)
ncfile = close_and_compare_netcdffiles(outfilepath, pfile)
trajs = ncfile.variables['trajectory'][:, 0]
assert np.all(
np.diff(trajs) > 0) # all particles written in order of traj ID
def test_multi_kernel_reuse_varnames(fieldset, mode):
# Testing for merging of two Kernels with the same variable declared
# Should throw a warning, but go ahead regardless
def MoveEast1(particle, fieldset, time, dt):
add_lon = 0.2
particle.lon += add_lon
def MoveEast2(particle, fieldset, time, dt):
particle.lon += add_lon # NOQA - no flake8 testing of this line
pset = ParticleSet(fieldset, pclass=ptype[mode], lon=[0.5], lat=[0.5])
pset.execute(pset.Kernel(MoveEast1) + pset.Kernel(MoveEast2),
starttime=0.,
endtime=1.,
dt=1.)
assert np.allclose([p.lon for p in pset], [0.9],
rtol=1e-5) # should be 0.5 + 0.2 + 0.2 = 0.9
def test_multi_kernel_duplicate_varnames(fieldset, mode):
# Testing for merging of two Kernels with the same variable declared
# Should throw a warning, but go ahead regardless
def MoveEast(particle, fieldset, time, dt):
add_lon = 0.1
particle.lon += add_lon
def MoveWest(particle, fieldset, time, dt):
add_lon = -0.3
particle.lon += add_lon
pset = ParticleSet(fieldset, pclass=ptype[mode], lon=[0.5], lat=[0.5])
pset.execute(pset.Kernel(MoveEast) + pset.Kernel(MoveWest),
starttime=0.,
endtime=1.,
dt=1.)
assert np.allclose([p.lon for p in pset], 0.3, rtol=1e-5)
def test_rectilinear_s_grids_advect1(mode):
# Constant water transport towards the east. check that the particle stays at the same relative depth (z/bath)
lon_g0 = np.linspace(0, 1e4, 21, dtype=np.float32)
lat_g0 = np.linspace(0, 1000, 2, dtype=np.float32)
depth_g0 = np.zeros((lon_g0.size, lat_g0.size, 5), dtype=np.float32)
def bath_func(lon):
return lon / 1000. + 10
bath = bath_func(lon_g0)
for i in range(depth_g0.shape[0]):
for k in range(depth_g0.shape[2]):
depth_g0[i, :, k] = bath[i] * k / (depth_g0.shape[2]-1)
grid = RectilinearSGrid('grid0py', lon_g0, lat_g0, depth=depth_g0)
u_data = np.zeros((lon_g0.size, lat_g0.size, depth_g0.shape[2]), dtype=np.float32)
v_data = np.zeros((lon_g0.size, lat_g0.size, depth_g0.shape[2]), dtype=np.float32)
w_data = np.zeros((lon_g0.size, lat_g0.size, depth_g0.shape[2]), dtype=np.float32)
for i in range(depth_g0.shape[0]):
u_data[i, :, :] = 1 * 10 / bath[i]
for k in range(depth_g0.shape[2]):
w_data[i, :, k] = u_data[i, :, k] * depth_g0[i, :, k] / bath[i] * 1e-3
u_field = Field('U', u_data, grid=grid, transpose=True)
v_field = Field('V', v_data, grid=grid, transpose=True)
w_field = Field('W', w_data, grid=grid, transpose=True)
field_set = FieldSet(u_field, v_field, fields={'W': w_field})
lon = np.zeros((11))
lat = np.zeros((11))
ratio = [min(i/10., .99) for i in range(11)]
depth = bath_func(lon)*ratio
pset = ParticleSet.from_list(field_set, ptype[mode], lon=lon, lat=lat, depth=depth)
pset.execute(AdvectionRK4_3D, runtime=10000, dt=500)
assert np.allclose([p.depth/bath_func(p.lon) for p in pset], ratio)
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=(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)
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 stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4):
fieldset = stommel_fieldset()
filename = 'stommel'
fieldset.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=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)
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 compute_MITgcm_particle_advection(field_set, mode, lonp, latp):
def periodicBC(particle, fieldSet, time):
dlon = -89.0 + 91.8
dlat = 0.7 + 1.4
if particle.lon > -89.0:
particle.lon -= dlon
if particle.lon < -91.8:
particle.lon += dlon
if particle.lat > 0.7:
particle.lat -= dlat
if particle.lat < -1.4:
particle.lat += dlat
pset = ParticleSet.from_list(field_set, ptype[mode], lon=lonp, lat=latp)
pfile = ParticleFile("MITgcm_particles_chunk",
pset,
outputdt=delta(days=1))
kernels = pset.Kernel(periodicBC) + pset.Kernel(AdvectionRK4)
pset.execute(kernels,
runtime=delta(days=30),
dt=delta(hours=12),
output_file=pfile)
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_variable_write_double(fieldset, mode, tmpdir):
filepath = tmpdir.join("pfile_variable_write_double.nc")
def Update_lon(particle, fieldset, time):
particle.lon += 0.1
pset = ParticleSet(fieldset,
pclass=ptype[mode],
lon=[0],
lat=[0],
lonlatdepth_dtype=np.float64)
ofile = pset.ParticleFile(name=filepath, outputdt=0.1)
pset.execute(pset.Kernel(Update_lon), endtime=1, dt=0.1, output_file=ofile)
ncfile = close_and_compare_netcdffiles(filepath, ofile)
lons = ncfile.variables['lon'][:]
assert (isinstance(lons[0, 0], np.float64))
ncfile.close()
def run_corefootprintparticles(dirwrite, outfile, lonss, latss, dep):
ufiles = sorted(glob(dirread_top + 'means/ORCA0083-N06_200?????d05U.nc'))
vfiles = sorted(glob(dirread_top + 'means/ORCA0083-N06_200?????d05V.nc'))
wfiles = sorted(glob(dirread_top + 'means/ORCA0083-N06_200?????d05W.nc'))
tfiles = sorted(glob(dirread_top + 'means/ORCA0083-N06_200?????d05T.nc'))
fieldset = set_nemo_fieldset(
ufiles, vfiles, wfiles, tfiles, dirread_pal +
'domain/coordinates.nc') # coordinates.nc is the mesh mask here
fieldset.add_constant('dwellingdepth', np.float(dd))
fieldset.add_constant('maxage', 300000. * 86400)
fieldset.add_constant('surface', 2.5)
class FixParticle(JITParticle):
temp = Variable('temp', dtype=np.float32, initial=np.nan)
age = Variable('age', dtype=np.float32, initial=0.)
salin = Variable('salin', dtype=np.float32, initial=np.nan)
pset = ParticleSet.from_list(fieldset=fieldset,
pclass=FixParticle,
lon=lonss.tolist(),
lat=latss.tolist(),
time=time)
pfile = ParticleFile(dirwrite + outfile, pset, outputdt=delta(days=6))
kernels = pset.Kernel(SampleSurf) + Age
pset.execute(kernels,
runtime=delta(days=2170),
dt=delta(days=-3),
output_file=pfile,
verbose_progress=False,
recovery={ErrorCode.ErrorOutOfBounds: DeleteParticle})
print 'Execution finished'
def test_variable_write_double(fieldset, mode, tmpdir):
filepath = tmpdir.join("pfile_variable_write_double")
def Update_lon(particle, fieldset, time):
particle.lon += 0.1
pset = ParticleSet(fieldset,
pclass=JITParticle,
lon=[0],
lat=[0],
lonlatdepth_dtype=np.float64)
pset.execute(pset.Kernel(Update_lon),
endtime=1,
dt=0.1,
output_file=pset.ParticleFile(name=filepath, outputdt=0.1))
ncfile = Dataset(filepath + ".nc", 'r', 'NETCDF4')
lons = ncfile.variables['lon'][:]
assert (isinstance(lons[0, 0], np.float64))
def test_fieldset_initialisation_kernel_dask(
time2, tmpdir, filename='test_parcels_defer_loading'):
filepath = tmpdir.join(filename)
data0, dims0 = generate_fieldset(3, 3, 4, 10)
data0['U'] = np.random.rand(10, 4, 3, 3)
dims0['time'] = np.arange(0, 10, 1)
dims0['depth'] = np.arange(0, 4, 1)
fieldset_out = FieldSet.from_data(data0, dims0)
fieldset_out.write(filepath)
fieldset = FieldSet.from_parcels(filepath,
chunksize={
'time': ('time_counter', 1),
'depth': ('depthu', 1),
'lat': ('y', 2),
'lon': ('x', 2)
})
def SampleField(particle, fieldset, time):
particle.u_kernel = fieldset.U[time, particle.depth, particle.lat,
particle.lon]
class SampleParticle(JITParticle):
u_kernel = Variable('u_kernel', dtype=np.float32, initial=0.)
u_scipy = Variable('u_scipy', dtype=np.float32, initial=fieldset.U)
pset = ParticleSet(fieldset,
pclass=SampleParticle,
time=[0, time2],
lon=[0.5, 0.5],
lat=[0.5, 0.5],
depth=[0.5, 0.5])
if time2 > 1:
failed = False
try:
pset.execute(SampleField, dt=0.)
except OutOfTimeError:
failed = True
assert failed
else:
pset.execute(SampleField, dt=0.)
assert np.allclose([p.u_kernel for p in pset],
[p.u_scipy for p in pset])
assert isinstance(fieldset.U.data, da.core.Array)
def test_fieldset_defer_loading_with_diff_time_origin(
tmpdir, fail, filename='test_parcels_defer_loading'):
filepath = tmpdir.join(filename)
data0, dims0 = generate_fieldset(10, 10, 1, 10)
dims0['time'] = np.arange(0, 10, 1) * 3600
fieldset_out = FieldSet.from_data(data0, dims0)
fieldset_out.U.grid.time_origin = TimeConverter(
np.datetime64('2018-04-20'))
fieldset_out.V.grid.time_origin = TimeConverter(
np.datetime64('2018-04-20'))
data1, dims1 = generate_fieldset(10, 10, 1, 10)
if fail:
dims1['time'] = np.arange(0, 10, 1) * 3600
else:
dims1['time'] = np.arange(0, 10, 1) * 1800 + (24 + 25) * 3600
if fail:
Wtime_origin = TimeConverter(np.datetime64('2018-04-22'))
else:
Wtime_origin = TimeConverter(np.datetime64('2018-04-18'))
gridW = RectilinearZGrid(dims1['lon'],
dims1['lat'],
dims1['depth'],
dims1['time'],
time_origin=Wtime_origin)
fieldW = Field('W', np.zeros(data1['U'].shape), grid=gridW)
fieldset_out.add_field(fieldW)
fieldset_out.write(filepath)
fieldset = FieldSet.from_parcels(filepath, extra_fields={'W': 'W'})
assert fieldset.U.creation_log == 'from_parcels'
pset = ParticleSet.from_list(fieldset,
pclass=JITParticle,
lon=[0.5],
lat=[0.5],
depth=[0.5],
time=[datetime.datetime(2018, 4, 20, 1)])
pset.execute(AdvectionRK4_3D, runtime=delta(hours=4), dt=delta(hours=1))
def test_mitgcm(mode, chunk_mode, using_add_field):
if chunk_mode in [
'auto',
]:
dask.config.set({'array.chunk-size': '256KiB'})
else:
dask.config.set({'array.chunk-size': '128MiB'})
field_set = fieldset_from_mitgcm(chunk_mode, using_add_field)
lons, lats = 5e5, 5e5
pset = ParticleSet.from_list(fieldset=field_set,
pclass=ptype[mode],
lon=lons,
lat=lats)
pset.execute(AdvectionRK4, runtime=delta(days=1), dt=delta(minutes=5))
# MITgcm sample file dimensions: time=10, XG=400, YG=200
if chunk_mode != 'specific_different':
assert (len(field_set.U.grid.load_chunk) == len(
field_set.V.grid.load_chunk))
if chunk_mode in [
False,
]:
assert (len(field_set.U.grid.load_chunk) == 1)
elif chunk_mode in [
'auto',
]:
assert (len(field_set.U.grid.load_chunk) != 1)
elif 'specific' in chunk_mode:
assert (len(
field_set.U.grid.load_chunk) == (1 * int(math.ceil(400.0 / 50.0)) *
int(math.ceil(200.0 / 100.0))))
if chunk_mode == 'specific_same':
assert field_set.gridset.size == 1
elif chunk_mode == 'specific_different':
assert field_set.gridset.size == 2
assert np.allclose(pset[0].lon, 5.27e5, atol=1e3)
def test_add_field_after_pset(fieldtype):
data, dimensions = generate_fieldset(100, 100)
fieldset = FieldSet.from_data(data, dimensions)
pset = ParticleSet(fieldset, ScipyParticle, lon=0,
lat=0) # noqa ; to trigger fieldset.check_complete
field1 = Field('field1',
fieldset.U.data,
lon=fieldset.U.lon,
lat=fieldset.U.lat)
field2 = Field('field2',
fieldset.U.data,
lon=fieldset.U.lon,
lat=fieldset.U.lat)
vfield = VectorField('vfield', field1, field2)
error_thrown = False
try:
if fieldtype == 'normal':
fieldset.add_field(field1)
elif fieldtype == 'vector':
fieldset.add_vector_field(vfield)
except RuntimeError:
error_thrown = True
assert error_thrown
def test_random_field(mode, k_sample_p, xdim=20, ydim=20, npart=100):
"""Sampling test that tests for overshoots by sampling a field of
random numbers between 0 and 1.
"""
np.random.seed(123456)
dimensions = {
'lon': np.linspace(0., 1., xdim, dtype=np.float32),
'lat': np.linspace(0., 1., ydim, dtype=np.float32)
}
data = {
'U': np.zeros((xdim, ydim), dtype=np.float32),
'V': np.zeros((xdim, ydim), dtype=np.float32),
'P': np.random.uniform(0, 1., size=(xdim, ydim)),
'start': np.ones((xdim, ydim), dtype=np.float32)
}
fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
pset = ParticleSet.from_field(fieldset,
size=npart,
pclass=pclass(mode),
start_field=fieldset.start)
pset.execute(k_sample_p, endtime=1., dt=1.0)
sampled = np.array([p.p for p in pset])
assert ((sampled >= 0.).all())
def test_moving_eddies_fwdbwd(mode, npart=2):
method = AdvectionRK4
grid = moving_eddies_grid()
# 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))
# Execte for 14 days, with 30sec timesteps and hourly output
endtime = delta(days=1)
dt = delta(minutes=5)
interval = delta(hours=1)
print("MovingEddies: Advecting %d particles for %s" %
(npart, str(endtime)))
pset.execute(method,
starttime=0,
endtime=endtime,
dt=dt,
interval=interval,
output_file=pset.ParticleFile(name="EddyParticlefwd"))
print("Now running in backward time mode")
pset.execute(method,
starttime=endtime,
endtime=0,
dt=-dt,
interval=-interval,
output_file=pset.ParticleFile(name="EddyParticlebwd"))
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_curvilinear_grids(mode):
x = np.linspace(0, 1e3, 7, dtype=np.float32)
y = np.linspace(0, 1e3, 5, dtype=np.float32)
(xx, yy) = np.meshgrid(x, y)
r = np.sqrt(xx * xx + yy * yy)
theta = np.arctan2(yy, xx)
theta = theta + np.pi / 6.
lon = r * np.cos(theta)
lat = r * np.sin(theta)
time = np.array([0, 86400], dtype=np.float64)
grid = CurvilinearZGrid(lon, lat, time=time)
u_data = np.ones((2, y.size, x.size), dtype=np.float32)
v_data = np.zeros((2, y.size, x.size), dtype=np.float32)
u_data[0, :, :] = lon[:, :] + lat[:, :]
u_field = Field('U', u_data, grid=grid, transpose=False)
v_field = Field('V', v_data, grid=grid, transpose=False)
field_set = FieldSet(u_field, v_field)
def sampleSpeed(particle, fieldset, time):
u = fieldset.U[time, particle.depth, particle.lat, particle.lon]
v = fieldset.V[time, particle.depth, particle.lat, particle.lon]
particle.speed = math.sqrt(u * u + v * v)
class MyParticle(ptype[mode]):
speed = Variable('speed', dtype=np.float32, initial=0.)
pset = ParticleSet.from_list(field_set,
MyParticle,
lon=[400, -200],
lat=[600, 600])
pset.execute(pset.Kernel(sampleSpeed), runtime=0, dt=0)
assert (np.allclose(pset[0].speed, 1000))