util.py

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'


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


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

    dim_0_name = 'x'
    dim_1_name = 'y'

    prob_s = []
    pred_s = []

    flt_lvls = list(preds_dct.keys())
    for flvl in flt_lvls:
        preds = preds_dct[flvl]
        pred_s.append(preds)
        icing_pred_ds = h5f_out.create_dataset('icing_prediction_level_'+flt_level_ranges_str[flvl], data=preds, dtype='i2')
        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)
        icing_pred_ds.dims[0].label = dim_0_name
        icing_pred_ds.dims[1].label = dim_1_name

    for flvl in flt_lvls:
        probs = probs_dct[flvl]
        prob_s.append(probs)
        icing_prob_ds = h5f_out.create_dataset('icing_probability_level_'+flt_level_ranges_str[flvl], data=probs, 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

    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

    max_lvl = np.argmax(prob_s, axis=2)

    icing_pred_ds = h5f_out.create_dataset('max_icing_probability_level', data=max_lvl, dtype='i2')
    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)
    icing_pred_ds.dims[0].label = dim_0_name
    icing_pred_ds.dims[1].label = dim_1_name

    lon_ds = h5f_out.create_dataset('longitude', data=lons, dtype='f4')
    lon_ds.attrs.create('units', data='degrees_east')
    lon_ds.attrs.create('long_name', data='icing prediction longitude')
    lon_ds.dims[0].label = dim_0_name
    lon_ds.dims[1].label = dim_1_name

    lat_ds = h5f_out.create_dataset('latitude', data=lats, dtype='f4')
    lat_ds.attrs.create('units', data='degrees_north')
    lat_ds.attrs.create('long_name', data='icing prediction latitude')
    lat_ds.dims[0].label = dim_0_name
    lat_ds.dims[1].label = dim_1_name

    proj_ds = h5f_out.create_dataset('Projection', data=0, dtype='b')
    proj_ds.attrs.create('long_name', data='Himawari Imagery Projection')
    proj_ds.attrs.create('grid_mapping_name', data='geostationary')
    proj_ds.attrs.create('sweep_angle_axis', data='y')
    proj_ds.attrs.create('units', data='rad')
    proj_ds.attrs.create('semi_major_axis', data=6378.137)
    proj_ds.attrs.create('semi_minor_axis', data=6356.7523)
    proj_ds.attrs.create('inverse_flattening', data=298.257)
    proj_ds.attrs.create('perspective_point_height', data=35785.863)
    proj_ds.attrs.create('latitude_of_projection_origin', data=0.0)
    proj_ds.attrs.create('longitude_of_projection_origin', data=140.7)
    proj_ds.attrs.create('CFAC', data=20466275)
    proj_ds.attrs.create('LFAC', data=20466275)
    proj_ds.attrs.create('COFF', data=2750.5)
    proj_ds.attrs.create('LOFF', data=2750.5)

    if x is not None:
        x_ds = h5f_out.create_dataset('x', data=x, dtype='f8')
        x_ds.dims[0].label = dim_0_name
        x_ds.attrs.create('units', data='rad')
        x_ds.attrs.create('standard_name', data='projection_x_coordinate')
        x_ds.attrs.create('long_name', data='GOES PUG W-E fixed grid viewing angle')
        x_ds.attrs.create('scale_factor', data=5.58879902955962e-05)
        x_ds.attrs.create('add_offset', data=-0.153719917308037)
        x_ds.attrs.create('CFAC', data=20466275)
        x_ds.attrs.create('COFF', data=2750.5)

        y_ds = h5f_out.create_dataset('y', data=y, dtype='f8')
        y_ds.dims[0].label = dim_1_name
        y_ds.attrs.create('units', data='rad')
        y_ds.attrs.create('standard_name', data='projection_y_coordinate')
        y_ds.attrs.create('long_name', data='GOES PUG S-N fixed grid viewing angle')
        y_ds.attrs.create('scale_factor', data=-5.58879902955962e-05)
        y_ds.attrs.create('add_offset', data=0.153719917308037)
        y_ds.attrs.create('LFAC', data=20466275)
        y_ds.attrs.create('LOFF', data=2750.5)

    if elems is not None:
        elem_ds = h5f_out.create_dataset('elems', data=elems, dtype='i2')
        elem_ds.dims[0].label = dim_0_name
        line_ds = h5f_out.create_dataset('lines', data=lines, dtype='i2')
        line_ds.dims[0].label = dim_1_name
        pass

    h5f_out.close()


def write_icing_file_nc4(clvrx_str_time, output_dir, preds_dct, probs_dct,
                         x, y, lons, lats, elems, lines, satellite='GOES16', domain='CONUS',
                         has_time=False, use_nan=False, prob_thresh=0.5, bt_10_4=None):
    outfile_name = output_dir + 'icing_prediction_'+clvrx_str_time+'.nc'
    rootgrp = Dataset(outfile_name, 'w', format='NETCDF4')

    rootgrp.setncattr('Conventions', 'CF-1.7')

    dim_0_name = 'x'
    dim_1_name = 'y'
    time_dim_name = 'time'
    geo_coords = 'time y x'

    dim_0 = rootgrp.createDimension(dim_0_name, size=x.shape[0])
    dim_1 = rootgrp.createDimension(dim_1_name, size=y.shape[0])
    dim_time = rootgrp.createDimension(time_dim_name, size=1)

    tvar = rootgrp.createVariable('time', 'f8', time_dim_name)
    tvar[0] = get_timestamp(clvrx_str_time)
    tvar.units = 'seconds since 1970-01-01 00:00:00'

    if not has_time:
        var_dim_list = [dim_1_name, dim_0_name]
    else:
        var_dim_list = [time_dim_name, dim_1_name, dim_0_name]

    prob_s = []

    flt_lvls = list(preds_dct.keys())
    for flvl in flt_lvls:
        preds = preds_dct[flvl]

        icing_pred_ds = rootgrp.createVariable('icing_prediction_level_'+flt_level_ranges_str[flvl], 'i2', var_dim_list)
        icing_pred_ds.setncattr('coordinates', geo_coords)
        icing_pred_ds.setncattr('grid_mapping', 'Projection')
        icing_pred_ds.setncattr('missing', -1)
        if has_time:
            preds = preds.reshape((1, y.shape[0], x.shape[0]))
        icing_pred_ds[:,] = preds

    for flvl in flt_lvls:
        probs = probs_dct[flvl]
        prob_s.append(probs)

        icing_prob_ds = rootgrp.createVariable('icing_probability_level_'+flt_level_ranges_str[flvl], 'f4', var_dim_list)
        icing_prob_ds.setncattr('coordinates', geo_coords)
        icing_prob_ds.setncattr('grid_mapping', 'Projection')
        if not use_nan:
            icing_prob_ds.setncattr('missing', -1.0)
        else:
            icing_prob_ds.setncattr('missing', np.nan)
        if has_time:
            probs = probs.reshape((1, y.shape[0], x.shape[0]))
        if use_nan:
            probs = np.where(probs < prob_thresh, np.nan, probs)
        icing_prob_ds[:,] = probs

    prob_s = np.stack(prob_s, axis=-1)
    max_prob = np.max(prob_s, axis=2)
    if use_nan:
        max_prob = np.where(max_prob < prob_thresh, np.nan, max_prob)
    if has_time:
        max_prob = max_prob.reshape(1, y.shape[0], x.shape[0])

    icing_prob_ds = rootgrp.createVariable('max_icing_probability_column', 'f4', var_dim_list)
    icing_prob_ds.setncattr('coordinates', geo_coords)
    icing_prob_ds.setncattr('grid_mapping', 'Projection')
    if not use_nan:
        icing_prob_ds.setncattr('missing', -1.0)
    else:
        icing_prob_ds.setncattr('missing', np.nan)
    icing_prob_ds[:,] = max_prob

    prob_s = np.where(prob_s < prob_thresh, -1.0, prob_s)
    max_lvl = np.where(np.all(prob_s == -1, axis=2), -1, np.argmax(prob_s, axis=2))
    if use_nan:
        max_lvl = np.where(max_lvl == -1, np.nan, max_lvl)
    if has_time:
        max_lvl = max_lvl.reshape((1, y.shape[0], x.shape[0]))

    icing_pred_ds = rootgrp.createVariable('max_icing_probability_level', 'i2', var_dim_list)
    icing_pred_ds.setncattr('coordinates', geo_coords)
    icing_pred_ds.setncattr('grid_mapping', 'Projection')
    icing_pred_ds.setncattr('missing', -1)
    icing_pred_ds[:,] = max_lvl

    if bt_10_4 is not None:
        bt_ds = rootgrp.createVariable('bt_10_4', 'f4', var_dim_list)
        bt_ds.setncattr('coordinates', geo_coords)
        bt_ds.setncattr('grid_mapping', 'Projection')
        bt_ds[:,] = bt_10_4

    lon_ds = rootgrp.createVariable('longitude', 'f4', [dim_1_name, dim_0_name])
    lon_ds.units = 'degrees_east'
    lon_ds[:,] = lons

    lat_ds = rootgrp.createVariable('latitude', 'f4', [dim_1_name, dim_0_name])
    lat_ds.units = 'degrees_north'
    lat_ds[:,] = lats

    cf_nav_dct = get_cf_nav_parameters(satellite, domain)

    if satellite == 'H08':
        long_name = 'Himawari Imagery Projection'
    elif satellite == 'GOES16':
        long_name = 'GOES-16/17 Imagery Projection'

    proj_ds = rootgrp.createVariable('Projection', 'b')
    proj_ds.setncattr('long_name', long_name)
    proj_ds.setncattr('grid_mapping_name', 'geostationary')
    proj_ds.setncattr('sweep_angle_axis', cf_nav_dct['sweep_angle_axis'])
    proj_ds.setncattr('semi_major_axis', cf_nav_dct['semi_major_axis'])
    proj_ds.setncattr('semi_minor_axis', cf_nav_dct['semi_minor_axis'])
    proj_ds.setncattr('inverse_flattening', cf_nav_dct['inverse_flattening'])
    proj_ds.setncattr('perspective_point_height', cf_nav_dct['perspective_point_height'])
    proj_ds.setncattr('latitude_of_projection_origin', cf_nav_dct['latitude_of_projection_origin'])
    proj_ds.setncattr('longitude_of_projection_origin', cf_nav_dct['longitude_of_projection_origin'])

    if x is not None:
        x_ds = rootgrp.createVariable(dim_0_name, 'f8', [dim_0_name])
        x_ds.units = 'rad'
        x_ds.setncattr('axis', 'X')
        x_ds.setncattr('standard_name', 'projection_x_coordinate')
        x_ds.setncattr('long_name', 'fixed grid viewing angle')
        x_ds.setncattr('scale_factor', cf_nav_dct['x_scale_factor'])
        x_ds.setncattr('add_offset', cf_nav_dct['x_add_offset'])
        x_ds[:] = x

        y_ds = rootgrp.createVariable(dim_1_name, 'f8', [dim_1_name])
        y_ds.units = 'rad'
        y_ds.setncattr('axis', 'Y')
        y_ds.setncattr('standard_name', 'projection_y_coordinate')
        y_ds.setncattr('long_name', 'fixed grid viewing angle')
        y_ds.setncattr('scale_factor', cf_nav_dct['y_scale_factor'])
        y_ds.setncattr('add_offset', cf_nav_dct['y_add_offset'])
        y_ds[:] = y

    if elems is not None:
        elem_ds = rootgrp.createVariable('elems', 'i2', [dim_0_name])
        elem_ds[:] = elems
        line_ds = rootgrp.createVariable('lines', 'i2', [dim_1_name])
        line_ds[:] = lines
        pass

    rootgrp.close()


def write_icing_file_nc4_viirs(clvrx_str_time, output_dir, preds_dct, probs_dct, lons, lats,
                               has_time=False, use_nan=False, prob_thresh=0.5, bt_10_4=None):
    outfile_name = output_dir + 'icing_prediction_'+clvrx_str_time+'.nc'
    rootgrp = Dataset(outfile_name, 'w', format='NETCDF4')

    rootgrp.setncattr('Conventions', 'CF-1.7')

    dim_0_name = 'x'
    dim_1_name = 'y'
    time_dim_name = 'time'
    geo_coords = 'longitude latitude'

    dim_1_len, dim_0_len = lons.shape

    dim_0 = rootgrp.createDimension(dim_0_name, size=dim_0_len)
    dim_1 = rootgrp.createDimension(dim_1_name, size=dim_1_len)
    dim_time = rootgrp.createDimension(time_dim_name, size=1)

    tvar = rootgrp.createVariable('time', 'f8', time_dim_name)
    tvar[0] = get_timestamp(clvrx_str_time)
    tvar.units = 'seconds since 1970-01-01 00:00:00'

    if not has_time:
        var_dim_list = [dim_1_name, dim_0_name]
    else:
        var_dim_list = [time_dim_name, dim_1_name, dim_0_name]

    prob_s = []

    flt_lvls = list(preds_dct.keys())
    for flvl in flt_lvls:
        preds = preds_dct[flvl]

        icing_pred_ds = rootgrp.createVariable('icing_prediction_level_'+flt_level_ranges_str[flvl], 'i2', var_dim_list)
        icing_pred_ds.setncattr('coordinates', geo_coords)
        icing_pred_ds.setncattr('grid_mapping', 'Projection')
        icing_pred_ds.setncattr('missing', -1)
        if has_time:
            preds = preds.reshape((1, dim_1_len, dim_0_len))
        icing_pred_ds[:,] = preds

    for flvl in flt_lvls:
        probs = probs_dct[flvl]
        prob_s.append(probs)

        icing_prob_ds = rootgrp.createVariable('icing_probability_level_'+flt_level_ranges_str[flvl], 'f4', var_dim_list)
        icing_prob_ds.setncattr('coordinates', geo_coords)
        icing_prob_ds.setncattr('grid_mapping', 'Projection')
        if not use_nan:
            icing_prob_ds.setncattr('missing', -1.0)
        else:
            icing_prob_ds.setncattr('missing', np.nan)
        if has_time:
            probs = probs.reshape((1, dim_1_len, dim_0_len))
        if use_nan:
            probs = np.where(probs < prob_thresh, np.nan, probs)
        icing_prob_ds[:,] = probs

    prob_s = np.stack(prob_s, axis=-1)
    max_prob = np.max(prob_s, axis=2)
    if use_nan:
        max_prob = np.where(max_prob < prob_thresh, np.nan, max_prob)
    if has_time:
        max_prob = max_prob.reshape(1, dim_1_len, dim_0_len)

    icing_prob_ds = rootgrp.createVariable('max_icing_probability_column', 'f4', var_dim_list)
    icing_prob_ds.setncattr('coordinates', geo_coords)
    icing_prob_ds.setncattr('grid_mapping', 'Projection')
    if not use_nan:
        icing_prob_ds.setncattr('missing', -1.0)
    else:
        icing_prob_ds.setncattr('missing', np.nan)
    icing_prob_ds[:,] = max_prob

    prob_s = np.where(prob_s < prob_thresh, -1.0, prob_s)
    max_lvl = np.where(np.all(prob_s == -1, axis=2), -1, np.argmax(prob_s, axis=2))
    if use_nan:
        max_lvl = np.where(max_lvl == -1, np.nan, max_lvl)
    if has_time:
        max_lvl = max_lvl.reshape((1, dim_1_len, dim_0_len))

    icing_pred_ds = rootgrp.createVariable('max_icing_probability_level', 'i2', var_dim_list)
    icing_pred_ds.setncattr('coordinates', geo_coords)
    icing_pred_ds.setncattr('grid_mapping', 'Projection')
    icing_pred_ds.setncattr('missing', -1)
    icing_pred_ds[:,] = max_lvl

    if bt_10_4 is not None:
        bt_ds = rootgrp.createVariable('bt_10_4', 'f4', var_dim_list)
        bt_ds.setncattr('coordinates', geo_coords)
        bt_ds.setncattr('grid_mapping', 'Projection')
        bt_ds[:,] = bt_10_4

    lon_ds = rootgrp.createVariable('longitude', 'f4', [dim_1_name, dim_0_name])
    lon_ds.units = 'degrees_east'
    lon_ds[:,] = lons

    lat_ds = rootgrp.createVariable('latitude', 'f4', [dim_1_name, dim_0_name])
    lat_ds.units = 'degrees_north'
    lat_ds[:,] = lats

    proj_ds = rootgrp.createVariable('Projection', 'b')
    proj_ds.setncattr('grid_mapping_name', 'latitude_longitude')

    rootgrp.close()


def downscale_2x(original, smoothing=False, samples_axis_first=False):
    # if smoothing:
    #     original = scipy.ndimage.gaussian_filter(original, sigma = 1/2)
    if not samples_axis_first:
        lr = np.nanmean(np.array([original[0::2,0::2],
              original[1::2,0::2],
              original[0::2,1::2],
              original[1::2,1::2]]),axis=0).squeeze()
    elif samples_axis_first:
        lr = np.nanmean(np.array([original[:,0::2,0::2],
              original[:,1::2,0::2],
              original[:,0::2,1::2],
              original[:,1::2,1::2]]),axis=0).squeeze()
    return lr


def resample(x, y, z, x_new, y_new):
    z_intrp = []
    for k in range(z.shape[0]):
        z_k = z[k, :, :]
        f = RectBivariateSpline(x, y, z_k)
        z_intrp.append(f(x_new, y_new))
    return np.stack(z_intrp)


def resample_one(x, y, z, x_new, y_new):
    f = RectBivariateSpline(x, y, z)
    return f(x_new, y_new)


def resample_2d_linear(x, y, z, x_new, y_new):
    z_intrp = []
    for k in range(z.shape[0]):
        z_k = z[k, :, :]
        f = interp2d(x, y, z_k)
        z_intrp.append(f(x_new, y_new))
    return np.stack(z_intrp)


def resample_2d_linear_one(x, y, z, x_new, y_new):
    f = interp2d(x, y, z)
    return f(x_new, y_new)


# Gaussian filter suitable for model training Data Pipeline
# z: input array. Must have dimensions: [BATCH_SIZE, Y, X]
# sigma: Standard deviation for Gaussian kernel
# returns stacked 2d arrays of same input dimension
def smooth_2d(z, sigma=1.0):
    z_smoothed = []
    for j in range(z.shape[0]):
        z_j = z[j, :, :]
        z_smoothed.append(gaussian_filter(z_j, sigma=sigma))
    return np.stack(z_smoothed)


# For [Y, X], see above
def smooth_2d_single(z, sigma=1.0):
    return gaussian_filter(z, sigma=sigma)


def median_filter_2d(z, kernel_size=3):
    z_filtered = []
    for j in range(z.shape[0]):
        z_j = z[j, :, :]
        z_filtered.append(medfilt2d(z_j, kernel_size=kernel_size))
    return np.stack(z_filtered)


def median_filter_2d_single(z, kernel_size=3):
    return medfilt2d(z, kernel_size=kernel_size)


def get_training_parameters(day_night='DAY', l1b_andor_l2='both', satellite='GOES16', use_dnb=False):
    if day_night == 'DAY':
        train_params_l2 = ['cld_height_acha', 'cld_geo_thick', 'cld_temp_acha', 'cld_press_acha', 'supercooled_cloud_fraction',
                           'cld_emiss_acha', 'conv_cloud_fraction', 'cld_reff_dcomp', 'cld_opd_dcomp', 'iwc_dcomp', 'lwc_dcomp']

        if satellite == 'GOES16':
            train_params_l1b = ['temp_10_4um_nom', 'temp_11_0um_nom', 'temp_12_0um_nom', 'temp_13_3um_nom', 'temp_3_75um_nom',
                                'temp_6_2um_nom', 'temp_6_7um_nom', 'temp_7_3um_nom', 'temp_9_7um_nom',
                                'refl_0_47um_nom', 'refl_0_65um_nom', 'refl_0_86um_nom', 'refl_1_38um_nom', 'refl_1_60um_nom']
                                # 'refl_2_10um_nom']
        elif satellite == 'H08':
            train_params_l1b = ['temp_10_4um_nom', 'temp_12_0um_nom', 'temp_8_5um_nom', 'temp_3_75um_nom', 'refl_2_10um_nom',
                                'refl_1_60um_nom', 'refl_0_86um_nom', 'refl_0_47um_nom']
    else:
        train_params_l2 = ['cld_height_acha', 'cld_geo_thick', 'cld_temp_acha', 'cld_press_acha', 'supercooled_cloud_fraction',
                           'cld_emiss_acha', 'conv_cloud_fraction', 'cld_reff_acha', 'cld_opd_acha']

        if use_dnb is True:
            train_params_l2 = ['cld_height_acha', 'cld_geo_thick', 'cld_temp_acha', 'cld_press_acha', 'supercooled_cloud_fraction',
                               'cld_emiss_acha', 'conv_cloud_fraction', 'cld_reff_dcomp', 'cld_opd_dcomp', 'iwc_dcomp', 'lwc_dcomp']

        if satellite == 'GOES16':
            train_params_l1b = ['temp_10_4um_nom', 'temp_11_0um_nom', 'temp_12_0um_nom', 'temp_13_3um_nom', 'temp_3_75um_nom',
                                'temp_6_2um_nom', 'temp_6_7um_nom', 'temp_7_3um_nom', 'temp_9_7um_nom']
        elif satellite == 'H08':
            train_params_l1b = ['temp_10_4um_nom', 'temp_12_0um_nom', 'temp_8_5um_nom', 'temp_3_75um_nom']

    if l1b_andor_l2 == 'both':
        train_params = train_params_l1b + train_params_l2
    elif l1b_andor_l2 == 'l1b':
        train_params = train_params_l1b
    elif l1b_andor_l2 == 'l2':
        train_params = train_params_l2

    return train_params, train_params_l1b, train_params_l2