num_neighbors pruning for CoilSetMinDistance

[TElder-thea]

Add optional num_neighbors parameter to the CoilSetMinDistance objective. This limits pairwise distance computation to only the num_neighbors nearest neighbors per coil, selected by centroid distance. Neighbor indices use jax.lax.stop_gradient so AD only traces through the fine-grained distances to nearby coils.

Depending on coilset used, we found:

• 16.5x speedup on CoilSetMinDistance Jacobian alone
• 5.8x full auglag optimizer speedup (468s -> 80s)
• 3.5x full lsq-exact optimizer speedup (464s -> 134s)
• Enables jac_chunk_size=32 (baseline OOMs at chunk=6)
• Machine-precision correctness (K=5 matches full at alpha=100)
  For more generic modular coil optimizations, pruning may decrease performance due to overhead.

Usage: CoilSetMinDistance(..., n_neighbors=20)
Default n_neighbors=None preserves existing behavior.

[ddudt]

Remember to also update CHANGELOG.md

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    6.13 %    |     4.040e+03      |     4.288e+03      |    247.67    |       43.28        |       41.39        |
  test_proximal_jac_w7x_with_eq_update   |   -0.81 %    |     6.579e+03      |     6.526e+03      |    -53.34    |       172.42       |       171.81       |
  test_proximal_freeb_jac                |    0.20 %    |     1.337e+04      |     1.340e+04      |    26.24     |       92.78        |       91.27        |
  test_proximal_freeb_jac_blocked        |    0.56 %    |     7.710e+03      |     7.753e+03      |    43.03     |       82.23        |       82.34        |
  test_proximal_freeb_jac_batched        |    0.51 %    |     7.671e+03      |     7.710e+03      |    39.00     |       81.44        |       82.45        |
  test_proximal_jac_ripple               |    1.01 %    |     3.573e+03      |     3.609e+03      |    36.03     |       67.22        |       66.60        |
  test_proximal_jac_ripple_bounce1d      |    0.26 %    |     3.821e+03      |     3.831e+03      |     9.98     |       81.37        |       81.89        |
  test_eq_solve                          |   -0.84 %    |     2.232e+03      |     2.214e+03      |    -18.64    |       101.07       |       101.48       |

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.45%. Comparing base (2471d55) to head (39aaca5).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2140      +/-   ##
==========================================
- Coverage   94.45%   94.45%   -0.01%     
==========================================
  Files         101      101              
  Lines       28593    28604      +11     
==========================================
+ Hits        27008    27018      +10     
- Misses       1585     1586       +1     

Files with missing lines	Coverage Δ
desc/objectives/_coils.py	99.34% <100.00%> (+<0.01%)	⬆️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ddudt]

Note for reviewers: Some of the original code was written by Claude, but I re-wrote almost everything as a human.

Changes made

[YigitElma]

I am not sure about the behavior when the neighboring coils change, but I guess it is similar to how jnp.min or softmin behaves under similar conditions.

It looks good to me.

[YigitElma]

Should this be true if one of the neighboring coils is about to change?

[ddudt]

Yes I believe so. Both the original method and this new pruning should have the same behavior for these derivatives, which doesn't fully capture what would happen if the closest coil changes.

[YigitElma]

If I get confused in the future again, here is a small example,

Details

from desc.objectives.utils import softmin
def fun(x):
    f = x**2
    g = x
    return softmin(jnp.array([f, g]), alpha=8)

x = jnp.linspace(-0.5, 1.5, 300)
y = jax.vmap(fun)(x)
dy = jax.vmap(jax.grad(fun))(x)
plt.title("softmin with alpha=8")
plt.plot(x, x, "--", label="g", lw=2)
plt.plot(x, x**2, "--", label="f", lw=2)
plt.plot(x, y, label="softmin(f, g)", lw=3)
plt.plot(x, dy, label="dy/dx")
plt.legend()

def fun(x):
    f = x**2
    g = x
    return jnp.min(jnp.array([f, g]))

x = jnp.linspace(-0.5, 1.5, 300)
y = jax.vmap(fun)(x)
dy = jax.vmap(jax.grad(fun))(x)
plt.title("min(f, g)")
plt.plot(x, x, "--", label="g", lw=2)
plt.plot(x, x**2, "--", label="f", lw=2)
plt.plot(x, y, label="min(f, g)", lw=3)
plt.plot(x, dy, label="dy/dx")
plt.legend()

<\details