Subversion Repositories lagranto.arpege

Rev

Details | Last modification | View Log | RSS feed

Rev Author Line No. Line
3 michaesp 1
#!/usr/bin/env python
2
# -*- coding:utf-8 -*-
3
 
4
 
5
import scipy
6
import netCDF4
7
import datetime
8
import magic
9
import copy
10
 
11
'''
12
TODO :
13
 
14
-Make write_netcdf work when usedatetime=False
15
 
16
 
17
'''
18
 
19
 
20
class Tra(object):
21
    """
22
    Class to work with LAGRANTO output
23
 
24
    Read trajectories from a LAGRANTO file and return a structured numpy array
25
 
26
    filename : file containing lagranto trajectories
27
    return :   structured numpy array traj(ntra,ntime) with variables as tuple.
28
 
29
    Example:
30
        trajs = Tra(filename)
31
        trajs['lon'][0,:] : return the longitudes for the first trajectory.
32
 
33
    Author : Nicolas Piaget, ETH Zurich , 2014
34
             Sebastiaan Crezee, ETH Zurich , 2014
35
    """
36
 
37
    read_format = {
38
        'NetCDF Data Format data': 'netcdf',
39
        'ASCII text': 'ascii',
40
    }
41
 
42
    def __init__(self, filename, typefile=None, usedatetime=True):
43
        """
44
        Read a file based on its type.
45
        By default typefile is determined automatically,
46
        but it can also be specified explicitly
47
        as a string argument:'netcdf' or 'ascii'
48
        """
49
        if typefile is None:
50
            typefile = magic.from_file(filename)
51
            error = "Unkown fileformat. Known formats are {}".format(
52
                    ":".join(self.read_format.keys())
53
            )
54
            typefile = self.read_format[typefile]
55
        else:
56
            error = "Unkown fileformat. Known formats are ascii or netcdf"
57
        try:
58
            function = '_read_{}'.format(typefile)
59
            self._array = globals()[function](filename,
60
                                              usedatetime=usedatetime)
61
        except IOError:
62
            raise IOError(error)
63
 
64
    def __len__(self):
65
        return len(self._array)
66
 
67
    def __getattr__(self, attr):
68
        if attr in self.__dict__:
69
            return getattr(self, attr)
70
        return getattr(self._array, attr)
71
 
72
    def __getitem__(self, key):
73
        return self._array[key]
74
 
75
    def __setitem__(self, key, item):
76
        if type(key) is slice:
77
            self._array = item
78
        else:
79
            formats = self._array.dtype.descr + [(key, item.dtype.descr[0][1])]
80
            formats = [(form[0].encode('ascii', 'ignore'),
81
                       form[1]) for form in formats]
82
 
83
            newarr = scipy.zeros(self._array.shape, dtype=formats)
84
            for var in self.variables:
85
                newarr[var] = self._array[var]
86
            newarr[key] = item
87
            self._array = newarr
88
 
89
    def __repr__(self):
90
        string = " \
91
        {} trajectories with {} time steps. \n \
92
        Available fields: {}\n \
93
        total duration: {} minutes".format(
94
            self.ntra, self.ntime,
95
            "/".join(self.variables),
96
            self.duration
97
        )
98
        return string
99
 
100
    @property
101
    def ntra(self):
102
        if self.ndim < 2:
103
            print(" \
104
                Be careful with the dimensions, \
105
                you may want to change the shape: \n \
106
                either shape + (1,) or (1,)+shape \
107
                ")
108
            return None
109
        return self.shape[0]
110
 
111
    @property
112
    def ntime(self):
113
        if self.ndim < 2:
114
            print(" \
115
                Be careful with the dimensions,\
116
                you may want to change the shape: \n \
117
                either shape + (1,) or (1,)+shape \
118
                ")
119
            return None
120
        return self.shape[1]
121
 
122
    @property
123
    def variables(self):
124
        return list(self.dtype.names)
125
 
126
    @property
127
    def duration(self):
128
        """ time duration in minutes """
129
        end = self['time'][0, -1]
130
        end = end.astype(datetime.datetime)
131
        start = self['time'][0, 0]
132
        start = start.astype(datetime.datetime)
133
        delta = end - start
134
        return delta.total_seconds() / 60.
135
 
136
    @property
137
    def initial(self):
138
        """ give the initial time of the trajectories """
139
        starttime = self['time'][0, 0]
140
        return starttime
141
 
142
    def set_array(self, array):
143
        """ To change the trajectories array """
144
        self._array = array
145
 
146
    def concatenate(self, trajs):
147
        """ To concatenate trajectories together
148
            return Tra object
149
        """
150
        if type(trajs) is not tuple:
151
            trajs = (trajs,)
152
        trajstuple = tuple(tra._array for tra in trajs)
153
        trajstuple += (self._array,)
154
        newtrajs = copy.copy(self)
155
        test = scipy.concatenate(trajstuple)
156
        newtrajs._array = test
157
        return newtrajs
158
 
159
    def write(self, filename, fileformat='netcdf'):
160
        globals()['_write_{}'.format(fileformat)](self, filename)
161
 
162
# -----------------------------------------------------------------------------
163
# read a netcdf lsl file
164
# -----------------------------------------------------------------------------
165
 
166
 
167
def _read_netcdf(filename, usedatetime):
168
    """ Read a netcdf lsl file """
169
 
170
    # read the netcdf, the formats and the variables names
171
 
172
    ncfile = netCDF4.Dataset(filename)
173
    if usedatetime:
174
        formats = [var[1].dtype if var[0] != 'time' else 'datetime64[s]'
175
                   for var in list(ncfile.variables.items())]
176
    else:
177
        formats = [var[1].dtype for var in list(ncfile.variables.items())]
178
 
179
    variables = list(ncfile.variables.keys())
180
 
181
    # set the numbers of trajectories and time step
182
    try:
183
        ntra = len(ncfile.dimensions['dimx_lon'])
184
    except:
185
        try:
186
            ntra = len(ncfile.dimensions['id'])
187
        except:
188
            try:
189
                ntra = len(ncfile.dimensions['ntra'])
190
            except:
191
                raise Exception('Cannot read the number of trajectories,\
192
                                not one of dimx_lon, id or ntra')
193
 
194
    try:
195
        ntime = len(ncfile.dimensions['time'])
196
    except:
197
        ntime = len(ncfile.dimensions['ntim'])
198
 
199
    # change names of the coordinates if necessary
200
 
201
    nvariables = list(variables)
202
 
203
    if "longitude" in variables:
204
        nvariables[variables.index("longitude")] = "lon"
205
 
206
    if "latitude" in variables:
207
        nvariables[variables.index("latitude")] = "lat"
208
 
209
    # create a structured array using numpy
210
 
211
    array = scipy.zeros((ntra, ntime),
212
                        dtype={'names': nvariables, 'formats': formats})
213
 
214
    for variable in variables:
215
        if variable == 'BASEDATE':
216
            continue
217
        nvariable = variable
218
        if variable == 'longitude':
219
            nvariable = 'lon'
220
        if variable == 'latitude':
221
            nvariable = 'lat'
222
        if len(ncfile.variables[variable].shape) == 4:
223
            vardata = ncfile.variables[variable][:, 0, 0, :][:].T
224
        else:
225
            vardata = ncfile.variables[variable][:].T
226
        vardata[vardata < -990] = scipy.nan
227
        array[nvariable] = vardata
228
 
229
    # read the starting date and duration
230
 
231
    try:
232
        date = [int(i) for i in ncfile.variables['BASEDATE'][0, 0, 0, :]]
233
        starttime = datetime.datetime(date[0], date[1],
234
                                      date[2], date[3], date[4])
235
    except:
236
        starttime = datetime.datetime(ncfile.ref_year,
237
                                      ncfile.ref_month,
238
                                      ncfile.ref_day,
239
                                      ncfile.ref_hour,
240
                                      ncfile.ref_min)
241
 
242
    # find characteristic of trajectories
243
    timestep = ncfile.variables['time'][1] - ncfile.variables['time'][0]
244
    period = ncfile.variables['time'][-1] - ncfile.variables['time'][0]
245
 
246
    # True if time = hh.mm
247
    cond1 = 0.00 < (ncfile.variables['time'][:] % 1).max() <= 0.60
248
 
249
    if cond1:
250
        timestep = scipy.floor(timestep) + ((timestep % 1) / .60)
251
        timestep = scipy.around(timestep, 6)
252
        period = scipy.floor(period) + ((period % 1) / .60)
253
        period = scipy.around(period, 6)
254
 
255
    # transform the times into datetime object
256
    # special treatment for online trajectories (time given in minutes)
257
    if usedatetime:
258
        try:
259
            time = scipy.array(
260
                [scipy.datetime64(starttime + datetime.timedelta(hours=t))
261
                 for t in scipy.arange(0, period + timestep, timestep)]
262
            )
263
        except AttributeError:
264
            time = scipy.array(
265
                [scipy.datetime64(starttime + datetime.timedelta(seconds=t))
266
                 for t in scipy.arange(0, period + timestep, timestep)]
267
            )
268
        time.shape = (1,) + time.shape
269
        time = time.repeat(array.shape[0], axis=0)
270
        array['time'] = time
271
 
272
    ncfile.close()
273
    return array
274
 
275
 
276
# ------------------------------------------------------------------------------
277
# read an ASCII lsl file
278
# -----------------------------------------------------------------------------
279
 
280
def _read_ascii(filename, usedatetime, nhead=5):
281
    """ Read a lsl file from ASCII """
282
    # open the file
283
    open_file = open(filename, 'r')
284
 
285
    # get the header
286
    file_lines = open_file.readlines()
287
    nvariables = file_lines[2].strip().split()
288
    head = file_lines[0].split()
289
 
290
    # read starting time new and old formats
291
    try:
292
        starttime = datetime.datetime.strptime(head[2], '%Y%m%d_%H%M')
293
    except ValueError:
294
        try:
295
            starttime = datetime.datetime.strptime(
296
                head[2] + '_' + head[3], '%Y%m%d_%H'
297
            )
298
        except:
299
            print("Warning: could not retrieve starttime from header,\
300
                  setting to default value ")
301
            starttime = datetime.datetime(1970, 1, 1)
302
    if usedatetime:
303
        dtypes = ['f8' if var != 'time' else datetime.datetime
304
                  for var in nvariables]
305
    else:
306
        dtypes = ['f8' for var in nvariables]
307
 
308
    # read the content as numpy array
309
    array = scipy.genfromtxt(filename,
310
                             dtype=dtypes,
311
                             names=nvariables,
312
                             skip_header=nhead,
313
                             missing_values=-999.99)
314
 
315
    # find characteristic of trajectories
316
    timestep = float(array[1][0]) - float(array[0][0])
317
    period = float(array[-1][0]) - float(array[0][0])
318
 
319
    # Convert minutes to decimal hours
320
    if max((array['time'].astype(float)) % 1) <= 0.60:
321
        timestep = scipy.floor(timestep) + ((timestep % 1) / .60)
322
        period = scipy.floor(period) + ((period % 1) / .60)
323
 
324
    # period/timestep gives strange offset (related to precision??)
325
    # so use scipy.around
326
    ntime = int(1 + scipy.around(period / timestep))
327
    ntra = int(array.size / ntime)
328
 
329
    # reshape traj file
330
    array = scipy.reshape(array, (ntra, ntime))
331
 
332
    if usedatetime:
333
        # change time into datetime object
334
        time = scipy.array(
335
            [scipy.datetime64(starttime + datetime.timedelta(hours=t))
336
             for t in scipy.arange(0, period + timestep, timestep)]
337
        )
338
 
339
        time.shape = (1,) + time.shape
340
        time = time.repeat(array.shape[0], axis=0)
341
        array['time'] = 0
342
        newdtypes = [descr if descr[0] != 'time' else ('time', 'datetime64[s]')
343
                     for descr in array.dtype.descr]
344
        array = array.astype(newdtypes)
345
        array['time'] = time
346
 
347
    return array
348
 
349
# ------------------------------------------------------------------------------
350
# write trajectories to a netcdf file
351
# ------------------------------------------------------------------------------
352
 
353
 
354
def _write_netcdf(Tra, filename):
355
 
356
    ncfile = netCDF4.Dataset(filename, 'w', format='NETCDF3_CLASSIC')
357
    ncfile.createDimension('ntra', Tra.ntra)
358
    ncfile.createDimension('ntim', Tra.ntime)
359
    ncfile.ref_year, ncfile.ref_month, ncfile.ref_day, ncfile.ref_hour,\
360
        ncfile.ref_min = Tra.initial.astype(datetime.datetime).timetuple()[0:5]
361
 
362
    vararray = ncfile.createVariable('time', 'f4', ('ntim', ))
363
    delta = Tra['time'][0, :] - Tra.initial
364
    time = [int(a.astype(datetime.datetime).total_seconds() / 3600) +
365
            (a.astype(datetime.datetime).total_seconds() -
366
             int(a.astype(datetime.datetime).total_seconds() / 3600) * 3600)
367
            / 60 * 0.01 for a in delta]
368
    vararray[:] = time
369
 
370
    # add variables
371
    for var in Tra.variables:
372
        if var == 'time':
373
            continue
374
        vararray = ncfile.createVariable(var, Tra[var].dtype, ('ntim', 'ntra'))
375
        vararray[:] = Tra[var].T
376
 
377
    ncfile.close()