eaf.py

import argparse
from matplotlib import pyplot as plt
import numpy as np
from snp import main as run
import multiprocessing
from joblib import Parallel, delayed

num_cores = multiprocessing.cpu_count()


def get_args_parse():
    can = argparse.ArgumentParser()

    can.add_argument("--nb-run", type=int, default=10, help="number of runs")

    can.add_argument(
        "--quality-threshold", type=float, default=660, nargs="+", help="Quality threshold"
    )

    can.add_argument(
        "--not-parallel", type=bool, default=False, help="launch runs sequentialy"
    )

    can.add_argument("--solver", type=str, default="num_random", help="Solver to use.")

    return can


def main(eaf_args=None):

    can = get_args_parse()

    the = can.parse_args(eaf_args)

    args = ["--solver", the.solver, "--iters", "2000", "--steady-delta", "500"]

    if the.not_parallel:
        results = []
        for i in range(the.nb_run):
            # print(f"{i}th run")
            results.append(run(args))
    else:

        def f(i):
            # print(f"{i}th run")
            return run(args)

        results = Parallel(n_jobs=num_cores)(delayed(f)(i) for i in range(the.nb_run))

    values = [(val, iter) for val, iter, _, _ in results]

    values = sorted(values, key=lambda el: el[1])

    t = [iter for _, iter in values]
    values = np.array([val for val, _ in values])
    buff = np.empty((1,the.nb_run))
    for quality_threshold in the.quality_threshold:
        curr_qualities = np.cumsum(values >= quality_threshold) / len(values) 
        curr_qualities = np.expand_dims(curr_qualities, axis=0)
        buff = np.concatenate((buff,curr_qualities))

    fig = plt.figure()

    plt.title(f"{the.solver}, threshold {the.quality_threshold}")
    for i in range(1, buff.shape[0]):
        plt.step(t, buff[i])
    
    # ax4 = fig.add_subplot(122, projection="3d")
    # plt.title(f"{the.solver}, thresholds : {the.quality_threshold}")

    # ax4.plot_surface(t, values, the.quality_threshold )
    plt.legend([f'threshold : {the.quality_threshold[i]}' for i in range(len(the.quality_threshold))])
    plt.show()


if __name__ == "__main__":
    main()