improve simulation code

Author: lanceXwq

src/base.jl

@@ -87,6 +87,9 @@ function _rand(prior::P, Δparams::NTuple{2,T}) where {T,P}
     return rand(P(postparams...))
 end
 
+_randn!(x::AbstractVecOrMat{T}, μ::AbstractVector{T}, σ::AbstractVector{T}) where {T} =
+    x .= muladd.(randn!(x), σ, μ)
+
 _randn!(x::AbstractArray{T}, σ::Union{T,AbstractMatrix{T}}) where {T} = randn!(x) .*= σ
 
 function _randn!(x::AbstractArray{T,3}, σ::T, σ₀::AbstractVecOrMat{T}) where {T}
@@ -142,5 +145,15 @@ function mrw_propose(P::Union{BetaPrime{T},Gamma{T}}, xᵒ::T) where {T<:Abstrac
     return xᵖ, log(xᵒ) + logpdf(P, 1 / ϵ) - log(xᵖ) - logpdf(P, ϵ) #* Needs test
 end
 
+function _resize3(
+    x::AbstractArray{T,3},
+    newsz::Integer;
+    fill = Union{Nothing,T} = nothing,
+) where {T}
+    y = similar(x, size(x, 1), size(x, 2), newsz)
+    copyto!(y, x)
+    !isnothing(fill) && fill!(view(y, :, :, size(x, 3)+1:newsz), fill)
+    return y
+end
 # get_Δlogprior(xᵖ::T, xᵒ::T, distr::Gamma{T}) where {T<:AbstractFloat} =
 #     (shape(distr) - 1) * log(xᵖ / xᵒ) - (xᵖ - xᵒ) / scale(distr)

src/chain.jl

@@ -93,7 +93,7 @@ function extend!(
         logposterior = get_logposterior(loglikelihood, tracks, msd)
         push!(
             chain.samples,
-            Sample(tracks.onpart, msd, brightness, iter, 𝑇, logposterior, loglikelihood),
+            Sample(tracks.onchunk, msd, brightness, iter, 𝑇, logposterior, loglikelihood),
         )
         isfull(chain) && shrink!(chain)
     end
@@ -110,8 +110,8 @@ function get_loglikelihood!(
     detector::Detector{T},
     psf::PointSpreadFunction{T},
 ) where {T}
-    seteffvalue!(tracks.onpart)
-    set_poisson_mean!(llarray, detector, tracks.onpart.effvalue, brightness.value, psf)
+    seteffvalue!(tracks.onchunk)
+    set_poisson_mean!(llarray, detector, tracks.onchunk.effvalue, brightness.value, psf)
     return get_loglikelihood!(llarray, detector)
 end
 
@@ -121,7 +121,7 @@ get_logposterior(
     msd::MeanSquaredDisplacement{T},
 ) where {T} =
     loglikelihood +
-    logprior(tracks.onpart, msd.value) +
+    logprior(tracks.onchunk, msd.value) +
     logprior(msd) +
     logprior(tracks.nemitters)
 
@@ -134,10 +134,10 @@ function parametricMCMC!(
     psf::PointSpreadFunction{T},
     𝑇::T,
 ) where {T}
-    update_onpart!(tracks, msd.value, brightness.value, llarray, detector, psf, 𝑇)
-    update!(brightness, tracks.onpart.effvalue, llarray, detector, psf, 𝑇)
-    setdisplacement²!(tracks.onpart)
-    update!(msd, tracks.onpart.displacement², 𝑇)
+    update_onchunk!(tracks, msd.value, brightness.value, llarray, detector, psf, 𝑇)
+    update!(brightness, tracks.onchunk.effvalue, llarray, detector, psf, 𝑇)
+    setdisplacement²!(tracks.onchunk)
+    update!(msd, tracks.onchunk.displacement², 𝑇)
     return tracks, msd
 end
 
@@ -150,9 +150,9 @@ function nonparametricMCMC!(
     psf::PointSpreadFunction{T},
     𝑇::T,
 ) where {T}
-    simulate!(tracks.offpart, msd.value)
+    simulate!(tracks.offchunk, msd.value)
     if any(tracks)
-        update_onpart!(tracks, msd.value, brightness.value, llarray, detector, psf, 𝑇)
+        update_onchunk!(tracks, msd.value, brightness.value, llarray, detector, psf, 𝑇)
         onshuffle!(tracks)
     end
 
@@ -168,7 +168,7 @@ function nonparametricMCMC!(
     )
     reassign!(tracks)
 
-    update!(brightness, tracks.onpart.effvalue, llarray, detector, psf, 𝑇)
+    update!(brightness, tracks.onchunk.effvalue, llarray, detector, psf, 𝑇)
 
     setdisplacement²!(tracks)
     update!(msd, tracks.displacement²[1], 𝑇)
@@ -209,7 +209,7 @@ function runMCMC(;
     parametric::Bool = false,
 ) where {T}
     isnothing(annealing) && (annealing = ConstantAnnealing{T}(1))
-    chain = Chain([Sample(tracks.onpart, msd, brightness)], sizelimit, annealing)
+    chain = Chain([Sample(tracks.onchunk, msd, brightness)], sizelimit, annealing)
     runMCMC!(chain, tracks, msd, brightness, detector, psf, niters, parametric)
     return chain
 end

src/distributions.jl

@@ -17,7 +17,9 @@ function Base.getproperty(d::DNormal, s::Symbol)
     end
 end
 
-Distributions.params(n::DNormal) = n.μ, n.σ
+Distributions.params(P::DNormal) = P.μ, P.σ
 
 logprior(ℕ::DNormal{T}, x::AbstractArray{T}) where {T} =
-    sum(vec(@. -(x - ℕ.μ) / (2 * ℕ.σ^2)))
+    sum(vec(@. -(x - ℕ.μ) / (2 * ℕ.σ^2)))
+
+Random.rand!(x::AbstractVecOrMat{T}, P::DNormal{T}) where {T} = _randn!(x, P.μ, P.σ)

src/permutation.jl

@@ -30,8 +30,8 @@ function _permute!(
 end
 
 function _permute!(tracks::Tracks{T}, p::AbstractVector{<:Integer}) where {T}
-    _permute!(tracks.onpart.value, p, tracks.proposals.value)
-    _permute!(tracks.onpart.active, p, tracks.proposals.active)
+    _permute!(tracks.onchunk.value, p, tracks.proposals.value)
+    _permute!(tracks.onchunk.active, p, tracks.proposals.active)
     return tracks
 end