1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
|
# coding=utf-8
import numpy as np
import sys
from . import series_expansion as se
class ConvergenceError(Exception):
pass
class MaxIterReached(ConvergenceError):
pass
class DivergenceException(ConvergenceError):
pass
def _nonlin_solve_1d_iter(prev, grid, basis_vals, Fs, Fs_args):
order = grid.shape[0]
N = len(Fs)
# evaluate the previous iteration on the grid
prev_vals = []
for diff_order in range(N - 1):
prev_vals.append(prev.eval(grid, diff_order))
# evaluate the RHS functions using the previous iteration
F_vals = []
for F in Fs:
args = [grid, prev_vals]
if Fs_args is not None:
args += Fs_args
F_vals.append(F(*args))
# TODO should be doable with fewer explicit loops
mat = np.copy(basis_vals[-1])
for idx in range(order):
for diff_order in range(N - 1):
mat[:, idx] -= basis_vals[diff_order][:, idx] * F_vals[diff_order + 1]
rhs = F_vals[0] - prev.eval(grid, N - 1)
return np.linalg.solve(mat, rhs)
def nonlin_solve_1d(initial_guess, Fs, args = None, maxiter = 100, atol = 1e-14, grid = None, verbose = True):
"""
Solve a non-linear ODE using a spectral method with Newton iteration.
The ODE is expected to be of the form
d^n u
----- = F(x, u(x), u'(x), u''(x), ... , u^(n-1)(x))
dx^n
Args:
initial_guess (SeriesExpansion): The initial guess, from which the iteration is started.
The spectral basis and its order are inferred from the
initial guess.
Fs (list of callables): The right-hand side of the ODE and its variational derivatives. I.e. Fs[0] is F,
Fs[i] is δF / δu^(i). The order of the equation n is determined by
the length of Fs minus one. Each callable must accept 3 parameters:
- a numpy array containing the values of x
- an iterable of n numpy arrays containing the values of u and
its derivatives at x
- a tuple of additional parameters passed to nonlin_solve_1d as args
args (tuple or None): If not None, extra parameters passed to the Fs.
maxiter: Maximum number of iterations done by the solver.
atol: the difference between the coefficients in subsequent iterations below which convergence
is assumed.
grid: If non-None, the collocation grid to use. If left as None, the default grid from the basis
object is used.
Returns:
The solution to the equation (a SeriesExpansion object).
Raises:
MaxIterReached: If the iteration fails to converge.
"""
N = len(Fs)
order = len(initial_guess.coeffs)
if grid is None:
grid = initial_guess.basis[0].colloc_grid(order)
# precompute the values of the basis functions and their derivatives
# at the grid
basis = initial_guess.basis[0]
basis_vals = []
for diff_order in range(N):
basis_val = np.empty((order, order))
for idx in range(order):
basis_val[:, idx] = basis.eval(idx, grid, diff_order)
basis_vals.append(basis_val)
solution = initial_guess
for i in range(maxiter):
delta = _nonlin_solve_1d_iter(solution, grid, basis_vals, Fs, args)
err = np.max(np.abs(delta))
if not np.isfinite(err):
raise DivergenceException('nan')
if err < atol:
return solution
solution = se.SeriesExpansion(solution.coeffs + delta, solution.basis)
if verbose:
sys.stderr.write('delta: %g, coeffs[0]: %g\n' % (err, solution.coeffs[0]))
raise MaxIterReached('The horizon finder failed to converge')
def nonlin_residual(solution, N, grid, F, args):
"""
Calculate the residual for the solution computed by nonlin_solve_1d() using
2nd order finite differences.
Args:
solution (SeriesExpansion): The solution, as returned by nonlin_solve_1d().
N (int): Order of the equation.
grid (ndarray): An equidistant array of coordinates on which to evaluate the solution.
F: The right-hand side of the ODE.
args: Extra args to pass to F.
"""
sol_vals = [solution.eval(grid)]
dx = abs(grid[1] - grid[0])
for diff_order in range(N):
sol_vals.append(np.gradient(sol_vals[-1], dx))
rhs = F(grid, sol_vals[:-1], args)
return sol_vals[-1] - rhs
|
