How can I use traits more efficiently in a solver for coupled differential equations?

Question

some context

I am writing a solver for systems of (delay-)coupled differential equations i.e. "integrating the equations". Similar to the Lorenz-System (no delay feedback) or the Mackey-Glass-system (with delay-feedback), but with (potentially) multiple of such systems that feedback into each other either directly or after a certain delay. The feedback is through a complex network topology. I started using traits in order to restrict how a certain system can be defined so that it can work with the rest of my code (using Runge-Kutta methods for example). In general I have two kinds of systems - with and without delay (for which I have some special cases each that I'd like to optimize the integration process a little). Right now I feel I have to write everything two or more times and this is inefficient hard to understand. Every "dynamical system" has a "state" (think of a vector of coordinates) and a function f(state) which gives the derivative for a state. Think about it like a vector-field that attaches an arrow to every point in space. Model is just a set of parameters important for f, that characterize a given system. My two traits I thought of:

pub trait DynamicalSystem {
    type StateT;
    type ModelT; // the model parameters
    fn f(state: &StateT, model: &ModelT) -> StateT;
}
pub trait DynamicalDelaySystem {
    type StateT;
    type ModelT;
    type DelayT;
    fn f(state: &StateT, model: &ModelT, delay: &DelayT) -> StateT; 
    fn keep_delay(state: &StateT) -> DelayT; // keep *something* for the feedback
}

If I am integrating this with the Euler method I need to define two integrations:

euler(state: &mut StateT, model: &ModelT, dt: &f64, f: fn(&StateT, &ModelT) -> StateT) {
    *state += f(&state, &model) * dt;
}
euler_delay(state: &mut StateT, model: &ModelT, delay: &DelayT, dt: &f64, f: fn(&StateT, &ModelT) -> StateT, keep_delay: fn(&StateT) -> DelayT) {
    *state += f(&state, &model, &delay) * dt;
    keep_delay(state); // some function that puts the delay into some container e.g. a ringbuffer for later use
}

My idea was to write an "Integrator"-struct that keeps a vector of StateT's, ModelT's and so on. But I need to do it twice - once with and once without delay.

struct Integrator<SystemT> {
    dt: f64,
    states: Vec<SystemT::StateT>,
    models: Vec<SystemT::ModelT>,
}
impl<SystemT> Integrator<SystemT> {
    fn integration_step(&mut self) {
        for (s, m) in &mut self.states.iter_mut().zip(self.models) {
            euler_step(&mut s, &m, &dt);
    }
struct DelayIntegrator<SystemT> {
    // the same as `Integrator<SystemT>`
    history: Vec<Vec<SystemT::DelayT>>,
}
impl<SystemT> DelayIntegrator<SystemT> {
    // everything again, but retrieve delayed values and save new values after integration step.
}

This seems kind of redundant. Even worse: Besides the general case I have special cases that I want to address (i.e. optimize for):

1. Only one single dynamical system (states.len() == models.len() == 1)
2. multiple dynamical systems with identical parameters (states.len() > 1 && models.len() == 1)
3. Feedback between systems, but no delay i.e. I don't need a complicated container to store delay-values in.
4. no (delay-)feedback My first try to solve 1 and 2 was to use a function pointer for integration_step that I set in the new-method of the Integrator. Depending on which case of the system is setup (number of systems, number of model parameters). But function pointers seem like a hack that would be slow.

question:

What general advice for refactorization / more "DRY" code would you recommend here? I am unsure if I understand the full potential of traits, yet. Especially combining multiple traits instead of writing everything into one is difficult for me to conceive. Could I split this out into multiple traits like DelayFeedback, InstantFeedback, NoFeedback, SingleSystem, MultipleIdenticalSystems and MultipleDistinctSystems? I seem unable to use traits to enforce the existence of a structs member-variable e.g. the delay-container.

Answer 1

It seems to me that a DynamicalSystem is a special case of a DynamicalDelaySystem where DelayT = (). That observation may let you eliminate the DynamicalSystem trait.

pub trait DynamicalSystem {
    type StateT;
    type ModelT;
    type DelayT;

    fn next(state: &Self::StateT, model: &Self::ModelT, delay: &Self::DelayT) -> Self::StateT;

    fn keep_delay(state: &Self::StateT) -> Self::DelayT;
}

impl DynamicalSystem for usize {
    type StateT = usize;
    type ModelT = usize;
    type DelayT = (); // there is no delay
    fn next(state: &Self::StateT, model: &Self::ModelT, delay: &Self::DelayT) -> Self::StateT {
        *state
    }

    fn keep_delay(state: &Self::StateT) -> Self::DelayT {
        ()
    }
}

Or maybe try the experimental trait alias feature:

#![feature(trait_alias)]
pub trait DynamicalSystem = DynamicalDelaySystem<DelayT = ()>;

pub trait DynamicalDelaySystem {
    type StateT;
    type ModelT;
    type DelayT;

    fn next(state: &Self::StateT, model: &Self::ModelT, delay: &Self::DelayT) -> Self::StateT;

    fn keep_delay(state: &Self::StateT) -> Self::DelayT;
}

Another possible way to do it is to make DynamicalSystem a super trait of DynamicalDelaySystem:

pub trait DynamicalSystem {
    type StateT;
    type ModelT;
    fn next(state: &Self::StateT, model: &Self::ModelT) -> Self::StateT;
}

pub trait DynamicalDelaySystem: DynamicalSystem {
    type DelayT;

    fn next_with_delay(
        state: &Self::StateT,
        model: &Self::ModelT,
        delay: &Self::DelayT,
    ) -> Self::StateT;

    fn keep_delay(state: &Self::StateT) -> Self::DelayT;
}

Finally, if nothing else is good enough, you could use macros to do the code duplication work for you.