util.py

import metpy
import numpy as np
import xarray as xr
import datetime
from datetime import timezone
from metpy.units import units
from metpy.calc import thickness_hydrostatic
from collections import namedtuple
import os
import h5py
import pickle
from netCDF4 import Dataset
from util.setup import ancillary_path

LatLonTuple = namedtuple('LatLonTuple', ['lat', 'lon'])

homedir = os.path.expanduser('~') + '/'


class GenericException(Exception):
    def __init__(self, message):
        self.message = message


class EarlyStop:
    def __init__(self, window_length=3, patience=5):
        self.patience = patience
        self.min = np.finfo(np.single).max
        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


def get_time_tuple_utc(timestamp):
    dt_obj = datetime.datetime.fromtimestamp(timestamp, timezone.utc)
    return dt_obj, dt_obj.timetuple()


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


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

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

    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

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


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


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


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


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


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

    if pres < prof_pres[0]:
        raise GenericException("target pressure less than top of pressure profile")

    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:]

    if sfc_pres is not None:
        if pres > sfc_pres:  # incoming pressure below surface
            return -1

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

        if sfc_pres > prof_pres[i_bot]:  # surface below profile bottom
            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]

    prof_pres = np.array(pres_s)