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 typeinit_f
for later use in the solver.u0_prototype
is the expected type for the initial conditionu0
.nlsolve(init_f,u0)
: Solves for the root units the initializedf
and the initial conditionu0
. 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