source: lmdz_wrf/trunk/tools/sfcVAR_global_modification.py

Last change on this file was 187, checked in by lfita, 10 years ago

Adding scripts for the generation of the aqua-planet initial conditions
File size: 15.6 KB
RevLine 
[187]1# Python program to modify any sfc variable from netCDF file to impose aprofile of a given variable
2# L. Fita Borrell, LMD-Jussieu, IPSL, UMPC, CNRS, Paris, France
3# July 2013
4
5from optparse import OptionParser
6import numpy as np
7from netCDF4 import Dataset as NetCDFFile
8import os
9
10## g.e. # python sfcVAR_global_modification.py -f met_em.d01_1979-01-01_00:00:00 -k control -v SST -s 300.15
11
12fname = 'sfcVAR_global_modification.py'
13
14errormsg = 'ERROR -- error -- ERROR -- error'
15
16kelvin=273.15
17pi3=np.pi/3.
18pi36=np.pi/36.
19
20def HS_kind(Ndims,dimx,dimy,dimt,lonvals,latvals,zero,amp,dl):
21    print '  Imposing "Held & Suarez, 94, BAMS" characteristics...'
22    if Ndims == 1:
23        varvals = np.zeros((dimy,dimx), dtype=np.float32)
24        for ilat in range(dimy):
25            varvals[ilat,:] = amp-dl*np.sin(latvals[ilat])**2.
26    elif Ndims==2:
27        varvals = np.zeros((dimy,dimx), dtype=np.float32)
28        for ilat in range(dimy):
29            for ilon in range(dimx):
30                varvals[ilat,ilon] = amp-dl*np.sin(latvals[ilat,ilon])**2.
31    elif Ndims==3:
32        varvals = np.zeros((dimt,dimy,dimx), dtype=np.float32)
33        for ilat in range(dimy):
34            for ilon in range(dimx):
35                varvals[:,ilat,ilon] = amp-dl*np.sin(latvals[0,ilat,ilon])**2.
36
37    return varvals + zero
38
39
40def control_kind(Ndims,dimx,dimy,dimt,lonvals,latvals,zero,amp):
41    print '  Imposing "control" characteristics...'
42    if Ndims == 1:
43        varvals = np.zeros((dimy,dimx), dtype=np.float32)
44        for ilat in range(dimy):
45            if np.abs(latvals[ilat]) < pi3:
46                varvals[ilat,:] = amp*(1.-np.sin(3.*latvals[ilat]/2.)**2.)
47            else:
48                varvals[ilat,:] = 0.
49    elif Ndims==2:
50        varvals = np.zeros((dimy,dimx), dtype=np.float32)
51        for ilat in range(dimy):
52            for ilon in range(dimx):
53                if np.abs(latvals[ilat,ilon]) < pi3:
54                    varvals[ilat,ilon] = amp*(1.-np.sin(3.*latvals[ilat,ilon]/2.)**2.)
55                else:
56                    varvals[ilat,ilon] = 0.
57    elif Ndims==3:
58        varvals = np.zeros((dimt,dimy,dimx), dtype=np.float32)
59        for ilat in range(dimy):
60            for ilon in range(dimx):
61                if np.abs(latvals[0,ilat,ilon]) < pi3:
62                    varvals[:,ilat,ilon] = amp*(1.-np.sin(3.*latvals[0,ilat,ilon]/2.)**2.)
63                else:
64                    varvals[:,ilat,ilon] = 0.
65
66    return varvals + zero
67
68def control5N_kind(Ndims,dimx,dimy,dimt,lonvals,latvals,zero,amp):
69    print '  Imposing "control5N" characteristics...'
70    if Ndims == 1:
71        varvals = np.zeros((dimy,dimx), dtype=np.float32)
72        for ilat in range(dimy):
73            if pi36 < latvals[ilat] and latvals[ilat] < pi3:
74                varvals[ilat,:] = amp*(1.-np.sin(90./55.*(latvals[ilat]-pi36))**2.)
75            elif -pi3 < latvals[ilat] and latvals[ilat] < pi36:
76                varvals[ilat,:] = amp*(1.-np.sin(90./65.*(latvals[ilat]-pi36))**2.)
77            else:
78                varvals[ilat,:] = 0.
79    elif Ndims==2:
80        varvals = np.zeros((dimy,dimx), dtype=np.float32)
81        for ilat in range(dimy):
82            for ilon in range(dimx):
83                if pi36 <= latvals[ilat,ilon] and latvals[ilat,ilon] < pi3:
84                    varvals[ilat,ilon] = amp*(1.-np.sin(90./55.*(latvals[ilat,ilon]-pi36))**2.)
85                elif -pi3 < latvals[ilat,ilon] and latvals[ilat,ilon] < pi36:
86                    varvals[ilat,ilon] = amp*(1.-np.sin(90./65.*(latvals[ilat,ilon]-pi36))**2.)
87                else:
88                    varvals[ilat,ilon] = 0.
89    elif Ndims==3:
90        varvals = np.zeros((dimt,dimy,dimx), dtype=np.float32)
91        for ilat in range(dimy):
92            for ilon in range(dimx):
93                if pi36 <= latvals[0,ilat,ilon] and latvals[0,ilat,ilon] < pi3:
94                    varvals[:,ilat,ilon] = amp*(1.-np.sin(90./55.*(latvals[0,ilat,ilon]-pi36))**2.)
95                elif -pi3 < latvals[0,ilat,ilon] and latvals[0,ilat,ilon] < pi36:
96                    varvals[:,ilat,ilon] = amp*(1.-np.sin(90./65.*(latvals[0,ilat,ilon]-pi36))**2.)
97                else:
98                    varvals[:,ilat,ilon] = 0.
99
100    return varvals + zero
101
102def flat_kind(Ndims,dimx,dimy,dimt,lonvals,latvals,zero,amp):
103    print '  Imposing "flat" characteristics...'
104    if Ndims == 1:
105        varvals = np.zeros((dimy,dimx), dtype=np.float32)
106        for ilat in range(dimy):
107            if np.abs(latvals[ilat]) < pi3:
108                varvals[ilat,:] = amp*(1.-np.sin(3.*latvals[ilat]/2.)**4.)
109            else:
110                varvals[ilat,:] = 0.
111    elif Ndims==2:
112        varvals = np.zeros((dimy,dimx), dtype=np.float32)
113        for ilat in range(dimy):
114            for ilon in range(dimx):
115                if np.abs(latvals[ilat,ilon]) < pi3:
116                    varvals[ilat,ilon] = amp*(1.-np.sin(3.*latvals[ilat,ilon]/2.)**4.)
117                else:
118                    varvals[ilat,ilon] = 0.
119    elif Ndims==3:
120        varvals = np.zeros((dimt,dimy,dimx), dtype=np.float32)
121        for ilat in range(dimy):
122            for ilon in range(dimx):
123                if np.abs(latvals[0,ilat,ilon]) < pi3:
124                    varvals[:,ilat,ilon] = amp*(1.-np.sin(3.*latvals[0,ilat,ilon]/2.)**4.)
125                else:
126                    varvals[:,ilat,ilon] = 0.
127
128    return varvals + zero
129
130def KEQ_kind(Ndims,dimx,dimy,dimt,lonvals,latvals,zero,amp,L0,dL,dl):
131    print '  Imposing "KEQ" characteristics...'
132    npoints=0
133    if Ndims == 1:
134        varvals = np.zeros((dimy,dimx), dtype=np.float32)
135        for ilat in range(dimy):
136            for ilon in range(dimx):
137                if L0-dL < lonvals[ilon] and lonvals[ilon] < L0+dL and np.abs(latvals[ilat]) < dl:
138                    npoints = npoints + 1
139                    A=np.cos(np.pi/2.*(lonvals[ilon]-L0)/dL)**2.
140                    B=np.cos(np.pi/2.*latvals[ilat]/dl)**2.
141                    varvals[ilat,ilon] = amp*A*B
142                else:
143                    varvals[ilat,ilon] = 0.
144    elif Ndims==2:
145        varvals = np.zeros((dimy,dimx), dtype=np.float32)
146        for ilat in range(dimy):
147            for ilon in range(dimx):
148                if L0-dL < lonvals[ilat,ilon] and lonvals[ilat,ilon] < L0+dL and np.abs(latvals[ilat,ilon]) < dl:
149                    npoints = npoints + 1
150                    A=np.cos(np.pi/2.*(lonvals[ilat,ilon]-L0)/dL)**2.
151                    B=np.cos(np.pi/2.*latvals[ilat,ilon]/dl)**2.
152                    varvals[ilat,ilon] = amp*A*B
153                else:
154                    varvals[ilat,ilon] = 0.
155    elif Ndims==3:
156        varvals = np.zeros((dimt,dimy,dimx), dtype=np.float32)
157        for ilat in range(dimy):
158            for ilon in range(dimx):
159                if L0-dL < lonvals[0,ilat,ilon] and lonvals[0,ilat,ilon] < L0+dL and np.abs(latvals[0,ilat,ilon]) < dl:
160                    npoints = npoints + 1
161                    A=np.cos(np.pi/2.*(lonvals[0,ilat,ilon]-L0)/dL)**2.
162                    B=np.cos(np.pi/2.*latvals[0,ilat,ilon]/dl)**2.
163                    varvals[:,ilat,ilon] = amp*A*B
164                else:
165                    varvals[:,ilat,ilon] = 0.
166
167    if npoints < 1:
168       print errormsg
169       print '    KEQ: not modified points generated!!!!'
170       print '       L0-dL: ',L0-dL
171       print '       L0+dL: ',L0+dL
172       print '       dl:', dl
173       print errormsg
174       quit(-1)
175
176    return varvals
177
178def KW1_kind(Ndims,dimx,dimy,dimt,lonvals,latvals,zero,amp,L0,dL,dl):
179    print '  Imposing "KW1 characteristics...'
180    npoints=0
181    if Ndims == 1:
182        varvals = np.zeros((dimy,dimx), dtype=np.float32)
183        for ilat in range(dimy):
184            for ilon in range(dimx):
185                if np.abs(latvals[ilat]) < dl:
186                    npoints = npoints + 1
187                    A=np.cos(lonvals[ilon]-L0)
188                    B=np.cos(np.pi/2.*(latvals[ilat])/dl)**2.
189                    varvals[ilat,ilon] =  amp*A*B
190                else:
191                    varvals[ilat,ilon] = 0.
192    elif Ndims==2:
193        varvals = np.zeros((dimy,dimx), dtype=np.float32)
194        for ilat in range(dimy):
195            for ilon in range(dimx):
196                if np.abs(latvals[ilat,ilon]) < dl:
197                    npoints = npoints + 1
198                    A=np.cos(lonvals[ilat,ilon]-L0)
199                    B=np.cos(np.pi/2.*(latvals[ilat,ilon])/dl)**2.
200                    varvals[ilat,ilon] =  amp*A*B
201                else:
202                    varvals[ilat,ilon] = 0.
203    elif Ndims==3:
204        varvals = np.zeros((dimt,dimy,dimx), dtype=np.float32)
205        for ilat in range(dimy):
206            for ilon in range(dimx):
207                if np.abs(latvals[0,ilat,ilon]) < dl:
208                    npoints = npoints + 1
209                    A=np.cos(lonvals[0,ilat,ilon]-L0)
210                    B=np.cos(np.pi/2.*(latvals[0,ilat,ilon])/dl)**2.
211                    varvals[:,ilat,ilon] = amp*A*B
212                else:
213                    varvals[:,ilat,ilon] = 0.
214
215    if npoints < 1:
216       print errormsg
217       print '    KW1: not modified points generated!!!!'
218       print '       dl:', dl
219       print errormsg
220       quit(-1)
221
222    return varvals
223
224def peaked_kind(Ndims,dimx,dimy,dimt,lonvals,latvals,zero,amp):
225    print '  Imposing "peaked" characteristics...'
226    if Ndims == 1:
227        varvals = np.zeros((dimy,dimx), dtype=np.float32)
228        for ilat in range(dimy):
229            if np.abs(latvals[ilat]) < pi3:
230                varvals[ilat,:] = amp*(1.-(3.*np.abs(latvals[ilat])/np.pi))
231            else:
232                varvals[ilat,:] = 0.
233    elif Ndims==2:
234        varvals = np.zeros((dimy,dimx), dtype=np.float32)
235        for ilat in range(dimy):
236            for ilon in range(dimx):
237                if np.abs(latvals[ilat,ilon]) < pi3:
238                    varvals[ilat,ilon] = amp*(1.-(3.*np.abs(latvals[ilat,ilon])/np.pi))
239                else:
240                    varvals[ilat,ilon] = 0.
241    elif Ndims==3:
242        varvals = np.zeros((dimt,dimy,dimx), dtype=np.float32)
243        for ilat in range(dimy):
244            for ilon in range(dimx):
245                if np.abs(latvals[0,ilat,ilon]) < pi3:
246                    varvals[:,ilat,ilon] = amp*(1.-(3.*np.abs(latvals[0,ilat,ilon])/np.pi))
247                else:
248                    varvals[:,ilat,ilon] = 0.
249
250    return varvals + zero
251
252####### ###### ##### #### ### ## #
253
254parser = OptionParser()
255parser.add_option("-f", "--netCDF_file", dest="ncfile", 
256                  help="file to use", metavar="FILE")
257parser.add_option("-k", "--kindMOD", type='choice', dest="modkind",
258 choices=['control', 'control5N', 'flat', 'peaked', 'Qobs', 'KEQ', 'KW1', 'HS'], 
259                  help="kind of modification from Neale and Hoskins, Atmos. Sci. Letters, (2001), HS: Held & Suarez, 1994 conditions (uses [ref]:[amp]:[perAmp, as /\Ty])", metavar="VALUE")
260parser.add_option("-L", "--longitude", dest="lonname",
261                  help="name of the longitude", metavar="VAR")
262parser.add_option("-l", "--latitude", dest="latname",
263                  help="name of the latitude", metavar="VAR")
264parser.add_option("-v", "--variable", dest="varname",
265                  help="variable to modify", metavar="VAR")
266parser.add_option("-s", "--values", dest="values",
267                  help="aditional values for 'KEQ' and 'KW1' [ref, in [units]]:[amp, with respect 'ref' in [units]]:[perAmp, in [units]]:[L0, in deg]:[dL, in deg]:[dl, in deg]", metavar="VAR")
268
269(opts, args) = parser.parse_args()
270
271#######    #######
272## MAIN
273    #######
274
275if not os.path.isfile(opts.ncfile):
276  print errormsg
277  print '  File ' + opts.ncfile + ' does not exist !!'
278  print errormsg
279  quit()   
280
281ncf = NetCDFFile(opts.ncfile,'a')
282
283if not ncf.variables.has_key(opts.varname):
284      print errormsg
285      print '   ' + fname + ': File "' + opts.ncfile + '" does not have variable "' + opts.varname + '" !!!!'
286      print errormsg
287      ncf.close()
288      quit(-1)
289
290additionalvalues = opts.values.split(':')
291Nvalues = len(additionalvalues)
292if opts.modkind == 'KEQ' or opts.modkind == 'KW1':
293    if Nvalues != 6:
294        print errormsg
295        print '  ' + fname + ': options "KEQ" and "KW1" require 4 additional arguments! ',Nvalues,' provided!'
296        print errormsg
297        ncf.close()
298        quit(-1)
299    else:
300        refval = np.float32(additionalvalues[0])
301        amplitude = np.float32(additionalvalues[1])
302        Peramplitude = np.float32(additionalvalues[2])
303        Lon0 = np.float32(additionalvalues[3])*np.pi/180.
304        dLon = np.float32(additionalvalues[4])*np.pi/180.
305        dlat = np.float32(additionalvalues[5])*np.pi/180.
306        print '  ' + fname +': additional values____'
307        print '    ref: ',refval
308        print '    amp: ',amplitude
309        print '    perAmp: ',Peramplitude
310        print '    L0: ',Lon0
311        print '    dL: ',dLon
312        print '    dl: ',dlat
313elif opts.modkind == 'HS':
314    if Nvalues != 3:
315        print errormsg
316        print '  ' + fname + ': option "HS" requires 3 additional arguments! ',Nvalues,' provided!'
317        print errormsg
318        ncf.close()
319        quit(-1)
320    else:
321        refval = np.float32(additionalvalues[0])
322        amplitude = np.float32(additionalvalues[1])
323        Peramplitude = np.float32(additionalvalues[2])
324        print '  ' + fname +': additional values____'
325        print '    ref: ',refval
326        print '    amp: ',amplitude
327        print '    perAmp: ',Peramplitude
328else:
329    refval = np.float32(additionalvalues[0])
330    amplitude = np.float32(additionalvalues[1])
331    print '  ' + fname +': values____'
332    print '    ref: ',refval
333    print '    amp: ',amplitude
334
335objvarmod = ncf.variables[opts.varname]
336objlon = ncf.variables[opts.lonname]
337objlat = ncf.variables[opts.latname]
338
339Nds = len(objvarmod.shape)
340if Nds == 1:
341    dx = objlon.shape[0]
342    dy = objlat.shape[0]
343else:
344    dx = objlon.shape[Nds-1]
345    dy = objlon.shape[Nds-2]
346    if Nds == 3:
347        dt = objvarmod.shape[0]
348    else:
349        dt = 0
350
351if Nds > 3:
352   print errormsg
353   print fname + ': Ndims= ',Nds,' not ready !!!'
354   print errormsg
355   ncf.close()
356   quit(-1)
357
358print fname + ': Ndims: ',Nds,' dimensions of the matrices: ', dt, ', ', dy, ',', dx
359
360lonV=objlon[:]
361lonV=np.where(lonV < 0., lonV + 360., lonV)
362lonV=lonV*np.pi/180.
363latV=objlat[:]*np.pi/180.
364
365if opts.modkind == 'HS':
366    varVals = HS_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude,Peramplitude)
367
368elif opts.modkind == 'control':
369    varVals = control_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude)
370
371elif opts.modkind == 'control5N':
372    varVals = control5N_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude)
373
374elif opts.modkind == 'flat':
375    varVals = flat_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude)
376
377elif opts.modkind == 'KEQ':
378    varvalues=objvarmod[:]
379    varValsA = varvalues.copy()
380    varValsB = varvalues.copy()
381    varValsA = control_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude)
382    varValsB = KEQ_kind(Nds,dx,dy,dt,lonV,latV,refval,Peramplitude,Lon0,dLon,dlat)
383    varVals= varValsA + varValsB
384
385elif opts.modkind == 'KW1':
386    varvalues=objvarmod[:]
387    varValsA = varvalues.copy()
388    varValsB = varvalues.copy()
389    varValsA = control_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude)
390    varValsB = KW1_kind(Nds,dx,dy,dt,lonV,latV,refval,Peramplitude,Lon0,dLon,dlat)
391    varVals= varValsA + varValsB
392
393elif opts.modkind == 'peaked':
394    varVals = peaked_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude)
395
396elif opts.modkind == 'Qobs':
397    varvalues=objvarmod[:]
398    varvals1=varvalues.copy()
399    varvals2=varvalues.copy()
400
401    varvals1 = control_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude)
402    varvals2 = flat_kind(Nds,dx,dy,dt,lonV,latV,refval,amplitude)
403
404    varVals = (varvals1 + varvals2)/2.
405
406else:
407    print errormsg
408    print '  ' + fname + ': kind "' + opts.modkind + '" does not exist!!!'
409    print errormsg
410    ncf.close()
411    quit(-1)
412
413objvarmod[:] = varVals
414
415ncf.sync()
416ncf.close()
417
418print fname + ': "' + opts.varname + '" of "' + opts.ncfile + '" has been modified with the "' + opts.modkind + '"'
Note: See TracBrowser for help on using the repository browser.