#!/usr/bin/env python3
# calculate the wind speed for u/v
# read U sequentially
import pywgrib2_s
import numpy
in_grib = "a.grb"
inv1='a.grb.inv1'
inv2='a.grb.inv2'
out='a_wind.grb'
# make two inventory files
err = pywgrib2_s.mk_inv(in_grib,inv1)
print("make inventory 1 err=", err)
err = pywgrib2_s.mk_inv(in_grib,inv2)
print("make inventory 2 err=", err)
count = 0
while pywgrib2_s.inq(in_grib,'UGRD',inv=inv1,Data=True, sequential=count, Matched=True) == 1:
count += 1
# save U
u = pywgrib2_s.data
# get match string for V
metadata = pywgrib2_s.matched[0]
# metadata UGRD -> VGRD, remove columns 1 and 2 (otherwise will not match)
metadata = metadata.replace(":UGRD:",":VGRD:")
i = metadata.find(":", metadata.find(":")+1)+1
metadata = metadata[i:]
# read corresponding V
nmatch = pywgrib2_s.inq(in_grib, metadata, inv=inv2, Data=True, grib='@mem:0')
if nmatch != 1:
print("error looking for V nmatch=",nmatch)
quit()
v = pywgrib2_s.data
wind = numpy.sqrt(u*u + v*v)
# write write wind speed, modify the V field
s=pywgrib2_s.write(out,'@mem:0',1,new_data=wind,var='WIND')
if s is None:
print("write failed")
count -= 1
pywgrib2_s.close(out)
print("wrote ",count," wind fields to ",out)
print("\n Note: if you get error message: RuntimeWarning: invalid value encountered in sqrt")
print(" You may be using a buggy anaconda-numpy.")
print(" See https://github.com/numpy/numpy/issues/11448")
print(" fix: install original numpy")