Parallel-in-Time (DEER)
Here we will illustrate the application of diffsol to a parallel-in-time method for solving a nonlinear ODE, the DEER algorithm introduced in Lim et. al.. The example will use the forced the forced Duffing oscillator, a nonlinear second order ODE, and utilize the DEER algorithm to solve it in parallel across multiple time chunks using CPU threads and the rayon crate.
First we will define a few types, we will use the nalgebra backend and matrix/vector types, and a f64 scalar type.
type M = NalgebraMat<f64>;
type V = <M as MatrixCommon>::V;
type T = <M as MatrixCommon>::T;
type C = <M as MatrixCommon>::C;The problem
The example used is the forced Duffing oscillator, a nonlinear second order ODE given by the equation
\[ \ddot{x} + \delta \dot{x} + x + x^3 = \gamma \cos(\omega t), \]
which, writing \(\mathbf{y} = (x, v)^T\) with \(v = \dot{x}\), is the first order system
\[ \begin{aligned} \frac{dx}{dt} &= v, \\ \frac{dv}{dt} &= -\delta v - x - x^3 + \gamma \cos(\omega t), \end{aligned} \]
or \(\dot{\mathbf{y}} = \mathbf{f}(\mathbf{y}, t)\), with \(\delta = 0.5\), \(\gamma = 0.5\), \(\omega = 1\) and \(\mathbf{y}(0) = (1, 0)^T\). The Jacobian of the right-hand side is
\[ \frac{\partial \mathbf{f}}{\partial \mathbf{y}} = \begin{pmatrix} 0 & 1 \\ -1 - 3x^2 & -\delta \end{pmatrix}. \]
The parameters are the initial state, \(\mathbf{p} = \mathbf{y}(t_0)\), so the right-hand side does not depend on them and the initial condition is the identity map:
\[ \frac{\partial \mathbf{f}}{\partial \mathbf{p}} = \begin{pmatrix} 0 & 0 \\ 0 & 0 \end{pmatrix}, \quad \frac{\partial \mathbf{y}(t_0)}{\partial \mathbf{p}} = \begin{pmatrix} 1 & 0 \\ 0 & 1 \end{pmatrix}. \]
const DELTA: f64 = 0.5;
const GAMMA: f64 = 0.5;
const OMEGA: f64 = 1.0;
const Y0: [f64; 2] = [1.0, 0.0];
fn problem(
rtol: f64,
atol: f64,
) -> OdeSolverProblem<impl OdeEquationsImplicitSens<M = M, V = V, T = T, C = C>> {
OdeBuilder::<M>::new()
.rtol(rtol)
.atol([atol])
.p(Y0)
.rhs_sens_implicit(
|y, _p, t, dy| {
dy[0] = y[1];
dy[1] = -DELTA * y[1] - y[0] - y[0] * y[0] * y[0] + GAMMA * (OMEGA * t).cos();
},
|y, _p, _t, v, jv| {
jv[0] = v[1];
jv[1] = -(1.0 + 3.0 * y[0] * y[0]) * v[0] - DELTA * v[1];
},
|_y, _p, _t, _v, jv| jv.fill(0.0),
)
.init_sens(
|p, _t, y| y.copy_from_slice(p),
|_p, _t, v, y| y.copy_from_slice(v),
2,
)
.build()
.unwrap()
}The DEER iteration
The time interval \([0, T]\) is split into \(N\) chunks, and each chunk is integrated independently, in parallel, from a guess of the state at its start time:
The states at the chunk boundaries \(\mathbf{y}_k \approx \mathbf{y}(t_k)\) satisfy the nonlinear recurrence
\[ \mathbf{y}_{k+1} = \boldsymbol{\phi}_k(\mathbf{y}_k), \quad k = 0, \dots, N - 1, \quad \mathbf{y}_0 = \mathbf{y}(0). \]
Linearising about the current iterate \(\mathbf{y}_k^{(i)}\) gives Newton's method for all boundary states at once,
\[ \mathbf{y}_{k+1}^{(i+1)} = \boldsymbol{\phi}_k(\mathbf{y}_k^{(i)}) + J_k \left( \mathbf{y}_k^{(i+1)} - \mathbf{y}_k^{(i)} \right), \quad J_k = \frac{\partial \boldsymbol{\phi}_k}{\partial \mathbf{y}} \left( \mathbf{y}_k^{(i)} \right), \]
which is a linear (affine) recurrence
\[ \mathbf{y}_{k+1}^{(i+1)} = J_k \mathbf{y}_k^{(i+1)} + \mathbf{b}_k, \quad \mathbf{b}_k = \boldsymbol{\phi}_k(\mathbf{y}_k^{(i)}) - J_k \mathbf{y}_k^{(i)}. \]
Each step of the recurrence is an affine map \(\mathbf{y} \mapsto A \mathbf{y} + \mathbf{b}\), fully described by the pair \((A, \mathbf{b})\). Applying step 0 and then step 1 gives
\[ \mathbf{y}_2 = J_1 (J_0 \mathbf{y}_0 + \mathbf{b}_0) + \mathbf{b}_1 = (J_1 J_0) \mathbf{y}_0 + (J_1 \mathbf{b}_0 + \mathbf{b}_1), \]
which is again an affine map. Writing this composition as an operator on pairs (read right to left, so the right-hand map is applied first),
\[ (A_2, \mathbf{b}_2) \bullet (A_1, \mathbf{b}_1) = (A_2 A_1, A_2 \mathbf{b}_1 + \mathbf{b}_2). \]
Like any composition of functions, this operator is associative, \((f \bullet g) \bullet h = f \bullet (g \bullet h)\), and it has the identity \((I, \mathbf{0})\), the map that leaves \(\mathbf{y}\) unchanged. Associativity means the steps can be grouped in any order, so different threads can combine different groups of steps at the same time. The prefix \((P_k, \mathbf{c}_k)\) is the single affine map equal to applying steps \(0, 1, \dots, k\) in order,
\[ (P_k, \mathbf{c}_k) = (J_k, \mathbf{b}_k) \bullet (J_{k-1}, \mathbf{b}_{k-1}) \bullet \dots \bullet (J_0, \mathbf{b}_0), \]
which takes \(\mathbf{y}_0\) straight to \(\mathbf{y}_{k+1}\). A sequential loop would build each prefix from the previous one, but a parallel scan computes all \(N\) prefixes in \(O(\log N)\) rounds of combining (given enough threads). Each new boundary state then depends only on \(\mathbf{y}_0\), so they can all be evaluated in parallel:
\[ \mathbf{y}_{k+1}^{(i+1)} = P_k \mathbf{y}_0 + \mathbf{c}_k. \]
The iteration stops when
\[ \max_k \left\Vert \mathbf{y}_k^{(i+1)} - \mathbf{y}_k^{(i)} \right\Vert_\infty < \text{tol}. \]
const RTOL: f64 = 1e-10;
const ATOL: f64 = 1e-12;
const JAC_TOL: f64 = 1e-4;
const MAX_ITERS: usize = 50;
const NEWTON_TOL: f64 = 1e-6;
fn deer() -> Vec<Vec<Vector2<f64>>> {
let y0 = Vector2::from(Y0);
let mut x = coarse();
let mut iterates = vec![x.clone()];
for _ in 0..MAX_ITERS {
let maps: Vec<(Matrix2<f64>, Vector2<f64>)> = x[..N_CHUNKS]
.par_iter()
.enumerate()
.map_init(
|| (problem(RTOL, ATOL), problem(JAC_TOL, JAC_TOL)),
|problems, (k, x_k)| {
let (f_k, j_k) = propagate(problems, k, x_k);
(j_k, f_k - j_k * x_k)
},
)
.collect();
let prefixes: Vec<(Matrix2<f64>, Vector2<f64>)> = maps
.into_par_iter()
.scan(
|(a1, b1), (a2, b2)| (a2 * a1, a2 * b1 + b2),
(Matrix2::identity(), Vector2::zeros()),
)
.collect();
let mut x_new = vec![y0];
x_new.par_extend(prefixes.par_iter().map(|(a, b)| a * y0 + b));
let delta = max_error(&x, &x_new);
x = x_new;
iterates.push(x.clone());
if delta < NEWTON_TOL {
break;
}
}
iterates
}Solving a single chunk
Let \(\boldsymbol{\phi}_k(\mathbf{y})\) be the flow map of chunk \(k\), the solution at \(t_{k+1}\) of
\[ \dot{\mathbf{y}} = \mathbf{f}(\mathbf{y}, t), \quad \mathbf{y}(t_k) = \mathbf{y}. \]
Its Jacobian \(J_k = \partial \boldsymbol{\phi}_k / \partial \mathbf{y}\) is given by the forward sensitivities \(S(t) = \partial \mathbf{y}(t) / \partial \mathbf{y}(t_k)\), which satisfy
\[ \dot{S} = \frac{\partial \mathbf{f}}{\partial \mathbf{y}} S, \quad S(t_k) = I, \quad J_k = S(t_{k+1}). \]
Newton's method still converges to the solution of the nonlinear recurrence if the \(J_k\) are only approximate, so we can use a much cheaper solve with loose tolerances and forward sensitivities to compute \(J_k\).
const T_FINAL: f64 = 3000.0;
const N_CHUNKS: usize = 96;
const DT_CHUNK: f64 = T_FINAL / N_CHUNKS as f64;
fn propagate(
problems: &mut (
OdeSolverProblem<impl OdeEquationsImplicitSens<M = M, V = V, T = T, C = C>>,
OdeSolverProblem<impl OdeEquationsImplicitSens<M = M, V = V, T = T, C = C>>,
),
k: usize,
y_start: &Vector2<f64>,
) -> (Vector2<f64>, Matrix2<f64>) {
let (accurate, loose) = problems;
let (t0, t1) = (k as f64 * DT_CHUNK, (k + 1) as f64 * DT_CHUNK);
let p = V::from_vec(vec![y_start[0], y_start[1]], *accurate.eqn.context());
accurate.eqn_mut().set_params(&p);
accurate.t0 = t0;
let mut solver = accurate.tsit45().unwrap();
solver.set_stop_time(t1).unwrap();
while !matches!(solver.step().unwrap(), OdeSolverStopReason::TstopReached) {}
let y = Vector2::from_fn(|i, _| solver.state().y.get_index(i));
loose.eqn_mut().set_params(&p);
loose.t0 = t0;
let mut solver = loose.tsit45_sens().unwrap();
solver.set_stop_time(t1).unwrap();
while !matches!(solver.step().unwrap(), OdeSolverStopReason::TstopReached) {}
let s = solver.state().s;
let jac = Matrix2::from_fn(|i, j| s.get_batch(j).get_index(i));
(y, jac)
}Initial guess
The initial guess \(\mathbf{y}_k^{(0)}\), \(k = 0, \dots, N\) also only needs to be approximate, so we use a cheap sequential solve of \(\dot{\mathbf{y}} = \mathbf{f}(\mathbf{y}, t)\) with loose tolerances, evaluated at the chunk boundaries \(t_k\).
const COARSE_TOL: f64 = 1e-3;
fn coarse() -> Vec<Vector2<f64>> {
let t: Vec<f64> = (0..=N_CHUNKS).map(|k| k as f64 * DT_CHUNK).collect();
let (ys, _) = problem(COARSE_TOL, COARSE_TOL)
.tsit45()
.unwrap()
.solve_dense(&t)
.unwrap();
(0..t.len())
.map(|k| Vector2::new(ys.column(k)[0], ys.column(k)[1]))
.collect()
}Convergence
With a reasonable initial guess the Newton iteration converges in a few iterations, and the final solution is accurate to the tolerance of the chunk solves.
Thread Scaling
We do not expect perfect linear scaling with the number of threads, since:
- the parallel scan requires \(O(\log N)\) rounds of combining
- the initial coarse solve is serial
- each solve chuck is an adaptive solve, so slower solves will hold up the faster ones
Here we see a maximum speedup of just over 5x with 25-30 threads and a total of 96 chunks. Some possible improvements that we could make to this example are:
- We are not getting much advantage out of the parallel scan, since its only 96 compositions of 2x2 matrices. A more expensive composition (more states or more chunks) could help here.
- The coarse solve is serial so is limiting the speedup, could investigate an even looser tolerance, or if the solve was more expensive (see first point), could try removing this rely on more newton iterations.
- We solve every chunk at each newton iteration, but the chunks that have already converged could be skipped in later iterations (e.g. the first chunk is always exact)
- Each chunk solve creates a new solver, but we could reuse the solvers between iterations and reset them to the new initial condition, which would save some setup time.
- Could reuse the first iteration's jacobians, since these only need to be approximate.