snapshot...

modules/contrail/util.py

@@ -7,9 +7,9 @@ import xarray as xr
 import pandas as pd
 import rasterio
 from PIL import Image
-import matplotlib.pyplot as plt
 import matplotlib.image as mpimg
 import h5py
+from util.GFSDataset import GFSData
 from util.util import get_grid_values_all
 from util.gfs_reader import get_volume, volume_np_to_xr
 from util.geos_nav import GEOSNavigation
@@ -48,8 +48,8 @@ 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)
 
+    # Open the CLAVRx file
     clvrx_h5f = h5py.File(clavrx_path, 'r')
     cloud_top_press = get_grid_values_all(clvrx_h5f, 'cld_press_acha').flatten()
     clvrx_lons = get_grid_values_all(clvrx_h5f, 'longitude').flatten()
@@ -78,30 +78,32 @@ def extract(mask_image, image_ts, clavrx_path):
     lon_range = [np.min(contrail_lons), np.max(contrail_lons)]
     lat_range = [np.min(contrail_lats), np.max(contrail_lats)]
 
-    uwind_3d = get_volume(xr_dataset, 'u-wind', 'm s-1', lon_range=lon_range, lat_range=lat_range)
-    vwind_3d = get_volume(xr_dataset, 'v-wind', 'm s-1', lon_range=lon_range, lat_range=lat_range)
-    temp_3d = get_volume(xr_dataset, 'temperature', 'degK', lon_range=lon_range, lat_range=lat_range)
-    rh_3d = get_volume(xr_dataset, 'rh', '%', lon_range=lon_range, lat_range=lat_range)
-
-    uwind_3d = uwind_3d.transpose('Pressure', 'Latitude', 'Longitude')
-    vwind_3d = vwind_3d.transpose('Pressure', 'Latitude', 'Longitude')
-    temp_3d = temp_3d.transpose('Pressure', 'Latitude', 'Longitude')
-    rh_3d = rh_3d.transpose('Pressure', 'Latitude', 'Longitude')
-
-    horz_shear_3d = shearing_deformation(uwind_3d, vwind_3d)
-    static_3d = static_stability(temp_3d.coords['Pressure'] * units.hPa, temp_3d)
-    horz_wind_spd_3d = wind_speed(uwind_3d, vwind_3d)
-
-    # This one's a bit more work: `first_derivative` only returns a ndarray with no units, so we use the
-    # helper function to create a DataArray and add units via metpy's pint support
-    vert_shear_3d = first_derivative(horz_wind_spd_3d, axis=0, x=temp_3d.coords['Pressure'])
-    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)
+    # create an instance of the file data accessor and initialize the region of interest
+    with GFSData(gfs_file, lon_range=lon_range, lat_range=lat_range) as gfs_dataset:
+        uwind_3d = gfs_dataset.get_volume('u-wind', 'm s-1')
+        vwind_3d = gfs_dataset.get_volume('v-wind', 'm s-1')
+        temp_3d = gfs_dataset.get_volume('temperature', 'degK')
+        rh_3d = gfs_dataset.get_volume('rh', '%')
+
+        uwind_3d = uwind_3d.transpose('Pressure', 'Latitude', 'Longitude')
+        vwind_3d = vwind_3d.transpose('Pressure', 'Latitude', 'Longitude')
+        temp_3d = temp_3d.transpose('Pressure', 'Latitude', 'Longitude')
+        rh_3d = rh_3d.transpose('Pressure', 'Latitude', 'Longitude')
+
+        horz_shear_3d = shearing_deformation(uwind_3d, vwind_3d)
+        static_3d = static_stability(temp_3d.coords['Pressure'] * units.hPa, temp_3d)
+        horz_wind_spd_3d = wind_speed(uwind_3d, vwind_3d)
+
+        # This one's a bit more work: `first_derivative` only returns a ndarray with no units, so we use the
+        # helper function to create a DataArray and add units via metpy's pint support
+        vert_shear_3d = first_derivative(horz_wind_spd_3d, axis=0, x=temp_3d.coords['Pressure'])
+        vert_shear_3d = gfs_dataset.volume_np_to_xr(vert_shear_3d, ['Pressure', 'Latitude', 'Longitude'])
+        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()}
@@ -147,13 +149,11 @@ def extract(mask_image, image_ts, clavrx_path):
                           columns=["pressure_level", "pressure", "lat", "lon", "temperature", "relative_humidity",
                                    "horz_shear_deform", "static_stability", "horz_wind_speed", "vert_wind_shear"])
 
-    xr_dataset.close()
-
     return all_df
 
 
-def analyze(dataFrame, column, value):
-    result_df = dataFrame[dataFrame[column] == value]  # get rows where column has a certain value
+def analyze(pd_df, column, value):
+    result_df = pd_df[pd_df[column] == value]  # get rows where column has a certain value
 
     mean = result_df.mean()  # calculate mean for other columns
     stddev = result_df.std()  # calculate standard deviation for other columns