1# Authors: Pearu Peterson, Pauli Virtanen, John Travers 2""" 3First-order ODE integrators. 4 5User-friendly interface to various numerical integrators for solving a 6system of first order ODEs with prescribed initial conditions:: 7 8 d y(t)[i] 9 --------- = f(t,y(t))[i], 10 d t 11 12 y(t=0)[i] = y0[i], 13 14where:: 15 16 i = 0, ..., len(y0) - 1 17 18class ode 19--------- 20 21A generic interface class to numeric integrators. It has the following 22methods:: 23 24 integrator = ode(f, jac=None) 25 integrator = integrator.set_integrator(name, **params) 26 integrator = integrator.set_initial_value(y0, t0=0.0) 27 integrator = integrator.set_f_params(*args) 28 integrator = integrator.set_jac_params(*args) 29 y1 = integrator.integrate(t1, step=False, relax=False) 30 flag = integrator.successful() 31 32class complex_ode 33----------------- 34 35This class has the same generic interface as ode, except it can handle complex 36f, y and Jacobians by transparently translating them into the equivalent 37real-valued system. It supports the real-valued solvers (i.e., not zvode) and is 38an alternative to ode with the zvode solver, sometimes performing better. 39""" 40# XXX: Integrators must have: 41# =========================== 42# cvode - C version of vode and vodpk with many improvements. 43# Get it from http://www.netlib.org/ode/cvode.tar.gz. 44# To wrap cvode to Python, one must write the extension module by 45# hand. Its interface is too much 'advanced C' that using f2py 46# would be too complicated (or impossible). 47# 48# How to define a new integrator: 49# =============================== 50# 51# class myodeint(IntegratorBase): 52# 53# runner = <odeint function> or None 54# 55# def __init__(self,...): # required 56# <initialize> 57# 58# def reset(self,n,has_jac): # optional 59# # n - the size of the problem (number of equations) 60# # has_jac - whether user has supplied its own routine for Jacobian 61# <allocate memory,initialize further> 62# 63# def run(self,f,jac,y0,t0,t1,f_params,jac_params): # required 64# # this method is called to integrate from t=t0 to t=t1 65# # with initial condition y0. f and jac are user-supplied functions 66# # that define the problem. f_params,jac_params are additional 67# # arguments 68# # to these functions. 69# <calculate y1> 70# if <calculation was unsuccessful>: 71# self.success = 0 72# return t1,y1 73# 74# # In addition, one can define step() and run_relax() methods (they 75# # take the same arguments as run()) if the integrator can support 76# # these features (see IntegratorBase doc strings). 77# 78# if myodeint.runner: 79# IntegratorBase.integrator_classes.append(myodeint) 80 81__all__ = ['ode', 'complex_ode'] 82__version__ = "$Id$" 83__docformat__ = "restructuredtext en" 84 85import re 86import warnings 87 88from numpy import asarray, array, zeros, isscalar, real, imag, vstack 89 90from . import vode as _vode 91from . import _dop 92from . import lsoda as _lsoda 93 94 95_dop_int_dtype = _dop.types.intvar.dtype 96_vode_int_dtype = _vode.types.intvar.dtype 97_lsoda_int_dtype = _lsoda.types.intvar.dtype 98 99 100# ------------------------------------------------------------------------------ 101# User interface 102# ------------------------------------------------------------------------------ 103 104 105class ode: 106 """ 107 A generic interface class to numeric integrators. 108 109 Solve an equation system :math:`y'(t) = f(t,y)` with (optional) ``jac = df/dy``. 110 111 *Note*: The first two arguments of ``f(t, y, ...)`` are in the 112 opposite order of the arguments in the system definition function used 113 by `scipy.integrate.odeint`. 114 115 Parameters 116 ---------- 117 f : callable ``f(t, y, *f_args)`` 118 Right-hand side of the differential equation. t is a scalar, 119 ``y.shape == (n,)``. 120 ``f_args`` is set by calling ``set_f_params(*args)``. 121 `f` should return a scalar, array or list (not a tuple). 122 jac : callable ``jac(t, y, *jac_args)``, optional 123 Jacobian of the right-hand side, ``jac[i,j] = d f[i] / d y[j]``. 124 ``jac_args`` is set by calling ``set_jac_params(*args)``. 125 126 Attributes 127 ---------- 128 t : float 129 Current time. 130 y : ndarray 131 Current variable values. 132 133 See also 134 -------- 135 odeint : an integrator with a simpler interface based on lsoda from ODEPACK 136 quad : for finding the area under a curve 137 138 Notes 139 ----- 140 Available integrators are listed below. They can be selected using 141 the `set_integrator` method. 142 143 "vode" 144 145 Real-valued Variable-coefficient Ordinary Differential Equation 146 solver, with fixed-leading-coefficient implementation. It provides 147 implicit Adams method (for non-stiff problems) and a method based on 148 backward differentiation formulas (BDF) (for stiff problems). 149 150 Source: http://www.netlib.org/ode/vode.f 151 152 .. warning:: 153 154 This integrator is not re-entrant. You cannot have two `ode` 155 instances using the "vode" integrator at the same time. 156 157 This integrator accepts the following parameters in `set_integrator` 158 method of the `ode` class: 159 160 - atol : float or sequence 161 absolute tolerance for solution 162 - rtol : float or sequence 163 relative tolerance for solution 164 - lband : None or int 165 - uband : None or int 166 Jacobian band width, jac[i,j] != 0 for i-lband <= j <= i+uband. 167 Setting these requires your jac routine to return the jacobian 168 in packed format, jac_packed[i-j+uband, j] = jac[i,j]. The 169 dimension of the matrix must be (lband+uband+1, len(y)). 170 - method: 'adams' or 'bdf' 171 Which solver to use, Adams (non-stiff) or BDF (stiff) 172 - with_jacobian : bool 173 This option is only considered when the user has not supplied a 174 Jacobian function and has not indicated (by setting either band) 175 that the Jacobian is banded. In this case, `with_jacobian` specifies 176 whether the iteration method of the ODE solver's correction step is 177 chord iteration with an internally generated full Jacobian or 178 functional iteration with no Jacobian. 179 - nsteps : int 180 Maximum number of (internally defined) steps allowed during one 181 call to the solver. 182 - first_step : float 183 - min_step : float 184 - max_step : float 185 Limits for the step sizes used by the integrator. 186 - order : int 187 Maximum order used by the integrator, 188 order <= 12 for Adams, <= 5 for BDF. 189 190 "zvode" 191 192 Complex-valued Variable-coefficient Ordinary Differential Equation 193 solver, with fixed-leading-coefficient implementation. It provides 194 implicit Adams method (for non-stiff problems) and a method based on 195 backward differentiation formulas (BDF) (for stiff problems). 196 197 Source: http://www.netlib.org/ode/zvode.f 198 199 .. warning:: 200 201 This integrator is not re-entrant. You cannot have two `ode` 202 instances using the "zvode" integrator at the same time. 203 204 This integrator accepts the same parameters in `set_integrator` 205 as the "vode" solver. 206 207 .. note:: 208 209 When using ZVODE for a stiff system, it should only be used for 210 the case in which the function f is analytic, that is, when each f(i) 211 is an analytic function of each y(j). Analyticity means that the 212 partial derivative df(i)/dy(j) is a unique complex number, and this 213 fact is critical in the way ZVODE solves the dense or banded linear 214 systems that arise in the stiff case. For a complex stiff ODE system 215 in which f is not analytic, ZVODE is likely to have convergence 216 failures, and for this problem one should instead use DVODE on the 217 equivalent real system (in the real and imaginary parts of y). 218 219 "lsoda" 220 221 Real-valued Variable-coefficient Ordinary Differential Equation 222 solver, with fixed-leading-coefficient implementation. It provides 223 automatic method switching between implicit Adams method (for non-stiff 224 problems) and a method based on backward differentiation formulas (BDF) 225 (for stiff problems). 226 227 Source: http://www.netlib.org/odepack 228 229 .. warning:: 230 231 This integrator is not re-entrant. You cannot have two `ode` 232 instances using the "lsoda" integrator at the same time. 233 234 This integrator accepts the following parameters in `set_integrator` 235 method of the `ode` class: 236 237 - atol : float or sequence 238 absolute tolerance for solution 239 - rtol : float or sequence 240 relative tolerance for solution 241 - lband : None or int 242 - uband : None or int 243 Jacobian band width, jac[i,j] != 0 for i-lband <= j <= i+uband. 244 Setting these requires your jac routine to return the jacobian 245 in packed format, jac_packed[i-j+uband, j] = jac[i,j]. 246 - with_jacobian : bool 247 *Not used.* 248 - nsteps : int 249 Maximum number of (internally defined) steps allowed during one 250 call to the solver. 251 - first_step : float 252 - min_step : float 253 - max_step : float 254 Limits for the step sizes used by the integrator. 255 - max_order_ns : int 256 Maximum order used in the nonstiff case (default 12). 257 - max_order_s : int 258 Maximum order used in the stiff case (default 5). 259 - max_hnil : int 260 Maximum number of messages reporting too small step size (t + h = t) 261 (default 0) 262 - ixpr : int 263 Whether to generate extra printing at method switches (default False). 264 265 "dopri5" 266 267 This is an explicit runge-kutta method of order (4)5 due to Dormand & 268 Prince (with stepsize control and dense output). 269 270 Authors: 271 272 E. Hairer and G. Wanner 273 Universite de Geneve, Dept. de Mathematiques 274 CH-1211 Geneve 24, Switzerland 275 e-mail: ernst.hairer@math.unige.ch, gerhard.wanner@math.unige.ch 276 277 This code is described in [HNW93]_. 278 279 This integrator accepts the following parameters in set_integrator() 280 method of the ode class: 281 282 - atol : float or sequence 283 absolute tolerance for solution 284 - rtol : float or sequence 285 relative tolerance for solution 286 - nsteps : int 287 Maximum number of (internally defined) steps allowed during one 288 call to the solver. 289 - first_step : float 290 - max_step : float 291 - safety : float 292 Safety factor on new step selection (default 0.9) 293 - ifactor : float 294 - dfactor : float 295 Maximum factor to increase/decrease step size by in one step 296 - beta : float 297 Beta parameter for stabilised step size control. 298 - verbosity : int 299 Switch for printing messages (< 0 for no messages). 300 301 "dop853" 302 303 This is an explicit runge-kutta method of order 8(5,3) due to Dormand 304 & Prince (with stepsize control and dense output). 305 306 Options and references the same as "dopri5". 307 308 Examples 309 -------- 310 311 A problem to integrate and the corresponding jacobian: 312 313 >>> from scipy.integrate import ode 314 >>> 315 >>> y0, t0 = [1.0j, 2.0], 0 316 >>> 317 >>> def f(t, y, arg1): 318 ... return [1j*arg1*y[0] + y[1], -arg1*y[1]**2] 319 >>> def jac(t, y, arg1): 320 ... return [[1j*arg1, 1], [0, -arg1*2*y[1]]] 321 322 The integration: 323 324 >>> r = ode(f, jac).set_integrator('zvode', method='bdf') 325 >>> r.set_initial_value(y0, t0).set_f_params(2.0).set_jac_params(2.0) 326 >>> t1 = 10 327 >>> dt = 1 328 >>> while r.successful() and r.t < t1: 329 ... print(r.t+dt, r.integrate(r.t+dt)) 330 1 [-0.71038232+0.23749653j 0.40000271+0.j ] 331 2.0 [0.19098503-0.52359246j 0.22222356+0.j ] 332 3.0 [0.47153208+0.52701229j 0.15384681+0.j ] 333 4.0 [-0.61905937+0.30726255j 0.11764744+0.j ] 334 5.0 [0.02340997-0.61418799j 0.09523835+0.j ] 335 6.0 [0.58643071+0.339819j 0.08000018+0.j ] 336 7.0 [-0.52070105+0.44525141j 0.06896565+0.j ] 337 8.0 [-0.15986733-0.61234476j 0.06060616+0.j ] 338 9.0 [0.64850462+0.15048982j 0.05405414+0.j ] 339 10.0 [-0.38404699+0.56382299j 0.04878055+0.j ] 340 341 References 342 ---------- 343 .. [HNW93] E. Hairer, S.P. Norsett and G. Wanner, Solving Ordinary 344 Differential Equations i. Nonstiff Problems. 2nd edition. 345 Springer Series in Computational Mathematics, 346 Springer-Verlag (1993) 347 348 """ 349 350 def __init__(self, f, jac=None): 351 self.stiff = 0 352 self.f = f 353 self.jac = jac 354 self.f_params = () 355 self.jac_params = () 356 self._y = [] 357 358 @property 359 def y(self): 360 return self._y 361 362 def set_initial_value(self, y, t=0.0): 363 """Set initial conditions y(t) = y.""" 364 if isscalar(y): 365 y = [y] 366 n_prev = len(self._y) 367 if not n_prev: 368 self.set_integrator('') # find first available integrator 369 self._y = asarray(y, self._integrator.scalar) 370 self.t = t 371 self._integrator.reset(len(self._y), self.jac is not None) 372 return self 373 374 def set_integrator(self, name, **integrator_params): 375 """ 376 Set integrator by name. 377 378 Parameters 379 ---------- 380 name : str 381 Name of the integrator. 382 integrator_params 383 Additional parameters for the integrator. 384 """ 385 integrator = find_integrator(name) 386 if integrator is None: 387 # FIXME: this really should be raise an exception. Will that break 388 # any code? 389 warnings.warn('No integrator name match with %r or is not ' 390 'available.' % name) 391 else: 392 self._integrator = integrator(**integrator_params) 393 if not len(self._y): 394 self.t = 0.0 395 self._y = array([0.0], self._integrator.scalar) 396 self._integrator.reset(len(self._y), self.jac is not None) 397 return self 398 399 def integrate(self, t, step=False, relax=False): 400 """Find y=y(t), set y as an initial condition, and return y. 401 402 Parameters 403 ---------- 404 t : float 405 The endpoint of the integration step. 406 step : bool 407 If True, and if the integrator supports the step method, 408 then perform a single integration step and return. 409 This parameter is provided in order to expose internals of 410 the implementation, and should not be changed from its default 411 value in most cases. 412 relax : bool 413 If True and if the integrator supports the run_relax method, 414 then integrate until t_1 >= t and return. ``relax`` is not 415 referenced if ``step=True``. 416 This parameter is provided in order to expose internals of 417 the implementation, and should not be changed from its default 418 value in most cases. 419 420 Returns 421 ------- 422 y : float 423 The integrated value at t 424 """ 425 if step and self._integrator.supports_step: 426 mth = self._integrator.step 427 elif relax and self._integrator.supports_run_relax: 428 mth = self._integrator.run_relax 429 else: 430 mth = self._integrator.run 431 432 try: 433 self._y, self.t = mth(self.f, self.jac or (lambda: None), 434 self._y, self.t, t, 435 self.f_params, self.jac_params) 436 except SystemError as e: 437 # f2py issue with tuple returns, see ticket 1187. 438 raise ValueError( 439 'Function to integrate must not return a tuple.' 440 ) from e 441 442 return self._y 443 444 def successful(self): 445 """Check if integration was successful.""" 446 try: 447 self._integrator 448 except AttributeError: 449 self.set_integrator('') 450 return self._integrator.success == 1 451 452 def get_return_code(self): 453 """Extracts the return code for the integration to enable better control 454 if the integration fails. 455 456 In general, a return code > 0 implies success, while a return code < 0 457 implies failure. 458 459 Notes 460 ----- 461 This section describes possible return codes and their meaning, for available 462 integrators that can be selected by `set_integrator` method. 463 464 "vode" 465 466 =========== ======= 467 Return Code Message 468 =========== ======= 469 2 Integration successful. 470 -1 Excess work done on this call. (Perhaps wrong MF.) 471 -2 Excess accuracy requested. (Tolerances too small.) 472 -3 Illegal input detected. (See printed message.) 473 -4 Repeated error test failures. (Check all input.) 474 -5 Repeated convergence failures. (Perhaps bad Jacobian 475 supplied or wrong choice of MF or tolerances.) 476 -6 Error weight became zero during problem. (Solution 477 component i vanished, and ATOL or ATOL(i) = 0.) 478 =========== ======= 479 480 "zvode" 481 482 =========== ======= 483 Return Code Message 484 =========== ======= 485 2 Integration successful. 486 -1 Excess work done on this call. (Perhaps wrong MF.) 487 -2 Excess accuracy requested. (Tolerances too small.) 488 -3 Illegal input detected. (See printed message.) 489 -4 Repeated error test failures. (Check all input.) 490 -5 Repeated convergence failures. (Perhaps bad Jacobian 491 supplied or wrong choice of MF or tolerances.) 492 -6 Error weight became zero during problem. (Solution 493 component i vanished, and ATOL or ATOL(i) = 0.) 494 =========== ======= 495 496 "dopri5" 497 498 =========== ======= 499 Return Code Message 500 =========== ======= 501 1 Integration successful. 502 2 Integration successful (interrupted by solout). 503 -1 Input is not consistent. 504 -2 Larger nsteps is needed. 505 -3 Step size becomes too small. 506 -4 Problem is probably stiff (interrupted). 507 =========== ======= 508 509 "dop853" 510 511 =========== ======= 512 Return Code Message 513 =========== ======= 514 1 Integration successful. 515 2 Integration successful (interrupted by solout). 516 -1 Input is not consistent. 517 -2 Larger nsteps is needed. 518 -3 Step size becomes too small. 519 -4 Problem is probably stiff (interrupted). 520 =========== ======= 521 522 "lsoda" 523 524 =========== ======= 525 Return Code Message 526 =========== ======= 527 2 Integration successful. 528 -1 Excess work done on this call (perhaps wrong Dfun type). 529 -2 Excess accuracy requested (tolerances too small). 530 -3 Illegal input detected (internal error). 531 -4 Repeated error test failures (internal error). 532 -5 Repeated convergence failures (perhaps bad Jacobian or tolerances). 533 -6 Error weight became zero during problem. 534 -7 Internal workspace insufficient to finish (internal error). 535 =========== ======= 536 """ 537 try: 538 self._integrator 539 except AttributeError: 540 self.set_integrator('') 541 return self._integrator.istate 542 543 def set_f_params(self, *args): 544 """Set extra parameters for user-supplied function f.""" 545 self.f_params = args 546 return self 547 548 def set_jac_params(self, *args): 549 """Set extra parameters for user-supplied function jac.""" 550 self.jac_params = args 551 return self 552 553 def set_solout(self, solout): 554 """ 555 Set callable to be called at every successful integration step. 556 557 Parameters 558 ---------- 559 solout : callable 560 ``solout(t, y)`` is called at each internal integrator step, 561 t is a scalar providing the current independent position 562 y is the current soloution ``y.shape == (n,)`` 563 solout should return -1 to stop integration 564 otherwise it should return None or 0 565 566 """ 567 if self._integrator.supports_solout: 568 self._integrator.set_solout(solout) 569 if self._y is not None: 570 self._integrator.reset(len(self._y), self.jac is not None) 571 else: 572 raise ValueError("selected integrator does not support solout," 573 " choose another one") 574 575 576def _transform_banded_jac(bjac): 577 """ 578 Convert a real matrix of the form (for example) 579 580 [0 0 A B] [0 0 0 B] 581 [0 0 C D] [0 0 A D] 582 [E F G H] to [0 F C H] 583 [I J K L] [E J G L] 584 [I 0 K 0] 585 586 That is, every other column is shifted up one. 587 """ 588 # Shift every other column. 589 newjac = zeros((bjac.shape[0] + 1, bjac.shape[1])) 590 newjac[1:, ::2] = bjac[:, ::2] 591 newjac[:-1, 1::2] = bjac[:, 1::2] 592 return newjac 593 594 595class complex_ode(ode): 596 """ 597 A wrapper of ode for complex systems. 598 599 This functions similarly as `ode`, but re-maps a complex-valued 600 equation system to a real-valued one before using the integrators. 601 602 Parameters 603 ---------- 604 f : callable ``f(t, y, *f_args)`` 605 Rhs of the equation. t is a scalar, ``y.shape == (n,)``. 606 ``f_args`` is set by calling ``set_f_params(*args)``. 607 jac : callable ``jac(t, y, *jac_args)`` 608 Jacobian of the rhs, ``jac[i,j] = d f[i] / d y[j]``. 609 ``jac_args`` is set by calling ``set_f_params(*args)``. 610 611 Attributes 612 ---------- 613 t : float 614 Current time. 615 y : ndarray 616 Current variable values. 617 618 Examples 619 -------- 620 For usage examples, see `ode`. 621 622 """ 623 624 def __init__(self, f, jac=None): 625 self.cf = f 626 self.cjac = jac 627 if jac is None: 628 ode.__init__(self, self._wrap, None) 629 else: 630 ode.__init__(self, self._wrap, self._wrap_jac) 631 632 def _wrap(self, t, y, *f_args): 633 f = self.cf(*((t, y[::2] + 1j * y[1::2]) + f_args)) 634 # self.tmp is a real-valued array containing the interleaved 635 # real and imaginary parts of f. 636 self.tmp[::2] = real(f) 637 self.tmp[1::2] = imag(f) 638 return self.tmp 639 640 def _wrap_jac(self, t, y, *jac_args): 641 # jac is the complex Jacobian computed by the user-defined function. 642 jac = self.cjac(*((t, y[::2] + 1j * y[1::2]) + jac_args)) 643 644 # jac_tmp is the real version of the complex Jacobian. Each complex 645 # entry in jac, say 2+3j, becomes a 2x2 block of the form 646 # [2 -3] 647 # [3 2] 648 jac_tmp = zeros((2 * jac.shape[0], 2 * jac.shape[1])) 649 jac_tmp[1::2, 1::2] = jac_tmp[::2, ::2] = real(jac) 650 jac_tmp[1::2, ::2] = imag(jac) 651 jac_tmp[::2, 1::2] = -jac_tmp[1::2, ::2] 652 653 ml = getattr(self._integrator, 'ml', None) 654 mu = getattr(self._integrator, 'mu', None) 655 if ml is not None or mu is not None: 656 # Jacobian is banded. The user's Jacobian function has computed 657 # the complex Jacobian in packed format. The corresponding 658 # real-valued version has every other column shifted up. 659 jac_tmp = _transform_banded_jac(jac_tmp) 660 661 return jac_tmp 662 663 @property 664 def y(self): 665 return self._y[::2] + 1j * self._y[1::2] 666 667 def set_integrator(self, name, **integrator_params): 668 """ 669 Set integrator by name. 670 671 Parameters 672 ---------- 673 name : str 674 Name of the integrator 675 integrator_params 676 Additional parameters for the integrator. 677 """ 678 if name == 'zvode': 679 raise ValueError("zvode must be used with ode, not complex_ode") 680 681 lband = integrator_params.get('lband') 682 uband = integrator_params.get('uband') 683 if lband is not None or uband is not None: 684 # The Jacobian is banded. Override the user-supplied bandwidths 685 # (which are for the complex Jacobian) with the bandwidths of 686 # the corresponding real-valued Jacobian wrapper of the complex 687 # Jacobian. 688 integrator_params['lband'] = 2 * (lband or 0) + 1 689 integrator_params['uband'] = 2 * (uband or 0) + 1 690 691 return ode.set_integrator(self, name, **integrator_params) 692 693 def set_initial_value(self, y, t=0.0): 694 """Set initial conditions y(t) = y.""" 695 y = asarray(y) 696 self.tmp = zeros(y.size * 2, 'float') 697 self.tmp[::2] = real(y) 698 self.tmp[1::2] = imag(y) 699 return ode.set_initial_value(self, self.tmp, t) 700 701 def integrate(self, t, step=False, relax=False): 702 """Find y=y(t), set y as an initial condition, and return y. 703 704 Parameters 705 ---------- 706 t : float 707 The endpoint of the integration step. 708 step : bool 709 If True, and if the integrator supports the step method, 710 then perform a single integration step and return. 711 This parameter is provided in order to expose internals of 712 the implementation, and should not be changed from its default 713 value in most cases. 714 relax : bool 715 If True and if the integrator supports the run_relax method, 716 then integrate until t_1 >= t and return. ``relax`` is not 717 referenced if ``step=True``. 718 This parameter is provided in order to expose internals of 719 the implementation, and should not be changed from its default 720 value in most cases. 721 722 Returns 723 ------- 724 y : float 725 The integrated value at t 726 """ 727 y = ode.integrate(self, t, step, relax) 728 return y[::2] + 1j * y[1::2] 729 730 def set_solout(self, solout): 731 """ 732 Set callable to be called at every successful integration step. 733 734 Parameters 735 ---------- 736 solout : callable 737 ``solout(t, y)`` is called at each internal integrator step, 738 t is a scalar providing the current independent position 739 y is the current soloution ``y.shape == (n,)`` 740 solout should return -1 to stop integration 741 otherwise it should return None or 0 742 743 """ 744 if self._integrator.supports_solout: 745 self._integrator.set_solout(solout, complex=True) 746 else: 747 raise TypeError("selected integrator does not support solouta," 748 + "choose another one") 749 750 751# ------------------------------------------------------------------------------ 752# ODE integrators 753# ------------------------------------------------------------------------------ 754 755def find_integrator(name): 756 for cl in IntegratorBase.integrator_classes: 757 if re.match(name, cl.__name__, re.I): 758 return cl 759 return None 760 761 762class IntegratorConcurrencyError(RuntimeError): 763 """ 764 Failure due to concurrent usage of an integrator that can be used 765 only for a single problem at a time. 766 767 """ 768 769 def __init__(self, name): 770 msg = ("Integrator `%s` can be used to solve only a single problem " 771 "at a time. If you want to integrate multiple problems, " 772 "consider using a different integrator " 773 "(see `ode.set_integrator`)") % name 774 RuntimeError.__init__(self, msg) 775 776 777class IntegratorBase: 778 runner = None # runner is None => integrator is not available 779 success = None # success==1 if integrator was called successfully 780 istate = None # istate > 0 means success, istate < 0 means failure 781 supports_run_relax = None 782 supports_step = None 783 supports_solout = False 784 integrator_classes = [] 785 scalar = float 786 787 def acquire_new_handle(self): 788 # Some of the integrators have internal state (ancient 789 # Fortran...), and so only one instance can use them at a time. 790 # We keep track of this, and fail when concurrent usage is tried. 791 self.__class__.active_global_handle += 1 792 self.handle = self.__class__.active_global_handle 793 794 def check_handle(self): 795 if self.handle is not self.__class__.active_global_handle: 796 raise IntegratorConcurrencyError(self.__class__.__name__) 797 798 def reset(self, n, has_jac): 799 """Prepare integrator for call: allocate memory, set flags, etc. 800 n - number of equations. 801 has_jac - if user has supplied function for evaluating Jacobian. 802 """ 803 804 def run(self, f, jac, y0, t0, t1, f_params, jac_params): 805 """Integrate from t=t0 to t=t1 using y0 as an initial condition. 806 Return 2-tuple (y1,t1) where y1 is the result and t=t1 807 defines the stoppage coordinate of the result. 808 """ 809 raise NotImplementedError('all integrators must define ' 810 'run(f, jac, t0, t1, y0, f_params, jac_params)') 811 812 def step(self, f, jac, y0, t0, t1, f_params, jac_params): 813 """Make one integration step and return (y1,t1).""" 814 raise NotImplementedError('%s does not support step() method' % 815 self.__class__.__name__) 816 817 def run_relax(self, f, jac, y0, t0, t1, f_params, jac_params): 818 """Integrate from t=t0 to t>=t1 and return (y1,t).""" 819 raise NotImplementedError('%s does not support run_relax() method' % 820 self.__class__.__name__) 821 822 # XXX: __str__ method for getting visual state of the integrator 823 824 825def _vode_banded_jac_wrapper(jacfunc, ml, jac_params): 826 """ 827 Wrap a banded Jacobian function with a function that pads 828 the Jacobian with `ml` rows of zeros. 829 """ 830 831 def jac_wrapper(t, y): 832 jac = asarray(jacfunc(t, y, *jac_params)) 833 padded_jac = vstack((jac, zeros((ml, jac.shape[1])))) 834 return padded_jac 835 836 return jac_wrapper 837 838 839class vode(IntegratorBase): 840 runner = getattr(_vode, 'dvode', None) 841 842 messages = {-1: 'Excess work done on this call. (Perhaps wrong MF.)', 843 -2: 'Excess accuracy requested. (Tolerances too small.)', 844 -3: 'Illegal input detected. (See printed message.)', 845 -4: 'Repeated error test failures. (Check all input.)', 846 -5: 'Repeated convergence failures. (Perhaps bad' 847 ' Jacobian supplied or wrong choice of MF or tolerances.)', 848 -6: 'Error weight became zero during problem. (Solution' 849 ' component i vanished, and ATOL or ATOL(i) = 0.)' 850 } 851 supports_run_relax = 1 852 supports_step = 1 853 active_global_handle = 0 854 855 def __init__(self, 856 method='adams', 857 with_jacobian=False, 858 rtol=1e-6, atol=1e-12, 859 lband=None, uband=None, 860 order=12, 861 nsteps=500, 862 max_step=0.0, # corresponds to infinite 863 min_step=0.0, 864 first_step=0.0, # determined by solver 865 ): 866 867 if re.match(method, r'adams', re.I): 868 self.meth = 1 869 elif re.match(method, r'bdf', re.I): 870 self.meth = 2 871 else: 872 raise ValueError('Unknown integration method %s' % method) 873 self.with_jacobian = with_jacobian 874 self.rtol = rtol 875 self.atol = atol 876 self.mu = uband 877 self.ml = lband 878 879 self.order = order 880 self.nsteps = nsteps 881 self.max_step = max_step 882 self.min_step = min_step 883 self.first_step = first_step 884 self.success = 1 885 886 self.initialized = False 887 888 def _determine_mf_and_set_bands(self, has_jac): 889 """ 890 Determine the `MF` parameter (Method Flag) for the Fortran subroutine `dvode`. 891 892 In the Fortran code, the legal values of `MF` are: 893 10, 11, 12, 13, 14, 15, 20, 21, 22, 23, 24, 25, 894 -11, -12, -14, -15, -21, -22, -24, -25 895 but this Python wrapper does not use negative values. 896 897 Returns 898 899 mf = 10*self.meth + miter 900 901 self.meth is the linear multistep method: 902 self.meth == 1: method="adams" 903 self.meth == 2: method="bdf" 904 905 miter is the correction iteration method: 906 miter == 0: Functional iteraton; no Jacobian involved. 907 miter == 1: Chord iteration with user-supplied full Jacobian. 908 miter == 2: Chord iteration with internally computed full Jacobian. 909 miter == 3: Chord iteration with internally computed diagonal Jacobian. 910 miter == 4: Chord iteration with user-supplied banded Jacobian. 911 miter == 5: Chord iteration with internally computed banded Jacobian. 912 913 Side effects: If either self.mu or self.ml is not None and the other is None, 914 then the one that is None is set to 0. 915 """ 916 917 jac_is_banded = self.mu is not None or self.ml is not None 918 if jac_is_banded: 919 if self.mu is None: 920 self.mu = 0 921 if self.ml is None: 922 self.ml = 0 923 924 # has_jac is True if the user provided a Jacobian function. 925 if has_jac: 926 if jac_is_banded: 927 miter = 4 928 else: 929 miter = 1 930 else: 931 if jac_is_banded: 932 if self.ml == self.mu == 0: 933 miter = 3 # Chord iteration with internal diagonal Jacobian. 934 else: 935 miter = 5 # Chord iteration with internal banded Jacobian. 936 else: 937 # self.with_jacobian is set by the user in the call to ode.set_integrator. 938 if self.with_jacobian: 939 miter = 2 # Chord iteration with internal full Jacobian. 940 else: 941 miter = 0 # Functional iteraton; no Jacobian involved. 942 943 mf = 10 * self.meth + miter 944 return mf 945 946 def reset(self, n, has_jac): 947 mf = self._determine_mf_and_set_bands(has_jac) 948 949 if mf == 10: 950 lrw = 20 + 16 * n 951 elif mf in [11, 12]: 952 lrw = 22 + 16 * n + 2 * n * n 953 elif mf == 13: 954 lrw = 22 + 17 * n 955 elif mf in [14, 15]: 956 lrw = 22 + 18 * n + (3 * self.ml + 2 * self.mu) * n 957 elif mf == 20: 958 lrw = 20 + 9 * n 959 elif mf in [21, 22]: 960 lrw = 22 + 9 * n + 2 * n * n 961 elif mf == 23: 962 lrw = 22 + 10 * n 963 elif mf in [24, 25]: 964 lrw = 22 + 11 * n + (3 * self.ml + 2 * self.mu) * n 965 else: 966 raise ValueError('Unexpected mf=%s' % mf) 967 968 if mf % 10 in [0, 3]: 969 liw = 30 970 else: 971 liw = 30 + n 972 973 rwork = zeros((lrw,), float) 974 rwork[4] = self.first_step 975 rwork[5] = self.max_step 976 rwork[6] = self.min_step 977 self.rwork = rwork 978 979 iwork = zeros((liw,), _vode_int_dtype) 980 if self.ml is not None: 981 iwork[0] = self.ml 982 if self.mu is not None: 983 iwork[1] = self.mu 984 iwork[4] = self.order 985 iwork[5] = self.nsteps 986 iwork[6] = 2 # mxhnil 987 self.iwork = iwork 988 989 self.call_args = [self.rtol, self.atol, 1, 1, 990 self.rwork, self.iwork, mf] 991 self.success = 1 992 self.initialized = False 993 994 def run(self, f, jac, y0, t0, t1, f_params, jac_params): 995 if self.initialized: 996 self.check_handle() 997 else: 998 self.initialized = True 999 self.acquire_new_handle() 1000 1001 if self.ml is not None and self.ml > 0: 1002 # Banded Jacobian. Wrap the user-provided function with one 1003 # that pads the Jacobian array with the extra `self.ml` rows 1004 # required by the f2py-generated wrapper. 1005 jac = _vode_banded_jac_wrapper(jac, self.ml, jac_params) 1006 1007 args = ((f, jac, y0, t0, t1) + tuple(self.call_args) + 1008 (f_params, jac_params)) 1009 y1, t, istate = self.runner(*args) 1010 self.istate = istate 1011 if istate < 0: 1012 unexpected_istate_msg = 'Unexpected istate={:d}'.format(istate) 1013 warnings.warn('{:s}: {:s}'.format(self.__class__.__name__, 1014 self.messages.get(istate, unexpected_istate_msg))) 1015 self.success = 0 1016 else: 1017 self.call_args[3] = 2 # upgrade istate from 1 to 2 1018 self.istate = 2 1019 return y1, t 1020 1021 def step(self, *args): 1022 itask = self.call_args[2] 1023 self.call_args[2] = 2 1024 r = self.run(*args) 1025 self.call_args[2] = itask 1026 return r 1027 1028 def run_relax(self, *args): 1029 itask = self.call_args[2] 1030 self.call_args[2] = 3 1031 r = self.run(*args) 1032 self.call_args[2] = itask 1033 return r 1034 1035 1036if vode.runner is not None: 1037 IntegratorBase.integrator_classes.append(vode) 1038 1039 1040class zvode(vode): 1041 runner = getattr(_vode, 'zvode', None) 1042 1043 supports_run_relax = 1 1044 supports_step = 1 1045 scalar = complex 1046 active_global_handle = 0 1047 1048 def reset(self, n, has_jac): 1049 mf = self._determine_mf_and_set_bands(has_jac) 1050 1051 if mf in (10,): 1052 lzw = 15 * n 1053 elif mf in (11, 12): 1054 lzw = 15 * n + 2 * n ** 2 1055 elif mf in (-11, -12): 1056 lzw = 15 * n + n ** 2 1057 elif mf in (13,): 1058 lzw = 16 * n 1059 elif mf in (14, 15): 1060 lzw = 17 * n + (3 * self.ml + 2 * self.mu) * n 1061 elif mf in (-14, -15): 1062 lzw = 16 * n + (2 * self.ml + self.mu) * n 1063 elif mf in (20,): 1064 lzw = 8 * n 1065 elif mf in (21, 22): 1066 lzw = 8 * n + 2 * n ** 2 1067 elif mf in (-21, -22): 1068 lzw = 8 * n + n ** 2 1069 elif mf in (23,): 1070 lzw = 9 * n 1071 elif mf in (24, 25): 1072 lzw = 10 * n + (3 * self.ml + 2 * self.mu) * n 1073 elif mf in (-24, -25): 1074 lzw = 9 * n + (2 * self.ml + self.mu) * n 1075 1076 lrw = 20 + n 1077 1078 if mf % 10 in (0, 3): 1079 liw = 30 1080 else: 1081 liw = 30 + n 1082 1083 zwork = zeros((lzw,), complex) 1084 self.zwork = zwork 1085 1086 rwork = zeros((lrw,), float) 1087 rwork[4] = self.first_step 1088 rwork[5] = self.max_step 1089 rwork[6] = self.min_step 1090 self.rwork = rwork 1091 1092 iwork = zeros((liw,), _vode_int_dtype) 1093 if self.ml is not None: 1094 iwork[0] = self.ml 1095 if self.mu is not None: 1096 iwork[1] = self.mu 1097 iwork[4] = self.order 1098 iwork[5] = self.nsteps 1099 iwork[6] = 2 # mxhnil 1100 self.iwork = iwork 1101 1102 self.call_args = [self.rtol, self.atol, 1, 1, 1103 self.zwork, self.rwork, self.iwork, mf] 1104 self.success = 1 1105 self.initialized = False 1106 1107 1108if zvode.runner is not None: 1109 IntegratorBase.integrator_classes.append(zvode) 1110 1111 1112class dopri5(IntegratorBase): 1113 runner = getattr(_dop, 'dopri5', None) 1114 name = 'dopri5' 1115 supports_solout = True 1116 1117 messages = {1: 'computation successful', 1118 2: 'computation successful (interrupted by solout)', 1119 -1: 'input is not consistent', 1120 -2: 'larger nsteps is needed', 1121 -3: 'step size becomes too small', 1122 -4: 'problem is probably stiff (interrupted)', 1123 } 1124 1125 def __init__(self, 1126 rtol=1e-6, atol=1e-12, 1127 nsteps=500, 1128 max_step=0.0, 1129 first_step=0.0, # determined by solver 1130 safety=0.9, 1131 ifactor=10.0, 1132 dfactor=0.2, 1133 beta=0.0, 1134 method=None, 1135 verbosity=-1, # no messages if negative 1136 ): 1137 self.rtol = rtol 1138 self.atol = atol 1139 self.nsteps = nsteps 1140 self.max_step = max_step 1141 self.first_step = first_step 1142 self.safety = safety 1143 self.ifactor = ifactor 1144 self.dfactor = dfactor 1145 self.beta = beta 1146 self.verbosity = verbosity 1147 self.success = 1 1148 self.set_solout(None) 1149 1150 def set_solout(self, solout, complex=False): 1151 self.solout = solout 1152 self.solout_cmplx = complex 1153 if solout is None: 1154 self.iout = 0 1155 else: 1156 self.iout = 1 1157 1158 def reset(self, n, has_jac): 1159 work = zeros((8 * n + 21,), float) 1160 work[1] = self.safety 1161 work[2] = self.dfactor 1162 work[3] = self.ifactor 1163 work[4] = self.beta 1164 work[5] = self.max_step 1165 work[6] = self.first_step 1166 self.work = work 1167 iwork = zeros((21,), _dop_int_dtype) 1168 iwork[0] = self.nsteps 1169 iwork[2] = self.verbosity 1170 self.iwork = iwork 1171 self.call_args = [self.rtol, self.atol, self._solout, 1172 self.iout, self.work, self.iwork] 1173 self.success = 1 1174 1175 def run(self, f, jac, y0, t0, t1, f_params, jac_params): 1176 x, y, iwork, istate = self.runner(*((f, t0, y0, t1) + 1177 tuple(self.call_args) + (f_params,))) 1178 self.istate = istate 1179 if istate < 0: 1180 unexpected_istate_msg = 'Unexpected istate={:d}'.format(istate) 1181 warnings.warn('{:s}: {:s}'.format(self.__class__.__name__, 1182 self.messages.get(istate, unexpected_istate_msg))) 1183 self.success = 0 1184 return y, x 1185 1186 def _solout(self, nr, xold, x, y, nd, icomp, con): 1187 if self.solout is not None: 1188 if self.solout_cmplx: 1189 y = y[::2] + 1j * y[1::2] 1190 return self.solout(x, y) 1191 else: 1192 return 1 1193 1194 1195if dopri5.runner is not None: 1196 IntegratorBase.integrator_classes.append(dopri5) 1197 1198 1199class dop853(dopri5): 1200 runner = getattr(_dop, 'dop853', None) 1201 name = 'dop853' 1202 1203 def __init__(self, 1204 rtol=1e-6, atol=1e-12, 1205 nsteps=500, 1206 max_step=0.0, 1207 first_step=0.0, # determined by solver 1208 safety=0.9, 1209 ifactor=6.0, 1210 dfactor=0.3, 1211 beta=0.0, 1212 method=None, 1213 verbosity=-1, # no messages if negative 1214 ): 1215 super().__init__(rtol, atol, nsteps, max_step, first_step, safety, 1216 ifactor, dfactor, beta, method, verbosity) 1217 1218 def reset(self, n, has_jac): 1219 work = zeros((11 * n + 21,), float) 1220 work[1] = self.safety 1221 work[2] = self.dfactor 1222 work[3] = self.ifactor 1223 work[4] = self.beta 1224 work[5] = self.max_step 1225 work[6] = self.first_step 1226 self.work = work 1227 iwork = zeros((21,), _dop_int_dtype) 1228 iwork[0] = self.nsteps 1229 iwork[2] = self.verbosity 1230 self.iwork = iwork 1231 self.call_args = [self.rtol, self.atol, self._solout, 1232 self.iout, self.work, self.iwork] 1233 self.success = 1 1234 1235 1236if dop853.runner is not None: 1237 IntegratorBase.integrator_classes.append(dop853) 1238 1239 1240class lsoda(IntegratorBase): 1241 runner = getattr(_lsoda, 'lsoda', None) 1242 active_global_handle = 0 1243 1244 messages = { 1245 2: "Integration successful.", 1246 -1: "Excess work done on this call (perhaps wrong Dfun type).", 1247 -2: "Excess accuracy requested (tolerances too small).", 1248 -3: "Illegal input detected (internal error).", 1249 -4: "Repeated error test failures (internal error).", 1250 -5: "Repeated convergence failures (perhaps bad Jacobian or tolerances).", 1251 -6: "Error weight became zero during problem.", 1252 -7: "Internal workspace insufficient to finish (internal error)." 1253 } 1254 1255 def __init__(self, 1256 with_jacobian=False, 1257 rtol=1e-6, atol=1e-12, 1258 lband=None, uband=None, 1259 nsteps=500, 1260 max_step=0.0, # corresponds to infinite 1261 min_step=0.0, 1262 first_step=0.0, # determined by solver 1263 ixpr=0, 1264 max_hnil=0, 1265 max_order_ns=12, 1266 max_order_s=5, 1267 method=None 1268 ): 1269 1270 self.with_jacobian = with_jacobian 1271 self.rtol = rtol 1272 self.atol = atol 1273 self.mu = uband 1274 self.ml = lband 1275 1276 self.max_order_ns = max_order_ns 1277 self.max_order_s = max_order_s 1278 self.nsteps = nsteps 1279 self.max_step = max_step 1280 self.min_step = min_step 1281 self.first_step = first_step 1282 self.ixpr = ixpr 1283 self.max_hnil = max_hnil 1284 self.success = 1 1285 1286 self.initialized = False 1287 1288 def reset(self, n, has_jac): 1289 # Calculate parameters for Fortran subroutine dvode. 1290 if has_jac: 1291 if self.mu is None and self.ml is None: 1292 jt = 1 1293 else: 1294 if self.mu is None: 1295 self.mu = 0 1296 if self.ml is None: 1297 self.ml = 0 1298 jt = 4 1299 else: 1300 if self.mu is None and self.ml is None: 1301 jt = 2 1302 else: 1303 if self.mu is None: 1304 self.mu = 0 1305 if self.ml is None: 1306 self.ml = 0 1307 jt = 5 1308 lrn = 20 + (self.max_order_ns + 4) * n 1309 if jt in [1, 2]: 1310 lrs = 22 + (self.max_order_s + 4) * n + n * n 1311 elif jt in [4, 5]: 1312 lrs = 22 + (self.max_order_s + 5 + 2 * self.ml + self.mu) * n 1313 else: 1314 raise ValueError('Unexpected jt=%s' % jt) 1315 lrw = max(lrn, lrs) 1316 liw = 20 + n 1317 rwork = zeros((lrw,), float) 1318 rwork[4] = self.first_step 1319 rwork[5] = self.max_step 1320 rwork[6] = self.min_step 1321 self.rwork = rwork 1322 iwork = zeros((liw,), _lsoda_int_dtype) 1323 if self.ml is not None: 1324 iwork[0] = self.ml 1325 if self.mu is not None: 1326 iwork[1] = self.mu 1327 iwork[4] = self.ixpr 1328 iwork[5] = self.nsteps 1329 iwork[6] = self.max_hnil 1330 iwork[7] = self.max_order_ns 1331 iwork[8] = self.max_order_s 1332 self.iwork = iwork 1333 self.call_args = [self.rtol, self.atol, 1, 1, 1334 self.rwork, self.iwork, jt] 1335 self.success = 1 1336 self.initialized = False 1337 1338 def run(self, f, jac, y0, t0, t1, f_params, jac_params): 1339 if self.initialized: 1340 self.check_handle() 1341 else: 1342 self.initialized = True 1343 self.acquire_new_handle() 1344 args = [f, y0, t0, t1] + self.call_args[:-1] + \ 1345 [jac, self.call_args[-1], f_params, 0, jac_params] 1346 y1, t, istate = self.runner(*args) 1347 self.istate = istate 1348 if istate < 0: 1349 unexpected_istate_msg = 'Unexpected istate={:d}'.format(istate) 1350 warnings.warn('{:s}: {:s}'.format(self.__class__.__name__, 1351 self.messages.get(istate, unexpected_istate_msg))) 1352 self.success = 0 1353 else: 1354 self.call_args[3] = 2 # upgrade istate from 1 to 2 1355 self.istate = 2 1356 return y1, t 1357 1358 def step(self, *args): 1359 itask = self.call_args[2] 1360 self.call_args[2] = 2 1361 r = self.run(*args) 1362 self.call_args[2] = itask 1363 return r 1364 1365 def run_relax(self, *args): 1366 itask = self.call_args[2] 1367 self.call_args[2] = 3 1368 r = self.run(*args) 1369 self.call_args[2] = itask 1370 return r 1371 1372 1373if lsoda.runner: 1374 IntegratorBase.integrator_classes.append(lsoda) 1375