tfp.math.ode.DormandPrince

View source on GitHub

Dormand-Prince explicit solver for non-stiff ODEs.

Inherits From: Solver

Used in the notebooks

Used in the tutorials

Implements 5th order Runge-Kutta with adaptive step size control and dense output, using the Dormand-Prince method. Similar to the 'dopri5' method of scipy.integrate.ode and MATLAB's ode45. For details see [1]. For solver API see tfp.math.ode.Solver.

References

[1]: Shampine, L. F. (1986). Some practical runge-kutta formulas. Mathematics of Computation, 46(173), 135-150, doi:10.2307/2008219

rtol Optional float Tensor specifying an upper bound on relative error, per element of the dependent variable. The error tolerance for the next step is tol = atol + rtol * abs(state) where state is the computed state at the current step (see also atol). The next step is rejected if it incurs a local truncation error larger than tol. Default value: 1e-3.
atol Optional float Tensor specifying an upper bound on absolute error, per element of the dependent variable (see also rtol). Default value: 1e-6.
first_step_size Scalar float Tensor specifying the size of the first step. Default value: 1e-3.
safety_factor Scalar positive float Tensor. At the end of every Runge Kutta step, the solver may choose to update the step size by applying a multiplicative factor to the current step size. This factor is factor = clamp(factor_unclamped, min_step_size_factor, max_step_size_factor) where factor_unclamped = error_ratio**(-1. / (order + 1)) * safety_factor (see also min_step_size_factor and max_step_size_factor). A small (respectively, large) value for the safety factor causes the solver to take smaller (respectively, larger) step sizes. A value larger than one, though not explicitly prohibited, is discouraged. Default value: 0.9.
min_step_size_factor Scalar float Tensor (see safety_factor). Default value: 0.1.
max_step_size_factor Scalar float Tensor (see safety_factor). Default value: 10..
max_num_steps Optional scalar integer Tensor specifying the maximum number of steps allowed (including rejected steps). If unspecified, there is no upper bound on the number of steps. Default value: None.
validate_args Whether to validate input with asserts. If validate_args is False and the inputs are invalid, correct behavior is not guaranteed. Default value: False.
name Python str name prefixed to Ops created by this function. Default value: None (i.e., 'dormand_prince').

Methods

solve

View source

Solves an initial value problem.

An initial value problem consists of a system of ODEs and an initial condition:

dy/dt(t) = ode_fn(t, y(t))
y(initial_time) = initial_state

Here, t (also called time) is a scalar float Tensor and y(t) (also called the state at time t) is an N-D float or complex Tensor.

Example

The ODE dy/dt(t) = dot(A, y(t)) is solved below.

t_init, t0, t1 = 0., 0.5, 1.
y_init = tf.constant([1., 1.], dtype=tf.float64)
A = tf.constant([[-1., -2.], [-3., -4.]], dtype=tf.float64)

def ode_fn(t, y):
  return tf.linalg.matvec(A, y)

results = tfp.math.ode.BDF().solve(ode_fn, t_init, y_init,
                                   solution_times=[t0, t1])
y0 = results.states[0]  # == dot(matrix_exp(A * t0), y_init)
y1 = results.states[1]  # == dot(matrix_exp(A * t1), y_init)

Using instead solution_times=tfp.math.ode.ChosenBySolver(final_time=1.) yields the state at various times between t_init and final_time chosen automatically by the solver. In this case, results.states[i] is the state at time results.times[i].

Gradient

The gradient of the result is computed using the adjoint sensitivity method described in [Chen et al. (2018)][1].

grad = tf.gradients(y1, y0) # == dot(e, J)
# J is the Jacobian of y1 with respect to y0. In this case, J = exp(A * t1).
# e = [1, ..., 1] is the row vector of ones.

References

[1]: Chen, Tian Qi, et al. "Neural ordinary differential equations." Advances in Neural Information Processing Systems. 2018.

Args
ode_fn Function of the form ode_fn(t, y). The input t is a scalar float Tensor. The input y and output are both Tensors with the same shape and dtype as initial_state.
initial_time Scalar float Tensor specifying the initial time.
initial_state N-D float or complex Tensor specifying the initial state. The dtype of initial_state must be complex for problems with complex-valued states (even if the initial state is real).
solution_times 1-D float Tensor specifying a list of times. The solver stores the computed state at each of these times in the returned Results object. Must satisfy initial_time <= solution_times[0] and solution_times[i] < solution_times[i+1]. Alternatively, the user can pass tfp.math.ode.ChosenBySolver(final_time) where final_time is a scalar float Tensor satisfying initial_time < final_time. Doing so requests that the solver automatically choose suitable times up to and including final_time at which to store the computed state.
jacobian_fn Optional function of the form jacobian_fn(t, y). The input t is a scalar float Tensor. The input y has the same shape and dtype as initial_state. The output is a 2N-D Tensor whose shape is initial_state.shape + initial_state.shape and whose dtype is the same as initial_state. In particular, the (i1, ..., iN, j1, ..., jN)-th entry of jacobian_fn(t, y) is the derivative of the (i1, ..., iN)-th entry of ode_fn(t, y) with respect to the (j1, ..., jN)-th entry of y. If this argument is left unspecified, the solver automatically computes the Jacobian if and when it is needed. Default value: None.
jacobian_sparsity Optional 2N-D boolean Tensor whose shape is initial_state.shape + initial_state.shape specifying the sparsity pattern of the Jacobian. This argument is ignored if jacobian_fn is specified. Default value: None.
batch_ndims Optional nonnegative integer. When specified, the first batch_ndims dimensions of initial_state are batch dimensions. Default value: None.
previous_solver_internal_state Optional solver-specific argument used to warm-start this invocation of solve. Default value: None.

Returns
Object of type Results.

Class Variables

  • ODE_FN_EVALS_PER_STEP = 6
  • ORDER = 5