diff --git a/pydda/cost_functions/__init__.py b/pydda/cost_functions/__init__.py index 44e908ee..b3a978c1 100644 --- a/pydda/cost_functions/__init__.py +++ b/pydda/cost_functions/__init__.py @@ -87,6 +87,7 @@ from ._cost_functions_numpy import calculate_radial_vel_cost_function from ._cost_functions_numpy import calculate_fall_speed from ._cost_functions_numpy import calculate_grad_radial_vel +from ._cost_functions_numpy import calculate_radial_vel_cost_and_gradient from ._cost_functions_numpy import calculate_mass_continuity from ._cost_functions_numpy import calculate_mass_continuity_gradient from ._cost_functions_numpy import calculate_smoothness_cost @@ -99,4 +100,4 @@ from ._cost_functions_numpy import calculate_model_gradient from ._cost_functions_numpy import calculate_point_cost, calculate_point_gradient from ._cost_functions_numpy import calculate_echo_top_mask -from .cost_functions import J_function, grad_J, grad_jax, J_function_jax +from .cost_functions import J_function, J_and_grad, grad_J, grad_jax, J_function_jax diff --git a/pydda/cost_functions/_cost_functions_numpy.py b/pydda/cost_functions/_cost_functions_numpy.py index 3e62729f..350241d1 100644 --- a/pydda/cost_functions/_cost_functions_numpy.py +++ b/pydda/cost_functions/_cost_functions_numpy.py @@ -235,6 +235,70 @@ def calculate_grad_radial_vel( return y.flatten() +def calculate_radial_vel_cost_and_gradient( + vrs, + els, + azs, + u, + v, + w, + wts, + weights, + rmsVr, + coeff=1.0, + upper_bc=1, + upper_bc_mask=None, + parallel=False, +): + """Calculate radar cost and gradient while sharing the radar projection. + + This is used by the combined SciPy optimizer callback. The existing + cost-only and gradient-only functions remain available for callers that + need just one result. + """ + 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) + 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) + v_ar = ( + x_projection * u[np.newaxis] + + y_projection * v[np.newaxis] + + sin_els * (w[np.newaxis] - np.abs(wts_arr)) + ) + residual = v_ar - vrs_arr + cost = lambda_o * np.sum(np.square(residual) * weights) + p_x1 = np.sum(2 * residual * x_projection * weights, axis=0) * lambda_o + p_y1 = np.sum(2 * residual * y_projection * weights, axis=0) * lambda_o + p_z1 = np.sum(2 * residual * sin_els * weights, axis=0) * lambda_o + else: + cost = 0.0 + 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)) + residual = v_ar - vr + cost += lambda_o * np.sum(np.square(residual) * weight) + p_x1 += 2 * residual * x_projection * weight * lambda_o + p_y1 += 2 * residual * y_projection * weight * lambda_o + p_z1 += 2 * residual * sin_el * weight * lambda_o + + p_z1[0, :, :] = 0 + p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask) + return cost, np.stack((p_x1, p_y1, p_z1), axis=0).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 diff --git a/pydda/cost_functions/_cost_functions_tensorflow.py b/pydda/cost_functions/_cost_functions_tensorflow.py index 9416ae32..533b7b71 100644 --- a/pydda/cost_functions/_cost_functions_tensorflow.py +++ b/pydda/cost_functions/_cost_functions_tensorflow.py @@ -95,6 +95,51 @@ def calculate_radial_vel_cost_function( return J_o +def calculate_radial_vel_cost_and_gradient( + vrs, + els, + azs, + u, + v, + w, + wts, + weights, + rmsVr, + coeff=1.0, + upper_bc=1, + upper_bc_mask=None, + lower_bc=True, +): + """Calculate radar cost and gradient while sharing each radar residual.""" + lambda_o = coeff / (rmsVr * rmsVr) + costs = [] + p_x1 = tf.zeros_like(u) + p_y1 = tf.zeros_like(v) + 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])) + ) + 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_z1 += 2 * residual * sin_el * weights[i] * lambda_o + + if lower_bc: + p_z1 = tf.concat( + [tf.zeros_like(p_z1[0:1]), p_z1[1:]], + axis=0, + ) + p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask) + return tf.add_n(costs), tf.reshape(tf.stack((p_x1, p_y1, p_z1)), [-1]) + + def calculate_grad_radial_vel( vrs, els, diff --git a/pydda/cost_functions/cost_functions.py b/pydda/cost_functions/cost_functions.py index 9ef5395f..037cf31f 100644 --- a/pydda/cost_functions/cost_functions.py +++ b/pydda/cost_functions/cost_functions.py @@ -10,6 +10,7 @@ TENSORFLOW_AVAILABLE = False try: + import jax import jax.numpy as jnp JAX_AVAILABLE = True @@ -66,6 +67,7 @@ def J_function(winds, parameters): raise ImportError( "Tensorflow 2.5 or greater is needed in order to use TensorFlow-based PyDDA!" ) + winds_input = winds winds = tf.reshape( winds, ( @@ -78,18 +80,22 @@ def J_function(winds, parameters): 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, - ) + radial_cache = getattr(parameters, "_radial_eval_cache", None) + if radial_cache is not None and radial_cache["source_winds"] is winds_input: + Jvel = radial_cache["cost"] + else: + 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: @@ -178,6 +184,7 @@ def J_function(winds, parameters): else: Jpoint = 0 elif parameters.engine == "scipy": + winds_input = winds winds = np.reshape( winds, ( @@ -188,19 +195,23 @@ def J_function(winds, parameters): ), ) # 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, - ) + radial_cache = getattr(parameters, "_radial_eval_cache", None) + if radial_cache is not None and radial_cache["source_winds"] is winds_input: + Jvel = radial_cache["cost"] + else: + 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: @@ -387,6 +398,7 @@ def grad_J(winds, parameters): raise ImportError( "Tensorflow 2.5 or greater is needed in order to use TensorFlow-based PyDDA!" ) + winds_input = winds winds = tf.reshape( winds, ( @@ -399,21 +411,25 @@ def grad_J(winds, parameters): 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, - upper_bc_mask=parameters.upper_bc_mask, - lower_bc=parameters.lower_bc, - ) + radial_cache = getattr(parameters, "_radial_eval_cache", None) + if radial_cache is not None and radial_cache["source_winds"] is winds_input: + grad = tf.identity(radial_cache["gradient"]) + else: + 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, + upper_bc_mask=parameters.upper_bc_mask, + lower_bc=parameters.lower_bc, + ) if parameters.Cm > 0: grad += _cost_functions_tensorflow.calculate_mass_continuity_gradient( @@ -557,6 +573,7 @@ def grad_J(winds, parameters): ) grad = tf.reshape(grad, [-1]) elif parameters.engine == "scipy": + winds_input = winds winds = np.reshape( winds, ( @@ -566,7 +583,11 @@ def grad_J(winds, parameters): parameters.grid_shape[2], ), ) - if parameters.parallel: + radial_cache = getattr(parameters, "_radial_eval_cache", None) + use_radial_cache = ( + radial_cache is not None and radial_cache["source_winds"] is winds_input + ) + if parameters.parallel and not use_radial_cache: futures = [] with ThreadPoolExecutor() as pool: futures.append( @@ -708,20 +729,23 @@ def grad_J(winds, parameters): ) 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, - upper_bc_mask=parameters.upper_bc_mask, - ) + if use_radial_cache: + grad = radial_cache["gradient"].copy() + 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, + upper_bc_mask=parameters.upper_bc_mask, + ) if parameters.Cm > 0: grad += _cost_functions_numpy.calculate_mass_continuity_gradient( @@ -865,6 +889,126 @@ def grad_J(winds, parameters): return grad +def J_and_grad(winds, parameters): + """Return the objective and gradient for a single optimizer callback. + + SciPy's L-BFGS-B accepts a callable that returns ``(f, g)`` when + ``fprime`` is omitted. Keeping this adapter here gives the optimizer one + callback per evaluation and provides a single place to add shared + objective/gradient intermediates in the future. + """ + if parameters.engine == "tensorflow": + shaped_winds = tf.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + radial_cost, radial_gradient = ( + _cost_functions_tensorflow.calculate_radial_vel_cost_and_gradient( + parameters.vrs, + parameters.els, + parameters.azs, + shaped_winds[0], + shaped_winds[1], + shaped_winds[2], + parameters.wts, + parameters.weights, + parameters.rmsVr, + coeff=parameters.Co, + upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, + lower_bc=parameters.lower_bc, + ) + ) + previous_cache = getattr(parameters, "_radial_eval_cache", None) + parameters._radial_eval_cache = { + "source_winds": winds, + "cost": radial_cost, + "gradient": radial_gradient, + } + try: + return J_function(winds, parameters), grad_J(winds, parameters) + finally: + if previous_cache is None: + del parameters._radial_eval_cache + else: + parameters._radial_eval_cache = previous_cache + + if parameters.engine == "jax": + objective = lambda objective_winds: J_function_jax(objective_winds, parameters) + value, gradient = jax.value_and_grad(objective)(winds) + gradient = jnp.reshape( + gradient, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + if parameters.lower_bc == 1: + gradient = gradient.at[2, 0].set(0) + if parameters.upper_bc == 1: + gradient = gradient.at[2, -1].set(0) + elif parameters.upper_bc == 2 and parameters.upper_bc_mask is not None: + gradient = gradient.at[2].set( + jnp.where(parameters.upper_bc_mask, 0, gradient[2]) + ) + if parameters.const_boundary_cond is True: + gradient = gradient.at[:, :, 0, :].set(0) + gradient = gradient.at[:, :, -1, :].set(0) + gradient = gradient.at[:, :, :, 0].set(0) + gradient = gradient.at[:, :, :, -1].set(0) + return value, gradient.flatten() + + if parameters.engine != "scipy": + return J_function(winds, parameters), grad_J(winds, parameters) + + shaped_winds = np.reshape( + winds, + ( + 3, + parameters.grid_shape[0], + parameters.grid_shape[1], + parameters.grid_shape[2], + ), + ) + radial_cost, radial_gradient = ( + _cost_functions_numpy.calculate_radial_vel_cost_and_gradient( + parameters.vrs, + parameters.els, + parameters.azs, + shaped_winds[0], + shaped_winds[1], + shaped_winds[2], + parameters.wts, + parameters.weights, + parameters.rmsVr, + coeff=parameters.Co, + upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, + parallel=parameters.parallel, + ) + ) + previous_cache = getattr(parameters, "_radial_eval_cache", None) + parameters._radial_eval_cache = { + "source_winds": winds, + "cost": radial_cost, + "gradient": radial_gradient, + } + try: + return J_function(winds, parameters), grad_J(winds, parameters) + finally: + if previous_cache is None: + del parameters._radial_eval_cache + else: + parameters._radial_eval_cache = previous_cache + + def J_function_jax(winds, parameters): if not JAX_AVAILABLE: raise ImportError("Jax is needed in order to use the Jax-based PyDDA!") diff --git a/pydda/retrieval/wind_retrieve.py b/pydda/retrieval/wind_retrieve.py index 0286e817..3934db0f 100644 --- a/pydda/retrieval/wind_retrieve.py +++ b/pydda/retrieval/wind_retrieve.py @@ -63,12 +63,9 @@ def _apply_low_pass_filter( # imports changed to local import path to run on computer from ..cost_functions import ( - J_function, - grad_J, + J_and_grad, calculate_fall_speed, calculate_echo_top_mask, - grad_jax, - J_function_jax, ) from copy import deepcopy from .angles import add_azimuth_as_field, add_elevation_as_field @@ -733,22 +730,19 @@ def _vert_velocity_callback(x): parameters.print_out = False if engine.lower() == "scipy": winds = fmin_l_bfgs_b( - J_function, + J_and_grad, 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 + x_loss, x_gradient = J_and_grad(x["winds"], parameters) + return x_loss, {"winds": x_gradient} bounds = ( {"winds": -max_wind_mag * jnp.ones(winds.shape)}, @@ -1325,7 +1319,7 @@ def _get_dd_wind_field_tensorflow( 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)) + loss_and_gradient = lambda x: J_and_grad(x, parameters) winds = tfp.optimizer.lbfgs_minimize( loss_and_gradient, diff --git a/pydda/tests/test_cost_functions.py b/pydda/tests/test_cost_functions.py index 840ea717..33ab7069 100644 --- a/pydda/tests/test_cost_functions.py +++ b/pydda/tests/test_cost_functions.py @@ -81,6 +81,30 @@ def test_calculate_rad_velocity_cost_parallel(): np.testing.assert_allclose(parallel_grad, serial_grad) +def test_calculate_rad_velocity_cost_and_gradient_shares_result(): + """The combined radar evaluator matches the separate evaluations.""" + vrs, azs, els, wts, weights = _make_radvel_inputs() + rng = np.random.default_rng(42) + u = rng.random((20, 20, 20)) + v = rng.random((20, 20, 20)) + w = rng.random((20, 20, 20)) + + combined_cost, combined_grad = ( + pydda.cost_functions.calculate_radial_vel_cost_and_gradient( + vrs, els, azs, u, v, w, wts, weights, 1.0 + ) + ) + separate_cost = pydda.cost_functions.calculate_radial_vel_cost_function( + vrs, azs, els, u, v, w, wts, 1.0, weights + ) + separate_grad = pydda.cost_functions.calculate_grad_radial_vel( + vrs, els, azs, u, v, w, wts, weights, 1.0 + ) + + np.testing.assert_allclose(combined_cost, separate_cost) + np.testing.assert_allclose(combined_grad, separate_grad) + + @pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") def test_calculate_rad_velocity_cost_jax(): """Test with a zero velocity field radar"""