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

• This PR contains AI-generated content.
  • I have tested any AI-generated content in my PR.
  • I take responsibility for any AI-generated content in my PR.
  • Describe how you used it (e.g., by pasting your prompt):

[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], ...)?