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)