source: anuga_work/development/Hinwood_2008/slope.py @ 5532

"""
Plot up files from the Hinwood project.
"""
from os import sep
import project
from copy import deepcopy
#from scipy import arange
from csv import writer

from Numeric import arange, array, zeros, Float, where, greater, less, \
     compress, argmin, choose, searchsorted

from anuga.fit_interpolate.interpolate import interpolate_sww2csv
from anuga.shallow_water.data_manager import csv2dict
from anuga.utilities.numerical_tools import ensure_numeric


SLOPE_STR = 'stage_slopes'
TIME_STR = 'times'

TIME_BORDER = 5
LOCATION_BORDER = .5

def load_sensors(quantity_file):
    """
    Load a csv file, where the first row is the column header and
    the first colum explains the rows.
    """
    #slope, _ = csv2dict(file_sim)
    
    # Read the depth file
    dfid = open(quantity_file)
    lines = dfid.readlines()
    dfid.close()

    title = lines.pop(0)
    n_time = len(lines)
    n_sensors = len(lines[0].split(','))-1  # -1 to remove time
    times = zeros(n_time, Float)  #Time
    depths = zeros(n_time, Float)  #
    sensors = zeros((n_time,n_sensors), Float)
    quantity_locations = title.split(',') #(',')
    quantity_locations.pop(0) # remove 'time'
    
    locations = [float(j.split(':')[0]) for j in quantity_locations]
    
    for i, line in enumerate(lines):
        fields = line.split(',') #(',')
        fields = [float(j) for j in fields]
        times[i] = fields[0]
        sensors[i] = fields[1:] # 1: to remove time

    #print "times",times
    #print "locations", locations
    #print "sensors", sensors
    return times, locations, sensors
   
def load_slopes(stage_file):
    """
    Finds the slope, wrt distance of a distance, time, quantity csv file.

    returns the times and slope_locations vectors and the slopes array.
    """
    times, locations, sensors = load_sensors(stage_file)
    n_slope_locations = len(locations)-1
    n_time = len(times)
    slope_locations = zeros(n_slope_locations, Float)  #
    slopes = zeros((n_time,n_slope_locations), Float)

    # An array of the sensor spacing values
    delta_locations = zeros(n_slope_locations, Float)
    
    for i in arange(n_slope_locations):
        delta_locations[i] = (locations[i+1] - locations[i])
        slope_locations[i] = locations[i] + 0.5*delta_locations[i]
    
    for j in arange(n_time):
        for i in arange(n_slope_locations):
            slopes[j,i] = (sensors[j,i+1] - sensors[j,i])/delta_locations[i]

    return times, slope_locations, slopes


def graph_contours(times, x_data, z_data,
                 y_label='Time, seconds',
                 plot_title="slope",
                 x_label='x location, m',
                 save_as=None,
                 is_interactive=False,
                 break_xs=None,
                 break_times=None):
    """
    Currently used to generate stage slope contour graphs.

    Has been generalised a bit.
    """
    # Do not move these imports.  Tornado doesn't have pylab
    from pylab import meshgrid, cm, contourf, contour, ion, plot, xlabel, \
         ylabel, close, legend, savefig, title, figure ,colorbar, show , axis

    origin = 'lower'

    if is_interactive:
        ion()
        
    # Can't seem to reshape this info once it is in the function
    CS = contourf(x_data, times, z_data, 10,
                  cmap=cm.bone,
                  origin=origin)
    
    #CS2 = contour(x_data, times, z_data, CS.levels[::1],
     #             colors = 'r',
      #            origin=origin,
       #           hold='on')
    
    title(plot_title)
    xlabel(x_label)
    ylabel(y_label)

    if break_times is not None and break_xs is not None:
        plot(break_xs, break_times, 'ro')
    
    
    # Make a colorbar for the ContourSet returned by the contourf call.
    cbar = colorbar(CS)

    # Add the contour line levels to the colorbar
    cbar.ax.set_ylabel('stage slope')
    #cbar.add_lines(CS2)
         
    if is_interactive:       
        raw_input() # Wait for enter pressed
        
    if save_as is not None:
        savefig(save_as)
    close()  #Need to close this plot

def graph_froude(times, x_data, z_data,
                 y_label='Time, seconds',
                 plot_title="Froude Number",
                 x_label='x location, m',
                 save_as=None,
                 is_interactive=False,
                 break_xs=None,
                 break_times=None):
    """
    Used to generate a froude Number contour graphs.

    """
    # Do not move these imports.  Tornado doesn't have pylab
    from pylab import meshgrid, cm, contourf, contour, ion, plot, xlabel, \
         ylabel, close, legend, savefig, title, figure ,colorbar, show , axis

    origin = 'lower'

    if is_interactive:
        ion()
        
    # Can't seem to reshape this info once it is in the function
    #CS = contourf(x_data, times, z_data, [-1,0.6,0.8,1,2,4],
     #             colors = ('black', 'r', 'g', 'b','r'),
      #            #cmap=cm.bone,
       #           origin=origin)
    CS = contourf(x_data, times, z_data, 10,
                  #colors = ('black', 'r', 'g', 'b','r'),
                  cmap=cm.bone,
                  origin=origin)
    
    #CS2 = contour(x_data, times, z_data, CS.levels[::1],
     #             colors = 'r',
      #            origin=origin,
       #           hold='on')
    
    title(plot_title)
    xlabel(x_label)
    ylabel(y_label)

    if break_times is not None and break_xs is not None:
        plot(break_xs, break_times, 'yo')
    
    
    # Make a colorbar for the ContourSet returned by the contourf call.
    cbar = colorbar(CS)

    # Add the contour line levels to the colorbar
    cbar.ax.set_ylabel('Froude Number')
    #cbar.add_lines(CS2)
         
    if is_interactive:       
        raw_input() # Wait for enter pressed
        
    if save_as is not None:
        savefig(save_as)
    close()  #Need to close this plot
    
def auto_graph_slopes(outputdir_tag, scenarios, is_interactive=False):
    """
    Used to generate all the stage slope contour graphs of a scenario list
    """
    plot_type = ".pdf"
    for run_data in scenarios:
        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"
        anuga_break_times = []
        for break_time in run_data['break_times']:
            anuga_break_times.append( \
                break_time -  run_data['ANUGA_start_time'])
        stage_file = pro_instance.outputdir + "fslope_stage_" + end
        plot_title = "Stage slope " + id + "\n file:" + \
                     outputdir_name + "_slope_stage" + plot_type
        print "Creating ", stage_file
        save_as = pro_instance.plots_dir + sep + \
                            outputdir_name + "_slope_stage" + plot_type
        times, locations, slopes = load_slopes(stage_file)
        #times, slopes = get_band(anuga_break_times[0]-TIME_BORDER,
         #                          100, times, slopes, 0)
        #locations, slopes = get_band(
         #   min(run_data['break_xs'])- 2*LOCATION_BORDER,
          #  100, locations, slopes, -1)
        graph_contours(times, locations, slopes,
                       plot_title=plot_title,
                       break_xs=run_data['break_xs'],
                       break_times=anuga_break_times,
                       save_as=save_as,
                       is_interactive=is_interactive)

def auto_graph_froudes(outputdir_tag, scenarios, is_interactive=False):
    """
    Used to generate all the Froude number contour graphs of a scenario list
    """
    
    plot_type = ".pdf"
    
    for run_data in scenarios:
        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"
        anuga_break_times = []
        for break_time in run_data['break_times']:
            anuga_break_times.append( \
                break_time -  run_data['ANUGA_start_time'])
        plot_title = "Froude Number" + id + "\n file:" + \
                     outputdir_name + "_froude" + plot_type
        froude_file = pro_instance.outputdir + "fslope_froude_" + end
        print "Creating ", froude_file
        save_as = pro_instance.plots_dir + sep + \
                            outputdir_name + "_froude" + plot_type
        dtimes, locations, sensors = load_sensors(froude_file)
        dtimes, sensors = get_band(anuga_break_times[0]-TIME_BORDER,
                                   100, dtimes, sensors, 0)
        locations, sensors = get_band(
            min(run_data['break_xs'])-LOCATION_BORDER,
            100, locations, sensors, -1)
        #print "dtimes", dtimes
        #print "sensors", sensors
        #times, slope_locations, slopes = load_slopes(stage_file)
        graph_froude(dtimes, locations, sensors,
                     plot_title=plot_title,
                     break_xs=run_data['break_xs'],
                     break_times=anuga_break_times,
                     save_as=save_as,
                     is_interactive=is_interactive)
        
def find_froude(times_froude, locations_froude, froudes_array,
                times, locations):
    """
    interpolate across location to find froude number values
    """
    
    if len(times) == 0:
        return []
    time_indexes = searchsorted(times_froude, times)
    location_indexes = searchsorted(locations_froude, locations)

    
    assert len(time_indexes) == len(location_indexes)

    froudes = []
    for time_i, loc_i, time, location in map(None, time_indexes,
                                             location_indexes,
                                             times, locations):
        # the time values should be the same
        assert times_froude[time_i] == time

        # The distance value should be half way between the froude locations
        midpoint = locations_froude[loc_i-1] + \
                   (locations_froude[loc_i]-locations_froude[loc_i-1])*0.5
        #print "location", location
        #print "midpoint", midpoint
        assert location == midpoint
        froude = froudes_array[time_i, loc_i-1] + \
                   (froudes_array[time_i, loc_i]- \
                    froudes_array[time_i, loc_i-1])*0.5
        froudes.append(froude)
                
    return froudes
    
def auto_find_min_slopes(slope_tag, outputdir_tag, scenarios):
    """
    Given stage and froude wrt time and location csv files,
    find the waves and get the froude number and stage slope
    at the wave face.

    For each wave write a csv file giving the location, stage slope, time and
    froude number.
    """
    
    for run_data in scenarios:
        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"
        anuga_break_times = []
        for break_time in run_data['break_times']:
            anuga_break_times.append( \
                break_time -  run_data['ANUGA_start_time'])
            
        stage_file = pro_instance.outputdir + slope_tag + "slope_stage_" + end
        froude_file = pro_instance.outputdir + slope_tag + "slope_froude_" + \
                      end
        
        times, slope_locations, slopes = load_slopes(stage_file)
        #print "slope_locations", slope_locations
        times_froude, locations_froude, froudes_a = load_sensors(froude_file)
        #print "locations_froude", locations_froude
        waves = find_min_slopes(times, slope_locations, slopes,
                                anuga_break_times,
                                run_data['band_offset'])
        
        # write the wave info here
        # and find the froude values
        for i, wave in enumerate(waves):
            
            id = "wave_" + str(i)
            wave_file = stage_file[:-4] + '_'+ id + ".csv"
            print "wave_file", wave_file
            froudes = find_froude(times_froude, locations_froude,
                             froudes_a, wave[TIME_STR],
                                  slope_locations)
            wave_writer = writer(file(wave_file, "wb"))
            wave_writer.writerow(["x location", "min slope", "Time", "Froude"])
            wave_writer.writerows(map(None,
                                      slope_locations,
                                      wave[SLOPE_STR],
                                      wave[TIME_STR],
                                      froudes))

def calc_wave_file_min_slope_max_froude(slope_tag, outputdir_tag, scenarios):
    """
    Calc the min slope and max froude number in the wave files
    Used so all graphs have the same axis.
    """
    min_slope = 0
    max_froude = 0
    
    for run_data in scenarios:
        for wave_file, save_as, wave_number in Get_file_name(
        run_data, outputdir_tag, slope_tag):
            simulation, _ = csv2dict(wave_file)
            slope = [float(x) for x in simulation['min slope']]
            froude = [float(x) for x in simulation['Froude']]
            
            min_slope = min(min(slope), min_slope)
            
            max_froude = max(max(froude), max_froude)
           
    
    return min_slope, max_froude

class Get_file_name:
    """
    Used to make the file names, and workout the wave number.
    """
    
    def __init__(self, run_data, outputdir_tag, slope_tag):
        
        self.plot_type = ".pdf"
        # The scenario data
        id = run_data['scenario_id']
        
        self.outputdir_name = id + outputdir_tag
        self.pro_instance = project.Project(['data','flumes','Hinwood_2008'],
                                            outputdir_name=self.outputdir_name)
        self.wave_number = -1
        self.max_waves = len(run_data['break_type'])
        self.slope_tag = slope_tag
        self.end = id + ".csv"

    def next(self):
        self.wave_number += 1
        if self.wave_number >= self.max_waves: raise StopIteration
        wave_tag = "wave_" + str(self.wave_number)
        stage_file = self.pro_instance.outputdir + self.slope_tag + \
                     "slope_stage_" + self.end
        wave_file = stage_file[:-4] + '_'+ wave_tag + ".csv"
        save_as = self.pro_instance.plots_dir + sep + \
                  self.outputdir_name + "_" + wave_tag + self.plot_type
        return wave_file, save_as, self.wave_number

    def __iter__(self):
        return self
        



def auto_plot_froude_slopes(slope_tag, outputdir_tag, scenarios):
    """
    Used to generate all the Froude number, stage slope, time graphs
    of a scenario list
    """

    slope_min, froude_max = calc_wave_file_min_slope_max_froude(
        slope_tag, outputdir_tag, scenarios)
    
    
    for run_data in scenarios:
        assert len(run_data['break_times']) == len(run_data['break_xs'])
        assert len(run_data['break_times']) == len(run_data['break_type'])
        
        anuga_break_times = []
        for break_time in run_data['break_times']:
            anuga_break_times.append( \
                break_time -  run_data['ANUGA_start_time'])
            
        for wave_file, save_as, wave_number in Get_file_name(
        run_data, outputdir_tag, slope_tag):
            print "wave_file", wave_file
            break_type = run_data['break_type'][wave_number]
            plot_title = run_data['scenario_id'] + \
                         ' Wave: ' + str(wave_number) + \
                         ' Break Type: ' + break_type + '\n' + \
                         'File: ' + wave_file[34:] # not good!
            plot_foude_slope_stage(wave_file,
                                   anuga_break_times[wave_number],
                                   run_data['break_xs'][wave_number],
                                   plot_title=plot_title,
                                   break_type=break_type,
                                   save_as=save_as,
                                   is_interactive=False,
                                   froude_min=0,
                                   froude_max=froude_max,
                                   slope_min=slope_min,
                                   slope_max=0)
        
        
    
def gauges_for_slope(slope_tag, outputdir_tag, scenarios):
    """
    This is used to create a stage file, using gauges relivent to
    finding a slope.

    It also create's a frounde file.
    """
    dx = 0.05
    for run_data in scenarios:
        point_x = arange(run_data['start_slope_x'],
                        run_data['finish_slope_x'],
                        dx).tolist()
        flume_y_middle = 0.5
        points = []
        for gauge_x in point_x:
            points.append([gauge_x, flume_y_middle])
        id = run_data['scenario_id']
    
        basename = 'zz_' + run_data['scenario_id']
        outputdir_name = id + outputdir_tag
        pro_instance = project.Project(['data','flumes','Hinwood_2008'],
                                       outputdir_name=outputdir_name)
        end = id + ".csv"
        interpolate_sww2csv( \
            pro_instance.outputdir + basename +".sww",
            points,
            pro_instance.outputdir + slope_tag + "slope_depth_" + end,
            pro_instance.outputdir + slope_tag + "slope_velocity_x_" + end,
            pro_instance.outputdir + slope_tag + "slope_velocity_y_" + end,
            pro_instance.outputdir + slope_tag + "slope_stage_" + end,
            pro_instance.outputdir + slope_tag + "slope_froude_" + end,
            time_thinning=1)

def find_min_slopes(times, slope_locations, slopes,
                    anuga_break_times, band_offset):
    """
    
    """
    bands = break_times2bands(anuga_break_times, band_offset)

    waves = []
    for i,_ in enumerate(bands[0:-1]):
        max_q, max_q_times = get_min_in_band(bands[i], bands[i+1],
                                             times, slopes)
        waves.append({SLOPE_STR:max_q, TIME_STR:max_q_times})
    return waves

    
def get_band(min, max, vector, quantity_array, axis):
    """
    Return a band of vector and quantity, within (not including) the
    min, max.

    For a time band, set the axis to 0.
    For a location band, set the axis to -1.
    
    """

    SMALL_MIN = -1e10  # Not that small, but small enough
    vector = ensure_numeric(vector)
    quantity_array = ensure_numeric(quantity_array)

    assert min > SMALL_MIN
    no_maxs = where(less(vector,max), vector, SMALL_MIN)
    #print "no_maxs", no_maxs
    band_condition = greater(no_maxs, min)
    band_vector = compress(band_condition, vector, axis=axis)
    #print "band_time", band_time
    #print "quantity_array", quantity_array.shape
    band_quantity = compress(band_condition, quantity_array, axis=axis)
    return band_vector, band_quantity

def get_min_in_band(min_time, max_time, time_vector, quantity_array):
    """
    given a quantity array, with the 2nd axis being
    time, represented by the time_vector, find the minimum within
    the time band.

    Assumes times are positive
    """
    
    time_vector = ensure_numeric(time_vector)
    quantity_array = ensure_numeric(quantity_array)

    band_time, band_quantity  = get_band(min_time, max_time,
                                         time_vector, quantity_array, 0)
    #print "band_quantity",band_quantity
    try:
        max_quantity_indices = argmin(band_quantity, axis=0)
    except:
        #print "time_vector", time_vector
        print "min_time",min_time 
        print "max_time", max_time
        return [],[]
        
    #print "max_quantity_indices", max_quantity_indices
    max_quantity_times = choose(max_quantity_indices, band_time)
    #print "max_quantity_times", max_quantity_times
    max_quantities = choose(max_quantity_indices, band_quantity)
    #print "max_quantities", max_quantities

    return max_quantities, max_quantity_times

def break_times2bands(break_times, band_offset):
    """
    Break_times is a list of times, ascending.
    bands is a list of times, being the midpoints of break_times, with a min
    and max band added.
    """
    assert len(break_times)>2
    
    bands = [] #deepcopy(break_times)
    bands.append(break_times[0]-0.5*(break_times[1]-break_times[0]))


    for i,break_x in enumerate(break_times[0:-1]):
        bands.append(break_times[i]+0.5*(break_times[i+1]-break_times[i]))
        
    bands.append(break_times[-1]+0.5*(break_times[-1]-break_times[-2]))
    bands = ensure_numeric(bands)
    bands += band_offset
    return bands

def plot_foude_slope_stage(wave_file,
                           break_time,
                           break_x,
                           save_as=None,
                           plot_title="",
                           is_interactive=False,
                           break_type="",
                           froude_min=None,
                           froude_max=None,
                           slope_min=None,
                           slope_max=None):
    """ 
    """
    from pylab import ion, plot, xlabel, ylabel, close, legend, \
         savefig, title, axis, setp, subplot, grid, axvspan
    from anuga.shallow_water.data_manager import csv2dict



    # Load in the csv files and convert info from strings to floats
    simulation, _ = csv2dict(wave_file)
    location = [float(x) for x in simulation['x location']]
    slope = [float(x) for x in simulation['min slope']]
    time = [float(x) for x in simulation['Time']]
    froude = [float(x) for x in simulation['Froude']]

    min_location = min(location)
    max_location = max(location)
    
    if is_interactive:
        ion()
    # The upper subplot
    subplot(311)
    l_froude = plot(location, froude)
    #setp(l_froude, color='r')

    # Add axis stuff
    title(plot_title)
    y_label = "Froude Number"
    ylabel(y_label)
    grid(True)
    axvspan(break_x-0.001,break_x+0.001, facecolor='g')
    if froude_min is not None and froude_max is not None:
        axis(ymin=froude_min, ymax=froude_max)
    
    # The slope subplot
    subplot(312)
    l_slope = plot(location, slope)
    setp(l_slope, color='r')

    # Add axis stuff and legend
    x_label = "X location, m"
    y_label = "Stage slope"
    #xlabel(x_label)
    ylabel(y_label)
    grid(True)
    axvspan(break_x-0.001,break_x+0.001, facecolor='g')
    if slope_min is not None and slope_max is not None:
        axis(ymin=slope_min, ymax=slope_max )

    # The time, x location subplot
    subplot(313)
    l_time = plot(location, time)
    setp(l_time, color='g')
    #print "break_x", break_x
    #print "break_time", break_time
    plot([break_x], [break_time], 'yo')
    #plot([break_x-1], [], 'yo')

    # Add axis stuff and legend
    x_label = "X location, m"
    y_label = "time, sec"
    xlabel(x_label)
    ylabel(y_label)
    grid(True)

    
    # The order defines the label
    #legend((legend_exp, legend_sim),'upper left')
    #legend(('Wave front'),'upper left')
    
    if is_interactive:
        # Wait for enter pressed
        raw_input()

    if save_as is not None:
        savefig(save_as)
    
    #Need to close this plot
    close()
    
#-------------------------------------------------------------
if __name__ == "__main__":
    """ 
    """
    from scenarios import scenarios
    #scenarios = [scenarios[0]]
    outputdir_tag = "_good_tri_area_0.01_limiterD"
    outputdir_tag = "_good_tri_area_0.001_limiterD"
    slope_tag = ""
    #outputdir_tag = "_test_limiterC"
    #scenarios = [scenarios[0]] # !!!!!!!!!!!!!!!!!!!!!!
    #scenarios = scenarios[4:] # !!!!!!!!!!!!!!!!!!!!!!
    
    #gauges_for_slope(slope_tag, outputdir_tag, scenarios)
    #auto_graph_slopes(outputdir_tag, scenarios) #, is_interactive=True) 
    #auto_find_min_slopes(slope_tag, outputdir_tag, scenarios)
    #auto_graph_froudes(outputdir_tag, scenarios) 
    auto_plot_froude_slopes(slope_tag, outputdir_tag, scenarios)
    #g = Get_file_name(scenarios[0], outputdir_tag, slope_tag)
    #for wave_file, save_as, wave_number in Get_file_name(
     #   scenarios[0], outputdir_tag, slope_tag):
      #  print "**************"