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_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 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_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_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 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)