FGCS/eval_runtimes.jl

using DataFrames, RCall, Statistics, Plots, Base.Threads, CSV, CategoricalArrays

@rlibrary qs2

prefix = "barite_fgcs"
input_dir = "barite/"

timing_file = "timings.qs2"
ema_file = "ema.out.csv"

relevant_timings = ["Chemistry", "Phreeqc", "PHT_read", "PHT_write", "DHT_gather", "Interpolation", "Energy"]

grid_size = 400 * 400
iterations = 5000
calculations = grid_size * iterations

output_dir = "timings"
mkdir(output_dir)

function read_timings(timing_file)
        res = nothing

        @rput timing_file
        R"res <- qs2::qs_read(timing_file)"
        @rget res

        return res
end

function read_ema(ema_file)
        if (!isfile(ema_file))
                return DataFrame()
        end
        return CSV.read(ema_file, DataFrame)
end

function eval_timings_interp(timings)
        chemistry = timings[:chemistry]

        time_vec = [chemistry[:simtime],
                sum(chemistry[:phreeqc_time]),
                sum(chemistry[:interp_w]),
                sum(chemistry[:interp_r]),
                sum(chemistry[:interp_g]),
                sum(chemistry[:interp_fc])]

        interp_df = DataFrame(Calls=chemistry[:interp_calls],
                Cached=chemistry[:interp_cached])

        return time_vec, interp_df
end

function eval_ema(ema_df)
        if (nrow(ema_df) == 0)
                return 0.0
        end
        only_package = filter(row -> occursin(r"package-\d+", row[:device_name]), ema_df)
        return sum(only_package[!, :energy]) * 1e-6 * (1 / 3600)
end

function eval_timings_ref(timings)
        chemistry = timings[:chemistry]

        time_vec = [chemistry[:simtime],
                sum(chemistry[:phreeqc_time]),
                0, 0, 0, 0]

        return time_vec
end

function find_all_directories(input_dir, prefix)
        result_dirs = readdir(input_dir, join=true)

        result_dirs = filter(x -> occursin(prefix, x), result_dirs)

        return result_dirs
end

function workhorse_loop(dirs)
        df_timings = DataFrame()
        df_interp = Dict{Tuple{Int,Int},DataFrame}()

        for dir in dirs

                metainfo = split(basename(dir), "_")

                type = nothing
                aqueous = 0
                minerals = 0

                energy = eval_ema(read_ema(joinpath(dir, ema_file)))

                if (metainfo[3] == "ref")
                        type = "ref"
                        timings = read_timings(joinpath(dir, timing_file))
                        time_vec = eval_timings_ref(timings)
                else
                        aqueous = parse(Int, metainfo[3])
                        minerals = parse(Int, metainfo[4])
                        type = "( " * string(aqueous) * ", " * string(minerals) * " )"

                        timings = read_timings(joinpath(dir, timing_file))
                        time_vec, interp_df = eval_timings_interp(timings)

                        df_interp[(aqueous, minerals)] = interp_df
                end

                push!(time_vec, energy)

                new_df = DataFrame(relevant_timings .=> time_vec)
                local_df = DataFrame(Type=type, Aqueous=aqueous, Minerals=minerals)
                local_df = hcat(local_df, new_df)
                append!(df_timings, local_df)
        end

        return df_timings, df_interp
end

function eval_additional_info(df, interp_dict, grid_size, iterations)
        interp_full = Vector()

        for row in eachrow(df)
                if (row[:Type] != "ref")
                        interp_count = sum(interp_dict[(row[:Aqueous], row[:Minerals])][!, :Calls])
                        push!(interp_full, interp_count)
                else
                        push!(interp_full, 0)
                end
        end

        calculations = grid_size * iterations
        phreeqc_calls = calculations .- interp_full

        time_per_phreeqc = (df[!, :Phreeqc] ./ phreeqc_calls)

        df = hcat(df, DataFrame(Phreeqc_calls=phreeqc_calls, Time_per_Phreeqc=time_per_phreeqc, Interp_calls=interp_full))

        return df
end

dirs = find_all_directories("barite/", "barite_fgcs")

df_timings, df_interp = workhorse_loop(dirs)
group = groupby(df_timings, [:Type, :Aqueous, :Minerals])
combine_df = combine(group,
        relevant_timings .=> mean .=> relevant_timings,
)

df_timings = eval_additional_info(combine_df, df_interp, grid_size, iterations)
CSV.write("$output_dir/timings_barite.csv", df_timings)

for (key, value) in df_interp
        CSV.write("$output_dir/interp_barite_$(key[1])_$(key[2]).csv", value)
end

# Plot 

bar(df_timings[!, :Type], df_timings[!, :Chemistry], legend=false, title="Barite Timings", ylabel="Time [s]", xlabel="Type")
savefig("$output_dir/timings_barite.pdf")

bar(df_timings[!, :Type], df_timings[!, :Energy], legend=false, title="Barite Energy", ylabel="Energy [Wh]", xlabel="Type")
savefig("$output_dir/energy_barite.pdf")

prefix = "dolo_inter"
input_dir = "dolomite/"

dirs = find_all_directories("dolomite/", "dolo_inter")

df_timings, df_interp = workhorse_loop(dirs)
group = groupby(df_timings, [:Type, :Aqueous, :Minerals])
combine_df = combine(group,
        relevant_timings .=> mean .=> relevant_timings,
)

df_timings = eval_additional_info(combine_df, df_interp, grid_size, iterations)
CSV.write("$output_dir/timings_dolomite.csv", df_timings)

for (key, value) in df_interp
        CSV.write("$output_dir/interp_dolomite_$(key[1])_$(key[2]).csv", value)
end

bar(df_timings[!, :Type], df_timings[!, :Chemistry], legend=false, title="Dolomite Timings", ylabel="Time [s]", xlabel="Type")
savefig("$output_dir/timings_dolomite.pdf")

bar(df_timings[!, :Type], df_timings[!, :Energy], legend=false, title="Dolomite Energy", ylabel="Energy [Wh]", xlabel="Type")
savefig("$output_dir/energy_dolomite.pdf")