ndarray/doc/ndarray_for_numpy_users/rk_step.rs
1//! Example of translating `rk_step` function from SciPy.
2//!
3//! This snippet is a [selection of lines from
4//! `rk.py`](https://github.com/scipy/scipy/blob/v1.0.0/scipy/integrate/_ivp/rk.py#L2-L78).
5//! See the [license for this snippet](#scipy-license).
6//!
7//! ```python
8//! import numpy as np
9//!
10//! def rk_step(fun, t, y, f, h, A, B, C, E, K):
11//! """Perform a single Runge-Kutta step.
12//! This function computes a prediction of an explicit Runge-Kutta method and
13//! also estimates the error of a less accurate method.
14//! Notation for Butcher tableau is as in [1]_.
15//! Parameters
16//! ----------
17//! fun : callable
18//! Right-hand side of the system.
19//! t : float
20//! Current time.
21//! y : ndarray, shape (n,)
22//! Current state.
23//! f : ndarray, shape (n,)
24//! Current value of the derivative, i.e. ``fun(x, y)``.
25//! h : float
26//! Step to use.
27//! A : list of ndarray, length n_stages - 1
28//! Coefficients for combining previous RK stages to compute the next
29//! stage. For explicit methods the coefficients above the main diagonal
30//! are zeros, so `A` is stored as a list of arrays of increasing lengths.
31//! The first stage is always just `f`, thus no coefficients for it
32//! are required.
33//! B : ndarray, shape (n_stages,)
34//! Coefficients for combining RK stages for computing the final
35//! prediction.
36//! C : ndarray, shape (n_stages - 1,)
37//! Coefficients for incrementing time for consecutive RK stages.
38//! The value for the first stage is always zero, thus it is not stored.
39//! E : ndarray, shape (n_stages + 1,)
40//! Coefficients for estimating the error of a less accurate method. They
41//! are computed as the difference between b's in an extended tableau.
42//! K : ndarray, shape (n_stages + 1, n)
43//! Storage array for putting RK stages here. Stages are stored in rows.
44//! Returns
45//! -------
46//! y_new : ndarray, shape (n,)
47//! Solution at t + h computed with a higher accuracy.
48//! f_new : ndarray, shape (n,)
49//! Derivative ``fun(t + h, y_new)``.
50//! error : ndarray, shape (n,)
51//! Error estimate of a less accurate method.
52//! References
53//! ----------
54//! .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential
55//! Equations I: Nonstiff Problems", Sec. II.4.
56//! """
57//! K[0] = f
58//! for s, (a, c) in enumerate(zip(A, C)):
59//! dy = np.dot(K[:s + 1].T, a) * h
60//! K[s + 1] = fun(t + c * h, y + dy)
61//!
62//! y_new = y + h * np.dot(K[:-1].T, B)
63//! f_new = fun(t + h, y_new)
64//!
65//! K[-1] = f_new
66//! error = np.dot(K.T, E) * h
67//!
68//! return y_new, f_new, error
69//! ```
70//!
71//! A direct translation to `ndarray` looks like this:
72//!
73//! ```
74//! use ndarray::prelude::*;
75//!
76//! fn rk_step<F>(
77//! mut fun: F,
78//! t: f64,
79//! y: ArrayView1<f64>,
80//! f: ArrayView1<f64>,
81//! h: f64,
82//! a: &[ArrayView1<f64>],
83//! b: ArrayView1<f64>,
84//! c: ArrayView1<f64>,
85//! e: ArrayView1<f64>,
86//! mut k: ArrayViewMut2<f64>,
87//! ) -> (Array1<f64>, Array1<f64>, Array1<f64>)
88//! where
89//! F: FnMut(f64, ArrayView1<f64>) -> Array1<f64>,
90//! {
91//! k.slice_mut(s![0, ..]).assign(&f);
92//! for (s, (a, c)) in a.iter().zip(c).enumerate() {
93//! let dy = k.slice(s![..s + 1, ..]).t().dot(a) * h;
94//! k.slice_mut(s![s + 1, ..])
95//! .assign(&(fun(t + c * h, (&y + &dy).view())));
96//! }
97//!
98//! let y_new = &y + &(h * k.slice(s![..-1, ..]).t().dot(&b));
99//! let f_new = fun(t + h, y_new.view());
100//!
101//! k.slice_mut(s![-1, ..]).assign(&f_new);
102//! let error = k.t().dot(&e) * h;
103//!
104//! (y_new, f_new, error)
105//! }
106//! #
107//! # fn main() { let _ = rk_step::<fn(_, ArrayView1<'_, _>) -> _>; }
108//! ```
109//!
110//! It's possible to improve the efficiency by doing the following:
111//!
112//! * Observe that `dy` is a temporary allocation. It's possible to allow the
113//! add operation to take ownership of `dy` to eliminate an extra allocation
114//! for the result of the addition. A similar situation occurs when computing
115//! `y_new`. See the comments in the example below.
116//!
117//! * Require the `fun` closure to mutate an existing view instead of
118//! allocating a new array for the result.
119//!
120//! * Don't return a newly allocated `f_new` array. If the caller wants this
121//! information, they can get it from the last row of `k`.
122//!
123//! * Use [`c.mul_add(h, t)`](f64::mul_add) instead of `t + c * h`. This is
124//! faster and reduces the floating-point error. It might also be beneficial
125//! to use [`.scaled_add()`](crate::ArrayRef::scaled_add) or a combination of
126//! [`azip!()`] and [`.mul_add()`](f64::mul_add) on the arrays in
127//! some places, but that's not demonstrated in the example below.
128//!
129//! ```
130//! use ndarray::prelude::*;
131//!
132//! fn rk_step<F>(
133//! mut fun: F,
134//! t: f64,
135//! y: ArrayView1<f64>,
136//! f: ArrayView1<f64>,
137//! h: f64,
138//! a: &[ArrayView1<f64>],
139//! b: ArrayView1<f64>,
140//! c: ArrayView1<f64>,
141//! e: ArrayView1<f64>,
142//! mut k: ArrayViewMut2<f64>,
143//! ) -> (Array1<f64>, Array1<f64>)
144//! where
145//! F: FnMut(f64, ArrayView1<f64>, ArrayViewMut1<f64>),
146//! {
147//! k.slice_mut(s![0, ..]).assign(&f);
148//! for (s, (a, c)) in a.iter().zip(c).enumerate() {
149//! let dy = k.slice(s![..s + 1, ..]).t().dot(a) * h;
150//! // Note that `dy` comes before `&y` in `dy + &y` in order to reuse the
151//! // `dy` allocation. (The addition operator will take ownership of `dy`
152//! // and assign the result to it instead of allocating a new array for the
153//! // result.) In contrast, you could use `&y + &dy`, but that would perform
154//! // an unnecessary memory allocation for the result, like NumPy does.
155//! fun(c.mul_add(h, t), (dy + &y).view(), k.slice_mut(s![s + 1, ..]));
156//! }
157//! // Similar case here — moving `&y` to the right hand side allows the addition
158//! // to reuse the allocated array on the left hand side.
159//! let y_new = h * k.slice(s![..-1, ..]).t().dot(&b) + &y;
160//! // Mutate the last row of `k` in-place instead of allocating a new array.
161//! fun(t + h, y_new.view(), k.slice_mut(s![-1, ..]));
162//!
163//! let error = k.t().dot(&e) * h;
164//!
165//! (y_new, error)
166//! }
167//! #
168//! # fn main() { let _ = rk_step::<fn(_, ArrayView1<'_, f64>, ArrayViewMut1<'_, f64>)>; }
169//! ```
170//!
171//! [`.scaled_add()`]: crate::ArrayRef::scaled_add
172//! [`azip!()`]: crate::azip!
173//!
174//! ### SciPy license
175//!
176//! ```text
177//! Copyright (c) 2001, 2002 Enthought, Inc.
178//! All rights reserved.
179//!
180//! Copyright (c) 2003-2017 SciPy Developers.
181//! All rights reserved.
182//!
183//! Redistribution and use in source and binary forms, with or without
184//! modification, are permitted provided that the following conditions are met:
185//!
186//! a. Redistributions of source code must retain the above copyright notice,
187//! this list of conditions and the following disclaimer.
188//! b. Redistributions in binary form must reproduce the above copyright
189//! notice, this list of conditions and the following disclaimer in the
190//! documentation and/or other materials provided with the distribution.
191//! c. Neither the name of Enthought nor the names of the SciPy Developers
192//! may be used to endorse or promote products derived from this software
193//! without specific prior written permission.
194//!
195//!
196//! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
197//! AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
198//! IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
199//! ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS
200//! BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
201//! OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
202//! SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
203//! INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
204//! CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
205//! ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
206//! THE POSSIBILITY OF SUCH DAMAGE.
207//! ```