Specifying (Non)Linear Solvers

One of the key features of DifferentialEquations.jl is its flexibility. Keeping with this trend, many of the native Julia solvers provided by DifferentialEquations.jl allow you to choose the method for linear and nonlinear solving. This section details how to make that choice.

Linear Solvers: linsolve Specification

For differential equation integrators which use linear solvers, an argument to the method linsolve determines the linear solver which is used. The signature is:

linsolve!(x,A,b,matrix_updated=false)

This is an in-place function which updates x by solving Ax=b. matrix_updated determines whether the matrix A has changed from the last call. This can be used to smartly cache factorizations.

Basic linsolve method: Factorization

The easiest way to specify a linsolve is by a factorization function which generates a type on which \ (or A_ldiv_B!) is called. This is done through the helper function LinSolveFactorize which makes the appropriate function. For example, the Rosenbrock23 takes in a linsolve function, which we can choose to be a QR-factorization by:

Rosenbrock23(linsolve=LinSolveFactorize(qrfact!))

LinSolveFactorize takes in a function which returns an object that can \. Direct methods like qrfact! will automatically cache the factorization, making it efficient for small dense problems.

However, for large sparse problems, you can let \ be an iterative method. For example, using PETSc.jl, we can define our factorization function to be:

linsolve = LinSolveFactorize((A) -> KSP(A, ksp_type="gmres", ksp_rtol=1e-6))

This function creates a KSP type which makes \ perform the GMRES iterative method provided by PETSc.jl. Thus if we pass this function into the algorithm as the factorization method, all internal linear solves will happen by PETSc.jl.

How LinSolveFactorize Was Created

In order to make your own linsolve functions, let's look at how the LinSolveFactorize function is created. For example, for an LU-Factorization, we would like to use lufact! to do our linear solving. We can directly write this as:

function linsolve!(x,A,b,update_matrix=false)
  _A = lufact!(A)
  A_ldiv_B!(x,_A,b)
end

This method works fine and you can pass it to the methods like

Rosenbrock23(linsolve=linsolve!)

and it will work, but this method does not cache _A, the factorization. This means that, even if A has not changed, it will re-factorize the matrix.

To change this, we can instead create a call-overloaded type. The generalized form of this is:

type LinSolveFactorize{F}
  factorization::F
  A
end
LinSolveFactorize(factorization) = LinSolveFactorize(factorization,nothing)
function (p::LinSolveFactorize)(x,A,b,matrix_updated=false)
  if matrix_updated
    p.A = p.factorization(A)
  end
  A_ldiv_B!(x,p.A,b)
end
linsolve = LinSolveFactorize(lufact!)

LinSolveFactorize is a type which holds the factorization method and the pre-factorized matrix. If matrix_updated is true, it will re-compute the factorization. Otherwise it just solves the linear system with the cached factorization. This general idea of using a call-overloaded type can be employed to do many other things.

Nonlinear Solvers: nlsolve Specification

Nonlinear solvers can be chosen via the nlsolve option. An nlsolve function should have two dispatches:

• nlsolve(Val{init},f,u0_prototype) : Does an initialization phase. Returns a type init_f for later use in the solver. u0_prototype is the expected type for the initial condition u0.

• nlsolve(init_f,u0) : Solves for the root units the initialized f and the initial condition u0. Returns the zeros of the equation.

Basic nlsolve mehtod: NLSOLVEJL_SETUP

By default, a basic nonlinear solver setup is given as NLSOLVEJL_SETUP. For example, the default nlsolve in Trapezoid is

Trapezoid(nlsolve=NLSOLVEJL_SETUP())

This will use NLsolve.jl with autodifferentiation to solve the nonlinear systems. NLSOLVEJL_SETUP has two options:

• chunk_size : The autodifferentiation chunk size. Integer. Defaults to ForwardDiff.jl's auto-detection.

• autodiff : Whether to use autodifferentiation. Defaults to true.

For example, to turn off autodifferentiation, use

Trapezoid(nlsolve=NLSOLVEJL_SETUP(autodiff=false))

How NLSOLVEJL_SETUP Was Created

To create a nonlinear solver, you need to define the two functions. Here we use a call-overloaded type so that way we can hold the chunk size and autodifferentiation information.

immutable NLSOLVEJL_SETUP{CS,AD} end
Base.@pure NLSOLVEJL_SETUP(;chunk_size=0,autodiff=true) = NLSOLVEJL_SETUP{chunk_size,autodiff}()

The solver function just calls NLsolve and returns the zeros

(p::NLSOLVEJL_SETUP)(f,u0) = (res=NLsolve.nlsolve(f,u0); res.zero)

while the initialization function has a different initialization for autodifferentiation or not:

function (p::NLSOLVEJL_SETUP{CS,AD}){CS,AD}(::Type{Val{:init}},f,u0_prototype)
  if AD
    return non_autodiff_setup(f,u0_prototype)
  else
    return autodiff_setup(f,u0_prototype,Val{determine_chunksize(initial_x,CS)})
  end
end

We need to declare the get_chunksize trait for the solver:

get_chunksize{CS,AD}(x::NLSOLVEJL_SETUP{CS,AD}) = CS

The initialization functions are directly for NLsolve. See the NLsolve.jl docs for the types of inputs it expects to see. This does exactly that:

function autodiff_setup{CS}(f!, initial_x::Vector,chunk_size::Type{Val{CS}})

    permf! = (fx, x) -> f!(x, fx)

    fx2 = copy(initial_x)
    jac_cfg = ForwardDiff.JacobianConfig{CS}(initial_x, initial_x)
    g! = (x, gx) -> ForwardDiff.jacobian!(gx, permf!, fx2, x, jac_cfg)

    fg! = (x, fx, gx) -> begin
        jac_res = DiffBase.DiffResult(fx, gx)
        ForwardDiff.jacobian!(jac_res, permf!, fx2, x, jac_cfg)
        DiffBase.value(jac_res)
    end

    return DifferentiableMultivariateFunction(f!, g!, fg!)
end

function non_autodiff_setup(f!, initial_x::Vector)
  DifferentiableMultivariateFunction(f!)
end