#!/usr/bin/env python3
# compute the average RH from the surface to 700 mb and write as a grib file
#
# input: grib file with
# surface pressure
# RH on 1000, 925, 850 and 700 mb levels
# (yes typical files has more levels of RH, but only example)
#
# output: grib file with pressure weighted RH from 700mb to surface
#
# for RH below 1000 mb, use 1000 mb value
#
import pywgrib2_s
import numpy
# can use regular files for inv and grib, but memory files are faster
# and don't have be cleaned up afterwards
input="c.grb"
inv="@mem:0"
out='c_ave_rh.grb'
grib='@mem:1'
level = [ 1000, 925, 850, 700 ]
# make inventory file
pywgrib2_s.mk_inv(input, inv)
print("made inventory file")
# read surface pressure
n = pywgrib2_s.inq(input,":PRES:surface:",inv=inv,Data=True,grib=grib,Matched=True)
if n != 1:
print("problem with reading surface pressure, number of matches=",n)
quit()
# save pywgrib2_s.data before it get overwritten by next read
# convert surface pressure from Pa to mb (hPa)
ps = 0.01 * pywgrib2_s.data
print("surface pressure ",pywgrib2_s.matched[0])
# find shape of ps
if ps.ndim != 2:
print("expected NxM array, ndim=",ps.ndim)
quit()
nx=ps.shape[0]
ny=ps.shape[1]
# allocate numpy 3d array
rh = numpy.zeros((nx, ny, len(level)), order='F')
print("number of levels=",len(level))
print("shape rh=",rh.shape)
for i, lev in enumerate(level):
# make lev compatible with wgrib2 search
slev=":"+str(lev)+" mb:"
n = pywgrib2_s.inq(input,":RH:",slev,inv=inv,Data=True)
if n != 1:
print("problem with RH number of matches=",n," level=",level)
quit()
rh[:,:,i] = pywgrib2_s.data[:,:]
print("read in RH lev=",lev)
print("read all RH fields")
print(ps > 1000.0)
# do the layer 1000mb - ps
layer_rh = rh[:,:,0]
ave_rh = numpy.maximum(ps-1000.0,0.0)*layer_rh
print('ave_rh')
print(ave_rh)
for layer in range(len(level)-1):
# make lev compatible with wgrib2 search
layer_rh = 0.5 * (rh[:,:,layer] + rh[:,:,layer+1])
level0=level[layer]
level1=level[layer+1]
print("layer=", layer, level0,' to ',level1)
ave_rh += numpy.minimum(numpy.maximum(ps-level1,0.0),level0-level1)*layer_rh
# divide by layer thickness, note: ave_rh is non-zero for thickness > 0
top = level[len(level)-1]
print("top=",top)
thickness = ps-top
print("thickness")
print(thickness)
thickness[thickness == 0.0] = 1.0
print("ave_rh")
ave_rh = ave_rh / thickness
print(ave_rh)