fit_routine cov added

src/juobs_tools.jl

@@ -240,14 +240,16 @@ Computes the results of a linear interpolation/extrapolation in the y axis
 y_lin_fit(par::Vector{uwreal}, x::Union{uwreal, Float64}) = par[1] + par[2] * x
 
 @doc raw"""
-fit_routine(model::Function, xdata::Array{<:Real}, ydata::Array{uwreal}, param::Int64=3)
+fit_routine(model::Function, xdata::Array{<:Real}, ydata::Array{uwreal}, param::Int64=3; wpm::Union{Dict{Int64,Vector{Float64}},Dict{String,Vector{Float64}}, Nothing}=nothing)
+
+fit_routine(model::Function, xdata::Array{uwreal}, ydata::Array{uwreal}, param::Int64=3; wpm::Union{Dict{Int64,Vector{Float64}},Dict{String,Vector{Float64}}, Nothing}=nothing)
 
 Given a model function with a number param of parameters and an array of uwreal,
 this function fit ydata with the given model and print  fit information
 The method return an array upar with the best fit parameters with their errors.
 '''@example
 @. model(x,p) = p[1] + p[2] * exp(-(p[3]-p[1])*x)
-fit_routine(model, ydata, param=3)
+fit_routine(model, xdata, ydata, param=3)
 """
 function fit_routine(model::Function, xdata::Array{<:Real}, ydata::Array{uwreal}, param::Int64=3; wpm::Union{Dict{Int64,Vector{Float64}},Dict{String,Vector{Float64}}, Nothing}=nothing)
     uwerr.(ydata)
@@ -269,8 +271,10 @@ function gen_chisq(f::Function, x::Array{<:Real}, err::Vector{Float64})
     return chisq
 end
 
-#TODO: COMBINED FITS, COVARIANCE X-Y
-function fit_routine(model::Function, xdata::Array{uwreal}, ydata::Array{uwreal}, param::Int64=3; wpm::Union{Dict{Int64,Vector{Float64}},Dict{String,Vector{Float64}}, Nothing}=nothing)
+#TODO: COMBINED FITS, UPDATE DOC
+function fit_routine(model::Function, xdata::Array{uwreal}, ydata::Array{uwreal}, param::Int64=3; 
+    wpm::Union{Dict{Int64,Vector{Float64}},Dict{String,Vector{Float64}}, Nothing}=nothing, covar::Bool=false)
+    
     uwerr.(ydata)
     yval = value.(ydata)
     yer = err.(ydata)
@@ -285,12 +289,23 @@ function fit_routine(model::Function, xdata::Array{uwreal}, ydata::Array{uwreal}
     #guess
     fit = curve_fit(model, xval, yval, 1.0 ./ yer.^2, fill(0.5, param))
 
-    chisq_full(p, d) = get_chi2(model, d, ddat, p)
-    min_fun(t) = chisq_full(t, dat)
-    sol = optimize(min_fun, vcat(fit.param, xval), method=LBFGS())
-    
-    (upar, chi2_exp) = isnothing(wpm) ? fit_error(chisq_full, Optim.minimizer(sol), data) : fit_error(chisq_full, Optim.minimizer(sol), data, wpm)
+    if covar
+        C = [ADerrors.cov([xdata[k], ydata[k]]) for k = 1:length(ydata)]
+        chisq_full_cov(p, d) = get_chi2_cov(model, d, C, p)
+        min_fun_cov(t) = chisq_full_cov(t, dat)
+        sol = optimize(min_fun_cov, vcat(fit.param, xval), method=LBFGS())
+        
+        (upar, chi2_exp) = isnothing(wpm) ? fit_error(chisq_full_cov, Optim.minimizer(sol), data) : fit_error(chisq_full_cov, Optim.minimizer(sol), data, wpm)
+    else
+        chisq_full(p, d) = get_chi2(model, d, ddat, p)
+        min_fun(t) = chisq_full(t, dat)
+        sol = optimize(min_fun, vcat(fit.param, xval), method=LBFGS())
+        
+        (upar, chi2_exp) = isnothing(wpm) ? fit_error(chisq_full, Optim.minimizer(sol), data) : fit_error(chisq_full, Optim.minimizer(sol), data, wpm)
+    end
 
+    #### chisq_full, min_fun out of conditional ->
+    #### COMPILER WARNING ** incremental compilation may be fatally broken for this module **
     for i = 1:length(upar)
         uwerr(upar[i])
         print("\n Fit parameter: ", i, ": ")
@@ -300,6 +315,7 @@ function fit_routine(model::Function, xdata::Array{uwreal}, ydata::Array{uwreal}
     return upar
     
 end
+
 function get_chi2(f::Function, x, dx, y, dy, par) #full
     chi2 = 0.0
     Ndata = length(x)
@@ -327,6 +343,40 @@ function get_chi2(f::Function, data, ddata, par) #full
     
     return get_chi2(f, x, dx, y, dy, par)
 end
+
+function get_chi2_cov(f::Function, x, y, C, par) #full + cov
+    chi2 = 0.0
+    Ndata = length(x)
+    Npar = length(par) - Ndata
+    p = par[1:Npar]
+
+    for k = 1:Ndata
+        xx = par[Npar + k]
+        yy = f(xx, p)
+
+        delta = [x[k] - xx, y[k] - yy]
+        
+        if det(C[k]) > 1e-12
+            Cinv = inv(C[k])
+        else
+            C11 = C[k][1, 1]
+            C22 = C[k][2, 2]
+            Cinv = [[1/C11, 0.0] [0.0, 1/C22]]
+        end
+        
+        chi2 += delta' * Cinv * delta
+    end
+    return chi2
+end
+
+function get_chi2_cov(f::Function, data, C, par) #full + cov
+    Ndata = div(length(data), 2)
+
+    x = data[1:Ndata]
+    y = data[Ndata+1:end]
+
+    return get_chi2_cov(f, x, y, C, par)
+end
 #=
 using LsqFit
 @doc raw"""