snapshot...

modules/aeolus/aeolus_amv.py

@@ -274,7 +274,33 @@ def run_best_fit_driver(amv_dir, source, raob_path, gfs_path, product_dir, produ
         if gfs_file is None:
             continue
 
-        bf_dct = run_best_fit(m_d, raob_dct, gfs_file)
+        if gfs_file is not None:
+            locs = np.array(keys)
+            do_gfs_best_fit = True
+            xr_dataset = xr.open_dataset(gfs_file)
+            gfs_press = xr_dataset['pressure levels']
+            gfs_press = gfs_press.values
+            gfs_press = gfs_press[::-1]
+
+            lon_idx = LatLonTuple._fields.index('lon')
+            lat_idx = LatLonTuple._fields.index('lat')
+
+            uv_wind = get_vert_profile_s(xr_dataset, ['u-wind', 'v-wind'], locs[:, lon_idx], locs[:, lat_idx],
+                                         method='nearest')
+            uv_wind = uv_wind.values
+            wspd, wdir = spd_dir_from_uv(uv_wind[0, :, :], uv_wind[1, :, :])
+            wspd = wspd.magnitude
+            wdir = wdir.magnitude
+            gfs_spd = wspd[:, ::-1]
+            gfs_dir = wdir[:, ::-1]
+
+            gfs_at_raob_dct = {}
+            keys = list(raob_dct.keys())
+            for key_idx, key in enumerate(keys):
+                gfs_at_raob_dct[key] = (gfs_spd[key_idx], gfs_dir[key_idx], gfs_press)
+
+        #bf_dct = run_best_fit(m_d, raob_dct, gfs_file)
+        bf_dct = run_best_fit(m_d, raob_dct, gfs_at_raob_dct)
 
         prd_dct = None
         if prd_files is not None:
@@ -282,6 +308,9 @@ def run_best_fit_driver(amv_dir, source, raob_path, gfs_path, product_dir, produ
             if prd_dct is None:
                 continue
 
+        # TODO: write output file here
+
+        # Skip this and below if writing an output file.
         out_list.append((bf_dct, prd_dct))
 
     for tup in out_list:
@@ -464,39 +493,41 @@ def get_product_at_lat_lons(files, ts, lons, lats, filepath=None):
     return np.transpose(aaa, axes=[1, 0])
 
 
-def run_best_fit(raob_to_amv_dct, raob_dct, gfs_filename=None):
+def run_best_fit(raob_to_amv_dct, raob_dct, gfs_at_raob_dct):
     out_dct = {}
     keys = list(raob_to_amv_dct.keys())
 
-    do_gfs_best_fit = False
+    do_gfs_best_fit = True
     gfs_press = None
     gfs_spd = None
     gfs_dir = None
 
-    if gfs_filename is not None:
-        locs = np.array(keys)
-        do_gfs_best_fit = True
-        xr_dataset = xr.open_dataset(gfs_filename)
-        gfs_press = xr_dataset['pressure levels']
-        gfs_press = gfs_press.values
-        gfs_press = gfs_press[::-1]
-
-        lon_idx = LatLonTuple._fields.index('lon')
-        lat_idx = LatLonTuple._fields.index('lat')
-
-        uv_wind = get_vert_profile_s(xr_dataset, ['u-wind', 'v-wind'], locs[:, lon_idx], locs[:, lat_idx], method='nearest')
-        uv_wind = uv_wind.values
-        wspd, wdir = spd_dir_from_uv(uv_wind[0, :, :], uv_wind[1, :, :])
-        wspd = wspd.magnitude
-        wdir = wdir.magnitude
-        gfs_spd = wspd[:, ::-1]
-        gfs_dir = wdir[:, ::-1]
+    # if gfs_filename is not None:
+    #     locs = np.array(keys)
+    #     do_gfs_best_fit = True
+    #     xr_dataset = xr.open_dataset(gfs_filename)
+    #     gfs_press = xr_dataset['pressure levels']
+    #     gfs_press = gfs_press.values
+    #     gfs_press = gfs_press[::-1]
+    #
+    #     lon_idx = LatLonTuple._fields.index('lon')
+    #     lat_idx = LatLonTuple._fields.index('lat')
+    #
+    #     uv_wind = get_vert_profile_s(xr_dataset, ['u-wind', 'v-wind'], locs[:, lon_idx], locs[:, lat_idx], method='nearest')
+    #     uv_wind = uv_wind.values
+    #     wspd, wdir = spd_dir_from_uv(uv_wind[0, :, :], uv_wind[1, :, :])
+    #     wspd = wspd.magnitude
+    #     wdir = wdir.magnitude
+    #     gfs_spd = wspd[:, ::-1]
+    #     gfs_dir = wdir[:, ::-1]
 
     for key_idx, key in enumerate(keys):
         raob = raob_dct.get(key)
         raob_prs = raob[:, 0]
         raob_spd = raob[:, 3]
         raob_dir = raob[:, 2]
+        gfs = gfs_at_raob_dct.get(key)
+
         amvs = raob_to_amv_dct.get(key)
         num_amvs = amvs.shape[1]
         bf_list = []
@@ -524,7 +555,8 @@ def run_best_fit(raob_to_amv_dct, raob_dct, gfs_filename=None):
             raob_match_list.append(tup)
 
             if do_gfs_best_fit:
-                bf = best_fit(amv_spd, amv_dir, amv_prs, amv_lat, amv_lon, gfs_spd[key_idx], gfs_dir[key_idx], gfs_press)
+                #bf = best_fit(amv_spd, amv_dir, amv_prs, amv_lat, amv_lon, gfs_spd[key_idx], gfs_dir[key_idx], gfs_press)
+                bf = best_fit(amv_spd, amv_dir, amv_prs, amv_lat, amv_lon, gfs[0], gfs[1], gfs[2])
                 bf_gfs_list.append(bf)
 
         out_dct[key] = (np.transpose(amvs, axes=[1, 0]), np.array(bf_list), np.array(raob_match_list), np.array(bf_gfs_list))