Implementation of jacobian_forward_ad

[DanielDoehring]

The current implementation of jacobian_forward_ad

Trixi.jl/src/semidiscretization/semidiscretization.jl

Lines 254 to 261 in a633f9c

	function _jacobian_ad_forward(semi, t0, u0_ode, du_ode, config)
	new_semi = remake(semi, uEltype = eltype(config))
	J = ForwardDiff.jacobian(du_ode, u0_ode, config) do du_ode, u_ode
	Trixi.rhs!(du_ode, u_ode, new_semi, t0)
	end
	return J
	end

uses the first proposed way of computing the Jacobian, i.e., this code

"""
    ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f, x), check=Val{true}())

Compute `J(f)` evaluated at `x` and store the result(s) in `result`, assuming `f` is called
as `f(x)`.

This method assumes that `isa(f(x), AbstractArray)`.

Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
"""
function jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T, CHK}
    result isa DiffResult ? require_one_based_indexing(x) : require_one_based_indexing(result, x)
    CHK && checktag(T, f, x)
    if chunksize(cfg) == length(x)
        vector_mode_jacobian!(result, f, x, cfg)
    else
        chunk_mode_jacobian!(result, f, x, cfg)
    end
    return result
end

Semantically, we use the second option, i.e., this code :

"""
    ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())

Compute `J(f!)` evaluated at `x` and store the result(s) in `result`, assuming `f!` is
called as `f!(y, x)` where the result is stored in `y`.

This method assumes that `isa(f(x), AbstractArray)`.

Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
"""
function jacobian!(result::Union{AbstractArray,DiffResult}, f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T,CHK}
    result isa DiffResult ? require_one_based_indexing(y, x) : require_one_based_indexing(result, y, x)
    CHK && checktag(T, f!, x)
    if chunksize(cfg) == length(x)
        vector_mode_jacobian!(result, f!, y, x, cfg)
    else
        chunk_mode_jacobian!(result, f!, y, x, cfg)
    end
    return result
end

What are the reasons for this? The additional arguments in Trixi.rhs!?

In any case, I think we should add some explaining comments to elaborate what is going on in the implementation of the Jacobian, especially given the use of the anonymous/lambda function.

[ranocha]

Actually, we use neither of the options you listed above. The do block syntax introduces an anonymous function that is inserted as first argument to the call to this method

"""
    ForwardDiff.jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())

Return `J(f!)` evaluated at `x`,  assuming `f!` is called as `f!(y, x)` where the result is
stored in `y`.

Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
"""
function jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T, CHK}
    CHK && checktag(T, f!, x)
    if chunksize(cfg) == length(x)
        return vector_mode_jacobian(f!, y, x, cfg)
    else
        return chunk_mode_jacobian(f!, y, x, cfg)
    end
end

Feel free to suggest some comments to improve this part.

[DanielDoehring]

Hm interesting that indeed the version with y, x is called. I guess I need to understand anonymous functions before being able to suggest something useful here.

[DanielDoehring]

Ah, I finally understood the Julia documentation of anonymous functions:

The do x syntax creates an anonymous function with argument x and passes it it as the first argument

I always read that as it referring to the argument x (which would be modified due to the anonymous function), not the anonymous function...

What do you think about

function _jacobian_ad_forward(semi, t0, u0_ode, du_ode, config)
    new_semi = remake(semi, uEltype = eltype(config))
    # Create anonymous function passed as first argument to `ForwardDiff.jacobian` to match
    # `ForwardDiff.jacobian(f!, y::AbstractArray, x::AbstractArray, 
    #                       cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())`
    J = ForwardDiff.jacobian(du_ode, u0_ode, config) do du_ode, u_ode
        Trixi.rhs!(du_ode, u_ode, new_semi, t0)
    end

    return J
end

[ranocha]

Would be fine with me