diff --git a/modules/aeolus/aeolus_amv.py b/modules/aeolus/aeolus_amv.py index fe4ef424cfddbb416bab0755256aef77c7da3e1c..415b9a3405ce03b5f082aa3f4d7f89cb370abe25 100644 --- a/modules/aeolus/aeolus_amv.py +++ b/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