From 81caba236fc0785e00c31d3d89d6e5efca537b2e Mon Sep 17 00:00:00 2001
From: tomrink <rink@ssec.wisc.edu>
Date: Tue, 31 Oct 2023 21:16:09 -0500
Subject: [PATCH] snapshot...

---
 modules/icing/util.py | 680 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 677 insertions(+), 3 deletions(-)

diff --git a/modules/icing/util.py b/modules/icing/util.py
index 318c492f..84c7ae49 100644
--- a/modules/icing/util.py
+++ b/modules/icing/util.py
@@ -3,15 +3,16 @@ import deeplearning.icing_fcn as icing_fcn
 import deeplearning.icing_cnn as icing_cnn
 from icing.pirep_goes import setup, time_filter_3
 from icing.moon_phase import moon_phase
-from util.util import get_time_tuple_utc, is_day, check_oblique, get_median, homedir, write_icing_file_nc4,\
-    write_icing_file_nc4_viirs, get_training_parameters,\
-    make_for_full_domain_predict, make_for_full_domain_predict_viirs_clavrx, prepare_evaluate
+from util.util import get_time_tuple_utc, is_day, check_oblique, get_median, homedir, \
+    get_cartopy_crs, \
+    get_fill_attrs, get_grid_values, GenericException, get_timestamp, get_cf_nav_parameters
 from util.plot import make_icing_image
 from util.geos_nav import get_navigation, get_lon_lat_2d_mesh
 from util.setup import model_path_day, model_path_night
 from aeolus.datasource import CLAVRx, CLAVRx_VIIRS, CLAVRx_H08, CLAVRx_H09
 import h5py
 import datetime
+from netCDF4 import Dataset
 import tensorflow as tf
 import os
 # from scipy.signal import medfilt2d
@@ -24,6 +25,679 @@ flt_level_ranges[2] = [4000.0, 6000.0]
 flt_level_ranges[3] = [6000.0, 8000.0]
 flt_level_ranges[4] = [8000.0, 15000.0]
 
+# Taiwan domain:
+# lon, lat = 120.955098, 23.834310
+# elem, line = (1789, 1505)
+# # UR from Taiwan
+# lon, lat = 135.0, 35.0
+# elem_ur, line_ur = (2499, 995)
+taiwan_i0 = 1079
+taiwan_j0 = 995
+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
+
+
+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_8_5um_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_8_5um_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
+
+
+def make_for_full_domain_predict(h5f, name_list=None, satellite='GOES16', domain='FD', res_fac=1,
+                                 data_extent=None):
+    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)
+
+    geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
+    geos_nav = get_navigation(satellite, domain)
+
+    if satellite == 'GOES16':
+        if data_extent is not None:
+            ll_pt_a = data_extent[0]
+            ll_pt_b = data_extent[1]
+            if domain == 'CONUS':
+                y_a, x_a = geos_nav.earth_to_lc(ll_pt_a['lon'], ll_pt_a['lat'])
+                y_b, x_b = geos_nav.earth_to_lc(ll_pt_b['lon'], ll_pt_b['lat'])
+
+    if satellite == 'H08':
+        xlen = taiwan_lenx
+        ylen = taiwan_leny
+        i_0 = taiwan_i0
+        j_0 = taiwan_j0
+    elif satellite == 'H09':
+        xlen = taiwan_lenx
+        ylen = taiwan_leny
+        i_0 = taiwan_i0
+        j_0 = taiwan_j0
+
+    grd_dct = {name: None for name in name_list}
+
+    cnt_a = 0
+    for ds_name in name_list:
+        fill_value, fill_value_name = get_fill_attrs(ds_name)
+        gvals = get_grid_values(h5f, ds_name, j_0, i_0, None, num_j=ylen, num_i=xlen, fill_value_name=fill_value_name, fill_value=fill_value)
+        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
+
+
+def make_for_full_domain_predict_viirs_clavrx(h5f, name_list=None, res_fac=1, day_night='DAY', use_dnb=False):
+    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]
+
+    use_nl_comp = False
+    if (day_night == 'NIGHT' or day_night == 'AUTO') and use_dnb:
+        use_nl_comp = True
+
+    grd_dct = {name: None for name in name_list}
+
+    cnt_a = 0
+    for ds_name in name_list:
+        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'
+
+        fill_value, fill_value_name = get_fill_attrs(name)
+        gvals = get_grid_values(h5f, name, j_0, i_0, None, num_j=ylen, num_i=xlen, fill_value_name=fill_value_name, fill_value=fill_value)
+        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')
+
+    # 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
+
+    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])
+
+    lats = get_grid_values(h5f, 'latitude', j_0, i_0, None, num_j=ylen, num_i=xlen)
+    lons = get_grid_values(h5f, 'longitude', j_0, i_0, None, num_j=ylen, num_i=xlen)
+
+    ll_2d, cc_2d = np.meshgrid(ll, cc, indexing='ij')
+
+    lats = lats[ll_2d, cc_2d]
+    lons = lons[ll_2d, cc_2d]
+
+    return grd_dct, ll, cc, lats, lons
+
+
+def make_for_full_domain_predict2(h5f, satellite='GOES16', domain='FD', res_fac=1):
+    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)
+
+    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
+
+    n_x = int(xlen/s_x)
+    n_y = int(ylen/s_y)
+
+    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)
+    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]
+
+    return solzen, satzen
+# -------------------------------------------------------------------------------------------
+
+
+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 = w_x // res_fac
+    s_y = w_y // res_fac
+
+    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
+    elif satellite == 'H09':
+        xlen = taiwan_lenx
+        ylen = taiwan_leny
+        i_0 = taiwan_i0
+        j_0 = taiwan_j0
+
+    n_x = (xlen // s_x)
+    n_y = (ylen // s_y)
+
+    # 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)]
+
+    grd_dct_n = {name: [] for name in name_list}
+
+    cnt_a = 0
+    for ds_name in name_list:
+        fill_value, fill_value_name = get_fill_attrs(ds_name)
+        gvals = get_grid_values(h5f, ds_name, j_0, i_0, None, num_j=ylen, num_i=xlen, fill_value_name=fill_value_name, fill_value=fill_value)
+        gvals = np.expand_dims(gvals, axis=0)
+        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)
+    solzen = solzen[0:n_y*s_y:s_y, 0:n_x*s_x:s_x]
+    satzen = satzen[0:n_y*s_y:s_y, 0:n_x*s_x:s_x]
+
+    return grd_dct_n, solzen, satzen, ll, cc
+
+
+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'
+
+
+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 == 'H09':
+        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 run_icing_predict(clvrx_dir='/Users/tomrink/data/clavrx/RadC/', output_dir=homedir,
                       day_model_path=model_path_day, night_model_path=model_path_night,
-- 
GitLab