From e67a2b0b29535fb47d73edcd741a1d6b53d4d9d4 Mon Sep 17 00:00:00 2001 From: Rob Thompson Date: Tue, 11 Aug 2026 16:25:31 +0000 Subject: [PATCH 1/4] attempt to add upper_bc == 2 cloud top boundary condition --- .../_cost_functions_numpy-checkpoint.py | 884 +++++++++ .../cost_functions-checkpoint.py | 1026 +++++++++++ pydda/cost_functions/_cost_functions_numpy.py | 15 +- pydda/cost_functions/cost_functions.py | 2 + .../wind_retrieve-checkpoint.py | 1596 +++++++++++++++++ pydda/retrieval/wind_retrieve.py | 9 +- 6 files changed, 3524 insertions(+), 8 deletions(-) create mode 100644 pydda/cost_functions/.ipynb_checkpoints/_cost_functions_numpy-checkpoint.py create mode 100644 pydda/cost_functions/.ipynb_checkpoints/cost_functions-checkpoint.py create mode 100644 pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py diff --git a/pydda/cost_functions/.ipynb_checkpoints/_cost_functions_numpy-checkpoint.py b/pydda/cost_functions/.ipynb_checkpoints/_cost_functions_numpy-checkpoint.py new file mode 100644 index 00000000..390e1cdf --- /dev/null +++ b/pydda/cost_functions/.ipynb_checkpoints/_cost_functions_numpy-checkpoint.py @@ -0,0 +1,884 @@ +import numpy as np +import scipy +import pyart + +from scipy.ndimage import _nd_image + +laplace_filter = np.asarray([1, -2, 1], dtype=np.float64) + + +def calculate_radial_vel_cost_function( + vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0, parallel=False +): + """ + Calculates the cost function due to difference of the wind field from + radar radial velocities. For more information on this cost function, see + Potvin et al. (2012) and Shapiro et al. (2009). + All arrays in the given lists must have the same dimensions and represent + the same spatial coordinates. + Parameters + ---------- + vrs: List of float arrays + List of radial velocities from each radar + els: List of float arrays + List of elevations from each radar + azs: List of float arrays + List of azimuths from each radar + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + wts: List of float arrays + Float array containing fall speed from radar. + rmsVr: float + The sum of squares of velocity/num_points. Use for normalization + of data weighting coefficient + weights: n_radars x_bins x y_bins float array + Data weights for each pair of radars + coeff: float + Constant for cost function + Returns + ------- + J_o: float + Observational cost function + References + ----------- + Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity + Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and + Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49, + https://doi.org/10.1175/JTECH-D-11-00019.1 + Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity + Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic + Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1 + """ + + lambda_o = coeff / (rmsVr * rmsVr) + if parallel: + vrs_arr = np.stack(vrs) + els_arr = np.stack(els) + azs_arr = np.stack(azs) + wts_arr = np.stack(wts) + v_ar = ( + np.cos(els_arr) * np.sin(azs_arr) * u[np.newaxis] + + np.cos(els_arr) * np.cos(azs_arr) * v[np.newaxis] + + np.sin(els_arr) * (w[np.newaxis] - np.abs(wts_arr)) + ) + return lambda_o * np.sum(np.square(vrs_arr - v_ar) * weights) + + J_o = 0 + for i in range(len(vrs)): + v_ar = ( + np.cos(els[i]) * np.sin(azs[i]) * u + + np.cos(els[i]) * np.cos(azs[i]) * v + + np.sin(els[i]) * (w - np.abs(wts[i])) + ) + J_o += lambda_o * np.sum(np.square(vrs[i] - v_ar) * weights[i]) + + return J_o + + +def calculate_grad_radial_vel( + vrs, + els, + azs, + u, + v, + w, + wts, + weights, + rmsVr, + coeff=1.0, + upper_bc=True, + parallel=False, +): + """ + Calculates the gradient of the cost function due to difference of wind + field from radar radial velocities. + All arrays in the given lists must have the same dimensions and represent + the same spatial coordinates. + Parameters + ---------- + vrs: List of float arrays + List of radial velocities from each radar + els: List of float arrays + List of elevations from each radar + azs: List of azimuths + List of azimuths from each radar + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + coeff: float + Constant for cost function + vel_name: str + Background velocity field name + weights: n_radars x_bins x y_bins float array + Data weights for each pair of radars + Returns + ------- + y: 1-D float array + Gradient vector of observational cost function. + + More information + ---------------- + The gradient is calculated by taking the functional derivative of the + cost function. For more information on functional derivatives, see the + Euler-Lagrange Equation: + https://en.wikipedia.org/wiki/Euler%E2%80%93Lagrange_equation + """ + + # Use zero for all masked values since we don't want to add them into + # the cost function + + lambda_o = coeff / (rmsVr * rmsVr) + + if parallel: + vrs_arr = np.stack(vrs) + els_arr = np.stack(els) + azs_arr = np.stack(azs) + wts_arr = np.stack(wts) + v_ar = ( + np.cos(els_arr) * np.sin(azs_arr) * u[np.newaxis] + + np.cos(els_arr) * np.cos(azs_arr) * v[np.newaxis] + + np.sin(els_arr) * (w[np.newaxis] - np.abs(wts_arr)) + ) + residual = 2 * (v_ar - vrs_arr) * lambda_o + p_x1 = np.sum(residual * np.cos(els_arr) * np.sin(azs_arr) * weights, axis=0) + p_y1 = np.sum(residual * np.cos(els_arr) * np.cos(azs_arr) * weights, axis=0) + p_z1 = np.sum(residual * np.sin(els_arr) * weights, axis=0) + else: + p_x1 = np.zeros(vrs[0].shape) + p_y1 = np.zeros(vrs[0].shape) + p_z1 = np.zeros(vrs[0].shape) + + for i in range(len(vrs)): + v_ar = ( + np.cos(els[i]) * np.sin(azs[i]) * u + + np.cos(els[i]) * np.cos(azs[i]) * v + + np.sin(els[i]) * (w - np.abs(wts[i])) + ) + + x_grad = ( + 2 * (v_ar - vrs[i]) * np.cos(els[i]) * np.sin(azs[i]) * weights[i] + ) * lambda_o + y_grad = ( + 2 * (v_ar - vrs[i]) * np.cos(els[i]) * np.cos(azs[i]) * weights[i] + ) * lambda_o + z_grad = (2 * (v_ar - vrs[i]) * np.sin(els[i]) * weights[i]) * lambda_o + + p_x1 += x_grad + p_y1 += y_grad + p_z1 += z_grad + + # Impermeability condition + p_z1[0, :, :] = 0 + if upper_bc is True: + p_z1[-1, :, :] = 0 + y = np.stack((p_x1, p_y1, p_z1), axis=0) + return y.flatten() + + +def calculate_smoothness_cost(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5): + """ + Calculates the smoothness cost function by taking the Laplacian of the + wind field. + All arrays in the given lists must have the same dimensions and represent + the same spatial coordinates. + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + Cx: float + Constant controlling smoothness in x-direction + Cy: float + Constant controlling smoothness in y-direction + Cz: float + Constant controlling smoothness in z-direction + Returns + ------- + Js: float + value of smoothness cost function + """ + dudx = np.gradient(u, dx, axis=2) + dudy = np.gradient(u, dy, axis=1) + dudz = np.gradient(u, dz, axis=0) + dvdx = np.gradient(v, dx, axis=2) + dvdy = np.gradient(v, dy, axis=1) + dvdz = np.gradient(v, dz, axis=0) + dwdx = np.gradient(w, dx, axis=2) + dwdy = np.gradient(w, dy, axis=1) + dwdz = np.gradient(w, dz, axis=0) + + x_term = ( + Cx + * ( + np.gradient(dudx, dx, axis=2) + + np.gradient(dvdx, dx, axis=2) + + np.gradient(dwdx, dx, axis=2) + ) + ** 2 + ) + y_term = ( + Cy + * ( + np.gradient(dudy, dy, axis=1) + + np.gradient(dvdy, dy, axis=1) + + np.gradient(dwdy, dy, axis=1) + ) + ** 2 + ) + z_term = ( + Cz + * ( + np.gradient(dudz, dz, axis=0) + + np.gradient(dvdz, dz, axis=0) + + np.gradient(dwdz, dz, axis=0) + ) + ** 2 + ) + return np.sum(np.nan_to_num(x_term + y_term + z_term)) + + +def calculate_smoothness_gradient( + u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5, upper_bc=True +): + """ + Calculates the gradient of the smoothness cost function + by taking the Laplacian of the Laplacian of the wind field. + All arrays in the given lists must have the same dimensions and represent + the same spatial coordinates. + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + Cx: float + Constant controlling smoothness in x-direction + Cy: float + Constant controlling smoothness in y-direction + Cz: float + Constant controlling smoothness in z-direction + Returns + ------- + y: float array + value of gradient of smoothness cost function + """ + du = np.zeros(w.shape) + dv = np.zeros(w.shape) + dw = np.zeros(w.shape) + grad_u = np.zeros(w.shape) + grad_v = np.zeros(w.shape) + grad_w = np.zeros(w.shape) + scipy.ndimage.laplace(u, du, mode="wrap") + scipy.ndimage.laplace(v, dv, mode="wrap") + scipy.ndimage.laplace(w, dw, mode="wrap") + du = du / dx + dv = dv / dy + dw = dw / dz + scipy.ndimage.laplace(du, grad_u, mode="wrap") + scipy.ndimage.laplace(dv, grad_v, mode="wrap") + scipy.ndimage.laplace(dw, grad_w, mode="wrap") + grad_u = grad_u / dx + grad_v = grad_v / dy + grad_w = grad_w / dz + + # Impermeability condition + grad_w[0, :, :] = 0 + if upper_bc is True: + grad_w[-1, :, :] = 0 + + y = np.stack([grad_u, grad_v, grad_w], axis=0) + + return y.flatten() + + +def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2): + """ + Calculates the cost function related to point observations. A mean square error cost + function term is applied to points that are within the sphere of influence + whose radius is determined by *roi*. + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + x: Float array + X coordinates of grid centers + y: Float array + Y coordinates of grid centers + z: Float array + Z coordinated of grid centers + point_list: list of dicts + List of point constraints. + Each member is a dict with keys of "u", "v", to correspond + to each component of the wind field and "x", "y", "z" + to correspond to the location of the point observation. + In addition, "site_id" gives the METAR code (or name) to the station. + Cp: float + The weighting coefficient of the point cost function. + roi: float + Radius of influence of observations + Returns + ------- + J: float + The cost function related to the difference between wind field and points. + """ + J = 0.0 + for the_point in point_list: + # Instead of worrying about whole domain, just find points in radius of influence + # Since we know that the weight will be zero outside the sphere of influence anyways + + dist = np.sqrt( + (x - the_point["x"]) ** 2 + + (y - the_point["y"]) ** 2 + + (z - the_point["z"]) ** 2 + ) + dist = np.maximum(dist, 1.0) + weight = 1 / dist**2 + weight = weight / np.max(weight) + + J += np.sum(weight * ((u - the_point["u"]) ** 2 + (v - the_point["v"]) ** 2)) + + return J * Cp + + +def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): + """ + Calculates the gradient of the cost function related to point observations. + A mean square error cost function term is applied to points that are within the sphere of influence + whose radius is determined by *roi*. + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + x: Float array + X coordinates of grid centers + y: Float array + Y coordinates of grid centers + z: Float array + Z coordinated of grid centers + point_list: list of dicts + List of point constraints. Each member is a dict with keys of "u", "v", + to correspond to each component of the wind field and "x", "y", "z" + to correspond to the location of the point observation. + In addition, "site_id" gives the METAR code (or name) to the station. + Cp: float + The weighting coefficient of the point cost function. + roi: float + Radius of influence of observations + Returns + ------- + gradJ: float array + The gradient of the cost function related to the difference between wind field and points. + """ + + gradJ_u = np.zeros_like(u) + gradJ_v = np.zeros_like(v) + gradJ_w = np.zeros_like(u) + + for the_point in point_list: + dist = np.sqrt( + (x - the_point["x"]) ** 2 + + (y - the_point["y"]) ** 2 + + (z - the_point["z"]) ** 2 + ) + dist = np.maximum(dist, 1.0) + weight = 1 / dist**2 + weight = weight / np.max(weight) + gradJ_u += 2 * weight * (u - the_point["u"]) + gradJ_v += 2 * weight * (v - the_point["v"]) + + gradJ = np.stack([gradJ_u, gradJ_v, gradJ_w], axis=0).flatten() + return gradJ * Cp + + +def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): + """ + Calculates the mass continuity cost function by taking the divergence + of the wind field. + All arrays in the given lists must have the same dimensions and represent + the same spatial coordinates. + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + dx: float + Grid spacing in x direction. + dy: float + Grid spacing in y direction. + dz: float + Grid spacing in z direction. + z: Float array (1D) + 1D Float array with heights of grid + coeff: float + Constant controlling contribution of mass continuity to cost function + anel: int + = 1 use anelastic approximation, 0=don't + Returns + ------- + J: float + value of mass continuity cost function + """ + dudx = np.gradient(u, dx, axis=2) + dvdy = np.gradient(v, dy, axis=1) + dwdz = np.gradient(w, dz, axis=0) + + if anel == 1: + rho = np.exp(-z / 10000.0) + drho_dz = np.gradient(rho, dz, axis=0) + anel_term = w / rho * drho_dz + else: + anel_term = np.zeros(w.shape) + div = dudx + dvdy + dwdz + anel_term + + return coeff * np.sum(np.square(div)) / 2.0 + + +def calculate_mass_continuity_gradient( + u, v, w, z, dx, dy, dz, vrs=0, coeff=1500.0, anel=1, upper_bc=True, above=2.0 +): + """ + Calculates the gradient of mass continuity cost function. This is done by + taking the negative gradient of the divergence of the wind field. + All grids must have the same grid specification. + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + z: Float array (1D) + 1D Float array with heights of grid + dx: float + Grid spacing in x direction. + dy: float + Grid spacing in y direction. + dz: float + Grid spacing in z direction. + coeff: float + Constant controlling contribution of mass continuity to cost function + anel: int + = 1 use anelastic approximation, 0=don't + Returns + ------- + y: float array + value of gradient of mass continuity cost function + """ + dudx = np.gradient(u, dx, axis=2) + dvdy = np.gradient(v, dy, axis=1) + dwdz = np.gradient(w, dz, axis=0) + if anel == 1: + rho = np.exp(-z / 10000.0) + drho_dz = np.gradient(rho, dz, axis=0) + anel_term = w / rho * drho_dz + else: + anel_term = 0 + + div = dudx + dvdy + dwdz + anel_term + + grad_u = -np.gradient(div, dx, axis=2) * coeff + grad_v = -np.gradient(div, dy, axis=1) * coeff + grad_w = -np.gradient(div, dz, axis=0) * coeff + + # Impermeability conditions + grad_w[0, :, :] = 0 # surface is impermeable + if upper_bc == 1:#is True: # impermeable at the grid top + grad_w[-1, :, :] = 0 + if upper_bc == 2: # impermeable at cloud top + N=np.sum(np.array(vrs)>-1000,axis=0) + z_mask = (z > above * 1000) + n_mask = (N == 0) + grad_w[z_mask & n_mask] = 0 + y = np.stack([grad_u, grad_v, grad_w], axis=0) + return y.flatten() + + +def calculate_fall_speed(grid, refl_field=None, frz=4500.0): + """ + Estimates fall speed based on reflectivity. + Uses methodology of Mike Biggerstaff and Dan Betten + Parameters + ---------- + Grid: Py-ART Grid + Py-ART Grid containing reflectivity to calculate fall speed from + refl_field: str + String containing name of reflectivity field. None will automatically + determine the name. + frz: float + Height of freezing level in m + Returns + ------- + 3D float array: + Float array of terminal velocities + """ + # Parse names of velocity field + if refl_field is None: + refl_field = pyart.config.get_field_name("reflectivity") + + refl = grid[refl_field].values + grid_z = grid["point_z"].values + A = np.zeros(refl.shape) + B = np.zeros(refl.shape) + rho = np.exp(-grid_z / 10000.0) + A[np.logical_and(grid_z < frz, refl < 55)] = -2.6 + B[np.logical_and(grid_z < frz, refl < 55)] = 0.0107 + A[np.logical_and(grid_z < frz, np.logical_and(refl >= 55, refl < 60))] = -2.5 + B[np.logical_and(grid_z < frz, np.logical_and(refl >= 55, refl < 60))] = 0.013 + A[np.logical_and(grid_z < frz, refl > 60)] = -3.95 + B[np.logical_and(grid_z < frz, refl > 60)] = 0.0148 + A[np.logical_and(grid_z >= frz, refl < 33)] = -0.817 + B[np.logical_and(grid_z >= frz, refl < 33)] = 0.0063 + A[np.logical_and(grid_z >= frz, np.logical_and(refl >= 33, refl < 49))] = -2.5 + B[np.logical_and(grid_z >= frz, np.logical_and(refl >= 33, refl < 49))] = 0.013 + A[np.logical_and(grid_z >= frz, refl > 49)] = -3.95 + B[np.logical_and(grid_z >= frz, refl > 49)] = 0.0148 + + fallspeed = A * np.power(10, refl * B) * np.power(1.2 / rho, 0.4) + del A, B, rho + return fallspeed + + +def calculate_background_cost(u, v, w, weights, u_back, v_back, Cb=0.01): + """ + Calculates the background cost function. The background cost function is + simply the sum of the squared differences between the wind field and the + background wind field multiplied by the weighting coefficient. + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + weights: Float array + Weights for each point to consider into cost function + u_back: 1D float array + Zonal winds vs height from sounding + w_back: 1D float array + Meridional winds vs height from sounding + Cb: float + Weight of background constraint to total cost function + Returns + ------- + cost: float + value of background cost function + """ + the_shape = u.shape + cost = 0 + for i in range(the_shape[0]): + cost += Cb * np.sum( + np.square(u[i] - u_back[i]) * (weights[i]) + + np.square(v[i] - v_back[i]) * (weights[i]) + ) + return cost + + +def calculate_background_gradient(u, v, w, weights, u_back, v_back, Cb=0.01): + """ + Calculates the gradient of the background cost function. For each u, v + this is given as 2*coefficent*(analysis wind - background wind). + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + weights: Float array + Weights for each point to consider into cost function + u_back: 1D float array + Zonal winds vs height from sounding + w_back: 1D float array + Meridional winds vs height from sounding + Cb: float + Weight of background constraint to total cost function + Returns + ------- + y: float array + value of gradient of background cost function + """ + the_shape = u.shape + u_grad = np.zeros(the_shape) + v_grad = np.zeros(the_shape) + w_grad = np.zeros(the_shape) + + for i in range(the_shape[0]): + u_grad[i] = Cb * 2 * (u[i] - u_back[i]) * (weights[i]) + v_grad[i] = Cb * 2 * (v[i] - v_back[i]) * (weights[i]) + + y = np.stack([u_grad, v_grad, w_grad], axis=0) + return y.flatten() + + +def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5): + """ + Calculates the cost function due to deviance from vertical vorticity + equation. For more information of the vertical vorticity cost function, + see Potvin et al. (2012) and Shapiro et al. (2009). + Parameters + ---------- + u: 3D array + Float array with u component of wind field + v: 3D array + Float array with v component of wind field + w: 3D array + Float array with w component of wind field + dx: float array + Spacing in x grid + dy: float array + Spacing in y grid + dz: float array + Spacing in z grid + coeff: float + Weighting coefficient + Ut: float + U component of storm motion + Vt: float + V component of storm motion + Returns + ------- + Jv: float + Value of vertical vorticity cost function. + References + ---------- + Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity + Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and + Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49, + https://doi.org/10.1175/JTECH-D-11-00019.1 + Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity + Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic + Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1 + """ + dvdz = np.gradient(v, dz, axis=0) + dudz = np.gradient(u, dz, axis=0) + dvdx = np.gradient(v, dx, axis=2) + dwdy = np.gradient(w, dy, axis=1) + dwdx = np.gradient(w, dx, axis=2) + dudx = np.gradient(u, dx, axis=2) + dvdy = np.gradient(v, dy, axis=2) + dudy = np.gradient(u, dy, axis=1) + zeta = dvdx - dudy + dzeta_dx = np.gradient(zeta, dx, axis=2) + dzeta_dy = np.gradient(zeta, dy, axis=1) + dzeta_dz = np.gradient(zeta, dz, axis=0) + jv_array = ( + (u - Ut) * dzeta_dx + + (v - Vt) * dzeta_dy + + w * dzeta_dz + + (dvdz * dwdx - dudz * dwdy) + + zeta * (dudx + dvdy) + ) + return np.sum(coeff * jv_array**2) + + +def calculate_vertical_vorticity_gradient( + u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5, upper_bc=True +): + """ + Calculates the gradient of the cost function due to deviance from vertical + vorticity equation. This is done by taking the functional derivative of + the vertical vorticity cost function. + Parameters + ---------- + u: 3D array + Float array with u component of wind field + v: 3D array + Float array with v component of wind field + w: 3D array + Float array with w component of wind field + dx: float array + Spacing in x grid + dy: float array + Spacing in y grid + dz: float array + Spacing in z grid + Ut: float + U component of storm motion + Vt: float + V component of storm motion + coeff: float + Weighting coefficient + Returns + ------- + Jv: 1D float array + Value of the gradient of the vertical vorticity cost function. + References + ---------- + Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity + Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and + Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49, + https://doi.org/10.1175/JTECH-D-11-00019.1 + Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity + Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic + Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1 + """ + + # First derivatives + dvdz = np.gradient(v, dz, axis=0) + dwdy = np.gradient(w, dy, axis=1) + dudx = np.gradient(u, dx, axis=2) + dvdy = np.gradient(v, dy, axis=1) + dvdx = np.gradient(v, dx, axis=2) + dwdx = np.gradient(w, dx, axis=2) + dudz = np.gradient(u, dz, axis=0) + dudy = np.gradient(u, dy, axis=1) + + zeta = dvdx - dudy + dzeta_dx = np.gradient(zeta, dx, axis=2) + dzeta_dy = np.gradient(zeta, dy, axis=1) + dzeta_dz = np.gradient(zeta, dz, axis=0) + + # Second deriviatives + dwdydz = np.gradient(dwdy, dz, axis=0) + dwdxdz = np.gradient(dwdx, dz, axis=0) + dudzdy = np.gradient(dudz, dy, axis=1) + dvdxdy = np.gradient(dvdx, dy, axis=1) + dudx2 = np.gradient(dudx, dx, axis=2) + dudxdy = np.gradient(dudx, dy, axis=1) + dudxdz = np.gradient(dudx, dz, axis=0) + dudy2 = np.gradient(dudx, dy, axis=1) + + dzeta_dt = ( + (u - Ut) * dzeta_dx + + (v - Vt) * dzeta_dy + + w * dzeta_dz + + (dvdz * dwdx - dudz * dwdy) + + zeta * (dudx + dvdy) + ) + + # Now we intialize our gradient value + u_grad = np.zeros(u.shape) + v_grad = np.zeros(v.shape) + w_grad = np.zeros(w.shape) + + # Vorticity Advection + u_grad += dzeta_dx + (Ut - u) * dudxdy + (Vt - v) * dudxdy + v_grad += dzeta_dy + (Vt - v) * dvdxdy + (Ut - u) * dvdxdy + w_grad += dzeta_dz + + # Tilting term + u_grad += dwdydz + v_grad += dwdxdz + w_grad += dudzdy - dudxdz + + # Stretching term + u_grad += -dudxdy + dudy2 - dzeta_dx + u_grad += -dudx2 + dudxdy - dzeta_dy + + # Multiply by 2*dzeta_dt according to chain rule + u_grad = u_grad * 2 * dzeta_dt * coeff + v_grad = v_grad * 2 * dzeta_dt * coeff + w_grad = w_grad * 2 * dzeta_dt * coeff + + # Impermeability condition + w_grad[0, :, :] = 0 + if upper_bc is True: + w_grad[-1, :, :] = 0 + y = np.stack([u_grad, v_grad, w_grad], axis=0) + return y.flatten() + + +def calculate_model_cost(u, v, w, weights, u_model, v_model, w_model, coeff=1.0): + """ + Calculates the cost function for the model constraint. + This is calculated simply as the sum of squares of the differences + between the model wind field and the analysis wind field. Vertical + velocities are not factored into this cost function as there is typically + a high amount of uncertainty in model derived vertical velocities. + Parameters + ---------- + u: 3D array + Float array with u component of wind field + v: 3D array + Float array with v component of wind field + w: 3D array + Float array with w component of wind field + weights: list of 3D arrays + Float array showing how much each point from model weighs into + constraint. + u_model: list of 3D arrays + Float array with u component of wind field from model + v_model: list of 3D arrays + Float array with v component of wind field from model + w_model: list of 3D arrays + Float array with w component of wind field from model + coeff: float + Weighting coefficient + Returns + ------- + Jv: float + Value of model cost function + """ + + cost = 0 + for i in range(len(u_model)): + cost += coeff * np.sum( + np.square(u - u_model[i]) * weights[i] + + np.square(v - v_model[i]) * weights[i] + ) + return cost + + +def calculate_model_gradient(u, v, w, weights, u_model, v_model, w_model, coeff=1.0): + """ + Calculates the cost function for the model constraint. + This is calculated simply as twice the differences + between the model wind field and the analysis wind field for each u, v. + Vertical velocities are not factored into this cost function as there is + typically a high amount of uncertainty in model derived vertical + velocities. Therefore, the gradient for all of the w's will be 0. + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + weights: list of 3D float arrays + Weights for each point to consider into cost function + u_model: list of 3D float arrays + Zonal wind field from model + v_model: list of 3D float arrays + Meridional wind field from model + w_model: list of 3D float arrays + Vertical wind field from model + coeff: float + Weight of background constraint to total cost function + Returns + ------- + y: float array + value of gradient of background cost function + """ + the_shape = u.shape + u_grad = np.zeros(the_shape) + v_grad = np.zeros(the_shape) + w_grad = np.zeros(the_shape) + for i in range(len(u_model)): + u_grad += coeff * 2 * (u - u_model[i]) * weights[i] + v_grad += coeff * 2 * (v - v_model[i]) * weights[i] + + y = np.stack([u_grad, v_grad, w_grad], axis=0) + return y.flatten() diff --git a/pydda/cost_functions/.ipynb_checkpoints/cost_functions-checkpoint.py b/pydda/cost_functions/.ipynb_checkpoints/cost_functions-checkpoint.py new file mode 100644 index 00000000..60c921ea --- /dev/null +++ b/pydda/cost_functions/.ipynb_checkpoints/cost_functions-checkpoint.py @@ -0,0 +1,1026 @@ +import numpy as np +from concurrent.futures import ThreadPoolExecutor + +# Adding jax import statements +try: + import tensorflow as tf + + TENSORFLOW_AVAILABLE = True +except ImportError: + TENSORFLOW_AVAILABLE = False + +try: + import jax.numpy as jnp + + JAX_AVAILABLE = True +except ImportError: + JAX_AVAILABLE = False + +import pyart + +# Added to incorpeate JAX within the cost functions +from . import _cost_functions_jax +from . import _cost_functions_numpy +from . import _cost_functions_tensorflow + + +def J_function(winds, parameters): + """ + Calculates the total cost function. This typically does not need to be + called directly as get_dd_wind_field is a wrapper around this function and + :py:func:`pydda.cost_functions.grad_J`. + In order to add more terms to the cost function, modify this + function and :py:func:`pydda.cost_functions.grad_J`. + + Parameters + ---------- + winds: 1-D float array + The wind field, flattened to 1-D for f_min. The total size of the + array will be a 1D array of 3*nx*ny*nz elements. + parameters: DDParameters + The parameters for the cost function evaluation as specified by the + :py:func:`pydda.retrieval.DDParameters` class. + + Returns + ------- + J: float + The value of the cost function + """ + if parameters.engine == "tensorflow": + if not TENSORFLOW_AVAILABLE: + raise ImportError( + "Tensorflow 2.5 or greater is needed in order to use TensorFlow-based PyDDA!" + ) + winds = tf.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + winds = tf.math.maximum(winds, tf.constant([-100.0])) + winds = tf.math.minimum(winds, tf.constant([100.0])) + # Had to change to float because Jax returns device array (use np.float_()) + Jvel = _cost_functions_tensorflow.calculate_radial_vel_cost_function( + parameters.vrs, + parameters.azs, + parameters.els, + winds[0], + winds[1], + winds[2], + parameters.wts, + rmsVr=parameters.rmsVr, + weights=parameters.weights, + coeff=parameters.Co, + ) + # print("apples Jvel", Jvel) + + if parameters.Cm > 0: + # Had to change to float because Jax returns device array (use np.float_()) + Jmass = _cost_functions_tensorflow.calculate_mass_continuity( + winds[0], + winds[1], + winds[2], + parameters.z, + parameters.dx, + parameters.dy, + parameters.dz, + coeff=parameters.Cm, + ) + else: + Jmass = 0 + + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: + Jsmooth = _cost_functions_tensorflow.calculate_smoothness_cost( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + Cx=parameters.Cx, + Cy=parameters.Cy, + Cz=parameters.Cz, + ) + else: + Jsmooth = 0 + + if parameters.Cb > 0: + Jbackground = _cost_functions_tensorflow.calculate_background_cost( + winds[0], + winds[1], + parameters.bg_weights, + parameters.u_back, + parameters.v_back, + parameters.Cb, + ) + else: + Jbackground = 0 + + if parameters.Cv > 0: + # Had to change to float because Jax returns device array (use np.float_()) + Jvorticity = _cost_functions_tensorflow.calculate_vertical_vorticity_cost( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + parameters.Ut, + parameters.Vt, + coeff=parameters.Cv, + ) + else: + Jvorticity = 0 + + if parameters.Cmod > 0: + Jmod = _cost_functions_tensorflow.calculate_model_cost( + winds[0], + winds[1], + winds[2], + parameters.model_weights, + parameters.u_model, + parameters.v_model, + parameters.w_model, + coeff=parameters.Cmod, + ) + else: + Jmod = 0 + + if parameters.Cpoint > 0: + Jpoint = _cost_functions_tensorflow.calculate_point_cost( + winds[0], + winds[1], + parameters.x, + parameters.y, + parameters.z, + parameters.point_list, + Cp=parameters.Cpoint, + roi=parameters.roi, + ) + else: + Jpoint = 0 + elif parameters.engine == "scipy": + winds = np.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + # Had to change to float because Jax returns device array (use np.float_()) + Jvel = _cost_functions_numpy.calculate_radial_vel_cost_function( + parameters.vrs, + parameters.azs, + parameters.els, + winds[0], + winds[1], + winds[2], + parameters.wts, + rmsVr=parameters.rmsVr, + weights=parameters.weights, + coeff=parameters.Co, + parallel=parameters.parallel, + ) + # print("apples Jvel", Jvel) + + if parameters.Cm > 0: + # Had to change to float because Jax returns device array (use np.float_()) + Jmass = _cost_functions_numpy.calculate_mass_continuity( + winds[0], + winds[1], + winds[2], + parameters.z, + parameters.dx, + parameters.dy, + parameters.dz, + coeff=parameters.Cm, + ) + else: + Jmass = 0 + + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: + Jsmooth = _cost_functions_numpy.calculate_smoothness_cost( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + Cx=parameters.Cx, + Cy=parameters.Cy, + Cz=parameters.Cz, + ) + else: + Jsmooth = 0 + + if parameters.Cb > 0: + Jbackground = _cost_functions_numpy.calculate_background_cost( + winds[0], + winds[1], + winds[2], + parameters.bg_weights, + parameters.u_back, + parameters.v_back, + parameters.Cb, + ) + else: + Jbackground = 0 + + if parameters.Cv > 0: + # Had to change to float because Jax returns device array (use np.float_()) + Jvorticity = _cost_functions_numpy.calculate_vertical_vorticity_cost( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + parameters.Ut, + parameters.Vt, + coeff=parameters.Cv, + ) + else: + Jvorticity = 0 + + if parameters.Cmod > 0: + Jmod = _cost_functions_numpy.calculate_model_cost( + winds[0], + winds[1], + winds[2], + parameters.model_weights, + parameters.u_model, + parameters.v_model, + parameters.w_model, + coeff=parameters.Cmod, + ) + else: + Jmod = 0 + + if parameters.Cpoint > 0: + Jpoint = _cost_functions_numpy.calculate_point_cost( + winds[0], + winds[1], + parameters.x, + parameters.y, + parameters.z, + parameters.point_list, + Cp=parameters.Cpoint, + roi=parameters.roi, + ) + else: + Jpoint = 0 + elif parameters.engine == "jax": + return J_function_jax(winds, parameters) + + if parameters.Nfeval % 10 == 0: + print( + ( + "Nfeval | Jvel | Jmass | Jsmooth | Jbg | Jvort | Jmodel | Jpoint |" + + " Max w " + ) + ) + print( + ( + "{:7d}".format(int(parameters.Nfeval)) + + "|" + + "{:9.4f}".format(float(Jvel)) + + "|" + + "{:9.4f}".format(float(Jmass)) + + "|" + + "{:9.4f}".format(float(Jsmooth)) + + "|" + + "{:9.4f}".format(float(Jbackground)) + + "|" + + "{:9.4f}".format(float(Jvorticity)) + + "|" + + "{:9.4f}".format(float(Jmod)) + + "|" + + "{:9.4f}".format(float(Jpoint)) + + "|" + + "{:9.4f}".format(np.ma.max(np.ma.abs(winds[2]))) + ) + ) + + parameters.Nfeval += 1 + # print("The cost functions print", Jvel + Jmass) + + return Jvel + Jmass + Jsmooth + Jbackground + Jvorticity + Jmod + Jpoint + + +def grad_J(winds, parameters): + """ + Calculates the gradient of the cost function. This typically does not need + to be called directly as get_dd_wind_field is a wrapper around this + function and :py:func:`pydda.cost_functions.J_function`. + In order to add more terms to the cost function, + modify this function and :py:func:`pydda.cost_functions.grad_J`. + + Parameters + ---------- + winds: 1-D float array + The wind field, flattened to 1-D for f_min + parameters: DDParameters + The parameters for the cost function evaluation as specified by the + :py:func:`pydda.retrieve.DDParameters` class. + + Returns + ------- + grad: 1D float array + Gradient vector of cost function + """ + if parameters.engine == "tensorflow": + if not TENSORFLOW_AVAILABLE: + raise ImportError( + "Tensorflow 2.5 or greater is needed in order to use TensorFlow-based PyDDA!" + ) + winds = tf.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + + winds = tf.math.maximum(winds, tf.constant([-100.0])) + winds = tf.math.minimum(winds, tf.constant([100.0])) + grad = _cost_functions_tensorflow.calculate_grad_radial_vel( + parameters.vrs, + parameters.els, + parameters.azs, + winds[0], + winds[1], + winds[2], + parameters.wts, + parameters.weights, + parameters.rmsVr, + coeff=parameters.Co, + upper_bc=parameters.upper_bc, + lower_bc=parameters.lower_bc, + ) + + if parameters.Cm > 0: + grad += _cost_functions_tensorflow.calculate_mass_continuity_gradient( + winds[0], + winds[1], + winds[2], + parameters.z, + parameters.dx, + parameters.dy, + parameters.dz, + coeff=parameters.Cm, + upper_bc=parameters.upper_bc, + lower_bc=parameters.lower_bc, + ) + + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: + grad += _cost_functions_tensorflow.calculate_smoothness_gradient( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + Cx=parameters.Cx, + Cy=parameters.Cy, + Cz=parameters.Cz, + upper_bc=parameters.upper_bc, + ) + + if parameters.Cb > 0: + grad += _cost_functions_tensorflow.calculate_background_gradient( + winds[0], + winds[1], + parameters.bg_weights, + parameters.u_back, + parameters.v_back, + parameters.Cb, + ) + + if parameters.Cv > 0: + grad += _cost_functions_tensorflow.calculate_vertical_vorticity_gradient( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + parameters.Ut, + parameters.Vt, + coeff=parameters.Cv, + upper_bc=parameters.upper_bc, + lower_bc=parameters.lower_bc, + ).numpy() + + if parameters.Cmod > 0: + grad += _cost_functions_tensorflow.calculate_model_gradient( + winds[0], + winds[1], + winds[2], + parameters.model_weights, + parameters.u_model, + parameters.v_model, + parameters.w_model, + coeff=parameters.Cmod, + ) + + if parameters.Cpoint > 0: + grad += _cost_functions_tensorflow.calculate_point_gradient( + winds[0], + winds[1], + parameters.x, + parameters.y, + parameters.z, + parameters.point_list, + Cp=parameters.Cpoint, + roi=parameters.roi, + upper_bc=parameters.upper_bc, + ) + if parameters.const_boundary_cond is True: + grad = tf.reshape( + grad, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + + grad = tf.concat( + [ + tf.zeros( + ( + 1, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + dtype=tf.float32, + ), + grad[:, :, 1:-1, :], + tf.zeros( + ( + 1, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + dtype=tf.float32, + ), + ], + axis=0, + ) + grad = tf.concat( + [ + tf.zeros( + ( + 1, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + dtype=tf.float32, + ), + grad[:, :, :, -1:1], + tf.zeros( + ( + 1, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + dtype=tf.float32, + ), + ], + axis=0, + ) + grad = tf.reshape(grad, [-1]) + elif parameters.engine == "scipy": + winds = np.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + if parameters.parallel: + futures = [] + with ThreadPoolExecutor() as pool: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_grad_radial_vel, + parameters.vrs, + parameters.els, + parameters.azs, + winds[0], + winds[1], + winds[2], + parameters.wts, + parameters.weights, + parameters.rmsVr, + parameters.Co, + parameters.upper_bc, + True, + ) + ) + if parameters.Cm > 0: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_mass_continuity_gradient, + winds[0], + winds[1], + winds[2], + parameters.z, + parameters.dx, + parameters.dy, + parameters.dz, + parameters.vrs, + parameters.Cm, + 1, + parameters.upper_bc, + above=parameters.above + ) + ) + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_smoothness_gradient, + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + parameters.Cx, + parameters.Cy, + parameters.Cz, + parameters.upper_bc, + ) + ) + if parameters.Cb > 0: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_background_gradient, + winds[0], + winds[1], + winds[2], + parameters.bg_weights, + parameters.u_back, + parameters.v_back, + parameters.Cb, + ) + ) + if parameters.Cv > 0: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_vertical_vorticity_gradient, + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + parameters.Ut, + parameters.Vt, + parameters.Cv, + parameters.upper_bc, + ) + ) + if parameters.Cmod > 0: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_model_gradient, + winds[0], + winds[1], + winds[2], + parameters.model_weights, + parameters.u_model, + parameters.v_model, + parameters.w_model, + parameters.Cmod, + ) + ) + if parameters.Cpoint > 0: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_point_gradient, + winds[0], + winds[1], + parameters.x, + parameters.y, + parameters.z, + parameters.point_list, + parameters.Cpoint, + parameters.roi, + ) + ) + grad = sum(f.result() for f in futures) + else: + grad = _cost_functions_numpy.calculate_grad_radial_vel( + parameters.vrs, + parameters.els, + parameters.azs, + winds[0], + winds[1], + winds[2], + parameters.wts, + parameters.weights, + parameters.rmsVr, + coeff=parameters.Co, + upper_bc=parameters.upper_bc, + ) + + if parameters.Cm > 0: + grad += _cost_functions_numpy.calculate_mass_continuity_gradient( + winds[0], + winds[1], + winds[2], + parameters.z, + parameters.dx, + parameters.dy, + parameters.dz, + coeff=parameters.Cm, + upper_bc=parameters.upper_bc, + ) + + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: + grad += _cost_functions_numpy.calculate_smoothness_gradient( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + Cx=parameters.Cx, + Cy=parameters.Cy, + Cz=parameters.Cz, + upper_bc=parameters.upper_bc, + ) + + if parameters.Cb > 0: + grad += _cost_functions_numpy.calculate_background_gradient( + winds[0], + winds[1], + winds[2], + parameters.bg_weights, + parameters.u_back, + parameters.v_back, + parameters.Cb, + ) + + if parameters.Cv > 0: + grad += _cost_functions_numpy.calculate_vertical_vorticity_gradient( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + parameters.Ut, + parameters.Vt, + coeff=parameters.Cv, + upper_bc=parameters.upper_bc, + ) + + if parameters.Cmod > 0: + grad += _cost_functions_numpy.calculate_model_gradient( + winds[0], + winds[1], + winds[2], + parameters.model_weights, + parameters.u_model, + parameters.v_model, + parameters.w_model, + coeff=parameters.Cmod, + ) + + if parameters.Cpoint > 0: + grad += _cost_functions_numpy.calculate_point_gradient( + winds[0], + winds[1], + parameters.x, + parameters.y, + parameters.z, + parameters.point_list, + Cp=parameters.Cpoint, + roi=parameters.roi, + upper_bc=parameters.upper_bc, + ) + + # Let's see if we need to enforce strong boundary conditions + if parameters.const_boundary_cond is True: + grad = np.reshape( + grad, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + grad[:, :, 0, :] = 0 + grad[:, :, -1, :] = 0 + grad[:, :, :, 0] = 0 + grad[:, :, :, -1] = 0 + grad = grad.flatten() + elif parameters.engine == "jax": + grad = grad_jax(winds, parameters) + if parameters.const_boundary_cond is True: + grad = jnp.reshape( + grad, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + grad.at[:, :, 0, :].set(0) + grad.at[:, :, -1, :].set(0) + grad.at[:, :, :, 0].set(0) + grad.at[:, :, :, -1].set(0) + grad = grad.flatten() + return grad + + if parameters.Nfeval % 10 == 0: + print("The gradient of the cost functions is", str(np.linalg.norm(grad, 2))) + return grad + + +def J_function_jax(winds, parameters): + if not JAX_AVAILABLE: + raise ImportError("Jax is needed in order to use the Jax-based PyDDA!") + + winds = jnp.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + # Had to change to float because Jax returns device array (use np.float_()) + Jvel = _cost_functions_jax.calculate_radial_vel_cost_function( + parameters.vrs, + parameters.azs, + parameters.els, + winds[0], + winds[1], + winds[2], + parameters.wts, + rmsVr=parameters.rmsVr, + weights=parameters.weights, + coeff=parameters.Co, + ) + + if parameters.Cm > 0: + # Had to change to float because Jax returns device array (use np.float_()) + Jmass = _cost_functions_jax.calculate_mass_continuity( + winds[0], + winds[1], + winds[2], + parameters.z, + parameters.dx, + parameters.dy, + parameters.dz, + coeff=parameters.Cm, + ) + else: + Jmass = 0 + + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: + Jsmooth = _cost_functions_jax.calculate_smoothness_cost( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + Cx=parameters.Cx, + Cy=parameters.Cy, + Cz=parameters.Cz, + ) + else: + Jsmooth = 0 + + if parameters.Cb > 0: + Jbackground = _cost_functions_jax.calculate_background_cost( + winds[0], + winds[1], + winds[2], + parameters.bg_weights, + parameters.u_back, + parameters.v_back, + parameters.Cb, + ) + else: + Jbackground = 0 + + if parameters.Cv > 0: + # Had to change to float because Jax returns device array (use np.float_()) + Jvorticity = _cost_functions_jax.calculate_vertical_vorticity_cost( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + parameters.Ut, + parameters.Vt, + coeff=parameters.Cv, + ) + else: + Jvorticity = 0 + + if parameters.Cmod > 0: + Jmod = _cost_functions_jax.calculate_model_cost( + winds[0], + winds[1], + winds[2], + parameters.model_weights, + parameters.u_model, + parameters.v_model, + parameters.w_model, + coeff=parameters.Cmod, + ) + else: + Jmod = 0 + + if parameters.Cpoint > 0: + Jpoint = _cost_functions_jax.calculate_point_cost( + winds[0], + winds[1], + parameters.x, + parameters.y, + parameters.z, + parameters.point_list, + Cp=parameters.Cpoint, + roi=parameters.roi, + ) + else: + Jpoint = 0 + + return Jvel + Jsmooth + Jmass + Jmod + Jpoint + Jvorticity + Jbackground + + +def grad_jax(winds, parameters): + winds = jnp.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + grad = _cost_functions_jax.calculate_grad_radial_vel( + parameters.vrs, + parameters.els, + parameters.azs, + winds[0], + winds[1], + winds[2], + parameters.wts, + parameters.weights, + parameters.rmsVr, + coeff=parameters.Co, + upper_bc=parameters.upper_bc, + ) + + if parameters.Cm > 0: + grad += _cost_functions_jax.calculate_mass_continuity_gradient( + winds[0], + winds[1], + winds[2], + parameters.z, + parameters.dx, + parameters.dy, + parameters.dz, + coeff=parameters.Cm, + upper_bc=parameters.upper_bc, + ) + + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: + grad += _cost_functions_jax.calculate_smoothness_gradient( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + Cx=parameters.Cx, + Cy=parameters.Cy, + Cz=parameters.Cz, + upper_bc=parameters.upper_bc, + ) + + if parameters.Cb > 0: + grad += _cost_functions_jax.calculate_background_gradient( + winds[0], + winds[1], + winds[2], + parameters.bg_weights, + parameters.u_back, + parameters.v_back, + parameters.Cb, + ) + + if parameters.Cv > 0: + grad += _cost_functions_jax.calculate_vertical_vorticity_gradient( + winds[0], + winds[1], + winds[2], + parameters.dx, + parameters.dy, + parameters.dz, + parameters.Ut, + parameters.Vt, + coeff=parameters.Cv, + upper_bc=parameters.upper_bc, + ).numpy() + + if parameters.Cmod > 0: + grad += _cost_functions_jax.calculate_model_gradient( + winds[0], + winds[1], + winds[2], + parameters.model_weights, + parameters.u_model, + parameters.v_model, + parameters.w_model, + coeff=parameters.Cmod, + ) + + if parameters.Cpoint > 0: + grad += _cost_functions_jax.calculate_point_gradient( + winds[0], + winds[1], + parameters.x, + parameters.y, + parameters.z, + parameters.point_list, + Cp=parameters.Cpoint, + roi=parameters.roi, + ) + return grad + + +def calculate_fall_speed(grid, refl_field=None, frz=4500.0): + """ + Estimates fall speed based on reflectivity. + + Uses methodology of Mike Biggerstaff and Dan Betten + + Parameters + ---------- + Grid: Py-ART Grid + Py-ART Grid containing reflectivity to calculate fall speed from + refl_field: str + String containing name of reflectivity field. None will automatically + determine the name. + frz: float + Height of freezing level in m + + Returns + ------- + 3D float array: + Float array of terminal velocities + + """ + # Parse names of velocity field + if refl_field is None: + refl_field = pyart.config.get_field_name("reflectivity") + + refl = grid[refl_field].values + grid_z = grid["point_z"].values + np.zeros(refl.shape) + A = np.zeros(refl.shape) + B = np.zeros(refl.shape) + rho = np.exp(-grid_z / 10000.0) + A[np.logical_and(grid_z < frz, refl < 55)] = -2.6 + B[np.logical_and(grid_z < frz, refl < 55)] = 0.0107 + A[np.logical_and(grid_z < frz, np.logical_and(refl >= 55, refl < 60))] = -2.5 + B[np.logical_and(grid_z < frz, np.logical_and(refl >= 55, refl < 60))] = 0.013 + A[np.logical_and(grid_z < frz, refl > 60)] = -3.95 + B[np.logical_and(grid_z < frz, refl > 60)] = 0.0148 + A[np.logical_and(grid_z >= frz, refl < 33)] = -0.817 + B[np.logical_and(grid_z >= frz, refl < 33)] = 0.0063 + A[np.logical_and(grid_z >= frz, np.logical_and(refl >= 33, refl < 49))] = -2.5 + B[np.logical_and(grid_z >= frz, np.logical_and(refl >= 33, refl < 49))] = 0.013 + A[np.logical_and(grid_z >= frz, refl > 49)] = -3.95 + B[np.logical_and(grid_z >= frz, refl > 49)] = 0.0148 + + fallspeed = A * np.power(10, refl * B) * np.power(1.2 / rho, 0.4) + print(fallspeed.max()) + del A, B, rho + return np.ma.masked_invalid(fallspeed) diff --git a/pydda/cost_functions/_cost_functions_numpy.py b/pydda/cost_functions/_cost_functions_numpy.py index 95150117..390e1cdf 100644 --- a/pydda/cost_functions/_cost_functions_numpy.py +++ b/pydda/cost_functions/_cost_functions_numpy.py @@ -453,7 +453,7 @@ def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): def calculate_mass_continuity_gradient( - u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1, upper_bc=True + u, v, w, z, dx, dy, dz, vrs=0, coeff=1500.0, anel=1, upper_bc=True, above=2.0 ): """ Calculates the gradient of mass continuity cost function. This is done by @@ -500,10 +500,15 @@ def calculate_mass_continuity_gradient( grad_v = -np.gradient(div, dy, axis=1) * coeff grad_w = -np.gradient(div, dz, axis=0) * coeff - # Impermeability condition - grad_w[0, :, :] = 0 - if upper_bc is True: - grad_w[-1, :, :] = 0 + # Impermeability conditions + grad_w[0, :, :] = 0 # surface is impermeable + if upper_bc == 1:#is True: # impermeable at the grid top + grad_w[-1, :, :] = 0 + if upper_bc == 2: # impermeable at cloud top + N=np.sum(np.array(vrs)>-1000,axis=0) + z_mask = (z > above * 1000) + n_mask = (N == 0) + grad_w[z_mask & n_mask] = 0 y = np.stack([grad_u, grad_v, grad_w], axis=0) return y.flatten() diff --git a/pydda/cost_functions/cost_functions.py b/pydda/cost_functions/cost_functions.py index 198e09b7..60c921ea 100644 --- a/pydda/cost_functions/cost_functions.py +++ b/pydda/cost_functions/cost_functions.py @@ -543,9 +543,11 @@ def grad_J(winds, parameters): parameters.dx, parameters.dy, parameters.dz, + parameters.vrs, parameters.Cm, 1, parameters.upper_bc, + above=parameters.above ) ) if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: diff --git a/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py b/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py new file mode 100644 index 00000000..e88b37b1 --- /dev/null +++ b/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py @@ -0,0 +1,1596 @@ +""" +Created on Mon Aug 7 09:17:40 2017 + +@author: rjackson +""" + +import pyart +import numpy as np +import time +import math +import xarray as xr + +from scipy.interpolate import interp1d +from scipy.ndimage import convolve1d +from scipy.optimize import fmin_l_bfgs_b +from scipy.signal import savgol_filter +from .auglag import auglag +from ..io import read_from_pyart_grid + +_LEISE_KERNEL = np.array([-1 / 16, 1 / 4, 5 / 8, 1 / 4, -1 / 16]) + + +def _apply_low_pass_filter( + winds, filter_type, filter_window, filter_order, leise_nstep +): + """Smooth the (3, nz, ny, nx) wind array in place along all spatial axes.""" + if filter_type not in ("savgol", "leise"): + raise ValueError( + "filter_type must be 'savgol' or 'leise', got %r" % filter_type + ) + for c in range(3): + for axis in range(winds[c].ndim): + if filter_type == "savgol": + winds[c] = savgol_filter( + winds[c], filter_window, filter_order, axis=axis + ) + else: + if winds[c].shape[axis] < 5: + continue + for _ in range(leise_nstep): + winds[c] = convolve1d( + winds[c], _LEISE_KERNEL, axis=axis, mode="mirror" + ) + return winds + + +try: + import tensorflow_probability as tfp + import tensorflow as tf + + TENSORFLOW_AVAILABLE = True +except (ImportError, AttributeError): + TENSORFLOW_AVAILABLE = False + +try: + import jax.numpy as jnp + import jax + import jaxopt + + JAX_AVAILABLE = True +except ImportError: + JAX_AVAILABLE = False + +# imports changed to local import path to run on computer +from ..cost_functions import ( + J_function, + grad_J, + calculate_fall_speed, + grad_jax, + J_function_jax, +) +from copy import deepcopy +from .angles import add_azimuth_as_field, add_elevation_as_field + +_wprevmax = np.empty(0) +_wcurrmax = np.empty(0) +iterations = 0 + + +class DDParameters(object): + """ + This is a helper class for inserting more arguments into the :func:`pydda.cost_functions.J_function` and + :func:`pydda.cost_functions.grad_J` function. Since these cost functions take numerous parameters, this class + will store the needed parameters as one positional argument for easier readability of the code. + + In addition, class members can be added here so that those contributing more constraints to the variational + framework can add any parameters they may need. + + Attributes + ---------- + vrs: List of float arrays + List of radial velocities from each radar + azs: List of float arrays + List of azimuths from each radar + els: List of float arrays + List of elevations from each radar + wts: List of float arrays + Float array containing fall speed from radar. + u_back: 1D float array (number of vertical levels) + Background u wind + v_back: 1D float array (number of vertical levels) + Background v wind + u_model: list of 3D float arrays + U from each model integrated into the retrieval + v_model: list of 3D float arrays + V from each model integrated into the retrieval + w_model: + W from each model integrated into the retrieval + Co: float + Weighting coefficient for data constraint. + Cm: float + Weighting coefficient for mass continuity constraint. + Cx: float + Smoothing coefficient for x-direction + Cy: float + Smoothing coefficient for y-direction + Cz: float + Smoothing coefficient for z-direction + Cb: float + Coefficient for sounding constraint + Cv: float + Weight for cost function related to vertical vorticity equation. + Cmod: float + Coefficient for model constraint + Cpoint: float + Coefficient for point constraint + Ut: float + Prescribed storm motion. This is only needed if Cv is not zero. + Vt: float + Prescribed storm motion. This is only needed if Cv is not zero. + grid_shape: + Shape of wind grid + dx: + Spacing of grid in x direction + dy: + Spacing of grid in y direction + dz: + Spacing of grid in z direction + x: + E-W grid levels in m + y: + N-S grid levels in m + z: + Grid vertical levels in m + rmsVr: float + The sum of squares of velocity/num_points. Use for normalization + of data weighting coefficient + weights: n_radars by z_bins by y_bins x x_bins float array + Data weights for each pair of radars + bg_weights: z_bins by y_bins x x_bins float array + Data weights for sounding constraint + model_weights: n_models by z_bins by y_bins by x_bins float array + Data weights for each model. + point_list: list or None + point_list: list of dicts + List of point constraints. Each member is a dict with keys of "u", "v", + to correspond to each component of the wind field and "x", "y", "z" + to correspond to the location of the point observation in the Grid's + Cartesian coordinates. + roi: float + The radius of influence of each point observation in m. + above: float + The altitude below which the cloud top impermeability would not apply (minimum impermeable height) + upper_bc: int + 0 to not enforce impermeability at top of domain + 1 to enforce w=0 at top of domain (impermeability condition), + 2 to enforce w=0 at cloud top (or "above" if higher) (impermeability condition), + """ + + def __init__(self): + self.Ut = np.nan + self.Vt = np.nan + self.rmsVr = np.nan + self.grid_shape = None + self.Cmod = np.nan + self.Cpoint = np.nan + self.u_back = None + self.v_back = None + self.wts = [] + self.vrs = [] + self.azs = [] + self.els = [] + self.weights = [] + self.bg_weights = [] + self.model_weights = [] + self.u_model = [] + self.v_model = [] + self.w_model = [] + self.x = None + self.y = None + self.z = None + self.dx = np.nan + self.dy = np.nan + self.dz = np.nan + self.Co = 1.0 + self.Cm = 1500.0 + self.Cx = 0.0 + self.Cy = 0.0 + self.Cz = 0.0 + self.Cb = 0.0 + self.Cv = 0.0 + self.Cmod = 0.0 + self.Cpoint = 0.0 + self.Ut = 0.0 + self.Vt = 0.0 + self.upper_bc = True + self.lower_bc = True + self.roi = 1000.0 + self.frz = 4500.0 + self.Nfeval = 0.0 + self.engine = "scipy" + self.point_list = [] + self.cvtol = 1e-2 + self.gtol = 1e-2 + self.Jveltol = 100.0 + self.const_boundary_cond = False + self.parallel = False + + +def _get_dd_wind_field_scipy( + Grids, + u_init, + v_init, + w_init, + engine, + points=None, + vel_name=None, + refl_field=None, + u_back=None, + v_back=None, + z_back=None, + frz=4500.0, + Co=1.0, + Cm=1500.0, + Cx=0.0, + Cy=0.0, + Cz=0.0, + Cb=0.0, + Cv=0.0, + Cmod=0.0, + Cpoint=0.0, + cvtol=1e-2, + gtol=1e-2, + Jveltol=100.0, + Ut=None, + Vt=None, + low_pass_filter=True, + mask_outside_opt=False, + weights_obs=None, + weights_model=None, + weights_bg=None, + max_iterations=1000, + mask_w_outside_opt=True, + filter_type="savgol", + filter_window=5, + filter_order=3, + leise_nstep=1, + min_bca=30.0, + max_bca=150.0, + upper_bc=True, + model_fields=None, + output_cost_functions=True, + roi=1000.0, + wind_tol=0.1, + tolerance=1e-8, + const_boundary_cond=False, + max_wind_mag=100.0, + parallel=True, +): + global _wcurrmax + global _wprevmax + global iterations + + # We have to have a prescribed storm motion for vorticity constraint + if Ut is None or Vt is None: + if Cv != 0.0: + raise ValueError( + ( + "Ut and Vt cannot be None if vertical " + + "vorticity constraint is enabled!" + ) + ) + + if not isinstance(Grids, list): + raise ValueError("Grids has to be a list!") + + parameters = DDParameters() + parameters.Ut = Ut + parameters.Vt = Vt + parameters.engine = engine + parameters.const_boundary_cond = const_boundary_cond + print(parameters.const_boundary_cond) + # Ensure that all Grids are on the same coordinate system + prev_grid = Grids[0] + for g in Grids: + if not np.allclose(g["x"].values, prev_grid["x"].values, atol=10): + raise ValueError("Grids do not have equal x coordinates!") + + if not np.allclose(g["y"].values, prev_grid["y"].values, atol=10): + raise ValueError("Grids do not have equal y coordinates!") + + if not np.allclose(g["z"].values, prev_grid["z"].values, atol=10): + raise ValueError("Grids do not have equal z coordinates!") + + if not np.allclose( + g["origin_latitude"].values, prev_grid["origin_latitude"].values + ): + raise ValueError(("Grids have unequal origin lat/lons!")) + + prev_grid = g + + if engine.lower() == "auglag" and not TENSORFLOW_AVAILABLE: + raise ModuleNotFoundError( + "Tensorflow 2.6+ needs to be installed for the Augmented Lagrangian solver." + ) + + # Disable background constraint if none provided + if u_back is None or v_back is None: + parameters.u_back = np.zeros(u_init.shape[0]) + parameters.v_back = np.zeros(v_init.shape[0]) + else: + # Interpolate sounding to radar grid + print("Interpolating sounding to radar grid") + + if isinstance(u_back, np.ma.MaskedArray): + u_back = u_back.filled(-9999.0) + if isinstance(v_back, np.ma.MaskedArray): + v_back = v_back.filled(-9999.0) + if isinstance(z_back, np.ma.MaskedArray): + z_back = z_back.filled(-9999.0) + valid_inds = np.logical_and.reduce( + (u_back > -9998, v_back > -9998, z_back > -9998) + ) + u_interp = interp1d(z_back[valid_inds], u_back[valid_inds], bounds_error=False) + v_interp = interp1d(z_back[valid_inds], v_back[valid_inds], bounds_error=False) + if isinstance(Grids[0]["z"].values, np.ma.MaskedArray): + parameters.u_back = u_interp(Grids[0]["z"].values.filled(np.nan)) + parameters.v_back = v_interp(Grids[0]["z"].values.filled(np.nan)) + else: + parameters.u_back = u_interp(Grids[0]["z"].values) + parameters.v_back = v_interp(Grids[0]["z"].values) + + print("Grid levels:") + print(Grids[0]["z"].values) + + # Parse names of velocity field + if refl_field is None: + refl_field = pyart.config.get_field_name("reflectivity") + + # Parse names of velocity field + if vel_name is None: + vel_name = pyart.config.get_field_name("corrected_velocity") + winds = np.stack([u_init, v_init, w_init]) + + # Set up wind fields and weights from each radar + parameters.weights = np.zeros( + (len(Grids), u_init.shape[0], u_init.shape[1], u_init.shape[2]) + ) + + parameters.bg_weights = np.zeros(v_init.shape) + if model_fields is not None: + parameters.model_weights = np.ones( + (len(model_fields), u_init.shape[0], u_init.shape[1], u_init.shape[2]) + ) + else: + parameters.model_weights = np.zeros( + (1, u_init.shape[0], u_init.shape[1], u_init.shape[2]) + ) + + if model_fields is None: + if Cmod != 0.0: + raise ValueError("Cmod must be zero if model fields are not specified!") + + bca = np.zeros((len(Grids), len(Grids), u_init.shape[1], u_init.shape[2])) + sum_Vr = np.zeros(len(Grids)) + + for i in range(len(Grids)): + parameters.wts.append( + np.ma.masked_invalid( + calculate_fall_speed(Grids[i], refl_field=refl_field, frz=frz).squeeze() + ) + ) + + parameters.vrs.append(np.ma.masked_invalid(Grids[i][vel_name].values.squeeze())) + parameters.azs.append( + np.ma.masked_invalid(Grids[i]["AZ"].values.squeeze() * np.pi / 180) + ) + parameters.els.append( + np.ma.masked_invalid(Grids[i]["EL"].values.squeeze() * np.pi / 180) + ) + + if len(Grids) > 1: + for i in range(len(Grids)): + for j in range(len(Grids)): + if i == j: + continue + print(("Calculating weights for radars " + str(i) + " and " + str(j))) + bca[i, j] = get_bca(Grids[i], Grids[j]) + + for k in range(parameters.vrs[i].shape[0]): + if weights_obs is None: + valid = np.logical_and.reduce( + ( + ~parameters.vrs[i][k].mask, + ~parameters.wts[i][k].mask, + ~parameters.azs[i][k].mask, + ~parameters.els[i][k].mask, + ) + ) + valid = np.logical_and.reduce( + ( + valid, + np.isfinite(parameters.vrs[i][k]), + np.isfinite(parameters.wts[i][k]), + np.isfinite(parameters.azs[i][k]), + np.isfinite(parameters.els[i][k]), + ) + ) + valid = np.logical_and.reduce( + ( + valid, + np.isfinite(parameters.vrs[j][k]), + np.isfinite(parameters.wts[j][k]), + np.isfinite(parameters.azs[j][k]), + np.isfinite(parameters.els[j][k]), + ) + ) + valid = np.logical_and.reduce( + ( + valid, + ~parameters.vrs[j][k].mask, + ~parameters.wts[j][k].mask, + ~parameters.azs[j][k].mask, + ~parameters.els[j][k].mask, + ) + ) + cur_array = parameters.weights[i, k].copy() + cur_array[ + np.logical_and( + valid, + np.logical_and( + bca[i, j] >= math.radians(min_bca), + bca[i, j] <= math.radians(max_bca), + ), + ) + ] = 1 + cur_array[~valid] = 0 + parameters.weights[i, k] += cur_array + else: + parameters.weights[i, k] = weights_obs[i][k, :, :] + + if weights_bg is None: + valid = np.logical_and.reduce( + ( + ~parameters.vrs[j][k].mask, + ~parameters.wts[j][k].mask, + ~parameters.azs[j][k].mask, + ~parameters.els[j][k].mask, + ) + ) + valid = np.logical_and.reduce( + ( + valid, + np.isfinite(parameters.vrs[j][k]), + np.isfinite(parameters.wts[j][k]), + np.isfinite(parameters.azs[j][k]), + np.isfinite(parameters.els[j][k]), + ) + ) + valid = np.logical_and.reduce( + ( + valid, + np.isfinite(parameters.vrs[j][k]), + np.isfinite(parameters.wts[j][k]), + np.isfinite(parameters.azs[j][k]), + np.isfinite(parameters.els[j][k]), + ) + ) + valid = np.logical_and.reduce( + ( + valid, + ~parameters.vrs[j][k].mask, + ~parameters.wts[j][k].mask, + ~parameters.azs[j][k].mask, + ~parameters.els[j][k].mask, + ) + ) + cur_array = parameters.bg_weights[k] + cur_array[ + np.logical_or.reduce( + ( + ~valid, + bca[i, j] < math.radians(min_bca), + bca[i, j] > math.radians(max_bca), + ) + ) + ] = 1 + cur_array[~valid] = 1 + parameters.bg_weights[i] += cur_array + else: + parameters.bg_weights[i] = weights_bg[i] + + print("Calculating weights for models...") + coverage_grade = parameters.weights.sum(axis=0) + coverage_grade = coverage_grade / coverage_grade.max() + + # Weigh in model input more when we have no coverage + # Model only weighs 1/(# of grids + 1) when there is full + # Coverage + if model_fields is not None: + if weights_model is None: + for i in range(len(model_fields)): + parameters.model_weights[i] = 1 - ( + coverage_grade / (len(Grids) + 1) + ) + else: + for i in range(len(model_fields)): + parameters.model_weights[i] = weights_model[i] + else: + if weights_obs is None: + parameters.weights[0] = np.where(~parameters.vrs[0].mask, 1, 0) + else: + parameters.weights[0] = weights_obs[0] + + if weights_bg is None: + parameters.bg_weights = np.where(~parameters.vrs[0].mask, 0, 1) + else: + parameters.bg_weights = weights_bg + + parameters.vrs = [x.filled(-9999.0) for x in parameters.vrs] + parameters.azs = [x.filled(-9999.0) for x in parameters.azs] + parameters.els = [x.filled(-9999.0) for x in parameters.els] + parameters.wts = [x.filled(-9999.0) for x in parameters.wts] + parameters.weights[~np.isfinite(parameters.weights)] = 0 + parameters.bg_weights[~np.isfinite(parameters.bg_weights)] = 0 + parameters.weights[parameters.weights > 0] = 1 + parameters.bg_weights[parameters.bg_weights > 0] = 1 + + # Zero out bg_weights at height levels where the interpolated background + # is NaN (i.e. outside the sounding's vertical range). Also replace NaN + # in u_back/v_back with 0 so those levels don't corrupt cost function + # arithmetic even though they carry zero weight. + nan_bg_levels = ~np.isfinite(parameters.u_back) | ~np.isfinite(parameters.v_back) + parameters.bg_weights[nan_bg_levels] = 0 + parameters.u_back = np.nan_to_num(parameters.u_back) + parameters.v_back = np.nan_to_num(parameters.v_back) + sum_Vr = np.nansum(np.square(parameters.vrs * parameters.weights)) + parameters.rmsVr = np.sqrt(np.nansum(sum_Vr) / np.nansum(parameters.weights)) + + del bca + parameters.grid_shape = u_init.shape + # Parse names of velocity field + + winds = winds.flatten() + + print("Starting solver ") + parameters.dx = np.diff(Grids[0]["x"].values, axis=0)[0] + parameters.dy = np.diff(Grids[0]["y"].values, axis=0)[0] + parameters.dz = np.diff(Grids[0]["z"].values, axis=0)[0] + print("rmsVR = " + str(parameters.rmsVr)) + print("Total points: %d" % parameters.weights.sum()) + parameters.z = Grids[0]["point_z"].values + parameters.x = Grids[0]["point_x"].values + parameters.y = Grids[0]["point_y"].values + bt = time.time() + + # First pass - no filter + wcurrmax = w_init.max() + print("The max of w_init is", wcurrmax) + iterations = 0 + bounds = [(-x, x) for x in max_wind_mag * np.ones(winds.shape)] + + if model_fields is not None: + for i, the_field in enumerate(model_fields): + u_field = "U_" + the_field + v_field = "V_" + the_field + w_field = "W_" + the_field + parameters.u_model.append(np.nan_to_num(Grids[0][u_field].values.squeeze())) + parameters.v_model.append(np.nan_to_num(Grids[0][v_field].values.squeeze())) + parameters.w_model.append(np.nan_to_num(Grids[0][w_field].values.squeeze())) + + # Don't weigh in where model data unavailable + where_finite_u = np.isfinite(Grids[0][u_field].values.squeeze()) + where_finite_v = np.isfinite(Grids[0][v_field].values.squeeze()) + where_finite_w = np.isfinite(Grids[0][w_field].values.squeeze()) + parameters.model_weights[i, :, :, :] = np.where( + np.logical_and.reduce((where_finite_u, where_finite_v, where_finite_w)), + 1, + 0, + ) + + print("Total number of model points: %d" % np.sum(parameters.model_weights)) + parameters.Co = Co + parameters.Cm = Cm + parameters.Cx = Cx + parameters.Cy = Cy + parameters.Cz = Cz + parameters.Cb = Cb + parameters.Cv = Cv + parameters.Cmod = Cmod + parameters.Cpoint = Cpoint + parameters.roi = roi + parameters.upper_bc = upper_bc + parameters.points = points + parameters.point_list = points + parameters.parallel = parallel + _wprevmax = np.zeros(parameters.grid_shape) + _wcurrmax = np.zeros(parameters.grid_shape) + iterations = 0 + if engine.lower() == "scipy" or engine.lower() == "jax": + + def _vert_velocity_callback(x): + global _wprevmax + global _wcurrmax + global iterations + + if iterations % 10 > 0: + iterations = iterations + 1 + return False + + wind = np.reshape( + x, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + _wcurrmax = wind[2] + if iterations == 0: + _wprevmax = _wcurrmax + iterations = iterations + 1 + return False + diff = np.abs(_wprevmax - _wcurrmax) + diff = np.where(parameters.bg_weights == 0, diff, np.nan) + delta = np.nanmax(diff) + if delta < wind_tol: + return True + _wprevmax = _wcurrmax + iterations = iterations + 1 + print("Max change in w: %4.3f" % delta) + return False + + parameters.print_out = False + if engine.lower() == "scipy": + winds = fmin_l_bfgs_b( + J_function, + winds, + args=(parameters,), + maxiter=max_iterations, + pgtol=tolerance, + bounds=bounds, + fprime=grad_J, + callback=_vert_velocity_callback, + ) + else: + + def loss_and_gradient(x): + x_loss = J_function_jax(x["winds"], parameters) + x_grad = {} + x_grad["winds"] = grad_jax(x["winds"], parameters) + return x_loss, x_grad + + bounds = ( + {"winds": -max_wind_mag * jnp.ones(winds.shape)}, + {"winds": max_wind_mag * jnp.ones(winds.shape)}, + ) + winds = jnp.array(winds) + # JIT-compile the cost function explicitly so the compilation + # delay is isolated and visible before the solver loop starts. + loss_and_gradient = jax.jit(loss_and_gradient) + print("Compiling JAX cost functions...") + loss_and_gradient({"winds": winds}) + print("Compilation complete.") + solver = jaxopt.LBFGSB( + loss_and_gradient, + True, + has_aux=False, + maxiter=max_iterations, + tol=tolerance, + jit=False, + implicit_diff=False, + verbose=True, + ) + winds = {"winds": winds} + winds, state = solver.run(winds, bounds=bounds) + winds = [np.asanyarray(winds["winds"])] + + winds = np.reshape( + winds[0], + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + parameters.print_out = True + + elif engine.lower() == "auglag": + if not TENSORFLOW_AVAILABLE: + raise ImportError( + "Tensorflow must be available to use the Augmented Lagrangian engine!" + ) + parameters.vrs = [tf.constant(x, dtype=tf.float32) for x in parameters.vrs] + parameters.azs = [tf.constant(x, dtype=tf.float32) for x in parameters.azs] + parameters.els = [tf.constant(x, dtype=tf.float32) for x in parameters.els] + parameters.wts = [tf.constant(x, dtype=tf.float32) for x in parameters.wts] + parameters.model_weights = tf.constant( + parameters.model_weights, dtype=tf.float32 + ) + parameters.weights[~np.isfinite(parameters.weights)] = 0 + parameters.weights[parameters.weights > 0] = 1 + parameters.weights = tf.constant(parameters.weights, dtype=tf.float32) + parameters.bg_weights[parameters.bg_weights > 0] = 1 + parameters.bg_weights = tf.constant(parameters.bg_weights, dtype=tf.float32) + parameters.z = tf.constant(Grids[0]["point_z"].values, dtype=tf.float32) + parameters.x = tf.constant(Grids[0]["point_x"].values, dtype=tf.float32) + parameters.y = tf.constant(Grids[0]["point_y"].values, dtype=tf.float32) + bounds = [(-x, x) for x in max_wind_mag * np.ones(winds.shape, dtype="float32")] + winds = winds.astype("float32") + winds, mult, AL_Filter, funcalls = auglag(winds, parameters, bounds) + + # """ + winds = np.stack([winds[0], winds[1], winds[2]]) + winds = winds.flatten() + if low_pass_filter is True: + print("Applying %s low pass filter to wind field..." % filter_type) + winds = np.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + winds = _apply_low_pass_filter( + winds, filter_type, filter_window, filter_order, leise_nstep + ) + winds = np.stack([winds[0], winds[1], winds[2]]) + winds = winds.flatten() + + print("Done! Time = " + "{:2.1f}".format(time.time() - bt)) + + # First pass - no filter + the_winds = np.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + u = the_winds[0] + v = the_winds[1] + w = the_winds[2] + where_mask = np.sum(parameters.weights, axis=0) + np.sum( + parameters.model_weights, axis=0 + ) + + u = np.ma.array(u) + w = np.ma.array(w) + v = np.ma.array(v) + + if mask_outside_opt is True: + u = np.ma.masked_where(where_mask < 1, u) + v = np.ma.masked_where(where_mask < 1, v) + w = np.ma.masked_where(where_mask < 1, w) + + if mask_w_outside_opt is True: + w = np.ma.masked_where(where_mask < 1, w) + + u_field = {} + u_field["standard_name"] = "u_wind" + u_field["long_name"] = "zonal component of wind velocity" + u_field["units"] = "m/s" + u_field["min_bca"] = min_bca + u_field["max_bca"] = max_bca + v_field = {} + v_field["standard_name"] = "v_wind" + v_field["long_name"] = "meridional component of wind velocity" + v_field["units"] = "m/s" + v_field["min_bca"] = min_bca + v_field["max_bca"] = max_bca + w_field = {} + w_field["standard_name"] = "w_wind" + w_field["long_name"] = "vertical component of wind velocity" + w_field["units"] = "m/s" + w_field["min_bca"] = min_bca + w_field["max_bca"] = max_bca + + new_grid_list = [] + + for grid in Grids: + grid["u"] = xr.DataArray( + np.expand_dims(u, 0), dims=("time", "z", "y", "x"), attrs=u_field + ) + grid["v"] = xr.DataArray( + np.expand_dims(v, 0), dims=("time", "z", "y", "x"), attrs=v_field + ) + grid["w"] = xr.DataArray( + np.expand_dims(w, 0), dims=("time", "z", "y", "x"), attrs=w_field + ) + new_grid_list.append(grid) + + return new_grid_list, parameters + + +def _get_dd_wind_field_tensorflow( + Grids, + u_init, + v_init, + w_init, + points=None, + vel_name=None, + refl_field=None, + u_back=None, + v_back=None, + z_back=None, + frz=4500.0, + Co=1.0, + Cm=1500.0, + Cx=0.0, + Cy=0.0, + Cz=0.0, + Cb=0.0, + Cv=0.0, + Cmod=0.0, + Cpoint=0.0, + Ut=None, + Vt=None, + low_pass_filter=True, + mask_outside_opt=False, + weights_obs=None, + weights_model=None, + weights_bg=None, + max_iterations=200, + mask_w_outside_opt=True, + filter_type="savgol", + filter_window=5, + filter_order=3, + leise_nstep=1, + min_bca=30.0, + max_bca=150.0, + upper_bc=True, + model_fields=None, + output_cost_functions=True, + roi=1000.0, + lower_bc=True, + parallel_iterations=1, + wind_tol=0.1, + tolerance=1e-8, + const_boundary_cond=False, + max_wind_mag=100.0, +): + if not TENSORFLOW_AVAILABLE: + raise ImportError( + "Tensorflow >=2.5 and tensorflow-probability " + + "need to be installed in order to use the tensorflow engine." + ) + # We have to have a prescribed storm motion for vorticity constraint + if Ut is None or Vt is None: + if Cv != 0.0: + raise ValueError( + ( + "Ut and Vt cannot be None if vertical " + + "vorticity constraint is enabled!" + ) + ) + + if not isinstance(Grids, list): + raise ValueError("Grids has to be a list!") + + parameters = DDParameters() + parameters.Ut = Ut + parameters.Vt = Vt + parameters.upper_bc = upper_bc + parameters.lower_bc = lower_bc + parameters.engine = "tensorflow" + parameters.const_boundary_cond = const_boundary_cond + + # Ensure that all Grids are on the same coordinate system + prev_grid = Grids[0] + for g in Grids: + if not np.allclose(g["x"].values, prev_grid["x"].values, atol=10): + raise ValueError("Grids do not have equal x coordinates!") + + if not np.allclose(g["y"].values, prev_grid["y"].values, atol=10): + raise ValueError("Grids do not have equal y coordinates!") + + if not np.allclose(g["z"].values, prev_grid["z"].values, atol=10): + raise ValueError("Grids do not have equal z coordinates!") + + if not np.allclose( + g["origin_latitude"].values, prev_grid["origin_latitude"].values + ): + raise ValueError(("Grids have unequal origin lat/lons!")) + + prev_grid = g + + # Disable background constraint if none provided + if u_back is None or v_back is None: + parameters.u_back = tf.zeros(u_init.shape[0]) + parameters.v_back = tf.zeros(v_init.shape[0]) + else: + # Interpolate sounding to radar grid + print("Interpolating sounding to radar grid") + + if isinstance(u_back, np.ma.MaskedArray): + u_back = u_back.filled(-9999.0) + if isinstance(v_back, np.ma.MaskedArray): + v_back = v_back.filled(-9999.0) + if isinstance(z_back, np.ma.MaskedArray): + z_back = z_back.filled(-9999.0) + valid_inds = np.logical_and.reduce( + (u_back > -9998, v_back > -9998, z_back > -9998) + ) + u_interp = interp1d(z_back[valid_inds], u_back[valid_inds], bounds_error=False) + v_interp = interp1d(z_back[valid_inds], v_back[valid_inds], bounds_error=False) + if isinstance(Grids[0]["z"].values, np.ma.MaskedArray): + parameters.u_back = tf.constant( + u_interp(Grids[0]["z"].values.filled(np.nan)), dtype=tf.float32 + ) + parameters.v_back = tf.constant( + v_interp(Grids[0]["z"].values.filled(np.nan)), dtype=tf.float32 + ) + else: + parameters.u_back = tf.constant( + u_interp(Grids[0]["z"].values), dtype=tf.float32 + ) + parameters.v_back = tf.constant( + v_interp(Grids[0]["z"].values), dtype=tf.float32 + ) + + print("Interpolated U field:") + print(parameters.u_back) + print("Interpolated V field:") + print(parameters.v_back) + print("Grid levels:") + print(Grids[0]["z"].values) + + # Parse names of velocity field + if refl_field is None: + refl_field = pyart.config.get_field_name("reflectivity") + + # Parse names of velocity field + if vel_name is None: + vel_name = pyart.config.get_field_name("corrected_velocity") + winds = np.stack([u_init, v_init, w_init]) + winds = winds.astype(np.float32) + + # Set up wind fields and weights from each radar + parameters.weights = np.zeros( + (len(Grids), u_init.shape[0], u_init.shape[1], u_init.shape[2]), + dtype=np.float32, + ) + + parameters.bg_weights = np.zeros(v_init.shape) + if model_fields is not None: + parameters.model_weights = np.ones( + (len(model_fields), u_init.shape[0], u_init.shape[1], u_init.shape[2]), + dtype=np.float32, + ) + else: + parameters.model_weights = np.zeros( + (1, u_init.shape[0], u_init.shape[1], u_init.shape[2]), dtype=np.float32 + ) + + if model_fields is None: + if Cmod != 0.0: + raise ValueError("Cmod must be zero if model fields are not specified!") + + bca = np.zeros( + (len(Grids), len(Grids), u_init.shape[1], u_init.shape[2]), dtype=np.float32 + ) + + for i in range(len(Grids)): + parameters.wts.append( + np.ma.masked_invalid( + calculate_fall_speed(Grids[i], refl_field=refl_field, frz=frz).squeeze() + ) + ) + parameters.vrs.append(np.ma.masked_invalid(Grids[i][vel_name].values.squeeze())) + parameters.azs.append( + np.ma.masked_invalid(Grids[i]["AZ"].values.squeeze() * np.pi / 180) + ) + parameters.els.append( + np.ma.masked_invalid(Grids[i]["EL"].values.squeeze() * np.pi / 180) + ) + + if len(Grids) > 1: + for i in range(len(Grids)): + for j in range(len(Grids)): + if i == j: + continue + print(("Calculating weights for radars " + str(i) + " and " + str(j))) + bca[i, j] = get_bca(Grids[i], Grids[j]) + + for k in range(parameters.vrs[i].shape[0]): + if weights_obs is None: + cur_array = parameters.weights[i, k].copy() + valid = np.logical_and.reduce( + ( + ~parameters.vrs[i][k].mask, + ~parameters.wts[i][k].mask, + ~parameters.azs[i][k].mask, + ~parameters.els[i][k].mask, + ) + ) + valid = np.logical_and.reduce( + ( + valid, + np.isfinite(parameters.vrs[i][k]), + np.isfinite(parameters.wts[i][k]), + np.isfinite(parameters.azs[i][k]), + np.isfinite(parameters.els[i][k]), + ) + ) + valid = np.logical_and.reduce( + ( + valid, + np.isfinite(parameters.vrs[j][k]), + np.isfinite(parameters.wts[j][k]), + np.isfinite(parameters.azs[j][k]), + np.isfinite(parameters.els[j][k]), + ) + ) + valid = np.logical_and.reduce( + ( + valid, + ~parameters.vrs[j][k].mask, + ~parameters.wts[j][k].mask, + ~parameters.azs[j][k].mask, + ~parameters.els[j][k].mask, + ) + ) + + cur_array[ + np.logical_and( + valid, + np.logical_and( + bca[i, j] >= math.radians(min_bca), + bca[i, j] <= math.radians(max_bca), + ), + ) + ] = 1 + cur_array[~valid] = 0 + parameters.weights[i, k] += cur_array + else: + parameters.weights[i, k] = weights_obs[i][k, :, :] + + if weights_bg is None: + cur_array = parameters.bg_weights[k] + valid = np.logical_and.reduce( + ( + ~parameters.vrs[i][k].mask, + ~parameters.wts[i][k].mask, + ~parameters.azs[i][k].mask, + ~parameters.els[i][k].mask, + ) + ) + valid = np.logical_and.reduce( + ( + valid, + np.isfinite(parameters.vrs[i][k]), + np.isfinite(parameters.wts[i][k]), + np.isfinite(parameters.azs[i][k]), + np.isfinite(parameters.els[i][k]), + ) + ) + valid = np.logical_and.reduce( + ( + valid, + np.isfinite(parameters.vrs[j][k]), + np.isfinite(parameters.wts[j][k]), + np.isfinite(parameters.azs[j][k]), + np.isfinite(parameters.els[j][k]), + ) + ) + valid = np.logical_and.reduce( + ( + valid, + ~parameters.vrs[j][k].mask, + ~parameters.wts[j][k].mask, + ~parameters.azs[j][k].mask, + ~parameters.els[j][k].mask, + ) + ) + cur_array[ + np.logical_or.reduce( + ( + ~valid, + bca[i, j] < math.radians(min_bca), + bca[i, j] > math.radians(max_bca), + ) + ) + ] = 1 + cur_array[~valid] = 1 + parameters.bg_weights[i] += cur_array + else: + parameters.bg_weights[i] = weights_bg[i] + + print("Calculating weights for models...") + coverage_grade = parameters.weights.sum(axis=0) + coverage_grade = coverage_grade / coverage_grade.max() + + # Weigh in model input more when we have no coverage + # Model only weighs 1/(# of grids + 1) when there is full + # Coverage + if model_fields is not None: + if weights_model is None: + for i in range(len(model_fields)): + parameters.model_weights[i] = 1 - ( + coverage_grade / (len(Grids) + 1) + ) + + else: + for i in range(len(model_fields)): + parameters.model_weights[i] = weights_model[i] + else: + if weights_obs is None: + parameters.weights[0] = np.where(np.isfinite(parameters.vrs[0]), 1, 0) + else: + parameters.weights[0] = weights_obs[0] + + if weights_bg is None: + parameters.bg_weights = np.where(np.isfinite(parameters.vrs[0]), 0, 1) + else: + parameters.bg_weights = weights_bg + + parameters.vrs = [ + tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.vrs + ] + parameters.azs = [ + tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.azs + ] + parameters.els = [ + tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.els + ] + parameters.wts = [ + tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.wts + ] + + parameters.weights[~np.isfinite(parameters.weights)] = 0 + parameters.weights[parameters.weights > 0] = 1 + for i in range(len(Grids)): + print("Points from Radar %d: %d" % (i, parameters.weights[i].sum())) + parameters.weights = tf.constant(parameters.weights, dtype=tf.float32) + parameters.bg_weights[parameters.bg_weights > 0] = 1 + parameters.bg_weights = tf.constant(parameters.bg_weights, dtype=tf.float32) + sum_Vr = tf.experimental.numpy.nansum( + tf.square(parameters.vrs * parameters.weights) + ) + parameters.rmsVr = np.sqrt( + np.nansum(sum_Vr) / tf.experimental.numpy.nansum(parameters.weights) + ) + + del bca + parameters.grid_shape = u_init.shape + # Parse names of velocity field + + winds = winds.flatten() + winds = tf.Variable(winds, name="winds") + + print("Starting solver ") + parameters.dx = np.diff(Grids[0]["x"].values, axis=0)[0] + parameters.dy = np.diff(Grids[0]["y"].values, axis=0)[0] + parameters.dz = np.diff(Grids[0]["z"].values, axis=0)[0] + print("rmsVR = " + str(parameters.rmsVr)) + print("Total points: %d" % tf.reduce_sum(parameters.weights)) + parameters.z = tf.constant(Grids[0]["point_z"].values, dtype=tf.float32) + parameters.x = tf.constant(Grids[0]["point_x"].values, dtype=tf.float32) + parameters.y = tf.constant(Grids[0]["point_y"].values, dtype=tf.float32) + bt = time.time() + + # First pass - no filter + wcurrmax = w_init.max() + print("The max of w_init is", wcurrmax) + [(-x, x) for x in 100.0 * np.ones(winds.shape)] + + if model_fields is not None: + for i, the_field in enumerate(model_fields): + u_field = "U_" + the_field + v_field = "V_" + the_field + w_field = "W_" + the_field + parameters.u_model.append( + tf.constant(np.nan_to_num(Grids[0][u_field].values.squeeze())) + ) + parameters.v_model.append( + tf.constant(np.nan_to_num(Grids[0][v_field].values.squeeze())) + ) + parameters.w_model.append( + tf.constant(np.nan_to_num(Grids[0][w_field].values.squeeze())) + ) + + # Don't weigh in where model data unavailable + where_finite_u = np.isfinite(Grids[0][u_field].values.squeeze()) + where_finite_v = np.isfinite(Grids[0][v_field].values.squeeze()) + where_finite_w = np.isfinite(Grids[0][w_field].values.squeeze()) + parameters.model_weights[i, :, :, :] = np.where( + np.logical_and.reduce((where_finite_u, where_finite_v, where_finite_w)), + 1, + 0, + ) + + parameters.model_weights = tf.constant(parameters.model_weights, dtype=tf.float32) + + parameters.Co = Co + parameters.Cm = Cm + parameters.Cx = Cx + parameters.Cy = Cy + parameters.Cz = Cz + parameters.Cb = Cb + parameters.Cv = Cv + parameters.Cmod = Cmod + parameters.Cpoint = Cpoint + parameters.roi = roi + parameters.upper_bc = upper_bc + parameters.points = points + parameters.point_list = points + loss_and_gradient = lambda x: (J_function(x, parameters), grad_J(x, parameters)) + + winds = tfp.optimizer.lbfgs_minimize( + loss_and_gradient, + initial_position=winds, + tolerance=tolerance, + x_tolerance=wind_tol, + max_iterations=max_iterations, + parallel_iterations=parallel_iterations, + max_line_search_iterations=20, + ) + winds = np.reshape( + winds.position.numpy(), + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + wcurrmax = winds[2].max() + winds = np.stack([winds[0], winds[1], winds[2]]) + winds = winds.flatten() + # """ + + if low_pass_filter: + print("Applying %s low pass filter to wind field..." % filter_type) + winds = np.asarray(winds) + winds = np.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + winds = _apply_low_pass_filter( + winds, filter_type, filter_window, filter_order, leise_nstep + ) + winds = np.stack([winds[0], winds[1], winds[2]]) + winds = winds.flatten() + + print("Done! Time = " + "{:2.1f}".format(time.time() - bt)) + + the_winds = np.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + u = the_winds[0] + v = the_winds[1] + w = the_winds[2] + where_mask = np.sum(parameters.weights, axis=0) + np.sum( + parameters.model_weights, axis=0 + ) + + u = np.ma.array(u) + w = np.ma.array(w) + v = np.ma.array(v) + + if mask_outside_opt is True: + u = np.ma.masked_where(where_mask < 1, u) + v = np.ma.masked_where(where_mask < 1, v) + w = np.ma.masked_where(where_mask < 1, w) + + if mask_w_outside_opt is True: + w = np.ma.masked_where(where_mask < 1, w) + + u_field = {} + u_field["standard_name"] = "u_wind" + u_field["long_name"] = "zonal component of wind velocity" + u_field["units"] = "m/s" + u_field["min_bca"] = min_bca + u_field["max_bca"] = max_bca + v_field = {} + v_field["standard_name"] = "v_wind" + v_field["long_name"] = "meridional component of wind velocity" + v_field["units"] = "m/s" + v_field["min_bca"] = min_bca + v_field["max_bca"] = max_bca + w_field = {} + w_field["standard_name"] = "w_wind" + w_field["long_name"] = "vertical component of wind velocity" + w_field["units"] = "m/s" + w_field["min_bca"] = min_bca + w_field["max_bca"] = max_bca + + new_grid_list = [] + + for grid in Grids: + grid["u"] = xr.DataArray( + np.expand_dims(u, 0), dims=("time", "z", "y", "x"), attrs=u_field + ) + grid["v"] = xr.DataArray( + np.expand_dims(v, 0), dims=("time", "z", "y", "x"), attrs=v_field + ) + grid["w"] = xr.DataArray( + np.expand_dims(w, 0), dims=("time", "z", "y", "x"), attrs=w_field + ) + new_grid_list.append(grid) + + return new_grid_list, parameters + + +def get_dd_wind_field( + Grids, u_init=None, v_init=None, w_init=None, engine="scipy", **kwargs +): + """ + This function takes in a list of Py-ART Grid objects and derives a + wind field. Every Py-ART Grid in Grids must have the same grid + specification. + + In order for the model data constraint to be used, + the model data must be added as a field to at least one of the + grids in Grids. This involves interpolating the model data to the + Grids' coordinates. There are helper functions for this for WRF + and HRRR data in :py:func:`pydda.constraints`: + + :py:func:`make_constraint_from_wrf` + + :py:func:`add_hrrr_constraint_to_grid` + + Parameters + ========== + + Grids: list of Py-ART/DDA Grids + The list of Py-ART or PyDDA grids to take in corresponding to each radar. + All grids must have the same shape, x coordinates, y coordinates + and z coordinates. + u_init: 3D ndarray + The initial guess for the zonal wind field, input as a 3D array + with the same shape as the fields in Grids. If this is None, + PyDDA will use the u field in the first Grid as the initalization. + v_init: 3D ndarray + The initial guess for the meridional wind field, input as a 3D array + with the same shape as the fields in Grids. If this is None, + PyDDA will use the v field in the first Grid as the initalization. + w_init: 3D ndarray + The initial guess for the vertical wind field, input as a 3D array + with the same shape as the fields in Grids. If this is None, + PyDDA will use the w field in the first Grid as the initalization. + engine: str (one of "scipy", "tensorflow", "jax") + Setting this flag will use the solver based off of SciPy, TensorFlow, or Jax. + Using TensorFlow or Jax expands PyDDA's capability to take advantage of GPU-based systems. + In addition, these two implementations use automatic differentation to calculate the gradient + of the cost function in order to optimize the gradient calculation. + TensorFlow 2.6 and tensorflow-probability are required for the TensorFlow-based engine. + The latest version of Jax is required for the Jax-based engine. + points: None or list of dicts + Point observations as returned by :func:`pydda.constraints.get_iem_obs`. Set + to None to disable. + vel_name: string + Name of radial velocity field. Setting to None will have PyDDA attempt + to automatically detect the velocity field name. + refl_field: string + Name of reflectivity field. Setting to None will have PyDDA attempt + to automatically detect the reflectivity field name. + u_back: 1D array + Background zonal wind field from a sounding as a function of height. + This should be given in the sounding's vertical coordinates. + v_back: 1D array + Background meridional wind field from a sounding as a function of + height. This should be given in the sounding's vertical coordinates. + z_back: 1D array + Heights corresponding to background wind field levels in meters. This + is given in the sounding's original coordinates. + frz: float + Freezing level used for fall speed calculation in meters. + Co: float + Weight for cost function related to observed radial velocities. + Cm: float + Weight for cost function related to the mass continuity equation. + Cx: float + Weight for cost function related to smoothness in x direction + Cy: float + Weight for cost function related to smoothness in y direction + Cz: float + Weight for cost function related to smoothness in z direction + Cv: float + Weight for cost function related to vertical vorticity equation. + Cmod: float + Weight for cost function related to custom constraints. + Cpoint: float + Weight for cost function related to point observations. + weights_obs: list of floating point arrays or None + List of weights for each point in grid from each radar in Grids. + Set to None to let PyDDA determine this automatically. + weights_model: list of floating point arrays or None + List of weights for each point in grid from each custom field in + model_fields. Set to None to let PyDDA determine this automatically. + weights_bg: list of floating point arrays or None + List of weights for each point in grid from the sounding. Set to None + to let PyDDA determine this automatically. + Ut: float + Prescribed storm motion in zonal direction. + This is only needed if Cv is not zero. + Vt: float + Prescribed storm motion in meridional direction. + This is only needed if Cv is not zero. + filter_winds: bool + If this is True, PyDDA will run a low pass filter on + the retrieved wind field. Set to False to disable the low pass filter. + mask_outside_opt: bool + If set to true, wind values outside the multiple doppler lobes will + be masked, i.e. if less than 2 radars provide coverage for a given + point. + max_iterations: int + The maximum number of iterations to run the optimization loop for. + mask_w_outside_opt: bool + If set to true, vertical winds outside the multiple doppler lobes will + be masked, i.e. if less than 2 radars provide coverage for a given + point. + filter_type: str (one of "savgol", "leise") + Which low-pass filter to apply after the optimization. ``"savgol"`` + (default) uses ``scipy.signal.savgol_filter`` along each axis with the + ``filter_window`` / ``filter_order`` parameters below. ``"leise"`` uses + the iterated 5-point Leise kernel ([-1/16, 1/4, 5/8, 1/4, -1/16]) with + mirror boundaries, controlled by ``leise_nstep``. + filter_window: int + Window size to use for the Savitzky-Golay low pass filter. A larger + window will increase the number of points factored into the polynomial + fit for the filter, and hence will increase the smoothness. Only used + when ``filter_type="savgol"``. + filter_order: int + The order of the polynomial to use for the Savitzky-Golay low pass + filter. Higher order polynomials allow for the retention of smaller + scale features but may also not remove enough noise. Only used when + ``filter_type="savgol"``. + leise_nstep: int + Number of Leise filter passes to apply along each spatial axis. Each + pass narrows the passband further. Only used when + ``filter_type="leise"``. + min_bca: float + Minimum beam crossing angle in degrees between two radars. 30.0 is the + typical value used in many publications. + max_bca: float + Maximum beam crossing angle in degrees between two radars. 150.0 is the + typical value used in many publications. + upper_bc: bool + Set this to true to enforce w = 0 at the top of the atmosphere. This is + commonly called the impermeability condition. + model_fields: list of strings + The list of fields in the first grid in Grids that contain the custom + data interpolated to the Grid's grid specification. Helper functions + to create such gridded fields for HRRR and NetCDF WRF data exist + in :py:func:`pydda.constraints`. PyDDA will look for fields named *U_(model + field name)*, *V_(model field name)*, and *W_(model field name)*. For + example, if you have *U_hrrr*, *V_hrrr*, and *W_hrrr*, then specify *["hrrr"]* + into model_fields. + output_cost_functions: bool + Set to True to output the value of each cost function every + 10 iterations. + roi: float + Radius of influence for the point observations. The point observation will + not hold any weight outside this radius. + parallel_iterations: int + The number of iterations to run in parallel in the optimization loop. + This is only for the TensorFlow-based engine. + wind_tol: float + Stop iterations after maximum change in winds is less than this value. + tolerance: float + Tolerance for :math:`L_{2}` norm of gradient before stopping. + max_wind_mag: float + Constrain the optimization to have :math:`|u|`, :math:`|v|`, and :math:`|w| < x` m/s. + parallel: bool + If True, enables parallelized cost and gradient computations for the scipy engine. + This vectorizes the radar loop in the radial velocity cost/gradient functions and + computes independent constraint gradients concurrently using a thread pool. + Default is False. + + Returns + ======= + new_grid_list: list + A list of Py-ART grids containing the derived wind fields. These fields + are displayable by the visualization module. + parameters: struct + The parameters used in the generation of the Multi-Doppler wind field. + """ + + if isinstance(Grids, list): + if isinstance(Grids[0], pyart.core.Grid): + for x in Grids: + new_grids = [read_from_pyart_grid(x) for x in Grids] + else: + new_grids = Grids + elif isinstance(Grids, pyart.core.Grid): + new_grids = [read_from_pyart_grid(Grids)] + elif isinstance(Grids, xr.Dataset): + new_grids = [Grids] + else: + raise TypeError( + "Input grids must be an xarray Dataset, Py-ART Grid, or a list of those." + ) + + if u_init is None: + u_init = new_grids[0]["u"].values.squeeze() + + if v_init is None: + v_init = new_grids[0]["v"].values.squeeze() + + if w_init is None: + w_init = new_grids[0]["w"].values.squeeze() + + if ( + engine.lower() == "scipy" + or engine.lower() == "jax" + or engine.lower() == "auglag" + ): + return _get_dd_wind_field_scipy( + new_grids, u_init, v_init, w_init, engine, **kwargs + ) + elif engine.lower() == "tensorflow": + return _get_dd_wind_field_tensorflow( + new_grids, u_init, v_init, w_init, **kwargs + ) + else: + raise NotImplementedError("Engine %s is not supported." % engine) + + +def get_bca(Grid1, Grid2): + """ + This function gets the beam crossing angle between two lat/lon pairs. + + Parameters + ========== + Grid1: xarray (PyDDA) Dataset + The PyDDA Dataset storing the first radar's Grid. + Grid2: PyDDA Dataset + The PyDDA Dataset storing the second radar's Grid. + + Returns + ======= + bca: nD float array + The beam crossing angle between the two radars in radians. + + """ + rad1_lon = Grid1["radar_longitude"].values + rad1_lat = Grid1["radar_latitude"].values + rad2_lon = Grid2["radar_longitude"].values + rad2_lat = Grid2["radar_latitude"].values + x = Grid1["point_x"].values + y = Grid1["point_y"].values + projparams = Grid1["projection"].attrs + if projparams["_include_lon_0_lat_0"] == "true": + projparams["lat_0"] = Grid1["origin_latitude"].values + projparams["lon_0"] = Grid1["origin_longitude"].values + + rad1 = pyart.core.geographic_to_cartesian(rad1_lon, rad1_lat, projparams) + rad2 = pyart.core.geographic_to_cartesian(rad2_lon, rad2_lat, projparams) + # Create grid with Radar 1 in center + + x = x - rad1[0] + y = y - rad1[1] + rad2 = np.array(rad2) - np.array(rad1) + a = np.sqrt(np.multiply(x, x) + np.multiply(y, y)) + b = np.sqrt(pow(x - rad2[0], 2) + pow(y - rad2[1], 2)) + c = np.sqrt(rad2[0] * rad2[0] + rad2[1] * rad2[1]) + inp_array1 = x / a + inp_array1 = np.where(inp_array1 < -1, -1, inp_array1) + inp_array1 = np.where(inp_array1 > 1, 1, inp_array1) + inp_array2 = (x - rad2[1]) / b + inp_array2 = np.where(inp_array2 < -1, -1, inp_array2) + inp_array2 = np.where(inp_array2 > 1, 1, inp_array2) + inp_array3 = (a * a + b * b - c * c) / (2 * a * b) + inp_array3 = np.where(inp_array3 < -1, -1, inp_array3) + inp_array3 = np.where(inp_array3 > 1, 1, inp_array3) + + return np.ma.masked_invalid(np.arccos(inp_array3))[0, :, :] diff --git a/pydda/retrieval/wind_retrieve.py b/pydda/retrieval/wind_retrieve.py index ff22c37c..e88b37b1 100644 --- a/pydda/retrieval/wind_retrieve.py +++ b/pydda/retrieval/wind_retrieve.py @@ -159,9 +159,12 @@ class DDParameters(object): Cartesian coordinates. roi: float The radius of influence of each point observation in m. - upper_bc: bool - True to enforce w=0 at top of domain (impermeability condition), - False to not enforce impermeability at top of domain + above: float + The altitude below which the cloud top impermeability would not apply (minimum impermeable height) + upper_bc: int + 0 to not enforce impermeability at top of domain + 1 to enforce w=0 at top of domain (impermeability condition), + 2 to enforce w=0 at cloud top (or "above" if higher) (impermeability condition), """ def __init__(self): From cefe348f3f6abe9d18afe4c014e014fe87458f14 Mon Sep 17 00:00:00 2001 From: Rob Thompson Date: Wed, 12 Aug 2026 11:47:50 +0000 Subject: [PATCH 2/4] minor fix where above had not been correctly implemented --- .../retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py | 4 +++- pydda/retrieval/wind_retrieve.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py b/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py index e88b37b1..e757a8df 100644 --- a/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py +++ b/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py @@ -206,6 +206,7 @@ def __init__(self): self.upper_bc = True self.lower_bc = True self.roi = 1000.0 + self.above = 2.0 self.frz = 4500.0 self.Nfeval = 0.0 self.engine = "scipy" @@ -257,10 +258,11 @@ def _get_dd_wind_field_scipy( leise_nstep=1, min_bca=30.0, max_bca=150.0, - upper_bc=True, + upper_bc=1, model_fields=None, output_cost_functions=True, roi=1000.0, + above=2.0, wind_tol=0.1, tolerance=1e-8, const_boundary_cond=False, diff --git a/pydda/retrieval/wind_retrieve.py b/pydda/retrieval/wind_retrieve.py index e88b37b1..e757a8df 100644 --- a/pydda/retrieval/wind_retrieve.py +++ b/pydda/retrieval/wind_retrieve.py @@ -206,6 +206,7 @@ def __init__(self): self.upper_bc = True self.lower_bc = True self.roi = 1000.0 + self.above = 2.0 self.frz = 4500.0 self.Nfeval = 0.0 self.engine = "scipy" @@ -257,10 +258,11 @@ def _get_dd_wind_field_scipy( leise_nstep=1, min_bca=30.0, max_bca=150.0, - upper_bc=True, + upper_bc=1, model_fields=None, output_cost_functions=True, roi=1000.0, + above=2.0, wind_tol=0.1, tolerance=1e-8, const_boundary_cond=False, From 617b55d47870936b9100e73adff352dbb78e6de1 Mon Sep 17 00:00:00 2001 From: Rob Thompson Date: Wed, 12 Aug 2026 12:14:19 +0000 Subject: [PATCH 3/4] Remove Jupyter checkpoint files --- .../_cost_functions_numpy-checkpoint.py | 884 --------- .../cost_functions-checkpoint.py | 1026 ----------- .../wind_retrieve-checkpoint.py | 1598 ----------------- 3 files changed, 3508 deletions(-) delete mode 100644 pydda/cost_functions/.ipynb_checkpoints/_cost_functions_numpy-checkpoint.py delete mode 100644 pydda/cost_functions/.ipynb_checkpoints/cost_functions-checkpoint.py delete mode 100644 pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py diff --git a/pydda/cost_functions/.ipynb_checkpoints/_cost_functions_numpy-checkpoint.py b/pydda/cost_functions/.ipynb_checkpoints/_cost_functions_numpy-checkpoint.py deleted file mode 100644 index 390e1cdf..00000000 --- a/pydda/cost_functions/.ipynb_checkpoints/_cost_functions_numpy-checkpoint.py +++ /dev/null @@ -1,884 +0,0 @@ -import numpy as np -import scipy -import pyart - -from scipy.ndimage import _nd_image - -laplace_filter = np.asarray([1, -2, 1], dtype=np.float64) - - -def calculate_radial_vel_cost_function( - vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0, parallel=False -): - """ - Calculates the cost function due to difference of the wind field from - radar radial velocities. For more information on this cost function, see - Potvin et al. (2012) and Shapiro et al. (2009). - All arrays in the given lists must have the same dimensions and represent - the same spatial coordinates. - Parameters - ---------- - vrs: List of float arrays - List of radial velocities from each radar - els: List of float arrays - List of elevations from each radar - azs: List of float arrays - List of azimuths from each radar - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - wts: List of float arrays - Float array containing fall speed from radar. - rmsVr: float - The sum of squares of velocity/num_points. Use for normalization - of data weighting coefficient - weights: n_radars x_bins x y_bins float array - Data weights for each pair of radars - coeff: float - Constant for cost function - Returns - ------- - J_o: float - Observational cost function - References - ----------- - Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity - Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and - Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49, - https://doi.org/10.1175/JTECH-D-11-00019.1 - Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity - Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic - Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1 - """ - - lambda_o = coeff / (rmsVr * rmsVr) - if parallel: - vrs_arr = np.stack(vrs) - els_arr = np.stack(els) - azs_arr = np.stack(azs) - wts_arr = np.stack(wts) - v_ar = ( - np.cos(els_arr) * np.sin(azs_arr) * u[np.newaxis] - + np.cos(els_arr) * np.cos(azs_arr) * v[np.newaxis] - + np.sin(els_arr) * (w[np.newaxis] - np.abs(wts_arr)) - ) - return lambda_o * np.sum(np.square(vrs_arr - v_ar) * weights) - - J_o = 0 - for i in range(len(vrs)): - v_ar = ( - np.cos(els[i]) * np.sin(azs[i]) * u - + np.cos(els[i]) * np.cos(azs[i]) * v - + np.sin(els[i]) * (w - np.abs(wts[i])) - ) - J_o += lambda_o * np.sum(np.square(vrs[i] - v_ar) * weights[i]) - - return J_o - - -def calculate_grad_radial_vel( - vrs, - els, - azs, - u, - v, - w, - wts, - weights, - rmsVr, - coeff=1.0, - upper_bc=True, - parallel=False, -): - """ - Calculates the gradient of the cost function due to difference of wind - field from radar radial velocities. - All arrays in the given lists must have the same dimensions and represent - the same spatial coordinates. - Parameters - ---------- - vrs: List of float arrays - List of radial velocities from each radar - els: List of float arrays - List of elevations from each radar - azs: List of azimuths - List of azimuths from each radar - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - coeff: float - Constant for cost function - vel_name: str - Background velocity field name - weights: n_radars x_bins x y_bins float array - Data weights for each pair of radars - Returns - ------- - y: 1-D float array - Gradient vector of observational cost function. - - More information - ---------------- - The gradient is calculated by taking the functional derivative of the - cost function. For more information on functional derivatives, see the - Euler-Lagrange Equation: - https://en.wikipedia.org/wiki/Euler%E2%80%93Lagrange_equation - """ - - # Use zero for all masked values since we don't want to add them into - # the cost function - - lambda_o = coeff / (rmsVr * rmsVr) - - if parallel: - vrs_arr = np.stack(vrs) - els_arr = np.stack(els) - azs_arr = np.stack(azs) - wts_arr = np.stack(wts) - v_ar = ( - np.cos(els_arr) * np.sin(azs_arr) * u[np.newaxis] - + np.cos(els_arr) * np.cos(azs_arr) * v[np.newaxis] - + np.sin(els_arr) * (w[np.newaxis] - np.abs(wts_arr)) - ) - residual = 2 * (v_ar - vrs_arr) * lambda_o - p_x1 = np.sum(residual * np.cos(els_arr) * np.sin(azs_arr) * weights, axis=0) - p_y1 = np.sum(residual * np.cos(els_arr) * np.cos(azs_arr) * weights, axis=0) - p_z1 = np.sum(residual * np.sin(els_arr) * weights, axis=0) - else: - p_x1 = np.zeros(vrs[0].shape) - p_y1 = np.zeros(vrs[0].shape) - p_z1 = np.zeros(vrs[0].shape) - - for i in range(len(vrs)): - v_ar = ( - np.cos(els[i]) * np.sin(azs[i]) * u - + np.cos(els[i]) * np.cos(azs[i]) * v - + np.sin(els[i]) * (w - np.abs(wts[i])) - ) - - x_grad = ( - 2 * (v_ar - vrs[i]) * np.cos(els[i]) * np.sin(azs[i]) * weights[i] - ) * lambda_o - y_grad = ( - 2 * (v_ar - vrs[i]) * np.cos(els[i]) * np.cos(azs[i]) * weights[i] - ) * lambda_o - z_grad = (2 * (v_ar - vrs[i]) * np.sin(els[i]) * weights[i]) * lambda_o - - p_x1 += x_grad - p_y1 += y_grad - p_z1 += z_grad - - # Impermeability condition - p_z1[0, :, :] = 0 - if upper_bc is True: - p_z1[-1, :, :] = 0 - y = np.stack((p_x1, p_y1, p_z1), axis=0) - return y.flatten() - - -def calculate_smoothness_cost(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5): - """ - Calculates the smoothness cost function by taking the Laplacian of the - wind field. - All arrays in the given lists must have the same dimensions and represent - the same spatial coordinates. - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - Cx: float - Constant controlling smoothness in x-direction - Cy: float - Constant controlling smoothness in y-direction - Cz: float - Constant controlling smoothness in z-direction - Returns - ------- - Js: float - value of smoothness cost function - """ - dudx = np.gradient(u, dx, axis=2) - dudy = np.gradient(u, dy, axis=1) - dudz = np.gradient(u, dz, axis=0) - dvdx = np.gradient(v, dx, axis=2) - dvdy = np.gradient(v, dy, axis=1) - dvdz = np.gradient(v, dz, axis=0) - dwdx = np.gradient(w, dx, axis=2) - dwdy = np.gradient(w, dy, axis=1) - dwdz = np.gradient(w, dz, axis=0) - - x_term = ( - Cx - * ( - np.gradient(dudx, dx, axis=2) - + np.gradient(dvdx, dx, axis=2) - + np.gradient(dwdx, dx, axis=2) - ) - ** 2 - ) - y_term = ( - Cy - * ( - np.gradient(dudy, dy, axis=1) - + np.gradient(dvdy, dy, axis=1) - + np.gradient(dwdy, dy, axis=1) - ) - ** 2 - ) - z_term = ( - Cz - * ( - np.gradient(dudz, dz, axis=0) - + np.gradient(dvdz, dz, axis=0) - + np.gradient(dwdz, dz, axis=0) - ) - ** 2 - ) - return np.sum(np.nan_to_num(x_term + y_term + z_term)) - - -def calculate_smoothness_gradient( - u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5, upper_bc=True -): - """ - Calculates the gradient of the smoothness cost function - by taking the Laplacian of the Laplacian of the wind field. - All arrays in the given lists must have the same dimensions and represent - the same spatial coordinates. - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - Cx: float - Constant controlling smoothness in x-direction - Cy: float - Constant controlling smoothness in y-direction - Cz: float - Constant controlling smoothness in z-direction - Returns - ------- - y: float array - value of gradient of smoothness cost function - """ - du = np.zeros(w.shape) - dv = np.zeros(w.shape) - dw = np.zeros(w.shape) - grad_u = np.zeros(w.shape) - grad_v = np.zeros(w.shape) - grad_w = np.zeros(w.shape) - scipy.ndimage.laplace(u, du, mode="wrap") - scipy.ndimage.laplace(v, dv, mode="wrap") - scipy.ndimage.laplace(w, dw, mode="wrap") - du = du / dx - dv = dv / dy - dw = dw / dz - scipy.ndimage.laplace(du, grad_u, mode="wrap") - scipy.ndimage.laplace(dv, grad_v, mode="wrap") - scipy.ndimage.laplace(dw, grad_w, mode="wrap") - grad_u = grad_u / dx - grad_v = grad_v / dy - grad_w = grad_w / dz - - # Impermeability condition - grad_w[0, :, :] = 0 - if upper_bc is True: - grad_w[-1, :, :] = 0 - - y = np.stack([grad_u, grad_v, grad_w], axis=0) - - return y.flatten() - - -def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2): - """ - Calculates the cost function related to point observations. A mean square error cost - function term is applied to points that are within the sphere of influence - whose radius is determined by *roi*. - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - x: Float array - X coordinates of grid centers - y: Float array - Y coordinates of grid centers - z: Float array - Z coordinated of grid centers - point_list: list of dicts - List of point constraints. - Each member is a dict with keys of "u", "v", to correspond - to each component of the wind field and "x", "y", "z" - to correspond to the location of the point observation. - In addition, "site_id" gives the METAR code (or name) to the station. - Cp: float - The weighting coefficient of the point cost function. - roi: float - Radius of influence of observations - Returns - ------- - J: float - The cost function related to the difference between wind field and points. - """ - J = 0.0 - for the_point in point_list: - # Instead of worrying about whole domain, just find points in radius of influence - # Since we know that the weight will be zero outside the sphere of influence anyways - - dist = np.sqrt( - (x - the_point["x"]) ** 2 - + (y - the_point["y"]) ** 2 - + (z - the_point["z"]) ** 2 - ) - dist = np.maximum(dist, 1.0) - weight = 1 / dist**2 - weight = weight / np.max(weight) - - J += np.sum(weight * ((u - the_point["u"]) ** 2 + (v - the_point["v"]) ** 2)) - - return J * Cp - - -def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): - """ - Calculates the gradient of the cost function related to point observations. - A mean square error cost function term is applied to points that are within the sphere of influence - whose radius is determined by *roi*. - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - x: Float array - X coordinates of grid centers - y: Float array - Y coordinates of grid centers - z: Float array - Z coordinated of grid centers - point_list: list of dicts - List of point constraints. Each member is a dict with keys of "u", "v", - to correspond to each component of the wind field and "x", "y", "z" - to correspond to the location of the point observation. - In addition, "site_id" gives the METAR code (or name) to the station. - Cp: float - The weighting coefficient of the point cost function. - roi: float - Radius of influence of observations - Returns - ------- - gradJ: float array - The gradient of the cost function related to the difference between wind field and points. - """ - - gradJ_u = np.zeros_like(u) - gradJ_v = np.zeros_like(v) - gradJ_w = np.zeros_like(u) - - for the_point in point_list: - dist = np.sqrt( - (x - the_point["x"]) ** 2 - + (y - the_point["y"]) ** 2 - + (z - the_point["z"]) ** 2 - ) - dist = np.maximum(dist, 1.0) - weight = 1 / dist**2 - weight = weight / np.max(weight) - gradJ_u += 2 * weight * (u - the_point["u"]) - gradJ_v += 2 * weight * (v - the_point["v"]) - - gradJ = np.stack([gradJ_u, gradJ_v, gradJ_w], axis=0).flatten() - return gradJ * Cp - - -def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): - """ - Calculates the mass continuity cost function by taking the divergence - of the wind field. - All arrays in the given lists must have the same dimensions and represent - the same spatial coordinates. - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - dx: float - Grid spacing in x direction. - dy: float - Grid spacing in y direction. - dz: float - Grid spacing in z direction. - z: Float array (1D) - 1D Float array with heights of grid - coeff: float - Constant controlling contribution of mass continuity to cost function - anel: int - = 1 use anelastic approximation, 0=don't - Returns - ------- - J: float - value of mass continuity cost function - """ - dudx = np.gradient(u, dx, axis=2) - dvdy = np.gradient(v, dy, axis=1) - dwdz = np.gradient(w, dz, axis=0) - - if anel == 1: - rho = np.exp(-z / 10000.0) - drho_dz = np.gradient(rho, dz, axis=0) - anel_term = w / rho * drho_dz - else: - anel_term = np.zeros(w.shape) - div = dudx + dvdy + dwdz + anel_term - - return coeff * np.sum(np.square(div)) / 2.0 - - -def calculate_mass_continuity_gradient( - u, v, w, z, dx, dy, dz, vrs=0, coeff=1500.0, anel=1, upper_bc=True, above=2.0 -): - """ - Calculates the gradient of mass continuity cost function. This is done by - taking the negative gradient of the divergence of the wind field. - All grids must have the same grid specification. - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - z: Float array (1D) - 1D Float array with heights of grid - dx: float - Grid spacing in x direction. - dy: float - Grid spacing in y direction. - dz: float - Grid spacing in z direction. - coeff: float - Constant controlling contribution of mass continuity to cost function - anel: int - = 1 use anelastic approximation, 0=don't - Returns - ------- - y: float array - value of gradient of mass continuity cost function - """ - dudx = np.gradient(u, dx, axis=2) - dvdy = np.gradient(v, dy, axis=1) - dwdz = np.gradient(w, dz, axis=0) - if anel == 1: - rho = np.exp(-z / 10000.0) - drho_dz = np.gradient(rho, dz, axis=0) - anel_term = w / rho * drho_dz - else: - anel_term = 0 - - div = dudx + dvdy + dwdz + anel_term - - grad_u = -np.gradient(div, dx, axis=2) * coeff - grad_v = -np.gradient(div, dy, axis=1) * coeff - grad_w = -np.gradient(div, dz, axis=0) * coeff - - # Impermeability conditions - grad_w[0, :, :] = 0 # surface is impermeable - if upper_bc == 1:#is True: # impermeable at the grid top - grad_w[-1, :, :] = 0 - if upper_bc == 2: # impermeable at cloud top - N=np.sum(np.array(vrs)>-1000,axis=0) - z_mask = (z > above * 1000) - n_mask = (N == 0) - grad_w[z_mask & n_mask] = 0 - y = np.stack([grad_u, grad_v, grad_w], axis=0) - return y.flatten() - - -def calculate_fall_speed(grid, refl_field=None, frz=4500.0): - """ - Estimates fall speed based on reflectivity. - Uses methodology of Mike Biggerstaff and Dan Betten - Parameters - ---------- - Grid: Py-ART Grid - Py-ART Grid containing reflectivity to calculate fall speed from - refl_field: str - String containing name of reflectivity field. None will automatically - determine the name. - frz: float - Height of freezing level in m - Returns - ------- - 3D float array: - Float array of terminal velocities - """ - # Parse names of velocity field - if refl_field is None: - refl_field = pyart.config.get_field_name("reflectivity") - - refl = grid[refl_field].values - grid_z = grid["point_z"].values - A = np.zeros(refl.shape) - B = np.zeros(refl.shape) - rho = np.exp(-grid_z / 10000.0) - A[np.logical_and(grid_z < frz, refl < 55)] = -2.6 - B[np.logical_and(grid_z < frz, refl < 55)] = 0.0107 - A[np.logical_and(grid_z < frz, np.logical_and(refl >= 55, refl < 60))] = -2.5 - B[np.logical_and(grid_z < frz, np.logical_and(refl >= 55, refl < 60))] = 0.013 - A[np.logical_and(grid_z < frz, refl > 60)] = -3.95 - B[np.logical_and(grid_z < frz, refl > 60)] = 0.0148 - A[np.logical_and(grid_z >= frz, refl < 33)] = -0.817 - B[np.logical_and(grid_z >= frz, refl < 33)] = 0.0063 - A[np.logical_and(grid_z >= frz, np.logical_and(refl >= 33, refl < 49))] = -2.5 - B[np.logical_and(grid_z >= frz, np.logical_and(refl >= 33, refl < 49))] = 0.013 - A[np.logical_and(grid_z >= frz, refl > 49)] = -3.95 - B[np.logical_and(grid_z >= frz, refl > 49)] = 0.0148 - - fallspeed = A * np.power(10, refl * B) * np.power(1.2 / rho, 0.4) - del A, B, rho - return fallspeed - - -def calculate_background_cost(u, v, w, weights, u_back, v_back, Cb=0.01): - """ - Calculates the background cost function. The background cost function is - simply the sum of the squared differences between the wind field and the - background wind field multiplied by the weighting coefficient. - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - weights: Float array - Weights for each point to consider into cost function - u_back: 1D float array - Zonal winds vs height from sounding - w_back: 1D float array - Meridional winds vs height from sounding - Cb: float - Weight of background constraint to total cost function - Returns - ------- - cost: float - value of background cost function - """ - the_shape = u.shape - cost = 0 - for i in range(the_shape[0]): - cost += Cb * np.sum( - np.square(u[i] - u_back[i]) * (weights[i]) - + np.square(v[i] - v_back[i]) * (weights[i]) - ) - return cost - - -def calculate_background_gradient(u, v, w, weights, u_back, v_back, Cb=0.01): - """ - Calculates the gradient of the background cost function. For each u, v - this is given as 2*coefficent*(analysis wind - background wind). - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - weights: Float array - Weights for each point to consider into cost function - u_back: 1D float array - Zonal winds vs height from sounding - w_back: 1D float array - Meridional winds vs height from sounding - Cb: float - Weight of background constraint to total cost function - Returns - ------- - y: float array - value of gradient of background cost function - """ - the_shape = u.shape - u_grad = np.zeros(the_shape) - v_grad = np.zeros(the_shape) - w_grad = np.zeros(the_shape) - - for i in range(the_shape[0]): - u_grad[i] = Cb * 2 * (u[i] - u_back[i]) * (weights[i]) - v_grad[i] = Cb * 2 * (v[i] - v_back[i]) * (weights[i]) - - y = np.stack([u_grad, v_grad, w_grad], axis=0) - return y.flatten() - - -def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5): - """ - Calculates the cost function due to deviance from vertical vorticity - equation. For more information of the vertical vorticity cost function, - see Potvin et al. (2012) and Shapiro et al. (2009). - Parameters - ---------- - u: 3D array - Float array with u component of wind field - v: 3D array - Float array with v component of wind field - w: 3D array - Float array with w component of wind field - dx: float array - Spacing in x grid - dy: float array - Spacing in y grid - dz: float array - Spacing in z grid - coeff: float - Weighting coefficient - Ut: float - U component of storm motion - Vt: float - V component of storm motion - Returns - ------- - Jv: float - Value of vertical vorticity cost function. - References - ---------- - Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity - Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and - Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49, - https://doi.org/10.1175/JTECH-D-11-00019.1 - Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity - Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic - Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1 - """ - dvdz = np.gradient(v, dz, axis=0) - dudz = np.gradient(u, dz, axis=0) - dvdx = np.gradient(v, dx, axis=2) - dwdy = np.gradient(w, dy, axis=1) - dwdx = np.gradient(w, dx, axis=2) - dudx = np.gradient(u, dx, axis=2) - dvdy = np.gradient(v, dy, axis=2) - dudy = np.gradient(u, dy, axis=1) - zeta = dvdx - dudy - dzeta_dx = np.gradient(zeta, dx, axis=2) - dzeta_dy = np.gradient(zeta, dy, axis=1) - dzeta_dz = np.gradient(zeta, dz, axis=0) - jv_array = ( - (u - Ut) * dzeta_dx - + (v - Vt) * dzeta_dy - + w * dzeta_dz - + (dvdz * dwdx - dudz * dwdy) - + zeta * (dudx + dvdy) - ) - return np.sum(coeff * jv_array**2) - - -def calculate_vertical_vorticity_gradient( - u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5, upper_bc=True -): - """ - Calculates the gradient of the cost function due to deviance from vertical - vorticity equation. This is done by taking the functional derivative of - the vertical vorticity cost function. - Parameters - ---------- - u: 3D array - Float array with u component of wind field - v: 3D array - Float array with v component of wind field - w: 3D array - Float array with w component of wind field - dx: float array - Spacing in x grid - dy: float array - Spacing in y grid - dz: float array - Spacing in z grid - Ut: float - U component of storm motion - Vt: float - V component of storm motion - coeff: float - Weighting coefficient - Returns - ------- - Jv: 1D float array - Value of the gradient of the vertical vorticity cost function. - References - ---------- - Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity - Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and - Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49, - https://doi.org/10.1175/JTECH-D-11-00019.1 - Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity - Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic - Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1 - """ - - # First derivatives - dvdz = np.gradient(v, dz, axis=0) - dwdy = np.gradient(w, dy, axis=1) - dudx = np.gradient(u, dx, axis=2) - dvdy = np.gradient(v, dy, axis=1) - dvdx = np.gradient(v, dx, axis=2) - dwdx = np.gradient(w, dx, axis=2) - dudz = np.gradient(u, dz, axis=0) - dudy = np.gradient(u, dy, axis=1) - - zeta = dvdx - dudy - dzeta_dx = np.gradient(zeta, dx, axis=2) - dzeta_dy = np.gradient(zeta, dy, axis=1) - dzeta_dz = np.gradient(zeta, dz, axis=0) - - # Second deriviatives - dwdydz = np.gradient(dwdy, dz, axis=0) - dwdxdz = np.gradient(dwdx, dz, axis=0) - dudzdy = np.gradient(dudz, dy, axis=1) - dvdxdy = np.gradient(dvdx, dy, axis=1) - dudx2 = np.gradient(dudx, dx, axis=2) - dudxdy = np.gradient(dudx, dy, axis=1) - dudxdz = np.gradient(dudx, dz, axis=0) - dudy2 = np.gradient(dudx, dy, axis=1) - - dzeta_dt = ( - (u - Ut) * dzeta_dx - + (v - Vt) * dzeta_dy - + w * dzeta_dz - + (dvdz * dwdx - dudz * dwdy) - + zeta * (dudx + dvdy) - ) - - # Now we intialize our gradient value - u_grad = np.zeros(u.shape) - v_grad = np.zeros(v.shape) - w_grad = np.zeros(w.shape) - - # Vorticity Advection - u_grad += dzeta_dx + (Ut - u) * dudxdy + (Vt - v) * dudxdy - v_grad += dzeta_dy + (Vt - v) * dvdxdy + (Ut - u) * dvdxdy - w_grad += dzeta_dz - - # Tilting term - u_grad += dwdydz - v_grad += dwdxdz - w_grad += dudzdy - dudxdz - - # Stretching term - u_grad += -dudxdy + dudy2 - dzeta_dx - u_grad += -dudx2 + dudxdy - dzeta_dy - - # Multiply by 2*dzeta_dt according to chain rule - u_grad = u_grad * 2 * dzeta_dt * coeff - v_grad = v_grad * 2 * dzeta_dt * coeff - w_grad = w_grad * 2 * dzeta_dt * coeff - - # Impermeability condition - w_grad[0, :, :] = 0 - if upper_bc is True: - w_grad[-1, :, :] = 0 - y = np.stack([u_grad, v_grad, w_grad], axis=0) - return y.flatten() - - -def calculate_model_cost(u, v, w, weights, u_model, v_model, w_model, coeff=1.0): - """ - Calculates the cost function for the model constraint. - This is calculated simply as the sum of squares of the differences - between the model wind field and the analysis wind field. Vertical - velocities are not factored into this cost function as there is typically - a high amount of uncertainty in model derived vertical velocities. - Parameters - ---------- - u: 3D array - Float array with u component of wind field - v: 3D array - Float array with v component of wind field - w: 3D array - Float array with w component of wind field - weights: list of 3D arrays - Float array showing how much each point from model weighs into - constraint. - u_model: list of 3D arrays - Float array with u component of wind field from model - v_model: list of 3D arrays - Float array with v component of wind field from model - w_model: list of 3D arrays - Float array with w component of wind field from model - coeff: float - Weighting coefficient - Returns - ------- - Jv: float - Value of model cost function - """ - - cost = 0 - for i in range(len(u_model)): - cost += coeff * np.sum( - np.square(u - u_model[i]) * weights[i] - + np.square(v - v_model[i]) * weights[i] - ) - return cost - - -def calculate_model_gradient(u, v, w, weights, u_model, v_model, w_model, coeff=1.0): - """ - Calculates the cost function for the model constraint. - This is calculated simply as twice the differences - between the model wind field and the analysis wind field for each u, v. - Vertical velocities are not factored into this cost function as there is - typically a high amount of uncertainty in model derived vertical - velocities. Therefore, the gradient for all of the w's will be 0. - Parameters - ---------- - u: Float array - Float array with u component of wind field - v: Float array - Float array with v component of wind field - w: Float array - Float array with w component of wind field - weights: list of 3D float arrays - Weights for each point to consider into cost function - u_model: list of 3D float arrays - Zonal wind field from model - v_model: list of 3D float arrays - Meridional wind field from model - w_model: list of 3D float arrays - Vertical wind field from model - coeff: float - Weight of background constraint to total cost function - Returns - ------- - y: float array - value of gradient of background cost function - """ - the_shape = u.shape - u_grad = np.zeros(the_shape) - v_grad = np.zeros(the_shape) - w_grad = np.zeros(the_shape) - for i in range(len(u_model)): - u_grad += coeff * 2 * (u - u_model[i]) * weights[i] - v_grad += coeff * 2 * (v - v_model[i]) * weights[i] - - y = np.stack([u_grad, v_grad, w_grad], axis=0) - return y.flatten() diff --git a/pydda/cost_functions/.ipynb_checkpoints/cost_functions-checkpoint.py b/pydda/cost_functions/.ipynb_checkpoints/cost_functions-checkpoint.py deleted file mode 100644 index 60c921ea..00000000 --- a/pydda/cost_functions/.ipynb_checkpoints/cost_functions-checkpoint.py +++ /dev/null @@ -1,1026 +0,0 @@ -import numpy as np -from concurrent.futures import ThreadPoolExecutor - -# Adding jax import statements -try: - import tensorflow as tf - - TENSORFLOW_AVAILABLE = True -except ImportError: - TENSORFLOW_AVAILABLE = False - -try: - import jax.numpy as jnp - - JAX_AVAILABLE = True -except ImportError: - JAX_AVAILABLE = False - -import pyart - -# Added to incorpeate JAX within the cost functions -from . import _cost_functions_jax -from . import _cost_functions_numpy -from . import _cost_functions_tensorflow - - -def J_function(winds, parameters): - """ - Calculates the total cost function. This typically does not need to be - called directly as get_dd_wind_field is a wrapper around this function and - :py:func:`pydda.cost_functions.grad_J`. - In order to add more terms to the cost function, modify this - function and :py:func:`pydda.cost_functions.grad_J`. - - Parameters - ---------- - winds: 1-D float array - The wind field, flattened to 1-D for f_min. The total size of the - array will be a 1D array of 3*nx*ny*nz elements. - parameters: DDParameters - The parameters for the cost function evaluation as specified by the - :py:func:`pydda.retrieval.DDParameters` class. - - Returns - ------- - J: float - The value of the cost function - """ - if parameters.engine == "tensorflow": - if not TENSORFLOW_AVAILABLE: - raise ImportError( - "Tensorflow 2.5 or greater is needed in order to use TensorFlow-based PyDDA!" - ) - winds = tf.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - winds = tf.math.maximum(winds, tf.constant([-100.0])) - winds = tf.math.minimum(winds, tf.constant([100.0])) - # Had to change to float because Jax returns device array (use np.float_()) - Jvel = _cost_functions_tensorflow.calculate_radial_vel_cost_function( - parameters.vrs, - parameters.azs, - parameters.els, - winds[0], - winds[1], - winds[2], - parameters.wts, - rmsVr=parameters.rmsVr, - weights=parameters.weights, - coeff=parameters.Co, - ) - # print("apples Jvel", Jvel) - - if parameters.Cm > 0: - # Had to change to float because Jax returns device array (use np.float_()) - Jmass = _cost_functions_tensorflow.calculate_mass_continuity( - winds[0], - winds[1], - winds[2], - parameters.z, - parameters.dx, - parameters.dy, - parameters.dz, - coeff=parameters.Cm, - ) - else: - Jmass = 0 - - if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: - Jsmooth = _cost_functions_tensorflow.calculate_smoothness_cost( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - Cx=parameters.Cx, - Cy=parameters.Cy, - Cz=parameters.Cz, - ) - else: - Jsmooth = 0 - - if parameters.Cb > 0: - Jbackground = _cost_functions_tensorflow.calculate_background_cost( - winds[0], - winds[1], - parameters.bg_weights, - parameters.u_back, - parameters.v_back, - parameters.Cb, - ) - else: - Jbackground = 0 - - if parameters.Cv > 0: - # Had to change to float because Jax returns device array (use np.float_()) - Jvorticity = _cost_functions_tensorflow.calculate_vertical_vorticity_cost( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - parameters.Ut, - parameters.Vt, - coeff=parameters.Cv, - ) - else: - Jvorticity = 0 - - if parameters.Cmod > 0: - Jmod = _cost_functions_tensorflow.calculate_model_cost( - winds[0], - winds[1], - winds[2], - parameters.model_weights, - parameters.u_model, - parameters.v_model, - parameters.w_model, - coeff=parameters.Cmod, - ) - else: - Jmod = 0 - - if parameters.Cpoint > 0: - Jpoint = _cost_functions_tensorflow.calculate_point_cost( - winds[0], - winds[1], - parameters.x, - parameters.y, - parameters.z, - parameters.point_list, - Cp=parameters.Cpoint, - roi=parameters.roi, - ) - else: - Jpoint = 0 - elif parameters.engine == "scipy": - winds = np.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - # Had to change to float because Jax returns device array (use np.float_()) - Jvel = _cost_functions_numpy.calculate_radial_vel_cost_function( - parameters.vrs, - parameters.azs, - parameters.els, - winds[0], - winds[1], - winds[2], - parameters.wts, - rmsVr=parameters.rmsVr, - weights=parameters.weights, - coeff=parameters.Co, - parallel=parameters.parallel, - ) - # print("apples Jvel", Jvel) - - if parameters.Cm > 0: - # Had to change to float because Jax returns device array (use np.float_()) - Jmass = _cost_functions_numpy.calculate_mass_continuity( - winds[0], - winds[1], - winds[2], - parameters.z, - parameters.dx, - parameters.dy, - parameters.dz, - coeff=parameters.Cm, - ) - else: - Jmass = 0 - - if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: - Jsmooth = _cost_functions_numpy.calculate_smoothness_cost( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - Cx=parameters.Cx, - Cy=parameters.Cy, - Cz=parameters.Cz, - ) - else: - Jsmooth = 0 - - if parameters.Cb > 0: - Jbackground = _cost_functions_numpy.calculate_background_cost( - winds[0], - winds[1], - winds[2], - parameters.bg_weights, - parameters.u_back, - parameters.v_back, - parameters.Cb, - ) - else: - Jbackground = 0 - - if parameters.Cv > 0: - # Had to change to float because Jax returns device array (use np.float_()) - Jvorticity = _cost_functions_numpy.calculate_vertical_vorticity_cost( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - parameters.Ut, - parameters.Vt, - coeff=parameters.Cv, - ) - else: - Jvorticity = 0 - - if parameters.Cmod > 0: - Jmod = _cost_functions_numpy.calculate_model_cost( - winds[0], - winds[1], - winds[2], - parameters.model_weights, - parameters.u_model, - parameters.v_model, - parameters.w_model, - coeff=parameters.Cmod, - ) - else: - Jmod = 0 - - if parameters.Cpoint > 0: - Jpoint = _cost_functions_numpy.calculate_point_cost( - winds[0], - winds[1], - parameters.x, - parameters.y, - parameters.z, - parameters.point_list, - Cp=parameters.Cpoint, - roi=parameters.roi, - ) - else: - Jpoint = 0 - elif parameters.engine == "jax": - return J_function_jax(winds, parameters) - - if parameters.Nfeval % 10 == 0: - print( - ( - "Nfeval | Jvel | Jmass | Jsmooth | Jbg | Jvort | Jmodel | Jpoint |" - + " Max w " - ) - ) - print( - ( - "{:7d}".format(int(parameters.Nfeval)) - + "|" - + "{:9.4f}".format(float(Jvel)) - + "|" - + "{:9.4f}".format(float(Jmass)) - + "|" - + "{:9.4f}".format(float(Jsmooth)) - + "|" - + "{:9.4f}".format(float(Jbackground)) - + "|" - + "{:9.4f}".format(float(Jvorticity)) - + "|" - + "{:9.4f}".format(float(Jmod)) - + "|" - + "{:9.4f}".format(float(Jpoint)) - + "|" - + "{:9.4f}".format(np.ma.max(np.ma.abs(winds[2]))) - ) - ) - - parameters.Nfeval += 1 - # print("The cost functions print", Jvel + Jmass) - - return Jvel + Jmass + Jsmooth + Jbackground + Jvorticity + Jmod + Jpoint - - -def grad_J(winds, parameters): - """ - Calculates the gradient of the cost function. This typically does not need - to be called directly as get_dd_wind_field is a wrapper around this - function and :py:func:`pydda.cost_functions.J_function`. - In order to add more terms to the cost function, - modify this function and :py:func:`pydda.cost_functions.grad_J`. - - Parameters - ---------- - winds: 1-D float array - The wind field, flattened to 1-D for f_min - parameters: DDParameters - The parameters for the cost function evaluation as specified by the - :py:func:`pydda.retrieve.DDParameters` class. - - Returns - ------- - grad: 1D float array - Gradient vector of cost function - """ - if parameters.engine == "tensorflow": - if not TENSORFLOW_AVAILABLE: - raise ImportError( - "Tensorflow 2.5 or greater is needed in order to use TensorFlow-based PyDDA!" - ) - winds = tf.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - - winds = tf.math.maximum(winds, tf.constant([-100.0])) - winds = tf.math.minimum(winds, tf.constant([100.0])) - grad = _cost_functions_tensorflow.calculate_grad_radial_vel( - parameters.vrs, - parameters.els, - parameters.azs, - winds[0], - winds[1], - winds[2], - parameters.wts, - parameters.weights, - parameters.rmsVr, - coeff=parameters.Co, - upper_bc=parameters.upper_bc, - lower_bc=parameters.lower_bc, - ) - - if parameters.Cm > 0: - grad += _cost_functions_tensorflow.calculate_mass_continuity_gradient( - winds[0], - winds[1], - winds[2], - parameters.z, - parameters.dx, - parameters.dy, - parameters.dz, - coeff=parameters.Cm, - upper_bc=parameters.upper_bc, - lower_bc=parameters.lower_bc, - ) - - if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: - grad += _cost_functions_tensorflow.calculate_smoothness_gradient( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - Cx=parameters.Cx, - Cy=parameters.Cy, - Cz=parameters.Cz, - upper_bc=parameters.upper_bc, - ) - - if parameters.Cb > 0: - grad += _cost_functions_tensorflow.calculate_background_gradient( - winds[0], - winds[1], - parameters.bg_weights, - parameters.u_back, - parameters.v_back, - parameters.Cb, - ) - - if parameters.Cv > 0: - grad += _cost_functions_tensorflow.calculate_vertical_vorticity_gradient( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - parameters.Ut, - parameters.Vt, - coeff=parameters.Cv, - upper_bc=parameters.upper_bc, - lower_bc=parameters.lower_bc, - ).numpy() - - if parameters.Cmod > 0: - grad += _cost_functions_tensorflow.calculate_model_gradient( - winds[0], - winds[1], - winds[2], - parameters.model_weights, - parameters.u_model, - parameters.v_model, - parameters.w_model, - coeff=parameters.Cmod, - ) - - if parameters.Cpoint > 0: - grad += _cost_functions_tensorflow.calculate_point_gradient( - winds[0], - winds[1], - parameters.x, - parameters.y, - parameters.z, - parameters.point_list, - Cp=parameters.Cpoint, - roi=parameters.roi, - upper_bc=parameters.upper_bc, - ) - if parameters.const_boundary_cond is True: - grad = tf.reshape( - grad, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - - grad = tf.concat( - [ - tf.zeros( - ( - 1, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - dtype=tf.float32, - ), - grad[:, :, 1:-1, :], - tf.zeros( - ( - 1, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - dtype=tf.float32, - ), - ], - axis=0, - ) - grad = tf.concat( - [ - tf.zeros( - ( - 1, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - dtype=tf.float32, - ), - grad[:, :, :, -1:1], - tf.zeros( - ( - 1, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - dtype=tf.float32, - ), - ], - axis=0, - ) - grad = tf.reshape(grad, [-1]) - elif parameters.engine == "scipy": - winds = np.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - if parameters.parallel: - futures = [] - with ThreadPoolExecutor() as pool: - futures.append( - pool.submit( - _cost_functions_numpy.calculate_grad_radial_vel, - parameters.vrs, - parameters.els, - parameters.azs, - winds[0], - winds[1], - winds[2], - parameters.wts, - parameters.weights, - parameters.rmsVr, - parameters.Co, - parameters.upper_bc, - True, - ) - ) - if parameters.Cm > 0: - futures.append( - pool.submit( - _cost_functions_numpy.calculate_mass_continuity_gradient, - winds[0], - winds[1], - winds[2], - parameters.z, - parameters.dx, - parameters.dy, - parameters.dz, - parameters.vrs, - parameters.Cm, - 1, - parameters.upper_bc, - above=parameters.above - ) - ) - if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: - futures.append( - pool.submit( - _cost_functions_numpy.calculate_smoothness_gradient, - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - parameters.Cx, - parameters.Cy, - parameters.Cz, - parameters.upper_bc, - ) - ) - if parameters.Cb > 0: - futures.append( - pool.submit( - _cost_functions_numpy.calculate_background_gradient, - winds[0], - winds[1], - winds[2], - parameters.bg_weights, - parameters.u_back, - parameters.v_back, - parameters.Cb, - ) - ) - if parameters.Cv > 0: - futures.append( - pool.submit( - _cost_functions_numpy.calculate_vertical_vorticity_gradient, - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - parameters.Ut, - parameters.Vt, - parameters.Cv, - parameters.upper_bc, - ) - ) - if parameters.Cmod > 0: - futures.append( - pool.submit( - _cost_functions_numpy.calculate_model_gradient, - winds[0], - winds[1], - winds[2], - parameters.model_weights, - parameters.u_model, - parameters.v_model, - parameters.w_model, - parameters.Cmod, - ) - ) - if parameters.Cpoint > 0: - futures.append( - pool.submit( - _cost_functions_numpy.calculate_point_gradient, - winds[0], - winds[1], - parameters.x, - parameters.y, - parameters.z, - parameters.point_list, - parameters.Cpoint, - parameters.roi, - ) - ) - grad = sum(f.result() for f in futures) - else: - grad = _cost_functions_numpy.calculate_grad_radial_vel( - parameters.vrs, - parameters.els, - parameters.azs, - winds[0], - winds[1], - winds[2], - parameters.wts, - parameters.weights, - parameters.rmsVr, - coeff=parameters.Co, - upper_bc=parameters.upper_bc, - ) - - if parameters.Cm > 0: - grad += _cost_functions_numpy.calculate_mass_continuity_gradient( - winds[0], - winds[1], - winds[2], - parameters.z, - parameters.dx, - parameters.dy, - parameters.dz, - coeff=parameters.Cm, - upper_bc=parameters.upper_bc, - ) - - if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: - grad += _cost_functions_numpy.calculate_smoothness_gradient( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - Cx=parameters.Cx, - Cy=parameters.Cy, - Cz=parameters.Cz, - upper_bc=parameters.upper_bc, - ) - - if parameters.Cb > 0: - grad += _cost_functions_numpy.calculate_background_gradient( - winds[0], - winds[1], - winds[2], - parameters.bg_weights, - parameters.u_back, - parameters.v_back, - parameters.Cb, - ) - - if parameters.Cv > 0: - grad += _cost_functions_numpy.calculate_vertical_vorticity_gradient( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - parameters.Ut, - parameters.Vt, - coeff=parameters.Cv, - upper_bc=parameters.upper_bc, - ) - - if parameters.Cmod > 0: - grad += _cost_functions_numpy.calculate_model_gradient( - winds[0], - winds[1], - winds[2], - parameters.model_weights, - parameters.u_model, - parameters.v_model, - parameters.w_model, - coeff=parameters.Cmod, - ) - - if parameters.Cpoint > 0: - grad += _cost_functions_numpy.calculate_point_gradient( - winds[0], - winds[1], - parameters.x, - parameters.y, - parameters.z, - parameters.point_list, - Cp=parameters.Cpoint, - roi=parameters.roi, - upper_bc=parameters.upper_bc, - ) - - # Let's see if we need to enforce strong boundary conditions - if parameters.const_boundary_cond is True: - grad = np.reshape( - grad, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - grad[:, :, 0, :] = 0 - grad[:, :, -1, :] = 0 - grad[:, :, :, 0] = 0 - grad[:, :, :, -1] = 0 - grad = grad.flatten() - elif parameters.engine == "jax": - grad = grad_jax(winds, parameters) - if parameters.const_boundary_cond is True: - grad = jnp.reshape( - grad, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - grad.at[:, :, 0, :].set(0) - grad.at[:, :, -1, :].set(0) - grad.at[:, :, :, 0].set(0) - grad.at[:, :, :, -1].set(0) - grad = grad.flatten() - return grad - - if parameters.Nfeval % 10 == 0: - print("The gradient of the cost functions is", str(np.linalg.norm(grad, 2))) - return grad - - -def J_function_jax(winds, parameters): - if not JAX_AVAILABLE: - raise ImportError("Jax is needed in order to use the Jax-based PyDDA!") - - winds = jnp.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - # Had to change to float because Jax returns device array (use np.float_()) - Jvel = _cost_functions_jax.calculate_radial_vel_cost_function( - parameters.vrs, - parameters.azs, - parameters.els, - winds[0], - winds[1], - winds[2], - parameters.wts, - rmsVr=parameters.rmsVr, - weights=parameters.weights, - coeff=parameters.Co, - ) - - if parameters.Cm > 0: - # Had to change to float because Jax returns device array (use np.float_()) - Jmass = _cost_functions_jax.calculate_mass_continuity( - winds[0], - winds[1], - winds[2], - parameters.z, - parameters.dx, - parameters.dy, - parameters.dz, - coeff=parameters.Cm, - ) - else: - Jmass = 0 - - if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: - Jsmooth = _cost_functions_jax.calculate_smoothness_cost( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - Cx=parameters.Cx, - Cy=parameters.Cy, - Cz=parameters.Cz, - ) - else: - Jsmooth = 0 - - if parameters.Cb > 0: - Jbackground = _cost_functions_jax.calculate_background_cost( - winds[0], - winds[1], - winds[2], - parameters.bg_weights, - parameters.u_back, - parameters.v_back, - parameters.Cb, - ) - else: - Jbackground = 0 - - if parameters.Cv > 0: - # Had to change to float because Jax returns device array (use np.float_()) - Jvorticity = _cost_functions_jax.calculate_vertical_vorticity_cost( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - parameters.Ut, - parameters.Vt, - coeff=parameters.Cv, - ) - else: - Jvorticity = 0 - - if parameters.Cmod > 0: - Jmod = _cost_functions_jax.calculate_model_cost( - winds[0], - winds[1], - winds[2], - parameters.model_weights, - parameters.u_model, - parameters.v_model, - parameters.w_model, - coeff=parameters.Cmod, - ) - else: - Jmod = 0 - - if parameters.Cpoint > 0: - Jpoint = _cost_functions_jax.calculate_point_cost( - winds[0], - winds[1], - parameters.x, - parameters.y, - parameters.z, - parameters.point_list, - Cp=parameters.Cpoint, - roi=parameters.roi, - ) - else: - Jpoint = 0 - - return Jvel + Jsmooth + Jmass + Jmod + Jpoint + Jvorticity + Jbackground - - -def grad_jax(winds, parameters): - winds = jnp.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - grad = _cost_functions_jax.calculate_grad_radial_vel( - parameters.vrs, - parameters.els, - parameters.azs, - winds[0], - winds[1], - winds[2], - parameters.wts, - parameters.weights, - parameters.rmsVr, - coeff=parameters.Co, - upper_bc=parameters.upper_bc, - ) - - if parameters.Cm > 0: - grad += _cost_functions_jax.calculate_mass_continuity_gradient( - winds[0], - winds[1], - winds[2], - parameters.z, - parameters.dx, - parameters.dy, - parameters.dz, - coeff=parameters.Cm, - upper_bc=parameters.upper_bc, - ) - - if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: - grad += _cost_functions_jax.calculate_smoothness_gradient( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - Cx=parameters.Cx, - Cy=parameters.Cy, - Cz=parameters.Cz, - upper_bc=parameters.upper_bc, - ) - - if parameters.Cb > 0: - grad += _cost_functions_jax.calculate_background_gradient( - winds[0], - winds[1], - winds[2], - parameters.bg_weights, - parameters.u_back, - parameters.v_back, - parameters.Cb, - ) - - if parameters.Cv > 0: - grad += _cost_functions_jax.calculate_vertical_vorticity_gradient( - winds[0], - winds[1], - winds[2], - parameters.dx, - parameters.dy, - parameters.dz, - parameters.Ut, - parameters.Vt, - coeff=parameters.Cv, - upper_bc=parameters.upper_bc, - ).numpy() - - if parameters.Cmod > 0: - grad += _cost_functions_jax.calculate_model_gradient( - winds[0], - winds[1], - winds[2], - parameters.model_weights, - parameters.u_model, - parameters.v_model, - parameters.w_model, - coeff=parameters.Cmod, - ) - - if parameters.Cpoint > 0: - grad += _cost_functions_jax.calculate_point_gradient( - winds[0], - winds[1], - parameters.x, - parameters.y, - parameters.z, - parameters.point_list, - Cp=parameters.Cpoint, - roi=parameters.roi, - ) - return grad - - -def calculate_fall_speed(grid, refl_field=None, frz=4500.0): - """ - Estimates fall speed based on reflectivity. - - Uses methodology of Mike Biggerstaff and Dan Betten - - Parameters - ---------- - Grid: Py-ART Grid - Py-ART Grid containing reflectivity to calculate fall speed from - refl_field: str - String containing name of reflectivity field. None will automatically - determine the name. - frz: float - Height of freezing level in m - - Returns - ------- - 3D float array: - Float array of terminal velocities - - """ - # Parse names of velocity field - if refl_field is None: - refl_field = pyart.config.get_field_name("reflectivity") - - refl = grid[refl_field].values - grid_z = grid["point_z"].values - np.zeros(refl.shape) - A = np.zeros(refl.shape) - B = np.zeros(refl.shape) - rho = np.exp(-grid_z / 10000.0) - A[np.logical_and(grid_z < frz, refl < 55)] = -2.6 - B[np.logical_and(grid_z < frz, refl < 55)] = 0.0107 - A[np.logical_and(grid_z < frz, np.logical_and(refl >= 55, refl < 60))] = -2.5 - B[np.logical_and(grid_z < frz, np.logical_and(refl >= 55, refl < 60))] = 0.013 - A[np.logical_and(grid_z < frz, refl > 60)] = -3.95 - B[np.logical_and(grid_z < frz, refl > 60)] = 0.0148 - A[np.logical_and(grid_z >= frz, refl < 33)] = -0.817 - B[np.logical_and(grid_z >= frz, refl < 33)] = 0.0063 - A[np.logical_and(grid_z >= frz, np.logical_and(refl >= 33, refl < 49))] = -2.5 - B[np.logical_and(grid_z >= frz, np.logical_and(refl >= 33, refl < 49))] = 0.013 - A[np.logical_and(grid_z >= frz, refl > 49)] = -3.95 - B[np.logical_and(grid_z >= frz, refl > 49)] = 0.0148 - - fallspeed = A * np.power(10, refl * B) * np.power(1.2 / rho, 0.4) - print(fallspeed.max()) - del A, B, rho - return np.ma.masked_invalid(fallspeed) diff --git a/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py b/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py deleted file mode 100644 index e757a8df..00000000 --- a/pydda/retrieval/.ipynb_checkpoints/wind_retrieve-checkpoint.py +++ /dev/null @@ -1,1598 +0,0 @@ -""" -Created on Mon Aug 7 09:17:40 2017 - -@author: rjackson -""" - -import pyart -import numpy as np -import time -import math -import xarray as xr - -from scipy.interpolate import interp1d -from scipy.ndimage import convolve1d -from scipy.optimize import fmin_l_bfgs_b -from scipy.signal import savgol_filter -from .auglag import auglag -from ..io import read_from_pyart_grid - -_LEISE_KERNEL = np.array([-1 / 16, 1 / 4, 5 / 8, 1 / 4, -1 / 16]) - - -def _apply_low_pass_filter( - winds, filter_type, filter_window, filter_order, leise_nstep -): - """Smooth the (3, nz, ny, nx) wind array in place along all spatial axes.""" - if filter_type not in ("savgol", "leise"): - raise ValueError( - "filter_type must be 'savgol' or 'leise', got %r" % filter_type - ) - for c in range(3): - for axis in range(winds[c].ndim): - if filter_type == "savgol": - winds[c] = savgol_filter( - winds[c], filter_window, filter_order, axis=axis - ) - else: - if winds[c].shape[axis] < 5: - continue - for _ in range(leise_nstep): - winds[c] = convolve1d( - winds[c], _LEISE_KERNEL, axis=axis, mode="mirror" - ) - return winds - - -try: - import tensorflow_probability as tfp - import tensorflow as tf - - TENSORFLOW_AVAILABLE = True -except (ImportError, AttributeError): - TENSORFLOW_AVAILABLE = False - -try: - import jax.numpy as jnp - import jax - import jaxopt - - JAX_AVAILABLE = True -except ImportError: - JAX_AVAILABLE = False - -# imports changed to local import path to run on computer -from ..cost_functions import ( - J_function, - grad_J, - calculate_fall_speed, - grad_jax, - J_function_jax, -) -from copy import deepcopy -from .angles import add_azimuth_as_field, add_elevation_as_field - -_wprevmax = np.empty(0) -_wcurrmax = np.empty(0) -iterations = 0 - - -class DDParameters(object): - """ - This is a helper class for inserting more arguments into the :func:`pydda.cost_functions.J_function` and - :func:`pydda.cost_functions.grad_J` function. Since these cost functions take numerous parameters, this class - will store the needed parameters as one positional argument for easier readability of the code. - - In addition, class members can be added here so that those contributing more constraints to the variational - framework can add any parameters they may need. - - Attributes - ---------- - vrs: List of float arrays - List of radial velocities from each radar - azs: List of float arrays - List of azimuths from each radar - els: List of float arrays - List of elevations from each radar - wts: List of float arrays - Float array containing fall speed from radar. - u_back: 1D float array (number of vertical levels) - Background u wind - v_back: 1D float array (number of vertical levels) - Background v wind - u_model: list of 3D float arrays - U from each model integrated into the retrieval - v_model: list of 3D float arrays - V from each model integrated into the retrieval - w_model: - W from each model integrated into the retrieval - Co: float - Weighting coefficient for data constraint. - Cm: float - Weighting coefficient for mass continuity constraint. - Cx: float - Smoothing coefficient for x-direction - Cy: float - Smoothing coefficient for y-direction - Cz: float - Smoothing coefficient for z-direction - Cb: float - Coefficient for sounding constraint - Cv: float - Weight for cost function related to vertical vorticity equation. - Cmod: float - Coefficient for model constraint - Cpoint: float - Coefficient for point constraint - Ut: float - Prescribed storm motion. This is only needed if Cv is not zero. - Vt: float - Prescribed storm motion. This is only needed if Cv is not zero. - grid_shape: - Shape of wind grid - dx: - Spacing of grid in x direction - dy: - Spacing of grid in y direction - dz: - Spacing of grid in z direction - x: - E-W grid levels in m - y: - N-S grid levels in m - z: - Grid vertical levels in m - rmsVr: float - The sum of squares of velocity/num_points. Use for normalization - of data weighting coefficient - weights: n_radars by z_bins by y_bins x x_bins float array - Data weights for each pair of radars - bg_weights: z_bins by y_bins x x_bins float array - Data weights for sounding constraint - model_weights: n_models by z_bins by y_bins by x_bins float array - Data weights for each model. - point_list: list or None - point_list: list of dicts - List of point constraints. Each member is a dict with keys of "u", "v", - to correspond to each component of the wind field and "x", "y", "z" - to correspond to the location of the point observation in the Grid's - Cartesian coordinates. - roi: float - The radius of influence of each point observation in m. - above: float - The altitude below which the cloud top impermeability would not apply (minimum impermeable height) - upper_bc: int - 0 to not enforce impermeability at top of domain - 1 to enforce w=0 at top of domain (impermeability condition), - 2 to enforce w=0 at cloud top (or "above" if higher) (impermeability condition), - """ - - def __init__(self): - self.Ut = np.nan - self.Vt = np.nan - self.rmsVr = np.nan - self.grid_shape = None - self.Cmod = np.nan - self.Cpoint = np.nan - self.u_back = None - self.v_back = None - self.wts = [] - self.vrs = [] - self.azs = [] - self.els = [] - self.weights = [] - self.bg_weights = [] - self.model_weights = [] - self.u_model = [] - self.v_model = [] - self.w_model = [] - self.x = None - self.y = None - self.z = None - self.dx = np.nan - self.dy = np.nan - self.dz = np.nan - self.Co = 1.0 - self.Cm = 1500.0 - self.Cx = 0.0 - self.Cy = 0.0 - self.Cz = 0.0 - self.Cb = 0.0 - self.Cv = 0.0 - self.Cmod = 0.0 - self.Cpoint = 0.0 - self.Ut = 0.0 - self.Vt = 0.0 - self.upper_bc = True - self.lower_bc = True - self.roi = 1000.0 - self.above = 2.0 - self.frz = 4500.0 - self.Nfeval = 0.0 - self.engine = "scipy" - self.point_list = [] - self.cvtol = 1e-2 - self.gtol = 1e-2 - self.Jveltol = 100.0 - self.const_boundary_cond = False - self.parallel = False - - -def _get_dd_wind_field_scipy( - Grids, - u_init, - v_init, - w_init, - engine, - points=None, - vel_name=None, - refl_field=None, - u_back=None, - v_back=None, - z_back=None, - frz=4500.0, - Co=1.0, - Cm=1500.0, - Cx=0.0, - Cy=0.0, - Cz=0.0, - Cb=0.0, - Cv=0.0, - Cmod=0.0, - Cpoint=0.0, - cvtol=1e-2, - gtol=1e-2, - Jveltol=100.0, - Ut=None, - Vt=None, - low_pass_filter=True, - mask_outside_opt=False, - weights_obs=None, - weights_model=None, - weights_bg=None, - max_iterations=1000, - mask_w_outside_opt=True, - filter_type="savgol", - filter_window=5, - filter_order=3, - leise_nstep=1, - min_bca=30.0, - max_bca=150.0, - upper_bc=1, - model_fields=None, - output_cost_functions=True, - roi=1000.0, - above=2.0, - wind_tol=0.1, - tolerance=1e-8, - const_boundary_cond=False, - max_wind_mag=100.0, - parallel=True, -): - global _wcurrmax - global _wprevmax - global iterations - - # We have to have a prescribed storm motion for vorticity constraint - if Ut is None or Vt is None: - if Cv != 0.0: - raise ValueError( - ( - "Ut and Vt cannot be None if vertical " - + "vorticity constraint is enabled!" - ) - ) - - if not isinstance(Grids, list): - raise ValueError("Grids has to be a list!") - - parameters = DDParameters() - parameters.Ut = Ut - parameters.Vt = Vt - parameters.engine = engine - parameters.const_boundary_cond = const_boundary_cond - print(parameters.const_boundary_cond) - # Ensure that all Grids are on the same coordinate system - prev_grid = Grids[0] - for g in Grids: - if not np.allclose(g["x"].values, prev_grid["x"].values, atol=10): - raise ValueError("Grids do not have equal x coordinates!") - - if not np.allclose(g["y"].values, prev_grid["y"].values, atol=10): - raise ValueError("Grids do not have equal y coordinates!") - - if not np.allclose(g["z"].values, prev_grid["z"].values, atol=10): - raise ValueError("Grids do not have equal z coordinates!") - - if not np.allclose( - g["origin_latitude"].values, prev_grid["origin_latitude"].values - ): - raise ValueError(("Grids have unequal origin lat/lons!")) - - prev_grid = g - - if engine.lower() == "auglag" and not TENSORFLOW_AVAILABLE: - raise ModuleNotFoundError( - "Tensorflow 2.6+ needs to be installed for the Augmented Lagrangian solver." - ) - - # Disable background constraint if none provided - if u_back is None or v_back is None: - parameters.u_back = np.zeros(u_init.shape[0]) - parameters.v_back = np.zeros(v_init.shape[0]) - else: - # Interpolate sounding to radar grid - print("Interpolating sounding to radar grid") - - if isinstance(u_back, np.ma.MaskedArray): - u_back = u_back.filled(-9999.0) - if isinstance(v_back, np.ma.MaskedArray): - v_back = v_back.filled(-9999.0) - if isinstance(z_back, np.ma.MaskedArray): - z_back = z_back.filled(-9999.0) - valid_inds = np.logical_and.reduce( - (u_back > -9998, v_back > -9998, z_back > -9998) - ) - u_interp = interp1d(z_back[valid_inds], u_back[valid_inds], bounds_error=False) - v_interp = interp1d(z_back[valid_inds], v_back[valid_inds], bounds_error=False) - if isinstance(Grids[0]["z"].values, np.ma.MaskedArray): - parameters.u_back = u_interp(Grids[0]["z"].values.filled(np.nan)) - parameters.v_back = v_interp(Grids[0]["z"].values.filled(np.nan)) - else: - parameters.u_back = u_interp(Grids[0]["z"].values) - parameters.v_back = v_interp(Grids[0]["z"].values) - - print("Grid levels:") - print(Grids[0]["z"].values) - - # Parse names of velocity field - if refl_field is None: - refl_field = pyart.config.get_field_name("reflectivity") - - # Parse names of velocity field - if vel_name is None: - vel_name = pyart.config.get_field_name("corrected_velocity") - winds = np.stack([u_init, v_init, w_init]) - - # Set up wind fields and weights from each radar - parameters.weights = np.zeros( - (len(Grids), u_init.shape[0], u_init.shape[1], u_init.shape[2]) - ) - - parameters.bg_weights = np.zeros(v_init.shape) - if model_fields is not None: - parameters.model_weights = np.ones( - (len(model_fields), u_init.shape[0], u_init.shape[1], u_init.shape[2]) - ) - else: - parameters.model_weights = np.zeros( - (1, u_init.shape[0], u_init.shape[1], u_init.shape[2]) - ) - - if model_fields is None: - if Cmod != 0.0: - raise ValueError("Cmod must be zero if model fields are not specified!") - - bca = np.zeros((len(Grids), len(Grids), u_init.shape[1], u_init.shape[2])) - sum_Vr = np.zeros(len(Grids)) - - for i in range(len(Grids)): - parameters.wts.append( - np.ma.masked_invalid( - calculate_fall_speed(Grids[i], refl_field=refl_field, frz=frz).squeeze() - ) - ) - - parameters.vrs.append(np.ma.masked_invalid(Grids[i][vel_name].values.squeeze())) - parameters.azs.append( - np.ma.masked_invalid(Grids[i]["AZ"].values.squeeze() * np.pi / 180) - ) - parameters.els.append( - np.ma.masked_invalid(Grids[i]["EL"].values.squeeze() * np.pi / 180) - ) - - if len(Grids) > 1: - for i in range(len(Grids)): - for j in range(len(Grids)): - if i == j: - continue - print(("Calculating weights for radars " + str(i) + " and " + str(j))) - bca[i, j] = get_bca(Grids[i], Grids[j]) - - for k in range(parameters.vrs[i].shape[0]): - if weights_obs is None: - valid = np.logical_and.reduce( - ( - ~parameters.vrs[i][k].mask, - ~parameters.wts[i][k].mask, - ~parameters.azs[i][k].mask, - ~parameters.els[i][k].mask, - ) - ) - valid = np.logical_and.reduce( - ( - valid, - np.isfinite(parameters.vrs[i][k]), - np.isfinite(parameters.wts[i][k]), - np.isfinite(parameters.azs[i][k]), - np.isfinite(parameters.els[i][k]), - ) - ) - valid = np.logical_and.reduce( - ( - valid, - np.isfinite(parameters.vrs[j][k]), - np.isfinite(parameters.wts[j][k]), - np.isfinite(parameters.azs[j][k]), - np.isfinite(parameters.els[j][k]), - ) - ) - valid = np.logical_and.reduce( - ( - valid, - ~parameters.vrs[j][k].mask, - ~parameters.wts[j][k].mask, - ~parameters.azs[j][k].mask, - ~parameters.els[j][k].mask, - ) - ) - cur_array = parameters.weights[i, k].copy() - cur_array[ - np.logical_and( - valid, - np.logical_and( - bca[i, j] >= math.radians(min_bca), - bca[i, j] <= math.radians(max_bca), - ), - ) - ] = 1 - cur_array[~valid] = 0 - parameters.weights[i, k] += cur_array - else: - parameters.weights[i, k] = weights_obs[i][k, :, :] - - if weights_bg is None: - valid = np.logical_and.reduce( - ( - ~parameters.vrs[j][k].mask, - ~parameters.wts[j][k].mask, - ~parameters.azs[j][k].mask, - ~parameters.els[j][k].mask, - ) - ) - valid = np.logical_and.reduce( - ( - valid, - np.isfinite(parameters.vrs[j][k]), - np.isfinite(parameters.wts[j][k]), - np.isfinite(parameters.azs[j][k]), - np.isfinite(parameters.els[j][k]), - ) - ) - valid = np.logical_and.reduce( - ( - valid, - np.isfinite(parameters.vrs[j][k]), - np.isfinite(parameters.wts[j][k]), - np.isfinite(parameters.azs[j][k]), - np.isfinite(parameters.els[j][k]), - ) - ) - valid = np.logical_and.reduce( - ( - valid, - ~parameters.vrs[j][k].mask, - ~parameters.wts[j][k].mask, - ~parameters.azs[j][k].mask, - ~parameters.els[j][k].mask, - ) - ) - cur_array = parameters.bg_weights[k] - cur_array[ - np.logical_or.reduce( - ( - ~valid, - bca[i, j] < math.radians(min_bca), - bca[i, j] > math.radians(max_bca), - ) - ) - ] = 1 - cur_array[~valid] = 1 - parameters.bg_weights[i] += cur_array - else: - parameters.bg_weights[i] = weights_bg[i] - - print("Calculating weights for models...") - coverage_grade = parameters.weights.sum(axis=0) - coverage_grade = coverage_grade / coverage_grade.max() - - # Weigh in model input more when we have no coverage - # Model only weighs 1/(# of grids + 1) when there is full - # Coverage - if model_fields is not None: - if weights_model is None: - for i in range(len(model_fields)): - parameters.model_weights[i] = 1 - ( - coverage_grade / (len(Grids) + 1) - ) - else: - for i in range(len(model_fields)): - parameters.model_weights[i] = weights_model[i] - else: - if weights_obs is None: - parameters.weights[0] = np.where(~parameters.vrs[0].mask, 1, 0) - else: - parameters.weights[0] = weights_obs[0] - - if weights_bg is None: - parameters.bg_weights = np.where(~parameters.vrs[0].mask, 0, 1) - else: - parameters.bg_weights = weights_bg - - parameters.vrs = [x.filled(-9999.0) for x in parameters.vrs] - parameters.azs = [x.filled(-9999.0) for x in parameters.azs] - parameters.els = [x.filled(-9999.0) for x in parameters.els] - parameters.wts = [x.filled(-9999.0) for x in parameters.wts] - parameters.weights[~np.isfinite(parameters.weights)] = 0 - parameters.bg_weights[~np.isfinite(parameters.bg_weights)] = 0 - parameters.weights[parameters.weights > 0] = 1 - parameters.bg_weights[parameters.bg_weights > 0] = 1 - - # Zero out bg_weights at height levels where the interpolated background - # is NaN (i.e. outside the sounding's vertical range). Also replace NaN - # in u_back/v_back with 0 so those levels don't corrupt cost function - # arithmetic even though they carry zero weight. - nan_bg_levels = ~np.isfinite(parameters.u_back) | ~np.isfinite(parameters.v_back) - parameters.bg_weights[nan_bg_levels] = 0 - parameters.u_back = np.nan_to_num(parameters.u_back) - parameters.v_back = np.nan_to_num(parameters.v_back) - sum_Vr = np.nansum(np.square(parameters.vrs * parameters.weights)) - parameters.rmsVr = np.sqrt(np.nansum(sum_Vr) / np.nansum(parameters.weights)) - - del bca - parameters.grid_shape = u_init.shape - # Parse names of velocity field - - winds = winds.flatten() - - print("Starting solver ") - parameters.dx = np.diff(Grids[0]["x"].values, axis=0)[0] - parameters.dy = np.diff(Grids[0]["y"].values, axis=0)[0] - parameters.dz = np.diff(Grids[0]["z"].values, axis=0)[0] - print("rmsVR = " + str(parameters.rmsVr)) - print("Total points: %d" % parameters.weights.sum()) - parameters.z = Grids[0]["point_z"].values - parameters.x = Grids[0]["point_x"].values - parameters.y = Grids[0]["point_y"].values - bt = time.time() - - # First pass - no filter - wcurrmax = w_init.max() - print("The max of w_init is", wcurrmax) - iterations = 0 - bounds = [(-x, x) for x in max_wind_mag * np.ones(winds.shape)] - - if model_fields is not None: - for i, the_field in enumerate(model_fields): - u_field = "U_" + the_field - v_field = "V_" + the_field - w_field = "W_" + the_field - parameters.u_model.append(np.nan_to_num(Grids[0][u_field].values.squeeze())) - parameters.v_model.append(np.nan_to_num(Grids[0][v_field].values.squeeze())) - parameters.w_model.append(np.nan_to_num(Grids[0][w_field].values.squeeze())) - - # Don't weigh in where model data unavailable - where_finite_u = np.isfinite(Grids[0][u_field].values.squeeze()) - where_finite_v = np.isfinite(Grids[0][v_field].values.squeeze()) - where_finite_w = np.isfinite(Grids[0][w_field].values.squeeze()) - parameters.model_weights[i, :, :, :] = np.where( - np.logical_and.reduce((where_finite_u, where_finite_v, where_finite_w)), - 1, - 0, - ) - - print("Total number of model points: %d" % np.sum(parameters.model_weights)) - parameters.Co = Co - parameters.Cm = Cm - parameters.Cx = Cx - parameters.Cy = Cy - parameters.Cz = Cz - parameters.Cb = Cb - parameters.Cv = Cv - parameters.Cmod = Cmod - parameters.Cpoint = Cpoint - parameters.roi = roi - parameters.upper_bc = upper_bc - parameters.points = points - parameters.point_list = points - parameters.parallel = parallel - _wprevmax = np.zeros(parameters.grid_shape) - _wcurrmax = np.zeros(parameters.grid_shape) - iterations = 0 - if engine.lower() == "scipy" or engine.lower() == "jax": - - def _vert_velocity_callback(x): - global _wprevmax - global _wcurrmax - global iterations - - if iterations % 10 > 0: - iterations = iterations + 1 - return False - - wind = np.reshape( - x, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - _wcurrmax = wind[2] - if iterations == 0: - _wprevmax = _wcurrmax - iterations = iterations + 1 - return False - diff = np.abs(_wprevmax - _wcurrmax) - diff = np.where(parameters.bg_weights == 0, diff, np.nan) - delta = np.nanmax(diff) - if delta < wind_tol: - return True - _wprevmax = _wcurrmax - iterations = iterations + 1 - print("Max change in w: %4.3f" % delta) - return False - - parameters.print_out = False - if engine.lower() == "scipy": - winds = fmin_l_bfgs_b( - J_function, - winds, - args=(parameters,), - maxiter=max_iterations, - pgtol=tolerance, - bounds=bounds, - fprime=grad_J, - callback=_vert_velocity_callback, - ) - else: - - def loss_and_gradient(x): - x_loss = J_function_jax(x["winds"], parameters) - x_grad = {} - x_grad["winds"] = grad_jax(x["winds"], parameters) - return x_loss, x_grad - - bounds = ( - {"winds": -max_wind_mag * jnp.ones(winds.shape)}, - {"winds": max_wind_mag * jnp.ones(winds.shape)}, - ) - winds = jnp.array(winds) - # JIT-compile the cost function explicitly so the compilation - # delay is isolated and visible before the solver loop starts. - loss_and_gradient = jax.jit(loss_and_gradient) - print("Compiling JAX cost functions...") - loss_and_gradient({"winds": winds}) - print("Compilation complete.") - solver = jaxopt.LBFGSB( - loss_and_gradient, - True, - has_aux=False, - maxiter=max_iterations, - tol=tolerance, - jit=False, - implicit_diff=False, - verbose=True, - ) - winds = {"winds": winds} - winds, state = solver.run(winds, bounds=bounds) - winds = [np.asanyarray(winds["winds"])] - - winds = np.reshape( - winds[0], - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - parameters.print_out = True - - elif engine.lower() == "auglag": - if not TENSORFLOW_AVAILABLE: - raise ImportError( - "Tensorflow must be available to use the Augmented Lagrangian engine!" - ) - parameters.vrs = [tf.constant(x, dtype=tf.float32) for x in parameters.vrs] - parameters.azs = [tf.constant(x, dtype=tf.float32) for x in parameters.azs] - parameters.els = [tf.constant(x, dtype=tf.float32) for x in parameters.els] - parameters.wts = [tf.constant(x, dtype=tf.float32) for x in parameters.wts] - parameters.model_weights = tf.constant( - parameters.model_weights, dtype=tf.float32 - ) - parameters.weights[~np.isfinite(parameters.weights)] = 0 - parameters.weights[parameters.weights > 0] = 1 - parameters.weights = tf.constant(parameters.weights, dtype=tf.float32) - parameters.bg_weights[parameters.bg_weights > 0] = 1 - parameters.bg_weights = tf.constant(parameters.bg_weights, dtype=tf.float32) - parameters.z = tf.constant(Grids[0]["point_z"].values, dtype=tf.float32) - parameters.x = tf.constant(Grids[0]["point_x"].values, dtype=tf.float32) - parameters.y = tf.constant(Grids[0]["point_y"].values, dtype=tf.float32) - bounds = [(-x, x) for x in max_wind_mag * np.ones(winds.shape, dtype="float32")] - winds = winds.astype("float32") - winds, mult, AL_Filter, funcalls = auglag(winds, parameters, bounds) - - # """ - winds = np.stack([winds[0], winds[1], winds[2]]) - winds = winds.flatten() - if low_pass_filter is True: - print("Applying %s low pass filter to wind field..." % filter_type) - winds = np.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - winds = _apply_low_pass_filter( - winds, filter_type, filter_window, filter_order, leise_nstep - ) - winds = np.stack([winds[0], winds[1], winds[2]]) - winds = winds.flatten() - - print("Done! Time = " + "{:2.1f}".format(time.time() - bt)) - - # First pass - no filter - the_winds = np.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - u = the_winds[0] - v = the_winds[1] - w = the_winds[2] - where_mask = np.sum(parameters.weights, axis=0) + np.sum( - parameters.model_weights, axis=0 - ) - - u = np.ma.array(u) - w = np.ma.array(w) - v = np.ma.array(v) - - if mask_outside_opt is True: - u = np.ma.masked_where(where_mask < 1, u) - v = np.ma.masked_where(where_mask < 1, v) - w = np.ma.masked_where(where_mask < 1, w) - - if mask_w_outside_opt is True: - w = np.ma.masked_where(where_mask < 1, w) - - u_field = {} - u_field["standard_name"] = "u_wind" - u_field["long_name"] = "zonal component of wind velocity" - u_field["units"] = "m/s" - u_field["min_bca"] = min_bca - u_field["max_bca"] = max_bca - v_field = {} - v_field["standard_name"] = "v_wind" - v_field["long_name"] = "meridional component of wind velocity" - v_field["units"] = "m/s" - v_field["min_bca"] = min_bca - v_field["max_bca"] = max_bca - w_field = {} - w_field["standard_name"] = "w_wind" - w_field["long_name"] = "vertical component of wind velocity" - w_field["units"] = "m/s" - w_field["min_bca"] = min_bca - w_field["max_bca"] = max_bca - - new_grid_list = [] - - for grid in Grids: - grid["u"] = xr.DataArray( - np.expand_dims(u, 0), dims=("time", "z", "y", "x"), attrs=u_field - ) - grid["v"] = xr.DataArray( - np.expand_dims(v, 0), dims=("time", "z", "y", "x"), attrs=v_field - ) - grid["w"] = xr.DataArray( - np.expand_dims(w, 0), dims=("time", "z", "y", "x"), attrs=w_field - ) - new_grid_list.append(grid) - - return new_grid_list, parameters - - -def _get_dd_wind_field_tensorflow( - Grids, - u_init, - v_init, - w_init, - points=None, - vel_name=None, - refl_field=None, - u_back=None, - v_back=None, - z_back=None, - frz=4500.0, - Co=1.0, - Cm=1500.0, - Cx=0.0, - Cy=0.0, - Cz=0.0, - Cb=0.0, - Cv=0.0, - Cmod=0.0, - Cpoint=0.0, - Ut=None, - Vt=None, - low_pass_filter=True, - mask_outside_opt=False, - weights_obs=None, - weights_model=None, - weights_bg=None, - max_iterations=200, - mask_w_outside_opt=True, - filter_type="savgol", - filter_window=5, - filter_order=3, - leise_nstep=1, - min_bca=30.0, - max_bca=150.0, - upper_bc=True, - model_fields=None, - output_cost_functions=True, - roi=1000.0, - lower_bc=True, - parallel_iterations=1, - wind_tol=0.1, - tolerance=1e-8, - const_boundary_cond=False, - max_wind_mag=100.0, -): - if not TENSORFLOW_AVAILABLE: - raise ImportError( - "Tensorflow >=2.5 and tensorflow-probability " - + "need to be installed in order to use the tensorflow engine." - ) - # We have to have a prescribed storm motion for vorticity constraint - if Ut is None or Vt is None: - if Cv != 0.0: - raise ValueError( - ( - "Ut and Vt cannot be None if vertical " - + "vorticity constraint is enabled!" - ) - ) - - if not isinstance(Grids, list): - raise ValueError("Grids has to be a list!") - - parameters = DDParameters() - parameters.Ut = Ut - parameters.Vt = Vt - parameters.upper_bc = upper_bc - parameters.lower_bc = lower_bc - parameters.engine = "tensorflow" - parameters.const_boundary_cond = const_boundary_cond - - # Ensure that all Grids are on the same coordinate system - prev_grid = Grids[0] - for g in Grids: - if not np.allclose(g["x"].values, prev_grid["x"].values, atol=10): - raise ValueError("Grids do not have equal x coordinates!") - - if not np.allclose(g["y"].values, prev_grid["y"].values, atol=10): - raise ValueError("Grids do not have equal y coordinates!") - - if not np.allclose(g["z"].values, prev_grid["z"].values, atol=10): - raise ValueError("Grids do not have equal z coordinates!") - - if not np.allclose( - g["origin_latitude"].values, prev_grid["origin_latitude"].values - ): - raise ValueError(("Grids have unequal origin lat/lons!")) - - prev_grid = g - - # Disable background constraint if none provided - if u_back is None or v_back is None: - parameters.u_back = tf.zeros(u_init.shape[0]) - parameters.v_back = tf.zeros(v_init.shape[0]) - else: - # Interpolate sounding to radar grid - print("Interpolating sounding to radar grid") - - if isinstance(u_back, np.ma.MaskedArray): - u_back = u_back.filled(-9999.0) - if isinstance(v_back, np.ma.MaskedArray): - v_back = v_back.filled(-9999.0) - if isinstance(z_back, np.ma.MaskedArray): - z_back = z_back.filled(-9999.0) - valid_inds = np.logical_and.reduce( - (u_back > -9998, v_back > -9998, z_back > -9998) - ) - u_interp = interp1d(z_back[valid_inds], u_back[valid_inds], bounds_error=False) - v_interp = interp1d(z_back[valid_inds], v_back[valid_inds], bounds_error=False) - if isinstance(Grids[0]["z"].values, np.ma.MaskedArray): - parameters.u_back = tf.constant( - u_interp(Grids[0]["z"].values.filled(np.nan)), dtype=tf.float32 - ) - parameters.v_back = tf.constant( - v_interp(Grids[0]["z"].values.filled(np.nan)), dtype=tf.float32 - ) - else: - parameters.u_back = tf.constant( - u_interp(Grids[0]["z"].values), dtype=tf.float32 - ) - parameters.v_back = tf.constant( - v_interp(Grids[0]["z"].values), dtype=tf.float32 - ) - - print("Interpolated U field:") - print(parameters.u_back) - print("Interpolated V field:") - print(parameters.v_back) - print("Grid levels:") - print(Grids[0]["z"].values) - - # Parse names of velocity field - if refl_field is None: - refl_field = pyart.config.get_field_name("reflectivity") - - # Parse names of velocity field - if vel_name is None: - vel_name = pyart.config.get_field_name("corrected_velocity") - winds = np.stack([u_init, v_init, w_init]) - winds = winds.astype(np.float32) - - # Set up wind fields and weights from each radar - parameters.weights = np.zeros( - (len(Grids), u_init.shape[0], u_init.shape[1], u_init.shape[2]), - dtype=np.float32, - ) - - parameters.bg_weights = np.zeros(v_init.shape) - if model_fields is not None: - parameters.model_weights = np.ones( - (len(model_fields), u_init.shape[0], u_init.shape[1], u_init.shape[2]), - dtype=np.float32, - ) - else: - parameters.model_weights = np.zeros( - (1, u_init.shape[0], u_init.shape[1], u_init.shape[2]), dtype=np.float32 - ) - - if model_fields is None: - if Cmod != 0.0: - raise ValueError("Cmod must be zero if model fields are not specified!") - - bca = np.zeros( - (len(Grids), len(Grids), u_init.shape[1], u_init.shape[2]), dtype=np.float32 - ) - - for i in range(len(Grids)): - parameters.wts.append( - np.ma.masked_invalid( - calculate_fall_speed(Grids[i], refl_field=refl_field, frz=frz).squeeze() - ) - ) - parameters.vrs.append(np.ma.masked_invalid(Grids[i][vel_name].values.squeeze())) - parameters.azs.append( - np.ma.masked_invalid(Grids[i]["AZ"].values.squeeze() * np.pi / 180) - ) - parameters.els.append( - np.ma.masked_invalid(Grids[i]["EL"].values.squeeze() * np.pi / 180) - ) - - if len(Grids) > 1: - for i in range(len(Grids)): - for j in range(len(Grids)): - if i == j: - continue - print(("Calculating weights for radars " + str(i) + " and " + str(j))) - bca[i, j] = get_bca(Grids[i], Grids[j]) - - for k in range(parameters.vrs[i].shape[0]): - if weights_obs is None: - cur_array = parameters.weights[i, k].copy() - valid = np.logical_and.reduce( - ( - ~parameters.vrs[i][k].mask, - ~parameters.wts[i][k].mask, - ~parameters.azs[i][k].mask, - ~parameters.els[i][k].mask, - ) - ) - valid = np.logical_and.reduce( - ( - valid, - np.isfinite(parameters.vrs[i][k]), - np.isfinite(parameters.wts[i][k]), - np.isfinite(parameters.azs[i][k]), - np.isfinite(parameters.els[i][k]), - ) - ) - valid = np.logical_and.reduce( - ( - valid, - np.isfinite(parameters.vrs[j][k]), - np.isfinite(parameters.wts[j][k]), - np.isfinite(parameters.azs[j][k]), - np.isfinite(parameters.els[j][k]), - ) - ) - valid = np.logical_and.reduce( - ( - valid, - ~parameters.vrs[j][k].mask, - ~parameters.wts[j][k].mask, - ~parameters.azs[j][k].mask, - ~parameters.els[j][k].mask, - ) - ) - - cur_array[ - np.logical_and( - valid, - np.logical_and( - bca[i, j] >= math.radians(min_bca), - bca[i, j] <= math.radians(max_bca), - ), - ) - ] = 1 - cur_array[~valid] = 0 - parameters.weights[i, k] += cur_array - else: - parameters.weights[i, k] = weights_obs[i][k, :, :] - - if weights_bg is None: - cur_array = parameters.bg_weights[k] - valid = np.logical_and.reduce( - ( - ~parameters.vrs[i][k].mask, - ~parameters.wts[i][k].mask, - ~parameters.azs[i][k].mask, - ~parameters.els[i][k].mask, - ) - ) - valid = np.logical_and.reduce( - ( - valid, - np.isfinite(parameters.vrs[i][k]), - np.isfinite(parameters.wts[i][k]), - np.isfinite(parameters.azs[i][k]), - np.isfinite(parameters.els[i][k]), - ) - ) - valid = np.logical_and.reduce( - ( - valid, - np.isfinite(parameters.vrs[j][k]), - np.isfinite(parameters.wts[j][k]), - np.isfinite(parameters.azs[j][k]), - np.isfinite(parameters.els[j][k]), - ) - ) - valid = np.logical_and.reduce( - ( - valid, - ~parameters.vrs[j][k].mask, - ~parameters.wts[j][k].mask, - ~parameters.azs[j][k].mask, - ~parameters.els[j][k].mask, - ) - ) - cur_array[ - np.logical_or.reduce( - ( - ~valid, - bca[i, j] < math.radians(min_bca), - bca[i, j] > math.radians(max_bca), - ) - ) - ] = 1 - cur_array[~valid] = 1 - parameters.bg_weights[i] += cur_array - else: - parameters.bg_weights[i] = weights_bg[i] - - print("Calculating weights for models...") - coverage_grade = parameters.weights.sum(axis=0) - coverage_grade = coverage_grade / coverage_grade.max() - - # Weigh in model input more when we have no coverage - # Model only weighs 1/(# of grids + 1) when there is full - # Coverage - if model_fields is not None: - if weights_model is None: - for i in range(len(model_fields)): - parameters.model_weights[i] = 1 - ( - coverage_grade / (len(Grids) + 1) - ) - - else: - for i in range(len(model_fields)): - parameters.model_weights[i] = weights_model[i] - else: - if weights_obs is None: - parameters.weights[0] = np.where(np.isfinite(parameters.vrs[0]), 1, 0) - else: - parameters.weights[0] = weights_obs[0] - - if weights_bg is None: - parameters.bg_weights = np.where(np.isfinite(parameters.vrs[0]), 0, 1) - else: - parameters.bg_weights = weights_bg - - parameters.vrs = [ - tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.vrs - ] - parameters.azs = [ - tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.azs - ] - parameters.els = [ - tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.els - ] - parameters.wts = [ - tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.wts - ] - - parameters.weights[~np.isfinite(parameters.weights)] = 0 - parameters.weights[parameters.weights > 0] = 1 - for i in range(len(Grids)): - print("Points from Radar %d: %d" % (i, parameters.weights[i].sum())) - parameters.weights = tf.constant(parameters.weights, dtype=tf.float32) - parameters.bg_weights[parameters.bg_weights > 0] = 1 - parameters.bg_weights = tf.constant(parameters.bg_weights, dtype=tf.float32) - sum_Vr = tf.experimental.numpy.nansum( - tf.square(parameters.vrs * parameters.weights) - ) - parameters.rmsVr = np.sqrt( - np.nansum(sum_Vr) / tf.experimental.numpy.nansum(parameters.weights) - ) - - del bca - parameters.grid_shape = u_init.shape - # Parse names of velocity field - - winds = winds.flatten() - winds = tf.Variable(winds, name="winds") - - print("Starting solver ") - parameters.dx = np.diff(Grids[0]["x"].values, axis=0)[0] - parameters.dy = np.diff(Grids[0]["y"].values, axis=0)[0] - parameters.dz = np.diff(Grids[0]["z"].values, axis=0)[0] - print("rmsVR = " + str(parameters.rmsVr)) - print("Total points: %d" % tf.reduce_sum(parameters.weights)) - parameters.z = tf.constant(Grids[0]["point_z"].values, dtype=tf.float32) - parameters.x = tf.constant(Grids[0]["point_x"].values, dtype=tf.float32) - parameters.y = tf.constant(Grids[0]["point_y"].values, dtype=tf.float32) - bt = time.time() - - # First pass - no filter - wcurrmax = w_init.max() - print("The max of w_init is", wcurrmax) - [(-x, x) for x in 100.0 * np.ones(winds.shape)] - - if model_fields is not None: - for i, the_field in enumerate(model_fields): - u_field = "U_" + the_field - v_field = "V_" + the_field - w_field = "W_" + the_field - parameters.u_model.append( - tf.constant(np.nan_to_num(Grids[0][u_field].values.squeeze())) - ) - parameters.v_model.append( - tf.constant(np.nan_to_num(Grids[0][v_field].values.squeeze())) - ) - parameters.w_model.append( - tf.constant(np.nan_to_num(Grids[0][w_field].values.squeeze())) - ) - - # Don't weigh in where model data unavailable - where_finite_u = np.isfinite(Grids[0][u_field].values.squeeze()) - where_finite_v = np.isfinite(Grids[0][v_field].values.squeeze()) - where_finite_w = np.isfinite(Grids[0][w_field].values.squeeze()) - parameters.model_weights[i, :, :, :] = np.where( - np.logical_and.reduce((where_finite_u, where_finite_v, where_finite_w)), - 1, - 0, - ) - - parameters.model_weights = tf.constant(parameters.model_weights, dtype=tf.float32) - - parameters.Co = Co - parameters.Cm = Cm - parameters.Cx = Cx - parameters.Cy = Cy - parameters.Cz = Cz - parameters.Cb = Cb - parameters.Cv = Cv - parameters.Cmod = Cmod - parameters.Cpoint = Cpoint - parameters.roi = roi - parameters.upper_bc = upper_bc - parameters.points = points - parameters.point_list = points - loss_and_gradient = lambda x: (J_function(x, parameters), grad_J(x, parameters)) - - winds = tfp.optimizer.lbfgs_minimize( - loss_and_gradient, - initial_position=winds, - tolerance=tolerance, - x_tolerance=wind_tol, - max_iterations=max_iterations, - parallel_iterations=parallel_iterations, - max_line_search_iterations=20, - ) - winds = np.reshape( - winds.position.numpy(), - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - wcurrmax = winds[2].max() - winds = np.stack([winds[0], winds[1], winds[2]]) - winds = winds.flatten() - # """ - - if low_pass_filter: - print("Applying %s low pass filter to wind field..." % filter_type) - winds = np.asarray(winds) - winds = np.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - winds = _apply_low_pass_filter( - winds, filter_type, filter_window, filter_order, leise_nstep - ) - winds = np.stack([winds[0], winds[1], winds[2]]) - winds = winds.flatten() - - print("Done! Time = " + "{:2.1f}".format(time.time() - bt)) - - the_winds = np.reshape( - winds, - ( - 3, - parameters.grid_shape[0], - parameters.grid_shape[1], - parameters.grid_shape[2], - ), - ) - u = the_winds[0] - v = the_winds[1] - w = the_winds[2] - where_mask = np.sum(parameters.weights, axis=0) + np.sum( - parameters.model_weights, axis=0 - ) - - u = np.ma.array(u) - w = np.ma.array(w) - v = np.ma.array(v) - - if mask_outside_opt is True: - u = np.ma.masked_where(where_mask < 1, u) - v = np.ma.masked_where(where_mask < 1, v) - w = np.ma.masked_where(where_mask < 1, w) - - if mask_w_outside_opt is True: - w = np.ma.masked_where(where_mask < 1, w) - - u_field = {} - u_field["standard_name"] = "u_wind" - u_field["long_name"] = "zonal component of wind velocity" - u_field["units"] = "m/s" - u_field["min_bca"] = min_bca - u_field["max_bca"] = max_bca - v_field = {} - v_field["standard_name"] = "v_wind" - v_field["long_name"] = "meridional component of wind velocity" - v_field["units"] = "m/s" - v_field["min_bca"] = min_bca - v_field["max_bca"] = max_bca - w_field = {} - w_field["standard_name"] = "w_wind" - w_field["long_name"] = "vertical component of wind velocity" - w_field["units"] = "m/s" - w_field["min_bca"] = min_bca - w_field["max_bca"] = max_bca - - new_grid_list = [] - - for grid in Grids: - grid["u"] = xr.DataArray( - np.expand_dims(u, 0), dims=("time", "z", "y", "x"), attrs=u_field - ) - grid["v"] = xr.DataArray( - np.expand_dims(v, 0), dims=("time", "z", "y", "x"), attrs=v_field - ) - grid["w"] = xr.DataArray( - np.expand_dims(w, 0), dims=("time", "z", "y", "x"), attrs=w_field - ) - new_grid_list.append(grid) - - return new_grid_list, parameters - - -def get_dd_wind_field( - Grids, u_init=None, v_init=None, w_init=None, engine="scipy", **kwargs -): - """ - This function takes in a list of Py-ART Grid objects and derives a - wind field. Every Py-ART Grid in Grids must have the same grid - specification. - - In order for the model data constraint to be used, - the model data must be added as a field to at least one of the - grids in Grids. This involves interpolating the model data to the - Grids' coordinates. There are helper functions for this for WRF - and HRRR data in :py:func:`pydda.constraints`: - - :py:func:`make_constraint_from_wrf` - - :py:func:`add_hrrr_constraint_to_grid` - - Parameters - ========== - - Grids: list of Py-ART/DDA Grids - The list of Py-ART or PyDDA grids to take in corresponding to each radar. - All grids must have the same shape, x coordinates, y coordinates - and z coordinates. - u_init: 3D ndarray - The initial guess for the zonal wind field, input as a 3D array - with the same shape as the fields in Grids. If this is None, - PyDDA will use the u field in the first Grid as the initalization. - v_init: 3D ndarray - The initial guess for the meridional wind field, input as a 3D array - with the same shape as the fields in Grids. If this is None, - PyDDA will use the v field in the first Grid as the initalization. - w_init: 3D ndarray - The initial guess for the vertical wind field, input as a 3D array - with the same shape as the fields in Grids. If this is None, - PyDDA will use the w field in the first Grid as the initalization. - engine: str (one of "scipy", "tensorflow", "jax") - Setting this flag will use the solver based off of SciPy, TensorFlow, or Jax. - Using TensorFlow or Jax expands PyDDA's capability to take advantage of GPU-based systems. - In addition, these two implementations use automatic differentation to calculate the gradient - of the cost function in order to optimize the gradient calculation. - TensorFlow 2.6 and tensorflow-probability are required for the TensorFlow-based engine. - The latest version of Jax is required for the Jax-based engine. - points: None or list of dicts - Point observations as returned by :func:`pydda.constraints.get_iem_obs`. Set - to None to disable. - vel_name: string - Name of radial velocity field. Setting to None will have PyDDA attempt - to automatically detect the velocity field name. - refl_field: string - Name of reflectivity field. Setting to None will have PyDDA attempt - to automatically detect the reflectivity field name. - u_back: 1D array - Background zonal wind field from a sounding as a function of height. - This should be given in the sounding's vertical coordinates. - v_back: 1D array - Background meridional wind field from a sounding as a function of - height. This should be given in the sounding's vertical coordinates. - z_back: 1D array - Heights corresponding to background wind field levels in meters. This - is given in the sounding's original coordinates. - frz: float - Freezing level used for fall speed calculation in meters. - Co: float - Weight for cost function related to observed radial velocities. - Cm: float - Weight for cost function related to the mass continuity equation. - Cx: float - Weight for cost function related to smoothness in x direction - Cy: float - Weight for cost function related to smoothness in y direction - Cz: float - Weight for cost function related to smoothness in z direction - Cv: float - Weight for cost function related to vertical vorticity equation. - Cmod: float - Weight for cost function related to custom constraints. - Cpoint: float - Weight for cost function related to point observations. - weights_obs: list of floating point arrays or None - List of weights for each point in grid from each radar in Grids. - Set to None to let PyDDA determine this automatically. - weights_model: list of floating point arrays or None - List of weights for each point in grid from each custom field in - model_fields. Set to None to let PyDDA determine this automatically. - weights_bg: list of floating point arrays or None - List of weights for each point in grid from the sounding. Set to None - to let PyDDA determine this automatically. - Ut: float - Prescribed storm motion in zonal direction. - This is only needed if Cv is not zero. - Vt: float - Prescribed storm motion in meridional direction. - This is only needed if Cv is not zero. - filter_winds: bool - If this is True, PyDDA will run a low pass filter on - the retrieved wind field. Set to False to disable the low pass filter. - mask_outside_opt: bool - If set to true, wind values outside the multiple doppler lobes will - be masked, i.e. if less than 2 radars provide coverage for a given - point. - max_iterations: int - The maximum number of iterations to run the optimization loop for. - mask_w_outside_opt: bool - If set to true, vertical winds outside the multiple doppler lobes will - be masked, i.e. if less than 2 radars provide coverage for a given - point. - filter_type: str (one of "savgol", "leise") - Which low-pass filter to apply after the optimization. ``"savgol"`` - (default) uses ``scipy.signal.savgol_filter`` along each axis with the - ``filter_window`` / ``filter_order`` parameters below. ``"leise"`` uses - the iterated 5-point Leise kernel ([-1/16, 1/4, 5/8, 1/4, -1/16]) with - mirror boundaries, controlled by ``leise_nstep``. - filter_window: int - Window size to use for the Savitzky-Golay low pass filter. A larger - window will increase the number of points factored into the polynomial - fit for the filter, and hence will increase the smoothness. Only used - when ``filter_type="savgol"``. - filter_order: int - The order of the polynomial to use for the Savitzky-Golay low pass - filter. Higher order polynomials allow for the retention of smaller - scale features but may also not remove enough noise. Only used when - ``filter_type="savgol"``. - leise_nstep: int - Number of Leise filter passes to apply along each spatial axis. Each - pass narrows the passband further. Only used when - ``filter_type="leise"``. - min_bca: float - Minimum beam crossing angle in degrees between two radars. 30.0 is the - typical value used in many publications. - max_bca: float - Maximum beam crossing angle in degrees between two radars. 150.0 is the - typical value used in many publications. - upper_bc: bool - Set this to true to enforce w = 0 at the top of the atmosphere. This is - commonly called the impermeability condition. - model_fields: list of strings - The list of fields in the first grid in Grids that contain the custom - data interpolated to the Grid's grid specification. Helper functions - to create such gridded fields for HRRR and NetCDF WRF data exist - in :py:func:`pydda.constraints`. PyDDA will look for fields named *U_(model - field name)*, *V_(model field name)*, and *W_(model field name)*. For - example, if you have *U_hrrr*, *V_hrrr*, and *W_hrrr*, then specify *["hrrr"]* - into model_fields. - output_cost_functions: bool - Set to True to output the value of each cost function every - 10 iterations. - roi: float - Radius of influence for the point observations. The point observation will - not hold any weight outside this radius. - parallel_iterations: int - The number of iterations to run in parallel in the optimization loop. - This is only for the TensorFlow-based engine. - wind_tol: float - Stop iterations after maximum change in winds is less than this value. - tolerance: float - Tolerance for :math:`L_{2}` norm of gradient before stopping. - max_wind_mag: float - Constrain the optimization to have :math:`|u|`, :math:`|v|`, and :math:`|w| < x` m/s. - parallel: bool - If True, enables parallelized cost and gradient computations for the scipy engine. - This vectorizes the radar loop in the radial velocity cost/gradient functions and - computes independent constraint gradients concurrently using a thread pool. - Default is False. - - Returns - ======= - new_grid_list: list - A list of Py-ART grids containing the derived wind fields. These fields - are displayable by the visualization module. - parameters: struct - The parameters used in the generation of the Multi-Doppler wind field. - """ - - if isinstance(Grids, list): - if isinstance(Grids[0], pyart.core.Grid): - for x in Grids: - new_grids = [read_from_pyart_grid(x) for x in Grids] - else: - new_grids = Grids - elif isinstance(Grids, pyart.core.Grid): - new_grids = [read_from_pyart_grid(Grids)] - elif isinstance(Grids, xr.Dataset): - new_grids = [Grids] - else: - raise TypeError( - "Input grids must be an xarray Dataset, Py-ART Grid, or a list of those." - ) - - if u_init is None: - u_init = new_grids[0]["u"].values.squeeze() - - if v_init is None: - v_init = new_grids[0]["v"].values.squeeze() - - if w_init is None: - w_init = new_grids[0]["w"].values.squeeze() - - if ( - engine.lower() == "scipy" - or engine.lower() == "jax" - or engine.lower() == "auglag" - ): - return _get_dd_wind_field_scipy( - new_grids, u_init, v_init, w_init, engine, **kwargs - ) - elif engine.lower() == "tensorflow": - return _get_dd_wind_field_tensorflow( - new_grids, u_init, v_init, w_init, **kwargs - ) - else: - raise NotImplementedError("Engine %s is not supported." % engine) - - -def get_bca(Grid1, Grid2): - """ - This function gets the beam crossing angle between two lat/lon pairs. - - Parameters - ========== - Grid1: xarray (PyDDA) Dataset - The PyDDA Dataset storing the first radar's Grid. - Grid2: PyDDA Dataset - The PyDDA Dataset storing the second radar's Grid. - - Returns - ======= - bca: nD float array - The beam crossing angle between the two radars in radians. - - """ - rad1_lon = Grid1["radar_longitude"].values - rad1_lat = Grid1["radar_latitude"].values - rad2_lon = Grid2["radar_longitude"].values - rad2_lat = Grid2["radar_latitude"].values - x = Grid1["point_x"].values - y = Grid1["point_y"].values - projparams = Grid1["projection"].attrs - if projparams["_include_lon_0_lat_0"] == "true": - projparams["lat_0"] = Grid1["origin_latitude"].values - projparams["lon_0"] = Grid1["origin_longitude"].values - - rad1 = pyart.core.geographic_to_cartesian(rad1_lon, rad1_lat, projparams) - rad2 = pyart.core.geographic_to_cartesian(rad2_lon, rad2_lat, projparams) - # Create grid with Radar 1 in center - - x = x - rad1[0] - y = y - rad1[1] - rad2 = np.array(rad2) - np.array(rad1) - a = np.sqrt(np.multiply(x, x) + np.multiply(y, y)) - b = np.sqrt(pow(x - rad2[0], 2) + pow(y - rad2[1], 2)) - c = np.sqrt(rad2[0] * rad2[0] + rad2[1] * rad2[1]) - inp_array1 = x / a - inp_array1 = np.where(inp_array1 < -1, -1, inp_array1) - inp_array1 = np.where(inp_array1 > 1, 1, inp_array1) - inp_array2 = (x - rad2[1]) / b - inp_array2 = np.where(inp_array2 < -1, -1, inp_array2) - inp_array2 = np.where(inp_array2 > 1, 1, inp_array2) - inp_array3 = (a * a + b * b - c * c) / (2 * a * b) - inp_array3 = np.where(inp_array3 < -1, -1, inp_array3) - inp_array3 = np.where(inp_array3 > 1, 1, inp_array3) - - return np.ma.masked_invalid(np.arccos(inp_array3))[0, :, :] From be7c08c113748be0d8fcb7500892f4fe03dd70b0 Mon Sep 17 00:00:00 2001 From: Robert Jackson Date: Wed, 12 Aug 2026 09:58:53 -0500 Subject: [PATCH 4/4] STY: Apply pre-commit hooks (black formatting) Co-Authored-By: Claude Sonnet 5 --- pydda/cost_functions/_cost_functions_numpy.py | 16 ++++++++-------- pydda/cost_functions/cost_functions.py | 2 +- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pydda/cost_functions/_cost_functions_numpy.py b/pydda/cost_functions/_cost_functions_numpy.py index 390e1cdf..9baf9c07 100644 --- a/pydda/cost_functions/_cost_functions_numpy.py +++ b/pydda/cost_functions/_cost_functions_numpy.py @@ -501,14 +501,14 @@ def calculate_mass_continuity_gradient( grad_w = -np.gradient(div, dz, axis=0) * coeff # Impermeability conditions - grad_w[0, :, :] = 0 # surface is impermeable - if upper_bc == 1:#is True: # impermeable at the grid top - grad_w[-1, :, :] = 0 - if upper_bc == 2: # impermeable at cloud top - N=np.sum(np.array(vrs)>-1000,axis=0) - z_mask = (z > above * 1000) - n_mask = (N == 0) - grad_w[z_mask & n_mask] = 0 + grad_w[0, :, :] = 0 # surface is impermeable + if upper_bc == 1: # is True: # impermeable at the grid top + grad_w[-1, :, :] = 0 + if upper_bc == 2: # impermeable at cloud top + N = np.sum(np.array(vrs) > -1000, axis=0) + z_mask = z > above * 1000 + n_mask = N == 0 + grad_w[z_mask & n_mask] = 0 y = np.stack([grad_u, grad_v, grad_w], axis=0) return y.flatten() diff --git a/pydda/cost_functions/cost_functions.py b/pydda/cost_functions/cost_functions.py index 60c921ea..7787ffaf 100644 --- a/pydda/cost_functions/cost_functions.py +++ b/pydda/cost_functions/cost_functions.py @@ -547,7 +547,7 @@ def grad_J(winds, parameters): parameters.Cm, 1, parameters.upper_bc, - above=parameters.above + above=parameters.above, ) ) if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: