Automatic truncation (comp_t0) + truncate_data! added

src/juobs.jl

@@ -8,7 +8,7 @@ include("juobs_reader.jl")
 include("juobs_tools.jl")
 include("juobs_obs.jl")
 
-export read_mesons, read_ms1, read_ms, read_md
+export read_mesons, read_ms1, read_ms, read_md, truncate_data!
 export get_matrix, uwgevp_tot, energies, uwdot
 export corr_obs, md_sea, plat_av, lin_fit, x_lin_fit, y_lin_fit, fit_routine
 export meff, dec_const_pcvc, comp_t0

src/juobs_obs.jl

@@ -138,10 +138,30 @@ t0_r = comp_t0(Y, [38, 58], L=32, rw=[rw1, rw2], pl=true)
 
 ```
 """
-function comp_t0(Y::YData, plat::Vector{Int64}; L::Int64, pl::Bool=false, rw::Union{Matrix{Float64}, Nothing}=nothing, npol::Int64=2)
+function comp_t0(Y::YData, plat::Vector{Int64}; L::Int64, pl::Bool=false, 
+    rw::Union{Matrix{Float64}, Nothing}=nothing, npol::Int64=2, ws::ADerrors.wspace=ADerrors.wsg)
+
     Ysl = Y.Ysl
     t = Y.t
     id = Y.id
+    replica = size.([Ysl], 1)
+
+    #Truncation
+    n_ws = findfirst(x-> x == id, ws.map_nob)
+    if !isnothing(n_ws)
+        ivrep_ws = ws.fluc[n_ws].ivrep
+
+        if length(ivrep_ws) != 1
+            error("Different number of replicas")
+        end
+
+        if replica[1] > ivrep_ws[1]
+            println("Automatic truncation in Ysl ", ivrep_ws[1], " / ", replica[1], ". R = 1")
+            Ysl = Ysl[1:ivrep_ws[1], :, :]
+        elseif replica[1] < ivrep_ws[1]
+            error("Automatic truncation failed. R = 1\nTry using truncate_data!")
+        end
+    end
 
     Ysl = isnothing(rw) ? Ysl : apply_rw(Ysl, rw)
     xmax = size(Ysl, 2)
@@ -190,18 +210,37 @@ function comp_t0(Y::YData, plat::Vector{Int64}; L::Int64, pl::Bool=false, rw::Un
     return t0
 end
 
-function comp_t0(Y::Vector{YData}, plat::Vector{Int64}; L::Int64, pl::Bool=false, rw::Union{Vector{Matrix{Float64}}, Nothing}=nothing, npol::Int64=2)
+function comp_t0(Y::Vector{YData}, plat::Vector{Int64}; L::Int64, pl::Bool=false, 
+    rw::Union{Vector{Matrix{Float64}}, Nothing}=nothing, npol::Int64=2, ws::ADerrors.wspace=ADerrors.wsg)
+
     nr = length(Y)
     Ysl = getfield.(Y, :Ysl)
     t = getfield.(Y, :t)
     t = t[1]
     id = getfield.(Y, :id)
-    vtr = getfield.(Y, :vtr)
-    replica = length.(vtr)
+    replica = size.(Ysl, 1)
 
     if !all(id .== id[1])
         error("IDs are not equal")
     end
+    #Truncation
+    n_ws = findfirst(x-> x == id[1], ws.map_nob)
+    if !isnothing(n_ws)
+        ivrep_ws = ws.fluc[n_ws].ivrep
+
+        if length(replica) != length(ivrep_ws)
+            error("Different number of replicas")
+        end
+
+        for k = 1:length(replica)
+            if replica[k] > ivrep_ws[k]
+                println("Automatic truncation in Ysl ", ivrep_ws[k], " / ", replica[k], ". R = ", k)
+                Ysl[k] = Ysl[k][1:ivrep_ws[k], :, :]
+            elseif replica[k] < ivrep_ws[k]
+                error("Automatic truncation failed. R = ", replica[k], "\nTry using truncate_data!")
+            end
+        end
+    end
     
     Ysl = isnothing(rw) ? Ysl : apply_rw(Ysl, rw)
     tmp = Ysl[1]

src/juobs_reader.jl

@@ -329,3 +329,60 @@ function read_ms(path::String; id::Union{Int64, Nothing}=nothing, dtr::Int64=1)
     t = Float64.(0:nn) .* dn .* eps
     return YData(vntr, t, Ysl, id)
 end
+
+@doc raw"""
+truncate_data!(data::YData, nc::Int64)
+
+truncate_data!(data::Vector{YData}, nc::Vector{Int64})
+
+truncate_data!(data::Vector{CData}, nc::Int64)
+
+truncate_data!(data::Vector{Vector{CData}}, nc::Vector{Int64})
+
+Truncates the output of read_mesons and read_ms taking the first nc configurations.
+
+Examples:
+```@example    
+#Single replica
+dat = read_mesons(path, "G5", "G5")
+Y = read_ms(path)
+truncate_data!(dat, nc)
+truncate_data!(Y, nc)
+
+#Two replicas
+dat = read_mesons([path1, path2], "G5", "G5")
+Y = read_ms.([path1_ms, path2_ms])
+truncate_data!(dat, [nc1, nc2])
+truncate_data!(Y, [nc1, nc2])
+```
+"""
+function truncate_data!(data::YData, nc::Int64)
+    data.vtr = data.vtr[1:nc]
+    data.Ysl = data.Ysl[1:nc, :, :]
+    return nothing
+end
+function truncate_data!(data::Vector{YData}, nc::Vector{Int64})
+    truncate_data!.(data, nc)
+    return nothing
+end
+function truncate_data!(data::Vector{CData}, nc::Int64)
+    N = length(data)
+    for k = 1:N
+        data[k].vcfg = data[k].vcfg[1:nc]
+        data[k].re_data = data[k].re_data[1:nc, :]
+        data[k].im_data = data[k].im_data[1:nc, :]
+    end
+    return nothing
+end
+function truncate_data!(data::Vector{Vector{CData}}, nc::Vector{Int64})
+    N = length(data)
+    R = length(data[1])
+    for k = 1:N
+        for r = 1:R
+            data[k][r].vcfg = data[k][r].vcfg[1:nc[r]]
+            data[k][r].re_data = data[k][r].re_data[1:nc[r], :]
+            data[k][r].im_data = data[k][r].im_data[1:nc[r], :]
+        end
+    end
+    return nothing
+end
\ No newline at end of file