Coverage for /usr/lib/python3/dist-packages/mpmath/calculus/odes.py: 9%
93 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1from bisect import bisect
2from ..libmp.backend import xrange
4class ODEMethods(object):
5 pass
7def ode_taylor(ctx, derivs, x0, y0, tol_prec, n):
8 h = tol = ctx.ldexp(1, -tol_prec)
9 dim = len(y0)
10 xs = [x0]
11 ys = [y0]
12 x = x0
13 y = y0
14 orig = ctx.prec
15 try:
16 ctx.prec = orig*(1+n)
17 # Use n steps with Euler's method to get
18 # evaluation points for derivatives
19 for i in range(n):
20 fxy = derivs(x, y)
21 y = [y[i]+h*fxy[i] for i in xrange(len(y))]
22 x += h
23 xs.append(x)
24 ys.append(y)
25 # Compute derivatives
26 ser = [[] for d in range(dim)]
27 for j in range(n+1):
28 s = [0]*dim
29 b = (-1) ** (j & 1)
30 k = 1
31 for i in range(j+1):
32 for d in range(dim):
33 s[d] += b * ys[i][d]
34 b = (b * (j-k+1)) // (-k)
35 k += 1
36 scale = h**(-j) / ctx.fac(j)
37 for d in range(dim):
38 s[d] = s[d] * scale
39 ser[d].append(s[d])
40 finally:
41 ctx.prec = orig
42 # Estimate radius for which we can get full accuracy.
43 # XXX: do this right for zeros
44 radius = ctx.one
45 for ts in ser:
46 if ts[-1]:
47 radius = min(radius, ctx.nthroot(tol/abs(ts[-1]), n))
48 radius /= 2 # XXX
49 return ser, x0+radius
51def odefun(ctx, F, x0, y0, tol=None, degree=None, method='taylor', verbose=False):
52 r"""
53 Returns a function `y(x) = [y_0(x), y_1(x), \ldots, y_n(x)]`
54 that is a numerical solution of the `n+1`-dimensional first-order
55 ordinary differential equation (ODE) system
57 .. math ::
59 y_0'(x) = F_0(x, [y_0(x), y_1(x), \ldots, y_n(x)])
61 y_1'(x) = F_1(x, [y_0(x), y_1(x), \ldots, y_n(x)])
63 \vdots
65 y_n'(x) = F_n(x, [y_0(x), y_1(x), \ldots, y_n(x)])
67 The derivatives are specified by the vector-valued function
68 *F* that evaluates
69 `[y_0', \ldots, y_n'] = F(x, [y_0, \ldots, y_n])`.
70 The initial point `x_0` is specified by the scalar argument *x0*,
71 and the initial value `y(x_0) = [y_0(x_0), \ldots, y_n(x_0)]` is
72 specified by the vector argument *y0*.
74 For convenience, if the system is one-dimensional, you may optionally
75 provide just a scalar value for *y0*. In this case, *F* should accept
76 a scalar *y* argument and return a scalar. The solution function
77 *y* will return scalar values instead of length-1 vectors.
79 Evaluation of the solution function `y(x)` is permitted
80 for any `x \ge x_0`.
82 A high-order ODE can be solved by transforming it into first-order
83 vector form. This transformation is described in standard texts
84 on ODEs. Examples will also be given below.
86 **Options, speed and accuracy**
88 By default, :func:`~mpmath.odefun` uses a high-order Taylor series
89 method. For reasonably well-behaved problems, the solution will
90 be fully accurate to within the working precision. Note that
91 *F* must be possible to evaluate to very high precision
92 for the generation of Taylor series to work.
94 To get a faster but less accurate solution, you can set a large
95 value for *tol* (which defaults roughly to *eps*). If you just
96 want to plot the solution or perform a basic simulation,
97 *tol = 0.01* is likely sufficient.
99 The *degree* argument controls the degree of the solver (with
100 *method='taylor'*, this is the degree of the Taylor series
101 expansion). A higher degree means that a longer step can be taken
102 before a new local solution must be generated from *F*,
103 meaning that fewer steps are required to get from `x_0` to a given
104 `x_1`. On the other hand, a higher degree also means that each
105 local solution becomes more expensive (i.e., more evaluations of
106 *F* are required per step, and at higher precision).
108 The optimal setting therefore involves a tradeoff. Generally,
109 decreasing the *degree* for Taylor series is likely to give faster
110 solution at low precision, while increasing is likely to be better
111 at higher precision.
113 The function
114 object returned by :func:`~mpmath.odefun` caches the solutions at all step
115 points and uses polynomial interpolation between step points.
116 Therefore, once `y(x_1)` has been evaluated for some `x_1`,
117 `y(x)` can be evaluated very quickly for any `x_0 \le x \le x_1`.
118 and continuing the evaluation up to `x_2 > x_1` is also fast.
120 **Examples of first-order ODEs**
122 We will solve the standard test problem `y'(x) = y(x), y(0) = 1`
123 which has explicit solution `y(x) = \exp(x)`::
125 >>> from mpmath import *
126 >>> mp.dps = 15; mp.pretty = True
127 >>> f = odefun(lambda x, y: y, 0, 1)
128 >>> for x in [0, 1, 2.5]:
129 ... print((f(x), exp(x)))
130 ...
131 (1.0, 1.0)
132 (2.71828182845905, 2.71828182845905)
133 (12.1824939607035, 12.1824939607035)
135 The solution with high precision::
137 >>> mp.dps = 50
138 >>> f = odefun(lambda x, y: y, 0, 1)
139 >>> f(1)
140 2.7182818284590452353602874713526624977572470937
141 >>> exp(1)
142 2.7182818284590452353602874713526624977572470937
144 Using the more general vectorized form, the test problem
145 can be input as (note that *f* returns a 1-element vector)::
147 >>> mp.dps = 15
148 >>> f = odefun(lambda x, y: [y[0]], 0, [1])
149 >>> f(1)
150 [2.71828182845905]
152 :func:`~mpmath.odefun` can solve nonlinear ODEs, which are generally
153 impossible (and at best difficult) to solve analytically. As
154 an example of a nonlinear ODE, we will solve `y'(x) = x \sin(y(x))`
155 for `y(0) = \pi/2`. An exact solution happens to be known
156 for this problem, and is given by
157 `y(x) = 2 \tan^{-1}\left(\exp\left(x^2/2\right)\right)`::
159 >>> f = odefun(lambda x, y: x*sin(y), 0, pi/2)
160 >>> for x in [2, 5, 10]:
161 ... print((f(x), 2*atan(exp(mpf(x)**2/2))))
162 ...
163 (2.87255666284091, 2.87255666284091)
164 (3.14158520028345, 3.14158520028345)
165 (3.14159265358979, 3.14159265358979)
167 If `F` is independent of `y`, an ODE can be solved using direct
168 integration. We can therefore obtain a reference solution with
169 :func:`~mpmath.quad`::
171 >>> f = lambda x: (1+x**2)/(1+x**3)
172 >>> g = odefun(lambda x, y: f(x), pi, 0)
173 >>> g(2*pi)
174 0.72128263801696
175 >>> quad(f, [pi, 2*pi])
176 0.72128263801696
178 **Examples of second-order ODEs**
180 We will solve the harmonic oscillator equation `y''(x) + y(x) = 0`.
181 To do this, we introduce the helper functions `y_0 = y, y_1 = y_0'`
182 whereby the original equation can be written as `y_1' + y_0' = 0`. Put
183 together, we get the first-order, two-dimensional vector ODE
185 .. math ::
187 \begin{cases}
188 y_0' = y_1 \\
189 y_1' = -y_0
190 \end{cases}
192 To get a well-defined IVP, we need two initial values. With
193 `y(0) = y_0(0) = 1` and `-y'(0) = y_1(0) = 0`, the problem will of
194 course be solved by `y(x) = y_0(x) = \cos(x)` and
195 `-y'(x) = y_1(x) = \sin(x)`. We check this::
197 >>> f = odefun(lambda x, y: [-y[1], y[0]], 0, [1, 0])
198 >>> for x in [0, 1, 2.5, 10]:
199 ... nprint(f(x), 15)
200 ... nprint([cos(x), sin(x)], 15)
201 ... print("---")
202 ...
203 [1.0, 0.0]
204 [1.0, 0.0]
205 ---
206 [0.54030230586814, 0.841470984807897]
207 [0.54030230586814, 0.841470984807897]
208 ---
209 [-0.801143615546934, 0.598472144103957]
210 [-0.801143615546934, 0.598472144103957]
211 ---
212 [-0.839071529076452, -0.54402111088937]
213 [-0.839071529076452, -0.54402111088937]
214 ---
216 Note that we get both the sine and the cosine solutions
217 simultaneously.
219 **TODO**
221 * Better automatic choice of degree and step size
222 * Make determination of Taylor series convergence radius
223 more robust
224 * Allow solution for `x < x_0`
225 * Allow solution for complex `x`
226 * Test for difficult (ill-conditioned) problems
227 * Implement Runge-Kutta and other algorithms
229 """
230 if tol:
231 tol_prec = int(-ctx.log(tol, 2))+10
232 else:
233 tol_prec = ctx.prec+10
234 degree = degree or (3 + int(3*ctx.dps/2.))
235 workprec = ctx.prec + 40
236 try:
237 len(y0)
238 return_vector = True
239 except TypeError:
240 F_ = F
241 F = lambda x, y: [F_(x, y[0])]
242 y0 = [y0]
243 return_vector = False
244 ser, xb = ode_taylor(ctx, F, x0, y0, tol_prec, degree)
245 series_boundaries = [x0, xb]
246 series_data = [(ser, x0, xb)]
247 # We will be working with vectors of Taylor series
248 def mpolyval(ser, a):
249 return [ctx.polyval(s[::-1], a) for s in ser]
250 # Find nearest expansion point; compute if necessary
251 def get_series(x):
252 if x < x0:
253 raise ValueError
254 n = bisect(series_boundaries, x)
255 if n < len(series_boundaries):
256 return series_data[n-1]
257 while 1:
258 ser, xa, xb = series_data[-1]
259 if verbose:
260 print("Computing Taylor series for [%f, %f]" % (xa, xb))
261 y = mpolyval(ser, xb-xa)
262 xa = xb
263 ser, xb = ode_taylor(ctx, F, xb, y, tol_prec, degree)
264 series_boundaries.append(xb)
265 series_data.append((ser, xa, xb))
266 if x <= xb:
267 return series_data[-1]
268 # Evaluation function
269 def interpolant(x):
270 x = ctx.convert(x)
271 orig = ctx.prec
272 try:
273 ctx.prec = workprec
274 ser, xa, xb = get_series(x)
275 y = mpolyval(ser, x-xa)
276 finally:
277 ctx.prec = orig
278 if return_vector:
279 return [+yk for yk in y]
280 else:
281 return +y[0]
282 return interpolant
284ODEMethods.odefun = odefun
286if __name__ == "__main__":
287 import doctest
288 doctest.testmod()