Reduce memory footprint

.github/workflows/CI.yml

@@ -20,7 +20,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          - '1.8'
           - '1'
           - 'pre'
         os:
@@ -68,7 +68,7 @@ jobs:
         include:
           - version: '1'
             downgrade: false
-          - version: '1.7'
+          - version: '1.8'
             downgrade: true
     steps:
       - uses: actions/checkout@v4

Project.toml

@@ -23,7 +23,6 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
-UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [weakdeps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -63,8 +62,7 @@ Statistics = "1.6"
 StatsBase = "0.33.7, 0.34"
 Transducers = "0.4.81"
 Turing = "0.31.4, 0.32, 0.33, 0.34, 0.35"
-UnPack = "1"
-julia = "1.6"
+julia = "1.8"
 
 [extras]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

src/Pathfinder.jl

@@ -19,7 +19,6 @@ using SciMLBase: SciMLBase
 using Statistics: Statistics
 using StatsBase: StatsBase
 using Transducers: Transducers
-using UnPack: @unpack
 
 export PathfinderResult, MultiPathfinderResult
 export pathfinder, multipathfinder
@@ -43,7 +42,7 @@ default_ad() = ADTypes.AutoForwardDiff()
 include("transducers.jl")
 include("woodbury.jl")
 include("optimize.jl")
-include("inverse_hessian.jl")
+include("lbfgs.jl")
 include("mvnormal.jl")
 include("elbo.jl")
 include("resample.jl")

src/elbo.jl

@@ -1,24 +1,30 @@
-function maximize_elbo(rng, logp, dists, ndraws, executor)
+function maximize_elbo(rng, logp, dists, ndraws, executor; save_samples::Bool=true)
+    dim = isempty(dists) ? 0 : length(first(dists))
+    Tdraws = Matrix{eltype(eltype(dists))}
+    draws = save_samples ? nothing : Tdraws(undef, (dim, ndraws))
     EE = Core.Compiler.return_type(
-        elbo_and_samples, Tuple{typeof(rng),typeof(logp),eltype(dists),Int}
+        elbo_and_samples!, Tuple{Tdraws,typeof(rng),typeof(logp),eltype(dists)}
     )
     estimates = similar(dists, EE)
     isempty(estimates) && return 0, estimates
+
     Folds.map!(estimates, dists, executor) do dist
-        return elbo_and_samples(rng, logp, dist, ndraws)
+        _draws = draws === nothing ? Tdraws(undef, (dim, ndraws)) : draws
+        return elbo_and_samples!(_draws, rng, logp, dist; save_samples)
     end
     _, iteration_opt = _findmax(estimates |> Transducers.Map(est -> est.value))
     return iteration_opt, estimates
 end
 
-function elbo_and_samples(rng, logp, dist, ndraws)
-    ϕ, logqϕ = rand_and_logpdf(rng, dist, ndraws)
+function elbo_and_samples!(ϕ, rng, logp, dist; save_samples::Bool=true)
+    ϕ, logqϕ = rand_and_logpdf!(rng, dist, ϕ)
     logpϕ = similar(logqϕ)
     logpϕ .= logp.(eachcol(ϕ))
     logr = logpϕ - logqϕ
     elbo = Statistics.mean(logr)
     elbo_se = sqrt(Statistics.var(logr) / length(logr))
-    return ELBOEstimate(elbo, elbo_se, ϕ, logpϕ, logqϕ, logr)
+    ϕ_save = save_samples ? copyto!(similar(ϕ), ϕ) : similar(ϕ, map(zero, size(ϕ)))
+    return ELBOEstimate(elbo, elbo_se, ϕ_save, logpϕ, logqϕ, logr)
 end
 
 struct ELBOEstimate{T,P,L<:AbstractVector{T}}