Merge pull request #3 from legend-exp/fit_components

Add `aoecorrectionplot!` and `lplot!` for fits with components and A/E correction plots

.github/workflows/ci.yml

@@ -33,9 +33,6 @@ jobs:
         arch:
           - x64
         include:
-          - version: 1
-            os: ubuntu-latest
-            arch: x86
           - version: 1
             os: macOS-latest
             arch: arm64

Project.toml

@@ -12,12 +12,14 @@ Format = "1fa38f19-a742-5d3f-a2b9-30dd87b9d5f8"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 MathTeXEngine = "0a4f8689-d25c-4efe-a92b-7142dfc1aa53"
+Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
 RadiationDetectorSignals = "bf2c0563-65cf-5db2-a620-ceb7de82658c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [extensions]
 LegendMakieExt = ["Makie", "Dates", "FileIO", "Format", "LaTeXStrings", "MathTeXEngine", "StatsBase", "Unitful"]
+LegendMakieMeasurementsExt = ["Makie", "LaTeXStrings", "Measurements"]
 LegendMakieRadiationDetectorSignalsExt = "RadiationDetectorSignals"
 
 [compat]
@@ -28,6 +30,7 @@ LaTeXStrings = "1.2"
 Makie = "0.22"
 MakieCore = "0.9"
 MathTeXEngine = "0.6.1"
+Measurements = "2.5"
 RadiationDetectorSignals = "0.3.8"
 StatsBase = "0.33.7, 0.34"
 Unitful = "1.6"

ext/LegendMakieMeasurementsExt.jl

@@ -0,0 +1,34 @@
+# This file is a part of LegendMakie.jl, licensed under the MIT License (MIT).
+
+module LegendMakieMeasurementsExt
+
+    import LegendMakie
+
+    import LaTeXStrings
+    import Makie
+    import Measurements
+
+    import LegendMakie: pt, aoecorrectionplot!
+
+    Makie.@recipe(AoECorrectionPlot, report) do scene
+        Makie.Attributes(
+            color = (LegendMakie.AchatBlue,0.5),
+            markercolor = :black
+        )
+    end
+
+    # Needed for creatings legend using Makie recipes
+    # https://discourse.julialang.org/t/makie-defining-legend-output-for-a-makie-recipe/121567
+    function Makie.get_plots(p::AoECorrectionPlot)
+        return p.plots
+    end
+
+    function Makie.plot!(p::AoECorrectionPlot{<:Tuple{NamedTuple{(:par, :f_fit, :x, :y, :gof, :e_unit, :label_y, :label_fit)}}})
+        report = p.report[]
+        Makie.lines!(p, 0:1:3000, report.f_fit ∘ Measurements.value, color = p.color, label = LaTeXStrings.latexstring("\\fontfamily{Roboto}" * report.label_fit))
+        Makie.errorbars!(p, report.x, Measurements.value.(report.y), Measurements.uncertainty.(report.y), color = p.markercolor)
+        Makie.scatter!(p, report.x, Measurements.value.(report.y), color = p.markercolor, label = "Compton band fits: Gaussian $(report.label_y)(A/E)")
+        p
+    end
+
+end

ext/recipes/lplot.jl

@@ -56,5 +56,154 @@ function LegendMakie.lplot!(
 
     watermark && LegendMakie.add_watermarks!(; kwargs...)
 
+    fig
+end
+
+
+function LegendMakie.lplot!(
+        report::NamedTuple{(:v, :h, :f_fit, :f_components, :gof)};
+        xlabel = "", ylabel = "Counts / bin",
+        title::AbstractString = "", legend_position = :lt,
+        xlims = extrema(first(report.h.edges)), ylims = (0.9,maximum(report.h.weights) * 2),
+        show_label::Bool = true, show_components::Bool = true, yticks = Makie.automatic,
+        watermark::Bool = true, show_residuals::Bool = true, col::Int = 1, kwargs...
+    )
+
+    fig = Makie.current_figure()
+
+    # create plot
+    g = Makie.GridLayout(fig[1,col])
+    ax = Makie.Axis(g[1,1], 
+        title = title, titlefont = :bold, titlesize = 16pt, 
+        xlabel = xlabel, ylabel = ylabel, yticks = yticks,
+        limits = (xlims, ylims),
+        yscale = Makie.log10
+    )
+
+    Makie.hist!(ax, StatsBase.midpoints(first(report.h.edges)), weights = report.h.weights, bins = first(report.h.edges), color = LegendMakie.DiamondGrey, label = "Data")
+    Makie.lines!(range(xlims..., length = 1000), x -> report.f_fit(x) * step(first(report.h.edges)), color = :black, label = ifelse(show_label, "Best Fit" * (!isempty(report.gof) ? " (p = $(round(report.gof.pvalue, digits=2)))" : ""), ""))
+
+    if show_components
+        for (idx, component) in enumerate(keys(report.f_components.funcs))
+            Makie.lines!(
+                range(extrema(first(report.h.edges))..., length = 1000), 
+                x -> report.f_components.funcs[component](x) * step(first(report.h.edges)), 
+                color = report.f_components.colors[component], 
+                label = ifelse(show_label, report.f_components.labels[component], ""),
+                linestyle = report.f_components.linestyles[component],
+                linewidth = 4
+            )
+        end
+    end
+
+    if legend_position != :none 
+        Makie.axislegend(ax, position = legend_position)
+    end
+
+    if !isempty(report.gof) && show_residuals
+
+        ax.xticklabelsize = 0
+        ax.xticksize = 0
+        ax.xlabel = ""
+
+        ax2 = Makie.Axis(g[2,1],
+                xlabel = xlabel, ylabel = "Residuals (σ)",
+                yticks = -3:3:3, limits = (xlims,(-5,5))
+            )
+        LegendMakie.residualplot!(ax2, (x = StatsBase.midpoints(first(report.h.edges)), residuals_norm = report.gof.residuals_norm))
+
+        Makie.linkxaxes!(ax, ax2)
+        Makie.rowgap!(g, 0)
+        Makie.rowsize!(g, 1, Makie.Auto(3))
+
+        # align ylabels
+        yspace = maximum(Makie.tight_yticklabel_spacing!, (ax, ax2))
+        ax.yticklabelspace = yspace
+        ax2.yticklabelspace = yspace
+    end
+
+    all = Makie.Axis(g[:,:])
+    Makie.hidedecorations!(all)
+    Makie.hidespines!(all)
+    Makie.current_axis!(all)
+
+    # add watermarks
+    watermark && LegendMakie.add_watermarks!(; kwargs...)
+
+    fig
+end
+
+
+function LegendMakie.lplot!( 
+        report::NamedTuple{(:par, :f_fit, :x, :y, :gof, :e_unit, :label_y, :label_fit)}; 
+        title::AbstractString = "", show_residuals::Bool = true,
+        xticks = 500:250:2250, xlims = (500,2300), ylims = nothing,
+        legend_position = :rt, col = 1, watermark::Bool = false, kwargs...
+    )
+
+    fig = Makie.current_figure()
+
+    # create plot
+    g = Makie.GridLayout(fig[1,col])
+    ax = Makie.Axis(g[1,1], 
+        title = title, titlefont = :bold, titlesize = 16pt, 
+        xlabel = "E ($(report.e_unit))", ylabel = report.label_y * " (a.u.)", 
+        xticks = xticks, limits = (xlims, ylims)
+    )
+
+    LegendMakie.aoecorrectionplot!(ax, report)
+    if legend_position != :none 
+        Makie.axislegend(ax, position = legend_position)
+    end
+
+    # add residuals
+    if !isempty(report.gof) && show_residuals
+
+        ax.xticklabelsize = 0
+        ax.xticksize = 0
+        ax.xlabel = ""
+
+        ax2 = Makie.Axis(g[2,1],
+            xlabel = "E ($(report.e_unit))", ylabel = "Residuals (σ)",
+            xticks = xticks, yticks = -3:3:3, limits = (xlims,(-5,5))
+        )
+        LegendMakie.residualplot!(ax2, (x = report.x, residuals_norm = report.gof.residuals_norm))
+
+        # link axis and put plots together
+        Makie.linkxaxes!(ax, ax2)
+        Makie.rowgap!(g, 0)
+        Makie.rowsize!(g, 1, Makie.Auto(3))
+
+        # align ylabels
+        yspace = maximum(Makie.tight_yticklabel_spacing!, (ax, ax2))
+        ax.yticklabelspace = yspace
+        ax2.yticklabelspace = yspace
+    end
+
+    all = Makie.Axis(g[:,:])
+    Makie.hidedecorations!(all)
+    Makie.hidespines!(all)
+    Makie.current_axis!(all)
+
+    # add watermarks
+    watermark && LegendMakie.add_watermarks!(; kwargs...)
+
+    fig 
+end
+
+
+function LegendMakie.lplot!( 
+        report::NamedTuple{(:par, :f_fit, :x, :y, :gof, :e_unit, :label_y, :label_fit)},
+        com_report::NamedTuple{(:values, :label_y, :label_fit, :energy)};
+        legend_position = :rt, col = 1, kwargs...
+    )
+
+    fig = LegendMakie.lplot!(report, legend_position = :none, col = col; kwargs...)
+
+    g = last(Makie.contents(fig[1,col]))
+    ax = Makie.contents(g)[1]
+    Makie.lines!(ax, com_report.energy, com_report.values, linewidth = 4, color = :red, linestyle = :dash, label = LaTeXStrings.latexstring("\\fontfamily{Roboto}" * com_report.label_fit))
+    Makie.axislegend(ax, position = legend_position)
+
     fig
 end

src/lplot.jl

@@ -25,6 +25,7 @@ export lsavefig
 
 # recipes
 function residualplot! end
+function aoecorrectionplot! end
 
 # watermark functions
 function add_logo! end

test/Project.toml

@@ -1,10 +1,12 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 LegendSpecFits = "18221496-77af-46cf-bab8-820da09f7f97"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 RadiationDetectorSignals = "bf2c0563-65cf-5db2-a620-ceb7de82658c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
 Documenter = "1"

test/test_lplot.jl

@@ -4,6 +4,8 @@ using LegendMakie
 using Makie
 
 import LegendSpecFits
+import Distributions
+import Unitful: @u_str
 
 using Test
 
@@ -28,5 +30,34 @@ using Test
             result, report = LegendSpecFits.fit_single_trunc_gauss(randn(10000), (low = -4.0, high = 4.0, max = NaN))
             @test_nowarn lplot(report, xlabel = "x")
         end
+
+        @testset "A/E correction" begin
+            # generate fake A/E distribution
+            e_cal = rand(Distributions.Exponential(300), 5_000_000) .+ 300
+            μA, μB, σA, σB = 1.01, 4e-6, 5e-3, 12.0
+            myμ(E) = μA - μB * E
+            myσ(E) = sqrt(σA^2 + σB^2/E^2)
+            aoe = [let _μ = myμ(E), _σ = myσ(E); (rand() < 0.2 ? -rand(Distributions.Exponential(5*_σ)) : 0) + _σ*randn() + _μ; end for E in e_cal]
+
+            compton_bands = collect((550:50:2350)u"keV")
+            compton_window = 20u"keV"
+            compton_band_peakhists = LegendSpecFits.generate_aoe_compton_bands(aoe, e_cal*u"keV", compton_bands, compton_window)
+            result_fit, report_fit = LegendSpecFits.fit_aoe_compton(compton_band_peakhists.peakhists, compton_band_peakhists.peakstats, compton_bands, uncertainty=true)
+            μs = [result_fit[band].μ for band in compton_bands]
+            σs = [result_fit[band].σ for band in compton_bands]
+            result_fit_single, report_fit_single = LegendSpecFits.fit_aoe_corrections(compton_bands, μs, σs)
+            result_fit_combined, report_fit_combined = LegendSpecFits.fit_aoe_compton_combined(compton_band_peakhists.peakhists, compton_band_peakhists.peakstats, compton_bands, result_fit_single, uncertainty=true)
+
+            # Compton band (individual) fit
+            @test_nowarn lplot(report_fit[first(compton_bands)])
+
+            # A/E correction plots
+            @test_nowarn lplot(report_fit_single.report_μ, col = 1, figsize = (1200,420))
+            @test_nowarn lplot!(report_fit_single.report_σ, col = 2)
+
+            # A/E combined fit plots
+            @test_nowarn lplot(report_fit_single.report_μ, report_fit_combined.report_μ, col = 1, figsize = (1200,420))
+            @test_nowarn lplot!(report_fit_single.report_σ, report_fit_combined.report_σ, col = 2)
+        end
     end
 end