Merge pull request #110 from fjebaker/fergus/better-plots

Better plots and background

Project.toml

@@ -11,8 +11,8 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
-FITSIO = "525bcba6-941b-5504-bd06-fd0dc1a4d2eb"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+FITSIO = "525bcba6-941b-5504-bd06-fd0dc1a4d2eb"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"

docs/src/examples/sherpa-example.md

@@ -69,7 +69,7 @@ The result card tells us a little bit about how successful the fit was. We furth
 
 ```@example sherpa
 plot(data, markersize = 3)
-plot!(data, result)
+plot!(result)
 ```
 
 We can create a contour plot of the fit statistic by evaluating the result everywhere on the grid and measuring the statistic:

docs/src/walkthrough.md

@@ -99,7 +99,7 @@ plot(data,
     xscale = :log10, 
     yscale = :log10
 )
-plot!(data, result)
+plot!(result)
 ```
 
 Our model does not account for the high energy range well. We can ignore that range for now, and select everything from 0 to 15 keV and refit:
@@ -115,7 +115,7 @@ plot(data,
     xscale = :log10, 
     yscale = :log10
 )
-plot!(data, result, label = "PowerLaw")
+plot!(result, label = "PowerLaw")
 ```
 
 The result is not yet baked into our model, and represents just the outcome of the fit. To update the parameters and errors in the model, we can use [`update_model!`](@ref)
@@ -179,7 +179,7 @@ plot(data,
     xscale = :log10, 
     yscale = :log10
 )
-plot!(data, flux_result)
+plot!(flux_result)
 vspan!([flux_model.c1.E_min.value, flux_model.c1.E_max.value], alpha = 0.5)
 ```
 
@@ -205,8 +205,8 @@ dp = plot(data,
     yscale = :log10,
     legend = :bottomleft,
 )
-plot!(dp, data, result, label = "PowerLaw $(round(result.χ2))")
-plot!(dp, data, result2, label = "BlackBody $(round(result2.χ2))")
+plot!(dp, result, label = "PowerLaw $(round(result.χ2))")
+plot!(dp, result2, label = "BlackBody $(round(result2.χ2))")
 ```
 
 Or a bremsstrahlung model:
@@ -218,13 +218,13 @@ result3 = fit(prob3, LevenbergMarquadt())
 ```
 
 ```@example walk
-plot!(dp, data, result3, label = "Brems $(round(result3.χ2))")
+plot!(dp, result3, label = "Brems $(round(result3.χ2))")
 ```
 
 Let's take a look at the residuals of these three models. There are utility methods for this in SpectralFitting.jl, but we can easily just interact with the result directly:
 
 ```@example walk
-function residuals(result)
+function calc_residuals(result)
     # select which result we want (only have one, but for generalisation to multi-model fits)
     r = result[1] 
     y = invoke_result(r)
@@ -234,9 +234,9 @@ end
 domain = SpectralFitting.plotting_domain(data)
 
 rp = hline([0], linestyle = :dash, legend = false)
-plot!(rp,domain, residuals(result), seriestype = :stepmid)
-plot!(rp, domain, residuals(result2), seriestype = :stepmid)
-plot!(rp, domain, residuals(result3), seriestype = :stepmid)
+plot!(rp,domain, calc_residuals(result), seriestype = :stepmid)
+plot!(rp, domain, calc_residuals(result2), seriestype = :stepmid)
+plot!(rp, domain, calc_residuals(result3), seriestype = :stepmid)
 rp
 ```
 
@@ -246,6 +246,19 @@ We can compose this figure with our previous one, and change to a linear x scale
 plot(dp, rp, layout = grid(2, 1, heights = [0.7, 0.3]), link = :x, xscale = :linear)
 ```
 
+We can do all that plotting work in one go with the [`plotresult`](@ref) recipe:
+
+```@example walk
+plotresult(
+    data,
+    [result, result2, result3],
+    ylims = (0.001, 2.0), 
+    xscale = :log10, 
+    yscale = :log10,
+    legend = :bottomleft,
+)
+```
+
 Let's modify the black body model with a continuum component
 
 ```@example walk
@@ -281,7 +294,7 @@ plot(data,
     yscale = :log10,
     legend = :bottomleft,
 )
-plot!(data, bbpl_result)
+plot!(bbpl_result)
 ```
 
 Update the model and fix the black body temperature to 2 keV:
@@ -306,7 +319,7 @@ bbpl_result2 = fit(
 Overplotting this new result:
 
 ```@example walk
-plot!(data, bbpl_result2)
+plot!(bbpl_result2)
 ```
 
 ## MCMC

src/SpectralFitting.jl

@@ -81,8 +81,6 @@ include("fitting/methods.jl")
 include("simulate.jl")
 include("fitting/goodness.jl")
 
-include("plotting-recipes.jl")
-
 # include xspec models
 include("xspec-models/additive.jl")
 include("xspec-models/multiplicative.jl")
@@ -94,11 +92,17 @@ include("julia-models/additive.jl")
 include("julia-models/multiplicative.jl")
 include("julia-models/convolutional.jl")
 
+include("plots-recipes.jl")
+
 function __init__()
     # check if we have the minimum model data already
     _check_model_directory_present()
     # init HEASOFT
-    ccall((:FNINIT, libXSFunctions), Cvoid, ())
+    if get(ENV, "SPECTRAL_FITTING_XSPEC_INIT", "") == ""
+        ccall((:FNINIT, libXSFunctions), Cvoid, ())
+        # set an environment variable so we don't accidentally init again
+        ENV["SPECTRAL_FITTING_XSPEC_INIT"] = "true"
+    end
 end
 
 end # module

src/datasets/spectraldata.jl

@@ -142,6 +142,19 @@ function SpectralData(
     return data
 end
 
+function Base.copy(data::SpectralData)
+    SpectralData(
+        data.spectrum,
+        data.response,
+        data.background,
+        data.ancillary,
+        copy(data.energy_low),
+        copy(data.energy_high),
+        copy(data.domain),
+        copy(data.data_mask),
+    )
+end
+
 supports(::Type{<:SpectralData}) = (ContiguouslyBinned(),)
 
 function _objective_to_units(dataset::SpectralData, obj, units)
@@ -346,7 +359,7 @@ function _adjust_by_unit_difference!(
 end
 
 function adjust_to_units!(data::SpectralData, s::Spectrum, x, units)
-    ΔE = bin_widths(data)
+    ΔE = unmasked_bin_widths(data)
     exposure_time = s.exposure_time
     _adjust_by_unit_difference!(ΔE, exposure_time, x, units / s.units)
     x
@@ -454,6 +467,47 @@ function match_domains!(data::SpectralData)
     )
 end
 
+function background_dataset(data::SpectralData)
+    new_data = copy(data)
+    new_data.spectrum = new_data.background
+    new_data.background = nothing
+    new_data
+end
+
+function rescale!(data::SpectralData)
+    @. data.spectrum.data = data.spectrum.data / data.spectrum.area_scale
+    @. data.spectrum.errors = data.spectrum.errors / data.spectrum.area_scale
+    data.spectrum.area_scale = 1
+    if has_background(data)
+        rescale_background!(data)
+    end
+    data
+end
+
+function rescale_background!(data::SpectralData)
+    if has_background(data)
+        data.background.data = _scaled_background(
+            data.background.data,
+            data.background.area_scale,
+            data.spectrum.background_scale,
+            data.background.background_scale,
+        )
+        data.background.errors = _scaled_background(
+            data.background.errors,
+            data.background.area_scale,
+            data.spectrum.background_scale,
+            data.background.background_scale,
+        )
+
+        data.spectrum.background_scale = 1
+        data.background.background_scale = 1
+        data.background.area_scale = 1
+    else
+        throw("No background to subtract")
+    end
+    data
+end
+
 macro _forward_SpectralData_api(args)
     if args.head !== :.
         error("Bad syntax")
@@ -511,6 +565,16 @@ macro _forward_SpectralData_api(args)
             SpectralFitting.error_statistic(getfield(t, $(field)))
         SpectralFitting.set_units!(t::$(T), args...) =
             SpectralFitting.set_units!(getfield(t, $(field)), args...)
+        SpectralFitting.background_dataset(t::$(T), args...; kwargs...) =
+            SpectralFitting.background_dataset(getfield(t, $(field)), args...; kwargs...)
+        SpectralFitting.rescale_background!(t::$(T), args...; kwargs...) =
+            SpectralFitting.rescale_background!(
+                getfield(t, $(field)),
+                args...;
+                kwargs...,
+            )
+        SpectralFitting.rescale!(t::$(T), args...; kwargs...) =
+            SpectralFitting.rescale!(getfield(t, $(field)), args...; kwargs...)
     end |> esc
 end
 
@@ -609,4 +673,7 @@ export SpectralData,
     normalize!,
     subtract_background!,
     set_domain!,
-    set_units!
+    set_units!,
+    background_dataset,
+    rescale!,
+    rescale_background!

src/datasets/spectrum.jl

@@ -150,7 +150,7 @@ end
 error_statistic(spec::Spectrum) = spec.error_statistics
 
 function subtract_background!(spectrum::Spectrum, background::Spectrum)
-    # should all already be rates
+    @assert spectrum.units == u"counts"
     # errors added in quadrature
     # TODO: this needs fixing to propagate errors properly
     data_variance = spectrum.errors .^ 2
@@ -163,6 +163,8 @@ function subtract_background!(spectrum::Spectrum, background::Spectrum)
         background.area_scale,
         spectrum.background_scale,
         background.background_scale,
+        spectrum.exposure_time,
+        background.exposure_time,
     )
     @. spectrum.errors = √abs(spectrum.errors)
     _subtract_background!(
@@ -173,12 +175,21 @@ function subtract_background!(spectrum::Spectrum, background::Spectrum)
         background.area_scale,
         spectrum.background_scale,
         background.background_scale,
+        spectrum.exposure_time,
+        background.exposure_time,
     )
     spectrum
 end
 
-_subtract_background!(output, spec, back, aD, aB, bD, bB) =
-    @. output = (spec / aD) - (bD / bB) * (back / aB)
+"""
+Does the background subtraction and returns units of counts. That means we have
+multiplied through by a factor ``t_D`` relative to the reference equation (2.3)
+in the XSPEC manual.
+"""
+_subtract_background!(output, spec, back, aD, aB, bD, bB, tD, tB) =
+    @. output = (spec / (aD)) - (tD / tB) * _scaled_background(back, aB, bD, bB)
+
+_scaled_background(back, aB, bD, bB) = (bD / bB) * (back / aB)
 
 
 export Spectrum

src/fitting/result.jl

@@ -35,13 +35,19 @@ end
 get_cache(f::FittingResultSlice) = f.parent.config.cache
 get_model(f::FittingResultSlice) = f.parent.config.prob.model.m[f.index]
 get_dataset(f::FittingResultSlice) = f.parent.config.prob.data.d[f.index]
+fit_statistic(f::FittingResultSlice) = fit_statistic(f.parent.config)
 
 estimated_error(r::FittingResultSlice) = r.σu
 estimated_params(r::FittingResultSlice) = r.u
 
-function invoke_result(slice::FittingResultSlice, u)
+function invoke_result(slice::FittingResultSlice{P}, u) where {P}
     @assert length(u) == length(slice.u)
-    _invoke_and_transform!(get_cache(slice), slice.domain, u)
+    cache = if P <: MultiFittingResult
+        get_cache(slice).caches[slice.index]
+    else
+        get_cache(slice)
+    end
+    _invoke_and_transform!(cache, slice.domain, u)
 end
 
 function _pretty_print(slice::FittingResultSlice)
@@ -77,7 +83,7 @@ function Base.getindex(result::FittingResult, i)
             result.config.objective[:],
             result.config.variance[:],
             result.u[:],
-            result.σu[:],
+            isnothing(result.σu) ? nothing : result.σu[:],
             result.χ2,
         )
     else
@@ -208,3 +214,18 @@ function finalize(
         config,
     )
 end
+
+function determine_layout(result::FittingResultSlice)
+    dataset = get_dataset(result)
+    with_units(
+        common_support(get_model(result), dataset),
+        preferred_units(dataset, fit_statistic(result)),
+    )
+end
+
+function residuals(result::FittingResultSlice)
+    y = invoke_result(result, result.u)
+    y_residual = @. (result.objective - y) / sqrt(result.variance)
+    y_residual
+end
+residuals(result::FittingResult; kwargs...) = residuals(result[1]; kwargs...)