Source code for quantecon.ivp

r"""
Base class for solving initial value problems (IVPs) of the form:

.. math::

    \frac{dy}{dt} = f(t,y),\ y(t_0) = y_0

using finite difference methods. The `quantecon.ivp` class uses various
integrators from the `scipy.integrate.ode` module to perform the
integration (i.e., solve the ODE) and parametric B-spline interpolation
from `scipy.interpolate` to approximate the value of the solution
between grid points. The `quantecon.ivp` module also provides a method
for computing the residual of the solution which can be used for
assessing the overall accuracy of the approximated solution.

@author : David R. Pugh
@date : 2014-09-09

"""
from __future__ import division

import numpy as np
from scipy import integrate, interpolate


[docs]class IVP(integrate.ode): def __init__(self, f, jac=None): r""" Creates an instance of the IVP class. Parameters ---------- f : callable ``f(t, y, *f_args)`` Right hand side of the system of equations defining the ODE. The independent variable, ``t``, is a ``scalar``; ``y`` is an ``ndarray`` of dependent variables with ``y.shape == (n,)``. The function `f` should return a ``scalar``, ``ndarray`` or ``list`` (but not a ``tuple``). jac : callable ``jac(t, y, *jac_args)``, optional(default=None) Jacobian of the right hand side of the system of equations defining the ODE. .. :math: \mathcal{J}_{i,j} = \bigg[\frac{\partial f_i}{\partial y_j}\bigg] """ super(IVP, self).__init__(f, jac) def _integrate_fixed_trajectory(self, h, T, step, relax): """Generates a solution trajectory of fixed length.""" # initialize the solution using initial condition solution = np.hstack((self.t, self.y)) while self.successful(): self.integrate(self.t + h, step, relax) current_step = np.hstack((self.t, self.y)) solution = np.vstack((solution, current_step)) if (h > 0) and (self.t >= T): break elif (h < 0) and (self.t <= T): break else: continue return solution def _integrate_variable_trajectory(self, h, g, tol, step, relax): """Generates a solution trajectory of variable length.""" # initialize the solution using initial condition solution = np.hstack((self.t, self.y)) while self.successful(): self.integrate(self.t + h, step, relax) current_step = np.hstack((self.t, self.y)) solution = np.vstack((solution, current_step)) if g(self.t, self.y, *self.f_params) < tol: break else: continue return solution def _initialize_integrator(self, t0, y0, integrator, **kwargs): """Initializes the integrator prior to integration.""" # set the initial condition self.set_initial_value(y0, t0) # select the integrator self.set_integrator(integrator, **kwargs)
[docs] def compute_residual(self, traj, ti, k=3, ext=2): r""" The residual is the difference between the derivative of the B-spline approximation of the solution trajectory and the right-hand side of the original ODE evaluated along the approximated solution trajectory. Parameters ---------- traj : array_like (float) Solution trajectory providing the data points for constructing the B-spline representation. ti : array_like (float) Array of values for the independent variable at which to interpolate the value of the B-spline. k : int, optional(default=3) Degree of the desired B-spline. Degree must satisfy :math:`1 \le k \le 5`. ext : int, optional(default=2) Controls the value of returned elements for outside the original knot sequence provided by traj. For extrapolation, set `ext=0`; `ext=1` returns zero; `ext=2` raises a `ValueError`. Returns ------- residual : array (float) Difference between the derivative of the B-spline approximation of the solution trajectory and the right-hand side of the ODE evaluated along the approximated solution trajectory. """ # B-spline approximations of the solution and its derivative soln = self.interpolate(traj, ti, k, 0, ext) deriv = self.interpolate(traj, ti, k, 1, ext) # rhs of ode evaluated along approximate solution T = ti.size rhs_ode = np.vstack(self.f(ti[i], soln[i, 1:], *self.f_params) for i in range(T)) rhs_ode = np.hstack((ti[:, np.newaxis], rhs_ode)) # should be roughly zero everywhere (if approximation is any good!) residual = deriv - rhs_ode return residual
[docs] def solve(self, t0, y0, h=1.0, T=None, g=None, tol=None, integrator='dopri5', step=False, relax=False, **kwargs): r""" Solve the IVP by integrating the ODE given some initial condition. Parameters ---------- t0 : float Initial condition for the independent variable. y0 : array_like (float, shape=(n,)) Initial condition for the dependent variables. h : float, optional(default=1.0) Step-size for computing the solution. Can be positive or negative depending on the desired direction of integration. T : int, optional(default=None) Terminal value for the independent variable. One of either `T` or `g` must be specified. g : callable ``g(t, y, f_args)``, optional(default=None) Provides a stopping condition for the integration. If specified user must also specify a stopping tolerance, `tol`. tol : float, optional (default=None) Stopping tolerance for the integration. Only required if `g` is also specifed. integrator : str, optional(default='dopri5') Must be one of 'vode', 'lsoda', 'dopri5', or 'dop853' step : bool, optional(default=False) Allows access to internal steps for those solvers that use adaptive step size routines. Currently only 'vode', 'zvode', and 'lsoda' support `step=True`. relax : bool, optional(default=False) Currently only 'vode', 'zvode', and 'lsoda' support `relax=True`. **kwargs : dict, optional(default=None) Dictionary of integrator specific keyword arguments. See the Notes section of the docstring for `scipy.integrate.ode` for a complete description of solver specific keyword arguments. Returns ------- solution: ndarray (float) Simulated solution trajectory. """ self._initialize_integrator(t0, y0, integrator, **kwargs) if (g is not None) and (tol is not None): soln = self._integrate_variable_trajectory(h, g, tol, step, relax) elif T is not None: soln = self._integrate_fixed_trajectory(h, T, step, relax) else: mesg = "Either both 'g' and 'tol', or 'T' must be specified." raise ValueError(mesg) return soln
[docs] def interpolate(self, traj, ti, k=3, der=0, ext=2): r""" Parametric B-spline interpolation in N-dimensions. Parameters ---------- traj : array_like (float) Solution trajectory providing the data points for constructing the B-spline representation. ti : array_like (float) Array of values for the independent variable at which to interpolate the value of the B-spline. k : int, optional(default=3) Degree of the desired B-spline. Degree must satisfy :math:`1 \le k \le 5`. der : int, optional(default=0) The order of derivative of the spline to compute (must be less than or equal to `k`). ext : int, optional(default=2) Controls the value of returned elements for outside the original knot sequence provided by traj. For extrapolation, set `ext=0`; `ext=1` returns zero; `ext=2` raises a `ValueError`. Returns ------- interp_traj: ndarray (float) The interpolated trajectory. """ # array of parameter values u = traj[:, 0] # build list of input arrays n = traj.shape[1] x = [traj[:, i] for i in range(1, n)] # construct the B-spline representation (s=0 forces interpolation!) tck, t = interpolate.splprep(x, u=u, k=k, s=0) # evaluate the B-spline (returns a list) out = interpolate.splev(ti, tck, der, ext) # convert to a 2D array interp_traj = np.hstack((ti[:, np.newaxis], np.array(out).T)) return interp_traj