Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

#19 document OdeBuilder #20

Merged
merged 2 commits into from
Apr 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ diffsl16-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffs
diffsl17-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm17-0"], optional = true }
petgraph = "0.6.4"
faer = "0.18.2"
sundials-sys = { version = "0.4.0", features = ["ida"], optional = true }
sundials-sys = { version = "0.4.0", features = ["ida", "static_libraries"], optional = true }


[dev-dependencies]
Expand Down
2 changes: 1 addition & 1 deletion src/jacobian/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ mod tests {

#[test]
fn build_coloring() {
let test_triplets = vec![
let test_triplets = [
vec![(0, 0, 1.0), (1, 1, 1.0)],
vec![(0, 0, 1.0), (0, 1, 1.0), (1, 1, 1.0)],
vec![(1, 1, 1.0)],
Expand Down
7 changes: 5 additions & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,8 @@ use matrix::{DenseMatrix, Matrix, MatrixViewMut};
pub use nonlinear_solver::newton::NewtonNonlinearSolver;
use nonlinear_solver::NonLinearSolver;
pub use ode_solver::{
bdf::Bdf, equations::OdeEquations, OdeSolverMethod, OdeSolverProblem, OdeSolverState,
bdf::Bdf, builder::OdeBuilder, equations::OdeEquations, OdeSolverMethod, OdeSolverProblem,
OdeSolverState,
};
use op::NonLinearOp;
use scalar::{IndexType, Scalar, Scale};
Expand All @@ -68,7 +69,9 @@ pub use scalar::scale;
#[cfg(test)]
mod tests {

use crate::{ode_solver::OdeBuilder, vector::Vector, Bdf, OdeSolverMethod, OdeSolverState};
use crate::{
ode_solver::builder::OdeBuilder, vector::Vector, Bdf, OdeSolverMethod, OdeSolverState,
};

// WARNING: if this test fails and you make a change to the code, you should update the README.md file as well!!!
#[test]
Expand Down
268 changes: 268 additions & 0 deletions src/ode_solver/builder.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,268 @@
use crate::{op::Op, Matrix, OdeSolverProblem, Vector};
use anyhow::Result;

use super::equations::{OdeSolverEquations, OdeSolverEquationsMassI};

/// Builder for ODE problems. Use methods to set parameters and then call one of the build methods when done.
pub struct OdeBuilder {
t0: f64,
h0: f64,
rtol: f64,
atol: Vec<f64>,
p: Vec<f64>,
use_coloring: bool,
}

impl Default for OdeBuilder {
fn default() -> Self {
Self::new()
}
}

impl OdeBuilder {
/// Create a new builder with default parameters:
/// - t0 = 0.0
/// - h0 = 1.0
/// - rtol = 1e-6
/// - atol = [1e-6]
/// - p = []
/// - use_coloring = false
pub fn new() -> Self {
Self {
t0: 0.0,
h0: 1.0,
rtol: 1e-6,
atol: vec![1e-6],
p: vec![],
use_coloring: false,
}
}

/// Set the initial time.
pub fn t0(mut self, t0: f64) -> Self {
self.t0 = t0;
self
}

/// Set the initial step size.
pub fn h0(mut self, h0: f64) -> Self {
self.h0 = h0;
self
}

/// Set the relative tolerance.
pub fn rtol(mut self, rtol: f64) -> Self {
self.rtol = rtol;
self
}

/// Set the absolute tolerance.
pub fn atol<V, T>(mut self, atol: V) -> Self
where
V: IntoIterator<Item = T>,
f64: From<T>,
{
self.atol = atol.into_iter().map(|x| f64::from(x)).collect();
self
}

/// Set the parameters.
pub fn p<V, T>(mut self, p: V) -> Self
where
V: IntoIterator<Item = T>,
f64: From<T>,
{
self.p = p.into_iter().map(|x| f64::from(x)).collect();
self
}

/// Set whether to use coloring when computing the Jacobian.
/// This can speed up the computation of the Jacobian for large sparse systems.
/// However, it relys on the sparsity of the Jacobian being constant,
/// and for certain systems it may detect the wrong sparsity pattern.
pub fn use_coloring(mut self, use_coloring: bool) -> Self {
self.use_coloring = use_coloring;
self
}

fn build_atol<V: Vector>(atol: Vec<f64>, nstates: usize) -> Result<V> {
if atol.len() == 1 {
Ok(V::from_element(nstates, V::T::from(atol[0])))
} else if atol.len() != nstates {
Err(anyhow::anyhow!(
"atol must have length 1 or equal to the number of states"
))
} else {
let mut v = V::zeros(nstates);
for (i, &a) in atol.iter().enumerate() {
v[i] = V::T::from(a);
}
Ok(v)
}
}

fn build_p<V: Vector>(p: Vec<f64>) -> V {
let mut v = V::zeros(p.len());
for (i, &p) in p.iter().enumerate() {
v[i] = V::T::from(p);
}
v
}

/// Build an ODE problem with a mass matrix.
///
/// # 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.
/// - `mass`: Function of type Fn(v: &V, p: &V, t: S, y: &mut V) that computes the multiplication of the mass matrix with the vector v.
/// - `init`: Function of type Fn(p: &V, t: S) -> V that computes the initial state.
///
/// # 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
/// // 0 = z - y
/// // y(0) = 0.1
/// // z(0) = 0.1
/// let problem = OdeBuilder::new()
/// .build_ode_with_mass::<M, _, _, _, _>(
/// |x, _p, _t, y| {
/// y[0] = x[0];
/// y[1] = x[1] - x[0];
/// },
/// |x, _p, _t, v, y| {
/// y[0] = v[0];
/// y[1] = v[1] - v[0];
/// },
/// |v, _p, _t, y| {
/// y[0] = v[0];
/// y[1] = 0.0;
/// },
/// |p, _t| DVector::from_element(2, 0.1),
/// );
/// ```
#[allow(clippy::type_complexity)]
pub fn build_ode_with_mass<M, F, G, H, I>(
self,
rhs: F,
rhs_jac: G,
mass: H,
init: I,
) -> Result<OdeSolverProblem<OdeSolverEquations<M, F, G, H, I>>>
where
M: Matrix,
F: Fn(&M::V, &M::V, M::T, &mut M::V),
G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
H: Fn(&M::V, &M::V, M::T, &mut M::V),
I: Fn(&M::V, M::T) -> M::V,
{
let p = Self::build_p(self.p);
let eqn = OdeSolverEquations::new_ode_with_mass(
rhs,
rhs_jac,
mass,
init,
p,
M::T::from(self.t0),
self.use_coloring,
);
let atol = Self::build_atol(self.atol, eqn.nstates())?;
Ok(OdeSolverProblem::new(
eqn,
M::T::from(self.rtol),
atol,
M::T::from(self.t0),
M::T::from(self.h0),
))
}

/// Build an ODE problem with a mass matrix that is the identity matrix.
///
/// # 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.
///
/// # 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) = 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),
/// |p, _t| DVector::from_element(1, 0.1),
/// );
/// ```
pub fn build_ode<M, F, G, I>(
self,
rhs: F,
rhs_jac: G,
init: I,
) -> Result<OdeSolverProblem<OdeSolverEquationsMassI<M, F, G, I>>>
where
M: Matrix,
F: Fn(&M::V, &M::V, M::T, &mut M::V),
G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
I: Fn(&M::V, M::T) -> M::V,
{
let p = Self::build_p(self.p);
let eqn = OdeSolverEquationsMassI::new_ode(
rhs,
rhs_jac,
init,
p,
M::T::from(self.t0),
self.use_coloring,
);
let atol = Self::build_atol(self.atol, eqn.nstates())?;
Ok(OdeSolverProblem::new(
eqn,
M::T::from(self.rtol),
atol,
M::T::from(self.t0),
M::T::from(self.h0),
))
}

/// Build an ODE problem using the DiffSL language (requires the `diffsl` feature).
/// The source code is provided as a string, please see the [DiffSL documentation](https://martinjrobins.github.io/diffsl/) for more information.
#[cfg(feature = "diffsl")]
pub fn build_diffsl(
self,
source: &str,
) -> Result<OdeSolverProblem<crate::ode_solver::diffsl::DiffSl>> {
type V = crate::ode_solver::diffsl::V;
type T = crate::ode_solver::diffsl::T;
let p = Self::build_p::<V>(self.p);
let eqn = crate::ode_solver::diffsl::DiffSl::new(source, p, self.use_coloring)?;
let atol = Self::build_atol::<V>(self.atol, eqn.nstates())?;
Ok(OdeSolverProblem::new(
eqn,
T::from(self.rtol),
atol,
T::from(self.t0),
T::from(self.h0),
))
}
}
4 changes: 2 additions & 2 deletions src/ode_solver/diffsl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -214,8 +214,8 @@ mod tests {
use nalgebra::DVector;

use crate::{
linear_solver::lu::LU, nonlinear_solver::newton::NewtonNonlinearSolver,
ode_solver::OdeBuilder, Bdf, OdeEquations, OdeSolverMethod, Vector,
linear_solver::lu::LU, nonlinear_solver::newton::NewtonNonlinearSolver, Bdf, OdeBuilder,
OdeEquations, OdeSolverMethod, Vector,
};

use super::DiffSl;
Expand Down
Loading
Loading