diff --git a/modules/contrail/util.py b/modules/contrail/util.py index 6d6b47c4d8a9551d0a0e3638ce578963b9f6bcc8..f31389b0e416df2f29b20d63e83643562fd7f2fd 100644 --- a/modules/contrail/util.py +++ b/modules/contrail/util.py @@ -40,7 +40,6 @@ def get_contrail_mask_image(image, thresh=0.157): def extract(mask_image, image_ts, clavrx_path): gfs_file, _, _ = gfs_files.get_file(image_ts) - gfs_h5f = h5py.File(gfs_file, 'r') xr_dataset = xr.open_dataset(gfs_file) clvrx_h5f = h5py.File(clavrx_path, 'r') @@ -66,7 +65,36 @@ def extract(mask_image, image_ts, clavrx_path): contrail_lons = contrail_lons[keep] contrail_lats = contrail_lats[keep] - wind = get_point_s(xr_dataset, ['u-wind','v-wind'], contrail_lons, contrail_lats, contrail_press) + bins = np.arange(100, 1000, 100) + + # Indexes of contrail_press for individual bins + binned_indexes = np.digitize(contrail_press, bins) + + # Store the indexes in a dictionary where the key is the bin number and value is the list of indexes + bins_dict = {} + for bin_num in np.unique(binned_indexes): + bins_dict[bin_num] = np.where(binned_indexes == bin_num)[0] + + # This does point by point computation of model parameters for each contrail pixel + voxel_dict = {key: [] for key in bins_dict.keys()} + for key in bins_dict.keys(): + print('working on pressure level: ', bins[key]) + for c_idx in bins_dict[key]: + lon = contrail_lons[c_idx] + lat = contrail_lats[c_idx] + press = contrail_press[c_idx] + + wind_3d = get_voxel_s(xr_dataset, ['u-wind','v-wind'], lon, lat, press) + if wind_3d is not None: + voxel_dict[key].append(wind_3d) + + # This section will compute model parameters in bulk for the region then pull for each contrail pixel + lon_range = [np.min(contrail_lons), np.max(contrail_lons)] + lat_range = [np.min(contrail_lats), np.max(contrail_lats)] + + wind_3d = get_volume(xr_dataset, ['u-wind','v-wind'], lon_range=[lon_range[0], lon_range[1]], lat_range=[lat_range[0], lat_range[1]]) + + xr_dataset.close()