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