Correct sampler conditionals, posterior scores and detector edge cases

src/brightness.jl

@@ -34,9 +34,11 @@ Brightness{T}(;
 
 Calculate the log density for `b`'s prior given its current value. `normalization = false` by default. 
 """
-logprior(b::Brightness{T,P}; normalization::Bool = false) where {T,P<:Gamma{T}} =
-    normalization ? logpdf(b.prior, b.value) :
-    (shape(b.prior) - 1) * log(b.value) - b.value / scale(b.prior)
+function logprior(b::Brightness{T,P}; normalization::Bool = false) where {T,P<:Gamma{T}}
+    normalization && return logpdf(b.prior, b.value)
+    logterm = isone(shape(b.prior)) ? zero(T) : (shape(b.prior) - 1) * log(b.value)
+    return logterm - b.value / scale(b.prior)
+end
 
 """
     Δlogprior(b::Brightness{T}, xᵖ::T)

src/chain.jl

@@ -84,15 +84,38 @@ savestride(chain::Chain) =
 
 tosave(chain::Chain, iter::Real) = iter % savestride(chain) == 0
 
-function findmap(chain::Chain; burn_in::Real = 0)
-    ~, i = findmax(chain.logposteriors[burn_in+1:end])
-    return chain.samples[i+burn_in], i + burn_in
+function _findbest(chain::Chain, scorefield::Symbol, burn_in::Real)
+    isinteger(burn_in) && 0 <= burn_in < length(chain) || throw(
+        ArgumentError("burn_in must be a nonnegative whole number smaller than the number of samples"),
+    )
+    best_index = 0
+    best_score = -Inf
+    for i in Int(burn_in)+1:length(chain)
+        score = getproperty(chain.samples[i], scorefield)
+        if isfinite(score) && (best_index == 0 || score > best_score)
+            best_index, best_score = i, score
+        end
+    end
+    best_index > 0 || throw(ArgumentError("no finite $(scorefield) scores remain after burn-in"))
+    return chain.samples[best_index], best_index
 end
 
-function findmle(chain::Chain; burn_in::Real = 0)
-    ~, i = findmax(chain.loglikelihoods[burn_in+1:end])
-    return chain.samples[i+burn_in], i + burn_in
-end
+"""
+    findmap(chain::Chain; burn_in::Real = 0)
+
+Return the sample and index with the largest finite posterior score. `burn_in`
+skips stored samples. Scores are densities in the input coordinate units, not
+marginal probabilities of emitter counts.
+"""
+findmap(chain::Chain; burn_in::Real = 0) = _findbest(chain, :log𝒫, burn_in)
+
+"""
+    findmle(chain::Chain; burn_in::Real = 0)
+
+Return the sample and index with the largest finite likelihood score. `burn_in`
+skips stored samples.
+"""
+findmle(chain::Chain; burn_in::Real = 0) = _findbest(chain, :logℒ, burn_in)
 
 jumpdistances(tracks::AbstractArray{<:Real}) = sqrt.(sum(diff(tracks, dims=1) .^ 2, dims=2))
 function jumpdistances(chain::Chain{T}; burn_in::Integer=0) where {T}
@@ -122,7 +145,7 @@ function extend!(
     𝑇::T,
 ) where {T}
     loglikelihood = get_loglikelihood!(llarray, tracks, brightness, detector, psf)
-    logposterior = get_logposterior(loglikelihood, tracks, msd)
+    logposterior = get_logposterior(loglikelihood, tracks, msd, brightness)
     push!(
         chain.samples,
         Sample(tracks.onchunk, msd, brightness, iter, 𝑇, logposterior, loglikelihood),
@@ -146,10 +169,12 @@ get_logposterior(
     loglikelihood::T,
     tracks::Tracks{T},
     msd::MeanSquaredDisplacement{T},
+    brightness::Brightness{T},
 ) where {T} =
     loglikelihood +
     logprior(tracks.onchunk, msd.value) +
     logprior(msd) +
+    logprior(brightness) +
     logprior(tracks.nemitters)
 
 function parametricMCMC!(
@@ -268,7 +293,14 @@ function runMCMC(;
     parametric::Bool = false,
 ) where {T}
     isnothing(annealing) && (annealing = ConstantAnnealing{T}(1))
-    chain = Chain([Sample(tracks.onchunk, msd, brightness)], sizelimit, annealing)
+    llarray = LogLikelihoodArray{T}(detector.readouts)
+    loglikelihood = get_loglikelihood!(llarray, tracks, brightness, detector, psf)
+    logposterior = get_logposterior(loglikelihood, tracks, msd, brightness)
+    initial = Sample(
+        tracks.onchunk, msd, brightness, 0, temperature(annealing, 0),
+        logposterior, loglikelihood,
+    )
+    chain = Chain([initial], sizelimit, annealing)
     runMCMC!(
         chain = chain,
         tracks = tracks,

src/distributions.jl

@@ -19,7 +19,12 @@ end
 
 Distributions.params(P::DNormal) = P.μ, P.σ
 
-logprior(ℕ::DNormal{T}, x::AbstractArray{T}) where {T} =
-    sum(vec(@. -(x - ℕ.μ)^2 / (2 * ℕ.σ^2)))
+function logprior(ℕ::DNormal{T}, x::AbstractArray{T}; normalization::Bool = false) where {T}
+    quadratic = sum(vec(@. -(x - ℕ.μ)^2 / (2 * ℕ.σ^2)))
+    normalization || return quadratic
+    ndims = length(ℕ.μ)
+    npoints = length(x) ÷ ndims
+    return quadratic - npoints * (sum(log, ℕ.σ) + ndims * log(T(2π)) / 2)
+end
 
 Random.rand!(x::AbstractVecOrMat{T}, P::DNormal{T}) where {T} = _randn!(x, P.μ, P.σ)

src/msd.jl

@@ -37,15 +37,16 @@ logprior(msd::MeanSquaredDisplacement; normalization::Bool = false) =
 """
     update!(msd::MeanSquaredDisplacement{T}, displacement²::AbstractArray{T,3}, 𝑇::T)
 
-Directly sample `msd`'s new value from its posterior given the `displacement²` array and the temperature `𝑇`.
+Directly sample `msd`'s new value from its posterior given the `displacement²` array.
+Only the detector likelihood is tempered, so this conditional does not depend on `𝑇`.
 """
 function update!(
     msd::MeanSquaredDisplacement{T},
     displacement²::AbstractArray{T,3},
     𝑇::T,
 ) where {T}
     isfixed(msd) && return msd
-    Δparams = (length(displacement²), sum(vec(displacement²))) ./ (2 * 𝑇) #* sum(vec(displacement²))) as summing vec is currently faster using CUDA.jl
+    Δparams = (length(displacement²), sum(vec(displacement²))) ./ (2 * one(T))
     msd.value = _rand(msd.prior, Δparams)
     return msd
-end
+end

src/pixeldetector.jl

@@ -48,6 +48,7 @@ function sum_pixel_loglikelihood!(
     llarray::LogLikelihoodArray{T},
     detector::PixelDetector{T},
 ) where {T}
+    @. llarray.pixel = ifelse(iszero(detector.filter), zero(T), llarray.pixel)
     framesum!(llarray.frame, llarray.pixel, detector.filter)
     return llarray
 end
@@ -182,8 +183,13 @@ function set_poisson_means!(
     psf::PointSpreadFunction{T},
 ) where {T}
     set_poisson_mean!(llarray, detector, tracksᵥ, brightnessᵥ₁, psf, 1)
-    llarray.means[2] .=
-        (llarray.means[1] .- detector.darkcounts) .* (brightnessᵥ₂ / brightnessᵥ₁)
+    if iszero(brightnessᵥ₁)
+        set_poisson_mean!(llarray, detector, tracksᵥ, brightnessᵥ₂, psf, 2)
+    else
+        llarray.means[2] .=
+            (llarray.means[1] .- detector.darkcounts) .* (brightnessᵥ₂ / brightnessᵥ₁) .+
+            detector.darkcounts
+    end
     return llarray
 end
 
@@ -213,4 +219,4 @@ function getincident(
         bounds,
         psf,
     )
-end
+end

src/spad.jl

@@ -35,23 +35,46 @@ function SPAD{T}(;
     )
 end
 
+# Keep the hit probability accurate near zero and at saturation.
+_spad_log_hit(mean::T) where {T} =
+    mean <= log(T(2)) ? log(-expm1(-mean)) : log1p(-exp(-mean))
+
+function _spad_loglikelihood(readout, batchsize, mean::T) where {T}
+    hits = iszero(readout) ? zero(T) : readout * _spad_log_hit(mean)
+    misses = readout == batchsize ? zero(T) : (batchsize - readout) * mean
+    return hits - misses
+end
+
+function _spad_Δloglikelihood(readout, batchsize, oldmean::T, newmean::T) where {T}
+    hits = iszero(readout) ? zero(T) :
+           readout * (_spad_log_hit(newmean) - _spad_log_hit(oldmean))
+    misses = readout == batchsize ? zero(T) :
+             (batchsize - readout) * (newmean - oldmean)
+    return hits - misses
+end
+
 function set_pixel_loglikelihood!(
     llarray::LogLikelihoodArray{T},
     detector::SPAD{T},
 ) where {T}
-    @. llarray.pixel =
-        detector.readouts * logexpm1(llarray.means[1]) -
-        detector.batchsize * llarray.means[1]
+    @. llarray.pixel = _spad_loglikelihood(
+        detector.readouts,
+        detector.batchsize,
+        llarray.means[1],
+    )
     return llarray
 end
 
 function set_pixel_Δloglikelihood!(
     Δllarray::LogLikelihoodArray{T},
     detector::SPAD{T},
 ) where {T}
-    @. Δllarray.pixel =
-        detector.readouts * (logexpm1(Δllarray.means[2]) - logexpm1(Δllarray.means[1])) -
-        detector.batchsize * (Δllarray.means[2] - Δllarray.means[1])
+    @. Δllarray.pixel = _spad_Δloglikelihood(
+        detector.readouts,
+        detector.batchsize,
+        Δllarray.means[1],
+        Δllarray.means[2],
+    )
     return Δllarray
 end
 
@@ -62,4 +85,4 @@ function simulate_readouts!(detector::SPAD{T}, means::AbstractArray{T,3}) where
         @. detector.readouts = rand(Binomial(detector.batchsize, -expm1(-means)))
     end
     return detector
-end
+end

src/tracks.jl

@@ -74,13 +74,22 @@ setacceptance!(t::MHTrackChunk; start::Integer, step::Integer) =
 """
     seteffvalue!(tracks::AbstractTrackChunk{T})
 
-Set the effective value of the track chunk. The effective value is calculated as `value / active`. Note that `effvalue == value` when `active` is 1, and `effvalue == Inf` (a particle is infinitely far away so cannot contribute any photons) when `active` is 0. Currently, `active` should ONLY be used in parametric runs.
+Set the effective positions used by the detector. Inactive emitters have infinite lateral coordinates and zero axial coordinates, so they contribute no photons. Currently, `active` should ONLY be used in parametric runs.
 """
 function seteffvalue!(tracks::AbstractTrackChunk{T}) where {T}
-    @. tracks.effvalue = tracks.value / tracks.active
+    seteffvalue!(tracks.effvalue, tracks.value, tracks.active)
     return tracks
 end
 
+function seteffvalue!(effective::AbstractArray{T,3}, value::AbstractArray{T,3}, active::AbstractArray{T,3}) where {T}
+    @. effective = ifelse(iszero(active), T(Inf), value)
+    # Keep the 3D PSF width finite when an emitter is inactive.
+    if size(value, 2) >= 3
+        @views @. effective[:, 3:end, :] = ifelse(iszero(active[:, 3:end, :]), zero(T), value[:, 3:end, :])
+    end
+    return effective
+end
+
 """
     Tracks{T<:AbstractFloat, A<:AbstractArray{T,3}, NT<:NEmitters{T}, TR<:TrackChunk{T}, MH<:MHTrackChunk{T}}
 
@@ -132,8 +141,10 @@ TrackChunk(
 
 function logprior(t::AbstractTrackChunk{T}, msdᵥ::T) where {T}
     setdisplacement²!(t)
-    return -(log(msdᵥ) * length(t.displacement²) + sum(vec(t.displacement²)) / msdᵥ) / 2 -
-           logprior(t.prior, view(t.value, 1, :, :))
+    # Keep Gaussian normalizers because the number of emitters can change.
+    return -((log(msdᵥ) + log(T(2π))) * length(t.displacement²) +
+             sum(vec(t.displacement²)) / msdᵥ) / 2 +
+           logprior(t.prior, view(t.value, 1, :, :); normalization = true)
 end
 
 function Tracks{T}(;
@@ -151,17 +162,17 @@ function Tracks{T}(;
 
     scaling isa Real && (scaling = fill!(similar(value), scaling))
 
-    active = if !isnothing(active)
-        dimsmatch(guess, active, dims = 1) || throw(
-            DimensionMismatch(
-                "size of guess dose not match size of presence in the first dimension",
-            ),
-        )
-        _resize3(active, max_ntracks)
-    else
-        fill!(similar(value, nframes, ndims, max_ntracks), true)
+    activity = fill!(similar(value), one(T))
+    if !isnothing(active)
+        size(active, 1) == nframes && size(active, 2) in (1, ndims) &&
+            size(active, 3) == nguesses || throw(DimensionMismatch("activity must match the guessed frames and emitters"))
+        all(x -> iszero(x) || isone(x), active) || throw(ArgumentError("activity must be binary"))
+        @views activity[:, :, 1:nguesses] .= active
+        @views all(activity[:, :, 1:nguesses] .== activity[:, 1:1, 1:nguesses]) ||
+            throw(ArgumentError("activity must agree across coordinates"))
     end
-    actives = (active, fill!(similar(active), true))
+    active = activity
+    actives = (active, copy(active))
     Δx² = similar(value, nframes - 1, ndims, max_ntracks)
     Δx²s = (Δx², similar(Δx²))
     effs = (similar(value), similar(value))
@@ -192,7 +203,7 @@ function setdisplacement²!(tracks::Tracks, i::Integer = 1)
 end
 
 function seteffvalue!(tracks::Tracks, i::Integer = 1)
-    @. tracks.effvalue[i] = tracks.value[i] / tracks.active[i]
+    seteffvalue!(tracks.effvalue[i], tracks.value[i], tracks.active[i])
     return tracks
 end
 
@@ -257,6 +268,7 @@ function bridge!(x::AbstractArray{T,3}, msd::T, xend::AbstractArray{T,3}) where
 end
 
 function propose!(proposals::MHTrackChunk{T}, tracks::TrackChunk{T}) where {T}
+    copyto!(proposals.active, tracks.active)
     arw_propose!(
         proposals.value,
         tracks.value,
@@ -313,6 +325,8 @@ function Δlogmotion!(
     idx2::StepRange,
     ΣΔΔx²::AbstractVector{T},
 ) where {T}
+    # The last frame has no outgoing edge to overwrite a reused likelihood buffer.
+    fill!(Δ, zero(T))
     @views copyto!(Δ[idx1], ΣΔΔx²[idx1])
     @views Δ[idx2.+1] .+= ΣΔΔx²[idx2]
     return Δ
@@ -388,5 +402,6 @@ function update_onchunk!(
     update!(tracksᵒ, tracksᵖ, msdᵥ, llarray.frame, 1) # update particle positions at odd frame indices
     update!(tracksᵒ, tracksᵖ, msdᵥ, llarray.frame, 2) # update particle positions at even frame indices
     countacceptance!(tracksᵖ)
+    seteffvalue!(tracksᵒ)
     return tracksᵒ
 end

test/activity_regressions.jl

@@ -0,0 +1,59 @@
+using SP2T, Test
+
+@testset "Inactive emitters contribute no signal" begin
+    for T in (Float32, Float64), dims in (2, 3)
+        guess = zeros(T, 3, dims, 1)
+        guess[2, :, :] .= T(0.5)
+        activity = reshape(Bool[false, true, false], 3, 1, 1)
+        tracks = Tracks{T}(
+            guess = guess, prior = DNormal{T}(zeros(T, dims), ones(T, dims)),
+            max_ntracks = 2, scaling = T(0.01), logonprob = T(-1), active = activity,
+        )
+        @test eltype(tracks.active[1]) == T
+        @test tracks.proposals.active == tracks.onchunk.active
+        SP2T.seteffvalue!(tracks.onchunk)
+        @test !any(isnan, tracks.onchunk.effvalue)
+        SP2T.seteffvalue!(tracks)
+        @test tracks.effvalue[1][:, :, 1:1] == tracks.onchunk.effvalue
+        tracks.proposals.active .= one(T)
+        SP2T.propose!(tracks.proposals, tracks.onchunk)
+        @test tracks.proposals.active == tracks.onchunk.active
+        SP2T.seteffvalue!(tracks.proposals)
+
+        detector = SPAD{T}(
+            period = T(1), pixel_size = T(1), darkcounts = fill(T(0.1), 2, 2),
+            cutoffs = (0, Inf), readouts = zeros(UInt16, 2, 2, 3),
+        )
+        psf = dims == 2 ? CircularGaussian{T}(T(0.4), T(1)) :
+              CircularGaussianLorentzian{T}(T(0.4), T(0.8), T(1))
+        ll = SP2T.LogLikelihoodArray{T}(detector.readouts)
+        for positions in (tracks.onchunk.effvalue, tracks.proposals.effvalue)
+            SP2T.set_poisson_mean!(ll, detector, positions, T(2), psf)
+            @test all(isfinite, ll.means[1])
+            @test ll.means[1][:, :, 1] == detector.darkcounts
+            @test ll.means[1][:, :, 3] == detector.darkcounts
+            @test all(ll.means[1][:, :, 2] .> detector.darkcounts)
+        end
+        bad = ones(T, 3, dims, 1)
+        bad[1, 1, 1] = 0
+        @test_throws ArgumentError Tracks{T}(
+            guess = guess, prior = tracks.onchunk.prior, max_ntracks = 1,
+            scaling = T(0.01), logonprob = T(-1), active = bad,
+        )
+    end
+end
+
+@testset "Excluded pixels cannot contaminate likelihood reductions" begin
+    for T in (Float32, Float64), invalid in (T(-Inf), T(NaN))
+        detector = SPAD{T}(
+            period = T(1), pixel_size = T(1), darkcounts = fill(T(0.1), 2, 1),
+            cutoffs = (0, Inf), readouts = zeros(UInt16, 2, 1, 2),
+        )
+        detector.filter[1, 1] = 0
+        ll = SP2T.LogLikelihoodArray{T}(detector.readouts)
+        ll.pixel[1, 1, :] .= invalid
+        ll.pixel[2, 1, :] .= T(-3)
+        SP2T.sum_pixel_loglikelihood!(ll, detector)
+        @test ll.frame == T[-3, -3]
+    end
+end