source: anuga_work/development/Hinwood_2008/calc_rmsd.py @ 5709

"""

"""


#----------------------------------------------------------------------------
# Import necessary modules
#----------------------------------------------------------------------------

# Standard modules
import os
from os import sep, path
from csv import writer
from time import localtime, strftime

# Related major packages
from Numeric import arange, array, zeros, Float, where, greater, less, \
     compress, argmin, choose, searchsorted, sqrt, sum

import project          
from os import sep
from anuga.shallow_water.data_manager import csv2dict
from anuga.utilities.numerical_tools import ensure_numeric, err, norm

from slope import load_sensors
from anuga.utilities.interp import interp

def get_max_min_condition_array(min, max, vector):
    
    SMALL_MIN = -1e10  # Not that small, but small enough
    vector = ensure_numeric(vector)
    assert min > SMALL_MIN
    no_maxs = where(less(vector,max), vector, SMALL_MIN)
    #print "no_maxs", no_maxs
    band_condition = greater(no_maxs, min)
    return band_condition
    
def auto_rrms(outputdir_tag, scenarios, quantity, y_location_tag=':0.0'):
    """
    Given a bunch of scenarios that have CSV guage files, calc the
    err, Number_of_samples and rmsd for all gauges in each scenario.
    Write this info to a file for each scenario.
    """
    for run_data in scenarios:
        # Don't need to do this.  The simulation .csv files are already
        # anuga_times
        #anuga_start_stop_times = []
        #for time in run_data['wave_start_stop_times']:
        #    anuga_start_stop_times.append( \
        #        time -  run_data['ANUGA_start_time'])
            
        location_sims = []
        location_exps = []
        for gauge_x in run_data['gauge_x']:
            gauge_x = str(gauge_x)
            location_sims.append(gauge_x + y_location_tag)
            location_exps.append(gauge_x)
        
        id = run_data['scenario_id']
        outputdir_name = id + outputdir_tag
        pro_instance = project.Project(['data','flumes','Hinwood_2008'],
                                       outputdir_name=outputdir_name)
        end = id + ".csv"        
        file_sim = pro_instance.outputdir + quantity + "_" + end
        file_exp = pro_instance.raw_data_dir + sep + id + 'pressfilt_exp_' \
                   + quantity + '.csv'
        file_err = pro_instance.rmsd_dir + sep + outputdir_name + "_" \
                   + quantity + "_err.csv"
        

        simulation, _ = csv2dict(file_sim)
        experiment, _ = csv2dict(file_exp)
        
        time_sim = [float(x) for x in simulation['time']]
        time_exp = [float(x) for x in experiment['Time']]
        time_sim = ensure_numeric(time_sim)
        time_exp = ensure_numeric(time_exp)
        #print "min(time_exp)", min(time_exp)
        #print "max(time_exp)", max(time_exp)

        # Trim the experimental data
        condition_exp = get_max_min_condition_array(run_data['wave_times'][0],
                                                run_data['wave_times'][1],
                                                time_exp)
        time_exp_cut = compress(condition_exp, time_exp) #, axis=axis)
        #print "min(time_exp_cut)", min(time_exp_cut)
        #print "max(time_exp_cut)", max(time_exp_cut)
        
        #assert min(time_sim) < min(time_exp)

        # Trim the simulation data


        
        print "Writing to ", file_err
        
        err_list = []
        points = []
        rmsd_list = []
        for location_sim, location_exp in map(None, location_sims,
                                              location_exps):
            quantity_sim = [float(x) for x in simulation[location_sim]]
            quantity_exp = [float(x) for x in experiment[location_exp]]

            quantity_exp_cut = compress(condition_exp, quantity_exp)

            # Now let's do interpolation
            quantity_sim_interp = interp(quantity_sim, time_sim, time_exp_cut)

            assert len(quantity_sim_interp) == len(quantity_exp_cut)
            norm = err(quantity_sim_interp,
                       quantity_exp_cut,
                       2, relative = False)  # 2nd norm (rel. RMS)
            err_list.append(norm)
            points.append(len(quantity_sim_interp))
            rmsd_list.append(norm/sqrt(len(quantity_sim_interp)))              
            #print "norm", norm
            #for i in range(len(quantity_sim_interp)):
                
                #print "quantity_sim_interp", quantity_sim_interp[i]
                #print "quantity_exp_cut", quantity_exp_cut[i]
        assert len(location_exps) == len(err_list)

        # Writing the file out for one scenario
        a_writer = writer(file(file_err, "wb"))
        a_writer.writerow(["x location", "err", "Number_of_samples", "rmsd"])
        a_writer.writerows(map(None,
                               location_exps,
                               err_list,
                               points,
                               rmsd_list))


                  
def plot_rrms_sensor_settings(run_data, outputdir_tags, quantity,
              save_as=True,
              is_interactive=False, max_rmsd=None):
    """
    For a scenario, do 
    """
    from pylab import ion, plot, xlabel, ylabel, close, legend, \
         savefig, title, axis, setp, subplot, grid, axvspan, figlegend
    from anuga.shallow_water.data_manager import csv2dict


    plot_type = ".pdf"
    
    id = run_data['scenario_id']
    
    if is_interactive:
        ion()
    time_date = strftime('plot date: %d/%m/%Y Time: %H:%M:%S',
                                      localtime()) 
    subplot(212)    
    plot_title = id + " Root Mean Square Deviation comparison" + '\n' \
                 + time_date
                 
    title(plot_title)
    y_label = "RMSD"
    ylabel(y_label)
    xlabel("x location, m")
    grid(True)

    lines = []
    for outputdir_tag in outputdir_tags:
            
        outputdir_name = id + outputdir_tag
        pro_instance = project.Project(['data','flumes','Hinwood_2008'],
                                       outputdir_name=outputdir_name)
        
        file_err = pro_instance.rmsd_dir + sep + outputdir_name + "_" \
                   + quantity + "_err.csv"
            
        simulation, _ = csv2dict(file_err)
        locations = [float(x) for x in simulation['x location']]
        rmsd_list = [float(x) for x in simulation['rmsd']]
        lines.append(plot(locations, rmsd_list))
        
        if max_rmsd is not None:
            #print "setting axis"
            axis(ymin=0, ymax=max_rmsd)

    for break_x in run_data['break_xs']:
        axvspan(break_x-0.001,break_x+0.001, facecolor='g')
        
    legend(lines, outputdir_tags,'upper left')
    
    if is_interactive:
        # Wait for enter pressed
        raw_input()

    save_as = pro_instance.plots_dir + sep + \
              id + "_rmsd" + plot_type
    if save_as is not None:
        savefig(save_as)
    
    #Need to close this plot
    close()

def auto_plot_test_rmsd(scenarios, outputdir_tag,
                        to_publish_indexes, quantity="stage",
              save_as=True,add_run_info=False,
              is_interactive=False, max_rmsd=0.012):
    """
    Produce the RMSD graphs for the anuga article.

    """
    
    tests = []
    file_names = []
    
    for i in range(0,len(scenarios),2):
        if to_publish_indexes.has_key(scenarios[i]['scenario_id']):            
            tests.append([scenarios[i], scenarios[i+1]])
        else:            
            tests.append([scenarios[i+1], scenarios[i]])

    for test in tests:
        file_name = plot_test_rmsd(test, outputdir_tag, quantity,
                                   to_publish_indexes,
                                   save_as=save_as, add_run_info=add_run_info,
                                   is_interactive=is_interactive,
                                   max_rmsd=max_rmsd)
        file_names.append(file_name)
    return file_names
        
def plot_test_rmsd(test, outputdir_tag, quantity,
                   to_publish_indexes,
                   save_as=None, add_run_info=False,
                   is_interactive=False, max_rmsd=None):
    """
    For a test, plot rmsd vs x location for both runs.
    Also, add the break measurement.
    """
    from pylab import ion, plot, xlabel, ylabel, close, legend, \
         savefig, title, axis, setp, subplot, grid, axvspan, figlegend
    from anuga.shallow_water.data_manager import csv2dict
        
    plot_type = ".pdf"   
    id = 'S' + test[0]['scenario_id'][1:2]
    
    if add_run_info is True:
        plot_title = "Title will not be in final draft \n" + \
                     id + outputdir_tag
        
        title(plot_title)
    else:
        plot_title = ""
            
    if is_interactive:
        ion()
    time_date = strftime('plot date: %d/%m/%Y Time: %H:%M:%S',
                                      localtime()) 
    #subplot(212)    
   
                 
    #title(plot_title)
    y_label = "RMSD"
    ylabel(y_label)
    xlabel("x location, m")
    grid(True)

    lines = []
    legend_name = []
    i = 0
    for run in test:
        i += 1
        legend_name.append("S" + test[0]['scenario_id'][1:2]+"R"+str(i))    
        outputdir_name = run['scenario_id'] + outputdir_tag
        pro_instance = project.Project(['data','flumes','Hinwood_2008'],
                                       outputdir_name=outputdir_name)
        
        file_err = pro_instance.rmsd_dir + sep + outputdir_name + "_" \
                   + quantity + "_err.csv"
            
        simulation, _ = csv2dict(file_err)
        locations = [float(x) for x in simulation['x location']]
        rmsd_list = [float(x) for x in simulation['rmsd']]

        
        lines.append(plot(locations, rmsd_list, 'k-x')) # k is black

        for i_gauge, (location, rmsd) in enumerate(map(None,
                                                     locations,
                                                     rmsd_list)):
            if to_publish_indexes.has_key(run['scenario_id']) and \
                   to_publish_indexes[run['scenario_id']].count(i_gauge) == 1:
                plot([location], [rmsd], 'ko', markeredgewidth=2,
                     markeredgecolor='k')
                

        
        for break_x in run['break_xs']:
            axvspan(break_x-0.001,break_x+0.001, facecolor='g')
        
    if max_rmsd is not None:
        #print "setting axis"
        axis(ymin=0, ymax=max_rmsd)
        
    if id == 'S2':
        # Hack to get the desired width
        axis(xmax=7.0)
        
    setp(lines[1], linestyle='--')    
    legend(lines, legend_name,'upper left')
    
    if is_interactive:
        raw_input()

    save_as = pro_instance.plots_dir + sep + \
              id + "-rmsd" + plot_type
    print "save_as", save_as
    if save_as is not None:
        savefig(save_as)
    
    close()

    return save_as

# Return a bunch of lists
# The err files, for all scenarios
def err_files(scenarios, outputdir_tag, rmsd_dir, quantity):
    """
    The err files, for a list of scenarios
    """
    file_errs = []
    for scenario in scenarios:
        id = scenario['scenario_id']
        outputdir_name = id + outputdir_tag
        file_err = rmsd_dir + sep + outputdir_name + "_" \
                   + quantity + "_err.csv"
        file_errs.append(file_err)
    return file_errs
    

def compare_different_settings(outputdir_tags, scenarios, quantity,
                   save_as=None,
                   is_interactive=False):

    # A bit hacky.  Getting a pro_instance to get the rmsd_dir.
    outputdir_name = scenarios[0]['scenario_id'] + outputdir_tags[0]
    pro_instance = project.Project(['data','flumes','Hinwood_2008'],
                                   outputdir_name=outputdir_name)
    
     
    settings = {} # keys are different settings.
    # For each setting there will be err and amount
        
    for outputdir_tag in outputdir_tags:
        files = err_files(scenarios, outputdir_tag, pro_instance.rmsd_dir,
                          quantity)
        sim = {}
        for run_data, file in map(None, scenarios, files):
            
            simulation, _ = csv2dict(file)
            #locations = [float(x) for x in simulation['x location']]
            err_list = [float(x) for x in simulation['err']]
            amount_list = [float(x) for x in simulation['Number_of_samples']]
            
            if sim.has_key('err'):
                err_list.append(sim['err'])
                amount_list.append(sim['amount'])
            sim['err'], sim['amount'] = err_addition(err_list,
                                                     amount_list)
        sim['rmsd'] = sim['err']/sqrt(sim['amount'])
        settings[outputdir_tag] = sim
    #print "settings", settings

    aux = [(settings[k]['rmsd'], k) for k in settings.keys()]
    aux.sort()
    for val, key in aux:
        print key + "   " + str(val)
    
    
    
def err_addition(err_list, amount_list):
    """
    err1 is the err value (sqrt(sum_over_x&y((xi - yi)^2))) for a set of values
    amount1 is the number of values associated with the err.
    
    If this function gets used alot, maybe pull this out and make it an object
    """
    err = norm(ensure_numeric(err_list))
    amount = sum(ensure_numeric(amount_list))

    return err, amount


def calc_max_rmsd(scenarios, outputdir_tags, quantity):
    
    max_rmsd = 0
    
    for run_data in scenarios:
        id = run_data['scenario_id']
        for outputdir_tag in outputdir_tags:
            
            outputdir_name = id + outputdir_tag
            pro_instance = project.Project(['data','flumes','Hinwood_2008'],
                                           outputdir_name=outputdir_name)
            
            file_err = pro_instance.rmsd_dir + sep + outputdir_name + "_" \
                       + quantity + "_err.csv"
            
            simulation, _ = csv2dict(file_err)
            rmsd_list = [float(x) for x in simulation['rmsd']]
            max_rmsd = max(max(rmsd_list), max_rmsd)
    return max_rmsd

        
def auto_plot_rrms_sensor_settings(outputdir_tags, scenarios, quantity,
                   save_as=None,
                   is_interactive=False):
    """
    The general auto, for comparing things.

    """
    max_rmsd = calc_max_rmsd(scenarios, outputdir_tags, quantity)
    print "max_rmsd", max_rmsd
    for run_data in scenarios:
        plot_rrms_sensor_settings(run_data, outputdir_tags, quantity,
                                  max_rmsd=max_rmsd) 
                 
#-------------------------------------------------------------
if __name__ == "__main__":
    """ 
    """
    from scenarios import scenarios
    
    outputdir_tags = []
    outputdir_tags.append("_lmts_wdth_0.1_z_0.0_ys_0.01_mta_0.01_I")
    outputdir_tags.append("_no_velocity_nolmts_wdth_0.1_z_0.0_ys_0.01_mta_0.01_I")
    outputdir_tags.append("_nolmts_wdth_0.1_z_0.0_ys_0.01_mta_0.0001_I")
    outputdir_tags.append("_nolmts_wdth_0.1_z_0.0_ys_0.01_mta_0.001_I")
    outputdir_tags.append("_nolmts_wdth_0.1_z_0.0_ys_0.01_mta_0.01_I")
    outputdir_tags.append("_nolmts_wdth_0.1_z_0.012_ys_0.01_mta_0.01_I")
    outputdir_tags.append("_nolmts_wdth_1.0_z_0.0_ys_0.01_mta_0.01_I")
    outputdir_tags.append("_nolmts_wdth_0.1_z_0.012_ys_0.01_mta_0.0001_I")
    outputdir_tags.append("_lmts_wdth_0.1_z_0.0_ys_0.01_mta_0.0001_I")
    outputdir_tags.append("_no_velocity_nolmts_wdth_0.1_z_0.0_ys_0.01_mta_0.0001_I")
    
    outputdir_tags = []
    outputdir_tags.append("_nolmts_wdth_0.1_z_0.0_ys_0.01_mta_0.01_I")
    #outputdir_tag = "_test_limiterC"
    #scenarios = [scenarios[0]] # !!!!!!!!!!!!!!!!!!!!!!
    #scenarios = scenarios[4:] # !!!!!!!!!!!!!!!!!!!!!!
    calc_rmsds = True
    #calc_rmsds = False
    if calc_rmsds:
        for outputdir_tag in outputdir_tags:
            auto_rrms(outputdir_tag, scenarios, "stage", y_location_tag=':0.0')
    
    #scenarios = [scenarios[0]] # !!!!!!!!!!!!!!!!!!!!!!
    #auto_plot_rrms_sensor_settings(outputdir_tags, scenarios, "stage")
    compare_different_settings(outputdir_tags, scenarios, "stage")

    outputdir_tag = "_nolmts_wdth_0.1_z_0.0_ys_0.01_mta_0.01_I" 
    #auto_plot_test_rmsd(scenarios, outputdir_tag)