diff --git a/modules/contrail/util.py b/modules/contrail/util.py index bb1e7d43256d216c8700c2f0d5cf5ff2081fd34f..4ab55fb9aec7eb4869a38bb9a2977c253ce97c76 100644 --- a/modules/contrail/util.py +++ b/modules/contrail/util.py @@ -4,6 +4,7 @@ from datetime import timezone import numpy as np import xarray as xr +import pandas as pd import rasterio from PIL import Image import matplotlib.pyplot as plt @@ -108,8 +109,10 @@ def extract(mask_image, image_ts, clavrx_path): 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 units via metpy's pint support vert_shear_3d = first_derivative(horz_wind_spd_3d, axis=0, x=temp_3d.coords['Pressure']) - print(horz_wind_spd_3d.shape, vert_shear_3d.shape) vert_shear_3d = volume_np_to_xr(vert_shear_3d, ['Pressure', 'Latitude', 'Longitude'], lon_range=[lon_range[0], lon_range[1]], lat_range=[lat_range[0], lat_range[1]]) vert_shear_3d = vert_shear_3d * units.meter / (units.second * units.hPa) @@ -126,7 +129,12 @@ def extract(mask_image, image_ts, clavrx_path): horz_wind_spd_value = horz_wind_spd_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest') vert_shear_value = vert_shear_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest') - voxel_dict[key].append((horz_shear_value, static_value, horz_wind_spd_value, vert_shear_value)) + voxel_dict[key].append((press, lat, lon, horz_shear_value, static_value, horz_wind_spd_value, vert_shear_value)) + + # Create pandas DataFrame for each list of tuples in voxel_dict + voxel_dict_df = {} + for k, v in voxel_dict.items(): + voxel_dict_df[k] = pd.DataFrame(v, columns=["press", "lat", "lon", "horz_shear_value", "static_value", "horz_wind_spd_value", "vert_shear_value"]) xr_dataset.close()