11#!/usr/bin/env python3
22'''
33Copyright 2023 ECMWF.
4-
4+
55This software is licensed under the terms of the Apache Licence Version 2.0
66which can be obtained at http://www.apache.org/licenses/LICENSE-2.0
7-
7+
88In applying this licence, ECMWF does not waive the privileges and immunities
99granted to it by virtue of its status as an intergovernmental organisation
1010nor does it submit to any jurisdiction.
11-
11+
1212**************************************************************************
1313Function : compute_geopotential_on_ml
14-
14+
1515Author (date) : Cristian Simarro (09/10/2015)
1616modified: Cristian Simarro (20/03/2017) - migrated to eccodes
1717 Xavi Abellan (03/12/2018) - compatibilty with Python 3
1818 Xavi Abellan (27/03/2020) - Corrected steps in output file
1919 Xavi Abellan (05/08/2020) - Better handling of levels
2020 Xavi Abellan (31/01/2023) - Better handling of steps
21-
21+
2222Category : COMPUTATION
23-
23+
2424OneLineDesc : Computes geopotential on model levels
25-
25+
2626Description : Computes geopotential on model levels.
2727 Based on code from Nils Wedi, the IFS documentation:
2828 https://software.ecmwf.int/wiki/display/IFS/CY41R1+Official+IFS+Documentation
2929 part III. Dynamics and numerical procedures
3030 optimised implementation by Dominique Lucas.
3131 ported to Python by Cristian Simarro
32-
32+
3333Parameters : tq.grib - grib file with all the levelist
3434 of t and q
3535 zlnsp.grib - grib file with levelist 1 for params
3838 to store in the output
3939 -o output (optional) - name of the output file
4040 (default='z_out.grib')
41-
41+
4242Return Value : output (default='z_out.grib')
4343 A fieldset of geopotential on model levels
44-
44+
4545Dependencies : None
46-
46+
4747Example Usage :
4848 compute_geopotential_on_ml.py tq.grib zlnsp.grib
4949'''
5656 codes_index_add_file , codes_get_array , codes_get_values ,
5757 codes_index_release , codes_release , codes_set_values ,
5858 codes_write )
59-
59+
6060R_D = 287.06
6161R_G = 9.80665
62-
63-
62+
63+
6464def parse_args ():
6565 ''' Parse program arguments using ArgumentParser'''
6666 parser = argparse .ArgumentParser (
@@ -90,20 +90,20 @@ def parse_args():
9090 args .levelist = list (range (int (args .levelist [0 ]),
9191 int (args .levelist [2 ]) + 1 ))
9292 else :
93- args .levelist = [int (l ) for l in args .levelist .split ('/' )]
93+ args .levelist = [int (level ) for level in args .levelist .split ('/' )]
9494 return args
95-
96-
95+
96+
9797def main ():
9898 '''Main function'''
9999 args = parse_args ()
100-
100+
101101 print ('Arguments: %s' % ", " .join (
102102 ['%s: %s' % (k , v ) for k , v in vars (args ).items ()]))
103-
103+
104104 fout = open (args .output , 'wb' )
105105 index_keys = ['date' , 'time' , 'shortName' , 'level' , 'step' ]
106-
106+
107107 idx = codes_index_new_from_file (args .z_lnsp , index_keys )
108108 codes_index_add_file (idx , args .t_q )
109109 if 'u_v' in args :
@@ -131,22 +131,22 @@ def main():
131131 except WrongStepError :
132132 if step != '0' :
133133 raise
134-
134+
135135 try :
136136 codes_release (values ['sample' ])
137137 except KeyError :
138138 pass
139-
139+
140140 codes_index_release (idx )
141141 fout .close ()
142-
143-
142+
143+
144144def get_initial_values (idx , keep_sample = False ):
145145 '''Get the values of surface z, pv and number of levels '''
146146 codes_index_select (idx , 'level' , 1 )
147147 codes_index_select (idx , 'shortName' , 'z' )
148148 gid = codes_new_from_index (idx )
149-
149+
150150 values = {}
151151 # surface geopotential
152152 values ['z' ] = codes_get_values (gid )
@@ -158,8 +158,8 @@ def get_initial_values(idx, keep_sample=False):
158158 else :
159159 codes_release (gid )
160160 return values
161-
162-
161+
162+
163163def check_max_level (idx , values ):
164164 '''Make sure we have all the levels required'''
165165 # how many levels are we computing?
@@ -169,8 +169,8 @@ def check_max_level(idx, values):
169169 (sys .argv [0 ], values ['nlevels' ], max_level ),
170170 file = sys .stderr )
171171 values ['nlevels' ] = max_level
172-
173-
172+
173+
174174def get_surface_pressure (idx ):
175175 '''Get the surface pressure for date-time-step'''
176176 codes_index_select (idx , 'level' , 1 )
@@ -186,17 +186,17 @@ def get_surface_pressure(idx):
186186 sfc_p = np .exp (codes_get_values (gid ))
187187 codes_release (gid )
188188 return sfc_p
189-
190-
189+
190+
191191def get_ph_levs (values , level ):
192192 '''Return the presure at a given level and the next'''
193193 a_coef = values ['pv' ][0 :values ['nlevels' ] + 1 ]
194194 b_coef = values ['pv' ][values ['nlevels' ] + 1 :]
195195 ph_lev = a_coef [level - 1 ] + (b_coef [level - 1 ] * values ['sp' ])
196196 ph_levplusone = a_coef [level ] + (b_coef [level ] * values ['sp' ])
197197 return ph_lev , ph_levplusone
198-
199-
198+
199+
200200def compute_z_level (idx , lev , values , z_h ):
201201 '''Compute z at half & full level for the given level, based on t/q/sp'''
202202 # select the levelist and retrieve the vaules of t and q
@@ -215,34 +215,34 @@ def compute_z_level(idx, lev, values, z_h):
215215 raise MissingLevelError ('Q at level {} missing from input' .format (lev ))
216216 q_level = codes_get_values (gid )
217217 codes_release (gid )
218-
218+
219219 # compute moist temperature
220220 t_level = t_level * (1. + 0.609133 * q_level )
221-
221+
222222 # compute the pressures (on half-levels)
223223 ph_lev , ph_levplusone = get_ph_levs (values , lev )
224-
224+
225225 if lev == 1 :
226226 dlog_p = np .log (ph_levplusone / 0.1 )
227227 alpha = np .log (2 )
228228 else :
229229 dlog_p = np .log (ph_levplusone / ph_lev )
230230 alpha = 1. - ((ph_lev / (ph_levplusone - ph_lev )) * dlog_p )
231-
231+
232232 t_level = t_level * R_D
233-
233+
234234 # z_f is the geopotential of this full level
235235 # integrate from previous (lower) half-level z_h to the
236236 # full level
237237 z_f = z_h + (t_level * alpha )
238-
238+
239239 # z_h is the geopotential of 'half-levels'
240240 # integrate z_h to next half level
241241 z_h = z_h + (t_level * dlog_p )
242-
242+
243243 return z_h , z_f
244-
245-
244+
245+
246246def production_step (idx , step , values , fout ):
247247 '''Compute z at half & full level for the given level, based on t/q/sp'''
248248 # We want to integrate up into the atmosphere, starting at the
@@ -251,10 +251,10 @@ def production_step(idx, step, values, fout):
251251 # See the IFS documentation, part III
252252 # For speed and file I/O, we perform the computations with
253253 # numpy vectors instead of fieldsets.
254-
254+
255255 z_h = values ['z' ]
256256 codes_set (values ['sample' ], 'step' , int (step ))
257-
257+
258258 for lev in sorted (values ['levelist' ], reverse = True ):
259259 try :
260260 z_h , z_f = compute_z_level (idx , lev , values , z_h )
@@ -265,17 +265,17 @@ def production_step(idx, step, values, fout):
265265 except MissingLevelError as e :
266266 print ('%s [WARN] %s' % (sys .argv [0 ], e ),
267267 file = sys .stderr )
268-
269-
268+
269+
270270class WrongStepError (Exception ):
271271 ''' Exception capturing wrong step'''
272272 pass
273-
274-
273+
274+
275275class MissingLevelError (Exception ):
276276 ''' Exception capturing missing levels in input'''
277277 pass
278-
279-
278+
279+
280280if __name__ == '__main__' :
281281 main ()
0 commit comments