Add a spectral solver of non-linear ODEs.

1 files changed, 106 insertions, 0 deletions

diff --git a/nonlin_ode.py b/nonlin_ode.py
new file mode 100644
index 0000000..4944541
--- /dev/null
+++ b/nonlin_ode.py
@@ -0,0 +1,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