Merge branch 'main' into add_Lic

bin/compute_geopotential_on_ml.py

@@ -1,35 +1,35 @@
 #!/usr/bin/env python3
 '''
 Copyright 2023 ECMWF.
- 
+
 This software is licensed under the terms of the Apache Licence Version 2.0
 which can be obtained at http://www.apache.org/licenses/LICENSE-2.0
- 
+
 In applying this licence, ECMWF does not waive the privileges and immunities
 granted to it by virtue of its status as an intergovernmental organisation
 nor does it submit to any jurisdiction.
- 
+
 **************************************************************************
 Function      : compute_geopotential_on_ml
- 
+
 Author (date) : Cristian Simarro (09/10/2015)
 modified:       Cristian Simarro (20/03/2017) - migrated to eccodes
                 Xavi Abellan     (03/12/2018) - compatibilty with Python 3
                 Xavi Abellan     (27/03/2020) - Corrected steps in output file
                 Xavi Abellan     (05/08/2020) - Better handling of levels
                 Xavi Abellan     (31/01/2023) - Better handling of steps
- 
+
 Category      : COMPUTATION
- 
+
 OneLineDesc   : Computes geopotential on model levels
- 
+
 Description   : Computes geopotential on model levels.
                 Based on code from Nils Wedi, the IFS documentation:
                 https://software.ecmwf.int/wiki/display/IFS/CY41R1+Official+IFS+Documentation
                 part III. Dynamics and numerical procedures
                 optimised implementation by Dominique Lucas.
                 ported to Python by Cristian Simarro
- 
+
 Parameters    : tq.grib                - grib file with all the levelist
                                          of t and q
                 zlnsp.grib             - grib file with levelist 1 for params
@@ -38,12 +38,12 @@
                                          to store in the output
                 -o output   (optional) - name of the output file
                                          (default='z_out.grib')
- 
+
 Return Value  : output (default='z_out.grib')
                 A fieldset of geopotential on model levels
- 
+
 Dependencies  : None
- 
+
 Example Usage :
                 compute_geopotential_on_ml.py tq.grib zlnsp.grib
 '''
@@ -56,11 +56,11 @@
                      codes_index_add_file, codes_get_array, codes_get_values,
                      codes_index_release, codes_release, codes_set_values,
                      codes_write)
- 
+
 R_D = 287.06
 R_G = 9.80665
- 
- 
+
+
 def parse_args():
     ''' Parse program arguments using ArgumentParser'''
     parser = argparse.ArgumentParser(
@@ -90,20 +90,20 @@ def parse_args():
             args.levelist = list(range(int(args.levelist[0]),
                                        int(args.levelist[2]) + 1))
     else:
-        args.levelist = [int(l) for l in args.levelist.split('/')]
+        args.levelist = [int(level) for level in args.levelist.split('/')]
     return args
- 
- 
+
+
 def main():
     '''Main function'''
     args = parse_args()
- 
+
     print('Arguments: %s' % ", ".join(
         ['%s: %s' % (k, v) for k, v in vars(args).items()]))
- 
+
     fout = open(args.output, 'wb')
     index_keys = ['date', 'time', 'shortName', 'level', 'step']
- 
+
     idx = codes_index_new_from_file(args.z_lnsp, index_keys)
     codes_index_add_file(idx, args.t_q)
     if 'u_v' in args:
@@ -131,22 +131,22 @@ def main():
                     except WrongStepError:
                         if step != '0':
                             raise
- 
+
         try:
             codes_release(values['sample'])
         except KeyError:
             pass
- 
+
     codes_index_release(idx)
     fout.close()
- 
- 
+
+
 def get_initial_values(idx, keep_sample=False):
     '''Get the values of surface z, pv and number of levels '''
     codes_index_select(idx, 'level', 1)
     codes_index_select(idx, 'shortName', 'z')
     gid = codes_new_from_index(idx)
- 
+
     values = {}
     # surface geopotential
     values['z'] = codes_get_values(gid)
@@ -158,8 +158,8 @@ def get_initial_values(idx, keep_sample=False):
     else:
         codes_release(gid)
     return values
- 
- 
+
+
 def check_max_level(idx, values):
     '''Make sure we have all the levels required'''
     # how many levels are we computing?
@@ -169,8 +169,8 @@ def check_max_level(idx, values):
               (sys.argv[0], values['nlevels'], max_level),
               file=sys.stderr)
         values['nlevels'] = max_level
- 
- 
+
+
 def get_surface_pressure(idx):
     '''Get the surface pressure for date-time-step'''
     codes_index_select(idx, 'level', 1)
@@ -186,17 +186,17 @@ def get_surface_pressure(idx):
     sfc_p = np.exp(codes_get_values(gid))
     codes_release(gid)
     return sfc_p
- 
- 
+
+
 def get_ph_levs(values, level):
     '''Return the presure at a given level and the next'''
     a_coef = values['pv'][0:values['nlevels'] + 1]
     b_coef = values['pv'][values['nlevels'] + 1:]
     ph_lev = a_coef[level - 1] + (b_coef[level - 1] * values['sp'])
     ph_levplusone = a_coef[level] + (b_coef[level] * values['sp'])
     return ph_lev, ph_levplusone
- 
- 
+
+
 def compute_z_level(idx, lev, values, z_h):
     '''Compute z at half & full level for the given level, based on t/q/sp'''
     # select the levelist and retrieve the vaules of t and q
@@ -215,34 +215,34 @@ def compute_z_level(idx, lev, values, z_h):
         raise MissingLevelError('Q at level {} missing from input'.format(lev))
     q_level = codes_get_values(gid)
     codes_release(gid)
- 
+
     # compute moist temperature
     t_level = t_level * (1. + 0.609133 * q_level)
- 
+
     # compute the pressures (on half-levels)
     ph_lev, ph_levplusone = get_ph_levs(values, lev)
- 
+
     if lev == 1:
         dlog_p = np.log(ph_levplusone / 0.1)
         alpha = np.log(2)
     else:
         dlog_p = np.log(ph_levplusone / ph_lev)
         alpha = 1. - ((ph_lev / (ph_levplusone - ph_lev)) * dlog_p)
- 
+
     t_level = t_level * R_D
- 
+
     # z_f is the geopotential of this full level
     # integrate from previous (lower) half-level z_h to the
     # full level
     z_f = z_h + (t_level * alpha)
- 
+
     # z_h is the geopotential of 'half-levels'
     # integrate z_h to next half level
     z_h = z_h + (t_level * dlog_p)
- 
+
     return z_h, z_f
- 
- 
+
+
 def production_step(idx, step, values, fout):
     '''Compute z at half & full level for the given level, based on t/q/sp'''
     # We want to integrate up into the atmosphere, starting at the
@@ -251,10 +251,10 @@ def production_step(idx, step, values, fout):
     # See the IFS documentation, part III
     # For speed and file I/O, we perform the computations with
     # numpy vectors instead of fieldsets.
- 
+
     z_h = values['z']
     codes_set(values['sample'], 'step', int(step))
- 
+
     for lev in sorted(values['levelist'], reverse=True):
         try:
             z_h, z_f = compute_z_level(idx, lev, values, z_h)
@@ -265,17 +265,17 @@ def production_step(idx, step, values, fout):
         except MissingLevelError as e:
             print('%s [WARN] %s' % (sys.argv[0], e),
                   file=sys.stderr)
- 
- 
+
+
 class WrongStepError(Exception):
     ''' Exception capturing wrong step'''
     pass
- 
- 
+
+
 class MissingLevelError(Exception):
     ''' Exception capturing missing levels in input'''
     pass
- 
- 
+
+
 if __name__ == '__main__':
     main()