Non-linear functions

To illustrate how to implement a custom problem struct, we will take the familar logistic equation:

\[\frac{dy}{dt} = r y (1 - y/K),\]

Our goal is to implement a custom struct that can evaluate the rhs function \(f(y, p, t)\) and the jacobian multiplied by a vector \(f'(y, p, t, v)\).

To start with, lets define a few linear algebra types that we will use in our function. We need four types:

T is the scalar type (e.g. f64)
V is the vector type (e.g. NalgebraVec<T>)
M is the matrix type (e.g. NalgebraMat<T>)
C is the context type for the rest (e.g. NalgebraContext)

use diffsol::{NalgebraContext, NalgebraMat, NalgebraVec};
pub type T = f64;
pub type V = NalgebraVec<T>;
pub type M = NalgebraMat<T>;
pub type C = NalgebraContext;

Next, we'll define a struct that we'll use to calculate our RHS equations \(f(y, p, t)\). We'll pretend that this struct has a reference to a parameter vector \(p\) that we'll use to calculate the rhs function. This makes sense since we'll have multiple functions that make up our systems of equations, and they will probably share some parameters.

use crate::V;

pub struct MyRhs<'a> {
    pub p: &'a V,
}

Now we will implement the base Op trait for our struct. The Op trait specifies the types of the vectors and matrices that will be used, as well as the number of states and outputs in the rhs function.

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

impl Op for MyRhs<'_> {
    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()
    }
}

Next we implement the NonLinearOp and NonLinearOpJacobian trait for our struct. This trait specifies the functions that will be used to evaluate the rhs function and the jacobian multiplied by a vector.

use crate::{MyRhs, T, V};
use diffsol::{NonLinearOp, NonLinearOpJacobian};

impl NonLinearOp for MyRhs<'_> {
    fn call_inplace(&self, x: &V, _t: T, y: &mut V) {
        y[0] = x[0] * x[0];
    }
}

impl NonLinearOpJacobian for MyRhs<'_> {
    fn jac_mul_inplace(&self, x: &V, _t: T, v: &V, y: &mut V) {
        y[0] = v[0] * (1.0 - 2.0 * x[0]);
    }
}

There we go, all done! This demonstrates how to implement a custom struct to specify a rhs function.

Diffsol

Non-linear functions