def test_sampling_multigrids_non_vectorfield(pset_mode, 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 = pset_type[pset_mode]['pset'].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':
if pset_mode == 'soa':
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)
elif pset_mode == 'aos':
assert np.alltrue([[pxi > 0 for pxi in p.xi] for p in pset])
assert np.alltrue([len(p.xi) == fieldset.gridset.size for p in pset])
assert np.alltrue([p.xi[fieldset.B.igrid] < xdim * 4 for p in pset])
assert np.alltrue([p.xi[0] < xdim for p in pset])
assert np.alltrue([[pyi > 0 for pyi in p.yi] for p in pset])
assert np.alltrue([len(p.yi) == fieldset.gridset.size for p in pset])
assert np.alltrue([p.yi[fieldset.B.igrid] < ydim * 3 for p in pset])
assert np.alltrue([p.yi[0] < ydim for p in 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)
interp_method='linear',
mesh='spherical',
allow_time_extrapolation=True)
fieldset.add_field(iso_psi_all)
fieldset.add_field(source_id_psi)
fieldset.add_field(sink_id_psi)
fieldset.add_field(cdist)
fieldset.add_field(cnormx)
fieldset.add_field(cnormy)
fieldset.cnormx_rho.units = GeographicPolar()
fieldset.cnormy_rho.units = Geographic()
# ADD THE PERIODIC BOUNDARY
fieldset.add_constant('halo_west', -180.)
fieldset.add_constant('halo_east', 180.)
fieldset.add_periodic_halo(zonal=True)
# ADD DIFFUSION
if param['Kh']:
fieldset.add_constant_field('Kh_zonal', param['Kh'], mesh='spherical')
fieldset.add_constant_field('Kh_meridional', param['Kh'], mesh='spherical')
# ADD MAXIMUM PARTICLE AGE (IF LIMITED AGE)
fieldset.add_constant('max_age', param['max_age'])
# SET SEED
ParcelsRandom.seed(690)
##############################################################################
'lat': 'lat_rho'
},
interp_method='linear',
mesh='spherical',
allow_time_extrapolation=True)
fieldset.add_field(iso_psi)
fieldset.add_field(cdist)
fieldset.add_field(cnormx)
fieldset.add_field(cnormy)
fieldset.cnormx_rho.units = GeographicPolar()
fieldset.cnormy_rho.units = Geographic()
# ADD THE PERIODIC BOUNDARY
fieldset.add_constant('halo_west', -180.)
fieldset.add_constant('halo_east', 180.)
fieldset.add_periodic_halo(zonal=True)
# ADD MAXIMUM PARTICLE AGE (IF LIMITED AGE)
if param['max_age']:
fieldset.add_constant('max_age', param['max_age'] * 3600 * 24 * 365.25)
##############################################################################
# KERNELS #
##############################################################################
class debris(JITParticle):
# Coast mask (if particle is in a coastal cell)
cm = Variable('cm', dtype=np.float64, initial=0, to_write=False)
V=NestedField('V', vfields)
if options['run3D']:
W=NestedField('W', wfields)
if options['vdiffusion']:
Kz=NestedField('Kz', Kzfields)
Kz_EVD=NestedField('Kz_EVD', KzEVDfields)
if options['beaching']>0:
tmask=NestedField('tmask', tmaskfields)
if options['run3D']:
if options['vdiffusion']:
if options['beaching']>0:
fieldset = FieldSet(U, V, {'W':W, 'Kz':Kz, 'Kz_EVD':Kz_EVD, 'tmask':tmask})
else:
fieldset = FieldSet(U, V, {'W':W, 'Kz':Kz, 'Kz_EVD':Kz_EVD})
fieldset.add_constant('dres', 0.01)
else:
if options['beaching']>0:
fieldset = FieldSet(U, V, {'W':W, 'tmask':tmask})
else:
fieldset = FieldSet(U, V, {'W':W})
else:
if options['beaching']>0:
fieldset = FieldSet(U, V, {'tmask':tmask})
else:
fieldset = FieldSet(U, V)
else:
print('No nesting, using just 1 grid')
fieldset=myfieldset[0]
del myfieldset
