snapshot...

modules/aeolus/aeolus_amv.py

@@ -309,6 +309,8 @@ def run_best_fit_driver(amv_dir, source, raob_path, gfs_path, product_dir, produ
                 continue
 
         # TODO: write output file here
+        filename = None
+        create_file_new2(filename, raob_dct, gfs_at_raob_dct, bf_dct, prd_dct, raob_filename, amv_file_name)
 
         # Skip this and below if writing an output file.
         out_list.append((bf_dct, prd_dct))
@@ -2010,6 +2012,150 @@ def create_file(filename, aeolus_to_amv_dct, aeolus_dct, amv_files, cld_lyr=Fals
 
     rootgrp.close()
 
+def create_file_new2(filename, raob_dct, gfs_at_raob_dct, bf_dct, prd_dct, raob_filename, amv_file_name):
+    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(raob_dct.keys())
+    for key in keys:
+        raob = raob_dct.get(key)
+        bf_tup = bf_dct.get(key)
+        amvs = bf_tup[0]
+        num_amvs = amvs.shape[0]
+        nlevs = raob.shape[0]
+        lat = key[0]
+        lon = key[1]
+        if nlevs > max_num_alevels:
+            max_num_alevels = nlevs
+        if num_amvs > max_num_amvs:
+            max_num_amvs = num_amvs
+
+    #     tup_s = match_dct.get(key)
+    #     for tup in tup_s:
+    #         num_aprofs += 1
+    #         prof = tup[3]
+    #         lat = prof[0, 0]
+    #         lon = prof[0, 1]
+    #         alons.append(lon)
+    #         alats.append(lat)
+    #         atimes.append(key)
+    #         elems.append(tup[0])
+    #         lines.append(tup[1])
+    #         fidxs.append(tup[2])
+    #
+    #         nlevs = prof.shape[0]
+    #         if nlevs > max_num_alevels:
+    #             max_num_alevels = nlevs
+    #
+    #         amvs = tup[4]
+    #         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)
+    print(max_num_alevels, max_num_amvs)
+    return
+
+    # 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')
+
+    # 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_aeolus_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_aeolus_profs'])
+    amv_file_names = rootgrp.createVariable('amv_file_names', str, ['num_amv_files'])
+
+    # ---- Profile variables ---------------
+    prf_lon = rootgrp.createVariable('prof_longitude', 'f4', ['num_aeolus_profs'])
+    prf_lon.units = 'degrees east'
+    prf_lat = rootgrp.createVariable('prof_latitude', 'f4', ['num_aeolus_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_aeolus_profs'])
+    prf_line = rootgrp.createVariable('FD_line', 'f4', ['num_aeolus_profs'])
+    prf_fidx = rootgrp.createVariable('amv_file_index', 'i4', ['num_aeolus_profs'])
+    prf_azm = rootgrp.createVariable('prof_azm', 'f4', ['num_aeolus_profs', 'max_num_aeolus_levels'])
+    prf_azm.units = 'degree'
+    prf_spd = rootgrp.createVariable('prof_spd', 'f4', ['num_aeolus_profs', 'max_num_aeolus_levels'])
+    prf_spd.units = 'm s-1'
+    prf_hht = rootgrp.createVariable('prof_hht', 'f4', ['num_aeolus_profs', 'max_num_aeolus_levels'])
+    prf_hht.units = 'meter'
+    prf_hhb = rootgrp.createVariable('prof_hhb', 'f4', ['num_aeolus_profs', 'max_num_aeolus_levels'])
+    prf_hhb.units = 'meter'
+
+    # ----- Product 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_aeolus_profs', 'max_num_amvs'], fill_value=v.getncattr('_FillValue'))
+        else:
+            var = rootgrp.createVariable(param, v.dtype, ['num_aeolus_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:
+            prof = tup[3]
+            param_nd = tup[4]
+
+            nlevs = prof.shape[0]
+            for k in range(nlevs):
+                prf_spd[idx,k] = prof[k,6]
+                prf_azm[idx,k] = prof[k,5]
+                prf_hht[idx,k] = prof[k,3]
+                prf_hhb[idx,k] = prof[k,4]
+
+            for pidx, param in enumerate(amv_params):
+                nda = param_nd[pidx+4,]
+                cnt = nda.shape[0]
+                nc4_vars[pidx][idx, 0:cnt,] = nda
+            idx += 1
+
+    rg_exmpl.close()
+    rootgrp.close()
+
+
 def create_file_new(match_dct, filename, amv_params, amv_file_s):
     num_aparams = 7
     num_aprofs = 0