jarzynskiEstimator

je.jl

using Plots 
using Statistics

function read_xvg(filename)
    file = open(filename, "r")
    xs, ys = [], []
    for line in eachline(file)
        if !startswith(line, "#") && !startswith(line, "@")
            x, y = split(line)
            push!(xs, x)
            push!(ys, y)            
        end
    end
    close(file)
    return parse.(Float32, xs), parse.(Float32, ys)
end


function work(xs, ts, k, v)
    """
    Args: 
        xs: positions
        ts: times
        k: force constant 
        v: velocity
    Return 
        cum_dw: cummulated work at each time
    """
    dws = []
    x0 = xs[1]
    Δx = xs[2:end] - xs[1:end-1]
    for i in 2:length(xs)-1
        dw = -k * (xs[i] - x0 - v*ts[i]) * Δx[i]
        push!(dws, dw)
    end
    cum_dw = cumsum(dws)
    return cum_dw
end

function work(xs, fs)
    dws = []
    Δx = xs[2:end] - xs[1:end-1]
    fs =  0.5*(fs[2:end] + fs[1:end-1])  
    for i in 2:length(xs)-1
        dw = fs[i]* Δx[i]
        push!(dws, dw)
    end
    cum_dw = cumsum(dws)
    return cum_dw
end

function jeEstimator(xs, fs)
    cum_dw = work(xs, fs)
    ΔF = - exp.(-cum_dw)
    return ΔF
end 

function plot_work(ts, cum_works, v)
    xs = v*ts[1:100:end]
    p = plot(xs, cum_works[1:100:end])
    xlims!(p, 0, 1.0)
    ylims!(p, 0, 20)
    display(p)
end

function jeEstimator(cum_works)
    #ΔF = -β * ln(<exp(-βW)>)
end

function main()
    @show "main"
    k, v = 600.0, 0.01 #k: kJ mol^-1 nm^-2; v: nm per ps
    last = 50000
    cum_works_for_trajs = []
    ΔFs = []
    n_traj = 100
    for i in 1:n_traj
        # ts, ds = read_xvg("data/nacl/data_tmd/preprocessed/smd_$(i)_pullx.xvg")
        # ts, fs = read_xvg("data/nacl/data_tmd/preprocessed/smd_$(i)_pullf.xvg")

        ts, ds = read_xvg("/Volumes/siida-ssd/00_hla/data/processed/data_smd_k10000/smd_$(i)_pullx.xvg")
        ts, fs = read_xvg("/Volumes/siida-ssd/00_hla/data/processed/data_smd_k10000/smd_$(i)_pullf.xvg")

        ds = ds[1:last]
        fs = fs[1:last]
        ts = ts[1:last]
        
        #cum_works = work(ds, ts, k, v)
        cum_works = work(ds, fs)
        push!(cum_works_for_trajs, cum_works)
        
    end
    Ws = mean(cum_works_for_trajs, dims=1)
    #W2 = var(cum_works_for_trajs**2, dims=1) 
    ts, ds = read_xvg("/Volumes/siida-ssd/00_hla/data/processed/data_smd_k10000/smd_1_pullx.xvg")
    #ts, ds = read_xvg("data/nacl/data_tmd/preprocessed/smd_1_pullx.xvg")
    
    p = plot(ds[1:last-2], Ws)
#    p = plot(v*ts[1:last-2].+0.25, Ws)

    #ylims!(p, 0, 20)
    display(p)
end

main()