nonlin_ode: make args optional

1 files 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.