Skip to content

Commit

Permalink
doc: OdeSolverMethod methods, Ops, events, and clarify matrix sparsit…
Browse files Browse the repository at this point in the history
…y detection (#51)
  • Loading branch information
martinjrobins authored May 9, 2024
1 parent 0ac6142 commit 80c758e
Show file tree
Hide file tree
Showing 4 changed files with 89 additions and 43 deletions.
20 changes: 17 additions & 3 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,19 @@
//! You will also need to choose a matrix type to use. DiffSol can use the [nalgebra](https://nalgebra.org) `DMatrix` type, the [faer](https://github.com/sarah-ek/faer-rs) `Mat` type, or any other type that implements the
//! [Matrix] trait. You can also use the [sundials](https://computation.llnl.gov/projects/sundials) library for the matrix and vector types (see [SundialsMatrix]).
//!
//! ## The solver
//!
//! To solve the problem, you need to choose a solver. DiffSol provides the following solvers:
//! - A Backwards Difference Formulae [Bdf] solver, suitable for stiff problems and singular mass matrices.
//! - A Singly Diagonally Implicit Runge-Kutta (SDIRK or ESDIRK) solver [Sdirk]. You can use your own butcher tableau using [Tableau] or use one of the provided ([Tableau::tr_bdf2], [Tableau::esdirk34]).
//! - A BDF solver that wraps the IDA solver solver from the sundials library ([SundialsIda], requires the `sundials` feature).
//!
//! See the [OdeSolverMethod] trait for a more detailed description of the available methods on each solver.
//! See the [OdeSolverMethod] trait for a more detailed description of the available methods on each solver. Possible workflows are:
//! - Initialise the problem using [OdeSolverState::new] or [OdeSolverState::new_consistent], and then use [OdeSolverMethod::set_problem] to setup the solver with the problem and [OdeSolverState] instance.
//! - Use the [OdeSolverMethod::step] method to step the solution forward in time with an internal time step chosen by the solver to meet the error tolerances.
//! - Use the [OdeSolverMethod::interpolate] method to interpolate the solution between the last two time steps.
//! - Use the [OdeSolverMethod::set_stop_time] method to stop the solver at a specific time (i.e. this will override the internal time step so that the solver stops at the specified time).
//! - Alternatively, use the convenience functions [OdeSolverMethod::solve] and [OdeSolverMethod::make_consistent_and_solve] that will both initialise the problem and solve the problem up to a specific time.
//!
//! ## DiffSL
//!
Expand All @@ -38,13 +45,20 @@
//!
//! Via an implementation of [OdeEquations], the user provides the action of the jacobian on a vector `J(x) v`. By default DiffSol uses this to generate a jacobian matrix for the ODE solver.
//! Generally this requires `n` evaluations of the jacobian action for a system of size `n`, so it is often more efficient if the user can provide the jacobian matrix directly
//! by implementing the [NonLinearOp::jacobian_inplace] and the [LinearOp::matrix_inplace] (if applicable) functions.
//! by also implementing the optional [NonLinearOp::jacobian_inplace] and the [LinearOp::matrix_inplace] (if applicable) functions.
//!
//! DiffSol also provides an experimental feature to calculate sparse jacobians more efficiently by automatically detecting the sparsity pattern of the jacobian and using
//! If this is not possible, DiffSol also provides an experimental feature to calculate sparse jacobians more efficiently by automatically detecting the sparsity pattern of the jacobian and using
//! colouring \[1\] to reduce the number of jacobian evaluations. You can enable this feature by enabling [OdeBuilder::use_coloring()] option when building the ODE problem.
//! Note that if your implementation of [NonLinearOp::jac_mul_inplace] uses any control flow that depends on the input vector (e.g. an if statement that depends on the value of `x`),
//! the sparsity detection may not be accurate and you may need to provide the jacobian matrix directly.
//!
//! \[1\] Gebremedhin, A. H., Manne, F., & Pothen, A. (2005). What color is your Jacobian? Graph coloring for computing derivatives. SIAM review, 47(4), 629-705.
//!
//! ## Events / Root finding
//!
//! DiffSol provides a simple way to detect user-provided events during the integration of the ODEs. You can use this by providing a closure that has a zero-crossing at the event you want to detect, using the [OdeBuilder::build_ode_with_root] builder,
//! or by providing a [NonLinearOp] that has a zero-crossing at the event you want to detect. To use the root finding feature while integrating with the solver, you can use the return value of [OdeSolverMethod::step] to check if an event has been detected.
//!
//! ## Nonlinear and linear solvers
//!
//! DiffSol provides generic nonlinear and linear solvers that are used internally by the ODE solver. You can use the solvers provided by DiffSol, or implement your own following the provided traits.
Expand Down
43 changes: 40 additions & 3 deletions src/ode_solver/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,11 +262,11 @@ impl OdeBuilder {
///
///
/// // dy/dt = y
/// // y(0) = 1
/// // y(0) = 0.1
/// let problem = OdeBuilder::new()
/// .build_ode::<M, _, _, _>(
/// |x, _p, _t, y| y.copy_from(x),
/// |x, _p, _t, v , y| y.copy_from(v),
/// |x, _p, _t, y| y[0] = x[0],
/// |x, _p, _t, v , y| y[0] = v[0],
/// |p, _t| DVector::from_element(1, 0.1),
/// );
/// ```
Expand Down Expand Up @@ -304,6 +304,43 @@ impl OdeBuilder {
))
}

/// Build an ODE problem with an event.
///
/// # Arguments
///
/// - `rhs`: Function of type Fn(x: &V, p: &V, t: S, y: &mut V) that computes the right-hand side of the ODE.
/// - `rhs_jac`: Function of type Fn(x: &V, p: &V, t: S, v: &V, y: &mut V) that computes the multiplication of the Jacobian of the right-hand side with the vector v.
/// - `init`: Function of type Fn(p: &V, t: S) -> V that computes the initial state.
/// - `root`: Function of type Fn(x: &V, p: &V, t: S, y: &mut V) that computes the root function.
/// - `nroots`: Number of roots (i.e. number of elements in the `y` arg in `root`), an event is triggered when any of the roots changes sign.
///
/// # Generic Arguments
///
/// - `M`: Type that implements the `Matrix` trait. Often this must be provided explicitly (i.e. `type M = DMatrix<f64>; builder.build_ode::<M, _, _, _, _>`).
///
/// # Example
///
///
///
/// ```
/// use diffsol::OdeBuilder;
/// use nalgebra::DVector;
/// type M = nalgebra::DMatrix<f64>;
///
///
/// // dy/dt = y
/// // y(0) = 0.1
/// // event at y = 0.5
/// let problem = OdeBuilder::new()
/// .build_ode_with_root::<M, _, _, _, _>(
/// |x, _p, _t, y| y[0] = x[0],
/// |x, _p, _t, v , y| y[0] = v[0],
/// |p, _t| DVector::from_element(1, 0.1),
/// |x, _p, _t, y| y[0] = x[0] - 0.5,
/// 1,
/// );
/// ```
#[allow(clippy::type_complexity)]
pub fn build_ode_with_root<M, F, G, I, H>(
self,
Expand Down
7 changes: 4 additions & 3 deletions src/ode_solver/method.rs
Original file line number Diff line number Diff line change
Expand Up @@ -103,8 +103,8 @@ pub struct OdeSolverState<V: Vector> {
}

impl<V: Vector> OdeSolverState<V> {
/// Create a new solver state from an ODE problem. Note that this does not make the state consistent with the algebraic constraints.
/// If you need to make the state consistent, use `new_consistent` instead.
/// Create a new solver state from an ODE problem. Note that this does not make the state consistent with any algebraic constraints.
/// If you need to make the state consistent, use [Self::new_consistent].
pub fn new<Eqn>(ode_problem: &OdeSolverProblem<Eqn>) -> Self
where
Eqn: OdeEquations<T = V::T, V = V>,
Expand All @@ -115,7 +115,8 @@ impl<V: Vector> OdeSolverState<V> {
Self { y, t, h }
}

/// Create a new solver state from an ODE problem, making the state consistent with the algebraic constraints.
/// Create a new solver state from an [OdeSolverProblem], making the state consistent with the algebraic constraints using a solver that implements [NonLinearSolver].
/// If there are no algebraic constraints, please use [Self::new] instead.
pub fn new_consistent<Eqn, S>(
ode_problem: &OdeSolverProblem<Eqn>,
root_solver: &mut S,
Expand Down
62 changes: 28 additions & 34 deletions src/op/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,38 +80,43 @@ impl OpStatistics {
}
}

// NonLinearOp is a trait for non-linear operators. It extends the Op trait with methods for
// computing the operator and its Jacobian. The operator is defined by the call_inplace method,
// which computes the operator at a given state and time. The Jacobian is defined by the
// jac_mul_inplace method, which computes the product of the Jacobian with a given vector.
// NonLinearOp is a trait that defines a nonlinear operator or function `F` that maps an input vector `x` to an output vector `y`, (i.e. `y = F(x, t)`).
// It extends the [Op] trait with methods for computing the operator and its Jacobian.
//
// The operator is defined by the [Self::call_inplace] method, which computes the function `F(x, t)` at a given state and time.
// The Jacobian is defined by the [Self::jac_mul_inplace] method, which computes the product of the Jacobian with a given vector `J(x, t) * v`.
pub trait NonLinearOp: Op {
/// Compute the operator at a given state and time.
/// Compute the operator `F(x, t)` at a given state and time.
fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V);

/// Compute the product of the Jacobian with a given vector.
/// Compute the product of the Jacobian with a given vector `J(x, t) * v`.
fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V);

/// Compute the operator at a given state and time, and return the result.
/// Compute the operator `F(x, t)` at a given state and time, and return the result.
/// Use `[Self::call_inplace]` to for a non-allocating version.
fn call(&self, x: &Self::V, t: Self::T) -> Self::V {
let mut y = Self::V::zeros(self.nout());
self.call_inplace(x, t, &mut y);
y
}

/// Compute the product of the Jacobian with a given vector, and return the result.
/// Compute the product of the Jacobian with a given vector `J(x, t) * v`, and return the result.
/// Use `[Self::jac_mul_inplace]` to for a non-allocating version.
fn jac_mul(&self, x: &Self::V, t: Self::T, v: &Self::V) -> Self::V {
let mut y = Self::V::zeros(self.nstates());
self.jac_mul_inplace(x, t, v, &mut y);
y
}

/// Compute the Jacobian of the operator and store it in the matrix `y`.
/// Compute the Jacobian matrix `J(x, t)` of the operator and store it in the matrix `y`.
/// `y` should have been previously initialised using the output of [`Op::sparsity`].
/// The default implementation of this method computes the Jacobian using [Self::jac_mul_inplace],
/// but it can be overriden for more efficient implementations.
fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
self._default_jacobian_inplace(x, t, y);
}

/// Default implementation of the Jacobian computation.
/// Default implementation of the Jacobian computation (this is the default for [Self::jacobian_inplace]).
fn _default_jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
let mut v = Self::V::zeros(self.nstates());
let mut col = Self::V::zeros(self.nout());
Expand All @@ -123,7 +128,8 @@ pub trait NonLinearOp: Op {
}
}

/// Compute the Jacobian of the operator and return it.
/// Compute the Jacobian matrix `J(x, t)` of the operator and return it.
/// See [Self::jacobian_inplace] for a non-allocating version.
fn jacobian(&self, x: &Self::V, t: Self::T) -> Self::M {
let n = self.nstates();
let mut y = Self::M::new_from_sparsity(n, n, self.sparsity());
Expand All @@ -132,31 +138,35 @@ pub trait NonLinearOp: Op {
}
}

/// LinearOp is a trait for linear operators (i.e. they only depend linearly on the input `x`). It extends the Op trait with methods for
/// calling the operator via a GEMV operation (i.e. `y = t * A * x + beta * y`), and for computing the matrix representation of the operator.
/// LinearOp is a trait for linear operators (i.e. they only depend linearly on the input `x`), see [NonLinearOp] for a non-linear op.
/// An example of a linear operator is a matrix-vector product `y = A(t) * x`, where `A(t)` is a matrix.
/// It extends the [Op] trait with methods for calling the operator via a GEMV-like operation (i.e. `y = t * A * x + beta * y`), and for computing the matrix representation of the operator.
pub trait LinearOp: Op {
/// Compute the operator at a given state and time
/// Compute the operator `y = A(t) * x` at a given state and time, the default implementation uses [Self::gemv_inplace].
fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) {
let beta = Self::T::zero();
self.gemv_inplace(x, t, beta, y);
}

/// Computer the operator via a GEMV operation (i.e. `y = t * A * x + beta * y`)
/// Compute the operator via a GEMV operation (i.e. `y = A(t) * x + beta * y`)
fn gemv_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V);

/// Compute the matrix representation of the operator and return it.
/// Compute the matrix representation of the operator `A(t)` and return it.
/// See [Self::matrix_inplace] for a non-allocating version.
fn matrix(&self, t: Self::T) -> Self::M {
let mut y = Self::M::new_from_sparsity(self.nstates(), self.nstates(), self.sparsity());
self.matrix_inplace(t, &mut y);
y
}

/// Compute the matrix representation of the operator and store it in the matrix `y`.
/// Compute the matrix representation of the operator `A(t)` and store it in the matrix `y`.
/// The default implementation of this method computes the matrix using [Self::gemv_inplace],
/// but it can be overriden for more efficient implementations.
fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) {
self._default_matrix_inplace(t, y);
}

/// Default implementation of the matrix computation.
/// Default implementation of the matrix computation, see [Self::matrix_inplace].
fn _default_matrix_inplace(&self, t: Self::T, y: &mut Self::M) {
let mut v = Self::V::zeros(self.nstates());
let mut col = Self::V::zeros(self.nout());
Expand Down Expand Up @@ -197,22 +207,6 @@ impl<C: Op> Op for &C {
}
}

//impl <C: LinearOp> NonLinearOp for C {
// fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) {
// C::call_inplace(self, x, t, y)
// }
// fn jac_mul_inplace(&self, _x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
// C::call_inplace(self, v, t, y)
// }
// fn jacobian_inplace(&self, _x: &Self::V, t: Self::T, y: &mut Self::M) {
// C::matrix_inplace(self, t, y)
// }
// fn sparsity(&self) -> Option<&<Self::M as Matrix>::Sparsity> {
// C::sparsity(self)
// }
//
//}

impl<C: NonLinearOp> NonLinearOp for &C {
fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) {
C::call_inplace(*self, x, t, y)
Expand Down

0 comments on commit 80c758e

Please sign in to comment.