snapshot...

modules/util/gfs_reader.py

+import datetime
+from datetime import timezone
+import glob
+import numpy as np
+import xarray as xr
+from util.util import value_to_index
+#import cartopy
+#from cartopy import *
+from metpy import *
+#import cartopy.crs as ccrs
+import random
+
+gfs_directory = '/apollo/cloud/Ancil_Data/clavrx_ancil_data/dynamic/gfs/'
+gfs_date_format = '%y%m%d'
+
+lon360 = True
+lat_coords = np.linspace(-90, 90, 361)
+lon_coords = np.linspace(0, 359.5, 720)
+
+
+class MyGenericException(Exception):
+    def __init__(self, message):
+        self.message = message
+
+
+def get_time_tuple_utc(timestamp):
+    dt_obj = datetime.datetime.fromtimestamp(timestamp, timezone.utc)
+    return dt_obj, dt_obj.timetuple()
+
+
+def get_bounding_gfs_files(timestamp):
+    dt_obj, time_tup = get_time_tuple_utc(timestamp)
+    date_str = dt_obj.strftime(gfs_date_format)
+    dt_obj0 = datetime.datetime.strptime(date_str, gfs_date_format).replace(tzinfo=timezone.utc)
+    dt_obj1 = dt_obj0 + datetime.timedelta(days=1)
+    date_str_1 = dt_obj1.strftime(gfs_date_format)
+
+    flist_0 = glob.glob(gfs_directory+'gfs.'+date_str+'??_F012.hdf')
+    flist_1 = glob.glob(gfs_directory+'gfs.'+date_str_1+'00_F012.hdf')
+    if (len(flist_0) == 0) or (len(flist_1) == 0):
+        return None, None, None, None
+    filelist = flist_0 + flist_1
+
+    ftimes = []
+    for pname in filelist:  # TODO: make better with regular expressions (someday)
+        fname = os.path.split(pname)[1]
+        toks = fname.split('.')
+        tstr = toks[1].split('_')[0]
+        dto = datetime.datetime.strptime(tstr, gfs_date_format+'%H').replace(tzinfo=timezone.utc)
+        ftimes.append(dto.timestamp())
+
+    tarr = np.array(ftimes)
+    sidxs = tarr.argsort()
+
+    farr = np.array(filelist)
+    farr = farr[sidxs]
+    ftimes = tarr[sidxs]
+
+    above = ftimes >= timestamp
+    if not above.any():
+        return None, None, None, None
+    tR = ftimes[above].min()
+
+    below = ftimes <= timestamp
+    if not below.any():
+        return None, None, None, None
+    tL = ftimes[below].max()
+
+    iL = np.searchsorted(ftimes, tL, 'left')
+    iR = iL + 1
+
+    fList = farr.tolist()
+
+    return fList[iL], ftimes[iL], fList[iR], ftimes[iR]
+
+
+def get_horz_layer(xr_dataset, fld_name, press, lon_lo, lon_hi, lat_lo, lat_hi):
+    plevs = xr_dataset['pressure levels']
+    fld = xr_dataset[fld_name]
+
+    if lon360:
+        if lon_lo < 0:
+            lon_lo += 360
+        if lon_hi < 0:
+            lon_hi += 360
+
+    x_lo = value_to_index(lon_coords, lon_lo)
+    x_hi = value_to_index(lon_coords, lon_hi)
+    y_lo = value_to_index(lat_coords, lat_lo)
+    y_hi = value_to_index(lat_coords, lat_hi)
+    p_idx = value_to_index(plevs.values, press)
+
+    sub_fld = fld[y_lo:y_hi, x_lo:x_hi, p_idx]
+    sub_fld = sub_fld.expand_dims('channel')
+
+    sub_fld = sub_fld.assign_coords(channel=[fld_name], fakeDim2=lon_coords[x_lo:x_hi], fakeDim1=lat_coords[y_lo:y_hi])
+
+    return sub_fld
+
+
+def get_horz_layer_s(xr_dataset, fld_names, press, lon_lo=None, lon_hi=None, lat_lo=None, lat_hi=None):
+    plevs = xr_dataset['pressure levels']
+
+    if lon360:
+        if lon_lo < 0:
+            lon_lo += 360
+        if lon_hi < 0:
+            lon_hi += 360
+
+    x_lo = value_to_index(lon_coords, lon_lo)
+    x_hi = value_to_index(lon_coords, lon_hi)
+    y_lo = value_to_index(lat_coords, lat_lo)
+    y_hi = value_to_index(lat_coords, lat_hi)
+    p_idx = value_to_index(plevs.values, press)
+
+    sub_fld_s = []
+    for fld_name in fld_names:
+        fld = xr_dataset[fld_name]
+        sub_fld = fld[y_lo:y_hi, x_lo:x_hi, p_idx]
+        sub_fld_s.append(sub_fld)
+
+    sub_fld = xr.concat(sub_fld_s, 'channel')
+    sub_fld = sub_fld.assign_coords(channel=fld_names, fakeDim2=lon_coords[x_lo:x_hi], fakeDim1=lat_coords[y_lo:y_hi])
+
+    return sub_fld
+
+
+def get_vert_profile(xr_dataset, fld_name, lons, lats):
+    if lon360:  # convert -180/+180 to 0,360
+        lons = np.where(lons < 0, lons + 360, lons)
+
+    plevs = xr_dataset['pressure levels']
+
+    fld = xr_dataset[fld_name]
+    fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords, fakeDim0=plevs)
+
+    dim2 = xr.DataArray(lons, dims='k')
+    dim1 = xr.DataArray(lats, dims='k')
+
+    intrp_fld = fld.interp(fakeDim1=dim1, fakeDim2=dim2, fakeDim0=plevs, method='linear')
+
+    return intrp_fld
+
+
+def get_vert_profile_s(xr_dataset, fld_name_s, lons, lats):
+    if lon360:  # convert -180,+180 to 0,360
+        lons = np.where(lons < 0, lons + 360, lons)
+
+    plevs = xr_dataset['pressure levels']
+
+    fld_s = []
+    for fld_name in fld_name_s:
+        fld = xr_dataset[fld_name]
+        fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords, fakeDim0=plevs)
+        fld_s.append(fld)
+
+    fld = xr.concat(fld_s, 'fld_dim')
+
+    dim2 = xr.DataArray(lons, dims='k')
+    dim1 = xr.DataArray(lats, dims='k')
+
+    intrp_fld = fld.interp(fakeDim1=dim1, fakeDim2=dim2, fakeDim0=plevs, method='linear')
+
+    return intrp_fld
+
+
+def get_point(xr_dataset, fld_name, lons, lats):
+    if lon360:  # convert -180/+180 to 0,360
+        lons = np.where(lons < 0, lons + 360, lons) # convert -180,180 to 0,360
+
+    lat_coords = np.linspace(-90, 90, xr_dataset.fakeDim1.size)
+    lon_coords = np.linspace(0, 359.5, xr_dataset.fakeDim2.size)
+
+    fld = xr_dataset[fld_name]
+    fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords)
+
+    dim1 = xr.DataArray(lats, dims='k')
+    dim2 = xr.DataArray(lons, dims='k')
+
+    intrp_fld = fld.interp(fakeDim1=dim1, fakeDim2=dim2, method='linear')
+
+    return intrp_fld
+
+
+def get_time_interpolated_vert_profile(ds_0, ds_1, time0, time1, fld_name, time, lons, lats):
+
+    prof_0 = get_vert_profile(ds_0, fld_name, lons, lats)
+    prof_1 = get_vert_profile(ds_1, fld_name, lons, lats)
+
+    prof = xr.concat([prof_0, prof_1], 'time')
+    prof = prof.assign_coords(time=[time0, time1])
+
+    intrp_prof = prof.interp(time=time, method='linear')
+
+    intrp_prof = intrp_prof.values
+
+    return intrp_prof
+
+
+def get_time_interpolated_vert_profile_s(ds_0, ds_1, time0, time1, fld_name_s, time, lons, lats):
+
+    prof_0 = get_vert_profile_s(ds_0, fld_name_s, lons, lats)
+    prof_1 = get_vert_profile_s(ds_1, fld_name_s, lons, lats)
+
+    prof = xr.concat([prof_0, prof_1], 'time')
+    prof = prof.assign_coords(time=[time0, time1])
+
+    intrp_prof = prof.interp(time=time, method='linear')
+
+    intrp_prof = intrp_prof.values
+
+    return intrp_prof
+
+
+def get_time_interpolated_point(ds_0, ds_1, time0, time1, fld_name, time, lons, lats):
+    vals_0 = get_point(ds_0, fld_name, lons, lats)
+    vals_1 = get_point(ds_1, fld_name, lons, lats)
+
+    vals = xr.concat([vals_0, vals_1], 'time')
+    vals = vals.assign_coords(time=[time0, time1])
+
+    intrp_vals = vals.interp(time=time, method='linear')
+
+    intrp_vals = intrp_vals.values
+
+    return intrp_vals