How does one turn off Autodiff +Jacobian in DifferentialEquations.jl solvers?

General Usage
1

EDIT

I’m trying to solve a simple mixed point BVP such as:

\partial_z \mathbf{a} = C_1 \mathbf{a} + g \mathbf{a} \mathbf{b} \mathbf{b}^{*}

\partial_z \mathbf{b} = C_2 \mathbf{b} + g \mathbf{b} \mathbf{a} \mathbf{a}^{*}

with BCs:

\mathbf{a}(z=0) = \mathbf{A}

\mathbf{b}(z=L) = \mathbf{B}

Note that \mathbf{a} and \mathbf{b} are complex valued vectors so this is a large system of boundary value problem. Usually, this problem can be solved via fixed-point iteration.

Below is a currently working code that I’ve converted completely to real values and cleaned up as a minimal example. Now, I have the following issues when increasing the magnitude of the nonlinear factor gg:

  • MIRK4 crashes Julia
  • Switching to a FIRK such as:

alg = RadauIIa5(;jac_alg = BVPJacobianAlgorithm(bc_diffmode = AutoFiniteDiff(), nonbc_diffmode = AutoFiniteDiff()), nested_nlsolve=true)

gives an error:

MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{NonlinearFunction{true, SciMLBase.FullSpecialize, BoundaryValueDiffEqFIRK.var"#80#94"{BVProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Nothing, Params, BVPFunction{true, SciMLBase.FullSpecialize, false, typeof(f_packed!), typeof(bc_packed!), Nothing, Nothing, Nothing, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing, Nothing, Dict{Any, Any}}, Nothing}, Nothing, Nothing, SciMLBase.StandardBVProblem, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, BVPFunction{true, SciMLBase.FullSpecialize, false, typeof(f_packed!), typeof(bc_packed!), Nothing, Nothing, Nothing, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing, Nothing, Dict{Any, Any}}, Nothing}, BoundaryValueDiffEqFIRK.FIRKTableau{true, Int64, Matrix{Float64}, Vector{Float64}, Vector{Float64}}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 12})
The type `Float64` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:265
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:900
  Float64(::IrrationalConstants.Sqrthalfπ)
   @ IrrationalConstants ~/.julia/packages/IrrationalConstants/lWTip/src/macro.jl:131
  ...

I want to be able to solve a large system where a and b are large (2^14) so I was hoping FIRKs would work.

Edit Working Code:

using BoundaryValueDiffEq
using DiffEqBase, OrdinaryDiffEq


# ---------- parameters ----------
mutable struct Params
    M::Int
    C::NTuple{2,ComplexF64}       # linear coefficients for a,b
    g::NTuple{2,ComplexF64}       # nonlinear coefficients
    zbc::NTuple{2,Float64}        # (z_left, z_right)
    A_re::Vector{Float64}         # target a(0): real part
    A_im::Vector{Float64}         # target a(0): imag part
    B_re::Vector{Float64}         # target b(L): real part
    B_im::Vector{Float64}         # target b(L): imag part
    tmp_b2::Vector{Float64}       # workspace: |b|^2
end

# ---------- RHS on packed real state ----------
# v = [a_r(1:M); a_i(1:M); b_r(1:M); b_i(1:M)]
function f_packed!(dv::AbstractVector{T}, v::AbstractVector{T}, P::Params, z) where {T}
    M = P.M
    ar = @view v[1:M];        ai = @view v[M+1:2M]
    br = @view v[2M+1:3M];    bi = @view v[3M+1:4M]

    dar = @view dv[1:M];      dai = @view dv[M+1:2M]
    dbr = @view dv[2M+1:3M];  dbi = @view dv[3M+1:4M]

    # |b|^2
    @inbounds @simd for i in 1:M
        P.tmp_b2[i] = br[i]*br[i] + bi[i]*bi[i]
    end

    C1, C2 = P.C
    g1, g2 = P.g

    @inbounds @simd for i in 1:M
        ar_i = ar[i]; ai_i = ai[i]
        br_i = br[i]; bi_i = bi[i]
        b2   = P.tmp_b2[i]                # real scalar

        # ---- da/dz = (C1 + g1*|b|^2)*a
        s1   = C1 + g1*b2                 # complex scalar
        s1r  = real(s1); s1i = imag(s1)
        dar[i] = s1r*ar_i - s1i*ai_i
        dai[i] = s1r*ai_i + s1i*ar_i

        # ---- db/dz = C2*b + (g2*|b|^2)*a
        # linear part on b
        c2r = real(C2); c2i = imag(C2)
        dbr[i] = c2r*br_i - c2i*bi_i
        dbi[i] = c2r*bi_i + c2i*br_i
        # nonlinear part proportional to a
        t2   = g2*b2
        t2r  = real(t2); t2i = imag(t2)
        dbr[i] += t2r*ar_i - t2i*ai_i
        dbi[i] += t2r*ai_i + t2i*ar_i
    end
    return nothing
end

# ---------- boundary conditions on packed real state ----------
# Require: a(0) = A, b(L) = B  (complex equalities → two real equations each)
function bc_packed!(res, u, P::Params, args...)
    M = P.M
    # left boundary (z = zbc[1]): a(0) = A
    uL = u(P.zbc[1])                     # Vector{Float64}, length 4M
    @inbounds @simd for i in 1:M
        res[i]      = uL[i]        - P.A_re[i]        # a_r(0) - A_r = 0
        res[M + i]  = uL[M + i]    - P.A_im[i]        # a_i(0) - A_i = 0
    end
    # right boundary (z = zbc[2]): b(L) = B
    uR = u(P.zbc[2])
    off = 2M
    @inbounds @simd for i in 1:M
        res[2M + i]     = uR[off + i]     - P.B_re[i] # b_r(L) - B_r = 0
        res[3M + i]     = uR[off + M + i] - P.B_im[i] # b_i(L) - B_i = 0
    end
    return nothing
end

# helper: split a packed state vector v (Float64, length 4M) into Complex a,b (length M each)
function unpack_state(v::AbstractVector{<:Real}, M::Int)
    a_re = @view v[1:M]
    a_im = @view v[M+1:2M]
    b_re = @view v[2M+1:3M]
    b_im = @view v[3M+1:4M]
    a = complex.(a_re, a_im)
    b = complex.(b_re, b_im)
    return a, b
end

# ---------- problem setup ----------
M = 2^5
@show M
L = 2.0
x = collect(range(-5.0, 5.0; length=M))

A  = @. exp(-x^2)                    # complex A(x) with 0 imag
B  = @. 1.0e-6 * exp(-x^2)

A_re = Float64.(real.(A));  A_im = Float64.(imag.(A))
B_re = Float64.(real.(B));  B_im = Float64.(imag.(B))

cc = -1.0e-1
C = (cc + 0im, cc + 0im)
gg = -1.0e-4
g = (gg + 0.0im, gg + 0.0im)

P = Params(M, C, g, (0.0, L), A_re, A_im, B_re, B_im, zeros(M))

# initial guess, packed real
y0 = vcat(A_re, A_im, B_re, B_im);

prob = BVProblem(f_packed!, bc_packed!, y0, (0.0, L), P)
alg = MIRK4(;jac_alg = BVPJacobianAlgorithm(bc_diffmode = AutoFiniteDiff(), nonbc_diffmode = AutoFiniteDiff()))
sol = solve(prob, alg;  dt= L/100)

OLD

Below I have a code where I collected both vectors into one and made it into real values u. The RHS function converts the single vector into the two complex vectors does the computation then repacks them back into u and similarly for the mixed point boundary conditions.

The main errors I’m getting come from the Autodiff and Jacobian computaitons which are holding me back. So I’m wondering the following:

  • Can BVProblem handle complex values?
  • How do I turn off Autodiff in the usage of hte Jacobian?
  • What other ways can I use to solve this?

CODE: Note that I have some fftw stuff that I’m using for future problems so ignor this



# Uses Params work buffers, no allocations.
# ---------- parameters & nonlinearity kernel ----------
mutable struct Params
   
    M::Int
    C::AbstractArray
    g::AbstractArray
    zbc::Vector{Float64}
    bcL::Vector{ComplexF64} # boundary condition at z=0
    bcL_pack::Vector{Float64} # packed boundary condition at z=0
    bcR::Vector{ComplexF64} # boundary condition at z=L
    bcR_pack::Vector{Float64} # packed boundary condition at z=L
    
    plan_f::Any #FFTW.cFFTWPlan #{ComplexF64,-1,false,1}  # fft!
    plan_i::Any #FFTW.cFFTWPlan #{ComplexF64, 1,false,1}  # ifft!
    xa::Vector{ComplexF64} # work: fft(a)
    xb::Vector{ComplexF64} # work: fft(b)
    tmp::Vector{ComplexF64}# work: pointwise
    
end

# Define Na, Nb from N; edit these to your physics.
function Na!(out, a, b, P::Params)
    # Nxy!(out, a, b, P)         # example: same as N(a,b)
    @inbounds @simd for i in 1:P.M
        out[i] = P.g[1]*a[i]*b[i]*conj(b[i])  # example: Na(a,b) = a*b
    end
end
function Nb!(out, a, b, P::Params)
    # Nxy!(out, b, a, P)         # example: N(b,a); change as needed
    @inbounds @simd for i in 1:P.M
        out[i] = P.g[2]*a[i]*b[i]*conj(b[i])  # example: Na(a,b) = a*b
    end
end

# ---------- ODE RHS ----------
# state u = [a; b] length 2M (Complex)
function f!(du, u, P::Params, z)
    M = P.M
    a = @view u[1:M]
    b = @view u[M+1:2M]
    da = @view du[1:M]
    db = @view du[M+1:2M]

    # da = R*a + Na(a,b)
    @inbounds @simd for i in 1:M
        da[i] = P.C[1] * a[i]
    end
    Na!(P.tmp, a, b, P)               # tmp reuse
    @inbounds @simd for i in 1:M
        da[i] += P.tmp[i]
    end

    # db = S*b + Nb(a,b)
    @inbounds @simd for i in 1:M
        db[i] = P.C[2]* b[i]
    end
    Nb!(P.tmp, b, a, P)               # example Nb = N(b,a); swap to Nb! if different
    @inbounds @simd for i in 1:M
        db[i] += P.tmp[i]
    end
    return nothing
end



# (a+ib) -> [Re; Im]
pack(u::AbstractVector{T}) where {T<:Complex} = reinterpret(real(T), u)

# [Re; Im] -> (a+ib)
unpack(v::AbstractVector{T}) where {T<:Real} = reinterpret(Complex{T}, v)


function f_real!(dv, v, p, z)
    u  = unpack(v)             # Complex view (no copy)
    du = similar(u)
    f!(du, u, p, z)            # your complex RHS
    dv .= pack(du)             # write back as real
    return nothing
end




function bc_real!(res,
                  u,
                  P::Params,
                  args...)
    M2 = 2 * P.M


    @inbounds @simd for i in 1:M2
        res[i] = u(P.zbc[1])[i] - P.bcL_pack[i]          # a(0) = A
    end
    off = M2
    @inbounds @simd for i in 1:M2
        res[off + i] = u(P.zbc[2])[off + i] - P.bcR_pack[i]  # b(L) = B
    end
    return nothing
end




M = 2^6
L = 10.0  # length of integration domain
zspan = (0.0, L)
zbc = [zspan[1], zspan[2]]  # boundary points
x = Array(range( -10.0, 10.0, length=M))
A = @. exp(-x^2) + 0.0im; # example: A(x) = exp(x)
A_pack = pack(A)
B = @. 1.0e-6 * exp(-x^2) + 0.0im; # example: B(x) = exp(x)
B_pack = pack(B)
C = [-1.0e-6, -1.0e-6]
g = [-1.0+1.0im, -1.0+1.0im]

# work buffers
xa  = similar(ComplexF64[], M); resize!(xa,M)
xb  = similar(ComplexF64[], M); resize!(xb,M)
tmp = similar(ComplexF64[], M); resize!(tmp,M)
# plans in-place (reuse buffers)
plan_f = FFTW.plan_fft!(xa; flags=FFTW.MEASURE)
plan_i = FFTW.plan_ifft!(xa; flags=FFTW.MEASURE)  # we’ll swap buffers as needed
P = Params(M,C,g, zbc, A, A_pack, B, B_pack, plan_f,plan_i,xa,xb,tmp);
initguess = vcat(A,B);
initguess_pack = pack(initguess)  ; # packed initial guess

prob = BVProblem(f_real!, bc_real!, initguess_pack, zspan, P)
sol  = solve(prob,  MIRK4();   dt = L/100)

and the following error:

{
	"name": "LoadError",
	"message": "MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{BoundaryValueDiffEqMIRK.var\"#46#52\"{BoundaryValueDiffEqMIRK.MIRKCache{true, Float64, false, BoundaryValueDiffEqCore.DiffCacheNeeded, false, Tuple{Int64}, BVPFunction{true, SciMLBase.FullSpecialize, false, typeof(f_real!), typeof(bc_real!), Nothing, Nothing, Nothing, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing, Nothing, Dict{Any, Any}}, Nothing}, typeof(bc_real!), BVProblem{Base.ReinterpretArray{Float64, 1, ComplexF64, Vector{ComplexF64}, false}, Tuple{Float64, Float64}, true, Nothing, Params, BVPFunction{true, SciMLBase.FullSpecialize, false, typeof(f_real!), typeof(bc_real!), Nothing, Nothing, Nothing, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing, Nothing, Dict{Any, Any}}, Nothing}, Nothing, Nothing, SciMLBase.StandardBVProblem, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, SciMLBase.StandardBVProblem, Params, MIRK4{Nothing, BVPJacobianAlgorithm{AutoSparse{AutoForwardDiff{nothing, Nothing}, SparseConnectivityTracer.TracerLocalSparsityDetector{SparseConnectivityTracer.GradientTracer{Int64, BitSet}, SparseConnectivityTracer.HessianTracer{Int64, BitSet, Dict{Int64, BitSet}, SparseConnectivityTracer.NotShared}}, SparseMatrixColorings.GreedyColoringAlgorithm{:direct, SparseMatrixColorings.NaturalOrder}}, AutoSparse{AutoForwardDiff{nothing, Nothing}, SparseConnectivityTracer.TracerLocalSparsityDetector{SparseConnectivityTracer.GradientTracer{Int64, BitSet}, SparseConnectivityTracer.HessianTracer{Int64, BitSet, Dict{Int64, BitSet}, SparseConnectivityTracer.NotShared}}, SparseMatrixColorings.GreedyColoringAlgorithm{:direct, SparseMatrixColorings.NaturalOrder}}, Nothing}, Float64}, BoundaryValueDiffEqMIRK.MIRKTableau{Int64, Vector{Float64}, Vector{Float64}, Vector{Float64}, Matrix{Float64}}, BoundaryValueDiffEqMIRK.MIRKInterpTableau{Int64, Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64, Int64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{PreallocationTools.DiffCache{Matrix{Float64}, Vector{Float64}}}, VectorOfArray{Float64, 3, Vector{Matrix{Float64}}}, Vector{PreallocationTools.DiffCache{Vector{Float64}, Vector{Float64}}}, VectorOfArray{Float64, 2, Vector{Vector{Float64}}}, Vector{PreallocationTools.DiffCache{Vector{Float64}, Vector{Float64}}}, PreallocationTools.DiffCache{Vector{Float64}, Vector{Float64}}, Vector{Float64}, VectorOfArray{Float64, 2, Vector{Vector{Float64}}}, VectorOfArray{Float64, 2, Vector{Vector{Float64}}}, Tuple{Int64}, @NamedTuple{abstol::Float64}, @NamedTuple{abstol::Float64, dt::Float64, adaptive::Bool, controller::BoundaryValueDiffEqCore.DefectControl{Float64}, fit_parameters::Bool}}, BoundaryValueDiffEqCore.DiffCacheNeeded}, Float64}, Float64, 12})\nThe type `Float64` exists, but no method is defined for this combination of argument types when trying to construct it.\n\n\u001b[0mClosest candidates are:\n\u001b[0m  (::Type{T})(::Real, \u001b[91m::RoundingMode\u001b[39m) where T<:AbstractFloat\n\u001b[0m\u001b[90m   @\u001b[39m \u001b[90mBase\u001b[39m \u001b[90m\u001b[4mrounding.jl:265\u001b[24m\u001b[39m\n\u001b[0m  (::Type{T})(::T) where T<:Number\n\u001b[0m\u001b[90m   @\u001b[39m \u001b[90mCore\u001b[39m \u001b[90m\u001b[4mboot.jl:900\u001b[24m\u001b[39m\n\u001b[0m  Float64(\u001b[91m::IrrationalConstants.Sqrthalfπ\u001b[39m)\n\u001b[0m\u001b[90m   @\u001b[39m \u001b[33mIrrationalConstants\u001b[39m \u001b[90m~/.julia/packages/IrrationalConstants/lWTip/src/\u001b[39m\u001b[90m\u001b[4mmacro.jl:131\u001b[24m\u001b[39m\n\u001b[0m  ...\n",
	"stack": "MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{BoundaryValueDiffEqMIRK.var\"#46#52\"{BoundaryValueDiffEqMIRK.MIRKCache{true, Float64, false, BoundaryValueDiffEqCore.DiffCacheNeeded, false, Tuple{Int64}, BVPFunction{true, SciMLBase.FullSpecialize, false, typeof(f_real!), typeof(bc_real!), Nothing, Nothing, Nothing, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing, Nothing, Dict{Any, Any}}, Nothing}, typeof(bc_real!), BVProblem{Base.ReinterpretArray{Float64, 1, ComplexF64, Vector{ComplexF64}, false}, Tuple{Float64, Float64}, true, Nothing, Params, BVPFunction{true, SciMLBase.FullSpecialize, false, typeof(f_real!), typeof(bc_real!), Nothing, Nothing, Nothing, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing, Nothing, Dict{Any, Any}}, Nothing}, Nothing, Nothing, SciMLBase.StandardBVProblem, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, SciMLBase.StandardBVProblem, Params, MIRK4{Nothing, BVPJacobianAlgorithm{AutoSparse{AutoForwardDiff{nothing, Nothing}, SparseConnectivityTracer.TracerLocalSparsityDetector{SparseConnectivityTracer.GradientTracer{Int64, BitSet}, SparseConnectivityTracer.HessianTracer{Int64, BitSet, Dict{Int64, BitSet}, SparseConnectivityTracer.NotShared}}, SparseMatrixColorings.GreedyColoringAlgorithm{:direct, SparseMatrixColorings.NaturalOrder}}, AutoSparse{AutoForwardDiff{nothing, Nothing}, SparseConnectivityTracer.TracerLocalSparsityDetector{SparseConnectivityTracer.GradientTracer{Int64, BitSet}, SparseConnectivityTracer.HessianTracer{Int64, BitSet, Dict{Int64, BitSet}, SparseConnectivityTracer.NotShared}}, SparseMatrixColorings.GreedyColoringAlgorithm{:direct, SparseMatrixColorings.NaturalOrder}}, Nothing}, Float64}, BoundaryValueDiffEqMIRK.MIRKTableau{Int64, Vector{Float64}, Vector{Float64}, Vector{Float64}, Matrix{Float64}}, BoundaryValueDiffEqMIRK.MIRKInterpTableau{Int64, Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64, Int64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{PreallocationTools.DiffCache{Matrix{Float64}, Vector{Float64}}}, VectorOfArray{Float64, 3, Vector{Matrix{Float64}}}, Vector{PreallocationTools.DiffCache{Vector{Float64}, Vector{Float64}}}, VectorOfArray{Float64, 2, Vector{Vector{Float64}}}, Vector{PreallocationTools.DiffCache{Vector{Float64}, Vector{Float64}}}, PreallocationTools.DiffCache{Vector{Float64}, Vector{Float64}}, Vector{Float64}, VectorOfArray{Float64, 2, Vector{Vector{Float64}}}, VectorOfArray{Float64, 2, Vector{Vector{Float64}}}, Tuple{Int64}, @NamedTuple{abstol::Float64}, @NamedTuple{abstol::Float64, dt::Float64, adaptive::Bool, controller::BoundaryValueDiffEqCore.DefectControl{Float64}, fit_parameters::Bool}}, BoundaryValueDiffEqCore.DiffCacheNeeded}, Float64}, Float64, 12})\nThe type `Float64` exists, but no method is defined for this combination of argument types when trying to construct it.\n\n\u001b[0mClosest candidates are:\n\u001b[0m  (::Type{T})(::Real, \u001b[91m::RoundingMode\u001b[39m) where T<:AbstractFloat\n\u001b[0m\u001b[90m   @\u001b[39m \u001b[90mBase\u001b[39m \u001b[90m\u001b[4mrounding.jl:265\u001b[24m\u001b[39m\n\u001b[0m  (::Type{T})(::T) where T<:Number\n\u001b[0m\u001b[90m   @\u001b[39m \u001b[90mCore\u001b[39m \u001b[90m\u001b[4mboot.jl:900\u001b[24m\u001b[39m\n\u001b[0m  Float64(\u001b[91m::IrrationalConstants.Sqrthalfπ\u001b[39m)\n\u001b[0m\u001b[90m
2

sol = solve(prob, MIRK4(autodiff=AutoFiniteDiff()); dt = L/100)

3

Gives error:

MethodError: no method matching MIRK4(; autodiff::AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool})
This error has been manually thrown, explicitly, so the method may exist but be intentionally marked as unimplemented.

Closest candidates are:
  MIRK4(; nlsolve, jac_alg, defect_threshold, max_num_subintervals) got unsupported keyword argument "autodiff"
   @ BoundaryValueDiffEqMIRK ~/.julia/packages/BoundaryValueDiffEqMIRK/wcd7Z/src/algorithms.jl:51
  MIRK4(::N, ::J, ::T, ::Int64) where {N, J<:BVPJacobianAlgorithm, T} got unsupported keyword argument "autodiff"
   @ BoundaryValueDiffEqMIRK ~/.julia/packages/BoundaryValueDiffEqMIRK/wcd7Z/src/algorithms.jl:52
4

I guess the answer is that basically all the main Julia pacakges for solving diffeqs are a joke and one needs to write their own.

5

Sorry was thinking of an earlier version. Easiest thing here is to just fix the autodiff. The true problem here is that

tmp_b2::Vector{Float64}       # workspace: |b|^2

is not an auto-diffable workspace because it does not have memory sizes that can match those from autodiff. This should be swapped out to:

So make tmp_b2::T where T <: DiffCache constructed from the initial condition, and then change

@inbounds @simd for i in 1:M
    P.tmp_b2[i] = br[i]*br[i] + bi[i]*bi[i]
end

to

_b2 = get_tmp(P.tmp_b2, bi)
@inbounds @simd for i in 1:M
    _b2[i] = br[i]*br[i] + bi[i]*bi[i]
end

so that it resizes based on the forward-mode chunk size. If you do that then ForwardDiff and automatic sparsity should do just fine in all of the functions, and that’ll do much better than forward-mode.

The full code for this is:

using BoundaryValueDiffEq
using DiffEqBase, OrdinaryDiffEq
using PreallocationTools

# ---------- parameters ----------
mutable struct Params{T <: DiffCache}
    M::Int
    C::NTuple{2,ComplexF64}       # linear coefficients for a,b
    g::NTuple{2,ComplexF64}       # nonlinear coefficients
    zbc::NTuple{2,Float64}        # (z_left, z_right)
    A_re::Vector{Float64}         # target a(0): real part
    A_im::Vector{Float64}         # target a(0): imag part
    B_re::Vector{Float64}         # target b(L): real part
    B_im::Vector{Float64}         # target b(L): imag part
    tmp_b2::T       # workspace: |b|^2
end

# ---------- RHS on packed real state ----------
# v = [a_r(1:M); a_i(1:M); b_r(1:M); b_i(1:M)]
function f_packed!(dv::AbstractVector{T}, v::AbstractVector{T}, P::Params, z) where {T}
    M = P.M
    ar = @view v[1:M];        ai = @view v[M+1:2M]
    br = @view v[2M+1:3M];    bi = @view v[3M+1:4M]

    dar = @view dv[1:M];      dai = @view dv[M+1:2M]
    dbr = @view dv[2M+1:3M];  dbi = @view dv[3M+1:4M]

    # |b|^2
    _b2 = PreallocationTools.get_tmp(P.tmp_b2, dv)
    @inbounds @simd for i in 1:M
        _b2[i] = br[i]*br[i] + bi[i]*bi[i]
    end

    C1, C2 = P.C
    g1, g2 = P.g

    @inbounds @simd for i in 1:M
        ar_i = ar[i]; ai_i = ai[i]
        br_i = br[i]; bi_i = bi[i]
        b2   = _b2[i]                # real scalar

        # ---- da/dz = (C1 + g1*|b|^2)*a
        s1   = C1 + g1*b2                 # complex scalar
        s1r  = real(s1); s1i = imag(s1)
        dar[i] = s1r*ar_i - s1i*ai_i
        dai[i] = s1r*ai_i + s1i*ar_i

        # ---- db/dz = C2*b + (g2*|b|^2)*a
        # linear part on b
        c2r = real(C2); c2i = imag(C2)
        dbr[i] = c2r*br_i - c2i*bi_i
        dbi[i] = c2r*bi_i + c2i*br_i
        # nonlinear part proportional to a
        t2   = g2*b2
        t2r  = real(t2); t2i = imag(t2)
        dbr[i] += t2r*ar_i - t2i*ai_i
        dbi[i] += t2r*ai_i + t2i*ar_i
    end
    return nothing
end

# ---------- boundary conditions on packed real state ----------
# Require: a(0) = A, b(L) = B  (complex equalities → two real equations each)
function bc_packed!(res, u, P::Params, args...)
    M = P.M
    # left boundary (z = zbc[1]): a(0) = A
    uL = u(P.zbc[1])                     # Vector{Float64}, length 4M
    @inbounds @simd for i in 1:M
        res[i]      = uL[i]        - P.A_re[i]        # a_r(0) - A_r = 0
        res[M + i]  = uL[M + i]    - P.A_im[i]        # a_i(0) - A_i = 0
    end
    # right boundary (z = zbc[2]): b(L) = B
    uR = u(P.zbc[2])
    off = 2M
    @inbounds @simd for i in 1:M
        res[2M + i]     = uR[off + i]     - P.B_re[i] # b_r(L) - B_r = 0
        res[3M + i]     = uR[off + M + i] - P.B_im[i] # b_i(L) - B_i = 0
    end
    return nothing
end

# helper: split a packed state vector v (Float64, length 4M) into Complex a,b (length M each)
function unpack_state(v::AbstractVector{<:Real}, M::Int)
    a_re = @view v[1:M]
    a_im = @view v[M+1:2M]
    b_re = @view v[2M+1:3M]
    b_im = @view v[3M+1:4M]
    a = complex.(a_re, a_im)
    b = complex.(b_re, b_im)
    return a, b
end

# ---------- problem setup ----------
M = 2^5
@show M
L = 2.0
x = collect(range(-5.0, 5.0; length=M))

A  = @. exp(-x^2)                    # complex A(x) with 0 imag
B  = @. 1.0e-6 * exp(-x^2)

A_re = Float64.(real.(A));  A_im = Float64.(imag.(A))
B_re = Float64.(real.(B));  B_im = Float64.(imag.(B))

cc = -1.0e-1
C = (cc + 0im, cc + 0im)
gg = -1.0e-4
g = (gg + 0.0im, gg + 0.0im)

P = Params(M, C, g, (0.0, L), A_re, A_im, B_re, B_im, DiffCache(zeros(M)))

# initial guess, packed real
y0 = vcat(A_re, A_im, B_re, B_im);

prob = BVProblem(f_packed!, bc_packed!, y0, (0.0, L), P)
alg = MIRK4(;jac_alg = BVPJacobianAlgorithm(bc_diffmode = AutoForwardDiff(), nonbc_diffmode = AutoForwardDiff()))
sol = solve(prob, alg;  dt= L/100)

Though one style thing, @. _b2 = br*br + bi*bi should SIMD just fine, or at least FastBroadcast.@.. _b2 = br*br + bi*bi definitely will (and will ivdep it), and so you shouldn’t need to scalarize the loops @inbounds @simd to get performance. At least that loop. The second one you might gain a little from that manual loop fusion, though it’s not actually clear how the caching will work out on that by eye.

But note, forward-mode AD here isn’t going to be the biggest of differences. It can be like 2x and it can improve numerical stability, but it’s not the biggest thing in the BVP solver algorithm to focus on. The automated sparsity detection + mixed mode AD is where most of the gains end up coming from on difficult problems, so

alg = MIRK4(;jac_alg = BVPJacobianAlgorithm(bc_diffmode = AutoForwardDiff(), nonbc_diffmode = AutoSparse(AutoForwardDiff())))
sol = solve(prob, alg;  dt= L/100)

is usually going to be a bit better if the problem is large / grows in time steps, and then mixing in Enzyme reverse-mode into the nonbc_diffmode as well if you really want to go ham (the non-BC part is more natural for reverse mode). But that PreallocationTools change is all that’s needed for the autosparse stuff to work, so this will put you on the right path.

And note that once that cache is setup to be AD-compatible the default method works fine too:

alg = MIRK4()
sol = solve(prob, alg;  dt= L/100)

and that will change strategy depending on size.

FFTs next. FFTs are linear operators so they are pretty straightforward to differentiate, the overloads had stalled though.

Those two PRs together add it. It might be interesting to just dev the two forks and see if they do it together.

1 Like