From 4fd6c9d8d94ccb5c466ea805350be4e31eabcffb Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 24 Apr 2024 10:10:51 +0200 Subject: [PATCH 1/7] Improving RK45 kernel by allowing users to change tolerance and timesteps --- parcels/application_kernels/advection.py | 10 +++++----- parcels/kernel.py | 16 +++++++++++++++- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/parcels/application_kernels/advection.py b/parcels/application_kernels/advection.py index ea76311457..be1801b23e 100644 --- a/parcels/application_kernels/advection.py +++ b/parcels/application_kernels/advection.py @@ -64,10 +64,10 @@ def AdvectionRK45(particle, fieldset, time): 1e-5 * dt by default. Note that this kernel requires a Particle Class that has an extra Variable 'next_dt' + and a FieldSet with 'RK45_tol', 'RK45_min_dt' and 'RK45_max_dt' constants. """ + particle.next_dt = min(particle.next_dt, fieldset.RK45_max_dt) particle.dt = particle.next_dt - rk45tol = 1e-5 - min_dt = 1e-3 c = [1./4., 3./8., 12./13., 1., 1./2.] A = [[1./4., 0., 0., 0., 0.], [3./32., 9./32., 0., 0., 0.], @@ -99,11 +99,11 @@ def AdvectionRK45(particle, fieldset, time): lon_5th = (u1 * b5[0] + u2 * b5[1] + u3 * b5[2] + u4 * b5[3] + u5 * b5[4] + u6 * b5[5]) * particle.dt lat_5th = (v1 * b5[0] + v2 * b5[1] + v3 * b5[2] + v4 * b5[3] + v5 * b5[4] + v6 * b5[5]) * particle.dt - kappa2 = math.pow(lon_5th - lon_4th, 2) + math.pow(lat_5th - lat_4th, 2) - if kappa2 <= math.pow(math.fabs(particle.dt * rk45tol), 2) or particle.dt < min_dt: + kappa = math.sqrt(math.pow(lon_5th - lon_4th, 2) + math.pow(lat_5th - lat_4th, 2)) + if kappa <= fieldset.RK45_tol or particle.dt < fieldset.RK45_min_dt: particle_dlon += lon_4th # noqa particle_dlat += lat_4th # noqa - if kappa2 <= math.pow(math.fabs(particle.dt * rk45tol / 10), 2): + if kappa <= fieldset.RK45_tol / 10 and particle.dt*2 <= fieldset.RK45_max_dt: particle.next_dt *= 2 else: particle.next_dt /= 2 diff --git a/parcels/kernel.py b/parcels/kernel.py index 38b8dec6e5..b9d6b9edcf 100644 --- a/parcels/kernel.py +++ b/parcels/kernel.py @@ -23,7 +23,11 @@ MPI = None import parcels.rng as ParcelsRandom # noqa -from parcels.application_kernels.advection import AdvectionAnalytical, AdvectionRK4_3D +from parcels.application_kernels.advection import ( + AdvectionAnalytical, + AdvectionRK4_3D, + AdvectionRK45, +) from parcels.compilation.codegenerator import KernelGenerator, LoopGenerator from parcels.field import Field, NestedField, VectorField from parcels.grid import GridCode @@ -321,6 +325,16 @@ def check_fieldsets_in_kernels(self, pyfunc): raise NotImplementedError('Analytical Advection only works with C-grids') if self._fieldset.U.grid.gtype not in [GridCode.CurvilinearZGrid, GridCode.RectilinearZGrid]: raise NotImplementedError('Analytical Advection only works with Z-grids in the vertical') + elif pyfunc is AdvectionRK45: + if not hasattr(self.fieldset, 'RK45_tol'): + logger.info("Setting RK45 tolerance to 10 m. Use fieldset.add_constant('RK45_tol', [distance]) to change.") + self.fieldset.add_constant('RK45_tol', 10) + if not hasattr(self.fieldset, 'RK45_min_dt'): + logger.info("Setting RK45 minimum timestep to 1 s. Use fieldset.add_constant('RK45_min_dt', [timestep]) to change.") + self.fieldset.add_constant('RK45_min_dt', 1) + if not hasattr(self.fieldset, 'RK45_max_dt'): + logger.info("Setting RK45 maximum timestep to 1000 s. Use fieldset.add_constant('RK45_max_dt', [timestep]) to change.") + self.fieldset.add_constant('RK45_max_dt', 1000) def check_kernel_signature_on_version(self): numkernelargs = 0 From 9ff2155dbcb88f5014b085ce00a359f18c511dfc Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 24 Apr 2024 10:37:46 +0200 Subject: [PATCH 2/7] Fixing RK45 timestepping in Scipy mode --- parcels/application_kernels/advection.py | 3 +-- parcels/kernel.py | 4 +++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/parcels/application_kernels/advection.py b/parcels/application_kernels/advection.py index be1801b23e..07db066875 100644 --- a/parcels/application_kernels/advection.py +++ b/parcels/application_kernels/advection.py @@ -66,8 +66,7 @@ def AdvectionRK45(particle, fieldset, time): Note that this kernel requires a Particle Class that has an extra Variable 'next_dt' and a FieldSet with 'RK45_tol', 'RK45_min_dt' and 'RK45_max_dt' constants. """ - particle.next_dt = min(particle.next_dt, fieldset.RK45_max_dt) - particle.dt = particle.next_dt + particle.dt = min(particle.next_dt, fieldset.RK45_max_dt) c = [1./4., 3./8., 12./13., 1., 1./2.] A = [[1./4., 0., 0., 0., 0.], [3./32., 9./32., 0., 0., 0.], diff --git a/parcels/kernel.py b/parcels/kernel.py index b9d6b9edcf..8129c259c2 100644 --- a/parcels/kernel.py +++ b/parcels/kernel.py @@ -637,7 +637,9 @@ def evaluate_particle(self, p, endtime): if sign_dt*p.time_nextloop >= sign_dt*endtime: return p - if abs(endtime - p.time_nextloop) < abs(p.dt)-1e-6: + if hasattr(p, 'next_dt') and abs(endtime - p.time_nextloop) < abs(p.next_dt)-1e-6: + p.next_dt = abs(endtime - p.time_nextloop) * sign_dt + elif abs(endtime - p.time_nextloop) < abs(p.dt)-1e-6: p.dt = abs(endtime - p.time_nextloop) * sign_dt res = self._pyfunc(p, self._fieldset, p.time_nextloop) From 59517a6e0d1fd3496e56123be1199211a4eee57c Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 24 Apr 2024 15:11:30 +0200 Subject: [PATCH 3/7] Fixing AdvectionRK45 for negative dt (by using math.fabs) --- parcels/application_kernels/advection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/parcels/application_kernels/advection.py b/parcels/application_kernels/advection.py index 07db066875..511e191567 100644 --- a/parcels/application_kernels/advection.py +++ b/parcels/application_kernels/advection.py @@ -99,10 +99,10 @@ def AdvectionRK45(particle, fieldset, time): lat_5th = (v1 * b5[0] + v2 * b5[1] + v3 * b5[2] + v4 * b5[3] + v5 * b5[4] + v6 * b5[5]) * particle.dt kappa = math.sqrt(math.pow(lon_5th - lon_4th, 2) + math.pow(lat_5th - lat_4th, 2)) - if kappa <= fieldset.RK45_tol or particle.dt < fieldset.RK45_min_dt: + if (kappa <= fieldset.RK45_tol) or (math.fabs(particle.dt) < math.fabs(fieldset.RK45_min_dt)): particle_dlon += lon_4th # noqa particle_dlat += lat_4th # noqa - if kappa <= fieldset.RK45_tol / 10 and particle.dt*2 <= fieldset.RK45_max_dt: + if (kappa <= fieldset.RK45_tol) / 10 and (math.fabs(particle.dt*2) <= math.fabs(fieldset.RK45_max_dt)): particle.next_dt *= 2 else: particle.next_dt /= 2 From eb37afd82040b02af8e266390d2d8b275d2a6e1a Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 25 Apr 2024 08:46:22 +0200 Subject: [PATCH 4/7] Fixing JIT version of AdvectionRK45 --- parcels/compilation/codegenerator.py | 7 +++++-- parcels/kernel.py | 10 ++++++---- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/parcels/compilation/codegenerator.py b/parcels/compilation/codegenerator.py index 45d5e2d8b9..19c602079d 100644 --- a/parcels/compilation/codegenerator.py +++ b/parcels/compilation/codegenerator.py @@ -925,13 +925,16 @@ def generate(self, funcname, field_args, const_args, kernel_ast, c_include): # ==== statement clusters use to compose 'body' variable and variables 'time_loop' and 'part_loop' ==== ## sign_dt = c.Assign("sign_dt", "dt > 0 ? 1 : -1") + # ==== check if next_dt is in the particle type ==== # + dtname = "next_dt" if "next_dt" in [v.name for v in self.ptype.variables] else "dt" + # ==== main computation body ==== # body = [] body += [c.Value("double", "pre_dt")] body += [c.Statement("pre_dt = particles->dt[pnum]")] body += [c.If("sign_dt*particles->time_nextloop[pnum] >= sign_dt*(endtime)", c.Statement("break"))] - body += [c.If("fabs(endtime - particles->time_nextloop[pnum]) < fabs(particles->dt[pnum])-1e-6", - c.Statement("particles->dt[pnum] = fabs(endtime - particles->time_nextloop[pnum]) * sign_dt"))] + body += [c.If(f"fabs(endtime - particles->time_nextloop[pnum]) < fabs(particles->{dtname}[pnum])-1e-6", + c.Statement(f"particles->{dtname}[pnum] = fabs(endtime - particles->time_nextloop[pnum]) * sign_dt"))] body += [c.Assign("particles->state[pnum]", f"{funcname}(particles, pnum, {fargs_str})")] body += [c.If("particles->state[pnum] == SUCCESS", c.Block([c.If("sign_dt*particles->time[pnum] < sign_dt*endtime", diff --git a/parcels/kernel.py b/parcels/kernel.py index 8129c259c2..cc28ba83ec 100644 --- a/parcels/kernel.py +++ b/parcels/kernel.py @@ -637,10 +637,12 @@ def evaluate_particle(self, p, endtime): if sign_dt*p.time_nextloop >= sign_dt*endtime: return p - if hasattr(p, 'next_dt') and abs(endtime - p.time_nextloop) < abs(p.next_dt)-1e-6: - p.next_dt = abs(endtime - p.time_nextloop) * sign_dt - elif abs(endtime - p.time_nextloop) < abs(p.dt)-1e-6: - p.dt = abs(endtime - p.time_nextloop) * sign_dt + try: # Use next_dt from AdvectionRK45 if it is set + if abs(endtime - p.time_nextloop) < abs(p.next_dt)-1e-6: + p.next_dt = abs(endtime - p.time_nextloop) * sign_dt + except KeyError: + if abs(endtime - p.time_nextloop) < abs(p.dt)-1e-6: + p.dt = abs(endtime - p.time_nextloop) * sign_dt res = self._pyfunc(p, self._fieldset, p.time_nextloop) if res is None: From eb2f2ff76084b7c7e84e801846cb7bf08b81b7ff Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 25 Apr 2024 09:03:36 +0200 Subject: [PATCH 5/7] Simple RK45_tol conversion from meters to degrees --- parcels/kernel.py | 2 ++ tests/test_advection.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/parcels/kernel.py b/parcels/kernel.py index cc28ba83ec..4d883a89e4 100644 --- a/parcels/kernel.py +++ b/parcels/kernel.py @@ -329,6 +329,8 @@ def check_fieldsets_in_kernels(self, pyfunc): if not hasattr(self.fieldset, 'RK45_tol'): logger.info("Setting RK45 tolerance to 10 m. Use fieldset.add_constant('RK45_tol', [distance]) to change.") self.fieldset.add_constant('RK45_tol', 10) + if self.fieldset.U.grid.mesh == 'spherical': + self.fieldset.RK45_tol /= (1852 * 60) # TODO does not account for zonal variation in meter -> degree conversion if not hasattr(self.fieldset, 'RK45_min_dt'): logger.info("Setting RK45 minimum timestep to 1 s. Use fieldset.add_constant('RK45_min_dt', [timestep]) to change.") self.fieldset.add_constant('RK45_min_dt', 1) diff --git a/tests/test_advection.py b/tests/test_advection.py index 7e11338452..e1ab2f0ee4 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -166,6 +166,26 @@ def SubmergeParticle(particle, fieldset, time): assert len(pset) == 0 +@pytest.mark.parametrize('mode', ['scipy', 'jit']) +@pytest.mark.parametrize('rk45_tol', [10, 100]) +def test_advection_RK45(lon, lat, mode, rk45_tol, npart=10): + data2D = {'U': np.ones((lon.size, lat.size), dtype=np.float32), + 'V': np.zeros((lon.size, lat.size), dtype=np.float32)} + dimensions = {'lon': lon, 'lat': lat} + fieldset = FieldSet.from_data(data2D, dimensions, mesh='spherical', transpose=True) + fieldset.add_constant('RK45_tol', rk45_tol) + + dt = delta(seconds=30).total_seconds() + RK45Particles = ptype[mode].add_variable('next_dt', dtype=np.float32, initial=dt) + pset = ParticleSet(fieldset, pclass=RK45Particles, + lon=np.zeros(npart) + 20., + lat=np.linspace(0, 80, npart)) + pset.execute(AdvectionRK45, runtime=delta(hours=2), dt=dt) + assert (np.diff(pset.lon) > 1.e-4).all() + assert np.isclose(fieldset.RK45_tol, rk45_tol/(1852*60)) + print(fieldset.RK45_tol) + + def periodicfields(xdim, ydim, uvel, vvel): dimensions = {'lon': np.linspace(0., 1., xdim+1, dtype=np.float32)[1:], # don't include both 0 and 1, for periodic b.c. 'lat': np.linspace(0., 1., ydim+1, dtype=np.float32)[1:]} From 27dc7e96da31b5633d485e2e91d66326c073dd6e Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 25 Apr 2024 09:04:16 +0200 Subject: [PATCH 6/7] Updating some tests for new RK45 implementation --- docs/examples/example_moving_eddies.py | 9 ++++++++- tests/test_advection.py | 7 ++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/docs/examples/example_moving_eddies.py b/docs/examples/example_moving_eddies.py index 24d47f2b96..25fafa9501 100644 --- a/docs/examples/example_moving_eddies.py +++ b/docs/examples/example_moving_eddies.py @@ -97,7 +97,14 @@ def cosd(x): data = {'U': U, 'V': V, 'P': P} dimensions = {'lon': lon, 'lat': lat, 'time': time} - return FieldSet.from_data(data, dimensions, transpose=True, mesh=mesh) + + fieldset = FieldSet.from_data(data, dimensions, transpose=True, mesh=mesh) + + # setting some constants for AdvectionRK45 kernel + fieldset.RK45_min_dt = 1e-3 + fieldset.RK45_max_dt = 1e2 + fieldset.RK45_tol = 1e-5 + return fieldset def moving_eddies_example(fieldset, outfile, npart=2, mode='jit', verbose=False, diff --git a/tests/test_advection.py b/tests/test_advection.py index e1ab2f0ee4..50b9cd3320 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -299,7 +299,12 @@ def fieldset_stationary(xdim=100, ydim=100, maxtime=delta(hours=6)): 'time': time} data = {'U': np.ones((xdim, ydim, 1), dtype=np.float32) * u_0 * np.cos(f * time), 'V': np.ones((xdim, ydim, 1), dtype=np.float32) * -u_0 * np.sin(f * time)} - return FieldSet.from_data(data, dimensions, mesh='flat', transpose=True) + fieldset = FieldSet.from_data(data, dimensions, mesh='flat', transpose=True) + # setting some constants for AdvectionRK45 kernel + fieldset.RK45_min_dt = 1e-3 + fieldset.RK45_max_dt = 1e2 + fieldset.RK45_tol = 1e-5 + return fieldset @pytest.mark.parametrize('mode', ['scipy', 'jit']) From 92c82633575f9e9bb28ef21f3b9714413c3ffc33 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 25 Apr 2024 10:52:46 +0200 Subject: [PATCH 7/7] Updating default value of RK45_max_dt --- parcels/application_kernels/advection.py | 3 ++- parcels/kernel.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/parcels/application_kernels/advection.py b/parcels/application_kernels/advection.py index 511e191567..c661b4610d 100644 --- a/parcels/application_kernels/advection.py +++ b/parcels/application_kernels/advection.py @@ -64,7 +64,8 @@ def AdvectionRK45(particle, fieldset, time): 1e-5 * dt by default. Note that this kernel requires a Particle Class that has an extra Variable 'next_dt' - and a FieldSet with 'RK45_tol', 'RK45_min_dt' and 'RK45_max_dt' constants. + and a FieldSet with constants 'RK45_tol' (in meters), 'RK45_min_dt' (in seconds) + and 'RK45_max_dt' (in seconds). """ particle.dt = min(particle.next_dt, fieldset.RK45_max_dt) c = [1./4., 3./8., 12./13., 1., 1./2.] diff --git a/parcels/kernel.py b/parcels/kernel.py index 4d883a89e4..1c9758440a 100644 --- a/parcels/kernel.py +++ b/parcels/kernel.py @@ -335,8 +335,8 @@ def check_fieldsets_in_kernels(self, pyfunc): logger.info("Setting RK45 minimum timestep to 1 s. Use fieldset.add_constant('RK45_min_dt', [timestep]) to change.") self.fieldset.add_constant('RK45_min_dt', 1) if not hasattr(self.fieldset, 'RK45_max_dt'): - logger.info("Setting RK45 maximum timestep to 1000 s. Use fieldset.add_constant('RK45_max_dt', [timestep]) to change.") - self.fieldset.add_constant('RK45_max_dt', 1000) + logger.info("Setting RK45 maximum timestep to 1 day. Use fieldset.add_constant('RK45_max_dt', [timestep]) to change.") + self.fieldset.add_constant('RK45_max_dt', 60*60*24) def check_kernel_signature_on_version(self): numkernelargs = 0