From 2cf7f0ee052894cd608cfa32acdb816368c71b19 Mon Sep 17 00:00:00 2001 From: fjebaker Date: Fri, 28 Jun 2024 00:54:18 +0200 Subject: [PATCH 1/7] feat: added utils with convolution functions --- Project.toml | 3 ++- src/SpectralFitting.jl | 2 ++ src/utils.jl | 28 ++++++++++++++++++++++++++++ 3 files changed, 32 insertions(+), 1 deletion(-) create mode 100644 src/utils.jl diff --git a/Project.toml b/Project.toml index abaf65fa..cd438c56 100644 --- a/Project.toml +++ b/Project.toml @@ -11,14 +11,15 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" -FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" FITSIO = "525bcba6-941b-5504-bd06-fd0dc1a4d2eb" +FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" LibXSPEC_jll = "50917eeb-3e3c-5f3a-99bd-93e1c30e1561" Libz = "2ec943e9-cfe8-584d-b93d-64dcb6d567b7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890" LsqFit = "2fda8390-95c7-5789-9bda-21331edee243" Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" diff --git a/src/SpectralFitting.jl b/src/SpectralFitting.jl index 1618c360..e837e454 100644 --- a/src/SpectralFitting.jl +++ b/src/SpectralFitting.jl @@ -22,6 +22,7 @@ import DataInterpolations using SpecialFunctions using PreallocationTools using EnumX +using LoopVectorization import Crayons # fitting backends @@ -43,6 +44,7 @@ struct Cash <: AbstractStatistic end include("units.jl") SpectralUnits.@reexport using .SpectralUnits +include("utils.jl") include("print-utilities.jl") include("support.jl") diff --git a/src/utils.jl b/src/utils.jl new file mode 100644 index 00000000..602e15c9 --- /dev/null +++ b/src/utils.jl @@ -0,0 +1,28 @@ +function _convolve_1d_same_domain!( + output::Vector{T}, + vec_A::Vector{T}, + kernel::Vector{T}, +) where {T<:Real} + # Based on https://discourse.julialang.org/t/97658/15 + @assert length(output) == length(vec_A) + @assert length(output) == length(kernel) + + fill!(output, 0) + + @turbo for i in eachindex(output) + total = zero(T) + for k in eachindex(output) + ib0 = (i >= k) + oa = ib0 ? vec_A[i-k+1] : zero(T) + total += kernel[k] * oa + end + output[i] = total + end + output +end + +convolve!(output, A, kernel) = _convolve_1d_same_domain!(output, A, kernel) +function convolve(A, kernel) + output = similar(A) + convolve!(output, A, kernel) +end From 726450cfaa7e3ea4aadad62da4a8f2e5d62972d6 Mon Sep 17 00:00:00 2001 From: fjebaker Date: Sat, 29 Jun 2024 12:56:11 +0200 Subject: [PATCH 2/7] feat(model): delta function Added a delta function model, which is in essence an alias to a very thin Gaussian. This can be at some point updated to a more faithful implementation, provided we find a good way of propagating gradients as the line moves through the domain. Intended use is for emission lines or as the basis for convolutions. --- src/julia-models/additive.jl | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/src/julia-models/additive.jl b/src/julia-models/additive.jl index 07e53a3b..d2b182ec 100644 --- a/src/julia-models/additive.jl +++ b/src/julia-models/additive.jl @@ -130,4 +130,27 @@ end end end -export PowerLaw, BlackBody, GaussianLine +struct DeltaLine{W<:Number,T} <: AbstractSpectralModel{T,Additive} + _width::W + "Normalisation." + K::T + "Energy at which the delta function spikes." + E::T +end + +function DeltaLine(; K = FitParam(1.0), E = FitParam(5.0), width = 1e-2) + DeltaLine{typeof(width),typeof(K)}(width, K, E) +end + +Reflection.get_closure_symbols(::Type{<:DeltaLine}) = (:_width,) + +Reflection.get_parameter_symbols(model::Type{<:DeltaLine}) = fieldnames(model)[2:end] + +@inline function invoke!(flux, energy, model::DeltaLine{T}) where {T} + # we can't actually have a diract delta because that would ruin + # the ability to run dual numbers through the system. What we can do instead + # is have a miniscule gaussian + invoke!(flux, energy, GaussianLine(promote(model.K, model.E, model._width)...)) +end + +export PowerLaw, BlackBody, GaussianLine, DeltaLine From 53fe1fa75a36cd2f4cd62ed4c4c3e98a3ab4ce04 Mon Sep 17 00:00:00 2001 From: fjebaker Date: Sat, 29 Jun 2024 12:58:59 +0200 Subject: [PATCH 3/7] chore: move autodiff kwarg to keywords --- src/fitting/methods.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/fitting/methods.jl b/src/fitting/methods.jl index 0b91ca49..8277eba6 100644 --- a/src/fitting/methods.jl +++ b/src/fitting/methods.jl @@ -53,6 +53,7 @@ function fit( alg::LevenbergMarquadt; verbose = false, max_iter = 1000, + autodiff = supports_autodiff(config) ? :forward : :finite, method_kwargs..., ) @assert fit_statistic(config) == ChiSquared() "Least squares only for χ2 statistics." @@ -66,7 +67,7 @@ function fit( alg; verbose = verbose, max_iter = max_iter, - autodiff = supports_autodiff(config) ? :forward : :finite, + autodiff = autodiff, method_kwargs..., ) params = LsqFit.coef(lsq_result) From 95aab2508a766fe479f5624b0a3354e0ec3404c6 Mon Sep 17 00:00:00 2001 From: fjebaker Date: Sat, 29 Jun 2024 12:59:26 +0200 Subject: [PATCH 4/7] feat: export AsConvolution --- src/julia-models/convolutional.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/julia-models/convolutional.jl b/src/julia-models/convolutional.jl index 78c98b85..0c4acff1 100644 --- a/src/julia-models/convolutional.jl +++ b/src/julia-models/convolutional.jl @@ -65,3 +65,5 @@ function _or_else(value::Union{Nothing,T}, v::T)::T where {T} value end end + +export AsConvolution From e19a8e661627bc54afc9f1333ab4c2fa21bf755a Mon Sep 17 00:00:00 2001 From: fjebaker Date: Sat, 29 Jun 2024 13:00:20 +0200 Subject: [PATCH 5/7] test: added AsConvolution tests --- test/models/test-as-convolution.jl | 72 ++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 test/models/test-as-convolution.jl diff --git a/test/models/test-as-convolution.jl b/test/models/test-as-convolution.jl new file mode 100644 index 00000000..35613dfd --- /dev/null +++ b/test/models/test-as-convolution.jl @@ -0,0 +1,72 @@ +using SpectralFitting +using Test + +include("../dummies.jl") + +# put a couple of delta emission lines together +lines = DeltaLine(; E = FitParam(3.0), K = FitParam(2.0)) + DeltaLine(; E = FitParam(7.0)) + +# construct the convolutional wrapper +base_model = GaussianLine(; μ = FitParam(1.0), σ = FitParam(0.3)) +conv = AsConvolution(base_model) + +model = conv(lines) + +domain = collect(range(0.0, 10.0, 150)) + +plot(domain[1:end-1], invokemodel(domain, lines)) +plot(domain[1:end-1], invokemodel(domain, model)) + +output = invokemodel(domain, model) + +@test sum(output) ≈ 3.2570820013702395 atol = 1e-4 +@test output[10] ≈ 0.0036345342427057687 atol = 1e-4 +@test output[40] ≈ 0.055218163108951814 atol = 1e-4 + +# simulate a model spectrum +dummy_data = make_dummy_dataset((E) -> (E^(-3.0)); units = u"counts / (s * keV)") +sim = simulate(model, dummy_data; seed = 42) + +model.μ_1.frozen = true +model.K_1.frozen = true +model.K_2.frozen = true +model.E_1.frozen = true +model.E_2.frozen = true + +# change the width +model.σ_1.value = 0.1 +model + +begin + prob = FittingProblem(model => sim) + result = fit(prob, LevenbergMarquadt()) +end +@test result.χ2 ≈ 76.15221077389369 atol = 1e-3 + +# put a couple of delta emission lines together +lines = + DeltaLine(; E = FitParam(3.0), K = FitParam(2.0), width = 0.1) + + DeltaLine(; E = FitParam(7.0)) +model = conv(lines) + +sim = simulate(model, dummy_data; seed = 42) + +# now see if we can fit the delta line +model.μ_1.frozen = true +model.K_1.frozen = true +model.K_2.frozen = true +model.E_1.frozen = true +model.E_2.frozen = true +model.σ_1.frozen = true + +model.E_2.frozen = false +model.E_2.value = 2.0 +model.K_2.frozen = true +# model.K_2.value = 2.0 + +model +begin + prob = FittingProblem(model => sim) + result = fit(prob, LevenbergMarquadt(); verbose = true) +end +@test result.χ2 ≈ 75.736 atol = 1e-3 From dc8d7bdeafdc92f8b9f6429771162e4478b30d70 Mon Sep 17 00:00:00 2001 From: fjebaker Date: Sat, 29 Jun 2024 13:01:48 +0200 Subject: [PATCH 6/7] feat(model): AsConvolution wrapper This wrapper can turn any AbstractSpectralModel into a convolutional model, and will remove additive normalisations if required. Convolution function is a modification of the Toeplitz matrix method, which performs the matrix multiplications without allocating the matrix. The benefit this has over, e.g., the FFT(W) method of taking convolutions, in that we can propagate gradients through the operation. --- src/SpectralFitting.jl | 1 + src/meta-models/caching.jl | 18 ++++++--- src/meta-models/functions.jl | 77 ++++++++++++++++++++++++++++++++++++ src/utils.jl | 49 ++++++++++++++++------- 4 files changed, 125 insertions(+), 20 deletions(-) create mode 100644 src/meta-models/functions.jl diff --git a/src/SpectralFitting.jl b/src/SpectralFitting.jl index e837e454..70eee1e1 100644 --- a/src/SpectralFitting.jl +++ b/src/SpectralFitting.jl @@ -62,6 +62,7 @@ include("meta-models/wrappers.jl") include("meta-models/table-models.jl") include("meta-models/surrogate-models.jl") include("meta-models/caching.jl") +include("meta-models/functions.jl") include("poisson.jl") diff --git a/src/meta-models/caching.jl b/src/meta-models/caching.jl index 9af10642..d1951408 100644 --- a/src/meta-models/caching.jl +++ b/src/meta-models/caching.jl @@ -30,15 +30,21 @@ function AutoCache(model::AbstractSpectralModel{T,K}; abstol = 1e-3) where {T,K} AutoCache(model, cache, abstol) end -function _reinterpret_dual(::Type, v::AbstractArray, n::Int) +function _reinterpret_dual( + M::Type{<:AbstractSpectralModel}, + ::Type, + v::AbstractArray, + n::Int, +) needs_resize = n > length(v) if needs_resize - @warn "AutoCache: Growing dual buffer..." + @warn "$(Base.typename(M).name): Growing dual buffer..." resize!(v, n) end view(v, 1:n), needs_resize end function _reinterpret_dual( + M::Type{<:AbstractSpectralModel}, DualType::Type{<:ForwardDiff.Dual}, v::AbstractArray{T}, n::Int, @@ -46,7 +52,7 @@ function _reinterpret_dual( n_elems = div(sizeof(DualType), sizeof(T)) * n needs_resize = n_elems > length(v) if needs_resize - @warn "AutoCache: Growing dual buffer..." + @warn "$(Base.typename(M).name): Growing dual buffer..." resize!(v, n_elems) end reinterpret(DualType, view(v, 1:n_elems)), needs_resize @@ -58,8 +64,10 @@ function invoke!(output, domain, model::AutoCache{M,T,K}) where {M,T,K} _new_params = parameter_tuple(model.model) _new_limits = (first(domain), last(domain)) - output_cache, out_resized = _reinterpret_dual(D, model.cache.cache, length(output)) - param_cache, _ = _reinterpret_dual(D, model.cache.params, length(_new_params)) + output_cache, out_resized = + _reinterpret_dual(typeof(model), D, model.cache.cache, length(output)) + param_cache, _ = + _reinterpret_dual(typeof(model), D, model.cache.params, length(_new_params)) same_domain = model.cache.domain_limits == _new_limits diff --git a/src/meta-models/functions.jl b/src/meta-models/functions.jl new file mode 100644 index 00000000..1589e449 --- /dev/null +++ b/src/meta-models/functions.jl @@ -0,0 +1,77 @@ +struct AsConvolution{M,T,V,P} <: AbstractModelWrapper{M,T,Convolutional} + model::M + # the domain on which we evaluate this model + domain::V + # an additional output cache + cache::NTuple{2,Vector{P}} + function AsConvolution( + model::AbstractSpectralModel{T}, + domain::V, + cache::NTuple{2,Vector{P}}, + ) where {T,V,P} + new{typeof(model),T,V,P}(model, domain, cache) + end +end + +function AsConvolution( + model::AbstractSpectralModel{T}; + domain = collect(range(0, 2, 100)), +) where {T} + output = invokemodel(domain, model) + AsConvolution(model, domain, (output, deepcopy(output))) +end + +function invoke!(output, domain, model::AsConvolution{M,T}) where {M,T} + D = promote_type(eltype(domain), T) + model_output, _ = + _reinterpret_dual(typeof(model), D, model.cache[1], length(model.domain) - 1) + convolution_cache, _ = _reinterpret_dual( + typeof(model), + D, + model.cache[2], + length(output) + length(model_output) - 1, + ) + + # invoke the child model + invoke!(model_output, model.domain, model.model) + + # do the convolution + convolve!(convolution_cache, output, model_output) + + # overwrite the output + shift = div(length(model_output), 2) + @views output .= convolution_cache[1+shift:length(output)+shift] +end + +function Reflection.get_parameter_symbols( + ::Type{<:AsConvolution{M}}, +) where {M<:AbstractSpectralModel{T,K}} where {T,K} + syms = Reflection.get_parameter_symbols(M) + if K === Additive + # we need to lose the normalisation parameter + (syms[2:end]...,) + else + syms + end +end + +function Reflection.make_constructor( + M::Type{<:AsConvolution{Model}}, + closures::Vector, + params::Vector, + T::Type, +) where {Model<:AbstractSpectralModel{Q,K}} where {Q,K} + num_closures = fieldcount(M) - 1 # ignore the `model` field + my_closures = closures[1:num_closures] + + model_params = if K === Additive + # insert a dummy normalisation to the constructor + vcat(:(one($T)), params) + else + params + end + + model_constructor = + Reflection.make_constructor(Model, closures[num_closures+1:end], model_params, T) + :($(Base.typename(M).name)($(model_constructor), $(my_closures...))) +end diff --git a/src/utils.jl b/src/utils.jl index 602e15c9..6236d999 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,19 +1,38 @@ -function _convolve_1d_same_domain!( - output::Vector{T}, - vec_A::Vector{T}, - kernel::Vector{T}, -) where {T<:Real} +function _convolve_implementation!( + output::AbstractVector{T}, + vec_A::AbstractVector{T}, + kernel::AbstractVector{T}, +) where {T<:Number} # Based on https://discourse.julialang.org/t/97658/15 - @assert length(output) == length(vec_A) - @assert length(output) == length(kernel) + J = length(vec_A) + K = length(kernel) + @assert length(output) == J + K - 1 "Ouput is $(length(output)); should be $(J + K - 1)" - fill!(output, 0) - - @turbo for i in eachindex(output) + # do the kernel's side first + for i = 1:K-1 + total = zero(T) + for k = 1:K + ib = (i >= k) + oa = ib ? vec_A[i-k+1] : zero(T) + total += kernel[k] * oa + end + output[i] = total + end + # now the middle + for i = K:J-1 + total = zero(T) + for k = 1:K + oa = vec_A[i-k+1] + total += kernel[k] * oa + end + output[i] = total + end + # and finally the end + for i = J:(J+K-1) total = zero(T) - for k in eachindex(output) - ib0 = (i >= k) - oa = ib0 ? vec_A[i-k+1] : zero(T) + for k = 1:K + ib = (i < J + k) + oa = ib ? vec_A[i-k+1] : zero(T) total += kernel[k] * oa end output[i] = total @@ -21,8 +40,8 @@ function _convolve_1d_same_domain!( output end -convolve!(output, A, kernel) = _convolve_1d_same_domain!(output, A, kernel) +convolve!(output, A, kernel) = _convolve_implementation!(output, A, kernel) function convolve(A, kernel) - output = similar(A) + output = zeros(eltype(A), length(A) + length(kernel) - 1) convolve!(output, A, kernel) end From eab089e888747c540a92ef81e832c79ba5cd6ac0 Mon Sep 17 00:00:00 2001 From: fjebaker Date: Sat, 29 Jun 2024 13:10:18 +0200 Subject: [PATCH 7/7] chore: remove LoopVectorizations --- Project.toml | 3 +-- src/SpectralFitting.jl | 1 - 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index cd438c56..abaf65fa 100644 --- a/Project.toml +++ b/Project.toml @@ -11,15 +11,14 @@ 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" LibXSPEC_jll = "50917eeb-3e3c-5f3a-99bd-93e1c30e1561" Libz = "2ec943e9-cfe8-584d-b93d-64dcb6d567b7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890" LsqFit = "2fda8390-95c7-5789-9bda-21331edee243" Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" diff --git a/src/SpectralFitting.jl b/src/SpectralFitting.jl index 70eee1e1..b70f1d01 100644 --- a/src/SpectralFitting.jl +++ b/src/SpectralFitting.jl @@ -22,7 +22,6 @@ import DataInterpolations using SpecialFunctions using PreallocationTools using EnumX -using LoopVectorization import Crayons # fitting backends