source: development/stochastic_study/read_realisations.py @ 3009

"""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 caching import cache

def read_realisations(subdir, max_realisations = None, gauge_number=0,
                      exclude=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},
                       verbose=verbose)
    else:
        result = _read_realisations(subdir,
                                    max_realisations = max_realisations,
                                    gauge_number=gauge_number,
                                    exclude=exclude)

    return result

def _read_realisations(subdir, max_realisations = None, gauge_number=0,
                       exclude=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)
    
    # 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