Combined fits + full chi2 fit

get_chi2 + fit_routine restructured 

src/juobs_tools.jl

@@ -249,17 +249,22 @@ 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)
+@. model2(x,p) = p[1] + p[2] * x[:, 1] + (p[3] + p[4] * x[:, 1]) * x[:, 2]
 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)
     isnothing(wpm) ? uwerr.(ydata) : uwerr.(ydata, wpm)
+    
     yval = value.(ydata)
     yer = err.(ydata)
+    
+    # Generate chi2 + solver
     chisq = gen_chisq(model, xdata, yer)
     min_fun(t) = chisq(t, yval)
     sol = optimize(min_fun, fill(0.5, param), method=LBFGS())
     par = Optim.minimizer(sol)
 
+    # Info
     (upar, chi_exp) = isnothing(wpm) ? fit_error(chisq, par, ydata) : fit_error(chisq, par, ydata, wpm)
     for i = 1:length(upar)
         isnothing(wpm) ? uwerr(upar[i]) : uwerr(upar[i], wpm)
@@ -269,15 +274,12 @@ function fit_routine(model::Function, xdata::Array{<:Real}, ydata::Array{uwreal}
     println("Chisq / chiexp: ", sol.minimum, " / ", chi_exp, " (dof: ", length(yval) - param,")")
     return upar
 end
-function gen_chisq(f::Function, x::Array{<:Real}, err::Vector{Float64})
-    chisq(par, dat) = sum((dat .- f(x,par)).^2 ./err.^2)
-    return chisq
-end
 
-#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)
     
+    Nalpha = size(xdata, 2) # number of x-variables
+    Ndata = size(xdata, 1) # number of datapoints
     if isnothing(wpm)
         uwerr.(ydata)
         uwerr.(xdata)
@@ -291,33 +293,53 @@ function fit_routine(model::Function, xdata::Array{uwreal}, ydata::Array{uwreal}
     xval = value.(xdata)
     xer = err.(xdata)
 
-    dat = vcat(xval, yval)
-    ddat = vcat(xer, yer)
-    data = vcat(xdata, ydata)
+    dat = Vector{Float64}(undef, Ndata * (Nalpha+1))
+    ddat = Vector{Float64}(undef, Ndata * (Nalpha+1))
+    data = Vector{uwreal}(undef, Ndata * (Nalpha+1)) 
 
-    #guess
+    for i = 1:Nalpha
+        dat[(i-1)*Ndata+1:i*Ndata] = xval[:, i]
+        ddat[(i-1)*Ndata+1:i*Ndata] = xer[:, i]
+        data[(i-1)*Ndata+1:i*Ndata] = xdata[:, i]
+    end
+    dat[Nalpha*Ndata+1:end] = yval
+    ddat[Nalpha*Ndata+1:end] = yer
+    data[Nalpha*Ndata+1:end] = ydata
+
+    # Guess
     chisq = gen_chisq(model, xval, yer)
     min_fun_cons(t) = chisq(t, yval)
     sol_cons = optimize(min_fun_cons, fill(0.5, param), method=LBFGS())
     par_cons = Optim.minimizer(sol_cons)
 
+    # Generate chi2 + solver
     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)
+        aux = Vector{Vector{uwreal}}(undef, Ndata)
+        for k = 1:Ndata
+            aux[k] = Vector{uwreal}(undef, Nalpha+1)
+            for i = 1:Nalpha
+                aux[k][i] = xdata[k, i]
+            end
+            aux[k][Nalpha+1] = ydata[k]
+        end
+
+        C = isnothing(wpm) ? [ADerrors.cov(aux[k]) for k = 1:Ndata] : [ADerrors.cov(aux[k], wpm) for k = 1:Ndata]
+        chisq_full_cov(p, d) = get_chi2_cov(model, d, C, p, Nalpha)
         min_fun_cov(t) = chisq_full_cov(t, dat)
-        sol = optimize(min_fun_cov, vcat(par_cons, xval), method=LBFGS())
+        sol = optimize(min_fun_cov, vcat(par_cons, dat[1:Nalpha*Ndata]), 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)
+        chisq_full(p, d) = get_chi2(model, d, ddat, p, Nalpha)
         min_fun(t) = chisq_full(t, dat)
-        sol = optimize(min_fun, vcat(par_cons, xval), method=LBFGS())
-        
+        sol = optimize(min_fun, vcat(par_cons, dat[1:Nalpha*Ndata]), 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 **
+
+    # Info
     for i = 1:length(upar)
         isnothing(wpm) ? uwerr(upar[i]) : uwerr(upar[i], wpm)
         print("\n Fit parameter: ", i, ": ")
@@ -328,113 +350,54 @@ function fit_routine(model::Function, xdata::Array{uwreal}, ydata::Array{uwreal}
     
 end
 
-function get_chi2(f::Function, x, dx, y, dy, par) #full
+function gen_chisq(f::Function, x::Array{<:Real}, err::Vector{Float64}) #constrained
+    chisq(par, dat) = sum((dat .- f(x,par)).^2 ./err.^2)
+    return chisq
+end
+
+function get_chi2(f::Function, data, ddata, par, Nalpha) #full
     chi2 = 0.0
-    Ndata = length(x)
-    Npar = length(par) - Ndata
+
+    Ndata = div(length(data), Nalpha+1)
+    Npar = length(par) - Ndata * Nalpha
     p = par[1:Npar]
 
     for k = 1:Ndata
-        xx = par[Npar + k]
-        yy = f(xx, p)
+        xx = [par[Npar + k + (i-1)*Ndata] for i = 1:Nalpha]
+        Cinv = zeros(Nalpha+1, Nalpha+1)
+        [Cinv[i, i] = 1 / ddata[k + (i-1)*Ndata]^2 for i = 1:Nalpha+1]
+
+        xx = [par[Npar + k + (i-1)*Ndata] for i = 1:Nalpha]
+        delta = [data[k + (i-1)*Ndata] - xx[i] for i = 1:Nalpha]
+        yy = f(xx', p)
+        push!(delta, data[k + Nalpha*Ndata] - yy[1])
 
-        delta = [x[k] - xx, y[k] - yy]
-        C = [[1/dx[k]^2, 0.0] [0.0, 1/dy[k]^2]]
-        chi2 += delta' * C * delta
+        chi2 += delta' * Cinv * delta
     end
     return chi2
 end
 
-function get_chi2(f::Function, data, ddata, par) #full
-    Ndata = div(length(data), 2)
-
-    x = data[1:Ndata]
-    y = data[Ndata+1:end]
-    dx = ddata[1:Ndata]
-    dy = ddata[Ndata+1:end]
-    
-    return get_chi2(f, x, dx, y, dy, par)
-end
-
-function get_chi2_cov(f::Function, x, y, C, par) #full + cov
+function get_chi2_cov(f::Function, data, C, par, Nalpha) # full + cov
     chi2 = 0.0
-    Ndata = length(x)
-    Npar = length(par) - Ndata
+
+    Ndata = div(length(data), Nalpha+1)
+    Npar = length(par) - Ndata * Nalpha
     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
+        if det(C[k]) / prod(diag(C[k])) > 1e-6
             Cinv = inv(C[k])
         else
-            C11 = C[k][1, 1]
-            C22 = C[k][2, 2]
-            Cinv = [[1/C11, 0.0] [0.0, 1/C22]]
+            Cinv = zeros(Nalpha+1, Nalpha+1)
+            [Cinv[i, i] = 1 / C[k][i, i] for i = 1:Nalpha+1]
         end
 
+        xx = [par[Npar + k + (i-1)*Ndata] for i = 1:Nalpha]
+        delta = [data[k + (i-1)*Ndata] - xx[i] for i = 1:Nalpha]
+        yy = f(xx', p)
+        push!(delta, data[k + Nalpha*Ndata] - yy[1])
+        
         chi2 += delta' * Cinv * delta 
     end
     return chi2
 end
\ No newline at end of file
-
-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"""
-find_xmin(obs::Vector{uwreal}, y0::Int64; pl::Bool=false)
-
-find_xmin(corr::Corr; pl::Bool=false)
-
-Find the platau starting point for a given correlator. 
-The starting point xmin is determined 
-|C_2|^2 * exp(-M_2 * (xmin-y0)) / (2 * M_2) < dC(xmin, y0) / 4
-where C_2 and M_2 are the matrix element and mass of the first excited state and dC is 
-the statistical error of the correlator
-"""
-function find_xmin(obs::Vector{uwreal}, y0::Int64; pl::Bool=false)
-    p0 = [0.5, 0.5, 1.0, 1.0]
-    @. model(t, p) = p[1] * exp(-p[2] * (t-y0)) + p[3] * exp(-p[4] * (t-y0))
-
-    x = Int64.(0:length(obs)-1)
-    uwerr.(obs)
-    y = value.(obs)
-    dy = err.(obs)
-    wt = 1 ./ dy.^2
-
-    fit = curve_fit(model, Float64.(x[y0+1:end]), y[y0+1:end], p0)
-    fit = curve_fit(model, Float64.(x[y0+1:end]), y[y0+1:end], wt[y0+1:end], fit.param)
-    par = fit.param
-    #println(par)
-    
-    if par[2] > par[4]
-        p1 = par[1]
-        p2 = par[2]
-    else
-        p1 = par[3]
-        p2 = par[4]
-    end
-    @. f(t) = p1^2 * exp(-p2 * (t-y0)) / (2 * p2)
-    xmin = findfirst(t-> f(t) < 0.25*dy[t + 1], x)
-
-    if pl
-        errorbar(x, y, dy, fmt="x")
-        t = Int64.(y0:length(obs))
-        plot(t, model(t, par))
-        display(gcf())    
-    end
-    return xmin
-end
-
-find_xmin(corr::Corr; pl::Bool=false) = find_xmin(corr.obs, corr.y0, pl=pl)
-=#
\ No newline at end of file