Skip to content

Commit

Permalink
[ode] [Important] Store the mass matrix in System. Use Fn instead of …
Browse files Browse the repository at this point in the history
…FnMut for F and J
  • Loading branch information
cpmech committed Mar 8, 2024
1 parent 717bac9 commit 377c61f
Show file tree
Hide file tree
Showing 26 changed files with 400 additions and 286 deletions.
6 changes: 3 additions & 3 deletions russell_ode/src/erk_dense_out.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ impl ErkDenseOut {
/// Updates the data and returns the number of function evaluations
pub(crate) fn update<'a, F, J, A>(
&mut self,
system: &mut System<'a, F, J, A>,
system: &System<F, J, A>,
x: f64,
y: &Vector,
h: f64,
Expand All @@ -69,8 +69,8 @@ impl ErkDenseOut {
args: &mut A,
) -> Result<usize, StrError>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
let mut n_function_eval = 0;

Expand Down
30 changes: 16 additions & 14 deletions russell_ode/src/euler_backward.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,14 @@ use russell_sparse::{CooMatrix, Genie, LinSolver, SparseMatrix};
/// Implements the backward Euler (implicit) solver
pub(crate) struct EulerBackward<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// Holds the parameters
params: Params,

/// ODE system
system: System<'a, F, J, A>,
system: &'a System<F, J, A>,

/// Vector holding the function evaluation
///
Expand All @@ -38,11 +38,11 @@ where

impl<'a, F, J, A> EulerBackward<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// Allocates a new instance
pub fn new(params: Params, system: System<'a, F, J, A>) -> Self {
pub fn new(params: Params, system: &'a System<F, J, A>) -> Self {
let ndim = system.ndim;
let jac_nnz = if params.newton.use_numerical_jacobian {
ndim * ndim
Expand All @@ -67,8 +67,8 @@ where

impl<'a, F, J, A> OdeSolverTrait<A> for EulerBackward<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// Enables dense output
fn enable_dense_output(&mut self) -> Result<(), StrError> {
Expand Down Expand Up @@ -119,7 +119,9 @@ where
let kk = self.kk.get_coo_mut().unwrap();
if self.params.newton.use_numerical_jacobian || !self.system.jac_available {
work.bench.n_function += ndim;
self.system.numerical_jacobian(kk, x_new, y_new, &self.k, h, args)?;
let aux = &mut self.dy; // using dy as a workspace
self.system
.numerical_jacobian(kk, x_new, y_new, &self.k, h, args, aux)?;
} else {
(self.system.jacobian)(kk, x_new, y_new, h, args)?;
}
Expand Down Expand Up @@ -262,12 +264,12 @@ mod tests {

// problem
let (system, data, mut args) = Samples::kreyszig_ex4_page920();
let mut yfx = data.y_analytical.unwrap();
let yfx = data.y_analytical.unwrap();
let ndim = system.ndim;

// allocate structs
let params = Params::new(Method::BwEuler);
let mut solver = EulerBackward::new(params, system);
let mut solver = EulerBackward::new(params, &system);
let mut work = Workspace::new(Method::BwEuler);

// check dense output availability
Expand Down Expand Up @@ -317,13 +319,13 @@ mod tests {

// problem
let (system, data, mut args) = Samples::kreyszig_ex4_page920();
let mut yfx = data.y_analytical.unwrap();
let yfx = data.y_analytical.unwrap();
let ndim = system.ndim;

// allocate structs
let mut params = Params::new(Method::BwEuler);
params.newton.use_numerical_jacobian = true;
let mut solver = EulerBackward::new(params, system);
let mut solver = EulerBackward::new(params, &system);
let mut work = Workspace::new(Method::BwEuler);

// numerical approximation
Expand Down Expand Up @@ -366,7 +368,7 @@ mod tests {
let mut params = Params::new(Method::BwEuler);
let (system, data, mut args) = Samples::kreyszig_ex4_page920();
params.newton.n_iteration_max = 0;
let mut solver = EulerBackward::new(params, system);
let mut solver = EulerBackward::new(params, &system);
let mut work = Workspace::new(Method::BwEuler);
assert_eq!(
solver.step(&mut work, data.x0, &data.y0, 0.1, &mut args).err(),
Expand Down
20 changes: 10 additions & 10 deletions russell_ode/src/euler_forward.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,11 @@ use russell_sparse::CooMatrix;

pub(crate) struct EulerForward<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// ODE system
system: System<'a, F, J, A>,
system: &'a System<F, J, A>,

/// Vector holding the function evaluation
///
Expand All @@ -22,11 +22,11 @@ where

impl<'a, F, J, A> EulerForward<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// Allocates a new instance
pub fn new(system: System<'a, F, J, A>) -> Self {
pub fn new(system: &'a System<F, J, A>) -> Self {
let ndim = system.ndim;
EulerForward {
system,
Expand All @@ -38,8 +38,8 @@ where

impl<'a, F, J, A> OdeSolverTrait<A> for EulerForward<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// Enables dense output
fn enable_dense_output(&mut self) -> Result<(), StrError> {
Expand Down Expand Up @@ -89,11 +89,11 @@ mod tests {

// problem
let (system, data, mut args) = Samples::kreyszig_eq6_page902();
let mut yfx = data.y_analytical.unwrap();
let yfx = data.y_analytical.unwrap();
let ndim = system.ndim;

// allocate structs
let mut solver = EulerForward::new(system);
let mut solver = EulerForward::new(&system);
let mut work = Workspace::new(Method::FwEuler);

// check dense output availability
Expand Down
28 changes: 14 additions & 14 deletions russell_ode/src/explicit_runge_kutta.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,14 @@ use russell_sparse::CooMatrix;

pub(crate) struct ExplicitRungeKutta<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// Holds the parameters
params: Params,

/// ODE system
system: System<'a, F, J, A>,
system: &'a System<F, J, A>,

/// Information such as implicit, embedded, etc.
info: Information,
Expand Down Expand Up @@ -65,11 +65,11 @@ where

impl<'a, F, J, A> ExplicitRungeKutta<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// Allocates a new instance
pub fn new(params: Params, system: System<'a, F, J, A>) -> Result<Self, StrError> {
pub fn new(params: Params, system: &'a System<F, J, A>) -> Result<Self, StrError> {
// Runge-Kutta coefficients
#[rustfmt::skip]
let (aa, bb, cc) = match params.method {
Expand Down Expand Up @@ -149,8 +149,8 @@ where

impl<'a, F, J, A> OdeSolverTrait<A> for ExplicitRungeKutta<'a, F, J, A>
where
F: Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
{
/// Enables dense output
fn enable_dense_output(&mut self) -> Result<(), StrError> {
Expand Down Expand Up @@ -385,7 +385,7 @@ mod tests {
None,
None,
);
let erk = ExplicitRungeKutta::new(params, system).unwrap();
let erk = ExplicitRungeKutta::new(params, &system).unwrap();
let nstage = erk.nstage;
assert_eq!(erk.aa.dims(), (nstage, nstage));
assert_eq!(erk.bb.dim(), nstage);
Expand Down Expand Up @@ -506,12 +506,12 @@ mod tests {

// problem
let (system, data, mut args) = Samples::kreyszig_eq6_page902();
let mut yfx = data.y_analytical.unwrap();
let yfx = data.y_analytical.unwrap();
let ndim = system.ndim;

// allocate structs
let params = Params::new(Method::MdEuler); // aka the Improved Euler in Kreyszig's book
let mut solver = ExplicitRungeKutta::new(params, system).unwrap();
let mut solver = ExplicitRungeKutta::new(params, &system).unwrap();
let mut work = Workspace::new(Method::FwEuler);

// check dense output availability
Expand Down Expand Up @@ -589,12 +589,12 @@ mod tests {

// problem
let (system, data, mut args) = Samples::kreyszig_eq6_page902();
let mut yfx = data.y_analytical.unwrap();
let yfx = data.y_analytical.unwrap();
let ndim = system.ndim;

// allocate structs
let params = Params::new(Method::Rk4); // aka the Classical RK in Kreyszig's book
let mut solver = ExplicitRungeKutta::new(params, system).unwrap();
let mut solver = ExplicitRungeKutta::new(params, &system).unwrap();
let mut work = Workspace::new(Method::FwEuler);

// check dense output availability
Expand Down Expand Up @@ -701,7 +701,7 @@ mod tests {

// allocate solver
let params = Params::new(Method::Fehlberg4);
let mut solver = ExplicitRungeKutta::new(params, system).unwrap();
let mut solver = ExplicitRungeKutta::new(params, &system).unwrap();
let mut work = Workspace::new(Method::FwEuler);

// perform one step (compute k)
Expand Down
38 changes: 20 additions & 18 deletions russell_ode/src/ode_solver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,10 @@ impl<'a, A> OdeSolver<'a, A> {
/// # Generics
///
/// See [System] for an explanation of the generic parameters.
pub fn new<F, J>(params: Params, system: System<'a, F, J, A>) -> Result<Self, StrError>
pub fn new<F, J>(params: Params, system: &'a System<F, J, A>) -> Result<Self, StrError>
where
F: 'a + Send + FnMut(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: 'a + Send + FnMut(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
F: 'a + Send + Fn(&mut Vector, f64, &Vector, &mut A) -> Result<(), StrError>,
J: 'a + Send + Fn(&mut CooMatrix, f64, &Vector, f64, &mut A) -> Result<(), StrError>,
A: 'a,
{
params.validate()?;
Expand Down Expand Up @@ -97,7 +97,7 @@ impl<'a, A> OdeSolver<'a, A> {
x0: f64,
x1: f64,
h_equal: Option<f64>,
mut output: Option<&mut Output<'a>>,
mut output: Option<&mut Output>,
args: &mut A,
) -> Result<(), StrError> {
// check data
Expand Down Expand Up @@ -284,15 +284,16 @@ mod tests {
#[test]
fn solve_with_step_output_works() {
// ODE system
let (system, data, mut args) = Samples::simple_constant();
let (system, data, mut args) = Samples::simple_equation_constant();

// output
let mut out = Output::new();
out.enable_step(&[0]).set_analytical(data.y_analytical.unwrap());
out.y_analytical = data.y_analytical;
out.enable_step(&[0]);

// params and solver
let params = Params::new(Method::FwEuler);
let mut solver = OdeSolver::new(params, system).unwrap();
let mut solver = OdeSolver::new(params, &system).unwrap();

// solve the ODE system (will run with N_EQUAL_STEPS)
let mut y = data.y0.clone();
Expand Down Expand Up @@ -328,15 +329,15 @@ mod tests {
#[test]
fn solve_with_h_equal_works() {
// ODE system
let (system, mut data, mut args) = Samples::simple_constant();
let (system, mut data, mut args) = Samples::simple_equation_constant();

// output
let mut out = Output::new();
out.enable_step(&[0]);

// params and solver
let params = Params::new(Method::FwEuler);
let mut solver = OdeSolver::new(params, system).unwrap();
let mut solver = OdeSolver::new(params, &system).unwrap();
let x1 = 1.2; // => will generate 4 steps

// capture error
Expand Down Expand Up @@ -368,16 +369,17 @@ mod tests {
#[test]
fn solve_with_variable_steps_works() {
// ODE system
let (system, mut data, mut args) = Samples::simple_constant();
let (system, mut data, mut args) = Samples::simple_equation_constant();

// output
let mut out = Output::new();
out.enable_step(&[0]).set_analytical(data.y_analytical.unwrap());
out.y_analytical = data.y_analytical;
out.enable_step(&[0]);

// params and solver
let mut params = Params::new(Method::MdEuler);
params.step.h_ini = 0.1;
let mut solver = OdeSolver::new(params, system).unwrap();
let mut solver = OdeSolver::new(params, &system).unwrap();

// solve the ODE system
solver
Expand All @@ -395,7 +397,7 @@ mod tests {
#[test]
fn solve_with_dense_output_works() {
// ODE system
let (system, data, mut args) = Samples::simple_constant();
let (system, data, mut args) = Samples::simple_equation_constant();

// output
let mut out = Output::new();
Expand All @@ -404,7 +406,7 @@ mod tests {
// params and solver
let mut params = Params::new(Method::DoPri5);
params.step.h_ini = 0.1;
let mut solver = OdeSolver::new(params, system).unwrap();
let mut solver = OdeSolver::new(params, &system).unwrap();

// solve the ODE system
let mut y = data.y0.clone();
Expand Down Expand Up @@ -434,7 +436,7 @@ mod tests {
#[test]
fn solve_handles_zero_stepsize() {
// ODE system
let (system, mut data, mut args) = Samples::simple_constant();
let (system, mut data, mut args) = Samples::simple_equation_constant();

// output
let mut out = Output::new();
Expand All @@ -443,7 +445,7 @@ mod tests {
// params and solver
let mut params = Params::new(Method::DoPri5);
params.step.h_ini = 20.0; // since the problem is linear, this stepsize will lead to a jump from x=0.0 to 1.0
let mut solver = OdeSolver::new(params, system).unwrap();
let mut solver = OdeSolver::new(params, &system).unwrap();

// solve the ODE system
solver
Expand All @@ -459,9 +461,9 @@ mod tests {

#[test]
fn solve_and_output_handle_errors() {
let (system, mut data, mut args) = Samples::simple_constant();
let (system, mut data, mut args) = Samples::simple_equation_constant();
let params = Params::new(Method::FwEuler);
let mut solver = OdeSolver::new(params, system).unwrap();
let mut solver = OdeSolver::new(params, &system).unwrap();
let mut out = Output::new();
assert_eq!(out.enable_dense(-0.1, &[0]).err(), Some("h_out must be positive"));
out.enable_dense(0.1, &[0]).unwrap();
Expand Down
Loading

0 comments on commit 377c61f

Please sign in to comment.