def test_multiple_grid_addlater_error():
xdim, ydim = 10, 20
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))
fieldset = FieldSet(U, V)
pset = pset_type['soa']['pset'](
fieldset, pclass=pclass('jit'), lon=[0.8],
lat=[0.9]) # noqa ; to trigger fieldset.check_complete
P = Field('P',
np.zeros((ydim * 10, xdim * 10), dtype=np.float32),
lon=np.linspace(0., 1., xdim * 10, dtype=np.float32),
lat=np.linspace(0., 1., ydim * 10, dtype=np.float32))
fail = False
try:
fieldset.add_field(P)
except RuntimeError:
fail = True
assert fail
def test_pseudo_interaction(runtime, dt):
# A linear field where advected particles are moving at
# 1 m/s in the longitudinal direction.
xdim, ydim = (10, 20)
Uflow = Field('U',
np.ones((ydim, xdim), dtype=np.float64),
lon=np.linspace(0., 1e3, xdim, dtype=np.float64),
lat=np.linspace(0., 1e3, ydim, dtype=np.float64))
Vflow = Field('V',
np.zeros((ydim, xdim), dtype=np.float64),
grid=Uflow.grid)
fieldset = FieldSet(Uflow, Vflow)
# Execute the advection kernel only
pset = ParticleSet(fieldset, pclass=ScipyParticle, lon=[2], lat=[2])
pset.execute(AdvectionRK4, runtime=runtime, dt=dt)
# Execute both the advection and interaction kernel.
pset2 = ParticleSet(fieldset,
pclass=ScipyInteractionParticle,
lon=[2],
lat=[2],
interaction_distance=1)
pyfunc_inter = pset2.InteractionKernel(ConstantMoveInteraction)
pset2.execute(AdvectionRK4,
pyfunc_inter=pyfunc_inter,
runtime=runtime,
dt=dt)
# Check to see whether they have moved as predicted.
assert np.all(pset.lon == pset2.lon)
assert np.all(pset2.lat == pset2.lon)
assert np.all(
pset2._collection.data["time"][0] == pset._collection.data["time"][0])
def test_avoid_repeated_grids():
lon_g0 = np.linspace(0, 1000, 11, dtype=np.float32)
lat_g0 = np.linspace(0, 1000, 11, dtype=np.float32)
time_g0 = np.linspace(0, 1000, 2, dtype=np.float64)
grid_0 = RectilinearZGrid(lon_g0, lat_g0, time=time_g0)
lon_g1 = np.linspace(0, 1000, 21, dtype=np.float32)
lat_g1 = np.linspace(0, 1000, 21, dtype=np.float32)
time_g1 = np.linspace(0, 1000, 2, dtype=np.float64)
grid_1 = RectilinearZGrid(lon_g1, lat_g1, time=time_g1)
u_data = np.zeros((lon_g0.size, lat_g0.size, time_g0.size),
dtype=np.float32)
u_field = Field('U', u_data, grid=grid_0, transpose=True)
v_data = np.zeros((lon_g1.size, lat_g1.size, time_g1.size),
dtype=np.float32)
v_field = Field('V', v_data, grid=grid_1, transpose=True)
temp0_field = Field('temp',
u_data,
lon=lon_g0,
lat=lat_g0,
time=time_g0,
transpose=True)
other_fields = {}
other_fields['temp'] = temp0_field
field_set = FieldSet(u_field, v_field, fields=other_fields)
assert field_set.gridset.size == 2
assert field_set.U.grid is field_set.temp.grid
assert field_set.V.grid is not field_set.U.grid
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, dt):
u = fieldset.U[time, particle.lon, particle.lat, particle.depth]
v = fieldset.V[time, particle.lon, particle.lat, particle.depth]
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], lat=[600])
pset.execute(pset.Kernel(sampleSpeed), runtime=0, dt=0)
assert (np.allclose(pset[0].speed, 1000))
def test_multiple_grid_addlater_error():
xdim, ydim = 10, 20
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))
fieldset = FieldSet(U, V)
pset = ParticleSet(fieldset, pclass=pclass('jit'), lon=[0.8], lat=[0.9])
P = Field('P',
np.zeros((ydim * 10, xdim * 10), dtype=np.float32),
lon=np.linspace(0., 1., xdim * 10, dtype=np.float32),
lat=np.linspace(0., 1., ydim * 10, dtype=np.float32))
fieldset.add_field(P)
fail = False
try:
pset.execute(AdvectionRK4, runtime=10, dt=1)
except:
fail = True
assert fail
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 test_pset_create_field_curvi(npart=100):
np.random.seed(123456)
r_v = np.linspace(.25, 2, 20)
theta_v = np.linspace(0, np.pi / 2, 200)
dtheta = theta_v[1] - theta_v[0]
dr = r_v[1] - r_v[0]
(r, theta) = np.meshgrid(r_v, theta_v)
x = -1 + r * np.cos(theta)
y = -1 + r * np.sin(theta)
grid = CurvilinearZGrid(x, y)
u = np.ones(x.shape)
v = np.where(np.logical_and(theta > np.pi / 4, theta < np.pi / 3), 1, 0)
ufield = Field('U', u, grid=grid)
vfield = Field('V', v, grid=grid)
fieldset = FieldSet(ufield, vfield)
pset = ParticleSet.from_field(fieldset,
size=npart,
pclass=ptype['scipy'],
start_field=fieldset.V)
lons = np.array([p.lon + 1 for p in pset])
lats = np.array([p.lat + 1 for p in pset])
thetas = np.arctan2(lats, lons)
rs = np.sqrt(lons * lons + lats * lats)
test = np.pi / 4 - dtheta < thetas
test *= thetas < np.pi / 3 + dtheta
test *= rs > .25 - dr
test *= rs < 2 + dr
assert np.all(test)
def test_nestedfields(mode, k_sample_p):
xdim = 10
ydim = 20
U1 = Field('U1',
0.1 * np.ones((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32))
V1 = Field('V1',
0.2 * np.ones((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32))
U2 = Field('U2',
0.3 * np.ones((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 2., xdim, dtype=np.float32),
lat=np.linspace(0., 2., ydim, dtype=np.float32))
V2 = Field('V2',
0.4 * np.ones((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 2., xdim, dtype=np.float32),
lat=np.linspace(0., 2., ydim, dtype=np.float32))
U = NestedField('U', [U1, U2])
V = NestedField('V', [V1, V2])
fieldset = FieldSet(U, V)
P1 = Field('P1',
0.1 * np.ones((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32))
P2 = Field('P2',
0.2 * np.ones((ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 2., xdim, dtype=np.float32),
lat=np.linspace(0., 2., ydim, dtype=np.float32))
P = NestedField('P', [P1, P2])
fieldset.add_field(P)
def Recover(particle, fieldset, time):
particle.lon = -1
particle.lat = -1
particle.p = 999
particle.time = particle.time + particle.dt
pset = ParticleSet(fieldset, pclass=pclass(mode), lon=[0], lat=[.3])
pset.execute(AdvectionRK4 + pset.Kernel(k_sample_p), runtime=1, dt=1)
assert np.isclose(pset.lat[0], .5)
assert np.isclose(pset.p[0], .1)
pset = ParticleSet(fieldset, pclass=pclass(mode), lon=[0], lat=[1.3])
pset.execute(AdvectionRK4 + pset.Kernel(k_sample_p), runtime=1, dt=1)
assert np.isclose(pset.lat[0], 1.7)
assert np.isclose(pset.p[0], .2)
pset = ParticleSet(fieldset, pclass=pclass(mode), lon=[0], lat=[2.3])
pset.execute(AdvectionRK4 + pset.Kernel(k_sample_p),
runtime=1,
dt=1,
recovery={ErrorCode.ErrorOutOfBounds: Recover})
assert np.isclose(pset.lat[0], -1)
assert np.isclose(pset.p[0], 999)
assert np.allclose(fieldset.UV[0][0, 0, 0, 0], [.1, .2])
def test_rectilinear_s_grids_advect2(mode):
# Move particle towards the east, check relative depth evolution
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_0 = RectilinearSGrid('grid0py', lon_g0, lat_g0, depth=depth_g0)
grid_1 = RectilinearSGrid('grid0py', lon_g0, lat_g0, depth=depth_g0)
grid_2 = 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)
rel_depth_data = np.zeros((lon_g0.size, lat_g0.size, depth_g0.shape[2]),
dtype=np.float32)
for k in range(1, depth_g0.shape[2]):
rel_depth_data[:, :, k] = k / (depth_g0.shape[2] - 1.)
u_field = Field('U', u_data, grid=grid_0, transpose=True)
v_field = Field('V', v_data, grid=grid_1, transpose=True)
rel_depth_field = Field('relDepth',
rel_depth_data,
grid=grid_2,
transpose=True)
field_set = FieldSet(u_field,
v_field,
fields={'relDepth': rel_depth_field})
class MyParticle(ptype[mode]):
relDepth = Variable('relDepth', dtype=np.float32, initial=20.)
def moveEast(particle, fieldset, time, dt):
particle.lon += 5 * dt
particle.relDepth = fieldset.relDepth[time, particle.lon, particle.lat,
particle.depth]
depth = .9
pset = ParticleSet.from_list(field_set,
MyParticle,
lon=[0],
lat=[0],
depth=[depth])
kernel = pset.Kernel(moveEast)
for _ in range(10):
pset.execute(kernel, starttime=pset[0].time, runtime=100, dt=50)
assert np.allclose(pset[0].relDepth, depth / bath_func(pset[0].lon))
def simulate_parcels(source_url,
output_filename,
lat,
lon,
wind_drift_factor=0.0,
velocity_average=True,
duration=23):
"""
source_url: local file or list of local files with fielddata in NetCDF-format
output_filename: name of file in which to save calculated trajectory
lat, lon: initial coordinates of single drifter or lists with coordinates for multiple drifters
wind_drift_factor: fraction of wind-speed at which objects will be advected. Default is 0 (no direct wind-drift)
velocity_average: Boolean variable deciding whether averaged horisontal velocities or surface velocities will be used.
Default is average which is consistent with GPU Ocean
duration: duration of the simulation in hours. Default is 24 hours.
TODO: Add functionality to start drifters at a later time in the simulation like in GPU Ocean.
"""
filenames = {'U': source_url, 'V': source_url}
dimensions = {'lat': 'lat', 'lon': 'lon', 'time': 'time'}
if velocity_average:
variables = {'U': 'ubar', 'V': 'vbar'}
else:
variables = {'U': 'u', 'V': 'v'}
fieldset = FieldSet.from_netcdf(filenames,
variables,
dimensions,
interp_method='cgrid_velocity')
if wind_drift_factor:
Uwind = Field.from_netcdf(source_url, ('U', 'Uwind'),
dimensions,
field_chunksize='auto',
interp_method='cgrid_velocity')
Vwind = Field.from_netcdf(source_url, ('V', 'Vwind'),
dimensions,
field_chunksize='auto',
interp_method='cgrid_velocity')
Uwind.set_scaling_factor(wind_drift_factor)
Vwind.set_scaling_factor(wind_drift_factor)
fieldset = FieldSet(U=fieldset.U + Uwind, V=fieldset.V + Vwind)
pset = ParticleSet.from_list(fieldset=fieldset,
pclass=JITParticle,
lon=lon,
lat=lat)
output_file = pset.ParticleFile(name=output_filename,
outputdt=timedelta(minutes=5))
pset.execute(AdvectionRK4,
runtime=timedelta(hours=duration),
dt=timedelta(minutes=5),
output_file=output_file)
output_file.export()
def test_rectilinear_s_grid_sampling(mode, z4d):
lon_g0 = np.linspace(-3e4, 3e4, 61, dtype=np.float32)
lat_g0 = np.linspace(0, 1000, 2, dtype=np.float32)
time_g0 = np.linspace(0, 1000, 2, dtype=np.float64)
if z4d:
depth_g0 = np.zeros((time_g0.size, 5, lat_g0.size, lon_g0.size), dtype=np.float32)
else:
depth_g0 = np.zeros((5, lat_g0.size, lon_g0.size), dtype=np.float32)
def bath_func(lon):
bath = (lon <= -2e4) * 20.
bath += (lon > -2e4) * (lon < 2e4) * (110. + 90 * np.sin(lon/2e4 * np.pi/2.))
bath += (lon >= 2e4) * 200.
return bath
bath = bath_func(lon_g0)
zdim = depth_g0.shape[-3]
for i in range(depth_g0.shape[-1]):
for k in range(zdim):
if z4d:
depth_g0[:, k, :, i] = bath[i] * k / (zdim-1)
else:
depth_g0[k, :, i] = bath[i] * k / (zdim-1)
grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0, time=time_g0)
u_data = np.zeros((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32)
v_data = np.zeros((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32)
temp_data = np.zeros((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32)
for k in range(1, zdim):
temp_data[:, k, :, :] = k / (zdim-1.)
u_field = Field('U', u_data, grid=grid)
v_field = Field('V', v_data, grid=grid)
temp_field = Field('temp', temp_data, grid=grid)
other_fields = {}
other_fields['temp'] = temp_field
field_set = FieldSet(u_field, v_field, fields=other_fields)
def sampleTemp(particle, fieldset, time):
particle.temp = fieldset.temp[time, particle.depth, particle.lat, particle.lon]
class MyParticle(ptype[mode]):
temp = Variable('temp', dtype=np.float32, initial=20.)
lon = 400
lat = 0
ratio = .3
pset = ParticleSet.from_list(field_set, MyParticle, lon=[lon], lat=[lat], depth=[bath_func(lon)*ratio])
pset.execute(pset.Kernel(sampleTemp), runtime=0, dt=0)
assert np.allclose(pset.temp[0], ratio, atol=1e-4)
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_0 = RectilinearSGrid('grid0py', lon_g0, lat_g0, depth=depth_g0)
grid_1 = RectilinearSGrid('grid0py', lon_g0, lat_g0, depth=depth_g0)
grid_2 = 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_0, transpose=True)
v_field = Field('V', v_data, grid=grid_1, transpose=True)
w_field = Field('W', w_data, grid=grid_2, 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,
starttime=pset[0].time,
runtime=10000,
dt=500)
assert np.allclose([p.depth / bath_func(p.lon) for p in pset], ratio)
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_rectilinear_s_grids_advect2(pset_mode, mode):
# Move particle towards the east, check relative depth evolution
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 = 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)
rel_depth_data = np.zeros((zdim, lat_g0.size, lon_g0.size),
dtype=np.float32)
for k in range(1, zdim):
rel_depth_data[k, :, :] = k / (zdim - 1.)
u_field = Field('U', u_data, grid=grid)
v_field = Field('V', v_data, grid=grid)
rel_depth_field = Field('relDepth', rel_depth_data, grid=grid)
field_set = FieldSet(u_field,
v_field,
fields={'relDepth': rel_depth_field})
class MyParticle(ptype[mode]):
relDepth = Variable('relDepth', dtype=np.float32, initial=20.)
def moveEast(particle, fieldset, time):
particle.lon += 5 * particle.dt
particle.relDepth = fieldset.relDepth[time, particle.depth,
particle.lat, particle.lon]
depth = .9
pset = pset_type[pset_mode]['pset'].from_list(field_set,
MyParticle,
lon=[0],
lat=[0],
depth=[depth])
kernel = pset.Kernel(moveEast)
for _ in range(10):
pset.execute(kernel, runtime=100, dt=50)
assert np.allclose(pset.relDepth[0], depth / bath_func(pset.lon[0]))
def set_fields():
datestr = '201[0-2]'
hycomfiles = sorted(
glob(
'/Volumes/data01/HYCOMdata/GLBu0.08_expt_19.1_surf/hycom_GLBu0.08_191_%s*'
% datestr))
dimensions = {'lat': 'lat', 'lon': 'lon', 'time': 'time'}
uhycom = Field.from_netcdf(hycomfiles,
'water_u',
dimensions,
fieldtype='U')
vhycom = Field.from_netcdf(hycomfiles,
'water_v',
dimensions,
fieldtype='V',
grid=uhycom.grid,
dataFiles=uhycom.dataFiles)
uhycom.vmin = -99.
uhycom.set_scaling_factor(0.001)
vhycom.set_scaling_factor(0.001)
vhycom.vmin = -99.
stokesfiles = sorted(
glob('/Volumes/data01/WaveWatch3data/WW3-GLOB-30M_%s*' % datestr))
dimensions = {'lat': 'latitude', 'lon': 'longitude', 'time': 'time'}
uuss = Field.from_netcdf(stokesfiles, 'uuss', dimensions, fieldtype='U')
vuss = Field.from_netcdf(stokesfiles,
'vuss',
dimensions,
fieldtype='V',
grid=uuss.grid,
dataFiles=uuss.dataFiles)
windfiles = sorted(
glob('/Users/erik/Desktop/WaveWatchWind/WaveWatchWind*.nc'))
dimensions = {'lat': 'latitude', 'lon': 'longitude', 'time': 'time'}
uwnd = Field.from_netcdf(windfiles, 'uwnd', dimensions, fieldtype='U')
vwnd = Field.from_netcdf(windfiles,
'vwnd',
dimensions,
fieldtype='V',
grid=uwnd.grid,
dataFiles=uwnd.dataFiles)
uwnd.set_scaling_factor(0.01)
vwnd.set_scaling_factor(0.01)
fieldset = FieldSet(U=uhycom + uuss + uwnd, V=vhycom + vuss + vwnd)
fieldset.add_periodic_halo(zonal=True, meridional=False, halosize=5)
return fieldset
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)
depth_g0 = depth_g0.transpose(
) # we don't change it on purpose, to check if the transpose op if fixed in jit
grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0)
zdim = depth_g0.shape[0]
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)
for i in range(lon_g0.size):
u_data[:, :, i] = 1 * 10 / bath[i]
for k in range(zdim):
w_data[k, :,
i] = u_data[k, :, i] * depth_g0[k, :, i] / bath[i] * 1e-3
u_field = Field('U', u_data, grid=grid)
v_field = Field('V', v_data, grid=grid)
w_field = Field('W', w_data, grid=grid)
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 test_sampling_multiple_grid_sizes(pset_mode, mode, ugridfactor):
xdim, ydim = 10, 20
U = Field('U', np.zeros((ydim*ugridfactor, xdim*ugridfactor), dtype=np.float32),
lon=np.linspace(0., 1., xdim*ugridfactor, dtype=np.float32),
lat=np.linspace(0., 1., ydim*ugridfactor, 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))
fieldset = FieldSet(U, V)
pset = pset_type[pset_mode]['pset'](fieldset, pclass=pclass(mode), lon=[0.8], lat=[0.9])
if ugridfactor > 1:
assert fieldset.U.grid is not fieldset.V.grid
else:
assert fieldset.U.grid is fieldset.V.grid
pset.execute(AdvectionRK4, runtime=10, dt=1)
assert np.isclose(pset.lon[0], 0.8)
assert np.all((0 <= pset.xi) & (pset.xi < xdim*ugridfactor))
def test_summedfields(mode, with_W, k_sample_p, mesh):
xdim = 10
ydim = 20
zdim = 4
gf = 10 # factor by which the resolution of grid1 is higher than of grid2
U1 = Field('U', 0.2*np.ones((zdim*gf, ydim*gf, xdim*gf), dtype=np.float32),
lon=np.linspace(0., 1., xdim*gf, dtype=np.float32),
lat=np.linspace(0., 1., ydim*gf, dtype=np.float32),
depth=np.linspace(0., 20., zdim*gf, dtype=np.float32),
mesh=mesh)
U2 = Field('U', 0.1*np.ones((zdim, ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32),
depth=np.linspace(0., 20., zdim, dtype=np.float32),
mesh=mesh)
V1 = Field('V', np.zeros((zdim*gf, ydim*gf, xdim*gf), dtype=np.float32), grid=U1.grid, fieldtype='V')
V2 = Field('V', np.zeros((zdim, ydim, xdim), dtype=np.float32), grid=U2.grid, fieldtype='V')
fieldsetS = FieldSet(U1+U2, V1+V2)
conv = 1852*60 if mesh == 'spherical' else 1.
assert np.allclose(fieldsetS.U[0, 0, 0, 0]*conv, 0.3)
P1 = Field('P', 30*np.ones((zdim*gf, ydim*gf, xdim*gf), dtype=np.float32), grid=U1.grid)
P2 = Field('P', 20*np.ones((zdim, ydim, xdim), dtype=np.float32), grid=U2.grid)
P3 = Field('P', 10*np.ones((zdim, ydim, xdim), dtype=np.float32), grid=U2.grid)
P4 = Field('P', 0*np.ones((zdim, ydim, xdim), dtype=np.float32), grid=U2.grid)
fieldsetS.add_field((P1+P4)+(P2+P3), name='P')
assert np.allclose(fieldsetS.P[0, 0, 0, 0], 60)
if with_W:
W1 = Field('W', 2*np.ones((zdim * gf, ydim * gf, xdim * gf), dtype=np.float32), grid=U1.grid)
W2 = Field('W', np.ones((zdim, ydim, xdim), dtype=np.float32), grid=U2.grid)
fieldsetS.add_field(W1+W2, name='W')
pset = ParticleSet(fieldsetS, pclass=pclass(mode), lon=[0], lat=[0.9])
pset.execute(AdvectionRK4_3D+pset.Kernel(k_sample_p), runtime=2, dt=1)
assert np.isclose(pset[0].depth, 6)
else:
pset = ParticleSet(fieldsetS, pclass=pclass(mode), lon=[0], lat=[0.9])
pset.execute(AdvectionRK4+pset.Kernel(k_sample_p), runtime=2, dt=1)
assert np.isclose(pset[0].p, 60)
assert np.isclose(pset[0].lon*conv, 0.6, atol=1e-3)
assert np.isclose(pset[0].lat, 0.9)
assert np.allclose(fieldsetS.UV[0][0, 0, 0, 0], [.2/conv, 0])
def test_summedfields_slipinterp_warning(boundaryslip):
xdim = 10
ydim = 20
zdim = 4
gf = 10 # factor by which the resolution of grid1 is higher than of grid2
U1 = Field('U', 0.2*np.ones((zdim*gf, ydim*gf, xdim*gf), dtype=np.float32),
lon=np.linspace(0., 1., xdim*gf, dtype=np.float32),
lat=np.linspace(0., 1., ydim*gf, dtype=np.float32),
depth=np.linspace(0., 20., zdim*gf, dtype=np.float32),
interp_method=boundaryslip)
U2 = Field('U', 0.1*np.ones((zdim, ydim, xdim), dtype=np.float32),
lon=np.linspace(0., 1., xdim, dtype=np.float32),
lat=np.linspace(0., 1., ydim, dtype=np.float32),
depth=np.linspace(0., 20., zdim, dtype=np.float32))
V1 = Field('V', np.zeros((zdim*gf, ydim*gf, xdim*gf), dtype=np.float32), grid=U1.grid, fieldtype='V')
V2 = Field('V', np.zeros((zdim, ydim, xdim), dtype=np.float32), grid=U2.grid, fieldtype='V')
fieldsetS = FieldSet(U1+U2, V1+V2)
with pytest.warns(UserWarning):
fieldsetS.check_complete()
def test_sampling_multiple_grid_sizes(mode):
"""Sampling test that tests for FieldSet with different grid sizes
While this currently works fine in Scipy mode, it fails in JIT mode with
an out_of_bounds_error because there is only one (xi, yi, zi) for each particle
A solution would be to define xi, yi, zi for each field separately
"""
xdim = 10
ydim = 20
gf = 10 # factor by which the resolution of U is higher than of V
U = Field('U', np.zeros((ydim*gf, xdim*gf), dtype=np.float32),
lon=np.linspace(0., 1., xdim*gf, dtype=np.float32),
lat=np.linspace(0., 1., ydim*gf, 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))
fieldset = FieldSet(U, V)
pset = ParticleSet(fieldset, pclass=pclass(mode), lon=[0.8], lat=[0.9])
pset.execute(AdvectionRK4, runtime=10, dt=1)
assert np.isclose(pset[0].lon, 0.8)
def set_fields(hycomfiles, stokesfiles):
dimensions = {
'lat': 'Latitude',
'lon': 'Longitude',
'time': 'MT',
'depth': 'Depth'
}
bvelfile = '/Users/erik/Codes/ParcelsRuns/Hydrodynamics_AuxiliaryFiles/Hycom/boundary_velocities.nc'
MaskUvel = Field.from_netcdf(bvelfile, 'MaskUvel', dimensions)
MaskVvel = Field.from_netcdf(bvelfile, 'MaskVvel', dimensions)
uhycom = Field.from_netcdf(hycomfiles, 'u', dimensions, fieldtype='U')
vhycom = Field.from_netcdf(hycomfiles,
'v',
dimensions,
fieldtype='V',
grid=uhycom.grid,
timeFiles=uhycom.timeFiles)
dimensions = {'lat': 'latitude', 'lon': 'longitude', 'time': 'time'}
uuss = Field.from_netcdf(stokesfiles, 'uuss', dimensions, fieldtype='U')
vuss = Field.from_netcdf(stokesfiles,
'vuss',
dimensions,
fieldtype='V',
grid=uuss.grid,
timeFiles=uuss.timeFiles)
fieldset = FieldSet(U=uhycom + uuss, V=vhycom + vuss)
fieldset.add_field(MaskUvel)
fieldset.add_field(MaskVvel)
fieldset.MaskUvel.units = fieldset.U[0].units
fieldset.MaskVvel.units = fieldset.V[0].units
fieldset.add_periodic_halo(zonal=True, meridional=False, halosize=5)
return fieldset
lon = dfile.variables['longitude'][:]
lat = dfile.variables['latitude'][:]
iy_min, ix_min = getclosest_ij(lat, lon, NS_domain[2], NS_domain[0])
iy_max, ix_max = getclosest_ij(lat, lon, NS_domain[3], NS_domain[1])
### add fields
files = sorted(glob(data_in + "/SMOC_2020*.nc"))
variables = {'U': 'uo', 'V': 'vo'}
dimensions = {'lon': 'longitude', 'lat': 'latitude', 'time': 'time'}
indices = {'lon': range(ix_min, ix_max), 'lat': range(iy_min, iy_max)}
fset_currents = FieldSet.from_netcdf(files,
variables,
dimensions,
indices=indices)
fieldset_currents = FieldSet(U=fset_currents.U, V=fset_currents.V)
if withstokes:
stokesfiles = sorted(glob(data_in + "/SMOC_2020*.nc"))
stokesdimensions = {'lon': 'longitude', 'lat': 'latitude', 'time': 'time'}
stokesvariables = {'U': 'vsdx', 'V': 'vsdy'}
fset_stokes = FieldSet.from_netcdf(stokesfiles,
stokesvariables,
stokesdimensions,
indices=indices)
fname += '_stokes'
if withtides:
tidesfiles = sorted(glob(data_in + "/SMOC_2020*.nc"))
tidesdimensions = {'lon': 'longitude', 'lat': 'latitude', 'time': 'time'}
tidesvariables = {'U': 'utide', 'V': 'vtide'}
#Stokes Field
if wstokes:
files_stokes = sorted(glob(data_in + "/WaveWatch3data/CFSR/WW3-GLOB-30M_200[8-9]*_uss.nc"))
variables_stokes = {'U': 'uuss',
'V': 'vuss'}
dimensions_stokes = {'U': {'lon': 'longitude', 'lat': 'latitude', 'time': 'time'},
'V': {'lon': 'longitude', 'lat': 'latitude', 'time': 'time'}}
indices_stokes = {'lon': range(120, 220), 'lat': range(142, 170)}
fieldset_stokes = FieldSet.from_netcdf(files_stokes,
variables_stokes,
dimensions_stokes,
indices=indices_stokes)
fieldset_stokes.add_periodic_halo(zonal=True, meridional=False, halosize=5)
fieldset = FieldSet(U=fieldset_nemo.U + fieldset_stokes.U,
V=fieldset_nemo.V + fieldset_stokes.V)
fU = fieldset.U[0]
fname = path.join(data_out, filename_out + "_wstokes.nc")
else:
fieldset = fieldset_nemo
fU = fieldset.U
fname = path.join(data_out, filename_out + ".nc")
#initialize where to start particles
GC = [-90.5,-0.5] #centre of the Galapagos islands
release_extent = [GC[0]-3.5, GC[0]+3.5, GC[1]-3.5, GC[1]+3.5]
width = 7 #thicknes (in number of particles) of release square around Galapagos
templon, templat = np.meshgrid(np.arange(release_extent[0],
release_extent[1],0.2),
},
}
fieldset = FieldSet.from_netcdf(filenames,
variables,
dimensions,
allow_time_extrapolation=True)
###############################################################################
# Adding the Stokes drift to the HYCOM currents #
###############################################################################
if stokes == 0:
fieldset.Ust.units = GeographicPolar()
fieldset.Vst.units = Geographic()
fieldset = FieldSet(
U=fieldset.U + fieldset.Ust,
V=fieldset.V + fieldset.Vst,
)
###############################################################################
# Adding the border current, which applies for all scenarios except for 0 #
###############################################################################
datasetBor = Dataset(dataInputdirec + 'boundary_velocities_HYCOM.nc')
borU = datasetBor.variables['MaskUvel'][:]
borU[borU != 0] = borU[borU != 0] / abs(borU[borU != 0])
borV = datasetBor.variables['MaskVvel'][:]
borV[borV != 0] = borV[borV != 0] / abs(borV[borV != 0])
borMag = np.sqrt(np.square(borU) + np.square(borV))
nonzeroMag = borMag > 0
borMag[borMag == 0] = 1
borU = np.divide(borU, borMag)
borV = np.divide(borV, borMag)
if withwind:
windfiles = sorted(
glob('WIND_GLO_WIND_L4_NRT_OBSERVATIONS_012_004/2019/*.nc'))
winddimensions = {'lat': 'lat', 'lon': 'lon', 'time': 'time'}
windvariables = {'U': 'eastward_wind', 'V': 'northward_wind'}
fset_wind = FieldSet.from_netcdf(windfiles, windvariables, winddimensions)
fset_wind.add_periodic_halo(zonal=True)
fset_wind.U.set_scaling_factor(withwind)
fset_wind.V.set_scaling_factor(withwind)
fname += '_wind%.4d' % (withwind * 1000)
fname += '.nc'
if withstokes and withwind:
fieldset = FieldSet(U=fset_currents.U + fset_stokes.U + fset_wind.U,
V=fset_currents.V + fset_stokes.V + fset_wind.V)
elif withstokes:
fieldset = FieldSet(U=fset_currents.U + fset_stokes.U,
V=fset_currents.V + fset_stokes.V)
elif withwind:
fieldset = FieldSet(U=fset_currents.U + fset_wind.U,
V=fset_currents.V + fset_wind.V)
else:
fieldset = FieldSet(U=fset_currents.U, V=fset_currents.V)
if withwind: # for sampling
for var, name in zip((fset_wind.U, fset_wind.V), ('uwind', 'vwind')):
fld = copy.deepcopy(var)
fld.name = name
fld.units = UnitConverter() # to sample in m/s
fieldset.add_field(fld)
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)
'time': 'time'
},
'V': {
'lon': 'longitude',
'lat': 'latitude',
'time': 'time'
}
}
indices_stokes = {'lon': range(120, 220), 'lat': range(142, 170)}
fieldset_stokes = FieldSet.from_netcdf(files_stokes,
variables_stokes,
dimensions_stokes,
indices=indices_stokes)
fieldset_stokes.add_periodic_halo(zonal=True, meridional=False, halosize=5)
fieldset = FieldSet(U=fieldset_MITgcm.U + fieldset_stokes.U,
V=fieldset_MITgcm.V + fieldset_stokes.V)
fname = path.join(data_out, filename_out + "_wstokes.nc")
else:
fieldset = fieldset_MITgcm
fname = path.join(data_out, filename_out + ".nc")
fU = fieldset_MITgcm.U
# get all lon, lat that are land
fieldset_MITgcm.computeTimeChunk(fU.grid.time[0], 1)
lon = np.array(fU.lon[:])
lat = np.array(fU.lat[:])
LandMask = fU.data[0, :, :]
LandMask = np.array(LandMask)
land = np.where(LandMask == 0)
'V': {
'lon': 'longitude',
'lat': 'latitude',
'time': 'time'
}
}
interp_method = {'U': 'linear', 'V': 'linear'}
fieldset_wave = FieldSet.from_netcdf(filenames,
variables,
dimensions,
interp_method=interp_method)
if param['windage']:
fieldset = FieldSet(U=fieldset_ocean.U + fieldset_windwave.U,
V=fieldset_ocean.V + fieldset_windwave.V)
elif param['stokes']:
fieldset = FieldSet(U=fieldset_ocean.U + fieldset_wave.U,
V=fieldset_ocean.V + fieldset_wave.V)
else:
fieldset = fieldset_ocean
# ADD ADDITIONAL FIELDS
# Country identifier grid (on psi grid, nearest)
iso_psi_all = Field.from_netcdf(fh['grid'],
variable='iso_psi_all',
dimensions={
'lon': 'lon_psi',
'lat': 'lat_psi'
},
interp_method='nearest',
def test_multi_structured_grids(pset_mode, mode):
def temp_func(lon, lat):
return 20 + lat / 1000. + 2 * np.sin(lon * 2 * np.pi / 5000.)
a = 10000
b = 10000
# Grid 0
xdim_g0 = 201
ydim_g0 = 201
# Coordinates of the test fieldset (on A-grid in deg)
lon_g0 = np.linspace(0, a, xdim_g0, dtype=np.float32)
lat_g0 = np.linspace(0, b, ydim_g0, dtype=np.float32)
time_g0 = np.linspace(0., 1000., 2, dtype=np.float64)
grid_0 = RectilinearZGrid(lon_g0, lat_g0, time=time_g0)
# Grid 1
xdim_g1 = 51
ydim_g1 = 51
# Coordinates of the test fieldset (on A-grid in deg)
lon_g1 = np.linspace(0, a, xdim_g1, dtype=np.float32)
lat_g1 = np.linspace(0, b, ydim_g1, dtype=np.float32)
time_g1 = np.linspace(0., 1000., 2, dtype=np.float64)
grid_1 = RectilinearZGrid(lon_g1, lat_g1, time=time_g1)
u_data = np.ones((lon_g0.size, lat_g0.size, time_g0.size),
dtype=np.float32)
u_data = 2 * u_data
u_field = Field('U', u_data, grid=grid_0, transpose=True)
temp0_data = np.empty((lon_g0.size, lat_g0.size, time_g0.size),
dtype=np.float32)
for i in range(lon_g0.size):
for j in range(lat_g0.size):
temp0_data[i, j, :] = temp_func(lon_g0[i], lat_g0[j])
temp0_field = Field('temp0', temp0_data, grid=grid_0, transpose=True)
v_data = np.zeros((lon_g1.size, lat_g1.size, time_g1.size),
dtype=np.float32)
v_field = Field('V', v_data, grid=grid_1, transpose=True)
temp1_data = np.empty((lon_g1.size, lat_g1.size, time_g1.size),
dtype=np.float32)
for i in range(lon_g1.size):
for j in range(lat_g1.size):
temp1_data[i, j, :] = temp_func(lon_g1[i], lat_g1[j])
temp1_field = Field('temp1', temp1_data, grid=grid_1, transpose=True)
other_fields = {}
other_fields['temp0'] = temp0_field
other_fields['temp1'] = temp1_field
field_set = FieldSet(u_field, v_field, fields=other_fields)
def sampleTemp(particle, fieldset, time):
# Note that fieldset.temp is interpolated at time=time+dt.
# Indeed, sampleTemp is called at time=time, but the result is written
# at time=time+dt, after the Kernel update
particle.temp0 = fieldset.temp0[time + particle.dt, particle.depth,
particle.lat, particle.lon]
particle.temp1 = fieldset.temp1[time + particle.dt, particle.depth,
particle.lat, particle.lon]
class MyParticle(ptype[mode]):
temp0 = Variable('temp0', dtype=np.float32, initial=20.)
temp1 = Variable('temp1', dtype=np.float32, initial=20.)
pset = pset_type[pset_mode]['pset'].from_list(field_set,
MyParticle,
lon=[3001],
lat=[5001],
repeatdt=1)
pset.execute(AdvectionRK4 + pset.Kernel(sampleTemp), runtime=3, dt=1)
# check if particle xi and yi are different for the two grids
# assert np.all([pset.xi[i, 0] != pset.xi[i, 1] for i in range(3)])
# assert np.all([pset.yi[i, 0] != pset.yi[i, 1] for i in range(3)])
assert np.alltrue([pset[i].xi[0] != pset[i].xi[1] for i in range(3)])
assert np.alltrue([pset[i].yi[0] != pset[i].yi[1] for i in range(3)])
# advect without updating temperature to test particle deletion
pset.remove_indices(np.array([1]))
pset.execute(AdvectionRK4, runtime=1, dt=1)
assert np.alltrue([np.isclose(p.temp0, p.temp1, atol=1e-3) for p in pset])
def test_multi_structured_grids(mode):
def temp_func(lon, lat):
return 20 + lat / 1000. + 2 * np.sin(lon * 2 * np.pi / 5000.)
a = 10000
b = 10000
# Grid 0
xdim_g0 = 201
ydim_g0 = 201
# Coordinates of the test fieldset (on A-grid in deg)
lon_g0 = np.linspace(0, a, xdim_g0, dtype=np.float32)
lat_g0 = np.linspace(0, b, ydim_g0, dtype=np.float32)
time_g0 = np.linspace(0., 1000., 2, dtype=np.float64)
grid_0 = RectilinearZGrid(lon_g0, lat_g0, time=time_g0)
# Grid 1
xdim_g1 = 51
ydim_g1 = 51
# Coordinates of the test fieldset (on A-grid in deg)
lon_g1 = np.linspace(0, a, xdim_g1, dtype=np.float32)
lat_g1 = np.linspace(0, b, ydim_g1, dtype=np.float32)
time_g1 = np.linspace(0., 1000., 2, dtype=np.float64)
grid_1 = RectilinearZGrid(lon_g1, lat_g1, time=time_g1)
u_data = np.ones((lon_g0.size, lat_g0.size, time_g0.size),
dtype=np.float32)
u_data = 2 * u_data
u_field = Field('U', u_data, grid=grid_0, transpose=True)
temp0_data = np.empty((lon_g0.size, lat_g0.size, time_g0.size),
dtype=np.float32)
for i in range(lon_g0.size):
for j in range(lat_g0.size):
temp0_data[i, j, :] = temp_func(lon_g0[i], lat_g0[j])
temp0_field = Field('temp0', temp0_data, grid=grid_0, transpose=True)
v_data = np.zeros((lon_g1.size, lat_g1.size, time_g1.size),
dtype=np.float32)
v_field = Field('V', v_data, grid=grid_1, transpose=True)
temp1_data = np.empty((lon_g1.size, lat_g1.size, time_g1.size),
dtype=np.float32)
for i in range(lon_g1.size):
for j in range(lat_g1.size):
temp1_data[i, j, :] = temp_func(lon_g1[i], lat_g1[j])
temp1_field = Field('temp1', temp1_data, grid=grid_1, transpose=True)
other_fields = {}
other_fields['temp0'] = temp0_field
other_fields['temp1'] = temp1_field
field_set = FieldSet(u_field, v_field, fields=other_fields)
def sampleTemp(particle, fieldset, time, dt):
# Note that fieldset.temp is interpolated at time=time+dt.
# Indeed, sampleTemp is called at time=time, but the result is written
# at time=time+dt, after the Kernel update
particle.temp0 = fieldset.temp0[time + dt, particle.lon, particle.lat,
particle.depth]
particle.temp1 = fieldset.temp1[time + dt, particle.lon, particle.lat,
particle.depth]
class MyParticle(ptype[mode]):
temp0 = Variable('temp0', dtype=np.float32, initial=20.)
temp1 = Variable('temp1', dtype=np.float32, initial=20.)
pset = ParticleSet.from_list(field_set, MyParticle, lon=[3001], lat=[5001])
pset.execute(AdvectionRK4 + pset.Kernel(sampleTemp), runtime=1, dt=1)
assert np.allclose(pset.particles[0].temp0,
pset.particles[0].temp1,
atol=1e-3)
