Partial summation in coordinate mapping #1826

unalmis · 2025-07-30T19:58:44Z

Progress toward addressing #1154.

The attached benchmark must be ran on ku/nufft branch.
250x speed improvement from ~15 seconds to ~50 milliseconds.
The improvement would be more if the FourierZernike basis is padded as discussed in #1243.
Then we would avoid the 3D spectral to real transform as well as N^2 FFTs of size N.
(Then this computation would likely be in microsecond range).

unalmis · 2025-07-30T20:37:06Z

desc/equilibrium/coords.py

+
+    TODO(#1243) Do proper partial summation once the DESC
+    DESC basis are improved to store the padded tensor product basis.
+    https://github.com/PlasmaControl/DESC/issues/1243#issuecomment-3131182128.


@PlasmaControl/desc-dev in #1508 header it says "figure out how to do FourierZernike". Basically until the basis is padded #1243 (comment) there is no efficient implementation without loops.

github-actions · 2025-07-30T21:00:47Z

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    0.76 %    |     3.949e+03      |     3.979e+03      |    30.18     |       33.55        |       29.94        |
  test_proximal_jac_w7x_with_eq_update   |   -1.19 %    |     6.832e+03      |     6.751e+03      |    -81.02    |       161.84       |       160.53       |
  test_proximal_freeb_jac                |    0.07 %    |     1.320e+04      |     1.321e+04      |     9.22     |       78.41        |       76.71        |
  test_proximal_freeb_jac_blocked        |    0.41 %    |     7.602e+03      |     7.633e+03      |    30.99     |       67.50        |       68.60        |
  test_proximal_freeb_jac_batched        |   -0.65 %    |     7.619e+03      |     7.570e+03      |    -49.36    |       69.21        |       68.00        |
  test_proximal_jac_ripple               |   -0.49 %    |     7.550e+03      |     7.513e+03      |    -37.08    |       69.00        |       70.33        |
  test_proximal_jac_ripple_spline        |   -0.32 %    |     3.480e+03      |     3.469e+03      |    -11.02    |       72.02        |       71.52        |
  test_eq_solve                          |   -0.05 %    |     2.024e+03      |     2.023e+03      |    -1.04     |       124.18       |       124.57       |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

review-notebook-app · 2025-07-31T07:01:37Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

desc/grid.py

desc/objectives/_fast_ion.py

Co-authored-by: Dario Panici <37969854+dpanici@users.noreply.github.com>

f0uriest

Just so I understand:

map_coordinates doesn't change, just internally we use different shortcut for dealing with clebsch coordinates
map_clebsch_coordinates is now specialized to tensor product grids and uses partial summation to avoid re-evaluating the radial polynomials?

Can this also be used to speed up get_rtz_grid? since that's a meshgrid in clebsch coordinates?

desc/equilibrium/coords.py

docs/notebooks/tutorials/EffectiveRipple.ipynb

f0uriest · 2025-08-13T21:33:10Z

desc/equilibrium/equilibrium.py

-    def compute_theta_coords(
-        self, flux_coords, L_lmn=None, tol=1e-6, maxiter=20, full_output=False, **kwargs
+    @staticmethod
+    def _map_clebsch_coordinates(


does this need to be a method here? I'd vote for just using the function directly where needed

It has to be an object method to avoid circular imports.

f0uriest · 2025-08-13T21:35:21Z

desc/equilibrium/coords.py

+        ResolutionWarning,
+        msg="High frequency lambda modes will be truncated in coordinate mapping.",
+    )
+    lmbda_minus_iota_omega = L.transform(L_lmn)


why is this lambda-iota*omega and not just lambda?

The root-finding from desc coords to clebsch coords involves both lambda and omega

desc/equilibrium/coords.py

f0uriest · 2025-08-13T21:41:59Z

desc/equilibrium/coords.py

-                **kwargs,
-            )
+    @partial(jnp.vectorize, signature="(),(),(m)->()")
+    def vecroot(theta0, alpha, c_m):


I think you'll still have the issue @YigitElma mentioned about recompilation here. The fix would be to make rootfun/jacfun/vecroot global (private) functions rather than define them locally within this function

In ku/nufft I jitted the outer function as well.

This function should be fine, but you are welcome to fix map_coordinates.

desc/equilibrium/coords.py

review comment

unalmis · 2025-08-13T23:26:23Z

Just so I understand:
* `map_coordinates` doesn't change, just internally we use different shortcut for dealing with clebsch coordinates

* `map_clebsch_coordinates` is now specialized to tensor product grids and uses partial summation to avoid re-evaluating the radial polynomials?
Can this also be used to speed up get_rtz_grid? since that's a meshgrid in clebsch coordinates?

Yes. Yes, but also avoid evaluating toroidal series. Yes also for the third, but I left that as is for the following reason. Tthe partial summation implemented here has a totally unnecessary FourierZernike spectral to real transform then again an unnecessary $N^2$ FFT's of size $N$. To avoid this, I suggested padding the FourierZernike basis modes in issue 1243 to make the partial summation trivial. Until the proper partial summation is implemented, I don't want to change the API of get_rtz_grid.

unalmis · 2025-08-13T23:59:40Z

desc/equilibrium/equilibrium.py

+            **kwargs,
+        )
+
+    def _compute_iota_under_jit(self, rho, params=None, profiles=None, **kwargs):


this is used in map_coordinates and Bounce2D.compute_theta

rahulgaur104

Mostly minor comments. Addressing them will help future development/ developers.

rahulgaur104 · 2025-08-14T02:14:32Z

desc/equilibrium/coords.py

+                iota = eq._compute_iota_under_jit(coords, params, profiles, **kwargs)
+            rho, alpha, zeta = coords.T
+            omega = 0  # TODO(#568)
+            coords = jnp.column_stack([rho, alpha + iota * (zeta + omega), zeta])


If zeta is the generalized toroidal angle, don't we assume zeta = phi + Omega where phi is the cylindrical toroidal angle?
So shouldn't theta_PEST = alpha + iota * (zeta - omega)?

Phi = zeta + omega, and theta_PEST = theta + lambda. The left hand side of these relations are defined quantities. Phi is the cylindrical toroidal angle and theta_PEST is the angle where the field lines are straight in the (theta_PEST, Phi) plane. When we mention generalizing angles, we refer to changing the meaning of the angles "zeta" and "theta". These relations must still hold and hence the stream functions must negate the change in zeta and theta.

alpha = theta_PEST - iota phi

desc/equilibrium/coords.py

rahulgaur104 · 2025-08-14T02:28:28Z

desc/equilibrium/coords.py

-    guess=None,
-    period=np.inf,
+    lmbda,
+    theta0=None,


theta0 is a very bad variable name. theta0 is dangerously close to theta_0 that people in the gyrokinetics communit use in flux tube codes to define the location of vanishing integrated magnetic shear.
Why did you change it from guess? I strongly recomment you change it back. guess is great because it immediately tells you what it is. theta0 requires user to figure out what is happening in the rest of the code.

As a rule of thumb, please don't make unnecessary changes to variable names.

This is a new function, so nothing is an unnecessary change because nothing is a change; it is new.¹

theta0 is the initial guess in the root finding for theta. I will change it to guess as you requested. Whatever the name is, the public documentation next to the parameter theta0: jnp.ndarray : Optional initial guess for the computational coordinates. is there. So no one has to read code to figure out what it should be.

Footnotes

The function _map_clebsch_coordinates is private and not used anywhere in DESC. If there is external code that was using this function, we are not responsible for breaking API convention. ↩

rahulgaur104 · 2025-08-14T02:32:39Z

desc/equilibrium/coords.py

-    out : ndarray
-        Shape (k, 3).
-        DESC computational coordinates [ρ, θ, ζ].
-    info : tuple


Why did you remove info? This is not a good coding practice. This is basically passive encryption.

What happens if root finding fails for some reason and I want to debug it.

If the residual is calculated elsewhere, ignore this comment.

This is a new function, see #1826 (comment), so I disagree with the statement about coding practice. Whatever code users have that uses root finding, nothing has changed. They can still get their info tuple.

In this new function, I have not added functionality to return auxiliary information about the root finding because that is impossible --- functions that are decorated with jnp.vectorize, such as this one, must return arrays. info is not an array.

I am not convinced that this root finding can fail by the way.

rahulgaur104 · 2025-08-14T02:40:32Z

desc/integrals/_bounce_utils.py

@@ -79,8 +79,8 @@ def bounce_points(pitch_inv, knots, B, dB_dz, num_well=None):

    """
    intersect = polyroot_vec(


This routine find the intersection of lines of constant pitch with the magnetic field B, i.e., the bounce points.
B is an array with a shape rho, alpha, 1/lambda (inv_pitch), something, num_wells.
Is that right?

I put the shape of all my stuff in the public documentation string (docstring for short) because I find it helpful. Here is the definition of B and its shape in that docstring.

B : jnp.ndarray Shape (..., N - 1, B.shape[-1]). Polynomial coefficients of the spline of B in local power basis. Last axis enumerates the coefficients of power series. Second to last axis enumerates the polynomials that compose a particular spline.

When you see Shape (..., N - 1, B.shape[-1]), in Python, that means the code is making the following contract with you, the developer:

The last two axes need to have that shape. All the leading axes ... don't matter, the code is agonstic to them. If there are leading axes, the code will simply perform whatever it does to the last two axes in a vectorized manner. That is the same contract that numpy makes with you. Numpy calls this contract its "broadcasting conventions".

Less experienced developers will mess up this broadcasting contract, and instead may promise broadcasting on the trailing axes of the form Shape (N - 1, B.shape[-1], ...). You should NEVER do that. If you broadcast on the trailing axes like that you break from numpy convention, and both your code and the user-facing code will have to have a million transposes and reshapes to get their stuff working with what you wrote.

So what is the shape of B for this problem?
(..., B.shape[-3], B.shape[-2], B.shape[-1]) is an answer but I meant in terms of coordinates. Docstring does not have the shape here.

N is documented here as the number of knots. B.shape[-1] is documented here as the number of coefficients in each polynomial of the spline. If you have a cubic spline that is 4. The last axis of pitch is documented here to be the number of pitch angles. The returned shape is documented here to have its last two axes be number of pitch angles, number of wells, respectively.

rahulgaur104 · 2025-08-14T02:42:29Z

desc/integrals/bounce_integral.py

    -----
    Magnetic field line with label α, defined by B = ∇ψ × ∇α, is determined from
-      α : ρ, θ, ζ ↦ θ + λ(ρ,θ,ζ) − ι(ρ) [ζ + ω(ρ,θ,ζ)]
+      α : ρ, θ, ζ ↦ θ + Λ(ρ,θ,ζ) − ι(ρ) [ζ + ω(ρ,θ,ζ)]


I am confused again by this equation. Specifically the + sign before omega.
Is the convention zeta = phi + omega or zeta = phi - omega, where phi is the cylindrical angle and zeta is the generalized angle.

The latter. One can control +f search desc/compute/_core.py for name="phi" and name="theta_PEST" to see the definitions.

unalmis · 2025-08-14T05:23:28Z

I addressed the comments. I think at developer meeting you all wanted me to add more of my pull request comments to code. Please make a new pull request and add whichever of my comments into the code that you would like to add.

Progress toward addressing #1154. The [attached benchmark](https://github.com/user-attachments/files/21745848/benchmark_partial_sum.zip) must be ran on `ku/nufft` branch. 250x speed improvement from ~15 seconds to ~50 milliseconds. The improvement would be more if the FourierZernike basis is padded as discussed in #1243. Then we would avoid the 3D spectral to real transform as well as N^2 FFTs of size N. (Then this computation would likely be in microsecond range).

unalmis added 4 commits July 29, 2025 15:56

Merge changes from 1360

45693b9

Add NFP warning to eq.compute

479c497

Merge branch 'ku/NFP' into ku/partialsum

7181f79

first pass at partial sum

a85895a

unalmis changed the base branch from master to ku/NFP July 30, 2025 20:02

unalmis added the skip_changelog No need to update changelog on this PR label Jul 30, 2025

Base automatically changed from ku/NFP to master July 30, 2025 20:15

Merge branch 'master' into ku/partialsum

c390d07

unalmis commented Jul 30, 2025

View reviewed changes

unalmis added the performance New feature or request to make the code faster label Jul 30, 2025

unalmis added 3 commits July 30, 2025 19:52

working commit

c2a0e09

Merge branch 'master' into ku/partialsum

c12c884

Remove old static attributes

9e87fd1

unalmis self-assigned this Jul 31, 2025

unalmis added 3 commits July 31, 2025 01:15

partial sum pass two

3151bc7

Reduce resolution

c1cab79

Updated notebook

7dc7916

Dummy wrapper to avoid circular import

1bef805

unalmis marked this pull request as ready for review July 31, 2025 09:08

unalmis requested review from a team, YigitElma, ddudt, dpanici, f0uriest and rahulgaur104 and removed request for a team July 31, 2025 09:08

unalmis added 2 commits July 31, 2025 04:19

Update _fast_ion.py

f22e7d4

Cast to array first

39d1912

unalmis added 2 commits August 6, 2025 21:03

Merge branch 'master' into ku/partialsum

bdd174a

Merge branch 'master' into ku/partialsum

74f2b12

unalmis removed the run_benchmarks Run timing benchmarks on this PR against current master branch label Aug 13, 2025

unalmis linked an issue Aug 13, 2025 that may be closed by this pull request

Improve coordinate mapping performance #1154

Closed

Merge branch 'master' into ku/partialsum

03174ba

dpanici reviewed Aug 13, 2025

View reviewed changes

desc/grid.py Outdated Show resolved Hide resolved

dpanici reviewed Aug 13, 2025

View reviewed changes

desc/objectives/_fast_ion.py Show resolved Hide resolved

unalmis requested a review from dpanici August 13, 2025 21:40

dario comment suggestion

7be33eb

Co-authored-by: Dario Panici <37969854+dpanici@users.noreply.github.com>

f0uriest reviewed Aug 13, 2025

View reviewed changes

unalmis added 2 commits August 13, 2025 17:29

Pulling changes down from #1834 which are necessary to address @f0uriest

775bdf1

review comment

Add comment to address Rory comment

d5ff809

unalmis requested a review from f0uriest August 13, 2025 23:30

unalmis commented Aug 13, 2025

View reviewed changes

rahulgaur104 reviewed Aug 14, 2025

View reviewed changes

Changing variable name for Rahul

2b20f8d

unalmis requested a review from rahulgaur104 August 14, 2025 05:20

rahulgaur104 approved these changes Aug 14, 2025

View reviewed changes

Merge branch 'master' into ku/partialsum

fd17cc6

unalmis added the run_benchmarks Run timing benchmarks on this PR against current master branch label Aug 14, 2025

unalmis requested review from f0uriest and removed request for f0uriest August 14, 2025 19:26

unalmis merged commit 771b03c into master Aug 14, 2025
32 checks passed

unalmis deleted the ku/partialsum branch August 14, 2025 23:15

unalmis linked an issue Aug 14, 2025 that may be closed by this pull request

Fix #1334 #1411

Closed

unalmis mentioned this pull request Aug 21, 2025

Recompilation #1873

Open

unalmis mentioned this pull request Nov 16, 2025

Make mapping used for gyrokinetic sims public #2017

Open

		@@ -79,8 +79,8 @@ def bounce_points(pitch_inv, knots, B, dB_dz, num_well=None):

		"""
		intersect = polyroot_vec(

Partial summation in coordinate mapping #1826

Partial summation in coordinate mapping #1826

Uh oh!

Conversation

unalmis commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Memory benchmark result

Uh oh!

review-notebook-app bot commented Jul 31, 2025

Uh oh!

Uh oh!

Uh oh!

f0uriest left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

unalmis commented Aug 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rahulgaur104 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

unalmis Aug 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Footnotes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

unalmis Aug 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

unalmis Aug 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

unalmis commented Jul 30, 2025 •

edited

Loading

github-actions bot commented Jul 30, 2025 •

edited

Loading

unalmis commented Aug 13, 2025 •

edited

Loading

unalmis Aug 14, 2025 •

edited

Loading

unalmis Aug 14, 2025 •

edited

Loading

unalmis Aug 14, 2025 •

edited

Loading

unalmis Aug 14, 2025 •

edited

Loading