Merge pull request #77 from verenaaur/aoe_fit

Implementation of A/E fit background function with two EMGs

ext/LegendSpecFitsRecipesBaseExt.jl

@@ -154,10 +154,14 @@ end
     end
 end
 
-@recipe function f(report::NamedTuple{(:v, :h, :f_fit, :f_components, :gof)}; show_label=true, show_fit=true, show_components=true, show_residuals=true, f_fit_x_step_scaling=1/100, _subplot=1)
+@recipe function f(report::NamedTuple{(:v, :h, :f_fit, :f_components, :gof)}; show_label=true, show_fit=true, show_components=true, show_residuals=true, f_fit_x_step_scaling=1/100, _subplot=1, x_label="Energy (keV)")
     f_fit_x_step = ustrip(value(report.v.σ)) * f_fit_x_step_scaling
     bin_centers = collect(report.h.edges[1])[1:end-1] .+ diff(collect(report.h.edges[1]))[1]/2 
-    legend := :topright
+    if x_label == "A/E"
+        legend := :bottomleft
+    else
+        legend := :topright
+    end
     foreground_color_legend := :silver
     background_color_legend := :white
     ylim_max = max(3*value(report.f_fit(report.v.μ)), 3*maximum(report.h.weights))
@@ -174,7 +178,11 @@ end
         label --> ifelse(show_label, "Data", "")
         yscale --> :log10
         bins --> :sqrt
-        xlabel --> "Energy (keV)"
+        if x_label == "A/E"
+            xlabel --> L"A/E\ (\sigma_{A/E}))"
+        else
+            xlabel --> "Energy (keV)"
+        end 
         subplot --> _subplot
         bin_centers, LinearAlgebra.normalize(report.h, mode = :density).weights#LinearAlgebra.normalize(report.h, mode = :density)
     end
@@ -195,11 +203,19 @@ end
                 xlabel := ""
                 xticks --> ([])
             else
-                xlabel --> "Energy (keV)"
+                if x_label == "A/E"
+                    xlabel --> L"A/E\ (\sigma_{A/E}))"
+                else
+                    xlabel --> "Energy (keV)"
+                end 
             end
             ylims --> (ylim_min, ylim_max)
             xlims := (minimum(report.h.edges[1]), maximum(report.h.edges[1]))
-            ylabel := "Counts / $(round(step(report.h.edges[1]), digits=2)) keV"
+            if x_label == "A/E"
+                ylabel := "Counts / $(round(step(report.h.edges[1]), digits=2))"
+            else
+                ylabel := "Counts / $(round(step(report.h.edges[1]), digits=2)) keV"
+            end
             subplot --> _subplot
             minimum(report.h.edges[1]):f_fit_x_step:maximum(report.h.edges[1]), value.(report.f_fit.(minimum(report.h.edges[1]):f_fit_x_step:maximum(report.h.edges[1])))
         end
@@ -218,11 +234,19 @@ end
                             xlabel := ""
                             xticks --> ([])
                         else
-                            xlabel --> "Energy (keV)"
+                            if x_label == "A/E"
+                                xlabel --> L"A/E\ (\sigma_{A/E}))"
+                            else
+                                xlabel --> "Energy (keV)"
+                            end 
                         end
                         ylims --> (ylim_min, ylim_max)
                         xlims := (minimum(report.h.edges[1]), maximum(report.h.edges[1]))
-                        ylabel := "Counts / $(round(step(report.h.edges[1]), digits=2)) keV"
+                        if x_label == "A/E"
+                            ylabel := "Counts / $(round(step(report.h.edges[1]), digits=2))"
+                        else
+                            ylabel := "Counts / $(round(step(report.h.edges[1]), digits=2)) keV"
+                        end
                     end
                     subplot --> _subplot
                     minimum(report.h.edges[1]):f_fit_x_step:maximum(report.h.edges[1]), report.f_components.funcs[component]
@@ -266,12 +290,17 @@ end
                 title := ""
                 markercolor --> :black
                 ylabel --> "Residuals (σ)"
-                xlabel --> "Energy (keV)"
+                if x_label == "A/E"
+                    xlabel --> L"A/E\ (\sigma_{A/E}))"
+                else
+                    xlabel --> "Energy (keV)"
+                end 
                 link --> :x
                 top_margin --> (0, :mm)
                 ylims := (ylims_res_min, ylims_res_max)
                 xlims := (minimum(report.h.edges[1]), maximum(report.h.edges[1]))
                 yscale --> :identity
+                markersize --> 3 #can be changed manually in the code
                 if ylims_res_max == 5
                     yticks := ([-3, 0, 3])
                 end

src/aoe_pseudo_prior.jl

@@ -1,13 +1,24 @@
 function get_aoe_standard_pseudo_prior(h::Histogram, ps::NamedTuple, fit_func::Symbol; fixed_position::Bool=false)
     pprior_base = NamedTupleDist(
         μ = ifelse(fixed_position, ConstValueDist(ps.peak_pos), Normal(ps.peak_pos, 0.5*ps.peak_sigma)),
-        σ = weibull_from_mx(ps.peak_sigma, 2*ps.peak_sigma),
-        n = LogUniform(0.01*ps.peak_counts, 5*ps.peak_counts),
+        σ = weibull_from_mx(ps.peak_sigma, 1.5*ps.peak_sigma),
+        n = weibull_from_mx(ps.peak_counts, 1.5*ps.peak_counts),
         B = LogUniform(0.1*ps.mean_background, 10*ps.mean_background),
-        δ = LogUniform(0.001, 10.0)
+        δ = LogUniform(0.001, 10.0),
+        μ2 = Normal(-15, 5),
+        σ2 = weibull_from_mx(10, 15),
+        B2 = weibull_from_mx(ps.mean_background, 1.5*ps.mean_background),
+        δ2 = weibull_from_mx(5.0, 10.0),
     )
-    if fit_func == :f_fit
-        return pprior_base
+
+    # extract single prior arguments
+    (; μ, σ, n, B, δ, µ2, σ2, B2, δ2) = pprior_base
+
+    # select prior based on fit function
+    if fit_func == :aoe_one_bck
+        NamedTupleDist(; μ, σ, n, B, δ)
+    elseif fit_func == :aoe_two_bck
+        NamedTupleDist(; μ, σ, n, B, δ, μ2, σ2, B2, δ2)
     else
         throw(ArgumentError("fit_func $fit_func not supported for aoe peakshapes"))
     end

src/aoefit.jl

@@ -117,8 +117,7 @@ Fit the A/E Compton bands using the `f_aoe_compton` function consisting of a gau
     * `result`: Dict of NamedTuples of the fit results containing values and errors for each compton band
     * `report`: Dict of NamedTuples of the fit report which can be plotted for each compton band
 """
-function fit_aoe_compton(peakhists::Vector{<:Histogram}, peakstats::StructArray, compton_bands::Array{T},; uncertainty::Bool=false) where T<:Unitful.Energy{<:Real}
-
+function fit_aoe_compton(peakhists::Vector{<:Histogram}, peakstats::StructArray, compton_bands::Array{T},; uncertainty::Bool=false, fit_func::Symbol=:aoe_one_bck) where T<:Unitful.Energy{<:Real}    
     # create return and result dicts
     v_result = Vector{NamedTuple}(undef, length(compton_bands))
     v_report = Vector{NamedTuple}(undef, length(compton_bands))
@@ -132,7 +131,7 @@ function fit_aoe_compton(peakhists::Vector{<:Histogram}, peakstats::StructArray,
         # fit peak
         result_band, report_band = nothing, nothing
         try
-            result_band, report_band = fit_single_aoe_compton(h, ps, ; uncertainty=uncertainty)
+            result_band, report_band = fit_single_aoe_compton(h, ps, ; uncertainty=uncertainty, fit_func=fit_func)
         catch e
             @warn "Error fitting band $band: $e"
             continue
@@ -160,7 +159,7 @@ Perform a fit of the peakshape to the data in `h` using the initial values in `p
     * `result`: NamedTuple of the fit results containing values and errors
     * `report`: NamedTuple of the fit report which can be plotted
 """
-function fit_single_aoe_compton(h::Histogram, ps::NamedTuple; uncertainty::Bool=true, pseudo_prior::NamedTupleDist=NamedTupleDist(empty = true), fit_func::Symbol=:f_fit, background_center::Union{Real,Nothing} = ps.peak_pos, fixed_position::Bool=false)
+function fit_single_aoe_compton(h::Histogram, ps::NamedTuple; uncertainty::Bool=true, pseudo_prior::NamedTupleDist=NamedTupleDist(empty = true), fit_func::Symbol=:aoe_one_bck, background_center::Union{Real,Nothing} = ps.peak_pos, fixed_position::Bool=false)
     # create pseudo priors
     pseudo_prior = get_aoe_pseudo_prior(h, ps, fit_func; pseudo_prior = pseudo_prior, fixed_position = fixed_position)
 

src/aoefit_combined.jl

@@ -9,7 +9,7 @@ Fit a single A/E Compton band using the `f_aoe_compton` function consisting of a
     * `neg_log_likelihood`: The negative log-likelihood of the likelihood fit
     * `report`: Dict of NamedTuples of the fit report which can be plotted for each compton band
 """
-function fit_single_aoe_compton_with_fixed_μ_and_σ(h::Histogram, μ::Number, σ::Number, ps::NamedTuple; fit_func::Symbol = :f_fit, background_center::Union{Real,Nothing} = μ, uncertainty::Bool=false)
+function fit_single_aoe_compton_with_fixed_μ_and_σ(h::Histogram, μ::Number, σ::Number, ps::NamedTuple; fit_func::Symbol = :aoe_one_bck, background_center::Union{Real,Nothing} = μ, uncertainty::Bool=false)
 
     # create pseudo priors
     pseudo_prior = get_aoe_pseudo_prior(h, ps, fit_func;
@@ -105,7 +105,7 @@ function fit_single_aoe_compton_with_fixed_μ_and_σ(h::Histogram, μ::Number, 
 end
 
 # This function calculates the same thing as fit_single_aoe_compton_with_fixed_μ_and_σ, but just returns the value of the negative log-likelihood
-function neg_log_likelihood_single_aoe_compton_with_fixed_μ_and_σ(h::Histogram, μ::Real, σ::Real, ps::NamedTuple; fit_func::Symbol = :f_fit, background_center::Union{Real,Nothing} = μ, optimize::Bool=true)
+function neg_log_likelihood_single_aoe_compton_with_fixed_μ_and_σ(h::Histogram, μ::Real, σ::Real, ps::NamedTuple; fit_func::Symbol = :aoe_one_bck, background_center::Union{Real,Nothing} = μ, optimize::Bool=true)
 
     # create pseudo priors
     pseudo_prior = get_aoe_pseudo_prior(h, ps, fit_func;
@@ -153,7 +153,7 @@ assuming `f_aoe_mu` for `μ` and `f_aoe_sigma` for `σ`.
     * `report_bands`: Dict of NamedTuples of the fit report which can be plotted for each compton band
 """
 function fit_aoe_compton_combined(peakhists::Vector{<:Histogram}, peakstats::StructArray, compton_bands::Array{T}, result_corrections::NamedTuple; 
-    fit_func::Symbol = :f_fit, aoe_expression::Union{String,Symbol}="a / e", e_expression::Union{String,Symbol}="e", uncertainty::Bool=false) where T<:Unitful.Energy{<:Real}
+    fit_func::Symbol = :aoe_one_bck, aoe_expression::Union{String,Symbol}="a / e", e_expression::Union{String,Symbol}="e", uncertainty::Bool=false) where T<:Unitful.Energy{<:Real}
 
     μA = ustrip(result_corrections.μ_compton.par[1])
     μB = ustrip(result_corrections.μ_compton.par[2])

src/aoefit_functions.jl

@@ -8,7 +8,9 @@ MaybeWithEnergyUnits = Union{Real, Unitful.Energy{<:Real}}
 
 function get_aoe_fit_functions(; background_center::Union{Real,Nothing} = nothing)
     merge( 
-        (f_fit = (x, v) -> aoe_compton_peakshape(x, v.μ, v.σ, v.n, v.B, v.δ),),
+        (aoe_one_bck = (x, v) -> aoe_compton_peakshape(x, v.μ, v.σ, v.n, v.B, v.δ),
+        aoe_two_bck = (x, v) -> two_emg_aoe_compton_peakshape(x, v.μ, v.σ, v.n, v.B, v.δ, v.μ2, v.σ2, v.B2, v.δ2),
+        ),
     if isnothing(background_center)
         NamedTuple()
     else
@@ -23,13 +25,20 @@ This function defines the components (signal, low/high-energy tail, backgrounds)
 These component functions are used in the fit-report and in plot receipes 
 """
 function aoe_compton_peakshape_components(fit_func::Symbol; background_center::Union{Real,Nothing} = nothing)
-    if fit_func == :f_fit
+    if fit_func == :aoe_one_bck
         funcs = (f_sig = (x, v) -> aoe_compton_signal_peakshape(x, v.μ, v.σ, v.n),
             f_bck = (x, v) -> aoe_compton_background_peakshape(x, v.μ, v.σ, v.B, v.δ))
+        labels = (f_sig = "Signal", f_bck = "Background")
+        colors = (f_sig = :orangered1, f_bck = :dodgerblue2)
+        linestyles = (f_sig = :solid, f_bck = :dash)
+    elseif fit_func == :aoe_two_bck
+        funcs = (f_sig = (x, v) -> aoe_compton_signal_peakshape(x, v.μ, v.σ, v.n),
+            f_bck_one = (x, v) -> aoe_compton_background_peakshape(x, v.μ, v.σ, v.B, v.δ),
+            f_bck_two = (x, v) -> aoe_compton_background_peakshape(x, v.μ2, v.σ2, v.B2, v.δ2)) #maybe use one function only
+        labels = (f_sig = "Signal", f_bck_one = "First EMG", f_bck_two = "Second EMG")
+        colors = (f_sig = :orangered1, f_bck_one = :dodgerblue2, f_bck_two = :green)
+        linestyles = (f_sig = :solid, f_bck_one = :dash, f_bck_two = :dashdot)
     end 
-    labels = (f_sig = "Signal", f_bck = "Background")
-    colors = (f_sig = :orangered1, f_bck = :dodgerblue2)
-    linestyles = (f_sig = :solid, f_bck = :dash)
     return (funcs = funcs, labels = labels, colors = colors, linestyles = linestyles)
 end
 

src/peakshapes.jl

@@ -250,3 +250,25 @@ function double_gaussian(
 end
 export double_gaussian
 
+
+###############################################################
+# additional peakshape functions for two EMG background
+###############################################################
+
+function two_emg_aoe_compton_peakshape( # total fit function with two EMGs
+    x::Real, 
+    μ::Real, σ::Real, n::Real, background::Real, δ::Real, 
+    μ2::Real, σ2::Real, background2::Real, δ2::Real
+)
+    return iszero(σ) || iszero(σ2) ? zero(x) : n * gauss_pdf(x, μ, σ) + background * ex_gauss_pdf(-x, -μ, σ, δ) + background2 * ex_gauss_pdf(-x, -μ2, σ2, δ2)
+end
+export two_emg_aoe_compton_peakshape
+
+function two_emg_aoe_compton_background_peakshape( # background function with two EMGs
+    x::Real, 
+    μ::Real, σ::Real, background::Real, δ::Real,
+    μ2::Real, σ2::Real, background2::Real, δ2::Real
+)
+    return background * ex_gauss_pdf(-x, -μ, σ, δ) + background2 * ex_gauss_pdf(-x, -μ2, σ2, δ2)
+end
+export two_aoe_compton_background_peakshape