Commit e353b7dc authored by Romain Casati's avatar Romain Casati
Browse files

Adding internal module _ivp to support scipy.integrate.odeint and scipy.integrate.solve_ivp.

parent 0d179278
......@@ -140,6 +140,80 @@ def hack_PIL():
ImageShow.register(basthonviewer)
def hack_scipy():
"""
Add support for:
* scipy.integrate.solve_ivp from the Basthon's ivp module
* scipy.integrate.odeint (wrapper to solve_ivp)
"""
import _ivp
import scipy.integrate
import numpy as np
import warnings
# add support to solve_ivp
assert not hasattr(scipy.integrate, "solve_ivp")
scipy.integrate.solve_ivp = _ivp.solve_ivp
def odeint(func, y0, t, args=(),
Dfun=None,
tfirst=False,
method='RK45',
**kwargs):
# conversion table for options names
options = {'first_step': 'h0',
'min_step': 'hmin',
'max_step': 'hmax',
'rtol': 'rtol',
'atol': 'atol',
'lband': 'ml',
'uband': 'mu'}
# copying options
options = {k: kwargs.pop(v) for k, v in options.items()
if v in kwargs}
# is there any that left in kwargs?
if kwargs:
warnings.warn(f"Parameters {','.join(kwargs.keys())} are not supported by this version of scipy.integrate.odeint.")
dt = np.diff(t)
if not((dt >= 0).all() or (dt <= 0).all()):
raise ValueError("The values in t must be monotonically increasing "
"or monotonically decreasing; repeated values are "
"allowed.")
y0 = np.atleast_1d(y0)
fun = func
jac = Dfun
# support for tfirst since solve_ivp wants tfirst set to True
if not tfirst:
fun = lambda t, y, *a: func(y, t, *a)
if Dfun is not None:
jac = lambda t, y, *a: Dfun(y, t, *a)
if Dfun is not None:
options['jac'] = jac
# default args
if args == ():
args = None
# effective solve
sol = scipy.integrate.solve_ivp(
fun=fun,
t_span=(min(t), max(t)),
t_eval=t,
y0=y0,
args=args,
method=method,
**options)
return sol.y.T
scipy.integrate._odeint = scipy.integrate.odeint
scipy.integrate.odeint = odeint
scipy.integrate.odeint.__doc__ = scipy.integrate._odeint.__doc__
def hack(modules):
""" Hack a list of modules. """
if isinstance(modules, str):
......
......@@ -60,6 +60,12 @@ _internal_pkgs_dict = {
"path": f"{_internal_modules_path}/tutor-0.0.1-py3-none-any.whl",
"deps": ["pkg_resources"],
},
"_ivp": {
"path": f"{_internal_modules_path}/ivp-0.0.1-py3-none-any.whl",
},
"scipy": { # this force scipy to load _ivp first
"deps": ["_ivp"],
},
}
# checking that 'path' is in dict allows to add dependencies to pyodide packages
......
build/
dist/
*.egg-info/
build:
python3 setup.py sdist bdist_wheel
clean:
rm -rf *.egg-info/ build/ dist/
.PHONY: build clean
"""Suite of ODE solvers implemented in Python."""
from .ivp import solve_ivp
from .rk import RK23, RK45, DOP853
from .radau import Radau
from .bdf import BDF
from .lsoda import LSODA
from .common import OdeSolution
from .base import DenseOutput, OdeSolver
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
import numpy as np
from scipy.linalg import lu_factor, lu_solve
from scipy.sparse import issparse, csc_matrix, eye
from scipy.sparse.linalg import splu
from scipy.optimize._numdiff import group_columns
from .common import (validate_max_step, validate_tol, select_initial_step,
norm, EPS, num_jac, validate_first_step,
warn_extraneous)
from .base import OdeSolver, DenseOutput
MAX_ORDER = 5
NEWTON_MAXITER = 4
MIN_FACTOR = 0.2
MAX_FACTOR = 10
def compute_R(order, factor):
"""Compute the matrix for changing the differences array."""
I = np.arange(1, order + 1)[:, None]
J = np.arange(1, order + 1)
M = np.zeros((order + 1, order + 1))
M[1:, 1:] = (I - 1 - factor * J) / I
M[0] = 1
return np.cumprod(M, axis=0)
def change_D(D, order, factor):
"""Change differences array in-place when step size is changed."""
R = compute_R(order, factor)
U = compute_R(order, 1)
RU = R.dot(U)
D[:order + 1] = np.dot(RU.T, D[:order + 1])
def solve_bdf_system(fun, t_new, y_predict, c, psi, LU, solve_lu, scale, tol):
"""Solve the algebraic system resulting from BDF method."""
d = 0
y = y_predict.copy()
dy_norm_old = None
converged = False
for k in range(NEWTON_MAXITER):
f = fun(t_new, y)
if not np.all(np.isfinite(f)):
break
dy = solve_lu(LU, c * f - psi - d)
dy_norm = norm(dy / scale)
if dy_norm_old is None:
rate = None
else:
rate = dy_norm / dy_norm_old
if (rate is not None and (rate >= 1 or
rate ** (NEWTON_MAXITER - k) / (1 - rate) * dy_norm > tol)):
break
y += dy
d += dy
if (dy_norm == 0 or
rate is not None and rate / (1 - rate) * dy_norm < tol):
converged = True
break
dy_norm_old = dy_norm
return converged, k + 1, y, d
class BDF(OdeSolver):
"""Implicit method based on backward-differentiation formulas.
This is a variable order method with the order varying automatically from
1 to 5. The general framework of the BDF algorithm is described in [1]_.
This class implements a quasi-constant step size as explained in [2]_.
The error estimation strategy for the constant-step BDF is derived in [3]_.
An accuracy enhancement using modified formulas (NDF) [2]_ is also implemented.
Can be applied in the complex domain.
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 the ndarray ``y``:
It can either have shape (n,); then ``fun`` must return array_like with
shape (n,). Alternatively it can have shape (n, k); then ``fun``
must return an array_like with shape (n, k), i.e. each column
corresponds to a single column in ``y``. The choice between the two
options is determined by `vectorized` argument (see below). The
vectorized implementation allows a faster approximation of the Jacobian
by finite differences (required for this solver).
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.
first_step : float or None, optional
Initial step size. Default is ``None`` which means that the algorithm
should choose.
max_step : float, optional
Maximum allowed step size. Default is np.inf, i.e., the step size is not
bounded and determined solely by the solver.
rtol, atol : float and array_like, optional
Relative and absolute tolerances. The solver keeps the local error
estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a
relative accuracy (number of correct digits). But if a component of `y`
is approximately below `atol`, the error only needs to fall within
the same `atol` threshold, and the number of correct digits is not
guaranteed. If components of y have different scales, it might be
beneficial to set different `atol` values for different components by
passing array_like with shape (n,) for `atol`. Default values are
1e-3 for `rtol` and 1e-6 for `atol`.
jac : {None, array_like, sparse_matrix, callable}, optional
Jacobian matrix of the right-hand side of the system with respect to y,
required by this method. The Jacobian matrix has shape (n, n) and its
element (i, j) is equal to ``d f_i / d y_j``.
There are three ways to define the Jacobian:
* If array_like or sparse_matrix, the Jacobian is assumed to
be constant.
* If callable, the Jacobian is assumed to depend on both
t and y; it will be called as ``jac(t, y)`` as necessary.
For the 'Radau' and 'BDF' methods, the return value might be a
sparse matrix.
* If None (default), the Jacobian will be approximated by
finite differences.
It is generally recommended to provide the Jacobian rather than
relying on a finite-difference approximation.
jac_sparsity : {None, array_like, sparse matrix}, optional
Defines a sparsity structure of the Jacobian matrix for a
finite-difference approximation. Its shape must be (n, n). This argument
is ignored if `jac` is not `None`. If the Jacobian has only few non-zero
elements in *each* row, providing the sparsity structure will greatly
speed up the computations [4]_. A zero entry means that a corresponding
element in the Jacobian is always zero. If None (default), the Jacobian
is assumed to be dense.
vectorized : bool, optional
Whether `fun` is implemented in a vectorized fashion. 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 evaluations of the right-hand side.
njev : int
Number of evaluations of the Jacobian.
nlu : int
Number of LU decompositions.
References
----------
.. [1] G. D. Byrne, A. C. Hindmarsh, "A Polyalgorithm for the Numerical
Solution of Ordinary Differential Equations", ACM Transactions on
Mathematical Software, Vol. 1, No. 1, pp. 71-96, March 1975.
.. [2] L. F. Shampine, M. W. Reichelt, "THE MATLAB ODE SUITE", SIAM J. SCI.
COMPUTE., Vol. 18, No. 1, pp. 1-22, January 1997.
.. [3] E. Hairer, G. Wanner, "Solving Ordinary Differential Equations I:
Nonstiff Problems", Sec. III.2.
.. [4] A. Curtis, M. J. D. Powell, and J. Reid, "On the estimation of
sparse Jacobian matrices", Journal of the Institute of Mathematics
and its Applications, 13, pp. 117-120, 1974.
"""
def __init__(self, fun, t0, y0, t_bound, max_step=np.inf,
rtol=1e-3, atol=1e-6, jac=None, jac_sparsity=None,
vectorized=False, first_step=None, **extraneous):
warn_extraneous(extraneous)
super().__init__(fun, t0, y0, t_bound, vectorized,
support_complex=True)
self.max_step = validate_max_step(max_step)
self.rtol, self.atol = validate_tol(rtol, atol, self.n)
f = self.fun(self.t, self.y)
if first_step is None:
self.h_abs = select_initial_step(self.fun, self.t, self.y, f,
self.direction, 1,
self.rtol, self.atol)
else:
self.h_abs = validate_first_step(first_step, t0, t_bound)
self.h_abs_old = None
self.error_norm_old = None
self.newton_tol = max(10 * EPS / rtol, min(0.03, rtol ** 0.5))
self.jac_factor = None
self.jac, self.J = self._validate_jac(jac, jac_sparsity)
if issparse(self.J):
def lu(A):
self.nlu += 1
return splu(A)
def solve_lu(LU, b):
return LU.solve(b)
I = eye(self.n, format='csc', dtype=self.y.dtype)
else:
def lu(A):
self.nlu += 1
return lu_factor(A, overwrite_a=True)