From 52c429db82eb32cea648295a741e49644f9b1b00 Mon Sep 17 00:00:00 2001 From: Anton Khirnov Date: Mon, 2 Dec 2019 10:21:38 +0100 Subject: nonlin_ode: make args optional --- nonlin_ode.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/nonlin_ode.py b/nonlin_ode.py index dd10385..4cb5c2e 100644 --- a/nonlin_ode.py +++ b/nonlin_ode.py @@ -5,7 +5,7 @@ import sys import series_expansion -def _nonlin_solve_1d_iter(prev, grid, basis_vals, Fs, args): +def _nonlin_solve_1d_iter(prev, grid, basis_vals, Fs, Fs_args): order = grid.shape[0] N = len(Fs) @@ -17,7 +17,10 @@ def _nonlin_solve_1d_iter(prev, grid, basis_vals, Fs, args): # evaluate the RHS functions using the previous iteration F_vals = [] for F in Fs: - F_vals.append(F(grid, prev_vals, args)) + 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]) @@ -31,7 +34,7 @@ def _nonlin_solve_1d_iter(prev, grid, basis_vals, Fs, args): 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, verbose = True): +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. @@ -51,7 +54,7 @@ def nonlin_solve_1d(initial_guess, Fs, args, maxiter = 100, atol = 1e-14, grid = - 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): Extra parameters passed to the Fs. + 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. -- cgit v1.2.3