snapshot...

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()