source: trunk/UTIL/PYTHON/planetoplot_v2/ppcompute.py @ 1280

Last change on this file since 1280 was 1279, checked in by aslmd, 11 years ago

PLANETOPLOT. added vorticity and divergence computations. corrected perturbation calculations which was not working in previous commit. added examples of dynamical analysis. simplified handling of save filename and extension. simplified handling of projections (plus changed a few settings). added a makesave method to ppplot objects.

File size: 12.5 KB
Line 
1###############################################
2## PLANETOPLOT                               ##
3## --> PPCOMPUTE                             ##
4###############################################
5# python built-in librairies
6import os
7# added librairies
8import numpy as np
9import math  as m
10import scipy.signal as sp_signal
11###############################################
12
13## first a useful function to find settings in a folder in PYTHONPATH
14def findset(whereset,string="planetoplot_v2"):
15    # ... set a default whereset if it was set to None
16    # ... default is in the planetoplot_v2 folder
17    if whereset is None:
18        for path in os.environ['PYTHONPATH'].split(os.pathsep):
19            if string in path: whereset = path
20        if whereset is None: 
21            print "!! ERROR !! "+ string + "not in $PYTHONPATH"
22            print "--> either put it in $PYTHONPATH or change whereset"
23            exit()
24    # ... if the last / is missing put it
25    if whereset[-1] != "/": whereset = whereset + "/"
26    return whereset
27
28##########################
29#### MAX MEAN MIN SUM ####
30#####################################
31#### WITH SUPPORT FOR NaN VALUES ####
32#####################################
33
34## compute min
35## author AS + AC
36def min(field,axis=None): 
37        if field is None: return None
38        if type(field).__name__=='MaskedArray':
39              field.set_fill_value(np.NaN)
40              return np.ma.array(field).min(axis=axis)
41        elif (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
42              return np.ma.masked_invalid(field).min(axis=axis)
43        else: return np.array(field).min(axis=axis)
44
45## compute max
46## author AS + AC
47def max(field,axis=None):
48        if field is None: return None
49        if type(field).__name__=='MaskedArray':
50              field.set_fill_value(np.NaN)
51              return np.ma.array(field).max(axis=axis)
52        elif (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
53              return np.ma.masked_invalid(field).max(axis=axis) 
54        else: return np.array(field).max(axis=axis)
55
56## compute mean
57## author AS + AC
58def mean(field,axis=None):
59        if field is None: return None
60        else: 
61           if type(field).__name__=='MaskedArray':
62              field.set_fill_value(np.NaN)
63              zout=np.ma.array(field).mean(axis=axis)
64              if axis is not None:
65                 zout.set_fill_value(np.NaN)
66                 return zout.filled()
67              else:return zout
68           elif (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
69              zout=np.ma.masked_invalid(field).mean(axis=axis)
70              if axis is not None:
71                 zout.set_fill_value([np.NaN])
72                 return zout.filled()
73              else:return zout
74           else: 
75              return np.array(field).mean(axis=axis)
76
77## compute sum
78## author AS + AC
79def sum(field,axis=None):
80        if field is None: return None
81        else:
82           if type(field).__name__=='MaskedArray':
83              field.set_fill_value(np.NaN)
84              zout=np.ma.array(field).sum(axis=axis)
85              if axis is not None:
86                 zout.set_fill_value(np.NaN)
87                 return zout.filled()
88              else:return zout
89           elif (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
90              zout=np.ma.masked_invalid(field).sum(axis=axis)
91              if axis is not None:
92                 zout.set_fill_value([np.NaN])
93                 return zout.filled()
94              else:return zout
95           else:
96              return np.array(field).sum(axis=axis)
97
98## compute mean over bins
99## author AS
100def meanbin(y,x,bins):
101    bins.sort() # sorting is needed for binning
102    meanvalues = []
103    for iii in range(len(bins)):
104        ## GET VALUES OVER RELEVANT INTERVALS
105        if iii == 0:
106          ind = x<bins[iii+1] ; ys = y[ind]
107        elif iii == len(bins)-1:
108          ind = x>=bins[iii-1] ; ys = y[ind]
109        else:
110          ind = x>=bins[iii-1] ; interm = x[ind] ; intermf = y[ind]
111          ind = interm<bins[iii+1] ; xs = interm[ind] ; ys = intermf[ind]
112        ## COMPUTE MEAN and TREAT NAN CASE
113        meanvalues.append(mean(ys))
114    ## RETURN A NUMPY ARRAY
115    meanvalues = np.array(meanvalues)
116    return meanvalues
117
118## compute perturbation to mean
119## -- the whole dimension must exist!
120## author AS
121def perturbation(field,axis=None):
122    # calculate mean (averaged dim is reduced)
123    mm = mean(field,axis=axis)
124    # include back the reduced dim
125    if field.ndim == 4:
126      if axis == 0:   mm = mm[np.newaxis,:,:,:]
127      elif axis == 1: mm = mm[:,np.newaxis,:,:]
128      elif axis == 2: mm = mm[:,:,np.newaxis,:]
129      elif axis == 3: mm = mm[:,:,:,np.newaxis]
130    else:
131      print "not supported. yet!"
132      exit()
133    # repeat the calculated mean on this dimension
134    mm = np.repeat(mm,field.shape[axis],axis=axis)
135    # compute perturbations
136    field = field - mm
137    return field
138
139################
140#### SMOOTH ####
141################
142### TBD: works with missing values
143
144def smooth2diter(field,n=1):
145        count = 0
146        result = field
147        while count < n:
148            result = smooth2d(result, window=2)
149            count = count + 1
150        return result
151
152## Author: AS. uses gauss_kern and blur_image.
153def smooth2d(field, window=10):
154        ## actually blur_image could work with different coeff on x and y
155        if True in np.isnan(field):
156            print "!! ERROR !! Smooth is a disaster with missing values. This will be fixed."
157            exit()
158        if window > 1:  result = blur_image(field,int(window))
159        else:           result = field
160        return result
161
162## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
163def gauss_kern(size, sizey=None):
164        # Returns a normalized 2D gauss kernel array for convolutions
165        size = int(size)
166        if not sizey:
167                sizey = size
168        else:
169                sizey = int(sizey)
170        x, y = np.mgrid[-size:size+1, -sizey:sizey+1]
171        g = np.exp(-(x**2/float(size)+y**2/float(sizey)))
172        return g / g.sum()
173
174## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
175def blur_image(im, n, ny=None) :
176        # blurs the image by convolving with a gaussian kernel of typical size n.
177        # The optional keyword argument ny allows for a different size in the y direction.
178        g = gauss_kern(n, sizey=ny)
179        improc = sp_signal.convolve(im, g, mode='same')
180        return improc
181
182## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
183def smooth1d(x,window=11,window_type='hanning'):
184    """smooth the data using a window with requested size.
185    This method is based on the convolution of a scaled window with the signal.
186    The signal is prepared by introducing reflected copies of the signal
187    (with the window size) in both ends so that transient parts are minimized
188    in the begining and end part of the output signal.
189    input:
190        x: the input signal
191        window: the dimension of the smoothing window; should be an odd integer
192        window_type: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
193            flat window will produce a moving average smoothing.
194    output:
195        the smoothed signal
196    example:
197    t=linspace(-2,2,0.1)
198    x=sin(t)+randn(len(t))*0.1
199    y=smooth(x)
200    see also:
201    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
202    scipy.signal.lfilter
203    TODO: the window parameter could be the window itself if an array instead of a string   
204    """
205    if True in np.isnan(field):
206        print "!! ERROR !! Smooth is a disaster with missing values. This will be fixed."
207        exit()   
208    x = np.array(x)
209    if x.ndim != 1:
210        raise ValueError, "smooth only accepts 1 dimension arrays."
211    if x.size < window:
212        raise ValueError, "Input vector needs to be bigger than window size."
213    if window<3:
214        return x
215    if not window_type in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
216        raise ValueError, "Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'"
217    s=np.r_[x[window-1:0:-1],x,x[-1:-window:-1]]
218    #print(len(s))
219    if window == 'flat': #moving average
220        w=np.ones(window,'d')
221    else:
222        w=eval('np.'+window_type+'(window)')
223    y=np.convolve(w/w.sum(),s,mode='valid')
224    return y
225
226########################
227#### TIME CONVERTER ####
228########################
229
230#  mars_sol2ls
231#  author T. Navarro
232#  -----------------
233#  convert a given martian day number (sol)
234#  into corresponding solar longitude, Ls (in degr.),
235#  where sol=0=Ls=0 is the
236#  northern hemisphere spring equinox.
237#  -----------------
238def mars_sol2ls(soltabin,forcecontinuity=True):
239      year_day  = 668.6
240      peri_day  = 485.35
241      e_elips   = 0.09340
242      radtodeg  = 57.2957795130823
243      timeperi  = 1.90258341759902
244      if type(soltabin).__name__ in ['int','float','float32','float64']:
245        soltab=[soltabin]
246        solout=np.zeros([1])
247      else:
248        soltab=soltabin
249        solout=np.zeros([len(soltab)])
250      i=0
251      for sol in soltab:
252         zz=(sol-peri_day)/year_day
253         zanom=2.*np.pi*(zz-np.floor(zz))
254         xref=np.abs(zanom)
255#  The equation zx0 - e * sin (zx0) = xref, solved by Newton
256         zx0=xref+e_elips*m.sin(xref)
257         iter=0
258         while iter <= 10:
259            iter=iter+1
260            zdx=-(zx0-e_elips*m.sin(zx0)-xref)/(1.-e_elips*m.cos(zx0))
261            if(np.abs(zdx) <= (1.e-7)):
262              continue
263            zx0=zx0+zdx
264         zx0=zx0+zdx
265         if(zanom < 0.): zx0=-zx0
266# compute true anomaly zteta, now that eccentric anomaly zx0 is known
267         zteta=2.*m.atan(m.sqrt((1.+e_elips)/(1.-e_elips))*m.tan(zx0/2.))
268# compute Ls
269         ls=zteta-timeperi
270         if(ls < 0.): ls=ls+2.*np.pi
271         if(ls > 2.*np.pi): ls=ls-2.*np.pi
272# convert Ls in deg.
273         ls=radtodeg*ls
274         solout[i]=ls
275         i=i+1
276      if forcecontinuity:
277         for iii in range(len(soltab)-1):
278           while solout[iii+1] - solout[iii] > 180. :  solout[iii+1] = solout[iii+1] - 360.
279           while solout[iii] - solout[iii+1] > 180. :  solout[iii+1] = solout[iii+1] + 360.
280      return solout
281
282# mars_date
283# author A. Spiga
284# ------------------------------
285# get Ls, sol, utc from a string with format yyyy-mm-dd_hh:00:00
286# -- argument timechar is a vector of such strings indicating dates
287# -- example: timechar = nc.variables['Times'][:] in mesoscale files
288# NB: uses mars_sol2ls function above
289# ------------------------------
290def mars_date(timechar):
291    # some preliminary information
292    days_in_month = [61, 66, 66, 65, 60, 54, 50, 46, 47, 47, 51, 56]
293    plus_in_month = [ 0, 61,127,193,258,318,372,422,468,515,562,613]
294    # get utc and sol from strings
295    utc = [] ; sol = []
296    for zetime in timechar:
297        dautc = float(zetime[11]+zetime[12]) + float(zetime[14]+zetime[15])/37.
298        dasol = dautc / 24.
299        dasol = dasol + plus_in_month[ int(zetime[5]+zetime[6])-1 ] + int(zetime[8]+zetime[9])
300        dasol = dasol - 1 ##les sols GCM commencent a 0
301        utc.append(dautc)
302        sol.append(dasol)
303    sol = np.array(sol)
304    utc = np.array(utc)
305    # get ls from sol
306    ls = mars_sol2ls(sol)
307    return ls, sol, utc
308
309# timecorrect
310# author A. Spiga
311# -----------------------------
312# ensure time axis is monotonic
313# correct negative values
314# -----------------------------
315def timecorrect(time):
316    for ind in range(len(time)-1):
317        if time[ind] < 0.: time[ind] = time[ind] + 24.
318        if time[ind+1] < time[ind]: time[ind+1] = time[ind+1] + 24.
319    return time
320
321#########################
322#### FLOW DIAGNOSTIC ####
323#########################
324
325def divort(u,v,lon,lat,rad):
326  ####################
327  # compute divergence and vorticity
328  # div,vorti = divort(u,v,lon,lat,rad)
329  ####################
330  # u: zonal wind
331  # v: meridional wind
332  # lon: longitude
333  # lat: latitude
334  # rad: radius of the planet
335  ####################
336
337  # compute gradients -- with normalized coordinates
338  du_y,du_x = np.gradient(u)
339  dv_y,dv_x = np.gradient(v)
340
341  # convert to radian
342  la = lon*np.pi/180.
343  ph = lat*np.pi/180.
344
345  # ensure 2D arrays for lon,lat
346  if lon.ndim == 1 and lat.ndim == 1:
347    [lar,phr] = np.meshgrid(la,ph)
348  elif lon.ndim == 2 and lat.ndim == 2:
349    lar,phr = la,ph
350  else:
351    print "ppcompute. there is a problem with lat/lon rank. stop."
352    exit()
353
354  # compute normalized gradients for lat/lon grid
355  dump,dla_x = np.gradient(lar)
356  dph_y,dump = np.gradient(phr)
357   
358  # compute cartesian differential coordinates
359  dx = dla_x*rad*np.cos(phr)
360  dy = dph_y*rad
361   
362  # eventually compute cartesian derivatives
363  du_dx = du_x / dx
364  du_dy = du_y / dy
365  dv_dx = dv_x / dx
366  dv_dy = dv_y / dy
367   
368  # vorticity
369  vorti = dv_dx - du_dy
370  div = du_dx + dv_dy
371
372  return div,vorti
373
374
Note: See TracBrowser for help on using the repository browser.