diff --git a/modules/contrail/util.py b/modules/contrail/util.py index 20233e5aaf9567dbe3ff6e8a444f90b59bf04f29..bb1e7d43256d216c8700c2f0d5cf5ff2081fd34f 100644 --- a/modules/contrail/util.py +++ b/modules/contrail/util.py @@ -108,7 +108,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) - vert_shear_3d = first_derivative(horz_wind_spd_3d, axis=0, x=temp_3d.coords['Pressure'] * units.hPa) + 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) voxel_dict = {key: [] for key in bins_dict.keys()} for key in bins_dict.keys(): @@ -118,6 +121,13 @@ def extract(mask_image, image_ts, clavrx_path): lat = contrail_lats[c_idx] press = contrail_press[c_idx] + 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') + + voxel_dict[key].append((horz_shear_value, static_value, horz_wind_spd_value, vert_shear_value)) + xr_dataset.close() diff --git a/modules/util/gfs_reader.py b/modules/util/gfs_reader.py index 6878c466759be361a4a1356e8a445d94448dda69..8a65ad3550e43a935d34012add1a9b4bc676d1a0 100644 --- a/modules/util/gfs_reader.py +++ b/modules/util/gfs_reader.py @@ -480,3 +480,39 @@ def get_volume(xr_dataset, fld_name, unit_str, press_range=None, lon_range=None, attrs={"description": fld_name, "units": unit_str}) return xra + + +def volume_np_to_xr(nda, dims, press_range=None, lon_range=None, lat_range=None): + + x_lo, x_hi = 0, NX + y_lo, y_hi = 0, NY + z_lo, z_hi = 0, NZ + + if lon_range is not None: + lon_lo = lon_range[0] + lon_hi = lon_range[1] + + if lon360: + if lon_lo < 0: + lon_lo += 360 + if lon_hi < 0: + lon_hi += 360 + + x_lo = value_to_index(lon_coords, lon_lo) + x_hi = value_to_index(lon_coords, lon_hi) + + if lat_range is not None: + lat_lo = lat_range[0] + lat_hi = lat_range[1] + + y_lo = value_to_index(lat_coords, lat_lo) + y_hi = value_to_index(lat_coords, lat_hi) + + if press_range is not None: + z_lo = value_to_index(plevs, press_range[0]) + z_hi = value_to_index(plevs, press_range[1]) + + xra = xr.DataArray(nda, dims=dims, + coords={"Latitude": lat_coords[y_lo:y_hi], "Longitude": lon_coords[x_lo:x_hi], "Pressure": plevs[z_lo:z_hi]}) + + return xra