Fixes in the JAX engine, parallel support for scipy engine, and background weight fixes#156
Merged
rcjackson merged 6 commits intoopenradar:mainfrom Apr 8, 2026
Merged
Conversation
…dients Adds a `parallel=True` flag to `get_dd_wind_field` (scipy engine) that enables two optimizations per cost function iteration: 1. Vectorized radar loop: stacks all radar arrays into (n_radars, nz, ny, nx) tensors and computes the radial velocity cost and gradient in a single NumPy broadcast instead of a sequential Python loop. 2. Threaded constraint gradients: submits each active constraint gradient (radial velocity, mass continuity, smoothness, background, vorticity, model, point) to a ThreadPoolExecutor concurrently, then sums the results. NumPy releases the GIL for array ops so threads provide real parallelism. Unit tests extended to cover both serial and parallel paths: - `test_calculate_rad_velocity_cost_parallel`: asserts numerical equivalence of cost and gradient for serial vs. vectorized implementations. - `test_twpice_case_parallel`: end-to-end TWP-ICE retrieval with parallel=True, checking physical sanity and close numerical match with serial results. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
With jit=True, jaxopt compiled the cost function silently during the first solver.run() call, causing a long unexplained pause. Fix by manually jit-compiling loss_and_gradient with jax.jit and triggering a warm-up call before the solver starts. This surfaces the delay with clear print messages and passes jit=False to jaxopt since the function is already compiled. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…round When the sounding doesn't cover the full vertical extent of the grid, interp1d (bounds_error=False) returns NaN for out-of-range levels. Previously bg_weights was never zeroed at those levels, so the background constraint was applied using NaN wind values, corrupting the cost function. Fix by masking bg_weights wherever u_back or v_back is NaN after interpolation, and replacing NaN values in u_back/v_back with 0 so they don't propagate through cost function arithmetic. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #156 +/- ##
==========================================
+ Coverage 68.85% 69.21% +0.35%
==========================================
Files 32 32
Lines 5141 5217 +76
==========================================
+ Hits 3540 3611 +71
- Misses 1601 1606 +5
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR has three additions:
* The JAX engine now does the JIT compilation of the cost function before the optimization loop. In addition, the output is now verbose, making it easier to use the JAX engine and diagnose issues in the retrieval.
* The SciPy engine now has the option to run the individual cost functions in parallel during calculation.
* The background weights were still set to 1, even in regions where the sounding data are unavailable. This has been fixed.