initial commit

modules/amv/caliop_clavrx_amv.py

+import numpy as np
+import xarray as xr
+from netCDF4 import Dataset
+from aeolus.datasource import get_datasource, CLAVRx_CALIPSO
+
+amv_file_duration = 60  # minutes
+half_width = 50  # search box centered on AEOLUS profile (FGF coordinates)
+num_elems = 5424
+num_lines = 5424
+
+
+def get_search_box(cc, ll):
+
+    c_rng = [cc - half_width, cc + half_width]
+    l_rng = [ll - half_width, ll + half_width]
+
+    if c_rng[0] < 0:
+        c_rng[0] = 0
+
+    if l_rng[0] < 0:
+        l_rng[0] = 0
+
+    if c_rng[1] >= num_elems:
+        c_rng[1] = num_elems - 1
+
+    if l_rng[1] >= num_lines:
+        l_rng[1] = num_lines - 1
+
+    return c_rng, l_rng
+
+
+def match_calipso_clavrx_to_amvs(calipso_clavrx_path, calipso_clavrx_file, amv_files_path, amv_source='OPS', band='14', amv_files=None):
+    if amv_files is None:
+        amv_files = get_datasource(amv_files_path, amv_source, band=band)
+    nav = amv_files.get_navigation()
+    amv_params = amv_files.get_parameters()
+    all_params = [amv_files.lon_name, amv_files.lat_name, amv_files.elem_name, amv_files.line_name] + amv_params
+    match_dict = {}
+    coords = {'num_params': all_params}
+    dims = ['num_params', 'num_amvs']
+
+    calipso_clavrx_ds = CLAVRx_CALIPSO(calipso_clavrx_path)
+
+    nom_time = calipso_clavrx_ds.get_datetime(calipso_clavrx_file).timestamp()
+    calipso_clavrx_nc4 = Dataset(calipso_clavrx_file)
+
+    fname, ftime, f_idx = amv_files.get_file_containing_time(nom_time)
+    if f_idx is None:
+        return
+
+    match_dict[nom_time] = []
+
+    amv_ds = Dataset(fname)
+    amv_lons = amv_ds[amv_files.lon_name][:]
+    amv_lats = amv_ds[amv_files.lat_name][:]
+    if amv_files.elem_name is not None:
+        amv_cc = amv_ds[amv_files.elem_name][:]
+        amv_ll = amv_ds[amv_files.line_name][:]
+    else:
+        amv_cc, amv_ll = nav.earth_to_lc_s(amv_lons, amv_lats)
+
+    param_s = []
+    param_s.append(amv_lons)
+    param_s.append(amv_lats)
+    param_s.append(amv_cc)
+    param_s.append(amv_ll)
+    for param in amv_params:
+        if param == 'V_3D':
+            param_s.append(amv_ds[param][:, 0])
+            param_s.append(amv_ds[param][:, 1])
+        else:
+            param_s.append(amv_ds[param][:])
+
+    amv_ds.close()
+
+    calipso_clavrx_params = ['closest_calipso_top_alt', 'closest_calipso_base_alt', 'cld_height_acha', 'cld_height_base_acha']
+    coords_a = {'num_calipso_clavrx_params': calipso_clavrx_params}
+    dims_a = ['num_calipso_clavrx_params']
+    calipso_clavrx_data = []
+    for pname in calipso_clavrx_params:
+        calipso_clavrx_data.append(calipso_clavrx_nc4[pname][:])
+    lat_s = calipso_clavrx_nc4['latitude'][:]
+    lon_s = calipso_clavrx_nc4['longitude'][:]
+    xs = calipso_clavrx_nc4['x'][:]
+    ys = calipso_clavrx_nc4['y'][:]
+    clvr_xy_s = zip(xs, ys)
+
+    for idx, clvr_xy in enumerate(clvr_xy_s):
+        cc, ll = clvr_xy
+
+        c_rng, l_rng = get_search_box(cc, ll)
+
+        in_cc = np.logical_and(amv_cc > c_rng[0], amv_cc < c_rng[1])
+        in_ll = np.logical_and(amv_ll > l_rng[0], amv_ll < l_rng[1])
+        in_box = np.logical_and(in_cc, in_ll)
+
+        num_amvs = np.sum(in_box)
+        if num_amvs == 0:
+            continue
+        # dist = haversine_np(lon, lat, amv_lons[in_box], amv_lats[in_box])
+        param_nd = np.vstack(param_s)
+        param_nd = param_nd[:, in_box]
+        amvs_da = xr.DataArray(param_nd, coords=coords, dims=dims)
+
+        data_nd = np.vstack(calipso_clavrx_data)
+        data_da = xr.DataArray(data_nd[:, idx], coords=coords_a, dims=dims_a)
+
+        match_dict[nom_time].append((cc, ll, lon_s[idx], lat_s[idx], f_idx, data_da, amvs_da))
+
+    return match_dict
+
+
+def create_file_new(match_dct, filename, amv_params=['pressure', 'wind_speed', 'wind_direction'],
+                    caliop_clavrx_params=['closest_calipso_top_alt', 'closest_calipso_base_alt', 'cld_height_acha', 'cld_height_base_acha']):
+    num_aparams = 7
+    num_aprofs = 0
+    max_num_alevels = 0
+    max_num_amvs = 0
+
+    alons = []
+    alats = []
+    atimes = []
+    elems = []
+    lines = []
+    fidxs = []
+
+    #amv_file_s = np.array(amv_file_s, dtype='object')
+
+    # scan to get max num levels, amvs
+    keys = list(match_dct.keys())
+    for key in keys:
+        tup_s = match_dct.get(key)
+        for tup in tup_s:
+            num_aprofs += 1
+            atimes.append(key)
+            elems.append(tup[0])
+            lines.append(tup[1])
+            alons.append(tup[2])
+            alats.append(tup[3])
+            fidxs.append(tup[4])
+
+            amvs = tup[6]
+            num_amvs = amvs.shape[1]
+            if num_amvs > max_num_amvs:
+                max_num_amvs = num_amvs
+
+    alons = np.array(alons)
+    alats = np.array(alats)
+    atimes = np.array(atimes)
+    elems = np.array(elems)
+    lines = np.array(lines)
+    fidxs = np.array(fidxs)
+
+    # Sample file to retrieve and copy variable attributes
+    # rg_exmpl = Dataset('/ships19/cloud/scratch/4TH_AMV_INTERCOMPARISON/FMWK2_AMV/GOES16_ABI_2KM_FD_2019293_0020_34_WINDS_AMV_EN-14CT.nc', 'r')
+    rg_exmpl = Dataset('/ships19/cloud/scratch/AMV_FAST_BIAS/OPS_AMVS/OR_ABI-L2-DMWF-M6C14_G16_s20201050000166_e20201050009474_c20201050023226.nc')
+
+    caliop_clavrx_exmpl = Dataset('/data/Personal/rink/clavrx_calipso/g16_s20201050200_06kmCLay.matchup.calipso.h5', 'r')
+
+    # the top level group for the output file
+    rootgrp = Dataset(filename, 'w', format='NETCDF4')
+
+    #dim_aparams = rootgrp.createDimension('num_aeolus_params', size=num_aparams)
+    #dim_alevs = rootgrp.createDimension('max_num_aeolus_levels', size=max_num_alevels)
+    dim_amvs = rootgrp.createDimension('max_num_amvs', size=max_num_amvs)
+    dim_num_aeolus_prof = rootgrp.createDimension('num_caliop_profs', size=num_aprofs)
+    #dim_num_files = rootgrp.createDimension('num_amv_files', size=amv_file_s.shape[0])
+
+    prf_time = rootgrp.createVariable('time', 'f4', ['num_caliop_profs'])
+    #amv_file_names = rootgrp.createVariable('amv_file_names', str, ['num_amv_files'])
+
+    # ---- Profile variables ---------------
+    prf_lon = rootgrp.createVariable('prof_longitude', 'f4', ['num_caliop_profs'])
+    prf_lon.units = 'degrees east'
+    prf_lat = rootgrp.createVariable('prof_latitude', 'f4', ['num_caliop_profs'])
+    prf_lat.units = 'degrees north'
+    prf_time.units = 'seconds since 1970-01-1 00:00:00'
+    prf_elem = rootgrp.createVariable('FD_elem', 'f4', ['num_caliop_profs'])
+    prf_line = rootgrp.createVariable('FD_line', 'f4', ['num_caliop_profs'])
+    prf_fidx = rootgrp.createVariable('amv_file_index', 'i4', ['num_caliop_profs'])
+
+    # ----- CALIOP+CLAVRx variables ------
+    nc4_vars_cc = []
+    var_s = caliop_clavrx_exmpl.variables
+    for pidx, param in enumerate(caliop_clavrx_params):
+        v = var_s[param]
+        attr_s = v.ncattrs()
+        if '_FillValue' in attr_s:
+            var = rootgrp.createVariable(param, v.dtype, ['num_caliop_profs'], fill_value=v.getncattr('_FillValue'))
+        else:
+            var = rootgrp.createVariable(param, v.dtype, ['num_caliop_profs'])
+        # copy attributes from example to new output variable of the same name
+        for attr in attr_s:
+            if attr != '_FillValue':
+                var.setncattr(attr, v.getncattr(attr))
+        nc4_vars_cc.append(var)
+
+    # ----- AMV variables ----------------
+    nc4_vars = []
+    var_s = rg_exmpl.variables
+    for pidx, param in enumerate(amv_params):
+        v = var_s[param]
+        attr_s = v.ncattrs()
+        if '_FillValue' in attr_s:
+            var = rootgrp.createVariable(param, v.dtype, ['num_caliop_profs', 'max_num_amvs'], fill_value=v.getncattr('_FillValue'))
+        else:
+            var = rootgrp.createVariable(param, v.dtype, ['num_caliop_profs', 'max_num_amvs'])
+        # copy attributes from example to new output variable of the same name
+        for attr in attr_s:
+            if attr != '_FillValue':
+                var.setncattr(attr, v.getncattr(attr))
+        nc4_vars.append(var)
+
+    # Write data to file  ---------------------
+    prf_lon[:] = alons
+    prf_lat[:] = alats
+    prf_elem[:] = elems
+    prf_line[:] = lines
+    prf_time[:] = atimes
+    prf_fidx[:] = fidxs
+
+    #amv_file_names[:] = amv_file_s
+
+    idx = 0
+    for key in keys:
+        tup_s = match_dct.get(key)
+        for tup in tup_s:
+            data_nd = tup[5]
+            amvs_nd = tup[6]
+
+            for pidx, param in enumerate(caliop_clavrx_params):
+                nda = data_nd[pidx]
+                nc4_vars_cc[pidx][idx] = nda
+
+            for pidx, param in enumerate(amv_params):
+                nda = amvs_nd[pidx+4,]
+                cnt = nda.shape[0]
+                nc4_vars[pidx][idx, 0:cnt] = nda
+            idx += 1
+
+    rg_exmpl.close()
+    rootgrp.close()
+
+
+# def run_clavrx_to_aeolus_match(aeolus_files_dir, clvrx_files_dir, outfile, chan='mie'):
+#     if os.path.isdir(aeolus_files_dir):
+#         a_dct = get_aeolus_time_dict_s(aeolus_files_dir, chan=chan)
+#     else:
+#         a_dct = get_aeolus_time_dict(aeolus_files_dir)
+#
+#     a_dct = time_dict_to_nd(a_dct)
+#     clvrx_files = CLAVRx(clvrx_files_dir)
+#     m_d = match_aeolus_to_clavrx(a_dct, clvrx_files)
+#     create_aeolus_clavrx_match_file(m_d, outfile, clvrx_files.get_parameters(), clvrx_files.flist)
+#
+#
+# # aeolus_files_dir: S4 NOAA txt output files
+# # amv_files_dir: G16/17 AMV product files
+# # outfile: pathname for the Netcdf match file
+# def create_amv_to_aeolus_match_file(aeolus_files_dir, amv_files_dir, outfile, amv_source='OPS', band='14', chan='mie'):
+#
+#     if os.path.isdir(aeolus_files_dir):
+#         a_d = get_aeolus_time_dict_s(aeolus_files_dir, chan=chan)
+#     else:
+#         a_d = get_aeolus_time_dict(aeolus_files_dir)
+#     a_d = time_dict_to_nd(a_d)
+#
+#     amv_files = None
+#     if amv_source == 'CARR':
+#         amv_files = get_datasource(amv_files_dir, 'CARR', file_time_span=60, band=band)
+#     else:
+#         amv_files = get_datasource(amv_files_dir, amv_source, band=band)
+#
+#     m_d = match_amvs_to_aeolus_fast(a_d, amv_files_dir, amv_source, band, amv_files)
+#
+#     create_file_new(m_d, outfile, amv_files.get_parameters(), amv_files.flist)