infer_cloud_products.py 2.43 KiB
from util.geos_nav import get_navigation
from util.setup_cloud_products import model_path_cld_frac, model_path_cld_opd
from aeolus.datasource import CLAVRx
import os
from deeplearning.cloud_fraction_fcn_abi import SRCNN as SRCNN_CLD_FRAC
from deeplearning.cloud_opd_fcn_abi import SRCNN as SRCNN_CLD_OPD
from util.util import get_cartopy_crs, write_cld_prods_file_nc4
import numpy as np
import time
def infer_cloud_products(clvrx_path, output_dir, full_disk=True, satellite='GOES16', domain='FD', pattern=None):
# -- location of the trained model
ckpt_dir_s = os.listdir(model_path_cld_frac)
ckpt_dir_cld_frac = model_path_cld_frac + ckpt_dir_s[0]
ckpt_dir_s = os.listdir(model_path_cld_opd)
ckpt_dir_cld_opd = model_path_cld_opd + ckpt_dir_s[0]
# -- Navigation parameters
geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
nav = get_navigation(satellite, domain)
cc = np.arange(xlen)
ll = np.arange(ylen)
x_rad = cc * nav.CFAC + nav.COFF
y_rad = ll * nav.LFAC + nav.LOFF
# -- Create a model instance and initialize with trained model above
nn_cld_frac = SRCNN_CLD_FRAC()
nn_cld_frac.setup_inference(ckpt_dir_cld_frac)
nn_cld_opd = SRCNN_CLD_OPD()
nn_cld_opd.setup_inference(ckpt_dir_cld_opd)
if pattern is not None:
clvrx_ds = CLAVRx(clvrx_path, pattern=pattern)
else:
clvrx_ds = CLAVRx(clvrx_path)
for pname, t_start, t_stop, dto in clvrx_ds:
clvrx_str_time = dto.strftime('%Y-%m-%d_%H:%M')
out_file = output_dir + 'cloud_fraction_opd_' + clvrx_str_time + '.nc'
t0 = time.time()
if full_disk:
cld_frac = nn_cld_frac.run_inference_full_disk(pname, None)
cld_opd, opd = nn_cld_opd.run_inference_full_disk(pname, None)
else:
cld_frac = nn_cld_frac.run_inference(pname, None)
cld_opd, opd = nn_cld_opd.run_inference(pname, None)
clr_cat = cld_frac == 0
cld_cat = cld_frac == 4
cld_opd[clr_cat] = opd[clr_cat]
cld_opd[cld_cat] = opd[cld_cat]
cld_opd[np.logical_and(cld_opd < 0.0, cld_opd > 160.0)] = -1.0
cld_opd[np.isnan(opd)] = -1.0
write_cld_prods_file_nc4(clvrx_str_time, out_file, cld_frac, cld_opd, x_rad, y_rad, None, None,
satellite=satellite, domain=domain, has_time=True)
t1 = time.time()
print('total time: ', (t1 - t0))
print('Done: ', pname)