Add GeneralLazyBufferCache

[ChrisRackauckas]

This is a slower caching method but is super general, used to solve a real user's question (https://discourse.julialang.org/t/declaring-forwarddiff-tag-directly-with-a-differentialequations-integrator-nested-function/83766) that otherwise would take much nastier tricks.

[DanielVandH]

Would it be surprising that the results when using this cache are not inferrable? e.g.

using Random, OrdinaryDiffEq, LinearAlgebra, Optimization, OptimizationOptimJL,
    PreallocationTools

function testprob()
    lbc = GeneralLazyBufferCache(function (p)
        SciMLBase.init(ODEProblem(ode_fnc, y₀,
                (0.0, T), p),
            Tsit5(); saveat=t)
    end)

    Random.seed!(2992999)
    λ, y₀, σ = -0.5, 15.0, 0.1
    T, n = 5.0, 200
    Δt = T / n
    t = [j * Δt for j in 0:n]
    y = y₀ * exp.(λ * t)
    yᵒ = y .+ [0.0, σ * randn(n)...]
    ode_fnc(u, p, t) = p * u
    function loglik(θ, data, integrator)
        yᵒ, n, ε = data
        λ, σ, u0 = θ
        integrator.p = λ
        reinit!(integrator, u0)
        solve!(integrator)
        ε = yᵒ .- integrator.sol.u
        ℓ = -0.5n * log(2π * σ^2) - 0.5 / σ^2 * sum(ε .^ 2)
    end
    θ₀ = [-1.0, 0.5, 19.73]
    negloglik = (θ, p) -> -loglik(θ, p, lbc[θ[1]])
    fnc = OptimizationFunction(negloglik, Optimization.AutoForwardDiff())
    ε = zeros(n)
    prob = OptimizationProblem(fnc, θ₀, (yᵒ, n, ε), lb=[-10.0, 1e-6, 0.5],
        ub=[10.0, 10.0, 25.0])
    return prob
end 

prob = testprob()

@code_warntype prob.f([-1.0,0.5,19.73], prob.p)

julia> @code_warntype prob.f([-1.0,0.5,19.73], prob.p)
MethodInstance for (::OptimizationFunction{true, Optimization.AutoForwardDiff{nothing}, var"#7#12"{var"#loglik#11", GeneralLazyBufferCache{var"#5#8"}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing})(::Vector{Float64}, ::Tuple{Vector{Float64}, Int64, Vector{Float64}})
  from (f::OptimizationFunction)(args...) in SciMLBase at C:\Users\licer\.julia\packages\SciMLBase\lGVlK\src\scimlfunctions.jl:3580
Arguments
  f::OptimizationFunction{true, Optimization.AutoForwardDiff{nothing}, var"#7#12"{var"#loglik#11", GeneralLazyBufferCache{var"#5#8"}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}
  args::Tuple{Vector{Float64}, Tuple{Vector{Float64}, Int64, Vector{Float64}}}
Body::Any
1 ─ %1 = Base.getproperty(f, :f)::var"#7#12"{var"#loglik#11", GeneralLazyBufferCache{var"#5#8"}}
│   %2 = Core._apply_iterate(Base.iterate, %1, args)::Any
└──      return %2

I imagine it's related to the Dict definition, though this is also present in LazyBufferCache.

[ChrisRackauckas]

See the README:

Note that LazyBufferCache does cause a dynamic dispatch and its return is not type-inferred. This means it's the slowest of the preallocation methods, but it's the most general.

Notice that the example uses a function barrier.

[DanielVandH]

Ah, missed that in the README sorry - thanks for the prompt response. Am I understanding correctly, then, that it won't actually be a big problem when using solve compared to what I see when evaluating it directly in the REPL?

[ChrisRackauckas]

It depends on how you use it. With a function barrier you get hit with the standard function barrier cost of around 50-100ns. If that's fine then you're good.