diff --git a/modules/util/util.py b/modules/util/util.py index c4475ff4ed7a39c277757ec716d01d37c08eb72c..47df416bd4a94efd24f02418597aa5a1e4c94d4a 100644 --- a/modules/util/util.py +++ b/modules/util/util.py @@ -776,7 +776,8 @@ def write_icing_file(clvrx_str_time, output_dir, preds_dct, probs_dct, x, y, lon def write_icing_file_nc4(clvrx_str_time, output_dir, preds_dct, probs_dct, - x, y, lons, lats, elems, lines, satellite='GOES16', domain='CONUS'): + x, y, lons, lats, elems, lines, satellite='GOES16', domain='CONUS', + has_time=False): outfile_name = output_dir + 'icing_prediction_'+clvrx_str_time+'.nc' rootgrp = Dataset(outfile_name, 'w', format='NETCDF4') @@ -795,40 +796,53 @@ def write_icing_file_nc4(clvrx_str_time, output_dir, preds_dct, probs_dct, 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', [dim_1_name, dim_0_name]) + 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.reshap((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', [dim_1_name, dim_0_name]) + 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') icing_prob_ds.setncattr('missing', -1.0) + if has_time: + probs = probs.reshap((1, y.shape[0], x.shape[0])) icing_prob_ds[:,] = probs prob_s = np.stack(prob_s, axis=-1) max_prob = np.max(prob_s, axis=2) + 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', [dim_1_name, dim_0_name]) + 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') icing_prob_ds.setncattr('missing', -1.0) icing_prob_ds[:,] = max_prob max_lvl = np.argmax(prob_s, axis=2) + if has_time: + max_lvl = max_lvl.reshap((1, y.shape[0], x.shape[0])) - icing_pred_ds = rootgrp.createVariable('max_icing_probability_level', 'i2', [dim_1_name, dim_0_name]) + 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)