legend-exp · apearsonn · Oct 16, 2024 · Oct 30, 2024 · fhagemann · Oct 30, 2024
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -1 +1,15 @@
 # LegendSpecFits.jl
+
+This package provides functionality to fit probability distribution functions to histograms.
+
+This package is part of the [LegendJuliaRegistry](https://github.com/legend-exp/LegendJuliaRegistry), meaning that you can download the package as follows:
+```julia
+import Pkg
+
+# add the LegendJuliaRegistry
+Pkg.Registry.add(Pkg.RegistrySpec(url = "https://github.com/legend-exp/LegendJuliaRegistry"))
+
+# add LegendSpecFits
+Pkg.add("LegendSpecFits")
+```
+
diff --git a/src/aoe_cut.jl b/src/aoe_cut.jl
@@ -1,11 +1,22 @@
 
 """
-    prepare_dep_peakhist(e::Array{T}, dep::T,; relative_cut::T=0.5, n_bins_cut::Int=500) where T<:Real
+    prepare_dep_peakhist(e::Array{T}, dep::Quantity{T},; relative_cut::T=0.5, n_bins_cut::Int=500, uncertainty::Bool=true) where T<:Real
 
 Prepare an array of uncalibrated DEP energies for parameter extraction and calibration.
+
+# Arguments
+    * 'e': Calibrated energies
+    * 'dep': DEP energies
+
+#Keywords
+    * 'relative_cut': 
+    * 'n_bins_cut': Number of histogram bins to cut from plot
+    * 'uncertainty': Energy uncertainty
+
 # Returns
-- `result`: Result of the initial fit
-- `report`: Report of the initial fit
+    * `result`: Result of the initial fit
+    * `report`: Report of the initial fit
+
 """
 function prepare_dep_peakhist(e::Array{T}, dep::Quantity{T},; relative_cut::T=0.5, n_bins_cut::Int=500, uncertainty::Bool=true) where T<:Real
     # get cut window around peak
@@ -25,13 +36,65 @@
 export prepare_dep_peakhist
 
 
+"""
+    get_n_after_aoe_cut(aoe_cut::Unitful.RealOrRealQuantity, aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peak::T, window::Vector{T}, bin_width::T, result_before::NamedTuple, peakstats::NamedTuple; uncertainty::Bool=true, fixed_position::Bool=true) where T<:Unitful.Energy{<:Real}
+
+Get the number of counts after a AoE cut value `aoe_cut` for a given `peak` and `window` size while performing a peak fit with fixed position. The number of counts is determined by fitting the peak with a pseudo prior for the peak position.
+
+# Arguments
+    * 'aoe_cut': A/E cut value
+    * 'aoe': A/E
+    * 'e': Calibrated energies
+    * 'peak': Data peak 
+    * 'window': histogram window size of counts used
+    * 'bin_width': Histogram bin widths
+    * 'result_before': 
+    * 'peakstats': Peak statistics 
+
+# Keywords
+    * 'uncertainty': Data uncertainty
+    * 'fixed_position': Position of peak 
+
+# Returns
+    * `n`: Number of counts after the cut
+
+"""
+function get_n_after_aoe_cut(aoe_cut::Unitful.RealOrRealQuantity, aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peak::T, window::Vector{T}, bin_width::T, result_before::NamedTuple, peakstats::NamedTuple; uncertainty::Bool=true, fixed_position::Bool=true) where T<:Unitful.Energy{<:Real}
+    # get energy after cut and create histogram
+    peakhist = fit(Histogram, ustrip.(e[aoe .> aoe_cut]), ustrip(peak-first(window):bin_width:peak+last(window)))
+    # create pseudo_prior with known peak sigma in signal for more stable fit
+    pseudo_prior = if fixed_position
+        NamedTupleDist(σ = Normal(result_before.σ, 0.1), μ = ConstValueDist(result_before.μ))
+    else
+        NamedTupleDist(σ = Normal(result_before.σ, 0.1), )
+    end
+    # fit peak and return number of signal counts
+    result, _ = fit_single_peak_th228(peakhist, peakstats,; uncertainty=uncertainty, low_e_tail=false, pseudo_prior=pseudo_prior)
+    return result.n
+end
+export get_n_after_aoe_cut
+
+
 """
     get_sf_after_aoe_cut(aoe_cut::Unitful.RealOrRealQuantity, aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peak::T, window::Vector{T}, bin_width::T, result_before::NamedTuple; uncertainty::Bool=true, fit_func::Symbol=:gamma_def) where T<:Unitful.Energy{<:Real}
 
 Get the survival fraction after a AoE cut value `aoe_cut` for a given `peak` and `window` size from a combined fit to the survived and cut histograms.
 
+# Arguments
+    * 'aoe_cut': A/E cut value
+    * 'aoe': A/E
+    * 'e': Calibrated energies
+    * 'peak': Peak name
+    * 'window': Data window in energy
+    * 'bin_width': Histogram bin widths
+    * 'result_before': Survival fraction before A/E cut
+
+# Keywords
+    * 'uncertainty': uncertainty of survival fraction 
+    * `fit_func`: Fit function
+
 # Returns
-- `sf`: Survival fraction after the cut
+    * `sf`: Survival fraction after the cut
 """
 function get_sf_after_aoe_cut(aoe_cut::Unitful.RealOrRealQuantity, aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peak::T, window::Vector{T}, bin_width::T, result_before::NamedTuple; uncertainty::Bool=true, fit_func::Symbol=:gamma_def) where T<:Unitful.Energy{<:Real}
     # get energy after cut and create histogram
@@ -51,12 +114,30 @@
             cut_search_interval::Tuple{<:Unitful.RealOrRealQuantity, <:Unitful.RealOrRealQuantity}=(-25.0*unit(first(aoe)), 1.0*unit(first(aoe))), 
             bin_width_window::T=3.0u"keV", max_e_plot::T=3000.0u"keV",  plot_window::Vector{<:T}=[12.0, 50.0]u"keV",
             fixed_position::Bool=true, fit_func::Symbol=:gamma_def, uncertainty::Bool=true) where T<:Unitful.Energy{<:Real}
+
 Get the AoE cut value for a given `dep` and `window` size while performing a peak fit with fixed position. The AoE cut value is determined by finding the cut value for which the number of counts after the cut is equal to `dep_sf` times the number of counts before the cut.
 The algorhithm utilizes a root search algorithm to find the cut value with a relative tolerance of `rtol`.
+
+# Arguments
+    * 'aoe': A/E
+    * 'e': Calibrated energies
+
+# Keywords
+    * 'dep': DEP energies
+    * 'window': Data window in energy
+    * 'dep_sf': Surrival fraction of DEP energies
+    * 'cut_search_interval': Data interval to search for a cut value
+    * 'rtol': Relative tolerance
+    * 'bin_width_window': Specified window around the peak in which the bin algorithm is applied
+    * 'fixed_position': Fixed position of cut
+    * 'sigma_high_sided': Fixed upper limit cut
+    * 'uncertainty': Uncertainty
+    * 'maxiters': Maximum iterations
+
 # Returns
-- `cut`: AoE cut value
-- `n0`: Number of counts before the cut
-- `nsf`: Number of counts after the cut
+    * `cut`: AoE cut value
+    * `n0`: Number of counts before the cut
+    * `nsf`: Number of counts after the cut
 """
 function get_low_aoe_cut(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T},; 
             dep::T=1592.53u"keV", window::Vector{<:T}=[12.0, 10.0]u"keV", dep_sf::Float64=0.9, rtol::Float64=0.001, maxiters::Int=300, sigma_high_sided::Float64=Inf,
@@ -117,9 +198,24 @@
     get_peaks_surrival_fractions(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peaks::Vector{<:T}, peak_names::Vector{Symbol}, windows::Vector{<:Tuple{T, T}}, aoe_cut::Unitful.RealOrRealQuantity,; uncertainty::Bool=true, bin_width_window::T=2.0u"keV", sigma_high_sided::Unitful.RealOrRealQuantity=Inf*unit(first(aoe)), fit_funcs::Vector{Symbol}=fill(:gamma_def, length(peaks))) where T<:Unitful.Energy{<:Real}
 
 Get the surrival fraction of a peak after a AoE cut value `aoe_cut` for a given `peak` and `window` size while performing a peak fit with fixed position.
+
+# Arguments
+    * 'aoe': A/E
+    * 'e': Calibrated energies
+    * 'peaks': Data range of several peaks
+    * 'peak_names': Name of the peak
+    * 'windows': Energy data window 
+    * 'aoe_cut': A/E cut value
+
+# Keywords
+    * 'Uncertainty':
+    * 'bin_width_window': Specified window around the peak where the binning algorithm is applied to
+    * 'low_e_tail': Low energy tail
+    * 'sigma_high_sided': Fixed upper limit cut
+
 # Return 
-- `result`: Dict of results for each peak
-- `report`: Dict of reports for each peak
+    * `result`: Dict of results for each peak
+    * `report`: Dict of reports for each peak
 """
 function get_peaks_surrival_fractions(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peaks::Vector{<:T}, peak_names::Vector{Symbol}, windows::Vector{<:Tuple{T, T}}, aoe_cut::Unitful.RealOrRealQuantity,; uncertainty::Bool=true, bin_width_window::T=2.0u"keV", sigma_high_sided::Unitful.RealOrRealQuantity=Inf*unit(first(aoe)), fit_funcs::Vector{Symbol}=fill(:gamma_def, length(peaks))) where T<:Unitful.Energy{<:Real}
     @assert length(peaks) == length(peak_names) == length(windows) "Length of peaks, peak_names and windows must be equal"
@@ -149,16 +245,30 @@
 
 
 """
-    get_peak_surrival_fraction(aoe::Array{T}, e::Array{T}, peak::T, window::Array{T}, aoe_cut::T,; uncertainty::Bool=true, low_e_tail::Bool=true) where T<:Real
+    get_peak_surrival_fraction(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peak::T, window::Vector{T}, aoe_cut::Unitful.RealOrRealQuantity,; uncertainty::Bool=true, lq_mode::Bool=false, low_e_tail::Bool=true, bin_width_window::T=2.0u"keV", sigma_high_sided::Unitful.RealOrRealQuantity=NaN) where T<:Unitful.Energy{<:Real}
 
 Get the surrival fraction of a peak after a AoE cut value `aoe_cut` for a given `peak` and `window` size while performing a peak fit with fixed position.
 
+# Arguments
+    * 'aoe': A/E
+    * 'e': Calibrated energies
+    * 'peak': Peak name
+    * 'window': Data window in energy
+    * 'aoe_cut': A/E cut value
+
+# Keywords
+    * 'Uncertainty': Uncertainty 
+    * 'lq_mode': Inverts the cut logic
+    * 'low_e_tail': Low energy tail
+    * 'bin_width_window': Specified window around the peak to apply the binning algorithm
+    * 'sigma_high_sided': Fixed upper limit cut
+
 # Returns
-- `peak`: Peak position
-- `n_before`: Number of counts before the cut
-- `n_after`: Number of counts after the cut
-- `sf`: Surrival fraction
-- `err`: Uncertainties
+    * `peak`: Peak position
+    * `n_before`: Number of counts before the cut
+    * `n_after`: Number of counts after the cut
+    * `sf`: Surrival fraction
+    * `err`: Uncertainties
 """
 function get_peak_surrival_fraction(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peak::T, window::Vector{T}, aoe_cut::Unitful.RealOrRealQuantity,; 
     uncertainty::Bool=true, lq_mode::Bool=false, bin_width_window::T=2.0u"keV", sigma_high_sided::Unitful.RealOrRealQuantity=Inf*unit(first(aoe)), fit_func::Symbol=:gamma_def) where T<:Unitful.Energy{<:Real}
@@ -211,16 +321,27 @@
 
 
 """
-    get_continuum_surrival_fraction(aoe:::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, center::T, window::T, aoe_cut::Unitful.RealOrRealQuantity,; sigma_high_sided::Unitful.RealOrRealQuantity=NaN) where T<:Unitful.Energy{<:Real}
+    get_continuum_surrival_fraction(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, center::T, window::T, aoe_cut::Unitful.RealOrRealQuantity,; lq_mode::Bool=false, sigma_high_sided::Unitful.RealOrRealQuantity=NaN) where T<:Unitful.Energy{<:Real}
 
 Get the surrival fraction of a continuum after a AoE cut value `aoe_cut` for a given `center` and `window` size.
 
+# Arguments
+    * 'aoe': A/E
+    * 'e': Calibrated energies
+    * 'center': Center of the fit
+    * 'window': Data window in energy
+    * 'aoe_cut': A/E cut value
+
+# Keywords
+    * 'lq_mode': Inverted cut logic
+    * 'sigma_high_sided': Fixed upper limit cut 
+
 # Returns
-- `center`: Center of the continuum
-- `window`: Window size
-- `n_before`: Number of counts before the cut
-- `n_after`: Number of counts after the cut
-- `sf`: Surrival fraction
+    * `window`: Window size
+    * `n_before`: Number of counts before the cut
+    * `n_after`: Number of counts after the cut
+    * `sf`: Surrival fraction
+
 """
 function get_continuum_surrival_fraction(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, center::T, window::T, aoe_cut::Unitful.RealOrRealQuantity,; lq_mode::Bool=false, sigma_high_sided::Unitful.RealOrRealQuantity=Inf*unit(first(aoe))) where T<:Unitful.Energy{<:Real}
     # scale unit

diff --git a/src/aoe_fit_calibration.jl b/src/aoe_fit_calibration.jl
@@ -1,17 +1,28 @@
 @. f_aoe_sigma(x, p) = sqrt(p[1]^2 + p[2]^2/x^2)
 f_aoe_mu(x, p) = p[1] .+ p[2].*x
 """
-fit_aoe_corrections(e::Array{<:Unitful.Energy{<:Real}}, μ::Array{<:Real}, σ::Array{<:Real})
+    fit_aoe_corrections(e::Array{<:Unitful.Energy{<:Real}}, μ::Array{<:Real}, σ::Array{<:Real}; aoe_expression::Union{String,Symbol}="a / e", e_expression::Union{String,Symbol}="e")
 
 Fit the corrections for the AoE value of the detector.
+
+# Arguments
+    * 'e': Calibrated energies
+    * 'μ': Mean values
+    * 'σ': Sigma values
+
+# Keywords
+    * 'aoe_expression': A/E expression
+    * 'e_expression': Calibrated energy expression
+
 # Returns
-- `e`: Energy values
-- `μ`: Mean values
-- `σ`: Sigma values
-- `μ_scs`: Fit result for the mean values
-- `f_μ_scs`: Fit function for the mean values
-- `σ_scs`: Fit result for the sigma values
-- `f_σ_scs`: Fit function for the sigma values
+    * `e`: Energy values
+    * `μ`: Mean values
+    * `σ`: Sigma values
+    * `μ_scs`: Fit result for the mean values
+    * `f_μ_scs`: Fit function for the mean values
+    * `σ_scs`: Fit result for the sigma values
+    * `f_σ_scs`: Fit function for the sigma values
+
 """
 function fit_aoe_corrections(e::Array{<:Unitful.Energy{<:Real}}, μ::Array{<:Real}, σ::Array{<:Real}; aoe_expression::Union{String,Symbol}="a / e", e_expression::Union{String,Symbol}="e")
     # fit compton band mus with linear function

diff --git a/src/aoe_pseudo_prior.jl b/src/aoe_pseudo_prior.jl
@@ -1,3 +1,22 @@
+"""
+    get_aoe_standard_pseudo_prior(h::Histogram, ps::NamedTuple, fit_func::Symbol; fixed_position::Bool=false)
+
+Gets the A/E of a histogram using a standard pseudo prior. **
+
+TO DO: add description of function and returns
+
+# Arguments
+    * 'h': Histogram data
+    * 'ps': Peak statistics 
+    * 'fit_func': Fit function
+
+# Keywords
+    * 'fixed_position': 
+
+# Returns
+    * 'pprior_base': Base pseudo prior
+"""
+
 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)),
@@ -24,6 +43,24 @@ function get_aoe_standard_pseudo_prior(h::Histogram, ps::NamedTuple, fit_func::S
     end
 end
 
+"""
+    get_aoe_pseudo_prior(h::Histogram, ps::NamedTuple, fit_func::Symbol; pseudo_prior::NamedTupleDist=NamedTupleDist(empty = true), kwargs...)
+
+Gets the A/E of a histogram using a pseudo prior. 
+# Arguments
+    * 'h': Histogram data
+    * 'ps': Peak statistics
+    * 'fit_func': Histogram fit function
+
+# Keywords
+    * 'pseudo_prior': Initial guess for parameters of histogram
+
+# Returns
+    * 'pseudo_prior': Initial guess for parameters of histogram
+
+
+"""
+
 function get_aoe_pseudo_prior(h::Histogram, ps::NamedTuple, fit_func::Symbol; pseudo_prior::NamedTupleDist=NamedTupleDist(empty = true), kwargs...)
     standard_pseudo_prior = get_aoe_standard_pseudo_prior(h, ps, fit_func; kwargs...)
     # use standard priors in case of no overwrites given

diff --git a/src/aoefit.jl b/src/aoefit.jl
@@ -4,9 +4,15 @@
     generate_aoe_compton_bands(aoe::Vector{<:Real}, e::Vector{<:T}, compton_bands::Vector{<:T}, compton_window::T) where T<:Unitful.Energy{<:Real}
 
 Generate histograms for the A/E Compton bands and estimate peak parameters. 
-The compton bands are cutted out of the A/E spectrum and then binned using the Freedman-Diaconis Rule. For better performance
+The compton bands are cut out of the A/E spectrum and then binned using the Freedman-Diaconis Rule. For better performance
 the binning is only done in the area around the peak. The peak parameters are estimated using the `estimate_single_peak_stats_psd` function.
 
+# Arguments
+    * 'aoe': A/E
+    * 'e': Calibrated energies
+    * 'compton_bands': Compton band
+    * 'compton_window': Data window with Compton Scattering events 
+
 # Returns
     * `peakhists`: Array of histograms for each compton band
     * `peakstats`: StructArray of peak parameters for each compton band
@@ -109,13 +115,22 @@ export generate_aoe_compton_bands
 
 
 """
-    fit_aoe_compton(peakhists::Array, peakstats::StructArray, compton_bands::Array{T}) where T<:Real
+    fit_aoe_compton(peakhists::Vector{<:Histogram}, peakstats::StructArray, compton_bands::Array{T},; uncertainty::Bool=false) where T<:Unitful.Energy{<:Real}
 
 Fit the A/E Compton bands using the `f_aoe_compton` function consisting of a gaussian SSE peak and a step like background for MSE events.
 
+# Arguments
+    * 'peakhists': Data range of the histogram peaks
+    * 'peakstats': Peak statistics
+    * 'compton_bands': Array of Compton bands
+
+# Keywords
+    * 'uncertainty': Fit uncertainty
+
 # Returns
     * `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, fit_func::Symbol=:aoe_one_bck) where T<:Unitful.Energy{<:Real}    
     # create return and result dicts
@@ -151,13 +166,25 @@ export fit_aoe_compton
 
 
 """
-    fit_single_aoe_compton(h::Histogram, ps::NamedTuple{(:peak_pos, :peak_fwhm, :peak_sigma, :peak_counts, :mean_background, :μ, :σ), NTuple{7, T}}; uncertainty::Bool=true) where T<:Real
+    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)
 
 Perform a fit of the peakshape to the data in `h` using the initial values in `ps` while using the `f_aoe_compton` function consisting of a gaussian SSE peak and a step like background for MSE events.
 
+# Arguments
+    * 'h': Histogram data
+    * 'ps': peak statistics 
+
+# Keywords
+    * 'uncertainty': Fit uncertainty
+    * 'pseudo_prior': Initial guess for histogram parameters 
+    * 'fit_func': Fit function name
+    * 'background_center': Center of background fit curve
+    * 'fixed_position': 
+
 # Returns
     * `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=:aoe_one_bck, background_center::Union{Real,Nothing} = ps.peak_pos, fixed_position::Bool=false)
     # create pseudo priors