snapshot..

modules/aeolus/aeolus_amv.py

@@ -2,6 +2,7 @@ import datetime, os
 from datetime import timezone
 import glob
 import numpy as np
+import xarray as xr
 from netCDF4 import Dataset, Dimension, Variable
 from aeolus.geos_nav import GEOSNavigation
 from util.util import haversine_np
@@ -444,8 +445,8 @@ def create_amv_to_aeolus_match_file(aeolus_files_dir, amv_files_dir, outfile=Non
 # dt_str_0: start time (YYYY-MM-DD_HH:MM)
 # dt_str_1: end time (YYYY-MM-DD_HH:MM)
 # returns:
-# amvs[ntimes, max_num_amvs_per_prof, num_of_params], profs[ntimes, max_num_levs_per_prof, num_of_params], times[ntimes]
-# longitude, latitude
+# amvs[nprofs, max_num_amvs_per_prof, num_of_params], profs[nprofs, max_num_levs_per_prof, num_of_params],
+# prof_times_locs[nprofs, (time, lon, lat)
 def subset_by_time(match_file, dt_str_0, dt_str_1):
     rootgrp = Dataset(match_file, 'r', format='NETCDF4')
     all_dims = rootgrp.dimensions
@@ -526,6 +527,16 @@ def subset_by_time(match_file, dt_str_0, dt_str_1):
 
     prof_times_locs = np.column_stack([times[time_idxs], lons[time_idxs], lats[time_idxs]])
 
-    return profs, amvs, prof_times_locs
+    coords = {'num_profs' : times[time_idxs], 'num_params' : ['speed', 'azimuth', 'layer_bot', 'layer_top']}
+    prof_da = xr.DataArray(profs, coords=coords, dims=['num_profs', 'max_num_levels', 'num_params'])
+
+    coords = {'num_profs': times[time_idxs], 'num_params': ['speed', 'azimuth', 'pressure', 'distance']}
+    amvs_da = xr.DataArray(amvs, coords=coords, dims=['num_profs', 'max_num_amvs', 'num_params'])
+
+    prof_locs_da = xr.DataArray(np.column_stack([lons[time_idxs], lats[time_idxs]]),
+                                coords=[times[time_idxs], ['longitude', 'latitude']],
+                                dims=['num_profs', 'space'])
+
+    return prof_da, amvs_da, prof_locs_da