def testPerturbation(iFile, oFile, var, level, mode):
'''
Test if perturbation was succesful
:param iFile: original input file [path]
:type iFile: string
:param oFile: perturbed output file [path]
:type oFile: string
:param var: variable name
:type var: string
:param level: perturbed model level
:type level: int
:param mode: perturbed mode
:type mode: int
:returns: success
:rtype: bool
'''
# unperturbed gribfile
grbindxIn = pygrib.index(iFile, 'shortName', 'level')
unPert = grbindxIn.select(shortName=var, level=level)[0]['values'][mode]
grbindxIn.close()
# perturbed gribfile
grbindxOut = pygrib.index(oFile, 'shortName', 'level')
pert = grbindxOut.select(shortName=var, level=level)[0]['values'][mode]
grbindxOut.close()
if not (unPert == pert):
return True
else:
return False
def read_write_grib(file_name, is_grib2):
exists, exists_file_name = _check_exists(file_name, path=True)
if not exists:
print('* Warning: file %s not found' % file_name)
return
if verbose:
print('Loading %s' % exists_file_name)
_unzip(exists_file_name)
if verbose:
print(' Reading')
member_index = member_coord.index(member)
time_index = forecast_hour_coord.index(forecast_hour)
grib_data = pygrib.open(file_name)
if is_grib2:
table = grib2_table
grib_index = pygrib.index(file_name, 'parameterCategory', 'parameterNumber', 'level')
else:
table = grib1_table
grib_index = pygrib.index(file_name, 'indicatorOfParameter', 'indicatorOfTypeOfLevel', 'level')
if verbose:
print('Variables to fetch: %s' % (variables,))
for row in range(table.shape[0]):
var = table[row, 0]
if var in variables:
if var not in nc_fid.variables.keys():
if verbose:
print('Creating variable %s' % var)
nc_var = nc_fid.createVariable(var, np.float32,
('time', 'member', 'fhour', 'south_north', 'west_east'),
zlib=True)
nc_var.setncatts({
'long_name': table[row, 5],
'units': table[row, 6],
'_FillValue': fill_value
})
try:
if verbose:
print('Writing %s' % var)
if is_grib2:
grib_list = grib_index.select(parameterCategory=int(table[row, 1]),
parameterNumber=int(table[row, 2]),
level=int(table[row, 3]))
else:
grib_list = grib_index.select(indicatorOfParameter=int(table[row, 1]),
indicatorOfTypeOfLevel=str(table[row, 2]),
level=int(table[row, 3]))
if verbose and len(grib_list) > 1:
print('* Warning: found multiple matches for %s; using the last (%s)' %
(var, grib_list[-1]))
data = np.array(grib_list[-1].values, dtype=np.float32)
data[data > 1.e30] = np.nan
nc_fid.variables[var][0, member_index, time_index, ...] = data
except (ValueError, OSError): # missing index gives an OS read error
print('* Warning: grib variable %s not found in file %s' % (var, file_name))
pass
except BaseException as e:
print("* Warning: failed to write %s to netCDF file ('%s')" % (var, str(e)))
grib_data.close()
return
def _valid_times(variable, path):
messages = pg.index(path, "name")
for message in messages.select(name=variable):
validTime = "{0:8d}{1:04d}".format(message["validityDate"],
message["validityTime"])
yield dt.datetime.strptime(validTime, "%Y%m%d%H%M")
messages.close()
def get_grib2_data(self, path, valid_time, variable, pressure):
cache = {}
validTime = dt.datetime.strptime(str(valid_time), "%Y-%m-%d %H:%M:%S")
vTime = "{0:d}{1:02d}".format(validTime.hour, validTime.minute)
messages = pg.index(path, "name", "scaledValueOfFirstFixedSurface",
"validityTime")
if len(path) > 0:
field = messages.select(
name=variable,
scaledValueOfFirstFixedSurface=int(pressure),
validityTime=vTime)[0]
cache["longitude"] = field.latlons()[1][0, :]
cache["latitude"] = field.latlons()[0][:, 0]
cache["data"] = field.values
cache["units"] = field.units
cache["name"] = field.name
cache["valid"] = "{0:02d}:{1:02d} UTC".format(
validTime.hour, validTime.minute)
cache["initial"] = "blah"
scaledLowerLevel = float(field.scaledValueOfFirstFixedSurface)
scaleFactorLowerLevel = float(field.scaleFactorOfFirstFixedSurface)
lowerSigmaLevel = str(
round(scaledLowerLevel * 10**-scaleFactorLowerLevel, 2))
scaledUpperLevel = float(field.scaledValueOfSecondFixedSurface)
scaleFactorUpperLevel = float(
field.scaleFactorOfSecondFixedSurface)
upperSigmaLevel = str(
round(scaledUpperLevel * 10**-scaleFactorUpperLevel, 2))
cache['layer'] = lowerSigmaLevel + "-" + upperSigmaLevel
messages.close()
return cache
def __init__(self, filename, filename_info, filetype_info):
"""Open grib file and do initial message parsing."""
super(GRIBFileHandler, self).__init__(filename, filename_info,
filetype_info)
self._msg_datasets = {}
self._start_time = None
self._end_time = None
try:
with pygrib.open(self.filename) as grib_file:
first_msg = grib_file.message(1)
last_msg = grib_file.message(grib_file.messages)
start_time = self._convert_datetime(first_msg, 'validityDate',
'validityTime')
end_time = self._convert_datetime(last_msg, 'validityDate',
'validityTime')
self._start_time = start_time
self._end_time = end_time
if 'keys' not in filetype_info:
self._analyze_messages(grib_file)
self._idx = None
else:
self._create_dataset_ids(filetype_info['keys'])
self._idx = pygrib.index(self.filename,
*filetype_info['keys'].keys())
except (RuntimeError, KeyError):
raise IOError("Unknown GRIB file format: {}".format(self.filename))
def _pressures(variable, path):
messages = pg.index(path, "name")
try:
for message in messages.select(name=variable):
yield message["scaledValueOfFirstFixedSurface"]
except ValueError:
# messages.select(name=variable) raises ValueError if not found
pass
messages.close()
def read_write_grib(file_name):
exists, exists_file_name = _check_exists(file_name, path=True)
if not exists:
print('* Warning: file %s not found' % file_name)
return
if verbose:
print('Loading %s' % exists_file_name)
if verbose:
print(' Reading')
member_index = member_coord.index(member)
grib_data = pygrib.open(file_name)
if level == '':
grib_index = pygrib.index(file_name, 'forecastTime')
else:
grib_index = pygrib.index(file_name, 'level', 'forecastTime')
if verbose:
print('Writing %s' % variable)
for forecast_hour in forecast_hours:
fhour_index = forecast_hour_coord.index(forecast_hour)
if forecast_hour not in self.forecast_hour_coord:
print('* Warning: I am only set up to retrieve forecast hours within %s' %
self.forecast_hour_coord)
continue
try:
if level == '':
grib_list = grib_index.select(forecastTime=forecast_hour)
else:
grib_list = grib_index.select(forecastTime=forecast_hour, level=int(level))
if verbose and len(grib_list) > 1:
print('* Warning: found multiple matches for fhour %s; using the last (%s)' %
(forecast_hour, grib_list[-1]))
elif verbose:
print('%s' % grib_list[0])
data = np.array(grib_list[-1].values, dtype=np.float32)
data[data > 1.e30] = np.nan
nc_fid.variables[variable][0, member_index, fhour_index, ...] = data
except (ValueError, OSError): # missing index gives an OS read error
print('* Warning: %s for fhour %s not found in file %s' % (variable, forecast_hour, file_name))
pass
except BaseException as e:
print("* Warning: failed to write %s to netCDF file ('%s')" % (variable, str(e)))
grib_data.close()
return
def get_grib_field(infile, var_name): # LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL
'''
Extracts one 2d or 3d-field from grib file
Parameters
----------
infile : str
input filename of the concerned grib1 file
var_name : str
short name of variable in infile
Returns
--------
var : numpy array
field of variable var_name
'''
# open and read from COSMOS file -------------------------------------
Indx = pygrib.index(infile, 'shortName')
# get certain variable ...............................................
Cont = Indx.select(shortName=var_name)
Nx = Cont[0].Ni
Ny = Cont[0].Nj
Nz = np.shape(Cont)[0]
if Nz > 1:
var = np.zeros([Nx, Ny, Nz])
lev = np.zeros([Nz])
n = 0
for layer in Cont:
var[:, :, n] = np.transpose(layer.values)
lev[n] = layer.level
n = n + 1
# sort fields in vertical
lev = np.array(lev)
ils = lev.argsort()
l = lev[ils]
v = var[..., ils]
else:
layer = Cont[0]
v = np.transpose(layer.values)
l = [
layer.level,
]
return v # TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT
def _valid_times(variable, path):
messages = pg.index(path, "name")
try:
for message in messages.select(name=variable):
validTime = "{0:8d}{1:04d}".format(message["validityDate"],
message["validityTime"])
yield dt.datetime.strptime(validTime, "%Y%m%d%H%M")
except ValueError:
# messages.select(name=variable) raises ValueError if not found
pass
messages.close()
def get_2d_vis_data(filename):
"""Get the 2D visibility data of a GFS grib file."""
msgs0 = pygrib.index(filename, 'shortName', 'typeOfLevel')
msgs_vis = msgs0.select(shortName='vis', typeOfLevel='surface')
assert (len(msgs_vis)) == 1
msg_vis = msgs_vis[0]
del msgs_vis
forecasttime = msg_vis['startStep']
msg_vis_dict = grbmsg_dict_subset_region(msg_vis)
lat = msg_vis_dict['lat']
lon = msg_vis_dict['lon']
return forecasttime, lat, lon, msg_vis_dict['values']
def test_pygrib_select_real_file(real_file):
name = "U component of wind"
messages = pygrib.index(real_file, "name", "validityTime",
"scaledValueOfFirstFixedSurface")
time = 0 # NOTE: zero-padded strings raise Runtime errors
# level = 899999976
level = 899
actual = []
for message in messages.select(name=name,
validityTime=time,
scaledValueOfFirstFixedSurface=level):
actual.append(message["scaledValueOfFirstFixedSurface"])
messages.close()
assert actual == [level]
def read_scs_oriented(self, vars_mode, file_path, dt_cursor):
# Open ERA5 grib data and read 6-hourly
grbidx = pygrib.index(file_path, 'dataTime')
for hourtime in range(0, 2400,
self.CONFIG['product']\
['temporal_resolution']):
# Create ERA5 table for SCS
SCSERA5 = self.create_scs_era5_table(dt_cursor, hourtime)
selected_grbs = grbidx.select(dataTime=hourtime)
# Generate frame of one 6-hour SCS ERA5 table
table_entity = self.gen_scs_era5_entity(SCSERA5)
total = len(selected_grbs)
for idx, grb in enumerate(selected_grbs):
info = (f"""\rReading grbs on hour """
f"""{int(hourtime/100)} {idx+1}/{total}""")
print(f'\r{info}', end='')
# Traverse all data point in ERA5 grib message,
# find corresponding row in SCS ERA5 table frame
# and fill its environmental variables
table_entity = self.fill_scs_era5_table_entity(
grb, table_entity)
# Temporarily not interpolate the space between
# ERA5 0.25 degree grid points
# Method_1: conventional interpolation methods
# Method_2: GAN
# Insert entity into database
utils.bulk_insert_avoid_duplicate_unique(
table_entity, self.CONFIG['database']\
['batch_size']['insert'],
SCSERA5, ['x_y'], self.session,
check_self=True)
utils.delete_last_lines()
print('Done')
def get_fc_fields(cfile, var_list, lonlat=True):
'''
Reads a list of variables from grib file.
Parameters
----------
cfile : str
grib filename (full path)
var_list : list
list of variable names (shortName)
lonlat : bool, optional, default = False
option if longitude and latitude are also retrieved (default: False)
Returns
--------
var : dict of numpy arrays
dictionary of variable fields
'''
indx = pygrib.index(cfile, 'shortName')
var = {}
for vname in var_list:
print('... reading %s' % vname)
try:
if lonlat and 'lon' not in var:
var['lon'], var['lat'], var[vname] = get_field_from_indx(
indx, vname, lonlat=True)
else:
var[vname] = get_field_from_indx(indx, vname, lonlat=False)
except:
print('%s not in %s' % (vname, cfile))
return var
def get_full_grib_field_info(infile, var_name): # LLLLLLLLLLLLLLLLLLLLL
'''
Prints all keys and their corresponding values for a grib1 field.
Parameters
----------
infile : str
input filename of the concerned grib1 file
var_name : str
short name of variable in infile
Returns
--------
grb_info : dict
dictionary of field keys and their values
'''
# open and read from COSMOS file -------------------------------------
Indx = pygrib.index(infile, 'shortName')
# get certain variable ...............................................
cont = Indx.select(shortName=var_name)
# set first field
f = cont[0]
# start with info dict
grb_info = {}
for key in list(f.keys()):
grb_info[key] = f[key]
print((key, ' ', f[key]))
return grb_info # TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT
def get_3d_data(filename):
"""Get 3D data from ONE grib file."""
from operator import attrgetter
import numpy as np
msgs0 = pygrib.index(filename, 'shortName', 'typeOfLevel')
# msgs0 = pygrib.open(filename)
# Select from grbmsgs by variable
msgst = msgs0.select(shortName='t', typeOfLevel='isobaricInhPa')
msgsh = msgs0.select(shortName='gh', typeOfLevel='isobaricInhPa')
# Should contain only one element
msgsh0 = msgs0.select(shortName='gh', typeOfLevel='isothermZero')
assert (len(msgsh0) == 1)
# Sort grbmsgs by lv
msgst = sorted(msgst, key=attrgetter('level'), reverse=True)
msgsh = sorted(msgsh, key=attrgetter('level'), reverse=True)
# Get latlon
# lat, lon = grbmsg_latlon(msgst[0])
# Get forecast hour, ['stepRange'] and ['endStep'] are same
forecasttime = msgst[0]['startStep']
# Get levels and assume levels of 3 variables are the same
lvs_t = [item['level'] for item in msgst]
lvs_h = [item['level'] for item in msgsh]
assert (lvs_t == lvs_h)
lv = lvs_t
# Resizing and converting msgs to dict
msgst = [grbmsg_dict_subset_region(msg) for msg in msgst]
msgsh = [grbmsg_dict_subset_region(msg) for msg in msgsh]
msgsh0 = grbmsg_dict_subset_region(msgsh0[0])
# Stacking levels together
msgst = np.stack([msg['values'] for msg in msgst])
msgsh = np.stack([msg['values'] for msg in msgsh])
msgsh0value = msgsh0['values']
lat = msgsh0['lat']
lon = msgsh0['lon']
return forecasttime, lv, lat, lon, msgst, msgsh, msgsh0value
# %%
import pygrib
grbindex = pygrib.index(file_name, 'Name', 'level')
grbindex.keys
selected_grbs = grbindex.select(Name='Geopotential Height', level=5000.0)
for grb in selected_grbs:
grb
grbindex.write('filename_inx')
# %% NOTES
def grib2_to_am_layers(grib_buffer, lat, lon, alt):
with tempfile.NamedTemporaryFile(mode='wb', prefix='temp-',
suffix='.grb') as f:
f.write(grib_buffer)
try:
grbindx = pygrib.index(f.name, "name", "level")
except Exception as e:
# stderr is being captured so we can't print to it
# example: RuntimeError: b'End of resource reached when reading message'
msg = 'pygrib raised {}, length of input was {}'.format(
str(e), len(grib_buffer))
raise RuntimeError(msg)
# in memory -- not sure what syntax actually works for this?
# need to .index() after creation
# grbindx = pygrib.fromstring(grib_buffer)
latlon_delta = float(
LATLON_GRID_STR[0:1]) + 0.01 * float(LATLON_GRID_STR[2:])
leftlon = math.floor(lon / latlon_delta) * latlon_delta
bottomlat = math.floor(lat / latlon_delta) * latlon_delta
u = (lat - bottomlat) / latlon_delta
v = (lon - leftlon) / latlon_delta
Pbase = []
z = []
T = []
o3_vmr = []
RH = []
cloud_lmr = []
cloud_imr = []
for i, lev in enumerate(LEVELS):
Pbase.append(lev)
try:
x = (grid_interp(
grbindx.select(name="Geopotential Height",
level=lev)[0].values, u, v))
z.append(x)
except:
z.append(BADVAL)
try:
x = (grid_interp(
grbindx.select(name="Temperature", level=lev)[0].values, u, v))
T.append(x)
except:
T.append(BADVAL)
try:
x = (grid_interp(
grbindx.select(name="Ozone mixing ratio", level=lev)[0].values,
u, v))
x *= M_AIR / M_O3 # convert mass mixing ratio to volume mixing ratio
o3_vmr.append(x)
except:
o3_vmr.append(0.0)
try:
x = (grid_interp(
grbindx.select(name="Relative humidity", level=lev)[0].values,
u, v))
if (lev >= RH_TOP_PLEVEL):
RH.append(x)
else:
RH.append(0.0)
except:
RH.append(0.0)
try:
x = (grid_interp(
grbindx.select(name="Cloud mixing ratio", level=lev)[0].values,
u, v))
cloud_lmr.append(x)
except:
cloud_lmr.append(0.0)
try:
x = (grid_interp(
grbindx.select(name="Ice water mixing ratio",
level=lev)[0].values, u, v))
cloud_imr.append(x)
except:
cloud_imr.append(0.0)
return Pbase, z, T, o3_vmr, RH, cloud_lmr, cloud_imr
def mlev2pres(infile, var): # LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL
'''
Calculates the pressure of a ECMWF hybrid model level given as
p = a + b * p_surf
Parameters
----------
infile : str
input filename of the concerned grib1 file
var : str
short name of variable in infile
Returns
--------
pres : numpy array
pressure levels
'''
# open and read from COSMOS file -------------------------------------
Indx = pygrib.index(infile, 'shortName')
# get surface pressure ...............................................
try:
sp_cont = Indx.select(shortName='sp')
sp = np.transpose(sp_cont[0].values)
except:
pass # no sp variable
try:
sp_cont = Indx.select(shortName='lnsp')
lnsp = np.transpose(sp_cont[0].values)
sp = np.exp(lnsp)
except:
pass # no lnsp variable
try:
sp_shape = sp.shape
except:
raise
# get level information for variable .................................
var_cont = Indx.select(shortName=var)
Nx = var_cont[0].Ni
Ny = var_cont[0].Nj
Nz = np.shape(var_cont)[0]
lev = {}
n = 0
for layer in var_cont:
lev[n] = layer.level
n = n + 1
# get level parameters ...............................................
ab = var_cont[0].pv
Nlev_half = ab.shape[0] / 2
# these are the corefficients of the half levels
a, b = ab[0:Nlev_half], ab[Nlev_half:]
# little check
if a[0] != 0. or a[-1] != 0. or b[0] != 0. or b[-1] != 1.:
print('Warning: Coefficient Check failed!')
print(('a = ', a))
print(('b = ', b))
pres = np.zeros([Nx, Ny, Nz])
if Nz != Nlev_half:
# assume the variable is given on full levels
a = 0.5 * (a[1:] + a[:-1])
b = 0.5 * (b[1:] + b[:-1])
# pressure on half levels ............................................
for n in range(Nz):
pres[:, :, n] = a[n] + b[n] * sp[:, :]
return pres
def _get_data(self) :
if '_data' not in self.__dict__ :
# dummy conditions to play with
newConditions = self.conditions.copy()
# scalar conditions only (input for the gribIndex)
subConditions = self.conditions.copy()
################
# TIME & LEVEL #
################
if 'time' not in self.conditions :
newConditions['time'] = self.axes['time'].data
else :
# gribIndex won't want lists of datetimes
# but rather individual year/month/day/hour
del subConditions['time']
# make sure time condition is iterable
if not isinstance(newConditions['time'], list) :
if not isinstance(newConditions['time'], np.ndarray) :
newConditions['time'] = [newConditions['time']]
# if data is 2D, it will have already have a level condition
# idem if it's 3D and has already been sliced
# if not, that means the user wants all available levels
if 'level' not in self.conditions :
newConditions['level'] = self.axes['level'].data
########################
# LATITUDE & LONGITUDE #
########################
### MASK ###
# mask is used to slice the netcdf array contained in gribMessages
mask = []
if 'latitude' in self.conditions :
del subConditions['latitude']
mask.append(self.conditions['latitude'])
else :
mask.append(slice(None))
twistedLongitudes = False
if 'longitude' in self.conditions :
del subConditions['longitude']
# twisted longitudes...
if type(self.conditions['longitude']) == tuple :
twistedLongitudes = True
secondMask = mask[:]
mask.append(self.conditions['longitude'][0])
#slice1 = slice(0, -mask[-1].start)
slice1 = slice(0, mask[-1].stop - mask[-1].start)
secondMask.append(self.conditions['longitude'][1])
slice2 = slice(-secondMask[-1].stop, None)
else :
mask.append(self.conditions['longitude'])
else :
mask.append(slice(None))
mask = tuple(mask)
### HORIZONTAL SHAPE ###
# shape of the output array : (time, level, horizontalShape)
horizontalShape = []
if hasattr(self, 'lats') :
horizontalShape.append(len(self.lats))
if hasattr(self, 'lons') :
horizontalShape.append(len(self.lons))
horizontalShape = tuple(horizontalShape)
#####################
# GET GRIB MESSAGES #
#####################
shape = ()
for axisName, axis in self.axes.iteritems() :
shape = shape + (len(axis),)
# build the output numpy array
self._data = np.empty(shape, dtype=float)
# flatten time and level dimensions
# that's in case there's neither time nor level dimension
self._data.shape = (-1,) + horizontalShape
# load the grib index
gribIndex = pygrib.index(self.fileName+'.idx')
lineIndex = 0
for instant in newConditions['time'] :
subConditions['year'] = instant.year
subConditions['month'] = instant.month
subConditions['day'] = instant.day
subConditions['hour'] = instant.hour
for level in newConditions['level'] :
subConditions['level'] = \
np.asscalar(np.array(level))
# converts numpy types to standard types
# standard types are converted to numpy
# normally, there should be only one line
# that answers our query
gribLine = gribIndex(**subConditions)[0]
if twistedLongitudes :
self._data[lineIndex, ..., slice1] = \
gribLine.values[mask]
self._data[lineIndex, ..., slice2] = \
gribLine.values[secondMask]
else :
self._data[lineIndex] = gribLine.values[mask]
lineIndex += 1
gribIndex.close()
self._data.shape = shape
return self._data
def __init__(self, filePath) :
super(File, self).__init__()
fileName = splitext(filePath)[0]
rawFile = pygrib.open(filePath)
# read the first line of the file
gribLine = rawFile.readline()
firstInstant = datetime(gribLine.year, gribLine.month, gribLine.day,
gribLine.hour, gribLine.minute, gribLine.second)
###################
# HORIZONTAL AXES #
###################
lats, lons = gribLine.latlons()
if lats[0, 0] == lats[0, 1] :
self.axes['latitude'] = aa.Meridian(lats[:, 0], 'degrees')
self.axes['longitude'] = aa.Parallel(lons[0, :], 'degrees')
else :
self.axes['latitude'] = aa.Meridian(lats[0, :], 'degrees')
self.axes['longitude'] = aa.Parallel(lons[:, 0], 'degrees')
#################
# VERTICAL AXIS #
#################
# sometimes there are several types of level
# 2D data is followed by 3D data e.g. jra25
variablesLevels = {} # variable - level type - level
variablesMetaData = {}
# loop through the variables and levels of the first time step
# default : grib has a time axis
timeDimension = True
while datetime(gribLine.year, gribLine.month, gribLine.day,
gribLine.hour, gribLine.minute, gribLine.second)\
== firstInstant :
# is it the first time this variable is met ?
if gribLine.shortName not in variablesLevels :
# create a dictionary for that variable
# that will contain different level types
variablesLevels[gribLine.shortName] = {}
variablesMetaData[gribLine.shortName] = {}
variablesMetaData[gribLine.shortName]['shortName'] = gribLine.shortName
variablesMetaData[gribLine.shortName]['units'] = gribLine.units
variablesMetaData[gribLine.shortName]['name'] = gribLine.name
# is this the first time this type of level is met ?
if gribLine.typeOfLevel not in \
variablesLevels[gribLine.shortName] :
# create a list that will contain the level labels
variablesLevels[gribLine.shortName][gribLine.typeOfLevel] = []
# append the level label to the variable / level type
variablesLevels[gribLine.shortName][gribLine.typeOfLevel]\
.append(gribLine.level)
# move to the next line
gribLine = rawFile.readline()
if gribLine == None :
timeDimension = False
break
# find the longest vertical axis
maxLevelNumber = 0
for variableName, levelKinds in variablesLevels.iteritems() :
for levelType, levels in levelKinds.iteritems() :
# does levels look like a proper axis ?
if len(levels) > 1 :
variablesLevels[variableName][levelType] \
= aa.Vertical(np.array(levels), levelType)
# is levels longer than the previous longest axis ?
if len(levels) > maxLevelNumber :
maxLevelNumber = len(levels)
mainLevels = aa.Vertical(np.array(levels), levelType)
# the longest vertical axis gets to be the file's vertical axis
self.axes['level'] = mainLevels
if timeDimension :
#############
# TIME AXIS #
#############
# "seek/tell" index starts with 1
# but we've moved on the next instant at the end of the while loop
# hence the minus one
linesPerInstant = rawFile.tell() - 1
# determine the interval between two samples
secondInstant = datetime(gribLine.year, gribLine.month, gribLine.day,
gribLine.hour, gribLine.minute, gribLine.second)
timeStep = secondInstant - firstInstant
# go to the end of the file
rawFile.seek(0, 2)
lastIndex = rawFile.tell()
# this index points at the last message
# e.g. f.message(lastIndex) returns the last message
# indices start at 1 meaning that lastIndex is also the
# number of messages in the file
self.axes['time'] = aa.TimeAxis(
np.array([firstInstant + timeIndex*timeStep
for timeIndex in range(lastIndex/linesPerInstant)]), None)
# check consistency
gribLine = rawFile.message(lastIndex)
lastInstant = datetime(gribLine.year, gribLine.month, gribLine.day,
gribLine.hour, gribLine.minute, gribLine.second)
if lastInstant != self.dts[-1] or \
lastIndex % linesPerInstant != 0 :
raise Exception, "Error in time axis"
rawFile.rewind()
#############
# VARIABLES #
#############
for variableName, levelKinds in variablesLevels.iteritems() :
for levelType, verticalAxis in levelKinds.iteritems() :
conditions = {'shortName' : variableName.encode('ascii'),
'typeOfLevel' : levelType.encode('ascii')}
axes = aa.Axes()
if timeDimension :
axes['time'] = self.axes['time']
else :
conditions['time'] = firstInstant
variableLabel = variableName + '_' + levelType
# does this variable have a vertical extension ?
# it may not be the file's vertical axis
if len(verticalAxis) > 1 :
axes['level'] = verticalAxis
# in case of homonyms, only the variable with the main
# vertical axis gets to keep the original shortname
if verticalAxis.units == mainLevels.units :
variableLabel = variableName
else :
# flat level i.e. 2D data
# the condition is a list to be iterable
conditions['level'] = verticalAxis
# no ambiguity
if len(levelKinds) == 1 :
variableLabel = variableName
axes['latitude'] = self.axes['latitude']
axes['longitude'] = self.axes['longitude']
self.variables[variableLabel] = \
Variable(axes, variablesMetaData[variableName],
conditions, fileName)
##################
# PICKLE & INDEX #
##################
rawFile.close()
pickleFile = open(fileName+'.p', 'w')
#import pdb ; pdb.set_trace()
pickle.dump(self, pickleFile)
pickleFile.close()
gribIndex = pygrib.index(filePath,
'shortName', 'level', 'typeOfLevel',
'year', 'month', 'day', 'hour')
gribIndex.write(fileName+'.idx')
gribIndex.close()
def plotter(fdict):
""" Go """
ctx = get_autoplot_context(fdict, get_description())
ts = ctx["ts"].replace(tzinfo=pytz.utc)
hour = int(ctx["hour"])
ilabel = ctx["ilabel"] == "yes"
plot = MapPlot(
sector=ctx["t"],
continentalcolor="white",
state=ctx["state"],
cwa=ctx["wfo"],
title=("NWS RFC %s Hour Flash Flood Guidance on %s UTC") %
(hour, ts.strftime("%-d %b %Y %H")),
subtitle=("Estimated amount of %s Rainfall "
"needed for non-urban Flash Flooding to commence") %
(HOURS[ctx["hour"]], ),
)
cmap = plt.get_cmap(ctx["cmap"])
bins = [
0.01,
0.6,
0.8,
1.0,
1.2,
1.4,
1.6,
1.8,
2.0,
2.25,
2.5,
2.75,
3.0,
3.5,
4.0,
5.0,
]
if ts.year < 2019:
column = "hour%02i" % (hour, )
pgconn = get_dbconn("postgis")
df = read_sql(
"""
WITH data as (
SELECT ugc, rank() OVER (PARTITION by ugc ORDER by valid DESC),
hour01, hour03, hour06, hour12, hour24
from ffg WHERE valid >= %s and valid <= %s)
SELECT *, substr(ugc, 3, 1) as ztype from data where rank = 1
""",
pgconn,
params=(ts - datetime.timedelta(hours=24), ts),
index_col="ugc",
)
df2 = df[df["ztype"] == "C"]
plot.fill_ugcs(
df2[column].to_dict(),
bins,
cmap=cmap,
plotmissing=False,
ilabel=ilabel,
)
df2 = df[df["ztype"] == "Z"]
plot.fill_ugcs(
df2[column].to_dict(),
bins,
cmap=cmap,
plotmissing=False,
units="inches",
ilabel=ilabel,
)
else:
# use grib data
ts -= datetime.timedelta(hours=(ts.hour % 6))
ts = ts.replace(minute=0)
fn = None
for offset in range(0, 24, 4):
ts2 = ts - datetime.timedelta(hours=offset)
testfn = ts2.strftime(("/mesonet/ARCHIVE/data/%Y/%m/%d/model/ffg/"
"5kmffg_%Y%m%d00.grib2"))
if os.path.isfile(testfn):
fn = testfn
break
if fn is None:
raise NoDataFound("No valid grib data found!")
grbs = pygrib.index(fn, "stepRange")
grb = grbs.select(stepRange="0-%s" % (hour, ))[0]
lats, lons = grb.latlons()
data = (masked_array(grb.values,
data_units=units("mm")).to(units("inch")).m)
plot.pcolormesh(lons, lats, data, bins, cmap=cmap)
if ilabel:
plot.drawcounties()
df = pd.DataFrame()
return plot.fig, df
param_list[l][1] = 'sfc'
elif param_list[l][1] == 109:
param_list[l][1] = 'ml'
searchlist = [param_list[item][0:3] for item in range(len(param_list))]
levellist = [param_list[item][3] for item in range(len(param_list))]
print searchlist, levellist
#sys.exit(1)
a = post_data()
print a
for file in gribec_files[0:3]:
print file
grbindx = pygrib.index(file, 'indicatorOfParameter',
'indicatorOfTypeOfLevel', 'level')
print grbindx
#selected_grbs=grbindx.select(indicatorOfParameter=61,indicatorOfTypeOfLevel='sfc',level=[456,457]) #,typeOfLevel='isobaricInhPa',level=500)
#for grb in selected_grbs:
# print grb
# print grb.indicatorOfTypeOfLevel
grbs = pygrib.open(file)
print grbs(indicatorOfParameter=61,
indicatorOfTypeOfLevel='sfc',
level=[456, 457])
sys.exit(1)
for grb in grbs:
#if (grb.indicatorOfParameter,grb.indicatorOfTypeOfLevel,grb.level) in param_list:
#if grb.indicatorOfParameter in (11,61) and grb.level in (2,457):
m.drawcountries(linewidth=.5, color='k')
# Map/figure has been set up here, save axes instances for use again later
if par == 1:
keep_ax_lst_1 = ax.get_children()[:]
elif par == 2:
keep_ax_lst_2 = ax.get_children()[:]
par += 1
par = 1
for j in range(len(date_list)):
ymd=dtime.strftime("%Y%m%d")
fhour=date_list[j].strftime("%H")
fhr = str(fhours[j]).zfill(2)
prodind = pygrib.index(DATA_DIR_1+'/fv3lam.t'+cyc+'z.conus.f'+fhr+'.grib2','name','level')
paraind = pygrib.index(DATA_DIR_2+'/fv3lam.t'+cyc+'z.conus.f'+fhr+'.grib2','name','level')
refprod=np.asarray(prodind.select(name='Maximum/Composite radar reflectivity',level=0)[0].values)
rainprod=np.asarray(prodind.select(name='Categorical rain',level=0)[0].values)
frprod=np.asarray(prodind.select(name='Categorical freezing rain',level=0)[0].values)
plprod=np.asarray(prodind.select(name='Categorical ice pellets',level=0)[0].values)
snprod=np.asarray(prodind.select(name='Categorical snow',level=0)[0].values)
refpara=np.asarray(paraind.select(name='Maximum/Composite radar reflectivity',level=0)[0].values)
rainpara=np.asarray(paraind.select(name='Categorical rain',level=0)[0].values)
frpara=np.asarray(paraind.select(name='Categorical freezing rain',level=0)[0].values)
plpara=np.asarray(paraind.select(name='Categorical ice pellets',level=0)[0].values)
snpara=np.asarray(paraind.select(name='Categorical snow',level=0)[0].values)
typespara=np.zeros(frpara.shape)
typespara[rainpara==1]=typespara[rainpara==1]+1
def main(out_folder=""):
out_folder = os.path.expanduser(out_folder)
if not os.path.isdir(out_folder):
os.makedirs(out_folder)
path_template = "/RECH/data/Driving_data/Offline/ERA-Interim_0.75/Grib_files/ERA_Interim_0.75d_3h_snowfall_{}.grib"
# ind = pygrib.index(path, "year", "month")
# fcst_hour_list = [3, 6, 9, 12]
fcst_hour_list = [6, 12]
start_year = 1979
end_year = 2012
varname = "SN"
dateo = datetime(1979, 1, 1)
npas = 1
for year in range(start_year, end_year + 1):
path = path_template.format(year)
ind = pygrib.index(path, "year", "month")
for month in range(1, 13):
fname = "ERAI_0.75_{}h_{}{:02d}.rpn".format(fcst_hour_list[0], year, month)
fpath = os.path.join(out_folder, fname)
if os.path.isfile(fpath):
print("{} -- already exists, delete to regenerate".format(fpath))
raise Exception("Unfortunately the program cannot be run in parallel for consistency of the data!"
" Please delete all previously generated files before proceeding.")
r_out = RPN(fpath, "w")
for i, grb_message in enumerate(ind.select(year=year, month=month)):
if grb_message.startStep not in fcst_hour_list:
print("Start step is: ", grb_message.startStep)
continue
print(grb_message)
# for k in grb_message.keys():
# print "{}: {}".format(k, grb_message[k])
print(grb_message.startStep)
print([grb_message[k] for k in ["year", "month", "day", "hour"]])
print(grb_message.validityDate, grb_message.validityTime)
if grb_message.startStep == fcst_hour_list[0]:
data_previous = 0
# print np.sum((grb_message.values < data_previous) & (grb_message.values >= 0)), \
# np.sum((grb_message.values >= data_previous) & (data_previous >= 0))
data = (grb_message.values - data_previous) / float(fcst_hour_list[0] * 3600)
data[data <= 1.0e-16] = 0
data_previous = grb_message.values[:, :]
# add a column at the end
data_ext = np.zeros((data.shape[0], data.shape[1] + 1))
data_ext[:, :-1] = data
data_ext[:, -1] = data[:, 0]
data = data_ext
data = np.flipud(data[:, :])
print(data.shape)
r_out.write_2D_field(name=varname,
data=data.transpose(), ip=[0, npas * fcst_hour_list[0], 0],
ig=[0] * 4,
npas=npas, deet=3600 * fcst_hour_list[0], label="ERAI075", dateo=dateo,
grid_type="B", typ_var="P",
nbits=-32, data_type=data_types.compressed_floating_point)
npas += 1
r_out.close()
# plt.figure()
# plt.plot(np.mean(month_data, axis=1).mean(axis=1), marker="s")
# plt.show()
ind.close()
def get_grib_field_lll(infile, var_name): # LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL
'''
Extracts longitude, latitude, level and one 2d or 3d-field from grib file
Parameters
----------
infile : str
input filename of the concerned grib1 file
var_name : str
short name of variable in infile
Returns
--------
lon : numpy array
longitude
lat : numpy array
latitude
lev : list
level list
var : numpy array
field of variable var_name
'''
# open and read from COSMOS file -------------------------------------
Indx = pygrib.index(infile, 'shortName')
# get certain variable ...............................................
try:
Cont = Indx.select(shortName=var_name)
except:
print(('ERROR: variable ' + var_name + ' not in ' + infile))
return
Nx = Cont[0].Ni
Ny = Cont[0].Nj
Nz = np.shape(Cont)[0]
lat, lon = Cont[0].latlons()
lat = np.transpose(lat)
lon = np.transpose(lon)
if Nz > 1:
var = np.zeros([Nx, Ny, Nz])
lev = np.zeros([Nz])
n = 0
for layer in Cont:
var[:, :, n] = np.transpose(layer.values)
lev[n] = layer.level
n = n + 1
# sort fields in vertical
lev = np.array(lev)
ils = lev.argsort()
l = lev[ils]
v = var[..., ils]
else:
layer = Cont[0]
v = np.transpose(layer.values)
l = [
layer.level,
]
return lon, lat, l, v # TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT
#!/usr/bin/env python
#Martin Iglesias Github SudestadaARG
import numpy as np
import pygrib
idx = pygrib.index('gfs_4_20151201_0000_000.grb2', 'name', 'typeOfLevel',
'level')
t = idx.select(name="Temperature", typeOfLevel="isobaricInhPa", level='500')
lat, lon = t[0].latlons()
np.save('lon2d_gfs05.npy', lon)
np.save('lat2d_gfs05.npy', lat)
# %%
import pygrib
# for ncep files, just use one to get variable names
cmc_file = '/Users/rsteinhart/DATA/test_data/NAEFS/cdo_test/merged_files/cmc_merged.t00z.pgrb2f384' #for example
# open the grib file
grbs = pygrib.open(cmc_file)
# grbindex = pygrib.index(cmc_file)
# create a new text file ("x" creates the file)
# switch to "a" to append text file or "w" to overwrite the content of the file
cmc_var_index = open("cmc_var_index.txt", "x")
grbindex = pygrib.index(cmc_file, 'shortName', 'typeOfLevel', 'level')
# print the inventory of the file to a text file
#grbindex.seek(0)
for grb in grbs:
# grbindex = pygrib.index(cmc_file,'grb')
selected_grbs = grbindex.select(shortName=grb)
grbindex = str(grbindex)
cmc_var_index.write(grbindex)
cmc_var_index.write("\n")
# %%
# Variables not included in ncep
if (errflag):
retry = retry - 1
if (retry):
print(" Retrying...", file=sys.stderr)
time.sleep(RETRY_DELAY)
else:
print(" Giving up.", file=sys.stderr)
print("Failed URL was: ", file=sys.stderr)
print(request_url, file=sys.stderr)
exit(1)
else:
break
f = open("temp.grb", 'wb')
f.write(r.content)
f.flush()
grbindx = pygrib.index("temp.grb", "name", "level")
#
# Turn the grib2 data into am layers. In a first pass, interpolate
# at each pressure level to lat,lon. For height and temperature,
# insert BADVAL if the variable is missing or not defined on the level.
# For other variables, set missing values to zero.
#
u = (args.lat - bottomlat) / latlon_delta
v = (args.lon - leftlon) / latlon_delta
Pbase = []
z = []
T = []
o3_vmr = []
RH = []
cloud_lmr = []
def runPrediction(self):
lat_meter = 200.0/22226349.0
lon_meter = 1.0/78846.80572069259
try:
start_time = time.time()
grb_p = pygrib.open(self.data.params.files[0])
grb_p_test = pygrib.index(self.data.params.files[0], 'name', 'level')
grb_s = pygrib.open(self.data.params.files[1])
print "Open Files - ", time.time() - start_time, "seconds"
except:
print "Could not open valid wind files"
grb_p.seek(0)
print self.pressureAltitude(100)
start_time = time.time()
d_now = DataManager.DataManager()
d_now.u = grb_p.select(name='U component of wind')
d_now.v = grb_p.select(name='V component of wind')
print "Grib Selection - ", time.time() - start_time, "seconds"
# Assuming all data has the same shape and lat/lon axis
start_time = time.time()
shape = {}
shape['latitude'] = {}
shape['latitude']['min'] = d_now.u[0].latlons()[0].min()
shape['latitude']['max'] = d_now.u[0].latlons()[0].max()
shape['longitude'] = {}
shape['longitude']['min'] = d_now.u[0].latlons()[1].min()
shape['longitude']['max'] = d_now.u[0].latlons()[1].max()
print "Min/Max Calculation - ", time.time() - start_time, "seconds"
self.coordinates.append((self.data.params.launch_time, self.grbLat(shape, self.data.params.launch_lat), self.grbLon(shape, self.data.params.launch_lon)))
self.position.append((self.data.params.launch_time, self.data.params.launch_lat, self.data.params.launch_lon))
cur_u = grb_p.select(name='U component of wind', level=1000)
cur_v = grb_p.select(name='V component of wind', level=1000)
print cur_u
print self.coordinates
print self.position
# S0, we've got a now dict holding heights, u, and v components
# 1) Iterate through heights, given our altitude, find correct index
# 2) Use the same index for u and v wind components
# 3) ....
# 4) Predictions?
print self.data.params.profile
print len(self.data.params.profile)
cur_u = grb_p_test(name='U component of wind', level=200/100)
cur_v = grb_p_test(name='V component of wind', level=200/100)
current_values = {}
current_values['pressure'] = 200
descent = 0
# Could remove altitude/Pressure steps in the future to speed this stage up...
for balloon_time in sorted(self.data.params.profile.iterkeys()):
if balloon_time > calendar.timegm(self.data.params.burst_time.timetuple()):
descent = 1
pressure = self.pressureAltitude(self.data.params.profile[balloon_time])
if pressure != current_values['pressure']:
print "New Level: ", pressure
start_time = time.time()
cur_u = grb_p_test(name='U component of wind', level=pressure/100)
cur_v = grb_p_test(name='V component of wind', level=pressure/100)
print "Select Pressure - ", time.time() - start_time, "seconds"
current_values['pressure'] = pressure
# If we change pressure level, throw away all previous data
current_values['grbLat'] = -1
current_values['grbLon'] = -1
# If we're looking at a new .5 degree index into the file, get new values, otherwise continue
if self.grbLat(shape, self.position[-1][1]) != current_values['grbLat'] or self.grbLon(shape, self.position[-1][2]) != current_values['grbLon']:
start_time = time.time()
# Update Current Latitude and Longitude index in grb file
current_values['grbLat'] = self.grbLat(shape, self.position[-1][1])
current_values['grbLon'] = self.grbLon(shape, self.position[-1][2])
current_values['u'] = cur_u[0].values[self.grbLat(shape, self.position[-1][1])-2:self.grbLat(shape, self.position[-1][1]),
self.grbLon(shape, self.position[-1][2])-2:self.grbLon(shape, self.position[-1][2])]
current_values['v'] = cur_v[0].values[self.grbLat(shape, self.position[-1][1])-2:self.grbLat(shape, self.position[-1][1]),
self.grbLon(shape, self.position[-1][2])-2:self.grbLon(shape, self.position[-1][2])]
# Update Current U and V components of wind
current_values['u_hh'] = current_values['u'][1,1]
current_values['v_hh'] = current_values['v'][1,1]
current_values['u_ll'] = current_values['u'][0,0]
current_values['v_ll'] = current_values['v'][0,0]
current_values['u_lh'] = current_values['u'][0,1]
current_values['v_lh'] = current_values['v'][0,1]
current_values['u_hl'] = current_values['u'][1,0]
current_values['v_hl'] = current_values['v'][1,0]
print "Select Wind Speed - ", time.time() - start_time, "seconds"
u_move = self.squareInterpolate(self.grbLat(shape, self.position[-1][1])-1 * .5,
self.grbLon(shape, self.position[-1][2])-1 * .5,
self.grbLat(shape, self.position[-1][1]) * .5,
self.grbLon(shape, self.position[-1][2]) * .5,
self.position[-1][1],
self.position[-1][2],
current_values['u_ll'],
current_values['u_lh'],
current_values['u_hl'],
current_values['u_hh'])
v_move = self.squareInterpolate(self.grbLat(shape, self.position[-1][1])-1 * .5,
self.grbLon(shape, self.position[-1][2])-1 * .5,
self.grbLat(shape, self.position[-1][1]) * .5,
self.grbLon(shape, self.position[-1][2]) * .5,
self.position[-1][1],
self.position[-1][2],
current_values['v_ll'],
current_values['v_lh'],
current_values['v_hl'],
current_values['v_hh'])
u_degrees = u_move * self.data.params.ascent_step * lat_meter
v_degrees = v_move * self.data.params.ascent_step * lon_meter
self.position.append((0, self.position[-1][1] + u_degrees, self.position[-1][2] + v_degrees))
for point in self.position:
print point[1], ',', point[2]
ts_dict[source['shortname']] = ts.stationlist()
for station in stations:
stat = stations[station]
ts_dict[source['shortname']].add_station(
stat['name'], stat['latlon'], stat['stationid'])
grbindx_prev = None
for grib_file in filelist[source['shortname']]:
print '************ open file ', grib_file
print 'processing ', source['name'], time.clock() - t0, int(
time.clock() - t0) * '='
grbindx = pygrib.index(grib_file, 'indicatorOfParameter',
'indicatorOfTypeOfLevel', 'level',
'timeRangeIndicator')
if TS:
plottypes.point_extract(grbindx,
source,
stations,
ts_dict[source['shortname']],
outdir=topdir)
#sys.exit(1)
print 'Current gribs: ', grbindx, 'previous gribs: ', grbindx_prev
plottypes.prec_plot(grbindx,
grbindx_prev,
source,
# OUTPUT: List of values for the requested parameters, separated by space # # USAGE: get_grid_extents.py <file name> <list of parameters requested> # # [email protected] ## if len(sys.argv) > 2: file_name = sys.argv[1] params_requested = sys.argv[2:] else: print 'USAGE: %s file params' % sys.argv[0] sys.exit(1) if not os.path.exists(file_name): print 'File not found: ', file_name sys.exit(1) grbindx = pygrib.index(file_name, 'shortName', 'typeOfLevel', 'level') selected_grbs = grbindx.select(shortName='10u', typeOfLevel='heightAboveGround', level=10) if len(selected_grbs) == 0: print 'Did not find a u10 field in the given file' sys.exit(1) if len(selected_grbs) > 1: print 'More than one match, this is unexpected' sys.exit(1) grbmsg = selected_grbs[0] for param in params_requested: sys.stdout.write("%s " % grbmsg[param])
def __init__(self, gribFile):
self.gribFile = gribFile
self.grbIndx = pygrib.index(gribFile, 'name', 'typeOfLevel', 'level')
def _pressures(variable, path):
messages = pg.index(path, "name")
for message in messages.select(name=variable):
yield message["scaledValueOfFirstFixedSurface"]
messages.close()
def load_data(self):
"""
Loads data from grib2 file objects or list of grib2 file objects. Handles specific grib2 variable names
and grib2 message numbers.
Returns:
Array of data loaded from files in (time, y, x) dimensions, Units
"""
if self.variable in ['nldn', 'entln']:
data, units = self.load_lightning_data()
return data, units
if not self.file_objects:
print("No {0} model runs on {1}".format(self.member, self.run_date))
units = None
return self.data, units
for f, g_file in enumerate(self.file_objects):
grib = pygrib.open(g_file)
data_values = None
if type(self.variable) is int:
data_values = grib[self.variable].values
print(grib[self.variable])
if grib[self.variable].units == 'unknown':
Id = grib[self.variable].parameterNumber
# units = self.unknown_units[Id]
# else:
# units = grib[self.variable].units
elif type(self.variable) is str:
if '_' in self.variable:
# Multiple levels
variable = self.variable.split('_')[0]
level = self.variable.split('_')[1]
else:
# Only single level
variable = self.variable
level = None
message_keys = np.array([[message.name, message.shortName,
message.level, message.typeOfLevel] for message in grib])
##################################
# Unknown string variables
##################################
grib_data = []
if variable in self.unknown_names.values():
Id, units = self.format_grib_name(variable)
if level is None:
grib_data = pygrib.index(g_file, 'parameterNumber')(parameterNumber=Id)
elif level in message_keys[:, 2]:
grib_data = pygrib.index(g_file,
'parameterNumber', 'level')(parameterNumber=Id, level=level)
elif level in message_keys[:, 3]:
grib_data = pygrib.index(g_file,
'parameterNumber', 'typeofLevel')(parameterNumber=Id,
typeOfLevel=level)
else:
print('No {0} {1} grib message found for {2} {3}'.format(
self.run_date, self.member, variable, level))
continue
##################################
# Known string variables
##################################
if variable in message_keys[:, 0]:
if level is None:
grib_data = pygrib.index(g_file, 'name')(name=variable)
elif level in message_keys[:, 2]:
grib_data = pygrib.index(g_file,
'name', 'level')(name=variable, level=level)
elif level in message_keys[:, 3]:
grib_data = pygrib.index(g_file,
'name', 'typeOfLevel')(name=variable, typeOfLevel=level)
else:
print('No {0} {1} grib message found for {2} {3}'.format(
self.run_date, self.member, variable, level))
continue
if variable in message_keys[:, 1]:
if level is None:
grib_data = pygrib.index(g_file, 'shortName')(shortName=variable)
elif level in message_keys[:, 2]:
grib_data = pygrib.index(g_file,
'shortName', 'level')(shortName=variable, level=level)
elif level in message_keys[:, 3]:
grib_data = pygrib.index(g_file,
'shortName', 'typeOfLevel')(shortName=variable, typeOfLevel=level)
else:
print('No {0} {1} grib message found for {2} {3}'.format(
self.run_date, self.member, variable, level))
continue
if len(grib_data) > 1:
if variable in ['u', 'v', '10u', '10v']:
grib_short_names = [message.shortName for message in grib_data]
u_v_ind = grib_short_names.index(str(variable))
data_values = grib_data[u_v_ind].values
elif variable in ['U component of wind', 'V component of wind',
'10 metre U wind component', '10 metre V wind component']:
grib_names = [message.name for message in grib_data]
u_v_ind = grib_names.index(str(variable))
data_values = grib_data[u_v_ind].values
else:
print()
raise NameError(
"Multiple '{0}' records found for {1} {2}.\n Please rename with more description'".format(
self.variable, self.run_date, self.member))
else:
data_values = grib_data[0].values
grib.close()
if self.data is None:
self.data = np.empty((len(self.valid_dates), data_values.shape[0], data_values.shape[1]), dtype=float)
self.data[f] = data_values[:]
else:
self.data[f] = data_values[:]
return self.data, None
def __init__(self, filename, filename_info, filetype_info):
super(GRIBFileHandler, self).__init__(filename, filename_info, filetype_info)
self._msg_datasets = {}
self._start_time = None
self._end_time = None
try:
with pygrib.open(filename) as grib_file:
first_msg = grib_file.message(1)
last_msg = grib_file.message(grib_file.messages)
start_time = self._convert_datetime(
first_msg, 'validityDate', 'validityTime')
end_time = self._convert_datetime(
last_msg, 'validityDate', 'validityTime')
self._start_time = start_time
self._end_time = end_time
if 'keys' not in filetype_info:
self._analyze_messages(grib_file)
self._idx = None
else:
self._create_dataset_ids(filetype_info['keys'])
self._idx = pygrib.index(filename,
*filetype_info['keys'].keys())
except (RuntimeError, KeyError):
raise IOError("Unknown GRIB file format: {}".format(filename))
