Skip to content
Snippets Groups Projects
util.py 42.8 KiB
Newer Older
tomrink's avatar
tomrink committed
import metpy
rink's avatar
rink committed
import numpy as np
tomrink's avatar
tomrink committed
import xarray as xr
tomrink's avatar
tomrink committed
import datetime
tomrink's avatar
tomrink committed
from datetime import timezone
rink's avatar
rink committed
from metpy.units import units
from metpy.calc import thickness_hydrostatic
tomrink's avatar
tomrink committed
from collections import namedtuple
tomrink's avatar
tomrink committed
import os
tomrink's avatar
tomrink committed
import h5py
import pickle
from netCDF4 import Dataset
from util.setup import ancillary_path
tomrink's avatar
tomrink committed
from scipy.interpolate import RectBivariateSpline
tomrink's avatar
tomrink committed

LatLonTuple = namedtuple('LatLonTuple', ['lat', 'lon'])
rink's avatar
rink committed

tomrink's avatar
tomrink committed
homedir = os.path.expanduser('~') + '/'

rink's avatar
rink committed

tomrink's avatar
tomrink committed
class GenericException(Exception):
rink's avatar
rink committed
    def __init__(self, message):
        self.message = message
rink's avatar
rink committed


tomrink's avatar
tomrink committed
class EarlyStop:
    def __init__(self, window_length=3, patience=5):
        self.patience = patience
tomrink's avatar
tomrink committed
        self.min = np.finfo(np.single).max
tomrink's avatar
tomrink committed
        self.cnt = 0
        self.cnt_wait = 0
        self.window = np.zeros(window_length, dtype=np.single)
        self.window.fill(np.nan)

    def check_stop(self, value):
        self.window[:-1] = self.window[1:]
        self.window[-1] = value

        if np.any(np.isnan(self.window)):
            return False

        ave = np.mean(self.window)
        if ave < self.min:
            self.min = ave
            self.cnt_wait = 0
            return False
        else:
            self.cnt_wait += 1

        if self.cnt_wait > self.patience:
            return True
        else:
            return False


tomrink's avatar
tomrink committed
def get_time_tuple_utc(timestamp):
    dt_obj = datetime.datetime.fromtimestamp(timestamp, timezone.utc)
    return dt_obj, dt_obj.timetuple()


def get_datetime_obj(dt_str, format_code='%Y-%m-%d_%H:%M'):
    dto = datetime.datetime.strptime(dt_str, format_code).replace(tzinfo=timezone.utc)
    return dto


tomrink's avatar
tomrink committed
def get_timestamp(dt_str, format_code='%Y-%m-%d_%H:%M'):
    dto = datetime.datetime.strptime(dt_str, format_code).replace(tzinfo=timezone.utc)
    ts = dto.timestamp()
    return ts


tomrink's avatar
tomrink committed
def add_time_range_to_filename(pathname, tstart, tend=None):
    filename = os.path.split(pathname)[1]
    w_o_ext, ext = os.path.splitext(filename)

tomrink's avatar
tomrink committed
    dt_obj, _ = get_time_tuple_utc(tstart)
    str_start = dt_obj.strftime('%Y%m%d%H')
tomrink's avatar
tomrink committed
    filename = w_o_ext+'_'+str_start
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    if tend is not None:
        dt_obj, _ = get_time_tuple_utc(tend)
        str_end = dt_obj.strftime('%Y%m%d%H')
        filename = filename+'_'+str_end
    filename = filename+ext
tomrink's avatar
tomrink committed

    path = os.path.split(pathname)[0]
    path = path+'/'+filename
    return path


rink's avatar
rink committed
def haversine_np(lon1, lat1, lon2, lat2):
    """
    Calculate the great circle distance between two points
    on the earth (specified in decimal degrees)

    (lon1, lat1) must be broadcastable with (lon2, lat2).

    """
    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])

    dlon = lon2 - lon1
    dlat = lat2 - lat1

    a = np.sin(dlat/2.0)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2.0)**2

    c = 2 * np.arcsin(np.sqrt(a))
    km = 6367 * c
    return km


def bin_data_by(a, b, bin_ranges):
    nbins = len(bin_ranges)
    binned_data = []

    for i in range(nbins):
        rng = bin_ranges[i]
        idxs = (b >= rng[0]) & (b < rng[1])
        binned_data.append(a[idxs])

    return binned_data


tomrink's avatar
tomrink committed
def bin_data_by_edges(a, b, edges):
    nbins = len(edges) - 1
    binned_data = []

    for i in range(nbins):
        idxs = (b >= edges[i]) & (b < edges[i+1])
        binned_data.append(a[idxs])

    return binned_data


rink's avatar
rink committed
def get_bin_ranges(lop, hip, bin_size=100):
    bin_ranges = []
    delp = hip - lop
    nbins = int(delp/bin_size)

    for i in range(nbins):
        rng = [lop + i*bin_size, lop + i*bin_size + bin_size]
        bin_ranges.append(rng)

    return bin_ranges


# t must be monotonic increasing
def get_breaks(t, threshold):
    t_0 = t[0:t.shape[0]-1]
    t_1 = t[1:t.shape[0]]
    d = t_1 - t_0
    idxs = np.nonzero(d > threshold)
    return idxs

rink's avatar
rink committed

tomrink's avatar
tomrink committed
# return indexes of ts where value is within ts[i] - threshold < value < ts[i] + threshold
tomrink's avatar
tomrink committed
# eventually, if necessary, fully vectorize (numpy) this is possible
tomrink's avatar
tomrink committed
# threshold units: seconds
tomrink's avatar
tomrink committed
def get_indexes_within_threshold(ts, value, threshold):
    idx_s = []
tomrink's avatar
tomrink committed
    t_s = []
tomrink's avatar
tomrink committed
    for k, v in enumerate(ts):
tomrink's avatar
tomrink committed
        if (ts[k] - threshold) <= value <= (ts[k] + threshold):
tomrink's avatar
tomrink committed
            idx_s.append(k)
tomrink's avatar
tomrink committed
            t_s.append(v)
    return idx_s, t_s
tomrink's avatar
tomrink committed


tomrink's avatar
tomrink committed
def pressure_to_altitude(pres, temp, prof_pres, prof_temp, sfc_pres=None, sfc_temp=None, sfc_elev=0):
rink's avatar
rink committed
    if not np.all(np.diff(prof_pres) > 0):
tomrink's avatar
tomrink committed
        raise GenericException("target pressure profile must be monotonic increasing")
rink's avatar
rink committed

    if pres < prof_pres[0]:
tomrink's avatar
tomrink committed
        raise GenericException("target pressure less than top of pressure profile")
rink's avatar
rink committed

    if temp is None:
        temp = np.interp(pres, prof_pres, prof_temp)

    i_top = np.argmax(np.extract(prof_pres <= pres, prof_pres)) + 1

    pres_s = prof_pres.tolist()
    temp_s = prof_temp.tolist()

    pres_s = [pres] + pres_s[i_top:]
    temp_s = [temp] + temp_s[i_top:]

tomrink's avatar
tomrink committed
    if sfc_pres is not None:
tomrink's avatar
tomrink committed
        if pres > sfc_pres:  # incoming pressure below surface
tomrink's avatar
tomrink committed
            return -1

        prof_pres = np.array(pres_s)
        prof_temp = np.array(temp_s)
        i_bot = prof_pres.shape[0] - 1

rink's avatar
rink committed
        if sfc_pres > prof_pres[i_bot]:  # surface below profile bottom
tomrink's avatar
tomrink committed
            pres_s = pres_s + [sfc_pres]
            temp_s = temp_s + [sfc_temp]
        else:
            idx = np.argmax(np.extract(prof_pres < sfc_pres, prof_pres))
            if sfc_temp is None:
                sfc_temp = np.interp(sfc_pres, prof_pres, prof_temp)
            pres_s = prof_pres.tolist()
            temp_s = prof_temp.tolist()
            pres_s = pres_s[0:idx+1] + [sfc_pres]
            temp_s = temp_s[0:idx+1] + [sfc_temp]
rink's avatar
rink committed

    prof_pres = np.array(pres_s)
    prof_temp = np.array(temp_s)

rink's avatar
rink committed
    prof_pres = prof_pres[::-1]
    prof_temp = prof_temp[::-1]

rink's avatar
rink committed
    prof_pres = prof_pres * units.hectopascal
    prof_temp = prof_temp * units.kelvin
    sfc_elev = sfc_elev * units.meter

    z = thickness_hydrostatic(prof_pres, prof_temp) + sfc_elev

    return z


# http://fourier.eng.hmc.edu/e176/lectures/NM/node25.html
def minimize_quadratic(xa, xb, xc, ya, yb, yc):
tomrink's avatar
tomrink committed
    x_m = xb + 0.5*(((ya-yb)*(xc-xb)*(xc-xb) - (yc-yb)*(xb-xa)*(xb-xa)) / ((ya-yb)*(xc-xb) + (yc-yb)*(xb-xa)))
tomrink's avatar
tomrink committed
    return x_m


tomrink's avatar
tomrink committed
# Return index of nda closest to value. nda must be 1d
tomrink's avatar
tomrink committed
def value_to_index(nda, value):
    diff = np.abs(nda - value)
    idx = np.argmin(diff)
tomrink's avatar
tomrink committed
    return idx


tomrink's avatar
tomrink committed
def find_bin_index(nda, value_s):
    idxs = np.arange(nda.shape[0])
tomrink's avatar
tomrink committed
    iL_s = np.zeros(value_s.shape[0])
    iL_s[:,] = -1
tomrink's avatar
tomrink committed
    for k, v in enumerate(value_s):
        above = v >= nda
        if not above.any():
            continue
tomrink's avatar
tomrink committed
        below = v < nda
        if not below.any():
            continue
tomrink's avatar
tomrink committed
        iL = idxs[above].max()
        iL_s[k] = iL
tomrink's avatar
tomrink committed
    return iL_s.astype(np.int32)
# array solzen must be degrees, missing values must NaN. For small roughly 50x50km regions only
def is_day(solzen, test_angle=80.0):
tomrink's avatar
tomrink committed
    solzen = solzen.flatten()
    solzen = solzen[np.invert(np.isnan(solzen))]
    if len(solzen) == 0 or np.sum(solzen <= test_angle) < len(solzen):
tomrink's avatar
tomrink committed
        return False
    else:
        return True


# array solzen must be degrees, missing values must NaN. For small roughly 50x50km regions only
def is_night(solzen, test_angle=100.0):
    solzen = solzen.flatten()
    solzen = solzen[np.invert(np.isnan(solzen))]
    if len(solzen) == 0 or np.sum(solzen >= test_angle) < len(solzen):
        return False
    else:
tomrink's avatar
tomrink committed
        return True


def check_oblique(satzen, test_angle=70.0):
    satzen = satzen.flatten()
    satzen = satzen[np.invert(np.isnan(satzen))]
tomrink's avatar
tomrink committed
    if len(satzen) == 0 or np.sum(satzen <= test_angle) < len(satzen):
tomrink's avatar
tomrink committed
def get_median(tile_2d):
tomrink's avatar
tomrink committed
    tile = tile_2d.flatten()
tomrink's avatar
tomrink committed
    return np.nanmedian(tile)
tomrink's avatar
tomrink committed


def get_grid_values_all(h5f, grid_name, scale_factor_name='scale_factor', add_offset_name='add_offset',
tomrink's avatar
tomrink committed
                        fill_value_name='_FillValue', range_name='actual_range', fill_value=None, stride=None):
tomrink's avatar
tomrink committed
    hfds = h5f[grid_name]
    attrs = hfds.attrs

    if attrs is None:
        raise GenericException('No attributes object for: '+grid_name)

tomrink's avatar
tomrink committed
    if stride is None:
        grd_vals = hfds[:,]
    else:
        grd_vals = hfds[::stride, ::stride]

    if fill_value is not None:
        grd_vals = np.where(grd_vals == fill_value, np.nan, grd_vals)
tomrink's avatar
tomrink committed

    if scale_factor_name is not None:
        attr = attrs.get(scale_factor_name)
        if attr is None:
            raise GenericException('Attribute: '+scale_factor_name+' not found for variable: '+grid_name)
tomrink's avatar
tomrink committed
        if np.isscalar(attr):
            scale_factor = attr
        else:
            scale_factor = attr[0]
tomrink's avatar
tomrink committed
        grd_vals = grd_vals * scale_factor

    if add_offset_name is not None:
        attr = attrs.get(add_offset_name)
        if attr is None:
            raise GenericException('Attribute: '+add_offset_name+' not found for variable: '+grid_name)
tomrink's avatar
tomrink committed
        if np.isscalar(attr):
            add_offset = attr
        else:
            add_offset = attr[0]
tomrink's avatar
tomrink committed
        grd_vals = grd_vals + add_offset

    if range_name is not None:
        attr = attrs.get(range_name)
tomrink's avatar
tomrink committed
        if attr is not None:
            low = attr[0]
            high = attr[1]
            grd_vals = np.where(grd_vals < low, np.nan, grd_vals)
            grd_vals = np.where(grd_vals > high, np.nan, grd_vals)
    elif fill_value_name is not None:
        attr = attrs.get(fill_value_name)
tomrink's avatar
tomrink committed
        if attr is not None:
            if np.isscalar(attr):
                fill_value = attr
            else:
                fill_value = attr[0]
            grd_vals = np.where(grd_vals == fill_value, np.nan, grd_vals)
tomrink's avatar
tomrink committed
    return grd_vals
tomrink's avatar
tomrink committed


# dt_str_0: start datetime string in format YYYY-MM-DD_HH:MM
tomrink's avatar
tomrink committed
# dt_str_1: stop datetime string, if not None num_steps is computed
tomrink's avatar
tomrink committed
# format_code: default '%Y-%m-%d_%H:%M'
tomrink's avatar
tomrink committed
# num_steps with increment of days, hours, minutes or seconds
tomrink's avatar
tomrink committed
# dt_str_1 and num_steps cannot both be None
tomrink's avatar
tomrink committed
# return num_steps+1 lists of datetime strings and timestamps (edges of a numpy histogram)
tomrink's avatar
tomrink committed
def make_times(dt_str_0, dt_str_1=None, format_code='%Y-%m-%d_%H:%M', num_steps=None, days=None, hours=None, minutes=None, seconds=None):
tomrink's avatar
tomrink committed
    if days is not None:
        inc = 86400*days
    elif hours is not None:
        inc = 3600*hours
    elif minutes is not None:
        inc = 60*minutes
    else:
        inc = seconds

    dt_obj_s = []
    ts_s = []
tomrink's avatar
tomrink committed
    dto_0 = datetime.datetime.strptime(dt_str_0, format_code).replace(tzinfo=timezone.utc)
tomrink's avatar
tomrink committed
    ts_0 = dto_0.timestamp()
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    if dt_str_1 is not None:
tomrink's avatar
tomrink committed
        dto_1 = datetime.datetime.strptime(dt_str_1, format_code).replace(tzinfo=timezone.utc)
tomrink's avatar
tomrink committed
        ts_1 = dto_1.timestamp()
        num_steps = int((ts_1 - ts_0)/inc)
tomrink's avatar
tomrink committed

    dt_obj_s.append(dto_0)
    ts_s.append(ts_0)
    dto_last = dto_0
    for k in range(num_steps):
        dt_obj = dto_last + datetime.timedelta(seconds=inc)
        dt_obj_s.append(dt_obj)
        ts_s.append(dt_obj.timestamp())
        dto_last = dt_obj

tomrink's avatar
tomrink committed
    return dt_obj_s, ts_s


def make_histogram(values, edges):
    h = np.histogram(values, bins=edges)
tomrink's avatar
tomrink committed
    return h


def normalize(data, param, mean_std_dict, add_noise=False, noise_scale=1.0, seed=None):

    if mean_std_dict.get(param) is None:
        return data

    shape = data.shape
    data = data.flatten()

    mean, std, lo, hi = mean_std_dict.get(param)
    data -= mean
    data /= std

    if add_noise:
        if seed is not None:
            np.random.seed(seed)
        rnd = np.random.normal(loc=0, scale=noise_scale, size=data.size)
        data += rnd

    not_valid = np.isnan(data)
    data[not_valid] = 0

    data = np.reshape(data, shape)

tomrink's avatar
tomrink committed
    return data

tomrink's avatar
tomrink committed
def denormalize(data, param, mean_std_dict):

    if mean_std_dict.get(param) is None:
        return data

    shape = data.shape
    data = data.flatten()

    mean, std, lo, hi = mean_std_dict.get(param)
    data *= std
    data += mean

    data = np.reshape(data, shape)

    return data


tomrink's avatar
tomrink committed
def scale(data, param, mean_std_dict):

    if mean_std_dict.get(param) is None:
        return data

    shape = data.shape
    data = data.flatten()

tomrink's avatar
tomrink committed
    if mean_std_dict is None:
        lo = np.nanmin(data)
        hi = np.nanmax(data)

    _, _, lo, hi = mean_std_dict.get(param)
tomrink's avatar
tomrink committed

    data -= lo
    data /= (hi - lo)

    not_valid = np.isnan(data)
    data[not_valid] = 0

    data = np.reshape(data, shape)

    return data


f = open(ancillary_path+'geos_crs_goes16_FD.pkl', 'rb')
geos_goes16_fd = pickle.load(f)
f.close()

f = open(ancillary_path+'geos_crs_goes16_CONUS.pkl', 'rb')
geos_goes16_conus = pickle.load(f)
f.close()

f = open(ancillary_path+'geos_crs_H08_FD.pkl', 'rb')
geos_h08_fd = pickle.load(f)
f.close()


def get_cartopy_crs(satellite, domain):
    if satellite == 'GOES16':
        if domain == 'FD':
            geos = geos_goes16_fd
            xlen = 5424
            xmin = -5433893.0
tomrink's avatar
tomrink committed
            xmax = 5433893.0
            ylen = 5424
            ymin = -5433893.0
            ymax = 5433893.0
        elif domain == 'CONUS':
            geos = geos_goes16_conus
            xlen = 2500
            xmin = -3626269.5
            xmax = 1381770.0
            ylen = 1500
            ymin = 1584175.9
            ymax = 4588198.0
    elif satellite == 'H08':
        geos = geos_h08_fd
        xlen = 5500
        xmin = -5498.99990119
        xmax = 5498.99990119
        ylen = 5500
        ymin = -5498.99990119
        ymax = 5498.99990119

    return geos, xlen, xmin, xmax, ylen, ymin, ymax

tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
def get_grid_values(h5f, grid_name, j_c, i_c, half_width, num_j=None, num_i=None, scale_factor_name='scale_factor', add_offset_name='add_offset',
                    fill_value_name='_FillValue', range_name='actual_range', fill_value=None):
    hfds = h5f[grid_name]
    attrs = hfds.attrs
    if attrs is None:
        raise GenericException('No attributes object for: '+grid_name)

    ylen, xlen = hfds.shape

    if half_width is not None:
        j_l = j_c-half_width
        i_l = i_c-half_width
        if j_l < 0 or i_l < 0:
            return None

        j_r = j_c+half_width+1
        i_r = i_c+half_width+1
        if j_r >= ylen or i_r >= xlen:
            return None
    else:
        j_l = j_c
        j_r = j_c + num_j + 1
        i_l = i_c
        i_r = i_c + num_i + 1

    grd_vals = hfds[j_l:j_r, i_l:i_r]
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    if fill_value is not None:
        grd_vals = np.where(grd_vals == fill_value, np.nan, grd_vals)

    if scale_factor_name is not None:
        attr = attrs.get(scale_factor_name)
tomrink's avatar
tomrink committed
        if attr is not None:
            if np.isscalar(attr):
                scale_factor = attr
            else:
                scale_factor = attr[0]
            grd_vals = grd_vals * scale_factor
tomrink's avatar
tomrink committed
        else:
tomrink's avatar
tomrink committed
            print('Attribute: ' + scale_factor_name + ' not found for dataset: ' + grid_name)
tomrink's avatar
tomrink committed

    if add_offset_name is not None:
        attr = attrs.get(add_offset_name)
tomrink's avatar
tomrink committed
        if attr is not None:
            if np.isscalar(attr):
                add_offset = attr
            else:
                add_offset = attr[0]
            grd_vals = grd_vals + add_offset
tomrink's avatar
tomrink committed
        else:
tomrink's avatar
tomrink committed
            print('Attribute: '+add_offset_name+' not found for dataset: '+grid_name)
tomrink's avatar
tomrink committed

    if range_name is not None:
        attr = attrs.get(range_name)
tomrink's avatar
tomrink committed
        if attr is not None:
            low = attr[0]
            high = attr[1]
            grd_vals = np.where(grd_vals < low, np.nan, grd_vals)
            grd_vals = np.where(grd_vals > high, np.nan, grd_vals)
        else:
            print('Attribute: '+range_name+' not found for dataset: '+grid_name)

tomrink's avatar
tomrink committed
    elif fill_value_name is not None:
        attr = attrs.get(fill_value_name)
tomrink's avatar
tomrink committed
        if attr is not None:
            if np.isscalar(attr):
                fill_value = attr
            else:
                fill_value = attr[0]
            grd_vals = np.where(grd_vals == fill_value, np.nan, grd_vals)
tomrink's avatar
tomrink committed
        else:
tomrink's avatar
tomrink committed
            print('Attribute: '+fill_value_name+' not found for dataset: '+grid_name)
tomrink's avatar
tomrink committed

    return grd_vals

tomrink's avatar
tomrink committed

def concat_dict_s(t_dct_0, t_dct_1):
    keys_0 = list(t_dct_0.keys())
    nda_0 = np.array(keys_0)

    keys_1 = list(t_dct_1.keys())
    nda_1 = np.array(keys_1)

    comm_keys, comm0, comm1 = np.intersect1d(nda_0, nda_1, return_indices=True)

    comm_keys = comm_keys.tolist()

    for key in comm_keys:
        t_dct_1.pop(key)
    t_dct_0.update(t_dct_1)

    return t_dct_0


tomrink's avatar
tomrink committed
# Example GOES file to retrieve GEOS parameters in MetPy form (CONUS)
exmp_file_conus = '/Users/tomrink/data/OR_ABI-L1b-RadC-M6C14_G16_s20193140811215_e20193140813588_c20193140814070.nc'
# Full Disk
exmp_file_fd = '/Users/tomrink/data/OR_ABI-L1b-RadF-M6C16_G16_s20212521800223_e20212521809542_c20212521809596.nc'

tomrink's avatar
tomrink committed
# keep for reference
# if domain == 'CONUS':
#     exmpl_ds = xr.open_dataset(exmp_file_conus)
# elif domain == 'FD':
#     exmpl_ds = xr.open_dataset(exmp_file_fd)
# mdat = exmpl_ds.metpy.parse_cf('Rad')
# geos = mdat.metpy.cartopy_crs
# xlen = mdat.x.values.size
# ylen = mdat.y.values.size
# exmpl_ds.close()

# Taiwan domain:
# lon, lat = 120.955098, 23.834310
# elem, line = (1789, 1505)
# # UR from Taiwan
tomrink's avatar
tomrink committed
# lon, lat = 135.0, 35.0
# elem_ur, line_ur = (2499, 995)
tomrink's avatar
tomrink committed
taiwan_i0 = 1079
taiwan_j0 = 995
tomrink's avatar
tomrink committed
taiwan_lenx = 1420
taiwan_leny = 1020
# geos.transform_point(135.0, 35.0, ccrs.PlateCarree(), False)
# geos.transform_point(106.61, 13.97, ccrs.PlateCarree(), False)
taiwain_extent = [-3342, -502, 1470, 3510]  # GEOS coordinates, not line, elem
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
# ------------ This code will not be needed when we implement a Fully Convolutional CNN -----------------------------------
tomrink's avatar
tomrink committed
# Generate and return tiles of name_list parameters
tomrink's avatar
tomrink committed
def make_for_full_domain_predict(h5f, name_list=None, satellite='GOES16', domain='FD', res_fac=1):
tomrink's avatar
tomrink committed
    w_x = 16
    w_y = 16
tomrink's avatar
tomrink committed
    i_0 = 0
    j_0 = 0
tomrink's avatar
tomrink committed
    s_x = int(w_x / res_fac)
    s_y = int(w_y / res_fac)
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
tomrink's avatar
tomrink committed
    if satellite == 'H08':
        xlen = taiwan_lenx
        ylen = taiwan_leny
tomrink's avatar
tomrink committed
        i_0 = taiwan_i0
        j_0 = taiwan_j0
tomrink's avatar
tomrink committed

    grd_dct = {name: None for name in name_list}

tomrink's avatar
tomrink committed
    cnt_a = 0
    for ds_name in name_list:
        gvals = get_grid_values(h5f, ds_name, j_0, i_0, None, num_j=ylen, num_i=xlen)
        if gvals is not None:
            grd_dct[ds_name] = gvals
            cnt_a += 1

    if cnt_a > 0 and cnt_a != len(name_list):
        raise GenericException('weirdness')

    grd_dct_n = {name: [] for name in name_list}

    n_x = int(xlen/s_x) - 1
    n_y = int(ylen/s_y) - 1

    r_x = xlen - (n_x * s_x)
    x_d = 0 if r_x >= w_x else int((w_x - r_x)/s_x)
    n_x -= x_d

    r_y = ylen - (n_y * s_y)
    y_d = 0 if r_y >= w_y else int((w_y - r_y)/s_y)
    n_y -= y_d

    ll = [j_0 + j*s_y for j in range(n_y)]
    cc = [i_0 + i*s_x for i in range(n_x)]

    for ds_name in name_list:
        for j in range(n_y):
            j_ul = j * s_y
            j_ul_b = j_ul + w_y
            for i in range(n_x):
                i_ul = i * s_x
                i_ul_b = i_ul + w_x
                grd_dct_n[ds_name].append(grd_dct[ds_name][j_ul:j_ul_b, i_ul:i_ul_b])

    grd_dct = {name: None for name in name_list}
    for ds_name in name_list:
        grd_dct[ds_name] = np.stack(grd_dct_n[ds_name])

    return grd_dct, ll, cc


# rho_water = 1.
# rho_ice = 0.917

rho_water = 1000000.0  # g m^-3
rho_ice = 917000.0  # g m^-3
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
# real(kind=real4), parameter:: Rho_Water = 1.0    !g / m ^ 3
# real(kind=real4), parameter:: Rho_Ice = 0.917    !g / m ^ 3
#
# !--- compute
# cloud
# water
# path
# if (Iphase == 0) then
# Cwp_Dcomp(Elem_Idx, Line_Idx) = 0.55 * Tau * Reff * Rho_Water
# Lwp_Dcomp(Elem_Idx, Line_Idx) = 0.55 * Tau * Reff * Rho_Water
# else
# Cwp_Dcomp(Elem_Idx, Line_Idx) = 0.667 * Tau * Reff * Rho_Ice
# Iwp_Dcomp(Elem_Idx, Line_Idx) = 0.667 * Tau * Reff * Rho_Ice
# endif

tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
def compute_lwc_iwc(iphase, reff, opd, geo_dz):
    xy_shape = iphase.shape

    iphase = iphase.flatten()
tomrink's avatar
tomrink committed
    keep_0 = np.invert(np.isnan(iphase))

    reff = reff.flatten()
tomrink's avatar
tomrink committed
    keep_1 = np.invert(np.isnan(reff))

    opd = opd.flatten()
tomrink's avatar
tomrink committed
    keep_2 = np.invert(np.isnan(opd))

tomrink's avatar
tomrink committed
    geo_dz = geo_dz.flatten()
tomrink's avatar
tomrink committed
    keep_3 = np.logical_and(np.invert(np.isnan(geo_dz)), geo_dz > 1.0)

    keep = keep_0 & keep_1 & keep_2 & keep_3
tomrink's avatar
tomrink committed
    lwp_dcomp = np.zeros(reff.shape[0])
    iwp_dcomp = np.zeros(reff.shape[0])
    lwp_dcomp[:] = np.nan
    iwp_dcomp[:] = np.nan

tomrink's avatar
tomrink committed
    ice = iphase == 1 & keep
    no_ice = iphase != 1 & keep
tomrink's avatar
tomrink committed

    # compute ice/liquid water path, g m-2
    reff *= 1.0e-06  # convert microns to meters

tomrink's avatar
tomrink committed
    iwp_dcomp[ice] = 0.667 * opd[ice] * rho_ice * reff[ice]
    lwp_dcomp[no_ice] = 0.55 * opd[no_ice] * rho_water * reff[no_ice]

tomrink's avatar
tomrink committed
    iwp_dcomp /= geo_dz
    lwp_dcomp /= geo_dz

    lwp_dcomp = lwp_dcomp.reshape(xy_shape)
    iwp_dcomp = iwp_dcomp.reshape(xy_shape)

tomrink's avatar
tomrink committed
    return lwp_dcomp, iwp_dcomp


tomrink's avatar
tomrink committed
def make_for_full_domain_predict_viirs_clavrx(h5f, name_list=None, res_fac=1, day_night='DAY', use_dnb=False):
tomrink's avatar
tomrink committed
    w_x = 16
    w_y = 16
    i_0 = 0
    j_0 = 0
    s_x = int(w_x / res_fac)
    s_y = int(w_y / res_fac)

    ylen = h5f['scan_lines_along_track_direction'].shape[0]
    xlen = h5f['pixel_elements_along_scan_direction'].shape[0]

tomrink's avatar
tomrink committed
    use_nl_comp = False
tomrink's avatar
tomrink committed
    if (day_night == 'NIGHT' or day_night == 'AUTO') and use_dnb:
tomrink's avatar
tomrink committed
        use_nl_comp = True

tomrink's avatar
tomrink committed
    grd_dct = {name: None for name in name_list}

tomrink's avatar
tomrink committed
    cnt_a = 0
tomrink's avatar
tomrink committed
    for ds_name in name_list:
tomrink's avatar
tomrink committed
        name = ds_name

        if use_nl_comp:
            if ds_name == 'cld_reff_dcomp':
                name = 'cld_reff_nlcomp'
            elif ds_name == 'cld_opd_dcomp':
                name = 'cld_opd_nlcomp'

        gvals = get_grid_values(h5f, name, j_0, i_0, None, num_j=ylen, num_i=xlen)
tomrink's avatar
tomrink committed
        if gvals is not None:
            grd_dct[ds_name] = gvals
            cnt_a += 1

    if cnt_a > 0 and cnt_a != len(name_list):
        raise GenericException('weirdness')

tomrink's avatar
tomrink committed
    # TODO: need to investigate discrepencies with compute_lwc_iwc
    # if use_nl_comp:
    #     cld_phase = get_grid_values(h5f, 'cloud_phase', j_0, i_0, None, num_j=ylen, num_i=xlen)
    #     cld_dz = get_grid_values(h5f, 'cld_geo_thick', j_0, i_0, None, num_j=ylen, num_i=xlen)
    #     reff = grd_dct['cld_reff_dcomp']
    #     opd = grd_dct['cld_opd_dcomp']
    #
    #     lwc_nlcomp, iwc_nlcomp = compute_lwc_iwc(cld_phase, reff, opd, cld_dz)
    #     grd_dct['iwc_dcomp'] = iwc_nlcomp
    #     grd_dct['lwc_dcomp'] = lwc_nlcomp
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    grd_dct_n = {name: [] for name in name_list}

tomrink's avatar
tomrink committed
    n_x = int(xlen/s_x) - 1
    n_y = int(ylen/s_y) - 1

    r_x = xlen - (n_x * s_x)
    x_d = 0 if r_x >= w_x else int((w_x - r_x)/s_x)
    n_x -= x_d

    r_y = ylen - (n_y * s_y)
    y_d = 0 if r_y >= w_y else int((w_y - r_y)/s_y)
    n_y -= y_d
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    ll = [j_0 + j*s_y for j in range(n_y)]
    cc = [i_0 + i*s_x for i in range(n_x)]
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    for ds_name in name_list:
tomrink's avatar
tomrink committed
        for j in range(n_y):
tomrink's avatar
tomrink committed
            j_ul = j * s_y
tomrink's avatar
tomrink committed
            j_ul_b = j_ul + w_y
tomrink's avatar
tomrink committed
            for i in range(n_x):
tomrink's avatar
tomrink committed
                i_ul = i * s_x
tomrink's avatar
tomrink committed
                i_ul_b = i_ul + w_x
                grd_dct_n[ds_name].append(grd_dct[ds_name][j_ul:j_ul_b, i_ul:i_ul_b])
tomrink's avatar
tomrink committed

    grd_dct = {name: None for name in name_list}
tomrink's avatar
tomrink committed
    for ds_name in name_list:
tomrink's avatar
tomrink committed
        grd_dct[ds_name] = np.stack(grd_dct_n[ds_name])

tomrink's avatar
tomrink committed
    lats = get_grid_values(h5f, 'latitude', j_0, i_0, None, num_j=ylen, num_i=xlen, range_name=None)
    lons = get_grid_values(h5f, 'longitude', j_0, i_0, None, num_j=ylen, num_i=xlen, range_name=None)

tomrink's avatar
tomrink committed
    ll_2d, cc_2d = np.meshgrid(ll, cc, indexing='ij')
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    lats = lats[ll_2d, cc_2d]
    lons = lons[ll_2d, cc_2d]
tomrink's avatar
tomrink committed

    return grd_dct, ll, cc, lats, lons
tomrink's avatar
tomrink committed


tomrink's avatar
tomrink committed
def make_for_full_domain_predict2(h5f, satellite='GOES16', domain='FD', res_fac=1):
tomrink's avatar
tomrink committed
    w_x = 16
    w_y = 16
    i_0 = 0
    j_0 = 0
tomrink's avatar
tomrink committed
    s_x = int(w_x / res_fac)
    s_y = int(w_y / res_fac)
tomrink's avatar
tomrink committed

    geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
    if satellite == 'H08':
        xlen = taiwan_lenx
        ylen = taiwan_leny
        i_0 = taiwan_i0
        j_0 = taiwan_j0

tomrink's avatar
tomrink committed
    n_x = int(xlen/s_x)
    n_y = int(ylen/s_y)
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    solzen = get_grid_values(h5f, 'solar_zenith_angle', j_0, i_0, None, num_j=ylen, num_i=xlen)
    satzen = get_grid_values(h5f, 'sensor_zenith_angle', j_0, i_0, None, num_j=ylen, num_i=xlen)
tomrink's avatar
tomrink committed
    solzen = solzen[0:(n_y-1)*s_y:s_y, 0:(n_x-1)*s_x:s_x]
    satzen = satzen[0:(n_y-1)*s_y:s_y, 0:(n_x-1)*s_x:s_x]
tomrink's avatar
tomrink committed

    return solzen, satzen
# -------------------------------------------------------------------------------------------


tomrink's avatar
tomrink committed
def prepare_evaluate(h5f, name_list, satellite='GOES16', domain='FD', res_fac=1, offset=0):
    w_x = 16
    w_y = 16
    i_0 = 0
    j_0 = 0
    s_x = int(w_x / res_fac)
    s_y = int(w_y / res_fac)
tomrink's avatar
tomrink committed

    geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
    if satellite == 'H08':
        xlen = taiwan_lenx
        ylen = taiwan_leny
        i_0 = taiwan_i0
        j_0 = taiwan_j0

tomrink's avatar
tomrink committed
    n_x = int(xlen/s_x) - 1
    n_y = int(ylen/s_y) - 1

    r_x = xlen - (n_x * s_x)
    x_d = 0 if r_x >= w_x else int((w_x - r_x)/s_x)
    n_x -= x_d

    r_y = ylen - (n_y * s_y)
    y_d = 0 if r_y >= w_y else int((w_y - r_y)/s_y)
    n_y -= y_d

    ll = [(offset+j_0) + j*s_y for j in range(n_y)]
    cc = [(offset+i_0) + i*s_x for i in range(n_x)]

tomrink's avatar
tomrink committed
    grd_dct_n = {name: [] for name in name_list}

    cnt_a = 0
    for ds_name in name_list:
        gvals = get_grid_values(h5f, ds_name, j_0, i_0, None, num_j=ylen, num_i=xlen)
        if gvals is not None:
            grd_dct_n[ds_name] = gvals
            cnt_a += 1

    if cnt_a > 0 and cnt_a != len(name_list):
        raise GenericException('weirdness')

    solzen = get_grid_values(h5f, 'solar_zenith_angle', j_0, i_0, None, num_j=ylen, num_i=xlen)
    satzen = get_grid_values(h5f, 'sensor_zenith_angle', j_0, i_0, None, num_j=ylen, num_i=xlen)
tomrink's avatar
tomrink committed
    solzen = solzen[0:(n_y-1)*s_y:s_y, 0:(n_x-1)*s_x:s_x]
    satzen = satzen[0:(n_y-1)*s_y:s_y, 0:(n_x-1)*s_x:s_x]
tomrink's avatar
tomrink committed

    grd_dct = {name: None for name in name_list}
    for ds_name in name_list:
        grd_dct[ds_name] = np.stack(grd_dct_n[ds_name])

    return grd_dct, solzen, satzen, ll, cc


tomrink's avatar
tomrink committed
flt_level_ranges_str = {k: None for k in range(5)}
flt_level_ranges_str[0] = '0_2000'
flt_level_ranges_str[1] = '2000_4000'
flt_level_ranges_str[2] = '4000_6000'
flt_level_ranges_str[3] = '6000_8000'
flt_level_ranges_str[4] = '8000_15000'

# flt_level_ranges_str = {k: None for k in range(1)}
# flt_level_ranges_str[0] = 'column'
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
def get_cf_nav_parameters(satellite='GOES16', domain='FD'):
    param_dct = None

    if satellite == 'H08':  # We presently only have FD
        param_dct = {'semi_major_axis': 6378.137,
                     'semi_minor_axis': 6356.7523,
                     'perspective_point_height': 35785.863,
                     'latitude_of_projection_origin': 0.0,
                     'longitude_of_projection_origin': 140.7,
                     'inverse_flattening': 298.257,
                     'sweep_angle_axis': 'y',
                     'x_scale_factor': 5.58879902955962e-05,
                     'x_add_offset': -0.153719917308037,
                     'y_scale_factor': -5.58879902955962e-05,
                     'y_add_offset': 0.153719917308037}
    elif satellite == 'GOES16':
        if domain == 'CONUS':
            param_dct = {'semi_major_axis': 6378137.0,
                         'semi_minor_axis': 6356752.31414,
                         'perspective_point_height': 35786023.0,
                         'latitude_of_projection_origin': 0.0,
                         'longitude_of_projection_origin': -75,
                         'inverse_flattening': 298.257,
                         'sweep_angle_axis': 'x',
                         'x_scale_factor': 5.6E-05,
                         'x_add_offset': -0.101332,
                         'y_scale_factor': -5.6E-05,
                         'y_add_offset': 0.128212}

    return param_dct

tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
def write_icing_file(clvrx_str_time, output_dir, preds_dct, probs_dct, x, y, lons, lats, elems, lines):
tomrink's avatar
tomrink committed
    outfile_name = output_dir + 'icing_prediction_'+clvrx_str_time+'.h5'
tomrink's avatar
tomrink committed
    h5f_out = h5py.File(outfile_name, 'w')
    dim_0_name = 'x'
    dim_1_name = 'y'
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    prob_s = []
    pred_s = []

tomrink's avatar
tomrink committed
    flt_lvls = list(preds_dct.keys())
    for flvl in flt_lvls:
        preds = preds_dct[flvl]
tomrink's avatar
tomrink committed
        pred_s.append(preds)
tomrink's avatar
tomrink committed
        icing_pred_ds = h5f_out.create_dataset('icing_prediction_level_'+flt_level_ranges_str[flvl], data=preds, dtype='i2')
tomrink's avatar
tomrink committed
        icing_pred_ds.attrs.create('coordinates', data='y x')
        icing_pred_ds.attrs.create('grid_mapping', data='Projection')
        icing_pred_ds.attrs.create('missing', data=-1)
tomrink's avatar
tomrink committed
        icing_pred_ds.dims[0].label = dim_0_name
        icing_pred_ds.dims[1].label = dim_1_name
tomrink's avatar
tomrink committed

    for flvl in flt_lvls:
        probs = probs_dct[flvl]
tomrink's avatar
tomrink committed
        prob_s.append(probs)
tomrink's avatar
tomrink committed
        icing_prob_ds = h5f_out.create_dataset('icing_probability_level_'+flt_level_ranges_str[flvl], data=probs, dtype='f4')
tomrink's avatar
tomrink committed
        icing_prob_ds.attrs.create('coordinates', data='y x')
        icing_prob_ds.attrs.create('grid_mapping', data='Projection')
        icing_prob_ds.attrs.create('missing', data=-1.0)
tomrink's avatar
tomrink committed
        icing_prob_ds.dims[0].label = dim_0_name
        icing_prob_ds.dims[1].label = dim_1_name
tomrink's avatar
tomrink committed
    prob_s = np.stack(prob_s, axis=-1)
    max_prob = np.max(prob_s, axis=2)

    icing_prob_ds = h5f_out.create_dataset('max_icing_probability_column', data=max_prob, dtype='f4')
    icing_prob_ds.attrs.create('coordinates', data='y x')
    icing_prob_ds.attrs.create('grid_mapping', data='Projection')
    icing_prob_ds.attrs.create('missing', data=-1.0)
    icing_prob_ds.dims[0].label = dim_0_name
    icing_prob_ds.dims[1].label = dim_1_name

tomrink's avatar
tomrink committed
    max_lvl = np.argmax(prob_s, axis=2)
tomrink's avatar
tomrink committed

tomrink's avatar
tomrink committed
    icing_pred_ds = h5f_out.create_dataset('max_icing_probability_level', data=max_lvl, dtype='i2')
tomrink's avatar
tomrink committed
    icing_pred_ds.attrs.create('coordinates', data='y x')
    icing_pred_ds.attrs.create('grid_mapping', data='Projection')
    icing_pred_ds.attrs.create('missing', data=-1)