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
|
import numpy as np
import series_expansion
def _nonlin_solve_1d_iter(prev, grid, basis_vals, Fs, args):
order = grid.shape[0]
N = len(Fs)
# evaluate the previous iteration on the grid
prev_vals = []
for diff_order in xrange(N - 1):
prev_vals.append(prev.eval(grid, diff_order))
# evaluate the RHS functions using the previous iteration
F_vals = []
for F in Fs:
F_vals.append(F(grid, prev_vals, args))
# TODO should be doable with fewer explicit loops
mat = np.copy(basis_vals[-1])
for idx in xrange(order):
for diff_order in xrange(N - 1):
mat[:, idx] -= basis_vals[diff_order][:, idx] * F_vals[diff_order + 1]
rhs = F_vals[0]
for diff_order in xrange(N - 1):
rhs -= F_vals[diff_order + 1] * prev_vals[diff_order]
return series_expansion.SeriesExpansion(np.linalg.solve(mat, rhs), prev.basis)
def nonlin_solve_1d(initial_guess, Fs, args, maxiter = 100, atol = 1e-14, grid = None):
"""
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 derivatives. I.e. Fs[0] is F,
Fs[i] is d^i F / dx^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 F and
its derivatives at x
- a tuple of additional parameters passed to nonlin_solve_1d as args
args (tuple): 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:
RuntimeError: 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 xrange(N):
basis_val = np.empty((order, order))
for idx in xrange(order):
basis_val[:, idx] = basis.eval(idx, grid, diff_order)
basis_vals.append(basis_val)
solution_old = initial_guess
for i in xrange(maxiter):
solution_new = _nonlin_solve_1d_iter(solution_old, grid, basis_vals, Fs, args)
delta = np.max(np.abs(solution_new.coeffs - solution_old.coeffs))
print delta
if delta < atol:
return solution_new
solution_old = solution_new
raise RuntimeError('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 xrange(N):
sol_vals.append(np.gradient(sol_vals[-1], dx))
rhs = F(grid, sol_vals[:-1], args)
return sol_vals[-1] - rhs
|
