Generic Experiment with plot and save option

experiment.py

"""
    Generic example of experiment that compares two methods
    COS, SIN in two different experiments EXPERIMENT1, EXPERIMENT2
    Experiments can be lengthy so one can choose to SAVE the
    files so that the plots can be generated again when the
    appropriate RUNTAG is set and the PLOT program option is chosen.
"""
import os

import matplotlib.pyplot as plt
import numpy as np


class Method:
    COS = "cosine"
    SIN = "sine"


class Program:
    EXPERIMENT1 = 1
    EXPERIMENT2 = 2
    PLOT = 3


RUNTAG = 'run1'
FILETAG = 'experiment'
OUTPUTFOLDER = 'outputs'
STEPS = 5000
STEP_SIZE = 1e-3
SAVE = True

PROGRAM = Program.EXPERIMENT1


def plot_result(y, label):
    plt.plot(y, label=label)


class NameGen(object):
    def __init__(self, filetag, output_folder):
        try:
            os.makedirs(output_folder, exist_ok=True)
        except OSError:
            raise

        self.output_folder = output_folder
        self.filetag = filetag

    def get_output_prefix(self, exptag, label):
        return os.path.join(self.output_folder,
                            '{}_{}{}'.format(self.filetag, label, exptag))

    def get_output_filename(self, exptag, label, ext):
        return '{}.{}'.format(self.get_output_prefix(exptag, label), ext)

    def save_plot(self, tag, label='', dpi=300, ext='png'):
        plt.draw()
        plt.savefig(self.get_output_filename(tag, label, ext), dpi=dpi)


class Simulator(object):
    def __init__(self, namegen, runtag, save):
        self.namegen = namegen
        self.runtag = runtag
        self.save = save



    def run(self, func, steps, step_size, label):
        y = func(np.arange(steps) * step_size)
        plot_result(y, label=label)

        if self.save:
            npz_filename = self.namegen.get_output_filename(self.runtag, label, 'npz')
            np.savez(npz_filename, y=y)

    def load(self, label):
        npz_filename = self.namegen.get_output_filename(self.runtag, label, 'npz')
        y = np.load(npz_filename)['y']
        plot_result(y, label=label)

    def finalize_plot(self):
        plt.legend(loc='best', prop={'size': 12})

        plt.xlabel('Simulation steps')
        plt.ylabel('Output')

        self.namegen.save_plot(self.runtag, "")
        plt.clf()


def factory(method_name, **kwargs):
    if method_name == Method.COS:
        t = kwargs.pop('term', 0)
        return lambda x: np.cos(x + t)
    elif method_name == Method.SIN:
        f = kwargs.pop('factor', 1)
        return lambda x: np.sin(f * x)
    else:
        raise ValueError("Method: '{}' is not valid.".format(method_name))


def main_example1(simulator):
    for method_name in [Method.COS, Method.SIN]:
        myfunc = factory(method_name)
        simulator.run(myfunc, STEPS, STEP_SIZE, label=method_name)


def main_example2(simulator):
    for method_name in [Method.COS, Method.SIN]:
        myfunc = factory(method_name, factor=2, term=1)
        simulator.run(myfunc, STEPS, STEP_SIZE, label=method_name)


def main_plot(simulator):
    for label in [Method.COS, Method.SIN]:
        simulator.load(label)


if __name__ == '__main__':
    mynamegen = NameGen(FILETAG, OUTPUTFOLDER)
    mysimulator = Simulator(mynamegen, RUNTAG, SAVE)

    if PROGRAM == Program.EXPERIMENT1:
        main_example1(mysimulator)
    elif PROGRAM == Program.EXPERIMENT2:
        main_example2(mysimulator)
    elif PROGRAM == Program.PLOT:
        main_plot(mysimulator)

    mysimulator.finalize_plot()