compute_wind_speed_height.py
#!/usr/bin/env python
'''
Copyright 2019 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_wind_speed_height
Author (date) : Cristian Simarro (20/11/2015)
modified: Xavi Abellan (03/12/2018) - compatibilty with Python 3
Category : COMPUTATION
OneLineDesc : Computes the u/v components of the wind at certain height
Description : Computes computes the u/v components of the wind at certain
height. First, it calculates the geopotential of each model
level. Once the requested height is found between two model
levels, the program will vertically interpolate the u/v
component values.
Based on code from Nils Wedi, the IFS documentation:
https://www.ecmwf.int/en/forecasts/documentation-and-support/changes-ecmwf-model/ifs-documentation
part III. Dynamics and numerical procedures
optimised implementation by Dominique Lucas.
ported to Python by Cristian Simarro
Parameters : -w wind - height in meters you want to know
u/v components
tq.grib - grib file with all the levelist of
t and q
uv.grib - grib file with all the levelists of
u/v
zlnsp.grib - grib file with levelist 1 for
params z and lnsp
-o output (optional) - name of the output file
(default='uv_out_<wind>.grib')
Return Value : output (default='uv_out_<wind>.grib')
A fieldset the u/v components values for at the specified
height
Dependencies : None
Example Usage :
compute_wind_speed_height.py tq.grib zlnsp.grib uv.grib -w 100
'''
from __future__ import print_function
import sys
import argparse
import numpy as np
from eccodes import (codes_index_new_from_file, codes_index_get, codes_get,
codes_index_select, codes_new_from_index, codes_set,
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(
description='Python tool to calculate the wind at certain height')
parser.add_argument('-w', '--height', required=True, type=int,
help='height to calculate the wind components')
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'))
parser.add_argument('u_v', metavar='uv.grib', type=str,
help=('grib file with u and v component of wind for',
'the model levels'))
args = parser.parse_args()
if not args.output:
args.output = 'uv_out_%sm.grib' % args.height
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:
codes_index_add_file(idx, args.u_v)
# iterate date
for date in codes_index_get(idx, 'date'):
codes_index_select(idx, 'date', date)
# iterate step
for time in codes_index_get(idx, 'time'):
codes_index_select(idx, 'time', time)
values = get_initial_values(idx, keep_sample=True)
if 'height' in args:
values['height'] = args.height
values['gh'] = args.height * R_G + values['z']
if 'levelist' in args:
values['levelist'] = args.levelist
# iterate step all but geopotential z which is always step 0 (an)
for step in codes_index_get(idx, 'step'):
codes_index_select(idx, 'step', step)
# surface pressure
try:
values['sp'] = get_surface_pressure(idx)
production_step(idx, values, fout)
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, 'step', 0)
codes_index_select(idx, 'shortName', 'z')
gid = codes_new_from_index(idx)
values = {}
# surface geopotential
values['z'] = codes_get_values(gid)
values['pv'] = codes_get_array(gid, 'pv')
values['nlevels'] = codes_get(gid, 'NV', int) // 2 - 1
check_max_level(idx, values)
if keep_sample:
values['sample'] = gid
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?
max_level = max(codes_index_get(idx, 'level', int))
if max_level != values['nlevels']:
print('%s [WARN] total levels should be: %d but it is %d' %
(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)
codes_index_select(idx, 'shortName', 'lnsp')
gid = codes_new_from_index(idx)
if gid is None:
raise WrongStepError()
if codes_get(gid, 'gridType', str) == 'sh':
print('%s [ERROR] fields must be gridded, not spectral' % sys.argv[0],
file=sys.stderr)
sys.exit(1)
# surface pressure
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
# t_level: values for t
# q_level: values for q
codes_index_select(idx, 'level', lev)
codes_index_select(idx, 'shortName', 't')
gid = codes_new_from_index(idx)
t_level = codes_get_values(gid)
codes_release(gid)
codes_index_select(idx, 'shortName', 'q')
gid = codes_new_from_index(idx)
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, values, fout):
'''Produce u/v interpolated from ML for a given height'''
# 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, part III
# For speed and file I/O, we perform the computations with
# numpy vectors instead of fieldsets.
out = {}
params = ('u', 'v')
# we need to get z and lnsp from the first level to do the
# calculations
z_h = values['z']
# height in geopotential
z_f_prev = z_h
# initialize values for the output
for param in params:
out[param] = np.zeros(values['sp'].size)
found = [False for i in range(values['sp'].size)]
for lev in list(reversed(list(range(1, values['nlevels'] + 1)))):
z_h, z_f = compute_z_level(idx, lev, values, z_h)
# retrieve u/v params for the current level
for param in params:
codes_index_select(idx, 'level', lev) # 136
codes_index_select(idx, 'shortName', param)
gid = codes_new_from_index(idx)
values[param] = codes_get_values(gid)
codes_release(gid)
if lev < values['nlevels']:
codes_index_select(idx, 'level', lev + 1) # 137
gid = codes_new_from_index(idx)
values['prev' + param] = codes_get_values(gid)
codes_release(gid)
else:
values['prev' + param] = np.zeros(values['sp'].size)
# search if the provided wind height converted to
# geopotential (my_z) is between the current level (z_f)
# and the previous one (z_f_prev)
for i in range(z_f_prev.size):
if found[i]:
continue
if values['gh'][i] >= z_f_prev[i] and values['gh'][i] < z_f[i]:
found[i] = True
# when found, interpolate vertically to get the
# value and store it in out[param] to be written
# at the end
for param in params:
res = (((float(values[param][i]) *
(values['gh'][i] - z_f_prev[i])) +
(float(values['prev' + param][i]) *
(z_f[i] - values['gh'][i]))) /
(z_f[i] - z_f_prev[i]))
out[param][i] = res
# update z_f_prev
z_f_prev = z_f
# simple error check
for i in range(values['sp'].size):
if not found[i]:
print('point ', i, 'not found...')
# write the values in the fout file
for param in params:
codes_set(values['sample'], 'shortName', param)
codes_set(values['sample'], 'typeOfLevel', 'heightAboveGround')
codes_set(values['sample'], 'level', values['height'])
codes_set_values(values['sample'], out[param])
codes_write(values['sample'], fout)
class WrongStepError(Exception):
''' Exception capturing wrong step'''
pass
if __name__ == '__main__':
main()