275 lines
9.3 KiB
Python
275 lines
9.3 KiB
Python
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
def check_arguments(fun, y0, support_complex):
|
||
|
|
"""Helper function for checking arguments common to all solvers."""
|
||
|
|
y0 = np.asarray(y0)
|
||
|
|
if np.issubdtype(y0.dtype, np.complexfloating):
|
||
|
|
if not support_complex:
|
||
|
|
raise ValueError("`y0` is complex, but the chosen solver does "
|
||
|
|
"not support integration in a complex domain.")
|
||
|
|
dtype = complex
|
||
|
|
else:
|
||
|
|
dtype = float
|
||
|
|
y0 = y0.astype(dtype, copy=False)
|
||
|
|
|
||
|
|
if y0.ndim != 1:
|
||
|
|
raise ValueError("`y0` must be 1-dimensional.")
|
||
|
|
|
||
|
|
def fun_wrapped(t, y):
|
||
|
|
return np.asarray(fun(t, y), dtype=dtype)
|
||
|
|
|
||
|
|
return fun_wrapped, y0
|
||
|
|
|
||
|
|
|
||
|
|
class OdeSolver:
|
||
|
|
"""Base class for ODE solvers.
|
||
|
|
|
||
|
|
In order to implement a new solver you need to follow the guidelines:
|
||
|
|
|
||
|
|
1. A constructor must accept parameters presented in the base class
|
||
|
|
(listed below) along with any other parameters specific to a solver.
|
||
|
|
2. A constructor must accept arbitrary extraneous arguments
|
||
|
|
``**extraneous``, but warn that these arguments are irrelevant
|
||
|
|
using `common.warn_extraneous` function. Do not pass these
|
||
|
|
arguments to the base class.
|
||
|
|
3. A solver must implement a private method `_step_impl(self)` which
|
||
|
|
propagates a solver one step further. It must return tuple
|
||
|
|
``(success, message)``, where ``success`` is a boolean indicating
|
||
|
|
whether a step was successful, and ``message`` is a string
|
||
|
|
containing description of a failure if a step failed or None
|
||
|
|
otherwise.
|
||
|
|
4. A solver must implement a private method `_dense_output_impl(self)`,
|
||
|
|
which returns a `DenseOutput` object covering the last successful
|
||
|
|
step.
|
||
|
|
5. A solver must have attributes listed below in Attributes section.
|
||
|
|
Note that ``t_old`` and ``step_size`` are updated automatically.
|
||
|
|
6. Use `fun(self, t, y)` method for the system rhs evaluation, this
|
||
|
|
way the number of function evaluations (`nfev`) will be tracked
|
||
|
|
automatically.
|
||
|
|
7. For convenience, a base class provides `fun_single(self, t, y)` and
|
||
|
|
`fun_vectorized(self, t, y)` for evaluating the rhs in
|
||
|
|
non-vectorized and vectorized fashions respectively (regardless of
|
||
|
|
how `fun` from the constructor is implemented). These calls don't
|
||
|
|
increment `nfev`.
|
||
|
|
8. If a solver uses a Jacobian matrix and LU decompositions, it should
|
||
|
|
track the number of Jacobian evaluations (`njev`) and the number of
|
||
|
|
LU decompositions (`nlu`).
|
||
|
|
9. By convention, the function evaluations used to compute a finite
|
||
|
|
difference approximation of the Jacobian should not be counted in
|
||
|
|
`nfev`, thus use `fun_single(self, t, y)` or
|
||
|
|
`fun_vectorized(self, t, y)` when computing a finite difference
|
||
|
|
approximation of the Jacobian.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
fun : callable
|
||
|
|
Right-hand side of the system. The calling signature is ``fun(t, y)``.
|
||
|
|
Here ``t`` is a scalar and there are two options for ndarray ``y``.
|
||
|
|
It can either have shape (n,), then ``fun`` must return array_like with
|
||
|
|
shape (n,). Or, alternatively, it can have shape (n, n_points), then
|
||
|
|
``fun`` must return array_like with shape (n, n_points) (each column
|
||
|
|
corresponds to a single column in ``y``). The choice between the two
|
||
|
|
options is determined by `vectorized` argument (see below).
|
||
|
|
t0 : float
|
||
|
|
Initial time.
|
||
|
|
y0 : array_like, shape (n,)
|
||
|
|
Initial state.
|
||
|
|
t_bound : float
|
||
|
|
Boundary time --- the integration won't continue beyond it. It also
|
||
|
|
determines the direction of the integration.
|
||
|
|
vectorized : bool
|
||
|
|
Whether `fun` is implemented in a vectorized fashion.
|
||
|
|
support_complex : bool, optional
|
||
|
|
Whether integration in a complex domain should be supported.
|
||
|
|
Generally determined by a derived solver class capabilities.
|
||
|
|
Default is False.
|
||
|
|
|
||
|
|
Attributes
|
||
|
|
----------
|
||
|
|
n : int
|
||
|
|
Number of equations.
|
||
|
|
status : string
|
||
|
|
Current status of the solver: 'running', 'finished' or 'failed'.
|
||
|
|
t_bound : float
|
||
|
|
Boundary time.
|
||
|
|
direction : float
|
||
|
|
Integration direction: +1 or -1.
|
||
|
|
t : float
|
||
|
|
Current time.
|
||
|
|
y : ndarray
|
||
|
|
Current state.
|
||
|
|
t_old : float
|
||
|
|
Previous time. None if no steps were made yet.
|
||
|
|
step_size : float
|
||
|
|
Size of the last successful step. None if no steps were made yet.
|
||
|
|
nfev : int
|
||
|
|
Number of the system's rhs evaluations.
|
||
|
|
njev : int
|
||
|
|
Number of the Jacobian evaluations.
|
||
|
|
nlu : int
|
||
|
|
Number of LU decompositions.
|
||
|
|
"""
|
||
|
|
TOO_SMALL_STEP = "Required step size is less than spacing between numbers."
|
||
|
|
|
||
|
|
def __init__(self, fun, t0, y0, t_bound, vectorized,
|
||
|
|
support_complex=False):
|
||
|
|
self.t_old = None
|
||
|
|
self.t = t0
|
||
|
|
self._fun, self.y = check_arguments(fun, y0, support_complex)
|
||
|
|
self.t_bound = t_bound
|
||
|
|
self.vectorized = vectorized
|
||
|
|
|
||
|
|
if vectorized:
|
||
|
|
def fun_single(t, y):
|
||
|
|
return self._fun(t, y[:, None]).ravel()
|
||
|
|
fun_vectorized = self._fun
|
||
|
|
else:
|
||
|
|
fun_single = self._fun
|
||
|
|
|
||
|
|
def fun_vectorized(t, y):
|
||
|
|
f = np.empty_like(y)
|
||
|
|
for i, yi in enumerate(y.T):
|
||
|
|
f[:, i] = self._fun(t, yi)
|
||
|
|
return f
|
||
|
|
|
||
|
|
def fun(t, y):
|
||
|
|
self.nfev += 1
|
||
|
|
return self.fun_single(t, y)
|
||
|
|
|
||
|
|
self.fun = fun
|
||
|
|
self.fun_single = fun_single
|
||
|
|
self.fun_vectorized = fun_vectorized
|
||
|
|
|
||
|
|
self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1
|
||
|
|
self.n = self.y.size
|
||
|
|
self.status = 'running'
|
||
|
|
|
||
|
|
self.nfev = 0
|
||
|
|
self.njev = 0
|
||
|
|
self.nlu = 0
|
||
|
|
|
||
|
|
@property
|
||
|
|
def step_size(self):
|
||
|
|
if self.t_old is None:
|
||
|
|
return None
|
||
|
|
else:
|
||
|
|
return np.abs(self.t - self.t_old)
|
||
|
|
|
||
|
|
def step(self):
|
||
|
|
"""Perform one integration step.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
message : string or None
|
||
|
|
Report from the solver. Typically a reason for a failure if
|
||
|
|
`self.status` is 'failed' after the step was taken or None
|
||
|
|
otherwise.
|
||
|
|
"""
|
||
|
|
if self.status != 'running':
|
||
|
|
raise RuntimeError("Attempt to step on a failed or finished "
|
||
|
|
"solver.")
|
||
|
|
|
||
|
|
if self.n == 0 or self.t == self.t_bound:
|
||
|
|
# Handle corner cases of empty solver or no integration.
|
||
|
|
self.t_old = self.t
|
||
|
|
self.t = self.t_bound
|
||
|
|
message = None
|
||
|
|
self.status = 'finished'
|
||
|
|
else:
|
||
|
|
t = self.t
|
||
|
|
success, message = self._step_impl()
|
||
|
|
|
||
|
|
if not success:
|
||
|
|
self.status = 'failed'
|
||
|
|
else:
|
||
|
|
self.t_old = t
|
||
|
|
if self.direction * (self.t - self.t_bound) >= 0:
|
||
|
|
self.status = 'finished'
|
||
|
|
|
||
|
|
return message
|
||
|
|
|
||
|
|
def dense_output(self):
|
||
|
|
"""Compute a local interpolant over the last successful step.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
sol : `DenseOutput`
|
||
|
|
Local interpolant over the last successful step.
|
||
|
|
"""
|
||
|
|
if self.t_old is None:
|
||
|
|
raise RuntimeError("Dense output is available after a successful "
|
||
|
|
"step was made.")
|
||
|
|
|
||
|
|
if self.n == 0 or self.t == self.t_old:
|
||
|
|
# Handle corner cases of empty solver and no integration.
|
||
|
|
return ConstantDenseOutput(self.t_old, self.t, self.y)
|
||
|
|
else:
|
||
|
|
return self._dense_output_impl()
|
||
|
|
|
||
|
|
def _step_impl(self):
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
def _dense_output_impl(self):
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
|
||
|
|
class DenseOutput:
|
||
|
|
"""Base class for local interpolant over step made by an ODE solver.
|
||
|
|
|
||
|
|
It interpolates between `t_min` and `t_max` (see Attributes below).
|
||
|
|
Evaluation outside this interval is not forbidden, but the accuracy is not
|
||
|
|
guaranteed.
|
||
|
|
|
||
|
|
Attributes
|
||
|
|
----------
|
||
|
|
t_min, t_max : float
|
||
|
|
Time range of the interpolation.
|
||
|
|
"""
|
||
|
|
def __init__(self, t_old, t):
|
||
|
|
self.t_old = t_old
|
||
|
|
self.t = t
|
||
|
|
self.t_min = min(t, t_old)
|
||
|
|
self.t_max = max(t, t_old)
|
||
|
|
|
||
|
|
def __call__(self, t):
|
||
|
|
"""Evaluate the interpolant.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
t : float or array_like with shape (n_points,)
|
||
|
|
Points to evaluate the solution at.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
y : ndarray, shape (n,) or (n, n_points)
|
||
|
|
Computed values. Shape depends on whether `t` was a scalar or a
|
||
|
|
1-D array.
|
||
|
|
"""
|
||
|
|
t = np.asarray(t)
|
||
|
|
if t.ndim > 1:
|
||
|
|
raise ValueError("`t` must be a float or a 1-D array.")
|
||
|
|
return self._call_impl(t)
|
||
|
|
|
||
|
|
def _call_impl(self, t):
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
|
||
|
|
class ConstantDenseOutput(DenseOutput):
|
||
|
|
"""Constant value interpolator.
|
||
|
|
|
||
|
|
This class used for degenerate integration cases: equal integration limits
|
||
|
|
or a system with 0 equations.
|
||
|
|
"""
|
||
|
|
def __init__(self, t_old, t, value):
|
||
|
|
super().__init__(t_old, t)
|
||
|
|
self.value = value
|
||
|
|
|
||
|
|
def _call_impl(self, t):
|
||
|
|
if t.ndim == 0:
|
||
|
|
return self.value
|
||
|
|
else:
|
||
|
|
ret = np.empty((self.value.shape[0], t.shape[0]))
|
||
|
|
ret[:] = self.value[:, None]
|
||
|
|
return ret
|