Improve stacktraces by storing non-nothing optional values in a named tuple

[CameronBieganek]

As discussed on Discourse, it is possible to improve the SciML stacktraces by storing non-nothing optional values in a named tuple, rather than having highly parameterized structs. Since most of the optional fields in SciML structs end up being set to nothing, using a small named tuple with the non-nothing values significantly cuts down on the length of type signatures.

This idea applies to any struct that contains many optional fields. Here is an example in the spirit of ODEFunction.

struct MyFunction{F,N,T}
    f::F
    more_funcs::NamedTuple{N,T}
end

function MyFunction(
        f;
        mass_matrix=nothing, analytic=nothing, tgrad=nothing,
        jac=nothing, jvp=nothing, vjp=nothing, jac_prototype=nothing,
        sparsity=nothing, Wfact=nothing, Wfact_t=nothing,
        paramjac=nothing, syms=nothing, indepsym=nothing,
        paramsyms=nothing, observed=nothing, colorvec=nothing,
        sys=nothing
    )
    # This particular syntax requires Julia 1.7.
    kwargs = (;
        mass_matrix, analytic, tgrad, jac, jvp, vjp, jac_prototype,
        sparsity, Wfact, Wfact_t, paramjac, syms, indepsym, paramsyms,
        observed, colorvec, sys
    )

    more_funcs = NamedTuple(k => v for (k, v) in pairs(kwargs) if v !== nothing)
    MyFunction(f, more_funcs)
end

For comparison, here is a type SciMLFunction that is essentially the same as ODEFunction.

struct SciMLFunction{iip, specialize, F, TMM, Ta, Tt,
                   TJ, JVP, VJP, JP, SP, TW, TWt,
                   TPJ, S, S2, S3, O, TCV, SYS}
    f::F
    mass_matrix::TMM
    analytic::Ta
    tgrad::Tt
    jac::TJ
    jvp::JVP
    vjp::VJP
    jac_prototype::JP
    sparsity::SP
    Wfact::TW
    Wfact_t::TWt
    paramjac::TPJ
    syms::S
    indepsym::S2
    paramsyms::S3
    observed::O
    colorvec::TCV
    sys::SYS
end

function SciMLFunction(f; mass_matrix=nothing, analytic=nothing, tgrad=nothing,
                       jac=nothing, jvp=nothing, vjp=nothing, jac_prototype=nothing,
                       sparsity=nothing, Wfact=nothing, Wfact_t=nothing,
                       paramjac=nothing, syms=nothing, indepsym=nothing,
                       paramsyms=nothing, observed=nothing, colorvec=nothing,
                       sys=nothing)
    SciMLFunction{1, 2, typeof(f), typeof(mass_matrix), typeof(analytic), typeof(tgrad), typeof(jac), typeof(jvp), typeof(vjp), typeof(jac_prototype),
    typeof(sparsity), typeof(Wfact), typeof(Wfact_t), typeof(paramjac), typeof(syms), typeof(indepsym), typeof(paramsyms),
    typeof(observed), typeof(colorvec), typeof(sys)}(f, mass_matrix, analytic, tgrad, jac, jvp, vjp, jac_prototype,
        sparsity, Wfact, Wfact_t, paramjac, syms, indepsym, paramsyms,
        observed, colorvec, sys)
end

And here are two functions foo that operate on a SciMLFunction or MyFunction.

function foo(funcs::SciMLFunction, x)
    g = isnothing(funcs.jac) ? identity : funcs.jac
    h = isnothing(funcs.jvp) ? identity : funcs.jvp
    k = isnothing(funcs.vjp) ? identity : funcs.vjp
    funcs.f(g(h(k(x))))
end

function foo(funcs::MyFunction, x)
    g = get(funcs.more_funcs, :jac, identity)
    h = get(funcs.more_funcs, :jvp, identity)
    k = get(funcs.more_funcs, :vjp, identity)
    funcs.f(g(h(k(x))))
end

Now let's initialize the data structures:

sciml_funcs = SciMLFunction(x -> x + 1)
my_funcs = MyFunction(x -> x + 1)

Now let's compare the stacktraces generated by foo for the two different types. Here's a screenshot of the stacktrace created by SciMLFunction:

And here's a screenshot of the stacktrace created by MyFunction:

If we compare frame 2 from each stacktrace, we see that the printing of the MyFunction type is much shorter than the printing of the SciMLFunction type. I think it's a considerable improvement.

Implementing this technique across the SciML ecosystem has the potential to significantly reduce the size and complexity of SciML stacktraces.

[CameronBieganek]

I've updated the example implementation of MyFunction in the original post. The constructor now explicitly lists out all the valid keyword arguments, so it will naturally check to ensure the user-provided keyword arguments match the expected arguments.