Skip to content

Commit

Permalink
Merge pull request #74 from DaGeibl/lq_updated
Browse files Browse the repository at this point in the history
Lq updated
  • Loading branch information
fhagemann authored Nov 20, 2024
2 parents 0d4c7c7 + d45a0c5 commit 35684a2
Show file tree
Hide file tree
Showing 6 changed files with 410 additions and 106 deletions.
183 changes: 183 additions & 0 deletions ext/LegendSpecFitsRecipesBaseExt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1270,4 +1270,187 @@ end
end
end


### lq recipe functions

# recipe for the lq_drift_time_correction report

@recipe function f(report::NamedTuple{(:lq_prehist, :lq_report, :drift_prehist, :drift_report, :lq_box, :drift_time_func, :DEP_left, :DEP_right)}, e_cal, dt_eff, lq_e_corr, plot_type::Symbol)

# Extract data from the report
DEP_left = report.DEP_left
DEP_right = report.DEP_right
box = report.lq_box

#cut data to DEP
dt_DEP = dt_eff[DEP_left .< e_cal .< DEP_right]
lq_DEP = lq_e_corr[DEP_left .< e_cal .< DEP_right]

# Plot configuration: 2D histogram
xlabel := "Drift Time"
ylabel := "LQ (A.U.)"
framestyle := :box
left_margin := -2Plots.mm
bottom_margin := -4Plots.mm
top_margin := -3Plots.mm
color := :viridis
formatter := :plain
thickness_scaling := 1.6
size := (1200, 900)


if plot_type == :DEP
# Create 2D histogram with filtered data based on DEP_left and DEP_right

# dynamic bin size dependant on fit constraint box
t_diff = box.t_upper - box.t_lower
lq_diff = box.lq_upper - box.lq_lower
xmin = box.t_lower - 1*t_diff
xmax = box.t_upper + 1*t_diff
xstep = (xmax - xmin) / 100
ymin = box.lq_lower - 1*lq_diff
ymax = box.lq_upper + 4*lq_diff
ystep = (ymax - ymin) / 100
nbins := (xmin:xstep:xmax, ymin:ystep:ymax)

@series begin
seriestype := :histogram2d
dt_DEP, lq_DEP
end
elseif plot_type == :whole
# Create 2D histogram with all data
colorbar_scale := :log10

# dynamic bin size
xmin = ustrip.(quantile(dt_eff, 0.001))
xmax = ustrip.(quantile(dt_eff, 0.999))
xstep = (xmax - xmin) / 200
ymin = ustrip.(quantile(lq_e_corr[.!isnan.(lq_e_corr)], 0.005)) #filters NaNs for quantile
ymax = ustrip.(quantile(lq_e_corr[.!isnan.(lq_e_corr)], 0.94)) #filters NaNs for quantile
ystep = (ymax - ymin) / 200
nbins := (xmin:xstep:xmax, ymin:ystep:ymax)

@series begin
seriestype := :histogram2d
dt_eff, lq_e_corr
end
end

# Add vertical and horizontal lines for the fit box limits
@series begin
seriestype := :vline
label := ""
linewidth := 1.5
color := :red
[box.t_lower, box.t_upper]
end

@series begin
seriestype := :hline
label := ""
linewidth := 1.5
color := :red
[box.lq_lower, box.lq_upper]
end

# Add linear fit plot
@series begin
label := "Linear Fit"
linewidth := 1.5
color := :blue

# Evaluate drift_time_func over the full range of drift time (x-axis)
dt_range = range(xmin, xmax, length=100) # 100 points across the x-axis
lq_fit = report.drift_time_func.(dt_range) # Apply the linear function to the full dt range

dt_range, lq_fit

end
end


# recipe for the lq_cut report

@recipe function f(report::NamedTuple{(:cut, :fit_result, :temp_hists, :fit_report)}, lq_class::Vector{Float64}, e_cal, plot_type::Symbol)

# Extract cutvalue from the report
cut_value = Measurements.value.(report.cut)

# Plot configuration for all types
left_margin := -2Plots.mm
bottom_margin := -4Plots.mm
top_margin := -3Plots.mm
thickness_scaling := 1.6
size := (1200, 900)
framestyle := :box
formatter := :plain

# Plot configuration for each specific type
if plot_type == :lq_cut
# 2D histogram for LQ Cut
xlabel := "Energy"
ylabel := "LQ (A.U.)"
nbins := (0:6:3000, -2.0:0.012:4.0)
colorbar_scale := :log10
color := :viridis
legend := :bottomright

@series begin
seriestype := :histogram2d
e_cal, lq_class
end

@series begin
seriestype := :hline
label := "3σ exclusion"
linewidth := 2
color := :red
[cut_value]
end

elseif plot_type == :energy_hist
# Energy histogram before/after LQ cut
xlabel := "Energy"
ylabel := "Counts"
nbins := 0:1:3000
yscale := :log10

@series begin
seriestype := :stephist
label := "Data before LQ Cut"
e_cal
end

@series begin
seriestype := :stephist
label := "Surviving LQ Cut"
e_cal[lq_class .< cut_value]
end

@series begin
seriestype := :stephist
label := "Cut by LQ Cut"
e_cal[lq_class .> cut_value]
end

elseif plot_type == :cut_fraction
# Percentual cut plot
xlabel := "Energy"
ylabel := "Fraction"
legend := :topleft

# Calculate fraction of events cut by LQ cut
h1 = fit(Histogram, ustrip.(e_cal), 0:3:3000)
h2 = fit(Histogram, ustrip.(e_cal[lq_class .> cut_value]), 0:3:3000)
h_diff = h2.weights ./ h1.weights

@series begin
label := "Cut Fraction"
0:3:3000, h_diff
end
end
end



end # module LegendSpecFitsRecipesBaseExt
54 changes: 27 additions & 27 deletions src/aoe_cut.jl
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ get_peaks_surrival_fractions(aoe, e, peaks, peak_names, left_window_sizes::Vecto


"""
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, 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}
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.
Expand All @@ -134,8 +134,7 @@ 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, 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}
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}
# 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
Expand All @@ -145,13 +144,21 @@ 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 energy after cuts
e_survived, e_cut = if !lq_mode
#normal aoe version
e[aoe_cut .< aoe .< sigma_high_sided], e[aoe .<= aoe_cut .|| aoe .>= sigma_high_sided]
# 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]
else
#lq version
e[aoe .< aoe_cut .|| aoe .> sigma_high_sided], e[aoe_cut .<= aoe .<= sigma_high_sided]
#inverted mode for lq cut
e_survived = e[aoe_cut .>= aoe]
e_cut = e[aoe_cut .< aoe]
end

# estimate bin width
Expand Down Expand Up @@ -185,7 +192,7 @@ export get_peak_surrival_fraction


"""
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,; inverted_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.
Expand All @@ -196,25 +203,18 @@ 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,; lq_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)
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}
# get number of events in window before cut
n_before = length(e)
# get energy after cuts
e_survived, e_cut = if !lq_mode
#normal 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]
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])
end
n_after = length(e_survived)

# calculate surrival fraction
sf = n_after / n_before
Expand Down
Loading

0 comments on commit 35684a2

Please sign in to comment.