diff --git a/pydda/cost_functions/_cost_functions_jax.py b/pydda/cost_functions/_cost_functions_jax.py index 8ac68a6e..b1bd313c 100644 --- a/pydda/cost_functions/_cost_functions_jax.py +++ b/pydda/cost_functions/_cost_functions_jax.py @@ -24,7 +24,7 @@ def _apply_upper_bc(grad_w, upper_bc, upper_bc_mask=None): def calculate_radial_vel_cost_function( - vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0 + vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0, radar_geometry=None ): """ Calculates the cost function due to difference of the wind field from @@ -80,11 +80,17 @@ def calculate_radial_vel_cost_function( J_o = 0 lambda_o = coeff / (rmsVr * rmsVr) for i in range(len(vrs)): - v_ar = ( - jnp.cos(els[i]) * jnp.sin(azs[i]) * u - + jnp.cos(els[i]) * jnp.cos(azs[i]) * v - + jnp.sin(els[i]) * (w - jnp.abs(wts[i])) - ) + if radar_geometry is None: + x_projection = jnp.cos(els[i]) * jnp.sin(azs[i]) + y_projection = jnp.cos(els[i]) * jnp.cos(azs[i]) + sin_el = jnp.sin(els[i]) + wts_abs = jnp.abs(wts[i]) + else: + x_projection = radar_geometry["x_projection"][i] + y_projection = radar_geometry["y_projection"][i] + sin_el = radar_geometry["sin_el"][i] + wts_abs = radar_geometry["wts_abs"][i] + v_ar = x_projection * u + y_projection * v + sin_el * (w - wts_abs) the_weight = jnp.asarray(weights[i]) J_o += lambda_o * jnp.sum(jnp.square(vrs[i] - v_ar) * the_weight) return J_o @@ -307,7 +313,7 @@ def calculate_smoothness_gradient( return y.flatten() -def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3): +def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, point_weights=None): """ 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 @@ -344,21 +350,26 @@ def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3): wind field and points. """ J = 0.0 - for the_point in point_list: - dist = jnp.sqrt( - (x - the_point["x"]) ** 2 - + (y - the_point["y"]) ** 2 - + (z - the_point["z"]) ** 2 - ) - dist = jnp.maximum(dist, 1.0) - weight = 1 / dist**2 - weight = weight / jnp.sum(weight) + for point_index, the_point in enumerate(point_list): + if point_weights is None: + dist = jnp.sqrt( + (x - the_point["x"]) ** 2 + + (y - the_point["y"]) ** 2 + + (z - the_point["z"]) ** 2 + ) + dist = jnp.maximum(dist, 1.0) + weight = 1 / dist**2 + weight = weight / jnp.sum(weight) + else: + weight = point_weights[point_index] J += jnp.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): +def calculate_point_gradient( + u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, point_weights=None +): """ 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 @@ -398,15 +409,18 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): gradJ_v = jnp.zeros_like(v) gradJ_w = jnp.zeros_like(u) - for the_point in point_list: - dist = jnp.sqrt( - (x - the_point["x"]) ** 2 - + (y - the_point["y"]) ** 2 - + (z - the_point["z"]) ** 2 - ) - dist = jnp.maximum(dist, 1.0) - weight = 1 / dist**2 - weight = weight / jnp.sum(weight) + for point_index, the_point in enumerate(point_list): + if point_weights is None: + dist = jnp.sqrt( + (x - the_point["x"]) ** 2 + + (y - the_point["y"]) ** 2 + + (z - the_point["z"]) ** 2 + ) + dist = jnp.maximum(dist, 1.0) + weight = 1 / dist**2 + weight = weight / jnp.sum(weight) + else: + weight = point_weights[point_index] gradJ_u += 2 * (u - the_point["u"]) * weight gradJ_v += 2 * (v - the_point["v"]) * weight diff --git a/pydda/cost_functions/_cost_functions_numpy.py b/pydda/cost_functions/_cost_functions_numpy.py index 350241d1..0e378b82 100644 --- a/pydda/cost_functions/_cost_functions_numpy.py +++ b/pydda/cost_functions/_cost_functions_numpy.py @@ -249,6 +249,7 @@ def calculate_radial_vel_cost_and_gradient( upper_bc=1, upper_bc_mask=None, parallel=False, + radar_geometry=None, ): """Calculate radar cost and gradient while sharing the radar projection. @@ -260,17 +261,24 @@ def calculate_radial_vel_cost_and_gradient( if parallel: vrs_arr = np.stack(vrs) - els_arr = np.stack(els) - azs_arr = np.stack(azs) - wts_arr = np.stack(wts) - cos_els = np.cos(els_arr) - sin_els = np.sin(els_arr) - x_projection = cos_els * np.sin(azs_arr) - y_projection = cos_els * np.cos(azs_arr) + if radar_geometry is None: + els_arr = np.stack(els) + azs_arr = np.stack(azs) + wts_arr = np.stack(wts) + cos_els = np.cos(els_arr) + sin_els = np.sin(els_arr) + x_projection = cos_els * np.sin(azs_arr) + y_projection = cos_els * np.cos(azs_arr) + wts_abs = np.abs(wts_arr) + else: + x_projection = radar_geometry["x_projection"] + y_projection = radar_geometry["y_projection"] + sin_els = radar_geometry["sin_el"] + wts_abs = radar_geometry["wts_abs"] v_ar = ( x_projection * u[np.newaxis] + y_projection * v[np.newaxis] - + sin_els * (w[np.newaxis] - np.abs(wts_arr)) + + sin_els * (w[np.newaxis] - wts_abs) ) residual = v_ar - vrs_arr cost = lambda_o * np.sum(np.square(residual) * weights) @@ -282,12 +290,21 @@ def calculate_radial_vel_cost_and_gradient( p_x1 = np.zeros(vrs[0].shape) p_y1 = np.zeros(vrs[0].shape) p_z1 = np.zeros(vrs[0].shape) - for vr, el, az, wt, weight in zip(vrs, els, azs, wts, weights): - cos_el = np.cos(el) - sin_el = np.sin(el) - x_projection = cos_el * np.sin(az) - y_projection = cos_el * np.cos(az) - v_ar = x_projection * u + y_projection * v + sin_el * (w - np.abs(wt)) + for index, (vr, wt, weight) in enumerate(zip(vrs, wts, weights)): + if radar_geometry is None: + el = els[index] + az = azs[index] + cos_el = np.cos(el) + sin_el = np.sin(el) + x_projection = cos_el * np.sin(az) + y_projection = cos_el * np.cos(az) + wts_abs = np.abs(wt) + else: + x_projection = radar_geometry["x_projection"][index] + y_projection = radar_geometry["y_projection"][index] + sin_el = radar_geometry["sin_el"][index] + wts_abs = radar_geometry["wts_abs"][index] + v_ar = x_projection * u + y_projection * v + sin_el * (w - wts_abs) residual = v_ar - vr cost += lambda_o * np.sum(np.square(residual) * weight) p_x1 += 2 * residual * x_projection * weight * lambda_o @@ -457,7 +474,9 @@ def _point_weights(x, y, z, the_point, roi): return weight / maximum if maximum > 0 else weight -def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0): +def calculate_point_cost( + u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0, point_weights=None +): """ 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 @@ -490,18 +509,24 @@ def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0) The cost function related to the difference between wind field and points. """ J = 0.0 - for the_point in point_list: + for point_index, the_point in enumerate(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 - weight = _point_weights(x, y, z, the_point, roi) + weight = ( + point_weights[point_index] + if point_weights is not None + else _point_weights(x, y, z, the_point, roi) + ) 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): +def calculate_point_gradient( + u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, point_weights=None +): """ 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 @@ -537,8 +562,12 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): gradJ_v = np.zeros_like(v) gradJ_w = np.zeros_like(u) - for the_point in point_list: - weight = _point_weights(x, y, z, the_point, roi) + for point_index, the_point in enumerate(point_list): + weight = ( + point_weights[point_index] + if point_weights is not None + else _point_weights(x, y, z, the_point, roi) + ) gradJ_u += 2 * weight * (u - the_point["u"]) gradJ_v += 2 * weight * (v - the_point["v"]) diff --git a/pydda/cost_functions/_cost_functions_tensorflow.py b/pydda/cost_functions/_cost_functions_tensorflow.py index 533b7b71..a8be4183 100644 --- a/pydda/cost_functions/_cost_functions_tensorflow.py +++ b/pydda/cost_functions/_cost_functions_tensorflow.py @@ -109,6 +109,7 @@ def calculate_radial_vel_cost_and_gradient( upper_bc=1, upper_bc_mask=None, lower_bc=True, + radar_geometry=None, ): """Calculate radar cost and gradient while sharing each radar residual.""" lambda_o = coeff / (rmsVr * rmsVr) @@ -118,17 +119,22 @@ def calculate_radial_vel_cost_and_gradient( p_z1 = tf.zeros_like(w) for i in range(len(vrs)): - cos_el = tf.math.cos(els[i]) - sin_el = tf.math.sin(els[i]) - v_ar = ( - cos_el * tf.math.sin(azs[i]) * u - + cos_el * tf.math.cos(azs[i]) * v - + sin_el * (w - tf.math.abs(wts[i])) - ) + if radar_geometry is None: + cos_el = tf.math.cos(els[i]) + sin_el = tf.math.sin(els[i]) + x_projection = cos_el * tf.math.sin(azs[i]) + y_projection = cos_el * tf.math.cos(azs[i]) + wts_abs = tf.math.abs(wts[i]) + else: + x_projection = radar_geometry["x_projection"][i] + y_projection = radar_geometry["y_projection"][i] + sin_el = radar_geometry["sin_el"][i] + wts_abs = radar_geometry["wts_abs"][i] + v_ar = x_projection * u + y_projection * v + sin_el * (w - wts_abs) residual = v_ar - vrs[i] costs.append(lambda_o * tf.reduce_sum(tf.math.square(residual) * weights[i])) - p_x1 += 2 * residual * cos_el * tf.math.sin(azs[i]) * weights[i] * lambda_o - p_y1 += 2 * residual * cos_el * tf.math.cos(azs[i]) * weights[i] * lambda_o + p_x1 += 2 * residual * x_projection * weights[i] * lambda_o + p_y1 += 2 * residual * y_projection * weights[i] * lambda_o p_z1 += 2 * residual * sin_el * weights[i] * lambda_o if lower_bc: @@ -383,7 +389,9 @@ def calculate_smoothness_gradient( return tf.reshape(y, (3 * np.prod(u.shape),)) -def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): +def calculate_point_cost( + u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, point_weights=None +): """ 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 @@ -421,26 +429,31 @@ def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): """ J = tf.Variable(0.0) - for the_point in point_list: + for point_index, the_point in enumerate(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 up = tf.ones_like(u) * the_point["u"] vp = tf.ones_like(v) * the_point["v"] - dist = tf.math.sqrt( - (x - the_point["x"]) ** 2 - + (y - the_point["y"]) ** 2 - + (z - the_point["z"]) ** 2 - ) - dist = tf.math.maximum(dist, 1.0) - weight = 1 / dist**2 - weight = weight / tf.reduce_max(weight) + if point_weights is None: + dist = tf.math.sqrt( + (x - the_point["x"]) ** 2 + + (y - the_point["y"]) ** 2 + + (z - the_point["z"]) ** 2 + ) + dist = tf.math.maximum(dist, 1.0) + weight = 1 / dist**2 + weight = weight / tf.reduce_max(weight) + else: + weight = point_weights[point_index] J.assign_add(tf.math.reduce_sum(((u - up) ** 2 + (v - vp) ** 2) * weight)) return J * Cp -def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): +def calculate_point_gradient( + u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, point_weights=None +): """ 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 @@ -477,20 +490,23 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): gradJ_u = tf.Variable(tf.zeros(u.shape, dtype=tf.float32)) gradJ_v = tf.Variable(tf.zeros(v.shape, dtype=tf.float32)) gradJ_w = tf.Variable(tf.zeros(u.shape, dtype=tf.float32)) - for the_point in point_list: + for point_index, the_point in enumerate(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 up = tf.ones_like(u, dtype=tf.float32) * the_point["u"] vp = tf.ones_like(v, dtype=tf.float32) * the_point["v"] - dist = tf.math.sqrt( - (x - the_point["x"]) ** 2 - + (y - the_point["y"]) ** 2 - + (z - the_point["z"]) ** 2 - ) - dist = tf.math.maximum(dist, 1.0) - weight = 1 / dist**2 - weight = weight / tf.reduce_max(weight) + if point_weights is None: + dist = tf.math.sqrt( + (x - the_point["x"]) ** 2 + + (y - the_point["y"]) ** 2 + + (z - the_point["z"]) ** 2 + ) + dist = tf.math.maximum(dist, 1.0) + weight = 1 / dist**2 + weight = weight / tf.reduce_max(weight) + else: + weight = point_weights[point_index] gradJ_u.assign_add((2 * (u - up) * weight)) gradJ_v.assign_add((2 * (v - vp) * weight)) gradJ = tf.stack([gradJ_u, gradJ_v, gradJ_w], axis=0) diff --git a/pydda/cost_functions/cost_functions.py b/pydda/cost_functions/cost_functions.py index 037cf31f..a1eae4e1 100644 --- a/pydda/cost_functions/cost_functions.py +++ b/pydda/cost_functions/cost_functions.py @@ -180,6 +180,7 @@ def J_function(winds, parameters): parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, + point_weights=parameters.point_weights, ) else: Jpoint = 0 @@ -309,6 +310,7 @@ def J_function(winds, parameters): parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, + point_weights=parameters.point_weights, ) else: Jpoint = 0 @@ -511,6 +513,7 @@ def grad_J(winds, parameters): parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, + point_weights=parameters.point_weights, ) if parameters.const_boundary_cond is True: grad = tf.reshape( @@ -711,6 +714,7 @@ def grad_J(winds, parameters): parameters.point_list, parameters.Cpoint, parameters.roi, + parameters.point_weights, ) ) if parameters.Cvad > 0: @@ -835,6 +839,7 @@ def grad_J(winds, parameters): parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, + point_weights=parameters.point_weights, ) if parameters.Cvad > 0: @@ -922,6 +927,7 @@ def J_and_grad(winds, parameters): upper_bc=parameters.upper_bc, upper_bc_mask=parameters.upper_bc_mask, lower_bc=parameters.lower_bc, + radar_geometry=parameters.radar_geometry, ) ) previous_cache = getattr(parameters, "_radial_eval_cache", None) @@ -992,6 +998,7 @@ def J_and_grad(winds, parameters): upper_bc=parameters.upper_bc, upper_bc_mask=parameters.upper_bc_mask, parallel=parameters.parallel, + radar_geometry=parameters.radar_geometry, ) ) previous_cache = getattr(parameters, "_radial_eval_cache", None) @@ -1034,6 +1041,7 @@ def J_function_jax(winds, parameters): rmsVr=parameters.rmsVr, weights=parameters.weights, coeff=parameters.Co, + radar_geometry=parameters.radar_geometry, ) if parameters.Cm > 0: @@ -1131,6 +1139,7 @@ def J_function_jax(winds, parameters): parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, + point_weights=parameters.point_weights, ) else: Jpoint = 0 @@ -1273,6 +1282,7 @@ def grad_jax(winds, parameters): parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, + point_weights=parameters.point_weights, ) if parameters.Cvad > 0: diff --git a/pydda/retrieval/wind_retrieve.py b/pydda/retrieval/wind_retrieve.py index 3934db0f..e0ff2b22 100644 --- a/pydda/retrieval/wind_retrieve.py +++ b/pydda/retrieval/wind_retrieve.py @@ -75,6 +75,54 @@ def _apply_low_pass_filter( iterations = 0 +def _precompute_point_weights(x, y, z, points, roi, engine): + """Precompute fixed point-constraint weights for optimizer iterations.""" + if not points: + return [] + cached = [] + for point in points: + distance = np.sqrt( + (x - point["x"]) ** 2 + (y - point["y"]) ** 2 + (z - point["z"]) ** 2 + ) + if engine == "jax": + weight = 1.0 / np.maximum(distance, 1.0) ** 2 + weight /= np.sum(weight) + else: + weight = np.where( + distance <= roi, 1.0 / np.maximum(distance, 1.0) ** 2, 0.0 + ) + maximum = np.max(weight) + if maximum > 0: + weight /= maximum + cached.append(weight) + return cached + + +def _precompute_radar_geometry(els, azs, wts, engine): + """Precompute trigonometric radar projection coefficients.""" + if engine == "tensorflow": + els = tf.stack(els) + azs = tf.stack(azs) + wts = tf.stack(wts) + cos_els = tf.math.cos(els) + return { + "x_projection": cos_els * tf.math.sin(azs), + "y_projection": cos_els * tf.math.cos(azs), + "sin_el": tf.math.sin(els), + "wts_abs": tf.math.abs(wts), + } + els = np.stack(els) + azs = np.stack(azs) + wts = np.stack(wts) + cos_els = np.cos(els) + return { + "x_projection": cos_els * np.sin(azs), + "y_projection": cos_els * np.cos(azs), + "sin_el": np.sin(els), + "wts_abs": np.abs(wts), + } + + class DDParameters(object): """ This is a helper class for inserting more arguments into the :func:`pydda.cost_functions.J_function` and @@ -232,6 +280,8 @@ def __init__(self): self.Jveltol = 100.0 self.const_boundary_cond = False self.parallel = False + self.radar_geometry = None + self.point_weights = [] def _get_dd_wind_field_scipy( @@ -689,6 +739,12 @@ def _get_dd_wind_field_scipy( parameters.points = points parameters.point_list = points parameters.parallel = parallel + parameters.radar_geometry = _precompute_radar_geometry( + parameters.els, parameters.azs, parameters.wts, engine.lower() + ) + parameters.point_weights = _precompute_point_weights( + parameters.x, parameters.y, parameters.z, points, roi, engine.lower() + ) _wprevmax = np.zeros(parameters.grid_shape) _wcurrmax = np.zeros(parameters.grid_shape) iterations = 0 @@ -1319,6 +1375,20 @@ def _get_dd_wind_field_tensorflow( parameters.upper_bc = upper_bc parameters.points = points parameters.point_list = points + parameters.radar_geometry = _precompute_radar_geometry( + parameters.els, parameters.azs, parameters.wts, "tensorflow" + ) + parameters.point_weights = [ + tf.constant(weight, dtype=tf.float32) + for weight in _precompute_point_weights( + Grids[0]["point_x"].values, + Grids[0]["point_y"].values, + Grids[0]["point_z"].values, + points, + roi, + "tensorflow", + ) + ] loss_and_gradient = lambda x: J_and_grad(x, parameters) winds = tfp.optimizer.lbfgs_minimize(