Major improvements in performance (#42)

* checkpoint

* picked series in csv

* fixed estimation

* estimation working for identity

* working

* not working for LogNormal 1 12

* version working for all scalings

* all tests are passing

* deleted unecessary csv

* changes in all dists

* fix gamma bug

* updated benchmarks

* fix bugs found on ENA

* take out Delimited Files

* remove unecessary line

Author: guilhermebodin

bench/benchmark_estimate.jl

@@ -17,15 +17,15 @@ num_seeds = 3
     estimate!(gas, $simulation; verbose = $verbose, opt_method = opt_method)
 end
 # BenchmarkTools.Trial:
-#   memory estimate:  7.33 GiB
-#   allocs estimate:  95576498
+#   memory estimate:  466.69 MiB
+#   allocs estimate:  40237
 #   --------------
-#   minimum time:     5.632 s (15.93% GC)
-#   median time:      5.632 s (15.93% GC)
-#   mean time:        5.632 s (15.93% GC)
-#   maximum time:     5.632 s (15.93% GC)
+#   minimum time:     2.553 s (0.57% GC)
+#   median time:      2.655 s (0.59% GC)
+#   mean time:        2.655 s (0.59% GC)
+#   maximum time:     2.756 s (0.60% GC)
 #   --------------
-#   samples:          1
+#   samples:          2
 #   evals/sample:     1
 
 scaling = 0.0
@@ -43,15 +43,15 @@ num_seeds = 3
     estimate!(gas, $simulation; verbose = $verbose, opt_method = opt_method)
 end
 # BenchmarkTools.Trial:
-#   memory estimate:  8.63 GiB
-#   allocs estimate:  112481074
+#   memory estimate:  353.02 MiB
+#   allocs estimate:  30598
 #   --------------
-#   minimum time:     5.418 s (19.27% GC)
-#   median time:      5.418 s (19.27% GC)
-#   mean time:        5.418 s (19.27% GC)
-#   maximum time:     5.418 s (19.27% GC)
+#   minimum time:     1.008 s (1.41% GC)
+#   median time:      1.265 s (1.36% GC)
+#   mean time:        1.281 s (1.35% GC)
+#   maximum time:     1.587 s (1.28% GC)
 #   --------------
-#   samples:          1
+#   samples:          4
 #   evals/sample:     1
 
 
@@ -62,13 +62,13 @@ scaling = 0.5
     estimate!(gas, $simulation; verbose = $verbose, opt_method = opt_method)
 end
 # BenchmarkTools.Trial:
-#   memory estimate:  17.39 GiB
-#   allocs estimate:  325680001
+#   memory estimate:  8.26 GiB
+#   allocs estimate:  159942290
 #   --------------
-#   minimum time:     9.760 s (23.28% GC)
-#   median time:      9.760 s (23.28% GC)
-#   mean time:        9.760 s (23.28% GC)
-#   maximum time:     9.760 s (23.28% GC)
+#   minimum time:     7.361 s (18.19% GC)
+#   median time:      7.361 s (18.19% GC)
+#   mean time:        7.361 s (18.19% GC)
+#   maximum time:     7.361 s (18.19% GC)
 #   --------------
 #   samples:          1
 #   evals/sample:     1
@@ -80,13 +80,13 @@ scaling = 1.0
     estimate!(gas, $simulation; verbose = $verbose, opt_method = opt_method)
 end
 # BenchmarkTools.Trial:
-#   memory estimate:  11.99 GiB
-#   allocs estimate:  172023821
+#   memory estimate:  666.69 MiB
+#   allocs estimate:  8744603
 #   --------------
-#   minimum time:     8.502 s (20.62% GC)
-#   median time:      8.502 s (20.62% GC)
-#   mean time:        8.502 s (20.62% GC)
-#   maximum time:     8.502 s (20.62% GC)
+#   minimum time:     2.908 s (2.08% GC)
+#   median time:      3.382 s (2.07% GC)
+#   mean time:        3.382 s (2.07% GC)
+#   maximum time:     3.855 s (2.07% GC)
 #   --------------
-#   samples:          1
+#   samples:          2
 #   evals/sample:     1

src/MLE.jl

@@ -1,12 +1,14 @@
 export estimate!
 
+const DEFAULT_INITIAL_PARAM = NaN.*ones(1, 1)
+
 function log_lik(psitilde::Vector{T}, y::Vector{T}, sdm::SDM{D, T}, 
                  initial_params::Vector{Vector{T}}, unknowns::Unknowns_SDM, n::Int) where {D, T}
     return error("log_lik not defined for a model of type ", typeof(sdm))
 end
 
 function estimate!(sdm::SDM{D, T}, y::Vector{T};
-                   initial_params::Vector{Vector{Float64}} = [[NaN]], # Means default initializations
+                   initial_params::Matrix{T} = DEFAULT_INITIAL_PARAM,
                    opt_method::AbstractOptimizationMethod = LBFGS(sdm, 3),
                    verbose::Int = 0) where {D, T}
 

src/ScoreDrivenModels.jl

@@ -5,6 +5,10 @@ using LinearAlgebra
 
 import Base.length
 
+const SCALINGS  = [0.0, 1/2, 1.0]
+const BIG_NUM = 1e8
+const SMALL_NUM = 1e-8
+
 # Core files
 include("abstracts.jl")
 include("utils.jl")

src/distributions/beta.jl

@@ -2,61 +2,56 @@
 Proof somewhere 
 parametrized in \\alpha and \\beta
 """
-function score(y::T, ::Type{Beta}, param::Vector{T}) where T
-    digamma_a_b = digamma(param[1] + param[2])
-    return [
-        log(y) + digamma_a_b - digamma(param[1]);
-        log(1 - y) + digamma_a_b - digamma(param[2])
-    ]
+function score!(score_til::Matrix{T}, y::T, ::Type{Beta}, param::Matrix{T}, t::Int) where T
+    score_til[t, 1] = log(y) + digamma(param[t, 1] + param[t, 2]) - digamma(param[t, 1])
+    score_til[t, 2] = log(1 - y) + digamma(param[t, 1] + param[t, 2]) - digamma(param[t, 2])
+    return
 end
 
 """
 Proof somewhere
 """
-function fisher_information(::Type{Beta}, param::Vector{T}) where T
-    minus_trigamma_a_b = -trigamma(param[1] + param[2])
-
-    return [
-        trigamma(param[1]) + minus_trigamma_a_b    minus_trigamma_a_b;
-        minus_trigamma_a_b                         trigamma(param[2]) + minus_trigamma_a_b
-    ]
+function fisher_information!(aux::AuxiliaryLinAlg{T}, ::Type{Beta}, param::Matrix{T}, t::Int) where T
+    minus_trigamma_a_b = -trigamma(param[t, 1] + param[t, 2])
+    aux.fisher[1, 1] = trigamma(param[t, 1]) + minus_trigamma_a_b
+    aux.fisher[2, 2] = trigamma(param[t, 2]) + minus_trigamma_a_b
+    aux.fisher[2, 1] = minus_trigamma_a_b
+    aux.fisher[1, 2] = minus_trigamma_a_b
+    return
 end
 
 """
 Proof somewhere
 """
-function log_likelihood(::Type{Beta}, y::Vector{T}, param::Vector{Vector{T}}, n::Int) where T
+function log_likelihood(::Type{Beta}, y::Vector{T}, param::Matrix{T}, n::Int) where T
     loglik = 0.0
-    for i in 1:n
-        loglik += (param[i][1] - 1)*log(y[i]) + (param[i][2] - 1)*log(1 - y[i]) - logbeta(param[i][1], param[i][2])
+    for t in 1:n
+        loglik += (param[t, 1] - 1)*log(y[t]) + (param[t, 2] - 1)*log(1 - y[t]) - logbeta(param[t, 1], param[t, 2])
     end
     return -loglik
 end
 
 # Links
-function link(::Type{Beta}, param::Vector{T}) where T 
-    return [
-        link(LogLink, param[1], zero(T));
-        link(LogLink, param[2], zero(T))
-    ]
+function link!(param_tilde::Matrix{T}, ::Type{Beta}, param::Matrix{T}, t::Int) where T 
+    param_tilde[t, 1] = link(LogLink, param[t, 1], zero(T))
+    param_tilde[t, 2] = link(LogLink, param[t, 2], zero(T))
+    return
 end
-function unlink(::Type{Beta}, param_tilde::Vector{T}) where T 
-    return [
-        unlink(LogLink, param_tilde[1], zero(T));
-        unlink(LogLink, param_tilde[2], zero(T))
-    ]
+function unlink!(param::Matrix{T}, ::Type{Beta}, param_tilde::Matrix{T}, t::Int) where T 
+    param[t, 1] = unlink(LogLink, param_tilde[t, 1], zero(T))
+    param[t, 2] = unlink(LogLink, param_tilde[t, 2], zero(T))
+    return
 end
-function jacobian_link(::Type{Beta}, param_tilde::Vector{T}) where T 
-    return Diagonal([
-        jacobian_link(LogLink, param_tilde[1], zero(T));
-        jacobian_link(LogLink, param_tilde[2], zero(T))
-    ])
+function jacobian_link!(aux::AuxiliaryLinAlg{T}, ::Type{Beta}, param::Matrix{T}, t::Int) where T 
+    aux.jac[1] = jacobian_link(LogLink, param[t, 1], zero(T))
+    aux.jac[2] = jacobian_link(LogLink, param[t, 2], zero(T))
+    return
 end
 
 # utils
-function update_dist(::Type{Beta}, param::Vector{T}) where T
-    small_threshold!(param, T(1e-8))
-    return Beta(param[1], param[2])
+function update_dist(::Type{Beta}, param::Matrix{T}, t::Int) where T
+    small_threshold!(param, SMALL_NUM, t)
+    return Beta(param[t, 1], param[t, 2])
 end 
 
 function num_params(::Type{Beta})

src/distributions/common_interface.jl

@@ -8,25 +8,28 @@ const DISTS = [
     Weibull
 ]
 
-function score(y::T, D::Type{<:Distribution}, param::Vector{T}) where T
-    return error("score not implemented for $D distribution")
+function score!(score_til::Matrix{T}, y::T, ::Type{<:Distribution}, param::Matrix{T}, t::Int) where T
+    return error("score! not implemented for $D distribution")
 end
-function score(y::Int, D::Type{<:Distribution}, param::Vector{T}) where T
-    return error("score not implemented for $D distribution")
+function score!(score_til::Matrix{T}, y::Int, ::Type{<:Distribution}, param::Matrix{T}, t::Int) where T
+    return error("score! not implemented for $D distribution")
 end
-function fisher_information(D::Type{<:Distribution}, param::Vector{T}) where T
-    return error("fisher_information not implemented for $D distribution")
+function fisher_information!(aux::AuxiliaryLinAlg{T},::Type{<:Distribution}, param::Matrix{T}, t::Int) where T
+    return error("fisher_information! not implemented for $D distribution")
 end
-function link(D::Type{<:Distribution}, param::Vector{T}) where T 
-    return error("link not implemented for $D distribution")
+function log_likelihood(::Type{<:Distribution}, y::Vector{T}, param::Matrix{T}, n::Int) where T
+    return error("log_likelihood not implemented for $D distribution")
 end
-function unlink(D::Type{<:Distribution}, param_tilde::Vector{T}) where T 
-    return error("unlink not implemented for $D distribution")
+function link!(param_tilde::Matrix{T}, ::Type{<:Distribution}, param::Matrix{T}, t::Int) where T 
+    return error("link! not implemented for $D distribution")
 end
-function jacobian_link(D::Type{<:Distribution}, param_tilde::Vector{T}) where T 
-    return error("jacobian_link not implemented for $D distribution")
+function unlink!(param::Matrix{T}, ::Type{<:Distribution}, param_tilde::Matrix{T}, t::Int) where T 
+    return error("unlink! not implemented for $D distribution")
 end
-function update_dist(D::Type{<:Distribution}, param::Vector{T}) where T
+function jacobian_link!(aux::AuxiliaryLinAlg{T}, ::Type{<:Distribution}, param::Matrix{T}, t::Int) where T 
+    return error("jacobian_link! not implemented for $D distribution")
+end
+function update_dist(::Type{<:Distribution}, param::Matrix{T}, t::Int) where T
     return error("update_dist not implemented for $D distribution")
 end 
 function num_params(D::Type{<:Distribution})

src/distributions/gamma.jl

@@ -2,57 +2,56 @@
 Proof somewhere 
 parametrized in \\alpha and \\theta
 """
-function score(y::T, ::Type{Gamma}, param::Vector{T}) where T
-    return [
-        log(y) - digamma(param[1]) - log(param[2]);
-        y/param[2]^2 - param[1]/param[2]
-    ]
+function score!(score_til::Matrix{T}, y::T, ::Type{Gamma}, param::Matrix{T}, t::Int) where T
+    score_til[t, 1] = log(y) - digamma(param[t, 1]) - log(param[t, 2])
+    score_til[t, 2] = y/param[t, 2]^2 - param[t, 1]/param[t, 2]
+    return 
 end
 
 """
 Proof somewhere
 """
-function fisher_information(::Type{Gamma}, param::Vector{T}) where T
-    [
-        -trigamma(param[1])     -1/param[2];
-        -1/param[2]             -2*y/param[2]^3 + param[1]/param[2]^2
-    ]
+function fisher_information!(aux::AuxiliaryLinAlg{T}, ::Type{Gamma}, param::Matrix{T}, t::Int) where T
+    error("Fisher information not implemented for Gamma distribution.")
+    aux.fisher[1, 1] = -trigamma(param[t, 1])
+    aux.fisher[2, 2] = -2 * y / param[t, 2] ^ 3 + param[t, 1] / param[t, 2]^2
+    aux.fisher[2, 1] = -1/param[t, 2]
+    aux.fisher[1, 2] = -1/param[t, 2]
+    return
 end
 
 """
 Proof somewhere
 """
-function log_likelihood(::Type{Gamma}, y::Vector{T}, param::Vector{Vector{T}}, n::Int) where T
+function log_likelihood(::Type{Gamma}, y::Vector{T}, param::Matrix{T}, n::Int) where T
     loglik = 0.0
-    for i in 1:n
-        loglik += (param[i][1] - 1)*log(y[i]) - y[i]/param[i][2] - loggamma(param[i][1]) - param[i][1]*log(param[i][2])
+    for t in 1:n
+        loglik += (param[t, 1] - 1)*log(y[t]) - y[t]/param[t, 2] - loggamma(param[t, 1]) - param[t, 1]*log(param[t, 2])
     end
     return -loglik
 end
 
 # Links
-function link(::Type{Gamma}, param::Vector{T}) where T 
-    return [
-        link(LogLink, param[1], zero(T));
-        link(LogLink, param[2], zero(T))
-    ]
+function link!(param_tilde::Matrix{T}, ::Type{Gamma}, param::Matrix{T}, t::Int) where T 
+    param_tilde[t, 1] = link(LogLink, param[t, 1], zero(T))
+    param_tilde[t, 2] = link(LogLink, param[t, 2], zero(T))
+    return
 end
-function unlink(::Type{Gamma}, param_tilde::Vector{T}) where T 
-    return [
-        unlink(LogLink, param_tilde[1], zero(T));
-        unlink(LogLink, param_tilde[2], zero(T))
-    ]
+function unlink!(param::Matrix{T}, ::Type{Gamma}, param_tilde::Matrix{T}, t::Int) where T 
+    param[t, 1] = unlink(LogLink, param_tilde[t, 1], zero(T))
+    param[t, 2] = unlink(LogLink, param_tilde[t, 2], zero(T))
+    return
 end
-function jacobian_link(::Type{Gamma}, param_tilde::Vector{T}) where T 
-    return Diagonal([
-        jacobian_link(LogLink, param_tilde[1], zero(T));
-        jacobian_link(LogLink, param_tilde[2], zero(T))
-    ])
+function jacobian_link!(aux::AuxiliaryLinAlg{T}, ::Type{Gamma}, param::Matrix{T}, t::Int) where T 
+    aux.jac[1] = jacobian_link(LogLink, param[t, 1], zero(T))
+    aux.jac[2] = jacobian_link(LogLink, param[t, 2], zero(T))
+    return
 end
 
 # utils 
-function update_dist(::Type{Gamma}, param::Vector{T}) where T
-    return Gamma(param[1], param[2])
+function update_dist(::Type{Gamma}, param::Matrix{T}, t::Int) where T
+    small_threshold!(param, SMALL_NUM, t)
+    return Gamma(param[t, 1], param[t, 2])
 end 
 
 function num_params(::Type{Gamma})