Moving RK2 and RK4 methods to separate integrators.py functions

[erikvansebille]

Description

This PR implements an idea by @michaeldenes (#2338) to provide users with the option to call an integrator (RK2 or RK4) when they themselves specify the right-hand-side of a tendency equation.

It also shows how to use this, in a new section at the end of the dt integrators tutorial

Things to discuss/explore

1. Is this a good idea in the first place?
2. How/where should we put the RK2 and RK4 functions of _integrators.py in the public API? Directly under kernels? or under integrators?
3. Any other suggestions/improvements?

Checklist

• Closes Writing kernels and integration schemes in a more 'function'-like way #2338
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

AI Disclosure

None

[michaeldenes]

As a preliminary thought, a way to generalise this even further might be to do something like this:

# Functions defining the RHS for advection in 2D and 3D
def advection_rhs_2D(fieldset, time, particles, positions):
    lon, lat = positions #evaluate the derivative at the provided positions/intermediate values
    return fieldset.UV[time,  particles.depth, lon, lat, particles]


def advection_rhs_3D(fieldset, time, particles, positions):
    depth, lon, lat = positions #evaluate the derivative at the provided positions/intermediate values
    return fieldset.UVW[time,  depth, lon, lat, particles]


# A single RK2 function that is generalised to any dimension
def RK2(fieldset, particles, rhs, variables):
    intermediateVals = np.zeros_like(variables)
    fields = rhs(fieldset, particles.time, particles, [particles[variable] for variable in variables])
    for i, variable in enumerate(variables):
        intermediateVals[i] = particles[variable] + fields[i]*0.5*particles.dt
    
    t = particles.time + 0.5 * particles.dt
    return rhs(fieldset, t, particles, intermediateVals)

# Advection kernels
def AdvectionRK2(particles, fieldset):  # pragma: no cover
    """Advection of particles using second-order Runge-Kutta integration."""
    dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt)

    u, v = RK2(fieldset, particles, advection_rhs_2D, ['U', 'V'])
    particles.dlon += u * dt
    particles.dlat += v * dt

def AdvectionRK2_3D(particles, fieldset):  # pragma: no cover
    """Advection of particles using second-order Runge-Kutta integration."""
    dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt)

    u, v, w = RK2(fieldset, particles, advection_rhs_3D, ['U', 'V', 'W'])
    particles.dlon += u * dt
    particles.dlat += v * dt
    particles.ddepth += w * dt

In this way, you can apply RK2, RK4, etc. to any variable. (Haven't tested this, need @VeckoTheGecko help for the pixi install stuff ;)

BTW, I think a cleaner way is to just define the RHS inside the kernel, and then apply the integrator against it (like in an earlier commit).

EDIT: Added a line that would create fields which was missing

[erikvansebille]

Nice idea, @michaeldenes. We'd need to double-check that this also works if particles is an array of depths, lats and lons (kernels are now vectored!) and think what the memory overhead is of the intermediateVals. But that will come when you try this out

Also, instead of RK2(fieldset, ... fields), would it not be more intuitive to call as `RK2([fieldset.U, fieldset.V, fieldset.W], ...)?

[wyatt-fluidnumerics]

This is a really cool change and I think will be really useful for users. Just a few comments, my main concern is the way RK2/RK4 check the dimensions/components of the field

[wyatt-fluidnumerics]

I like the _validate_rhs_output function, now that it exists, I don't believe any of the len(..) > 1 checks are needed? Other than that this looks good to me.

[wyatt-fluidnumerics]

Now that we have _validate_rhs_output I don't believe we need any if len(...) > 1 checks?

[erikvansebille]

Again, good catch! Now fixed in 85af999

[wyatt-fluidnumerics]

Great, looks good to me now!