global_fit_routine using LsqFit added

plots modification in juobs_obs

src/juobs.jl

@@ -11,7 +11,7 @@ include("juobs_obs.jl")
 
 export read_mesons, read_ms1, read_ms, read_md, truncate_data!
 export get_matrix, energies, uwdot, uweigvals, uweigvecs, uweigen, invert, getall_eigvals, getall_eigvecs, hess_reduce, uwcholesky, transpose, tridiag_reduction
-export corr_obs, corr_sym, md_sea, md_val, plat_av, lin_fit, x_lin_fit, y_lin_fit, fit_routine
+export corr_obs, corr_sym, md_sea, md_val, plat_av, lin_fit, x_lin_fit, y_lin_fit, fit_routine, global_fit_routine
 export meff, mpcac, dec_const, dec_const_pcvc, comp_t0
 
 end # module

src/juobs_obs.jl

@@ -331,6 +331,9 @@ function dec_const(vv::Vector{uwreal}, plat::Vector{Int64}, m::uwreal, y0::Int64
     R_av = plat_av(R, plat, wpm)
     f = sqrt(2 / m)  * R_av
     if pl
+        R .*= sqrt.(2 ./ [m for i in 1:length(R)])
+        uwerr.(R)
+        R_av *= sqrt(2/m)
         isnothing(wpm) ? uwerr(R_av) : uwerr(R_av, wpm)
         isnothing(wpm) ? uwerr(f) : uwerr(f, wpm)
         x = 1:length(R)
@@ -344,24 +347,29 @@ function dec_const(vv::Vector{uwreal}, plat::Vector{Int64}, m::uwreal, y0::Int64
         fill_between(plat[1]:plat[2], v-e, v+e, color="green", alpha=0.75, label=L"$R$")
         errorbar(x, y, dy, fmt="x", color="black", label=lbl)
         legend()
-        ylabel(L"$R_\mathrm{av}$")
+        xlim(y0, y0 +plat[2] +15)
+        ylim(v-10*e, v+10*e)
+        ylabel(L"$af_{B}$")
+        #ylabel(L"$R_\mathrm{av}$")
         xlabel(L"$x_0$")
-    
+        #title(L"$f_{B^*}$")
         if !isnothing(kappa)
-            title(string(L"$\kappa_1 = $", kappa[1], L" $\kappa_2 = $", kappa[2]))
+            title(string( L"$\kappa_1 = $", kappa[1], L" $\kappa_2 = $", kappa[2]))
         end
         display(gcf())
+        #t = "ps_decay_$(kappa[1])_$(kappa[2]).pdf"
+        #savefig(joinpath("/Users/ale/Il mio Drive/phd/secondment/3pf test/analysis/plots",t))
     end
     if !data
         return f
     else
-        return (f, R)
+        return f, R
     end
 end
 function dec_const(vv::Corr, plat::Vector{Int64}, m::uwreal; pl::Bool=true, data::Bool=false,
     wpm::Union{Dict{Int64,Vector{Float64}},Dict{String,Vector{Float64}}, Nothing}=nothing) 
     
-    return dec_const(vv.obs, plat, m, vv.y0, kappa=vv.kappa, pl=pl, data=data, wpm=wpm)
+    return dec_const(vv.obs, plat, m, vv.y0+1, kappa=vv.kappa, pl=pl, data=data, wpm=wpm)
 end
 
 @doc raw"""

src/juobs_tools.jl

@@ -387,7 +387,6 @@ function fit_routine(model::Function, xdata::Array{<:Real}, ydata::Array{uwreal}
     
     yval = value.(ydata)
     yer = err.(ydata)
-    
     # Generate chi2 + solver
     chisq = gen_chisq(model, xdata, yer)
     fit = curve_fit(model, xdata, yval, 1.0 ./ yer.^2, fill(0.5, param))
@@ -482,6 +481,74 @@ function fit_routine(model::Function, xdata::Array{uwreal}, ydata::Array{uwreal}
     
 end
 
+function global_fit_routine(models::Vector{Function}, xdata::Vector{Array{Float64, 1}}, ydata::Vector{Array{uwreal, 1}}, n_par::Int64; guess_param::Vector{Float64}=fill(0.5, n_par), log::Bool=true, wpm::Union{Dict{Int64,Vector{Float64}},Dict{String,Vector{Float64}}, Nothing}=nothing)
+    if !(length(models) == length(ydata))
+        error("Dimension mismatch")
+    end
+    N_sets = length(models)             # number of datasets
+    x_flat = xdata[1]                   # unrolling xdata 
+    y_flat = ydata[1]                   # unrolling ydata
+
+    chunk_idx = Vector{Vector{Int64}}(undef, N_sets)
+    j = 1
+    for i = 1:N_sets
+        isnothing(wpm) ? uwerr.(ydata[i]) :  [uwerr(yaux, wpm) for yaux in ydata[i]]
+        if i > 1 
+            y_flat = vcat(y_flat, ydata[i])
+            x_flat =vcat(x_flat, xdata[i])
+        end
+        stp = j + length(ydata[i]) - 1
+        chunk_idx[i] = collect(j:stp)
+        j = stp + 1
+    end
+    function flat_model(x_flat::Vector{Float64}, param::Vector{Float64})
+        fx = similar(x_flat)
+        for ii in 1:N_sets
+            mod = models[ii]
+            #par = param[idx[ii]]             
+            fx[chunk_idx[ii]] = mod(x_flat[chunk_idx[ii]], param)
+        end
+        return fx
+    end
+    function flat_chisq(param, data)
+        chi = 0.0
+        for ii in 1:N_sets
+            mod = models[ii]
+            #par = param[idx[ii]]
+            # uncorrelated chi2
+            chi += sum((data[chunk_idx[ii]] .- mod(x_flat[chunk_idx[ii]], param)).^2 ./ err.(y_flat[chunk_idx[ii]]).^2 )        
+            # correlated chi2
+            #cov_mat = cov(data[chunk_idx[ii]])
+            #chi += sum(
+            #    juobs.transpose(data[chunk_idx[ii]] .- mod(x_flat[chunk_idx[ii]], param)) * cov_mat^(-1) * (data[chunk_idx[ii]] .- mod(x_flat[chunk_idx[ii]], param))
+            #)
+        end
+        return chi
+    end
+    #tmp_par = [0.1, 0.6, 1, 1, 0.5, 0.75, 1, 1, 1, 1, 1, 1]
+    #tmp_par = [1, 0.1, 1, 0.6]
+    #tmp_par = [1, 0.5, 1, 0.75]
+    #tmp = flat_chisq(tmp_par, y_flat)
+    #println("chi2 with correct values = ", tmp)
+    fit = curve_fit(flat_model, x_flat, value.(y_flat), 1.0 ./ err.(y_flat).^2, guess_param )
+    #W = cov(y_flat)^(-1)
+    #println(typeof(W))
+    #fit = curve_fit(flat_model, x_flat, value.(y_flat), W, guess_param )
+    #uncorrelated chi2exp
+    (upar, chi2_exp) = isnothing(wpm) ? fit_error(flat_chisq, coef(fit) , y_flat) : fit_error(flat_chisq, coef(fit) , y_flat, wpm)
+    #correlated chi2exp
+    #(upar, chi2_exp) = isnothing(wpm) ? fit_error(flat_chisq, coef(fit) , y_flat, W=cov(y_flat)^(-1)) : fit_error(flat_chisq, coef(fit) , y_flat, wpm)
+    chisq = value(flat_chisq(coef(fit), y_flat))  
+    if log
+        println("Converged param: ", fit.converged)
+        #println("wt param: ", fit.wt)
+        println("Chisq / chiexp: ", chisq, " / ", chi2_exp, " (dof: ", dof(fit),")")
+        chis2_corrected = (dof(fit)) * chisq / chi2_exp
+        println("Chisq corrected: ", chis2_corrected)
+    end
+    return upar, chisq, chi2_exp, dof(fit)
+end
+
 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