Handling of coordinate conversions outside LCFS in eq.map_coordinates

[amicitas]

I am trying to use desc to help me convert real-space coordinates to flux-space coordinates for a synthetic diagnostic implementation. For points that fall within the last closed flux surface (LCFS) desc seems to work quite well using eq.map_coordinates. However, since I am starting from real-space, I need to sample points from a much wider area and currently desc is both inefficient and unpredictable for points outside the LCFS (see attached figure for W7-X).

Are there any known solutions to improve handling of eq.map_coordinates for handling of inverse coordinate transforms outside the LCFS? What I would like is a fast calculation of NaN or a mask for requests outside of the LCFS, rather than the slow and unpredictable result from the solver.

Sam Lazerson and I were able to implement such a filter within the LIBSTELL library (stelltools.f90/mirtools.f90), but I don't know if that approach works within desc.

-- Novimir

[YigitElma]

eq.map_coordinates does a small optimization to find $\rho, \theta, \phi$ flux coordinates that minimize the error between the evaluated and desired $R, \phi, Z$ points. Since this doesn't have an actual minimum for the points outside of the LCFS, the poor performance is expected.

I don't think we currently have a mask for this purpose, but I think you can create one using the distance to the magnetic axis. At points $\phi$ coordinate, evaluate $\sqrt{(R-R0)^2 + (Z-Z0)^2}$ and if it is bigger than some threshold, it can be masked. This is kind of tricky to do exactly, because in general DESC's native coordinate $\theta$ doesn't correspond to the polar angle with the Z=0 plane, so a threshold can only remove most of the points outside of the LCFS but not all without special care.

Given that, you can use maxiter and the output of the map_coordinates to have some idea whether point is inside or outside. For example, here we use similar approach to deduce if the point is in or out. Usually, points that are inside converges quite fast, so by capping the maxiter, long compute time can be reduced.

[dpanici]

Can you share a sample of the code you tried that lead to this? We then can help more with the implementation of the solution

[YigitElma]

@amicitas, any update on this?

[amicitas]

Through the processes of making a 'minimal example', I've realized this issue should probably be split into two parts:
1.) Return nan (or a masked array, or rho=1.0, etc.) for values outside of the LCFS.
2.) Optimize performance for attempted coordinate transformations outside of the LCFS.

For part 1, it's easy to create a 'minimal code' example, and an 'example fix':

• do a round-trip check (which is cheap compared to the inverse calculation).

For part 2, I'm still working on understanding performance impacts in way that is appropriate for a desc issue.

[amicitas]

Here is a simplified Jupyter notebook that captures part 1) of this issue, "unexpected values outside of the LCFS", and includes a round-trip check (which is close but not 100% successful in this demonstration code).

2026-07-23 - Desc Coordinate Transform Performance Minimal Example - Part 3d - preparation for desc github issue 2259.ipynb

I include some timing information, but this is not particularly useful at the moment.

Without round-trip filtering:

With filtering (not perfect):

Timing (MacBook Pro, M1; jax but no GPU):

step                            seconds    % total
--------------------------------------------------
2. load equilibrium               1.160      24.5%
3a. sampling                      0.001       0.0%
3b. map_coordinates (cold)        1.749      36.9%
3c. round-trip check              0.452       9.5%
3b. map_coordinates (warm)        1.377      29.1%
--------------------------------------------------
TOTAL                             4.738     100.0%

[dpanici]

Ok I think I understand the issue a little better now. So, the map_coordinates function is really meant for within-the-equilibrium coordinate conversions right now. I am not sure of a good way to increase the performance when a point outside the LCFS is requested to be mapped to flux coordinates, not without just doing like what you've done and filtering. Doing so within map_coordinates would basically require one to check if the real-space-coordinate input is within the LCFS or not, which in 3D is a generally difficult problem (it is possible, just involves computation which might be more expensibe than just letting the rootfind fail). If the points are all within the same or in just a few Phi planes, we have a utility to find if point is within the boundary at that phi point. I'll look later today or tomorrow at how expensive this is to do for a disordered array of inputs (not just at a few Phi points), and we could return NaN if it is outside the lcfs there or use it as a filter in your own work.

Perhaps a slightly better version of the way you shared could be to filter based off the residual of the rootfind (which can be returned from map_coordinates by passing full_output=True) and assign NaN to failed rootfinds defined by those with large residuals (this assumes that if a rootfind fails, it is due to the point not being within the LCFS). This of course does not remove the cost of the fruitless rootfinds, but it would allow you to flag the otherwise unexpected values you've found outside the LCFS and post-process remove them or assign them as NaN (since the residual in the rootfind will be very large compared to successful rootfinds within the LCFS)

How fast do you need this for your application?

Something similar to what STELLOPT has implemented is probably possible, assuming (from my skim of the source) what was done there is a spline map from (R,phi,Z) -> (s, theta, zeta) was pre-computed, with a linear extrapolation out from the boundary in s, and if s >1 from the mapping then one takes the point as outside the boundary. It would have a similar pre-computation cost as the STELLOPT method does though

[amicitas]

I had a video chat with Dario today, and he provided many helpful suggestions on how to accelerate coordinate conversions for my specific workflow.

I have tried to do a bit more profiling to see if my hypothesis that the points outside rho=1 were slowing the calculation down, and my conclusion (based on minimal testing) is that this is not the case beyond having a lot of points outside rho=1. While all the points outside rho=1 do take the full 30 iterations, they don't actually take much more time per point.

Timing (now accounting for jax async execution by using block_until_ready(). This matches my real production experience:

step                            seconds    % total
--------------------------------------------------
2. load equilibrium               0.491       3.2%
3a. sampling                      0.002       0.0%
3b. map_coordinates (cold)        7.352      48.2%
3c. round-trip check              0.504       3.3%
3b. map_coordinates (warm)        6.888      45.2%
--------------------------------------------------
TOTAL                            15.237     100.0%

Here is the separate timing for the 'inside' and 'outside' points:

set            n   cold [s]   warm [s]   warm ms/pt  mean niter  max niter
--------------------------------------------------------------------------
inside      2564      2.125      1.770        0.690         2.9         11
outside     7436      6.072      5.429        0.730        30.0         30
--------------------------------------------------------------------------

Reducing the number of iteration does improve performance.

MAXITER = 15

step                            seconds    % total
--------------------------------------------------
2. load equilibrium               0.482       4.9%
3a. sampling                      0.001       0.0%
3b. map_coordinates (cold)        4.874      49.1%
3c. round-trip check              0.432       4.3%
3b. map_coordinates (warm)        4.142      41.7%
--------------------------------------------------
TOTAL                             9.931     100.0%

MAXITER = 15

set            n   cold [s]   warm [s]   warm ms/pt  mean niter  max niter  # at cap
------------------------------------------------------------------------------------
inside      2564      2.034      1.708        0.666         2.9         11         0
outside     7436      3.460      3.343        0.450        15.0         15      7436
------------------------------------------------------------------------------------

Here are some plots of the number of iterations and the residual (note log-scale on the residual):

[dpanici]

@dpanici
see how expensive it is to check per-toroidal angle (with potentially some max number of toroidal bins as we don't want 10000 unique toroidal planes, and snap input pts to nearest bin for inside v outside lcfs purposes) if the pts are inside v outside the LCFS. If outside, dont rootfind there and just return nan or something like that