Initial package functionality

examples/calfit_example.jl

@@ -0,0 +1,129 @@
+# This file is a part of LegendSpecFits.jl, licensed under the MIT License (MIT).
+
+using LegendSpecFits, RadiationSpectra
+using LegendHDF5IO
+using StatsBase, Distributions
+using LinearAlgebra
+using StructArrays
+using ValueShapes, InverseFunctions, BAT, Optim
+using Plots
+
+# Some LH5-file with uncalibrated calibration data enery depositions:
+input_filename = ENV["LEGEND_CALTEST_EDEP_UNCAL"]
+
+detector_no = 40
+
+lhd = LHDataStore(input_filename)
+E_uncal = lhd[string(detector_no)][:]
+
+h_uncal = fit(Histogram, E_uncal, nbins = 20000)
+
+#th228_lines = [583.191, 727.330, 860.564, 1592.53, 1620.50, 2103.53, 2614.50]
+th228_lines = [583.191, 727.330, 860.564, 2103.53, 2614.50]
+h_cal, h_deconv, peakpos, threshold, c, c_precal = RadiationSpectra.calibrate_spectrum(
+    h_uncal, th228_lines,
+    σ = 2.0, threshold = 5.0
+)
+
+peakhists = RadiationSpectra.subhist.(Ref(h_cal), (x -> (x-25, x+25)).(th228_lines))
+peakstats = StructArray(estimate_single_peak_stats.(peakhists))
+
+plot(
+    (
+        plot(normalize(h_uncal, mode = :density), st = :stepbins, yscale = :log10);
+        vline!(peakpos)
+    ),
+    plot.(normalize.(peakhists, mode = :density), st = :stepbins, yscale = :log10)...
+)
+
+
+# Peak-by-peak fit:
+
+peak_fit_plots = Plots.Plot[]
+
+for i in eachindex(peakhists)
+    h = peakhists[i]
+    ps = peakstats[i]
+
+    pseudo_prior = NamedTupleDist(
+        μ = Uniform(ps.peak_pos-10, ps.peak_pos+10),
+        σ = weibull_from_mx(ps.peak_sigma, 2*ps.peak_sigma),
+        n = weibull_from_mx(ps.peak_counts, 2*ps.peak_counts),
+        step_amplitude = weibull_from_mx(ps.mean_background, 2*ps.mean_background),
+        skew_fraction = Uniform(0.01, 0.25),
+        skew_width = LogUniform(0.001, 0.1),
+        background = weibull_from_mx(ps.mean_background, 2*ps.mean_background),
+    )
+
+    f_trafo = BAT.DistributionTransform(Normal, pseudo_prior)
+
+    v_init = mean(pseudo_prior)
+
+    f_fit(x, v) = gamma_peakshape(x, v.μ, v.σ, v.n, v.step_amplitude, v.skew_fraction, v.skew_width) + v.background
+
+    f_loglike = let f_fit=f_fit, h=h
+        v -> hist_loglike(Base.Fix2(f_fit, v), h)
+    end
+
+    opt_r = optimize((-) ∘ f_loglike ∘ inverse(f_trafo), f_trafo(v_init))
+    v_ml = inverse(f_trafo)(Optim.minimizer(opt_r))
+    plt = plot(normalize(h, mode = :density), st = :stepbins, yscale = :log10)
+    plot!(minimum(h.edges[1]):0.1:maximum(h.edges[1]), Base.Fix2(f_fit, v_ml))
+    push!(peak_fit_plots, plt)
+end
+
+plot(peak_fit_plots...)
+
+
+# Global fit over all calibration gamma lines:
+
+th228_lines = [583.191, 727.330, 860.564, 2614.50]
+peakhists = RadiationSpectra.subhist.(Ref(h_cal), (x -> (x-25, x+25)).(th228_lines))
+peakstats = StructArray(estimate_single_peak_stats.(peakhists))
+
+function f_fit(x, v)
+    μ = v.cal_offs + v.cal_slope * v.expected_μ + v.cal_sqr * v.expected_μ^2
+    σ = sqrt(v.σ_enc^2 + (sqrt(μ) * v.σ_fano)^2)
+    gamma_peakshape(
+        x, μ, σ,
+        v.n, v.step_amplitude, v.skew_fraction, v.skew_width
+    ) + v.background
+end
+
+function empirical_prior_from_peakstats(peakstats::StructArray{<:NamedTuple})
+    ps = peakstats
+    mean_rel_sigma = mean(peakstats.peak_sigma ./ sqrt.(peakstats.peak_pos))
+    NamedTupleDist(
+        expected_μ = ConstValueDist(th228_lines),
+        cal_offs = Exponential(5.0),
+        cal_slope = weibull_from_mx.(1.0, 1.01),
+        cal_sqr = Exponential(0.000001),
+        σ_fano = weibull_from_mx.(mean_rel_sigma, 2 .* mean_rel_sigma),
+        σ_enc = Exponential(0.5),
+        n = product_distribution(weibull_from_mx.(ps.peak_counts, 2 .* ps.peak_counts)),
+        step_amplitude = product_distribution(weibull_from_mx.(ps.mean_background, 2 .* ps.mean_background)),
+        skew_fraction = product_distribution(fill(Uniform(0.01, 0.25), length(peakstats))),
+        skew_width = product_distribution(fill(LogUniform(0.001, 0.1), length(peakstats))),
+        background = product_distribution(weibull_from_mx.(ps.mean_background, 2 .* ps.mean_background)),
+    )
+end
+
+pseudo_prior = empirical_prior_from_peakstats(peakstats)
+
+f_trafo = BAT.DistributionTransform(Normal, pseudo_prior)
+
+v_init = mean(pseudo_prior)
+
+f_loglike = let f_fit=f_fit, peakhists=peakhists
+    v -> sum(hist_loglike.(Base.Fix2.(f_fit, expand_vars(v)), peakhists))
+end
+
+opt_r = optimize((-) ∘ f_loglike ∘ inverse(f_trafo), f_trafo(v_init))
+v_ml = inverse(f_trafo)(Optim.minimizer(opt_r))
+
+plot([begin
+    h = peakhists[i]
+    v = expand_vars(v_ml)[i]
+    plot(normalize(h, mode = :density), st = :stepbins, yscale = :log10)
+    plot!(minimum(h.edges[1]):0.1:maximum(h.edges[1]), Base.Fix2(f_fit, v))
+end; for i in eachindex(peakhists)]...)

src/LegendSpecFits.jl

@@ -25,6 +25,10 @@ using Tables
 using Unitful
 using ValueShapes
 
+include("utils.jl")
+include("peakshapes.jl")
+include("likelihoods.jl")
+include("priors.jl")
 include("specfit.jl")
 
 @static if !isdefined(Base, :get_extension)

src/likelihoods.jl

@@ -0,0 +1,22 @@
+# This file is a part of LegendSpecFits.jl, licensed under the MIT License (MIT).
+
+
+"""
+    hist_loglike(f_fit::Base.Callable, h::Histogram{<:Real,1})
+
+Calculate the Poisson log-likelihood of a fit function `f_fit(x)` and a
+histogram `h`. `f_fit` must accept all values `x` on the horizontal axis
+of the histogram.
+
+Currently uses a simple midpoint-rule integration of `f_fit` over the
+bins of `h`.
+"""
+function hist_loglike(f_fit::Base.Callable, h::Histogram{<:Real,1})
+    bin_edges = first(h.edges)
+    counts = h.weights
+    bin_centers = (bin_edges[begin:end-1] .+ bin_edges[begin+1:end]) ./ 2
+    bin_widths = bin_edges[begin+1:end] .- bin_edges[begin:end-1]
+    bin_ll(x, bw, k) = logpdf(Poisson(bw * f_fit(x)), k)
+    sum(Base.Broadcast.broadcasted(bin_ll, bin_centers, bin_widths, counts))
+end
+export hist_loglike

src/peakshapes.jl

@@ -0,0 +1,70 @@
+# This file is a part of LegendSpecFits.jl, licensed under the MIT License (MIT).
+
+"""
+    LegendSpecFits.gauss_pdf(x::Real, μ::Real, σ::Real)
+
+Equivalent to `pdf(Normal(μ, σ), x)`
+"""
+gauss_pdf(x::Real, μ::Real, σ::Real) = inv(σ * sqrt2π) * exp(-((x - μ) / σ)^2  / 2)
+
+
+"""
+    ex_gauss_pdf(x::Real, μ::Real, σ::Real, θ::Real)
+
+The PDF of an
+[Exponentially modified Gaussian distribution](https://en.wikipedia.org/wiki/Exponentially_modified_Gaussian_distribution)
+with Gaussian parameters `μ`, `σ` and exponential scale `θ` at `x`.
+
+It is the PDF of the distribution that descibes the random process
+`rand(Normal(μ, σ)) + rand(Exponential(θ))`.
+"""
+function ex_gauss_pdf(x::Real, μ::Real, σ::Real, θ::Real)
+    R = float(promote_type(typeof(x), typeof(σ), typeof(θ)))
+    x_μ = x - μ
+    gauss_pdf_value = inv(σ * sqrt2π) * exp(-(x_μ/σ)^2 / 2)
+
+    y = if θ < σ * R(10^-6)
+        # Use asymptotic form for very small θ - necessary?
+        R(gauss_pdf_value / (1 + x_μ * θ / σ^2))
+    elseif σ/θ - x_μ/σ < 0
+        # Original:
+        R(inv(2*θ) * exp((σ/θ)^2/2 - x_μ/θ) * erfc(invsqrt2 * (σ/θ - x_μ/σ)))
+    else
+        # More stable, numerically, for small values of θ:
+        R(gauss_pdf_value * σ/θ * sqrthalfπ * erfcx(invsqrt2 * (σ/θ - x_μ/σ)))
+    end
+    @assert !isnan(y) && !isinf(y)
+    return y
+end
+
+
+"""
+    step_gauss(x::Real, μ::Real, σ::Real)
+
+Evaluates the convulution of a Heaviside step function and the
+PDF of `Normal(μ, σ)` at `x`.
+
+The result does not correspond to a PDF as it is not normalizable.
+"""
+step_gauss(x::Real, μ::Real, σ::Real) = erfc( (μ-x) / (sqrt2 * σ) ) / 2
+
+
+"""
+    gamma_peakshape(
+        x::Real, μ::Real, σ::Real, n::Real,
+        step_amplitude::Real, skew_fraction::Real, skew_width::Real
+    )
+    
+Describes the shape of a typical gamma peak in a detector.
+"""
+function gamma_peakshape(
+    x::Real, μ::Real, σ::Real, n::Real,
+    step_amplitude::Real, skew_fraction::Real, skew_width::Real
+)
+    skew = skew_width * μ
+    return n * (
+            (1 - skew_fraction) * gauss_pdf(x, μ, σ) +
+            skew_fraction * ex_gauss_pdf(-x, -μ, σ, skew)
+        ) + step_amplitude * step_gauss(-x, -μ, σ);
+end
+export gamma_peakshape

src/priors.jl

@@ -0,0 +1,17 @@
+# This file is a part of LegendSpecFits.jl, licensed under the MIT License (MIT).
+
+
+"""
+    weibull_from_mx(m::Real, x::Real, p_x::Real = 0.6827)::Weibull
+
+Construct a Weibull distribution with a given median `m` and a given
+`p_x`-quantile `x`.
+
+Useful to construct priors for positive quantities.
+"""
+function weibull_from_mx(m::Real, x::Real, p_x::Real = 0.6827)
+    α = log(-log(1-p_x) / log(2)) / log(x/m)
+    θ = m / log(2)^(1/α)
+    Weibull(α, θ)
+end
+export weibull_from_mx

src/specfit.jl

@@ -1,3 +1,41 @@
 # This file is a part of LegendSpecFits.jl, licensed under the MIT License (MIT).
 
 
+"""
+    estimate_single_peak_stats(h::Histogram)
+
+Estimate statistics/parameters for a single peak in the given histogram `h`.
+
+`h` must only contain a single peak. The peak should have a Gaussian-like
+shape.
+
+Returns a `NamedTuple` with the fields
+
+    * `peak_pos`: estimated position of the peak (in the middle of the peak)
+    * `peak_fwhm`: full width at half maximum (FWHM) of the peak
+    * `peak_sigma`: estimated standard deviation of the peak
+    * `peak_counts`: estimated number of counts in the peak
+    * `mean_background`: estimated mean background value
+"""
+function estimate_single_peak_stats(h::Histogram)
+    W = h.weights
+    E = first(h.edges)
+    peak_amplitude, peak_idx = findmax(W)
+    fwhm_idx_left = findfirst(w -> w >= (first(W) + peak_amplitude) /2, W)
+    fwhm_idx_right = findlast(w -> w >= (last(W) + peak_amplitude) /2, W)
+    peak_max_pos = (E[peak_idx] + E[peak_idx+1]) / 2
+    peak_mid_pos = (E[fwhm_idx_right] + E[fwhm_idx_left]) / 2
+    peak_pos = (peak_max_pos + peak_mid_pos) / 2
+    peak_fwhm = E[fwhm_idx_right] - E[fwhm_idx_left]
+    peak_sigma = peak_fwhm * inv(2*√(2log(2)))
+    #peak_area = peak_amplitude * peak_sigma * sqrt(2*π)
+    mean_background = (first(W) + last(W)) / 2
+    peak_counts = inv(0.761) * (sum(view(W,fwhm_idx_left:fwhm_idx_right)) - mean_background * peak_fwhm)
+
+    (
+        peak_pos = peak_pos, peak_fwhm = peak_fwhm,
+        peak_sigma, peak_counts, mean_background
+    )
+end
+export estimate_single_peak_stats
+

src/utils.jl

@@ -0,0 +1,15 @@
+# This file is a part of LegendSpecFits.jl, licensed under the MIT License (MIT).
+
+"""
+    expand_vars(v::NamedTuple)::StructArray
+
+Expand all fields in `v` (scalars or arrays) to same array size and return
+a `StructArray`.
+"""
+function expand_vars(v::NamedTuple)
+    sz = Base.Broadcast.broadcast_shape((1,), map(size, values(v))...)
+    _expand(x::Real) = Fill(x, sz)
+    _expand(x::AbstractArray) = broadcast((a,b) -> b, Fill(nothing, sz), x)
+    StructArray(map(_expand, v))
+end
+export expand_vars