Subversion Repositories lagranto.ecmwf

#!/usr/bin/env python
# -*- coding:utf-8 -*-


import scipy
import netCDF4
import datetime
import magic
import copy

'''
TODO :

-Make write_netcdf work when usedatetime=False


'''


class Tra(object):
    """
    Class to work with LAGRANTO output

    Read trajectories from a LAGRANTO file and return a structured numpy array

    filename : file containing lagranto trajectories
    return :   structured numpy array traj(ntra,ntime) with variables as tuple.

    Example:
        trajs = Tra(filename)
        trajs['lon'][0,:] : return the longitudes for the first trajectory.

    Author : Nicolas Piaget, ETH Zurich , 2014
             Sebastiaan Crezee, ETH Zurich , 2014
    """

    read_format = {
        'NetCDF Data Format data': 'netcdf',
        'ASCII text': 'ascii',
    }

    def __init__(self, filename, typefile=None, usedatetime=True):
        """
        Read a file based on its type.
        By default typefile is determined automatically,
        but it can also be specified explicitly
        as a string argument:'netcdf' or 'ascii'
        """
        if typefile is None:
            typefile = magic.from_file(filename)
            error = "Unkown fileformat. Known formats are {}".format(
                    ":".join(self.read_format.keys())
            )
            typefile = self.read_format[typefile]
        else:
            error = "Unkown fileformat. Known formats are ascii or netcdf"
        try:
            function = '_read_{}'.format(typefile)
            self._array = globals()[function](filename,
                                              usedatetime=usedatetime)
        except IOError:
            raise IOError(error)

    def __len__(self):
        return len(self._array)

    def __getattr__(self, attr):
        if attr in self.__dict__:
            return getattr(self, attr)
        return getattr(self._array, attr)

    def __getitem__(self, key):
        return self._array[key]

    def __setitem__(self, key, item):
        if type(key) is slice:
            self._array = item
        else:
            formats = self._array.dtype.descr + [(key, item.dtype.descr[0][1])]
            formats = [(form[0].encode('ascii', 'ignore'),
                       form[1]) for form in formats]

            newarr = scipy.zeros(self._array.shape, dtype=formats)
            for var in self.variables:
                newarr[var] = self._array[var]
            newarr[key] = item
            self._array = newarr

    def __repr__(self):
        string = " \
        {} trajectories with {} time steps. \n \
        Available fields: {}\n \
        total duration: {} minutes".format(
            self.ntra, self.ntime,
            "/".join(self.variables),
            self.duration
        )
        return string

    @property
    def ntra(self):
        if self.ndim < 2:
            print(" \
                Be careful with the dimensions, \
                you may want to change the shape: \n \
                either shape + (1,) or (1,)+shape \
                ")
            return None
        return self.shape[0]

    @property
    def ntime(self):
        if self.ndim < 2:
            print(" \
                Be careful with the dimensions,\
                you may want to change the shape: \n \
                either shape + (1,) or (1,)+shape \
                ")
            return None
        return self.shape[1]

    @property
    def variables(self):
        return list(self.dtype.names)

    @property
    def duration(self):
        """ time duration in minutes """
        end = self['time'][0, -1]
        end = end.astype(datetime.datetime)
        start = self['time'][0, 0]
        start = start.astype(datetime.datetime)
        delta = end - start
        return delta.total_seconds() / 60.

    @property
    def initial(self):
        """ give the initial time of the trajectories """
        starttime = self['time'][0, 0]
        return starttime

    def set_array(self, array):
        """ To change the trajectories array """
        self._array = array

    def concatenate(self, trajs):
        """ To concatenate trajectories together
            return Tra object
        """
        if type(trajs) is not tuple:
            trajs = (trajs,)
        trajstuple = tuple(tra._array for tra in trajs)
        trajstuple += (self._array,)
        newtrajs = copy.copy(self)
        test = scipy.concatenate(trajstuple)
        newtrajs._array = test
        return newtrajs

    def write(self, filename, fileformat='netcdf'):
        globals()['_write_{}'.format(fileformat)](self, filename)

# -----------------------------------------------------------------------------
# read a netcdf lsl file
# -----------------------------------------------------------------------------


def _read_netcdf(filename, usedatetime):
    """ Read a netcdf lsl file """

    # read the netcdf, the formats and the variables names

    ncfile = netCDF4.Dataset(filename)
    if usedatetime:
        formats = [var[1].dtype if var[0] != 'time' else 'datetime64[s]'
                   for var in list(ncfile.variables.items())]
    else:
        formats = [var[1].dtype for var in list(ncfile.variables.items())]

    variables = list(ncfile.variables.keys())

    # set the numbers of trajectories and time step
    try:
        ntra = len(ncfile.dimensions['dimx_lon'])
    except:
        try:
            ntra = len(ncfile.dimensions['id'])
        except:
            try:
                ntra = len(ncfile.dimensions['ntra'])
            except:
                raise Exception('Cannot read the number of trajectories,\
                                not one of dimx_lon, id or ntra')

    try:
        ntime = len(ncfile.dimensions['time'])
    except:
        ntime = len(ncfile.dimensions['ntim'])

    # change names of the coordinates if necessary

    nvariables = list(variables)

    if "longitude" in variables:
        nvariables[variables.index("longitude")] = "lon"

    if "latitude" in variables:
        nvariables[variables.index("latitude")] = "lat"

    # create a structured array using numpy

    array = scipy.zeros((ntra, ntime),
                        dtype={'names': nvariables, 'formats': formats})

    for variable in variables:
        if variable == 'BASEDATE':
            continue
        nvariable = variable
        if variable == 'longitude':
            nvariable = 'lon'
        if variable == 'latitude':
            nvariable = 'lat'
        if len(ncfile.variables[variable].shape) == 4:
            vardata = ncfile.variables[variable][:, 0, 0, :][:].T
        else:
            vardata = ncfile.variables[variable][:].T
        vardata[vardata < -990] = scipy.nan
        array[nvariable] = vardata

    # read the starting date and duration

    try:
        date = [int(i) for i in ncfile.variables['BASEDATE'][0, 0, 0, :]]
        starttime = datetime.datetime(date[0], date[1],
                                      date[2], date[3], date[4])
    except:
        starttime = datetime.datetime(ncfile.ref_year,
                                      ncfile.ref_month,
                                      ncfile.ref_day,
                                      ncfile.ref_hour,
                                      ncfile.ref_min)

    # find characteristic of trajectories
    timestep = ncfile.variables['time'][1] - ncfile.variables['time'][0]
    period = ncfile.variables['time'][-1] - ncfile.variables['time'][0]

    # True if time = hh.mm
    cond1 = 0.00 < (ncfile.variables['time'][:] % 1).max() <= 0.60

    if cond1:
        timestep = scipy.floor(timestep) + ((timestep % 1) / .60)
        timestep = scipy.around(timestep, 6)
        period = scipy.floor(period) + ((period % 1) / .60)
        period = scipy.around(period, 6)

    # transform the times into datetime object
    # special treatment for online trajectories (time given in minutes)
    if usedatetime:
        try:
            time = scipy.array(
                [scipy.datetime64(starttime + datetime.timedelta(hours=t))
                 for t in scipy.arange(0, period + timestep, timestep)]
            )
        except AttributeError:
            time = scipy.array(
                [scipy.datetime64(starttime + datetime.timedelta(seconds=t))
                 for t in scipy.arange(0, period + timestep, timestep)]
            )
        time.shape = (1,) + time.shape
        time = time.repeat(array.shape[0], axis=0)
        array['time'] = time

    ncfile.close()
    return array


# ------------------------------------------------------------------------------
# read an ASCII lsl file
# -----------------------------------------------------------------------------

def _read_ascii(filename, usedatetime, nhead=5):
    """ Read a lsl file from ASCII """
    # open the file
    open_file = open(filename, 'r')

    # get the header
    file_lines = open_file.readlines()
    nvariables = file_lines[2].strip().split()
    head = file_lines[0].split()

    # read starting time new and old formats
    try:
        starttime = datetime.datetime.strptime(head[2], '%Y%m%d_%H%M')
    except ValueError:
        try:
            starttime = datetime.datetime.strptime(
                head[2] + '_' + head[3], '%Y%m%d_%H'
            )
        except:
            print("Warning: could not retrieve starttime from header,\
                  setting to default value ")
            starttime = datetime.datetime(1970, 1, 1)
    if usedatetime:
        dtypes = ['f8' if var != 'time' else datetime.datetime
                  for var in nvariables]
    else:
        dtypes = ['f8' for var in nvariables]

    # read the content as numpy array
    array = scipy.genfromtxt(filename,
                             dtype=dtypes,
                             names=nvariables,
                             skip_header=nhead,
                             missing_values=-999.99)

    # find characteristic of trajectories
    timestep = float(array[1][0]) - float(array[0][0])
    period = float(array[-1][0]) - float(array[0][0])

    # Convert minutes to decimal hours
    if max((array['time'].astype(float)) % 1) <= 0.60:
        timestep = scipy.floor(timestep) + ((timestep % 1) / .60)
        period = scipy.floor(period) + ((period % 1) / .60)

    # period/timestep gives strange offset (related to precision??)
    # so use scipy.around
    ntime = int(1 + scipy.around(period / timestep))
    ntra = int(array.size / ntime)

    # reshape traj file
    array = scipy.reshape(array, (ntra, ntime))

    if usedatetime:
        # change time into datetime object
        time = scipy.array(
            [scipy.datetime64(starttime + datetime.timedelta(hours=t))
             for t in scipy.arange(0, period + timestep, timestep)]
        )

        time.shape = (1,) + time.shape
        time = time.repeat(array.shape[0], axis=0)
        array['time'] = 0
        newdtypes = [descr if descr[0] != 'time' else ('time', 'datetime64[s]')
                     for descr in array.dtype.descr]
        array = array.astype(newdtypes)
        array['time'] = time

    return array

# ------------------------------------------------------------------------------
# write trajectories to a netcdf file
# ------------------------------------------------------------------------------


def _write_netcdf(Tra, filename):

    ncfile = netCDF4.Dataset(filename, 'w', format='NETCDF3_CLASSIC')
    ncfile.createDimension('ntra', Tra.ntra)
    ncfile.createDimension('ntim', Tra.ntime)
    ncfile.ref_year, ncfile.ref_month, ncfile.ref_day, ncfile.ref_hour,\
        ncfile.ref_min = Tra.initial.astype(datetime.datetime).timetuple()[0:5]

    vararray = ncfile.createVariable('time', 'f4', ('ntim', ))
    delta = Tra['time'][0, :] - Tra.initial
    time = [int(a.astype(datetime.datetime).total_seconds() / 3600) +
            (a.astype(datetime.datetime).total_seconds() -
             int(a.astype(datetime.datetime).total_seconds() / 3600) * 3600)
            / 60 * 0.01 for a in delta]
    vararray[:] = time

    # add variables
    for var in Tra.variables:
        if var == 'time':
            continue
        vararray = ncfile.createVariable(var, Tra[var].dtype, ('ntim', 'ntra'))
        vararray[:] = Tra[var].T

    ncfile.close()