Merge branch 'lq_updated' of https://github.com/DaGeibl/LegendSpecFit…

…s.jl into lq_updated
legend-exp · Oct 29, 2024 · 1b50ccf · 1b50ccf
2 parents f2fcd48 + 829d0f7
commit 1b50ccf
Showing 1 changed file with 29 additions and 24 deletions.
diff --git a/src/aoe_cut.jl b/src/aoe_cut.jl
@@ -160,7 +160,10 @@ Get the surrival fraction of a peak after a AoE cut value `aoe_cut` for a given
 - `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, inverted_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}
+
+function get_peak_surrival_fraction(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, peak::T, window::Vector{T}, aoe_cut::Unitful.RealOrRealQuantity,; 
+    uncertainty::Bool=true, inverted_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}
+
     # estimate bin width
     bin_width = get_friedman_diaconis_bin_width(e[e .> peak - bin_width_window .&& e .< peak + bin_width_window])
     # get energy before cut and create histogram
@@ -170,21 +173,15 @@ function get_peak_surrival_fraction(aoe::Vector{<:Unitful.RealOrRealQuantity}, e
     # fit peak and return number of signal counts
     result_before, report_before = fit_single_peak_th228(peakhist, peakstats,; uncertainty=uncertainty, fit_func=fit_func)
 
-    # get e after cut
-    if !isnan(sigma_high_sided)
-        # TODO: decide how to deal with the high sided cut!
-        e = e[aoe .< sigma_high_sided]
-        aoe = aoe[aoe .< sigma_high_sided]
-    end
 
-    if inverted_mode == false
-        #normal aoe mode
-        e_survived = e[aoe_cut .<= aoe]
-        e_cut = e[aoe_cut .> aoe]
+    # get energy after cuts
+    e_survived, e_cut = if !inverted_mode
+        #normal aoe version
+        e[aoe_cut .< aoe .< sigma_high_sided], e[aoe .<= aoe_cut .|| aoe .>= sigma_high_sided]
     else
-        #inverted mode for lq cut
-        e_survived = e[aoe_cut .>= aoe]
-        e_cut = e[aoe_cut .< aoe]
+        #lq version
+        e[aoe .< aoe_cut .|| aoe .> sigma_high_sided], e[aoe_cut .<= aoe .<= sigma_high_sided]
+
     end
 
     # estimate bin width
@@ -229,18 +226,26 @@ Get the surrival fraction of a continuum after a AoE cut value `aoe_cut` for a g
 - `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,; inverted_mode::Bool=false, sigma_high_sided::Unitful.RealOrRealQuantity=NaN) where T<:Unitful.Energy{<:Real}
+
+function get_continuum_surrival_fraction(aoe::Vector{<:Unitful.RealOrRealQuantity}, e::Vector{<:T}, center::T, window::T, aoe_cut::Unitful.RealOrRealQuantity,; inverted_mode::Bool=false, sigma_high_sided::Unitful.RealOrRealQuantity=Inf*unit(first(aoe))) where T<:Unitful.Energy{<:Real}
+    # scale unit
+    e_unit = u"keV"
+    # get energy around center
+    aoe = aoe[center - window .< e .< center + window]
+    e = e[center - window .< e .< center + window]
+    # get bin width
+    bin_width = get_friedman_diaconis_bin_width(e)
     # get number of events in window before cut
-    n_before = length(e[center - window .< e .< center + window])
-    # get number of events after cut
-    n_after = length(e[aoe .> aoe_cut .&& center - window .< e .< center + window])
-    if !isnan(sigma_high_sided)
-        n_after = length(e[aoe_cut .< aoe .< sigma_high_sided .&& center - window .< e .< center + window])
-    end
-    #inverted mode for lq cut
-    if inverted_mode == true
-        n_after = length(e[aoe .< aoe_cut .&& center - window .< e .< center + window])
+    n_before = length(e)
+    # get energy after cuts
+    e_survived, e_cut = if !inverted_mode
+        #aoe version
+        e[aoe_cut .< aoe .< sigma_high_sided], e[aoe .<= aoe_cut .|| aoe .>= sigma_high_sided]
+    else
+        #lq version
+        e[aoe .< aoe_cut .|| aoe .> sigma_high_sided], e[aoe_cut .<= aoe .<= sigma_high_sided]
     end
+    n_after = length(e_survived)
 
     # calculate surrival fraction
     sf = n_after / n_before