refactor + added plotting fucntion test

src/fitting.jl

@@ -76,6 +76,16 @@ function get_range(fit_range)
 end
 
 
+"""
+    get_deltaE(fit_range)
+
+Function that returns the net width of the fit range.
+"""
+function get_deltaE(fit_range)
+    return sum([arr[2] - arr[1] for arr in fit_range])
+end
+
+
 """
     norm_uniform(x::Real,p::NamedTuple,b_name::Symbol,fit_range)
 

src/likelihood.jl

@@ -43,6 +43,28 @@ function get_signal_pdf(evt_energy::Float64, Qbb::Float64, part_k::NamedTuple)
 end
 
 
+"""
+    get_mu_b(deltaE, exposure, bkg)
+
+Get the expected number of background counts in a partition
+"""
+function get_mu_b(deltaE, exposure, bkg_index)
+    return deltaE * exposure * bkg_index
+end
+
+"""
+    get_mu_s(deltaE, exposure, bkg)
+
+Get the expected number of signal counts in a partition
+"""
+function get_mu_s(exposure, eff, signal)
+    N_A = constants.N_A
+    m_76 = constants.m_76
+    sig_units = constants.sig_units
+    return log(2) * N_A * exposure * (eff) * (signal * sig_units) / m_76
+end
+
+
 """
     get_mu_s_b(p::NamedTuple,part_k::NamedTuple,idx_part_with_events::Int,settings::Dict,fit_range)
 
@@ -55,21 +77,18 @@ function get_mu_s_b(
     settings::Dict,
     fit_range,
 )
-    N_A = constants.N_A
-    m_76 = constants.m_76
-    sig_units = constants.sig_units
 
-    deltaE = sum([arr[2] - arr[1] for arr in fit_range])
+    deltaE = get_deltaE(fit_range)
     eff = get_efficiency(p, part_k, idx_part_with_events, settings)
 
     if (settings[:bkg_only] == false)
-        model_s_k = log(2) * N_A * part_k.exposure * (eff) * (p.S * sig_units) / m_76
+        model_s_k = get_mu_s(part_k.exposure, eff, p.S)
     else
         model_s_k = 0
     end
 
     b_name = part_k.bkg_name
-    model_b_k = deltaE * part_k.exposure * p[b_name]
+    model_b_k = get_mu_b(deltaE, part_k.exposure, p[b_name])
 
     return model_s_k, model_b_k
 end
@@ -129,7 +148,6 @@ function build_likelihood_per_partition(
     end
 
     ll_value += logpdf(Poisson(λ), length(events_k))
-    println("ll_value ", ll_value)
 
     for evt_energy in events_k
 

src/plotting.jl

@@ -34,45 +34,24 @@ function fit_model(
     settings::Dict,
     bkg_shape::Symbol,
     fit_range,
-    x,
+    x::Float64,
 )
     Qbb = constants.Qbb
     N_A = constants.N_A
     m_76 = constants.m_76
     sig_units = constants.sig_units
 
-    deltaE = sum([arr[2] - arr[1] for arr in fit_range])
-    eff = nothing
-
-    # logic for efficiency it can be either correlated, uncorrelated or fixed
-    if settings[:bkg_only] == true
-        eff = 0
-
-    elseif (settings[:eff_correlated] == true)
-        eff_group = part_k.eff_name
-        eff = part_k.eff_tot + p[eff_group] * part_k.eff_tot_sigma
-
-    elseif (
-        idx_part_with_events != 0 &&
-        settings[:eff_correlated] == false &&
-        settings[:eff_fixed] == false
-    )
-
-        eff = p.ε[idx_part_with_events]
-    else
-        eff = part_k.eff_tot
-    end
+    deltaE = get_deltaE(fit_range)
+    eff = get_efficiency(p, part_k, idx_part_with_events, settings)
 
     b_name = part_k.bkg_name
-    model_b_k = deltaE * part_k.exposure * p[b_name]
+    model_b_k = get_mu_b(deltaE, part_k.exposure, p[b_name])
     term1 = model_b_k * get_bkg_pdf(bkg_shape, x, p, b_name, fit_range)
 
     if (settings[:bkg_only] == false)
         reso, bias = get_energy_scale_pars(part_k, p, settings, idx_part_with_events)
-        model_s_k = log(2) * N_A * part_k.exposure * (eff) * (p.S * sig_units) / m_76
-        term2 =
-            model_s_k *
-            get_signal_pdf(part_k.signal_name, x, Qbb, bias, reso, part_k, fit_range)
+        model_s_k = get_mu_s(part_k.exposure, eff, p.S)
+        term2 = model_s_k * get_signal_pdf(x, Qbb, part_k)
     else
         term2 = 0
     end
@@ -187,6 +166,7 @@ function plot_data(
     end
 
     # exclude the gamma lines
+    ### TO DO -> GENERALIZE TO WHATEVER REMOVED INTERVAL WITHIN THE MIN-MAX OF fit_ranges 
     shape_x = [
         constants.gamma_2113_keV,
         constants.gamma_2113_keV,

src/utils.jl

@@ -340,13 +340,15 @@ function get_efficiency(
     if (settings[:eff_correlated] == true)
         eff_group = part_k.eff_name
         eff = part_k.eff_tot + p[eff_group] * part_k.eff_tot_sigma
+
         # UNCORRELATED efficiency (eff follows a pdf, different for each partition)
     elseif (
         idx_part_with_events != 0 &&
         settings[:eff_correlated] == false &&
         settings[:eff_fixed] == false
     )
         eff = p.ε[idx_part_with_events]
+
         # FIXED efficiency
     else
         eff = part_k.eff_tot
@@ -676,7 +678,7 @@ function save_outputs(
     )
 
     if config["light_output"] == false
-        @info "plot 2D posterior"
+        @info "... now we plot 2D posterior"
         plot_two_dim_posteriors(
             samples,
             free_pars,

test/plotting/test_all.jl

@@ -0,0 +1,5 @@
+using Test
+
+Test.@testset "plotting" begin
+    include("test_fit_model.jl")
+end

test/plotting/test_fit_model.jl

@@ -0,0 +1,98 @@
+using Pkg
+Pkg.activate(".") # activate the environment
+Pkg.instantiate() # instantiate the environment
+include("../../src/ZeroNuFit.jl")
+using .ZeroNuFit
+include("../../main.jl")
+include("../../src/plotting.jl")
+
+@testset "test_fit_model" begin
+
+    @info "Testing function to retrieve fit function added to the final plot evaluated at an energy of 2039.05 keV (function 'fit_model' in src/plotting.jl)"
+
+    fit_ranges =
+        Dict("l200a" => [[1930.0, 2098.511], [2108.511, 2113.513], [2123.513, 2190.0]])
+    part_event_index = [1]
+    partitions = Table(
+        experiment = Array(["L200"]),
+        fit_group = Array(["l200a"]),
+        bkg_name = Array([:B_l200a_all]),
+        signal_name = Array([:gaussian]),
+        energy_reso_name = Array([:αr_all]),
+        energy_bias_name = Array([:αb_all]),
+        eff_name = Array([:αe_all]),
+        detector = Array(["V09374A"]),
+        part_name = Array(["part0002"]),
+        start_ts = Array([1686522477]),
+        end_ts = Array([1690190843]),
+        eff_tot = Array([0.5]),
+        eff_tot_sigma = Array([0.01]),
+        width = Array([1.3]),
+        width_sigma = Array([0.01]),
+        exposure = Array([1]),
+        bias = Array([0.1]),
+        bias_sigma = Array([0.01]),
+        frac = Array([nothing]),
+        tau = Array([nothing]),
+        sigma = Array([nothing]),
+    )
+
+    settings = Dict()
+    settings[:energy_scale_fixed] = true
+    settings[:energy_scale_correlated] = false
+    settings[:eff_fixed] = true
+    settings[:eff_correlated] = false
+    settings[:bkg_only] = true
+
+    sqrt_prior = false
+    s_max = nothing
+    bkg_shape = :uniform
+    p = (S = 100, αe_all = 0.1, ω = [1.1], 𝛥 = [0.1], B_l200a_all = 2E-4)
+    x = 2039.05
+
+    total = nothing
+    try
+        total = ZeroNuFit.fit_model(
+            part_event_index[1],
+            partitions[1],
+            p,
+            settings,
+            bkg_shape,
+            fit_ranges[partitions[1].fit_group],
+            x,
+        )
+    catch e
+        @error "Error in fit_model: $e"
+        throw(e)
+    end
+
+    @testset "Check total is valid" begin
+        @test !isnothing(total)
+    end
+
+    expected_total = 0.0002
+    tolerance = 1e-3
+    @testset "Check events accuracy" begin
+        diff = abs(total - expected_total)
+        @test diff <= tolerance
+    end
+
+
+    settings[:bkg_only] = false
+    total = nothing
+    total = ZeroNuFit.fit_model(
+        part_event_index[1],
+        partitions[1],
+        p,
+        settings,
+        bkg_shape,
+        fit_ranges[partitions[1].fit_group],
+        x,
+    )
+    expected_total = 0.0845283663997091
+    tolerance = 1e-3
+    @testset "Check events accuracy" begin
+        diff = abs(total - expected_total)
+        @test diff <= tolerance
+    end
+end