Inspyred post-processing of two reservoir problem in Pywr

two_reservoir_inspyred_pp.py

import csv
from matplotlib import pyplot as plt
import pandas
import numpy as np


def quoter(rows):
    for row in rows:
        yield row.replace('[', '"').replace(']', '"')


def load_individuals(filename):

    with open(filename) as fh:
        reader = csv.reader(quoter(fh.readlines()), skipinitialspace=True)
        for row in reader:
            name, i, objectives, variables = row

            objectives = [float(v) for v in objectives.split(',')]
            variables = [float(v) for v in variables.split(',')]

            yield objectives, variables


def plot_individuals(filename, ax_var=None, ax_obj=None, color='k', alpha=0.2):

    if ax_var is None:
        fig, ax_var = plt.subplots()
    if ax_obj is None:
        fig, ax_obj = plt.subplots()

    dfs = []
    for obj, vars in load_individuals(filename):
        if len(vars) == 12:
            df = monthly_profile(vars)
        else:
            df = harmonic_profile(vars)
        #df.plot(ax=ax_var, c=color, alpha=0.1)
        dfs.append(df)

        ax_obj.scatter(obj[0], obj[1], color=color)

    dfs = pandas.concat(dfs, axis=1)
    dfs.mean(axis=1).plot(ax=ax_var, c=color, zorder=1)


def monthly_profile(variables):

    index = pandas.date_range('2017-01', periods=13, freq='MS')
    df = pandas.Series([v for v in variables] + [variables[-1]], index=index)
    return df.resample('D').pad()


def harmonic_profile(variables):

    index = pandas.date_range('2017-01', periods=365, freq='D')
    doy = index.dayofyear

    n = (len(variables) - 1) // 2

    values = variables[0]

    for i in range(n):
        amp = variables[1 + i]
        phase = variables[1 + n + i]
        values += amp*np.cos(doy*(i+1)*np.pi*2/365 + phase)

    return pandas.Series(values, index=index)


if __name__ == '__main__':

    fig, ax_var = plt.subplots()
    fig, ax_obj = plt.subplots()
    #plot_individuals('two_reservoir_moea-monthly-final-individuals-file.csv', ax_var=ax_var, ax_obj=ax_obj, color='b')
    plot_individuals('two_reservoir_moea-monthly-individuals-file.csv', ax_var=ax_var, ax_obj=ax_obj, color='b')

    #plot_individuals('two_reservoir_moea-harmonic-final-individuals-file.csv', ax_var=ax_var, ax_obj=ax_obj, color='r')
    plot_individuals('two_reservoir_moea-harmonic-individuals-file.csv', ax_var=ax_var, ax_obj=ax_obj, color='r')

    ax_var.set_ylim(0, 1)
    plt.show()