diff --git a/src/brightness.jl b/src/brightness.jl index 23ee976..dc83842 100644 --- a/src/brightness.jl +++ b/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) diff --git a/src/chain.jl b/src/chain.jl index cc35dc0..5256517 100644 --- a/src/chain.jl +++ b/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, diff --git a/src/distributions.jl b/src/distributions.jl index dd774b3..015e33c 100644 --- a/src/distributions.jl +++ b/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.σ) diff --git a/src/msd.jl b/src/msd.jl index 71be82f..7246231 100644 --- a/src/msd.jl +++ b/src/msd.jl @@ -37,7 +37,8 @@ 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}, @@ -45,7 +46,7 @@ function update!( 𝑇::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 \ No newline at end of file +end diff --git a/src/pixeldetector.jl b/src/pixeldetector.jl index ac2d87f..8ed0982 100644 --- a/src/pixeldetector.jl +++ b/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 \ No newline at end of file +end diff --git a/src/spad.jl b/src/spad.jl index aba40eb..021a082 100644 --- a/src/spad.jl +++ b/src/spad.jl @@ -35,13 +35,33 @@ 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 @@ -49,9 +69,12 @@ 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 \ No newline at end of file +end diff --git a/src/tracks.jl b/src/tracks.jl index 8e739a7..b5a5ecd 100644 --- a/src/tracks.jl +++ b/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 diff --git a/test/activity_regressions.jl b/test/activity_regressions.jl new file mode 100644 index 0000000..af4d821 --- /dev/null +++ b/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 diff --git a/test/detector_regressions.jl b/test/detector_regressions.jl new file mode 100644 index 0000000..d54abe4 --- /dev/null +++ b/test/detector_regressions.jl @@ -0,0 +1,57 @@ +using Test +using SP2T + +function spad_reference(readout, batchsize, mean) + setprecision(BigFloat, 2048) do + photonmean = BigFloat(mean) + hit = -expm1(-photonmean) + miss = exp(-photonmean) + probability = hit^readout * miss^(batchsize - readout) + log(probability) + end +end + +@testset "SPAD likelihood against binomial hit/miss probabilities" begin + for T in (Float32, Float64) + n = UInt16(8) + detector = SPAD{T}( + period = 1, + pixel_size = 1, + darkcounts = ones(T, 1, 1), + cutoffs = (0, Inf), + readouts = zeros(UInt16, 1, 1, 1), + batchsize = n, + ) + llarray = SP2T.LogLikelihoodArray{T}(detector.readouts) + for k in UInt16.((0, 1, 4, 8)), mean in T.((0, floatmin(T), 1e-8, 0.1, 1, 10, 100, 1e5)) + detector.readouts[1] = k + llarray.means[1][1] = mean + SP2T.set_pixel_loglikelihood!(llarray, detector) + expected = T(spad_reference(k, n, mean)) + @test !isnan(llarray.pixel[1]) + @test llarray.pixel[1] ≈ expected rtol=32eps(T) atol=16nextfloat(zero(T)) + end + for k in UInt16.((0, 1, 4, 8)), (oldmean, newmean) in ( + (T(1e-8), T(2e-8)), + (T(0.1), T(0.2)), + (T(10), T(11)), + (T(100), T(101)), + (T(1e5), T(1e5 + 1)), + ) + detector.readouts[1] = k + llarray.means[1][1] = oldmean + llarray.means[2][1] = newmean + SP2T.set_pixel_Δloglikelihood!(llarray, detector) + expected = T(spad_reference(k, n, newmean) - spad_reference(k, n, oldmean)) + @test llarray.pixel[1] ≈ expected rtol=64eps(T) atol=16nextfloat(zero(T)) + end + @test SP2T._spad_loglikelihood(UInt16(0), n, zero(T)) == zero(T) + @test SP2T._spad_loglikelihood(n, n, T(Inf)) == zero(T) + @test SP2T._spad_loglikelihood(UInt16(0), n, T(Inf)) == -Inf + @test SP2T._spad_Δloglikelihood(UInt16(0), n, zero(T), T(0.2)) ≈ -T(n) * T(0.2) + @test SP2T._spad_Δloglikelihood(UInt16(0), n, zero(T), zero(T)) == zero(T) + @test SP2T._spad_Δloglikelihood(n, n, T(Inf), T(Inf)) == zero(T) + @test SP2T._spad_Δloglikelihood(n, n, zero(T), one(T)) == Inf + @test SP2T._spad_Δloglikelihood(n, n, one(T), zero(T)) == -Inf + end +end diff --git a/test/runtests.jl b/test/runtests.jl index fab7d18..da3ce33 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,11 @@ using SP2T using Test +include("sampler_regressions.jl") +include("score_regressions.jl") +include("detector_regressions.jl") +include("activity_regressions.jl") + println("Running PSF-related tests...") @test SP2T.gaussianpeak(0.09890858623710756, 0.133) ≈ 0.2486333227348026 @@ -18,4 +23,4 @@ end @testset let (σ₀, z₀) = SP2T.getσ₀z₀(1.45, 1.515, 0.665) @test σ₀ ≈ 0.09890858623710756 @test z₀ ≈ 0.2800714501487841 -end \ No newline at end of file +end diff --git a/test/sampler_regressions.jl b/test/sampler_regressions.jl new file mode 100644 index 0000000..b185890 --- /dev/null +++ b/test/sampler_regressions.jl @@ -0,0 +1,114 @@ +using Distributions +using Random +using SP2T +using Test + +@testset "Brightness proposals preserve detector background" begin + for T in (Float32, Float64) + darkcounts = T[0.12 0.07; 0.03 0.19] + readouts = reshape(UInt16[0, 1, 2, 3, 1, 0, 2, 1, 0, 2, 3, 1], 2, 2, 3) + common = ( + period = T(1), + pixel_size = T(1), + darkcounts = darkcounts, + cutoffs = (T(0), T(Inf)), + readouts = readouts, + ) + detectors = ( + SPAD{T}(; common..., batchsize = 3), + EMCCD{T}(; common..., offset = T(0), gain = T(2), variance = T(1)), + ) + psf = CircularGaussian{T}(T(0.4), T(1)) + positions = reshape(T[0.4, 0.6, 0.8, 0.5, 0.7, 0.9], 3, 2, 1) + for detector in detectors, (old, proposed) in ((2, 2), (2, 3), (2, 0), (0, 3), (0, 0)) + llarray = SP2T.LogLikelihoodArray{T}(readouts) + reference = SP2T.LogLikelihoodArray{T}(readouts) + SP2T.set_poisson_mean!(reference, detector, positions, T(old), psf) + expected_old = copy(reference.means[1]) + old_loglikelihood = SP2T.get_loglikelihood!(reference, detector) + SP2T.set_poisson_mean!(reference, detector, positions, T(proposed), psf) + expected_proposed = copy(reference.means[1]) + proposed_loglikelihood = SP2T.get_loglikelihood!(reference, detector) + + SP2T.set_poisson_means!(llarray, detector, positions, T(old), T(proposed), psf) + @test llarray.means[1] ≈ expected_old + @test llarray.means[2] ≈ expected_proposed + @test all(isfinite, llarray.means[2]) + @test SP2T.get_Δloglikelihood!(llarray, detector) ≈ + proposed_loglikelihood - old_loglikelihood atol = 100 * eps(T) + if old == proposed + @test llarray.means[1] ≈ llarray.means[2] + end + end + end +end + +@testset "Motion contributions do not retain endpoint likelihoods" begin + for T in (Float32, Float64), nframes in (1, 2, 5, 6), parity in (1, 2) + edges = T[(-1)^i * (i + 0.25) for i in 1:nframes-1] + expected = zeros(T, nframes) + for frame in parity:2:nframes + frame > 1 && (expected[frame] += edges[frame-1]) + frame < nframes && (expected[frame] += edges[frame]) + end + scratch = fill(T(17), nframes) + SP2T.Δlogmotion!( + scratch, + parity:2:nframes-1, + (3-parity):2:nframes-1, + edges, + ) + @test scratch == expected + end +end + +@testset "Fixed-count sampling refreshes effective positions" begin + for T in (Float32, Float64) + initial = fill(T(0.5), 4, 2, 1) + tracks = Tracks{T}( + guess = initial, + prior = DNormal{T}(fill(T(0.5), 2), fill(T(10), 2)), + max_ntracks = 1, + scaling = T(0.05), + logonprob = T(-1), + ) + detector = SPAD{T}( + period = T(1), + pixel_size = T(1), + darkcounts = fill(T(0.1), 2, 2), + cutoffs = (T(0), T(Inf)), + readouts = zeros(UInt16, 2, 2, 4), + ) + llarray = SP2T.LogLikelihoodArray{T}(detector.readouts) + psf = CircularGaussian{T}(T(0.4), T(1)) + msd = MeanSquaredDisplacement{T}(guess = T(1e6), fixed = true) + brightness = Brightness{T}(guess = T(0), fixed = true) + Random.seed!(932) + SP2T.parametricMCMC!(tracks, msd, brightness, llarray, detector, psf, T(1)) + @test any(>(0), tracks.proposals.accepted) + @test tracks.onchunk.value != initial + @test tracks.onchunk.effvalue == tracks.onchunk.value ./ tracks.onchunk.active + end +end + +@testset "MSD conditional uses the untempered motion prior" begin + for T in (Float32, Float64), nsteps in (0, 3), temperature in (1, 4) + displacement² = reshape(T.(1:2nsteps) ./ T(10), nsteps, 2, 1) + priorparams = (T(3), T(0.7)) + posterior = InverseGamma( + priorparams[1] + T(length(displacement²)) / T(2), + priorparams[2] + sum(displacement²) / T(2), + ) + msd = MeanSquaredDisplacement{T}(guess = T(1), priorparams = priorparams) + Random.seed!(467) + expected = rand(posterior) + Random.seed!(467) + SP2T.update!(msd, displacement², T(temperature)) + @test msd.value == expected + + msd.fixed = true + fixed_value = msd.value + SP2T.update!(msd, displacement², T(temperature)) + @test msd.value == fixed_value + end +end diff --git a/test/score_regressions.jl b/test/score_regressions.jl new file mode 100644 index 0000000..082b63d --- /dev/null +++ b/test/score_regressions.jl @@ -0,0 +1,134 @@ +using Distributions: Normal, logpdf + +function score_test_tracks(::Type{T}, nemitters::Integer) where {T} + positions = Array{T}(undef, 3, 2, nemitters) + for m in 1:nemitters, d in 1:2, f in 1:3 + positions[f, d, m] = T(0.1 * (f + d + m)) + end + return Tracks{T}( + guess = positions, + prior = DNormal{T}(T[0.2, 0.5], T[0.4, 0.7]), + max_ntracks = max(nemitters, 1), + scaling = T(0.1), + logonprob = -2, + ) +end + +@testset "Normalized track scores" begin + for T in (Float32, Float64), nemitters in 0:2 + tracks = score_test_tracks(T, nemitters) + x = tracks.onchunk.value + prior = tracks.onchunk.prior + msd = T(0.15) + initial = zero(T) + motion = zero(T) + for m in 1:nemitters, d in 1:2 + initial += logpdf(Normal(prior.μ[d], prior.σ[d]), x[1, d, m]) + for f in 2:3 + motion += logpdf(Normal(zero(T), sqrt(msd)), x[f, d, m] - x[f-1, d, m]) + end + end + @test SP2T.logprior(prior, view(x, 1, :, :); normalization = true) ≈ initial + @test SP2T.logprior(tracks.onchunk, msd) ≈ initial + motion + end + + for T in (Float32, Float64) + tracks = score_test_tracks(T, 1) + prior = tracks.onchunk.prior + tracks.onchunk.value .= reshape(prior.μ, 1, 2, 1) + centered = SP2T.logprior(tracks.onchunk, T(0.15)) + tracks.onchunk.value[:, 1, :] .+= prior.σ[1] + shifted = SP2T.logprior(tracks.onchunk, T(0.15)) + @test shifted - centered ≈ T(-0.5) + end +end + +@testset "Posterior includes brightness prior" begin + for T in (Float32, Float64) + tracks = score_test_tracks(T, 1) + msd = MeanSquaredDisplacement{T}(guess = 0.15, priorparams = (2, 0.1)) + brightness = Brightness{T}(guess = 1.4, priorparams = (2, 3)) + first_value = brightness.value + first_score = SP2T.get_logposterior(T(-3), tracks, msd, brightness) + brightness.value = T(3.1) + second_score = SP2T.get_logposterior(T(-3), tracks, msd, brightness) + expected = logpdf(brightness.prior, brightness.value) - + logpdf(brightness.prior, first_value) + @test second_score - first_score ≈ expected + end +end + +@testset "Zero-brightness prior boundary" begin + for T in (Float32, Float64) + brightness = Brightness{T}(guess = 0, priorparams = (1, 3), fixed = true) + @test SP2T.logprior(brightness) == zero(T) + @test SP2T.logprior(brightness; normalization = true) ≈ -log(T(3)) + for (shape, expected) in ((2, -Inf), (0.5, Inf)) + brightness = Brightness{T}(guess = 0, priorparams = (shape, 3), fixed = true) + @test SP2T.logprior(brightness) == T(expected) + @test SP2T.logprior(brightness; normalization = true) == T(expected) + end + end +end + +@testset "MAP and MLE select finite scores" begin + for T in (Float32, Float64) + scores = [(NaN, NaN), (3, 5), (Inf, Inf), (4, 2), (4, 7), (-Inf, -Inf), (NaN, NaN)] + samples = [ + Sample(zeros(T, 2, 2, 0), T(0.15), one(T), i - 1, one(T), T(p), T(l)) + for (i, (p, l)) in enumerate(scores) + ] + chain = Chain(samples, 20, ConstantAnnealing{T}(1)) + @test findmap(chain)[2] == 4 + @test findmap(chain)[1] === samples[4] + @test findmle(chain)[2] == 5 + @test findmle(chain)[1] === samples[5] + @test findmap(chain; burn_in = 4)[2] == 5 + @test findmle(chain; burn_in = 2.0)[2] == 5 + for select in (findmap, findmle) + for invalid in (-1, 0.5, length(chain), Inf, NaN) + @test_throws ArgumentError select(chain; burn_in = invalid) + end + @test_throws ArgumentError select(chain; burn_in = 5) + @test_throws ArgumentError select(Chain(samples[1:1], 20, chain.annealing)) + @test_throws ArgumentError select(Chain(samples[1:0], 20, chain.annealing)) + end + end +end + +@testset "Initial sample has computed scores" begin + for T in (Float32, Float64), nemitters in 0:1, brightness_value in (0, 1.4) + tracks = score_test_tracks(T, nemitters) + msd = MeanSquaredDisplacement{T}(guess = 0.15, priorparams = (2, 0.1)) + brightness = Brightness{T}(guess = brightness_value, priorparams = (1, 3), fixed = true) + detector = SPAD{T}( + period = 0.01, + pixel_size = 0.1, + darkcounts = fill(T(0.05), 2, 2), + cutoffs = (0, Inf), + readouts = zeros(UInt16, 2, 2, 3), + ) + psf = CircularGaussian{T}(0.15, 0.1) + chain = runMCMC( + tracks = tracks, + msd = msd, + brightness = brightness, + detector = detector, + psf = psf, + niters = 0, + annealing = ConstantAnnealing{T}(2), + ) + sample = only(chain.samples) + @test sample.iteration == 0 + @test sample.𝑇 == T(2) + @test isfinite(sample.logℒ) + @test isfinite(sample.log𝒫) + incident = SP2T.getincident( + tracks.onchunk.value, brightness.value, detector.darkcounts, detector.pxbounds, psf, + ) + @test sample.logℒ ≈ -sum(incident) + @test sample.log𝒫 ≈ SP2T.get_logposterior(sample.logℒ, tracks, msd, brightness) + @test findmap(chain)[2] == 1 + @test findmle(chain)[2] == 1 + end +end