ODE systems

So far we've focused on using custom structs to specify individual equations, now we need to put these together to specify the entire system of equations.

Implementing the OdeEquations trait

To specify the entire system of equations, you need to implement the Op, OdeEquations and OdeEquationsRef traits for your struct.

Getting all your traits in order

The OdeEquations trait requires methods that return objects corresponding to the right-hand side function, mass matrix, root function, initial condition, and output functions. Therefore, you need to already have structs that implement the NonLinearOp, LinearOp, and ConstantOp traits for these functions. For the purposes of this example, we will assume that you have already implemented these traits for your structs.

Often, the structs that implement these traits will have to use data defined in the struct that implements the OdeEquations trait. For example, they might wish to have a reference to the same parameter vector p. Therefore, you will often need to define lifetimes for these structs to ensure that they can access the data they need.

Note that these struct will need to be lightweight and should not contain a significant amount of data. The data should be stored in the struct that implements the OdeEquations trait. This is because these structs will be created and destroyed many times during the course of the simulation (e.g. every time the right-hand side function is called).

My custom ode equations

We'll define a central struct MyEquations that will hold the data for the entire system of equations. In this example, MyEquations will own a parameter vector p that is used by all the equations.

use crate::{C, V};
use diffsol::Vector;
pub struct MyEquations {
    pub p: V,
}

impl MyEquations {
    pub fn new() -> Self {
        MyEquations {
            p: V::zeros(2, C::default()),
        }
    }
}

As with individual functions, we also need to implement the Op trait for this struct.

use crate::{MyEquations, C, M, T, V};
use diffsol::{Op, Vector};

impl Op for MyEquations {
    type T = T;
    type V = V;
    type M = M;
    type C = C;
    fn nstates(&self) -> usize {
        1
    }
    fn nout(&self) -> usize {
        1
    }
    fn nparams(&self) -> usize {
        2
    }
    fn context(&self) -> &Self::C {
        self.p.context()
    }
}

Implementing the OdeEquations traits

OdeEquations and OdeEquationsRef are a pair of traits that together define the system of equations using the individual function structs that we defined earlier. The OdeEquationsRef trait is designed to allow individual function structs to be able to reference the data in the main OdeEquations struct, allowing them to share data.

use crate::{MyEquations, MyInit, MyMass, MyOut, MyRhs, MyRoot, V};
use diffsol::{OdeEquations, OdeEquationsRef, Vector};

impl<'a> OdeEquationsRef<'a> for MyEquations {
    type Rhs = MyRhs<'a>;
    type Mass = MyMass<'a>;
    type Init = MyInit<'a>;
    type Root = MyRoot<'a>;
    type Out = MyOut<'a>;
}

impl OdeEquations for MyEquations {
    fn rhs(&self) -> <MyEquations as OdeEquationsRef<'_>>::Rhs {
        MyRhs { p: &self.p }
    }
    fn mass(&self) -> Option<<MyEquations as OdeEquationsRef<'_>>::Mass> {
        Some(MyMass { p: &self.p })
    }
    fn init(&self) -> <MyEquations as OdeEquationsRef<'_>>::Init {
        MyInit { p: &self.p }
    }
    fn root(&self) -> Option<<MyEquations as OdeEquationsRef<'_>>::Root> {
        Some(MyRoot { p: &self.p })
    }
    fn out(&self) -> Option<<MyEquations as OdeEquationsRef<'_>>::Out> {
        Some(MyOut { p: &self.p })
    }
    fn set_params(&mut self, p: &V) {
        self.p.copy_from(p);
    }
    fn get_params(&self, p: &mut Self::V) {
        p.copy_from(&self.p);
    }
}

Creating the problem

Now that we have our custom OdeEquations struct, we can use it in an OdeBuilder to create the problem.

use crate::{MyEquations, M};
use diffsol::OdeBuilder;

pub fn build() {
    OdeBuilder::<M>::new()
        .p(vec![1.0, 10.0])
        .build_from_eqn(MyEquations::new())
        .unwrap();
}