123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274 |
- 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
|