Versions Compared

Key

  • This line was added.
  • This line was removed.
  • Formatting was changed.

...

Code Block
languagepy
#!/usr/bin/env python

#
# Copyright 2015 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)
#
# 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    : -l levelist - list of levelist to store in the output 
#                  -o output   - name of the output file (default='z_out.grib')
#
# Files         : tq      - grib file with all the levelist of t and q
#                 zlnsp   - grib file with levelist 1 for params z and lnsp
#
# Return Value  : output (default='z_out.grib')
#                 A fieldset of geopotential on model levels
#
# Dependencies  : None
#
# Example Usage : 
#                 python compute_geopotential_on_ml.py tq.grib zlnsp.grib

from numpy import * 
import sys,math,os
import argparse 
from gribapi import *

def main(t_q,z_lnsp,out_name,levelist):


    print "Using as input files:\n   ",t_q,z_lnsp
    print "The result will be stored in:\n   ",out_name

    if levelist != "":
        print "Will only store these levels: "+levelist
        levelist_selected=levelist.split("/")
    fout = open(out_name,'w')
    ftmp = open(t_q)
    total_messages=grib_count_in_file(ftmp)/2
    ftmp.close()
    Rd = 287.06 
    index_keys = ["date","time","shortName","level","step"]

    values= {}
    pv = {}
    
    out_names = {'fc' : 'deterministic_pre.grib',
                 'cf' : 'control_pre.grib',
                 'pf' : 'ensemble_pre.grib',
                }

    zlnsp = grib_index_new_from_file(z_lnsp,index_keys)
    iidtq = grib_index_new_from_file(t_q,index_keys)

    #we need to get z and lnsp from the first level to do the calculations
    counter=0
    for date in grib_index_get(zlnsp,'date'):
        grib_index_select(zlnsp,'date',date)
        grib_index_select(iidtq,'date',date)
        for time in grib_index_get(zlnsp,'time'):
            grib_index_select(zlnsp,'time',time)
            grib_index_select(iidtq,'time',time)
        
            grib_index_select(zlnsp,'level',1)
   
            grib_index_select(zlnsp,'step',0)
            grib_index_select(zlnsp,'shortName','z')
            gid = grib_new_from_index(zlnsp)
            
            # surface geopotential
            values["z"] = grib_get_values(gid)
            z_h= values["z"]
            pv = grib_get_array(gid,'pv')
            levelSizeNV = grib_get(gid,'NV',int)/2 -1
            grib_release(gid)

            for step in grib_index_get(iidtq,'step'):
                z_h= values["z"]
                grib_index_select(iidtq,'step',step)
                grib_index_select(zlnsp,'step',step)
                for shortName in ["lnsp"]:
                   grib_index_select(zlnsp,'shortName',shortName)
                   gid = grib_new_from_index(zlnsp)
                #   if grib_get(gid,"gridType",str) == "sh":
                #       print(sys.argv[0]+' [ERROR] fields must be gridded, not spectral')
                #       sys.exit(1)
                   values[shortName] = grib_get_values(gid)
                   pv = grib_get_array(gid,'pv')
                   levelSizeNV = grib_get(gid,'NV',int)/2 -1
                   grib_release(gid)
   
                # surface pressure
                sp=exp(values["lnsp"])

                # get the coefficients for computing the pressures
                # how many levels are we computing?
                grib_index_select(iidtq,'shortName',"t")
                levelSize=max(grib_index_get(iidtq,"level",int))
                if levelSize != levelSizeNV:
                    print(sys.argv[0]+' [WARN] total levels should be: '+str(levelSizeNV)+' but it is '+str(levelSize))
                A = pv[0:levelSize+1]
                B = pv[levelSize+1:]

                Ph_levplusone = A[levelSize] + (B[levelSize]*sp)

                # We want to integrate up into the atmosphere, starting at the ground
                # so we start at the lowest level (highest number) and keep
                # accumulating the height as we go.
                # See the IFS documentation:
                # http://www.ecmwf.int/research/ifsdocs/DYNAMICS/Chap2_Discretization4.html
                # For speed and file I/O, we perform the computations with vectors instead
                # of fieldsets.

                for lev in list(reversed(range(1,levelSize+1))):
                    grib_index_select(iidtq,'level',lev)
                    grib_index_select(iidtq,'shortName',"t")
                    gid = grib_new_from_index(iidtq)
                    T_level = grib_get_values(gid)
                    gid2 = grib_clone(gid)
                    grib_release(gid)
                    grib_index_select(iidtq,'shortName',"q")
                    gid = grib_new_from_index(iidtq)
                    q_level = grib_get_values(gid)
                    grib_release(gid)

                    # compute moist temperature
                    T_level = T_level * (1.+0.609133*q_level)

                    # compute the pressures (on half-levels)
                    Ph_lev = A[lev-1] + (B[lev-1] * sp)

                    if lev == 1:
                        dlogP = log(Ph_levplusone/0.1)
                        #alpha = log(ones(len(Ph_levplusone))+1)
                        alpha = log(2)
                    else:
                        dlogP = log(Ph_levplusone/Ph_lev)
                        dP    = Ph_levplusone-Ph_lev
                        alpha = 1. - ((Ph_lev/dP)*dlogP)

                    TRd = T_level*Rd

                    # 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 + (TRd*alpha)

                    # z_h is the geopotential of 'half-levels'
                    # integrate z_h to next half level
                    z_h=z_h+(TRd*dlogP)

                    Ph_levplusone = Ph_lev
                    # store the result (z_f) in a field and add to the output fieldset
                    # (add it to the front, not the end, because we are going 'backwards'
                    # through the fields)
                    if levelist == "" or str(lev) in levelist_selected:
                        grib_set(gid2,"paramId",129)
                        grib_set(gid2,'generatingProcessIdentifier',128)
                        grib_set(gid2,'level', lev)
                        #grib_set(gid2,"bitsPerValue",24)
                        grib_set_values(gid2,z_f)
                        grib_write(gid2,fout)
                    counter += 1
                    if counter >= int((total_messages+1)/20):
                        sys.stdout.write('.')
                        sys.stdout.flush()
                        counter=0
                    grib_release(gid2)
           
    fout.close() 
    print("Done")

if __name__ == "__main__":
    request_date=0
    request_time=0
    levelist = ""
    out_name = 'z_out.grib'

    parser = argparse.ArgumentParser(
        description='Python tool to calculate the Z of the model levels')
    parser.add_argument("-l","--levelist", help="levelist to store")
    parser.add_argument("-o","--output", help="name of the output file")
    parser.add_argument('t_q', metavar='tq', type=str, 
                   help='grib file with temperature(t) and humidity(q) for the model levels')
    parser.add_argument('z_lnsp', metavar='zlnsp', type=str, 
                   help='grib file with geopotential(z) and Logarithm of surface pressure(lnsp) for the ml=1')
    args = parser.parse_args()
    for fname in (args.t_q,args.z_lnsp):
        if not os.path.isfile(fname):
            print "[ERROR] file %s does not exist" %(fname)
            sys.exit(1)
    if args.levelist:
        if args.levelist != "all":
            levelist=args.levelist 
    if args.output:
        out_name=args.output

    main(args.t_q,args.z_lnsp,out_name,levelist)