<reponame>Julian-Theis/AVATAR
import argparse
import os
import tensorflow as tf
import numpy as np

from avatar.util.LoadRelgan import LoadRelgan
from avatar.util.MHGAN import MHGAN
from avatar.util.util import writeToFile, readTraces
from conf.settings import DATA_PATH

WORK_PATH = os.path.abspath(os.getcwd())

if __name__ == "__main__":
    np.random.seed(seed=1234)
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--system', help='System to evaluate (e.g. pa_system_2_3)', required=True)
    parser.add_argument('-j', '--job', help='0 = beta 100, 1 = beta 1000', required=True)
    parser.add_argument('-sfx', '--suffix', help='Which suffix, i.e. final epoch of the trained SGAN to use?', required=True)
    parser.add_argument('-gpu', '--gpu', help='GPU on which the training is performed. For example 0.', required=True)

    """ Selector """
    parser.add_argument('-strategy', '--strategy', help='select "naive" or "mh"', required=True)

    """ Parameter for Naively Sampling """
    parser.add_argument('-n_n', '--n_samples', help='(NAIVE ONLY) Number of samples to generate? (Default: 10000)', default=10000)

    """ Parameter for MH Sampling """
    parser.add_argument('-mh_c', '--mh_count', help='(MH ONLY) Number of samples per batch? (Default: 50)',
                        default=50)
    parser.add_argument('-mh_p', '--mh_patience',
                        help='(MH ONLY) Patience parameter (Default: 5)',
                        default=5)
    parser.add_argument('-mh_k', '--mh_k',
                        help='(MH ONLY) Length of Markov chain (Default: 500)',
                        default=500)
    parser.add_argument('-mh_mi', '--mh_maxiter',
                        help='(MH ONLY) Max sampling iterations? (Default: 200)',
                        default=200)

    args = parser.parse_args()
    system = args.system
    job = int(args.job)
    suffix = args.suffix
    strategy = args.strategy
    n_samples = int(args.n_samples)
    mh_count = int(args.mh_count)
    mh_patience = int(args.mh_patience)
    mh_k = int(args.mh_count)
    mh_maxiter = int(args.mh_maxiter)

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    if strategy == "naive":
        tf.reset_default_graph()
        print("****** SAMPLE FOR SUFFIX ", suffix, " ******")
        relgan = LoadRelgan(system=system, suffix=suffix, job=job)

        if DATA_PATH is None:
            f_out = os.path.join(WORK_PATH, "data", "avatar", "variants", system + "_relgan_" + str(suffix) + "_j" + str(job) + "_naive.txt")
        else:
            f_out = os.path.join(DATA_PATH, "avatar", "variants",system + "_relgan_" + str(suffix) + "_j" + str(job) + "_naive.txt")

        print("Start NAIVE SAMPLING")
        gen_samples = relgan.generate(n_samples=n_samples)
        print("Generated samples - shape:", gen_samples.shape)
        print("Writing to file", f_out)
        writeToFile(relgan, f_out, gen_samples)

    elif strategy == "mh":
        if DATA_PATH is None:
            eval_path = os.path.join(WORK_PATH, "data", "avatar", "train_data", system + "_eval.txt")
            f_out = os.path.join(WORK_PATH, "data", "avatar", "variants", system + "_relgan_" + str(suffix) + "_j" + str(job) + "_mh.txt")
        else:
            eval_path = os.path.join(DATA_PATH, "avatar", "train_data", system + "_eval.txt")
            f_out = os.path.join(DATA_PATH, "avatar", "variants", system + "_relgan_" + str(suffix) + "_j" + str(job) + "_mh.txt")

        tf.reset_default_graph()
        print("****** SAMPLE FOR SUFFIX ", suffix, " ******")
        relgan = LoadRelgan(system=system, suffix=suffix, job=job)

        calibrate = readTraces(eval_path)
        calibrate = relgan.prep(calibrate)

        mhgan = MHGAN(relgan, c=mh_count, k=mh_k, real_samples=calibrate)
        samples = None
        gen_size = 0
        iter = 1
        cnt_patience = 0
        continue_sampling = True

        print("Start MH SAMPLING")
        while continue_sampling:
            print("**** MH-GAN Iteration", iter, ":")
            gen_samples, accepts, rejects = mhgan.generate_enhanced(
                relgan.sess,
                count=mh_count,
                k=mh_k
            )

            if samples is None:
                samples = gen_samples
            else:
                samples = np.concatenate([samples, gen_samples], axis=0)

            samples = np.unique(samples, axis=0)

            if gen_size != samples.shape[0]:
                cnt_patience = 0
            else:
                cnt_patience += 1

            gen_size = samples.shape[0]
            print("Generated samples (cumulative): ", gen_size)
            iter += 1

            if cnt_patience >= mh_patience:
                continue_sampling = False

            if mh_maxiter != -1 and iter >= mh_maxiter:
                continue_sampling = False

        print("Generated samples - shape:", samples.shape)
        print("Writing to file", f_out)
        writeToFile(relgan, f_out, samples)

    else:
        raise ValueError("Unknown sampling strategy.")

import argparse
import os
import tensorflow as tf
import numpy as np

from avatar.util.LoadRelgan import LoadRelgan
from avatar.util.MHGAN import MHGAN
from avatar.util.util import writeToFile, readTraces
from conf.settings import DATA_PATH

WORK_PATH = os.path.abspath(os.getcwd())

if __name__ == "__main__":
    np.random.seed(seed=1234)
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--system', help='System to evaluate (e.g. pa_system_2_3)', required=True)
    parser.add_argument('-j', '--job', help='0 = beta 100, 1 = beta 1000', required=True)
    parser.add_argument('-sfx', '--suffix', help='Which suffix, i.e. final epoch of the trained SGAN to use?', required=True)
    parser.add_argument('-gpu', '--gpu', help='GPU on which the training is performed. For example 0.', required=True)

    """ Selector """
    parser.add_argument('-strategy', '--strategy', help='select "naive" or "mh"', required=True)

    """ Parameter for Naively Sampling """
    parser.add_argument('-n_n', '--n_samples', help='(NAIVE ONLY) Number of samples to generate? (Default: 10000)', default=10000)

    """ Parameter for MH Sampling """
    parser.add_argument('-mh_c', '--mh_count', help='(MH ONLY) Number of samples per batch? (Default: 50)',
                        default=50)
    parser.add_argument('-mh_p', '--mh_patience',
                        help='(MH ONLY) Patience parameter (Default: 5)',
                        default=5)
    parser.add_argument('-mh_k', '--mh_k',
                        help='(MH ONLY) Length of Markov chain (Default: 500)',
                        default=500)
    parser.add_argument('-mh_mi', '--mh_maxiter',
                        help='(MH ONLY) Max sampling iterations? (Default: 200)',
                        default=200)

    args = parser.parse_args()
    system = args.system
    job = int(args.job)
    suffix = args.suffix
    strategy = args.strategy
    n_samples = int(args.n_samples)
    mh_count = int(args.mh_count)
    mh_patience = int(args.mh_patience)
    mh_k = int(args.mh_count)
    mh_maxiter = int(args.mh_maxiter)

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    if strategy == "naive":
        tf.reset_default_graph()
        print("****** SAMPLE FOR SUFFIX ", suffix, " ******")
        relgan = LoadRelgan(system=system, suffix=suffix, job=job)

        if DATA_PATH is None:
            f_out = os.path.join(WORK_PATH, "data", "avatar", "variants", system + "_relgan_" + str(suffix) + "_j" + str(job) + "_naive.txt")
        else:
            f_out = os.path.join(DATA_PATH, "avatar", "variants",system + "_relgan_" + str(suffix) + "_j" + str(job) + "_naive.txt")

        print("Start NAIVE SAMPLING")
        gen_samples = relgan.generate(n_samples=n_samples)
        print("Generated samples - shape:", gen_samples.shape)
        print("Writing to file", f_out)
        writeToFile(relgan, f_out, gen_samples)

    elif strategy == "mh":
        if DATA_PATH is None:
            eval_path = os.path.join(WORK_PATH, "data", "avatar", "train_data", system + "_eval.txt")
            f_out = os.path.join(WORK_PATH, "data", "avatar", "variants", system + "_relgan_" + str(suffix) + "_j" + str(job) + "_mh.txt")
        else:
            eval_path = os.path.join(DATA_PATH, "avatar", "train_data", system + "_eval.txt")
            f_out = os.path.join(DATA_PATH, "avatar", "variants", system + "_relgan_" + str(suffix) + "_j" + str(job) + "_mh.txt")

        tf.reset_default_graph()
        print("****** SAMPLE FOR SUFFIX ", suffix, " ******")
        relgan = LoadRelgan(system=system, suffix=suffix, job=job)

        calibrate = readTraces(eval_path)
        calibrate = relgan.prep(calibrate)

        mhgan = MHGAN(relgan, c=mh_count, k=mh_k, real_samples=calibrate)
        samples = None
        gen_size = 0
        iter = 1
        cnt_patience = 0
        continue_sampling = True

        print("Start MH SAMPLING")
        while continue_sampling:
            print("**** MH-GAN Iteration", iter, ":")
            gen_samples, accepts, rejects = mhgan.generate_enhanced(
                relgan.sess,
                count=mh_count,
                k=mh_k
            )

            if samples is None:
                samples = gen_samples
            else:
                samples = np.concatenate([samples, gen_samples], axis=0)

            samples = np.unique(samples, axis=0)

            if gen_size != samples.shape[0]:
                cnt_patience = 0
            else:
                cnt_patience += 1

            gen_size = samples.shape[0]
            print("Generated samples (cumulative): ", gen_size)
            iter += 1

            if cnt_patience >= mh_patience:
                continue_sampling = False

            if mh_maxiter != -1 and iter >= mh_maxiter:
                continue_sampling = False

        print("Generated samples - shape:", samples.shape)
        print("Writing to file", f_out)
        writeToFile(relgan, f_out, samples)

    else:
        raise ValueError("Unknown sampling strategy.")
Selector Parameter for Naively Sampling Parameter for MH Sampling