Add GeneralLazyBufferCache

Project.toml

@@ -1,7 +1,7 @@
 name = "PreallocationTools"
 uuid = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "0.4.7"
+version = "0.4.8"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -21,11 +21,12 @@ Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["LabelledArrays", "LinearAlgebra", "OrdinaryDiffEq", "Test", "RecursiveArrayTools", "Pkg", "SafeTestsets", "Optimization", "OptimizationOptimJL", "SparseArrays", "Symbolics"]
+test = ["Random", "LabelledArrays", "LinearAlgebra", "OrdinaryDiffEq", "Test", "RecursiveArrayTools", "Pkg", "SafeTestsets", "Optimization", "OptimizationOptimJL", "SparseArrays", "Symbolics"]

README.md

@@ -234,11 +234,77 @@ prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), LazyBufferCache()))
 solve(prob, TRBDF2())
 ```
 
-## Note About ReverseDiff Support for LazyBuffer
+### Note About ReverseDiff Support for LazyBuffer
 
 ReverseDiff support is done in SciMLSensitivity.jl to reduce the AD requirements on this package.
 Load that package if ReverseDiff overloads are required.
 
+## GeneralLazyBufferCache
+
+```julia
+GeneralLazyBufferCache(f=identity)
+```
+
+A `GeneralLazyBufferCache` is a `Dict`-like type for the caches which automatically defines
+new caches on demand when they are required. The function `f` generates the cache matching
+for the type of `u`, and subsequent indexing reuses that cache if that type of `u` has
+already ben seen.
+
+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.
+
+### Example
+
+In all of the previous cases our cache was an array. However, in this case we want to preallocate
+a DifferentialEquations `ODEIntegrator` object. This object is the one created via
+`DifferentialEquations.init(ODEProblem(ode_fnc, y₀, (0.0, T), p), Tsit5(); saveat = t)`, and we
+want to optimize `p` in a way that changes its type to ForwardDiff. Thus what we can do is make a
+GeneralLazyBufferCache which holds these integrator objects, defined by `p`, and indexing it with
+`p` in order to retrieve the cache. The first time it's called it will build the integrator, and 
+in subsequent calls it will reuse the cache.
+
+Defining the cache as a function of `p` to build an integrator thus looks like:
+
+```julia
+lbc = GeneralLazyBufferCache(function (p)
+    DifferentialEquations.init(ODEProblem(ode_fnc, y₀, (0.0, T), p), Tsit5(); saveat = t)
+end)
+```
+
+then `lbc[p]` will be smart and reuse the caches. A full example looks like the following:
+
+```julia
+using Random, DifferentialEquations, LinearAlgebra, Optimization, OptimizationNLopt, OptimizationOptimJL, PreallocationTools
+
+lbc = GeneralLazyBufferCache(function (p)
+    DifferentialEquations.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])
+solve(prob, LBFGS())
+```
+
 ## Similar Projects
 
 [AutoPreallocation.jl](https://github.com/oxinabox/AutoPreallocation.jl) tries

docs/src/index.md

@@ -224,11 +224,77 @@ prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), LazyBufferCache()))
 solve(prob, TRBDF2())
 ```
 
-## Note About ReverseDiff Support for LazyBuffer
+### Note About ReverseDiff Support for LazyBuffer
 
 ReverseDiff support is done in SciMLSensitivity.jl to reduce the AD requirements on this package.
 Load that package if ReverseDiff overloads are required.
 
+## GeneralLazyBufferCache
+
+```julia
+GeneralLazyBufferCache(f=identity)
+```
+
+A `GeneralLazyBufferCache` is a `Dict`-like type for the caches which automatically defines
+new caches on demand when they are required. The function `f` generates the cache matching
+for the type of `u`, and subsequent indexing reuses that cache if that type of `u` has
+already ben seen.
+
+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.
+
+### Example
+
+In all of the previous cases our cache was an array. However, in this case we want to preallocate
+a DifferentialEquations `ODEIntegrator` object. This object is the one created via
+`DifferentialEquations.init(ODEProblem(ode_fnc, y₀, (0.0, T), p), Tsit5(); saveat = t)`, and we
+want to optimize `p` in a way that changes its type to ForwardDiff. Thus what we can do is make a
+GeneralLazyBufferCache which holds these integrator objects, defined by `p`, and indexing it with
+`p` in order to retrieve the cache. The first time it's called it will build the integrator, and 
+in subsequent calls it will reuse the cache.
+
+Defining the cache as a function of `p` to build an integrator thus looks like:
+
+```julia
+lbc = GeneralLazyBufferCache(function (p)
+    DifferentialEquations.init(ODEProblem(ode_fnc, y₀, (0.0, T), p), Tsit5(); saveat = t)
+end)
+```
+
+then `lbc[p]` will be smart and reuse the caches. A full example looks like the following:
+
+```julia
+using Random, DifferentialEquations, LinearAlgebra, Optimization, OptimizationNLopt, OptimizationOptimJL, PreallocationTools
+
+lbc = GeneralLazyBufferCache(function (p)
+    DifferentialEquations.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])
+solve(prob, LBFGS())
+```
+
 ## Similar Projects
 
 [AutoPreallocation.jl](https://github.com/oxinabox/AutoPreallocation.jl) tries

src/PreallocationTools.jl

@@ -183,7 +183,35 @@ function Base.getindex(b::LazyBufferCache, u::T) where {T <: AbstractArray}
     return buf
 end
 
-export FixedSizeDiffCache, DiffCache, LazyBufferCache, dualcache
+# GeneralLazyBufferCache
+
+"""
+    b = GeneralLazyBufferCache(f=identity)
+
+A lazily allocated buffer object.  Given an array `u`, `b[u]` returns a cache object
+generated by `f(u)`, but the generator is only ran the first time (and all subsequent
+times it reuses the same cache)
+
+## Limitation
+
+The main limitation of this method is that its return is not type-inferred, and thus
+it can be slower than some of the other preallocation techniques. However, if used
+correct using things like function barriers, then this is a general technique that
+is sufficiently fast.
+"""
+struct GeneralLazyBufferCache{F <: Function}
+    bufs::Dict # a dictionary mapping types to buffers
+    f::F
+    GeneralLazyBufferCache(f::F = identity) where {F <: Function} = new{F}(Dict(), f) # start with empty dict
+end
+
+function Base.getindex(b::GeneralLazyBufferCache, u::T) where {T}
+    get!(b.bufs, T) do
+        b.f(u)
+    end
+end
+
+export GeneralLazyBufferCache, FixedSizeDiffCache, DiffCache, LazyBufferCache, dualcache
 export get_tmp
 
 end

test/general_lbc.jl

@@ -0,0 +1,33 @@
+using Random, OrdinaryDiffEq, LinearAlgebra, Optimization, OptimizationOptimJL,
+      PreallocationTools
+
+lbc = GeneralLazyBufferCache(function (p)
+                                 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])
+solve(prob, LBFGS())

test/runtests.jl

@@ -16,6 +16,7 @@ if GROUP == "All" || GROUP == "Core"
     @safetestset "DiffCache Resizing" begin include("core_resizing.jl") end
     @safetestset "DiffCache Nested Duals" begin include("core_nesteddual.jl") end
     @safetestset "DiffCache Sparsity Support" begin include("sparsity_support.jl") end
+    @safetestset "GeneralLazyBufferCache" begin include("general_lbc.jl") end
 end
 
 if !is_APPVEYOR && GROUP == "GPU"