snapshot...

modules/aeolus/aeolus_amv.py

@@ -81,8 +81,18 @@ class OPS(AMVFiles):
     def __init__(self, files_path, file_time_span, band='14'):
         super().__init__(files_path, file_time_span, 'OR_ABI-L2-DMWF*'+'C'+band+'*.nc', band)
 
+        self.elem_name = None
+        self.line_name = None
+        self.lon_name = 'lon'
+        self.lat_name = 'lat'
+
+        self.out_params = {'Lon', 'Lat', 'Element', 'Line', 'pressure', 'wind_speed', 'wind_direction'}
+        self.params = {'pressure', 'wind_speed', 'wind_direction'}
+        self.meta_dict = {'Lon': ('degrees east', 'f4'), 'Lat': ('degrees north', 'f4'), 'Element': (None, 'i4'), 'Line': (None, 'i4'),
+                          'pressure': ('hPa', 'f4'), 'wind_speed': ('m s-1', 'f4'), 'wind_direction': ('degrees', 'f4')}
+
     def get_navigation(self):
-        return GEOSNavigation(sub_lon=-75.0)
+        return GEOSNavigation(sub_lon=-75.2)
 
     def get_datetime(self, pathname):
         fname = os.path.split(pathname)[1]
@@ -93,6 +103,15 @@ class OPS(AMVFiles):
 
         return dto
 
+    def get_parameters(self):
+        return self.params
+
+    def get_out_parameters(self):
+        return self.out_params
+
+    def get_meta_dict(self):
+        return self.meta_dict
+
 
 class CarrStereo(AMVFiles):
     def __init__(self, files_path, file_time_span, band='14'):
@@ -137,6 +156,165 @@ class CarrStereo(AMVFiles):
         return self.meta_dict
 
 
+# raob_dict: time -> profiles
+# amv_files_path: directory containing AMVs, '/home/user/amvdir/'
+# return dict: aeolus time -> tuple (amv_lon, amv_lat, amv_pres, amv_spd, amv_dir)
+def match_amvs_to_raobs(raob_dict,  raob_time, amv_files):
+    nav = amv_files.get_navigation()
+    amv_params = amv_files.get_parameters()
+    match_dict = {}
+
+    #fname, ftime, f_idx = amv_files.get_file_containing_time(raob_time)
+    fname = '/Users/tomrink/data/OR_ABI-L2-DMWF-M6C14_G16_s20201190000156_e20201190009464_c20201190023107.nc'
+
+    ds = Dataset(fname)
+
+    amv_lons = ds[amv_files.lon_name][:]
+    amv_lats = ds[amv_files.lat_name][:]
+    if amv_files.elem_name is None or amv_files.line_name is None:
+        cc, ll = nav.earth_to_lc_s(amv_lons, amv_lats)
+    else:
+        cc = ds[amv_files.elem_name][:]
+        ll = ds[amv_files.line_name][:]
+
+    param_s = []
+    param_s.append(amv_lons)
+    param_s.append(amv_lats)
+    param_s.append(cc)
+    param_s.append(ll)
+    for param in amv_params:
+        if param == 'V_3D':
+            param_s.append(ds[param][:, 0])
+            param_s.append(ds[param][:, 1])
+        else:
+            param_s.append(ds[param][:])
+
+    ds.close()
+
+    keys = list(raob_dict.keys())
+
+    for key in keys:
+
+        raob = raob_dict.get(key)
+        if raob is None:
+            continue
+
+        lat = key[0]
+        lon = key[1]
+
+        c_rng, l_rng = get_search_box(nav, lon, lat)
+        if c_rng is None:
+            continue
+
+        in_cc = np.logical_and(cc > c_rng[0], cc < c_rng[1])
+        in_ll = np.logical_and(ll > l_rng[0], 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]
+        match_dict[key] = param_nd
+
+    return match_dict
+
+
+def create_file2(filename, raob_to_amv_dct, raob_dct, amv_files):
+    keys = list(raob_to_amv_dct.keys())
+
+    num_amvs = []
+    num_levs = []
+    times = []
+
+    namvs = 0
+    nlevs = 0
+    for key in keys:
+        param_nd = raob_to_amv_dct.get(key)
+        num_amvs.append(param_nd.shape[1])
+        namvs += param_nd.shape[1]
+
+        prof = raob_dct.get(key)
+        num_levs.append(prof.shape[0])
+        nlevs += prof.shape[0]
+
+        # TODO: need a time
+        times.append(0.0)
+
+    # amv_per_alus = len(aeolus_to_amv_dct)
+    rootgrp = Dataset(filename, 'w', format='NETCDF4')
+    dim_amvs = rootgrp.createDimension('amvs', size=namvs)
+    dim_alus = rootgrp.createDimension('raobs', size=nlevs)
+    dim_num_aeolus_prof = rootgrp.createDimension('num_raobs', size=len(raob_to_amv_dct))
+
+    nc4_vars = []
+    out_params = amv_files.get_out_parameters()
+    meta_dict = amv_files.get_meta_dict()
+
+    for pidx, param in enumerate(out_params):
+        u, t = meta_dict.get(param)
+        var = rootgrp.createVariable(param, t, ['amvs'])
+        if u is not None:
+            var.units = u
+        nc4_vars.append(var)
+
+    dist = rootgrp.createVariable('dist_to_raob', 'f4', ['amvs'])
+    dist.units = 'km'
+
+    num_amvs_per_raob = rootgrp.createVariable('num_amvs_per_raob', 'i4', ['num_raobs'])
+    num_levs_per_raob = rootgrp.createVariable('num_levs_per_raob', 'i4', ['num_raobs'])
+    prof_time = rootgrp.createVariable('time', 'f4', ['num_raobs'])
+    # ---- Profile variables ---------------
+    prf_lon = rootgrp.createVariable('raob_longitude', 'f4', ['num_raobs'])
+    prf_lon.units = 'degrees east'
+    prf_lat = rootgrp.createVariable('raob_latitude', 'f4', ['num_raobs'])
+    prf_lat.units = 'degrees north'
+    prof_time.units = 'seconds since 1970-01-1 00:00:00'
+
+    prf_azm = rootgrp.createVariable('prof_azm', 'f4', ['raobs'])
+    prf_azm.units = 'degree'
+    prf_spd = rootgrp.createVariable('prof_spd', 'f4', ['raobs'])
+    prf_spd.units = 'm s-1'
+    prf_prs = rootgrp.createVariable('prof_pres', 'f4', ['raobs'])
+    prf_prs.units = 'hPa'
+    prf_tmp = rootgrp.createVariable('prof_temp', 'f4', ['raobs'])
+    prf_tmp.units = 'K'
+    # --------------------------------------
+
+    i_a = 0
+    i_c = 0
+    for idx, key in enumerate(keys):
+        namvs = num_amvs[idx]
+        nlevs = num_levs[idx]
+        i_b = i_a + namvs
+        i_d = i_c + nlevs
+
+        prof = raob_dct.get(key)
+        prf_azm[i_c:i_d] = prof[:, 2]
+        prf_spd[i_c:i_d] = prof[:, 3]
+        prf_prs[i_c:i_d] = prof[:, 0]
+        prf_tmp[i_c:i_d] = prof[:, 1]
+        i_c += nlevs
+
+        plat = key[0]
+        plon = key[1]
+        prf_lat[idx::] = plat
+        prf_lon[idx::] = plon
+
+        param_nd = raob_to_amv_dct.get(key)
+        for pidx, param in enumerate(out_params):
+            nc4_vars[pidx][i_a:i_b] = param_nd[pidx, :]
+        dist[i_a:i_b] = haversine_np(plon, plat, param_nd[0, :], param_nd[1, :])
+        i_a += namvs
+
+    num_amvs_per_raob[:] = num_amvs
+    num_levs_per_raob[:] = num_levs
+    prof_time[:] = times
+
+    rootgrp.close()
+
+
 # imports the S4 NOAA output
 # filename: full path as a string, '/home/user/filename'
 # returns a dict: time -> list of profiles (a profile is a list of levels)