compute_geopotential_on_ml.py
#!/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 : tq.grib - grib file with all the levelist of t and q # zlnsp.grib - grib file with levelist 1 for params z and lnsp # -l levelist (optional) - list of levelist 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 : # 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): #some checks and information printing 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 = {} 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: # https://software.ecmwf.int/wiki/display/IFS/CY41R1+Official+IFS+Documentation # part III # For speed and file I/O, we perform the computations with numpy vectors instead # of fieldsets. for lev in list(reversed(range(1,levelSize+1))): # select the levelist and retrieve the vaules of t and q # t_level: values for t # q_level: values for q grib_index_select(iidtq,'level',lev) grib_index_select(iidtq,'shortName',"t") gid = grib_new_from_index(iidtq) t_level = grib_get_values(gid) #gid_out will be used as output, cloning to get the attributes gid_out = 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(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(gid_out,"paramId",129) grib_set(gid_out,'generatingProcessIdentifier',128) grib_set(gid_out,'level', lev) #grib_set(gid_out,"bitsPerValue",24) grib_set_values(gid_out,z_f) grib_write(gid_out,fout) counter += 1 if counter >= int((total_messages+1)/20): sys.stdout.write('.') sys.stdout.flush() counter=0 grib_release(gid_out) 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.grib', type=str, help='grib file with temperature(t) and humidity(q) for the model levels') parser.add_argument('z_lnsp', metavar='zlnsp.grib', 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 #calling main function main(args.t_q,args.z_lnsp,out_name,levelist)