Superficial changes for clarity in Bounce functions#2147
Superficial changes for clarity in Bounce functions#2147
Conversation
…intainer interpax_fft.
When drafting a reply to a reviewer comment, I realized that the atomic derivative computed by the autodiff tool for the `spline=True` option is not correct if the bounce point lies near a local maxima. The `spline=False` option is fine. It is unlikely this would have affected optimization. See section 3 of [autodiff.pdf](https://github.com/user-attachments/files/24988182/autodiff.pdf)
…able names) for improving readability. working commit still as have some outstanding things to understand further
- [x] resolves #1303 . - [x] Activates newton step to find bounce points. - It can be shown there exists O(sqrt(epsilon)) error for bounce integrals with 1/v_ll where epsilon is error of bounce point. For v_|| integrals error is conveniently O(epsilon times sqrt(epsilon)). Hence, the `spline` method would require thousands of knots per transit for just a couple digits of accuracy, and it would stop convergence at epsilon<=1e-5 (so sqrt(epsilon) <=3 digits) error due to condition number. Of course, the `spline=False` method has always computed the points with spectral accuracy and has very fast convergence after #1919 ; that method converges to epsilon = machine precision without the newton step. With the newton step, fast convergence is achieved with the spline method as well. - I suspect fast ion confinement optimization will be easier now. I did this on a couple lunch breaks, so it would be very weird if clicking the approve button to merge into `master` took a year. ## Benchmarks Here is a timing benchmark on my CPU with `nufft_eps=1e-6`. Prior to this PR, every adjoint call to nufft1 took `>= 1 second` and the full computation was `34 seconds`. Now every adjoint call to nufft1 is `250 milliseconds`, and the full computation is `14 seconds`. These `improvements become larger` as the `sparsity` grows and error tolerance parameter for the nuffts epsilon tends to 0. Likewise, the _improvement_ grows linearly with the `problem size`. As this is called within an optimization loop where time and memory are tight, the improvement is significant. ### Before <img width="510" height="135" alt="Screenshot From 2026-03-29 15-22-56" src="proxy.php?url=https%3A%2F%2Fgithub.com%2F%3Ca+href%3D"https://github.com/user-attachments/assets/0baf9a88-b775-4a9c-bb39-db568099aae7">https://github.com/user-attachments/assets/0baf9a88-b775-4a9c-bb39-db568099aae7" /> <img width="510" height="135" alt="Screenshot From 2026-03-29 15-22-36" src="proxy.php?url=https%3A%2F%2Fgithub.com%2F%3Ca+href%3D"https://github.com/user-attachments/assets/2bc1f730-56af-43ad-9c5e-45e9991c0bd1">https://github.com/user-attachments/assets/2bc1f730-56af-43ad-9c5e-45e9991c0bd1" /> ### After <img width="510" height="135" alt="Screenshot From 2026-03-29 15-12-35" src="proxy.php?url=https%3A%2F%2Fgithub.com%2F%3Ca+href%3D"https://github.com/user-attachments/assets/306a6c34-32a3-46e3-8b8c-c4f259a32169">https://github.com/user-attachments/assets/306a6c34-32a3-46e3-8b8c-c4f259a32169" /> <img width="510" height="135" alt="Screenshot From 2026-03-29 15-15-59" src="proxy.php?url=https%3A%2F%2Fgithub.com%2F%3Ca+href%3D"https://github.com/user-attachments/assets/db12ecb3-210a-4732-a5fc-666ba47475d8">https://github.com/user-attachments/assets/db12ecb3-210a-4732-a5fc-666ba47475d8" />
…ents' into dp/bounce-with-comments
unalmis
left a comment
There was a problem hiding this comment.
so it seems that the changes are mostly adding an internal implementation detail describing the angle to the public docstrings. That goes against API abstraction, and makes it more difficult to parse out the things I actually need to read to know how to use the functions
| # TODO: should eventually be a passable argument, we should not expect users to | ||
| # potentially modify source code to use the bounce integrals if this assert ever | ||
| # gets triggered | ||
| assert knots[0] > _sentinel, "Reduce sentinel in desc/integrals/_bounce_utils.py." |
There was a problem hiding this comment.
they would run out of memory way b4 this could be reached, we don't want to add more parameters, the feedback was that makes using it difficult.
The sentinel thing is documented in Bounce1D.
| min_B : jnp.ndarray | ||
| Minimum B value. | ||
| Minimum B value. May be a 1-D array, in which case the inverse pitch (1/λ) | ||
| values and weights will be returned for each pair of min_B and max_B values. |
There was a problem hiding this comment.
this change makes the function confusing. the shape is specified in the output already. existing style was numpy style - e.g. numpy docs don't point out numpy style vectorization since thats assumed by every function
| Fourier-Chebyshev representation used by ``theta_on_fieldlines`` to | ||
| reconstruct ``theta(zeta)`` on followed field lines. | ||
| This is required input to ``Bounce2D`` and is not computed inside | ||
| ``__init__`` so callers can precompute/reuse it across batched calls. |
There was a problem hiding this comment.
These are implementation details that should not be placed in a public API.
| interpolation, effectively obtaining the map delta(rho,alpha,zeta) so one may | ||
| compute theta on a given field line (given by alpha) at a given point along | ||
| that field line (given by zeta) by simply calculating | ||
| theta = alpha + delta(rho,alpha,zeta). |
There was a problem hiding this comment.
these are implementation details that should not be in the public API.
| not kwargs.pop("ignore_lambda_guard", False), | ||
| "The 'lambda' option for the 'name' kwarg in Bounce2D.angle " | ||
| "is deprecated and should not be used If you are using this," | ||
| " please update the test to use 'delta' instead.", |
There was a problem hiding this comment.
The lambda option was never made publically available, so this is not necessary.
| Note this dictionary will be modified. | ||
| returned by the registered compute functions in ``desc.compute`` e.g. | ||
| pointwise quantities are 1-D arrays of length grid.num_nodes and | ||
| vector quantities are 2-D arrays of shape (grid.num_nodes, 3). |
There was a problem hiding this comment.
we don't support that input shape
| "Got None for angle. Must pass in angle " | ||
| "as a kwarg when computing bounce-averaged quantities, with angle being " | ||
| "the output of Bounce2D.angle, which is the stream-map angle used to " | ||
| "reconstruct theta(zeta) on field lines. ", |
There was a problem hiding this comment.
this is implementation detail that need not be publicized to users
| 7. Return per-well bounce integrals, then reduce/average in callers as needed. | ||
|
|
||
| See "effective ripple 3/2" in ``desc/compute/_neoclassical.py`` for an example | ||
| of how the Bounce2D method is used in a compute function, and |
There was a problem hiding this comment.
i already say this in the Bounce1D2D docstrings
| ι (iota) | ||
| Rotational transform (normalized by 2π); relates toroidal and poloidal | ||
| winding of a field line. Used in ``theta_on_fieldlines`` and in the | ||
| upper bound on the number of magnetic wells per toroidal transit. |
There was a problem hiding this comment.
the above is defined in the desc/compute/data index, with all the proper units etc. so redefining all that here just increases the maintenance surface.
| Lref | ||
| Reference length scale used for non-dimensionalization. | ||
| xₖ, wₖ | ||
| Quadrature nodes and weights for approximating ∫₋₁¹ g(x) dx ≈ ∑ₖ wₖ g(xₖ). |
There was a problem hiding this comment.
Aren't these all defined in the Bounce1d2d docs
ideally, it would be computable quantity in data_index and then simply a required_names just like "|B|". But i think i got push back when I wanted to make map coordinates not require an equilibrium
These functions do not assume you have a block of data on every surface simaltenously therefore, the input data must be reshaped.
Interpolating a vector requires more care to ensure it retains vector invariances such as div free, curl free, rotation etc. Therefore, it is less accurate to interpolate a vector-valued quantity and we don't do it. |
Comments for future consideration/improved usability:
Bounce2D.integratecannot contain a vector quantity (at least not according to my one attempt), so if there ever were a use case where one wanted say a basis vector or vector B or something, it would not be possible. This is just a point for future improvement to note.