source: trunk/anuga_work/development/stochastic_study/read_realisations.py @ 7884

"""Read realisations and return MxN Numeric array: data, where M is the number of timesteps and N the number of realisations
"""

import os
import project
from Numeric import Float, zeros
from RandomArray import randint
from caching import cache

def read_realisations(subdir, max_realisations = None, gauge_number=0,
                      exclude=None,
                      sorting=None,
                      use_cache=False, 
                      verbose=True):
    """Read realisations
    """

    if use_cache is True:
        result = cache(_read_realisations,
                       (subdir,),
                       {'max_realisations': max_realisations,
                        'gauge_number': gauge_number,
                        'exclude': exclude,
                        'sorting': sorting},
                       verbose=verbose)
    else:
        result = _read_realisations(subdir,
                                    max_realisations = max_realisations,
                                    gauge_number=gauge_number,
                                    exclude=exclude,
                                    sorting=sorting)

    return result


def _read_realisations(subdir, max_realisations = None, gauge_number=0,
                       exclude=None, sorting=None,
                       use_cache=False, verbose=True):
    """Read realisations
    """

    
    gauge = '%s.txt' %project.gauge_names[gauge_number]

    if subdir[-1] != os.sep:
        subdir += os.sep
    
    # Establish dimensions and record filenames
    filenames = []
    for filename in os.listdir(project.working_dir + subdir):
        if filename.startswith(project.basename) and filename.endswith(gauge):
            if exclude is not None and filename.find(exclude) >= 0:
                print 'Excluded: %s' %filename
            else:    
                filenames.append(project.working_dir + subdir + filename)

        if max_realisations is not None and len(filenames) == max_realisations:
            break

    time = zeros(project.number_of_timesteps, Float)
    data = zeros((project.number_of_timesteps, len(filenames)), Float)

    if sorting is None:
        pass
    elif sorting == 'numerical':
        filenames = sort_numerically(filenames)
    elif sorting == 'randomised':    
        filenames = randomise(filenames)
    else:
        print 'Invalid value for sorting: %s' %sorting
        import sys; sys.exit() 


    # Read data        

    for j, filename in enumerate(filenames):
        print 'Reading filename %s (column %d)' %(filename, j)
        fid = open(filename)
        for i, line in enumerate(fid.readlines()):
            if i < data.shape[0]:
                fields = line.strip().split()
                time[i] = float(fields[0])
                data[i,j] = float(fields[1])
            
        fid.close()


    return time, data, filenames


def sort_numerically(filenames):
    """Sort filenames numerically
    simulation_realisation_n_ch.... (sort by n)
    """

    print 'Sorting'
    sorted_filenames = []

    numbers = []
    for filename in filenames:
        fields = filename.split('_')
        n = int(fields[2])
        numbers.append(n)

    A = zip(numbers, filenames)
    A.sort()
    
    sorted_filenames = [a[1] for a in A]    
    
    return sorted_filenames


def randomise(filenames):
    """Randomise filenames
    """

    print 'Randomising'
    randomised_filenames = []


    while len(filenames) > 0:
        n = randint(0, len(filenames))
        filename = filenames[n]
        del filenames[n]
        randomised_filenames.append(filename)
     
    return randomised_filenames