add new method

modules/aeolus/aeolus_amv.py

@@ -1671,6 +1671,79 @@ def get_search_box(nav, lon, lat):
 
     return c_rng, l_rng
 
+# aeolus_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_aeolus_fast(aeolus_dict, 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()
+    match_dict = {}
+
+    keys = list(aeolus_dict.keys())
+
+    amv_lons = None
+    amv_lats = None
+    cc = None
+    ll = None
+
+    last_f_idx = -1
+    for key in keys:
+        fname, ftime, f_idx = amv_files.get_file_containing_time(key)
+        if f_idx is None:
+            continue
+
+        profs = aeolus_dict.get(key)
+        if profs is None:
+            continue
+
+        match_dict[key] = []
+
+        if f_idx != last_f_idx:
+            last_f_idx = f_idx
+            ds = Dataset(fname)
+
+            amv_lons = ds[amv_files.lon_name][:]
+            amv_lats = ds[amv_files.lat_name][:]
+            cc = ds[amv_files.elem_name][:]
+            ll = ds[amv_files.line_name][:]
+            # cc, 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(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()
+
+        for prof in profs:
+            lat = prof[0, 0]
+            lon = prof[0, 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].append(param_nd)
+
+    return match_dict
 
 # aeolus_dict: time -> profiles
 # amv_files_path: directory containing AMVs, '/home/user/amvdir/'