import re
import datetime
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
import matplotlib.image as mpimg
import h5py
from util.util import get_grid_values_all
from util.gfs_reader import *
from util.geos_nav import GEOSNavigation
from aeolus.datasource import GFSfiles
import cartopy.crs as ccrs
from metpy.calc import shearing_deformation, static_stability, wind_speed, first_derivative
from metpy.units import units
gfs_files = GFSfiles('/Users/tomrink/data/contrail/gfs/')
# GEOSNavigation needs to be updated to support GOES-18
# geos_nav = GEOSNavigation()
def load_image(image_path):
# Extract date time string from image path
datetime_regex = '_\\d{8}_\\d{6}'
datetime_string = re.search(datetime_regex, image_path)
if datetime_string:
datetime_string = datetime_string.group()
dto = datetime.datetime.strptime(datetime_string, '_%Y%m%d_%H%M%S').replace(tzinfo=timezone.utc)
ts = dto.timestamp()
img = mpimg.imread(image_path)
return img, ts
def get_contrail_mask_image(image, thresh=0.157):
image = np.where(image > thresh,1, 0)
return image
def extract(mask_image, image_ts, clavrx_path):
gfs_file, _, _ = gfs_files.get_file(image_ts)
xr_dataset = xr.open_dataset(gfs_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()
clvrx_lats = get_grid_values_all(clvrx_h5f, 'latitude').flatten()
contrail_idxs = (mask_image == 1).flatten()
# Assuming GOES FD for now -------------------
# elems, lines = np.meshgrid(np.arange(5424), np.arange(5424))
# lines, elems = lines.flatten(), elems.flatten()
# See note above regarding support for GOES-18
# contrail_lines, contrail_elems = lines[contrail_idxs], elems[contrail_idxs]
# contrail_lons, contrail_lats = geos_nav.lc_to_earth(contrail_elems, contrail_lines)
contrail_press = cloud_top_press[contrail_idxs]
contrail_lons, contrail_lats = clvrx_lons[contrail_idxs], clvrx_lats[contrail_idxs]
keep = np.invert(np.isnan(contrail_press))
contrail_press = contrail_press[keep]
contrail_lons = contrail_lons[keep]
contrail_lats = contrail_lats[keep]
# Indexes of contrail_press for individual bins
bins = np.arange(100, 1000, 50)
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)]
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
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]:
press = contrail_press[c_idx]
lat = contrail_lats[c_idx]
lon = contrail_lons[c_idx]
if lon < 0: # Match GFS convention
lon += 360.0
horz_shear_value = horz_shear_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest')
static_value = static_3d.interp(Pressure=press, Longitude=lon, Latitude=lat, method='nearest')
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')
# tmp = horz_shear_3d.sel(Pressure=press, method='nearest')
# tmp = tmp.sel(Longitude=lon, Latitude=lat, method='nearest')
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():
df = pd.DataFrame(v, columns=["pressure", "lat", "lon", "horz_wind_shear", "static_stability", "horz_wind_speed", "vert_wind_shear"])
voxel_dict_df[k] = df
xr_dataset.close()