diff --git a/modules/contrail/util.py b/modules/contrail/util.py index ab43a249ceef38fa0b35c3bb0bb5dd1f58fd9970..c6977a065bfa2175878b908f02714e6dffd3ac12 100644 --- a/modules/contrail/util.py +++ b/modules/contrail/util.py @@ -1,5 +1,6 @@ import re import datetime +import time from datetime import timezone import numpy as np @@ -46,6 +47,7 @@ def get_contrail_mask_image(image, thresh=0.157): def extract(mask_image, image_ts, clavrx_path): + # The GFS file closest to image_ts gfs_file, _, _ = gfs_files.get_file(image_ts) xr_dataset = xr.open_dataset(gfs_file) @@ -54,11 +56,11 @@ def extract(mask_image, image_ts, clavrx_path): clvrx_lons = get_grid_values_all(clvrx_h5f, 'longitude').flatten() clvrx_lats = get_grid_values_all(clvrx_h5f, 'latitude').flatten() - contrail_idxs = (mask_image == 1).flatten() - - contrail_press = cloud_top_press[contrail_idxs] - contrail_lons, contrail_lats = clvrx_lons[contrail_idxs], clvrx_lats[contrail_idxs] + contrail_mask = (mask_image == 1).flatten() + contrail_press = cloud_top_press[contrail_mask] + contrail_lons, contrail_lats = clvrx_lons[contrail_mask], clvrx_lats[contrail_mask] + # Use only the contrail obs where we have valid CLAVRx keep = np.invert(np.isnan(contrail_press)) contrail_press = contrail_press[keep] contrail_lons = contrail_lons[keep] @@ -97,6 +99,11 @@ def extract(mask_image, image_ts, clavrx_path): vert_shear_3d = volume_np_to_xr(vert_shear_3d, ['Pressure', 'Latitude', 'Longitude'], lon_range=lon_range, lat_range=lat_range) vert_shear_3d = vert_shear_3d / units.hPa + # Form a new Dataset for interpolation below + da_dct = {'horz_shear_3d': horz_shear_3d, 'static_3d': static_3d, 'horz_wind_spd_3d': horz_wind_spd_3d, + 'vert_shear_3d': vert_shear_3d, 'temp_3d': temp_3d, 'rh_3d': rh_3d} + xr_ds = xr.Dataset(da_dct) + all_list = [] levels_dict = {press_bins[key]: [] for key in bins_dict.keys()} levels_dict_cidx = {press_bins[key]: [] for key in bins_dict.keys()} @@ -110,15 +117,10 @@ def extract(mask_image, image_ts, clavrx_path): if lon < 0: # Match GFS convention lon += 360.0 - horz_shear_value = horz_shear_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest').item(0) - static_value = static_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest').item(0) - horz_wind_spd_value = horz_wind_spd_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest').item(0) - vert_shear_value = vert_shear_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest').item(0) - temp_value = temp_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest').item(0) - rh_value = rh_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest').item(0) - - # tmp = horz_shear_3d.sel(Pressure=press, method='nearest') - # tmp = tmp.sel(Longitude=lon, Latitude=lat, method='nearest') + intrp_ds = xr_ds.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest') + horz_shear_value, static_value, horz_wind_spd_value, vert_shear_value, temp_value, rh_value = ( + intrp_ds['horz_shear_3d'].item(0), intrp_ds['static_3d'].item(0), intrp_ds['horz_wind_spd_3d'].item(0), + intrp_ds['vert_shear_3d'].item(0), intrp_ds['temp_3d'].item(0), intrp_ds['rh_3d'].item(0)) levels_dict_cidx[press_level].append(c_idx) levels_dict[press_level].append((press, lat, lon, temp_value, rh_value, horz_shear_value, static_value, horz_wind_spd_value, vert_shear_value))