Increase number of concentric grid nodes near axis

[dpanici]

• Makes the innermost ring and the 2nd-innermost ring of points in ConcentriGrid have the same points (i.e. make the inner ring poloidal points match those of the 2nd most inner ring) so that flux surface averages are more correct there.
  Tests seem to show that this improves equilibria and free boundary in most cases

Resolves #1876
Resolves #1985

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    3.81 %    |     3.835e+03      |     3.981e+03      |    145.94    |       40.27        |       36.66        |
  test_proximal_jac_w7x_with_eq_update   |   -2.78 %    |     6.517e+03      |     6.336e+03      |   -180.88    |       162.69       |       163.02       |
  test_proximal_freeb_jac                |   -0.19 %    |     1.318e+04      |     1.316e+04      |    -24.64    |       83.83        |       84.13        |
  test_proximal_freeb_jac_blocked        |    0.04 %    |     7.506e+03      |     7.509e+03      |     2.84     |       73.60        |       72.49        |
  test_proximal_freeb_jac_batched        |    0.15 %    |     7.431e+03      |     7.442e+03      |    11.46     |       72.78        |       72.70        |
  test_proximal_jac_ripple               |   -2.18 %    |     3.546e+03      |     3.468e+03      |    -77.38    |       65.10        |       65.90        |
  test_proximal_jac_ripple_bounce1d      |   -0.49 %    |     3.572e+03      |     3.554e+03      |    -17.45    |       76.35        |       75.75        |
  test_eq_solve                          |   -0.72 %    |     2.023e+03      |     2.009e+03      |    -14.55    |       94.09        |       94.07        |

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

[codecov]

Codecov Report

❌ Patch coverage is 83.33333% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 95.77%. Comparing base (a846c9b) to head (2648b7a).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/grid.py	83.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1877      +/-   ##
==========================================
- Coverage   95.77%   95.77%   -0.01%     
==========================================
  Files         101      101              
  Lines       27796    27799       +3     
==========================================
+ Hits        26622    26624       +2     
- Misses       1174     1175       +1     

Files with missing lines	Coverage Δ
desc/grid.py	94.60% <83.33%> (-0.12%)	⬇️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[dpanici]

running an analyti solovev current constrained that shuold have iota=2 with different grids, seems the 3 points is better but not as good as quadgrid or as using a separate iota grid, and notably took twice as long as teh 2x quadgrid to get down to the gtol that I set these to solve to:

the 1.5x quadgrid took less time than 2x quadgrid to get to the gtol I set, and is only slightly less accurate as measured by iota.

Running some stellarators now to see how this holds there.

times (roughly timed on my laptop so maybe useless tbh)

conc 1 pt with separate iota grid: 8.5s (hit 100 iters, not gtol, if allowed to go longer hits gtol at ~16s)
conc 1 pt: 4 (hit gtol)
conc 3 pt: 5s (hit 100 iters, if allowed to go longer hits gtol at ~16s)
quadgrid 2x res: 7s (hit gtol)
quadgrid 1.5x res: 5s (hit gtol)

[dpanici]

Plot dump here, for a few current constrained equilibria I ran the following solve scripts to solve with various grids and resolutions, where I set initial guess (and lambda to zero, as set initial guess does not zero out lambda), then ran with differing grids.

Then loaded each, plotted the fsa of |F|, |B|, and then plotted iota, all computed at high resolution of like 2*eq.X_grid (so twice as high as usual resolution). (pics btwn this and next post).

tl;dr is that it seems that quadgrid at 1.25x res (i.e. like 1.25 x eq.L) performs as good if not better than concentric grid at 2x res for these stellarator cases, for sym=False this is half the nodes roughly, for sym=True these are roughly the same amount of nodes (not sure how the sym results are yet though)

for sym=False though seems that quadgrid at 1.25 or 1.5x res would be worth using

I did just realize I ran all of these with sym=False for concentric, so will re-run with sym=True and see how things change

         name|  conc 1.00xres | quad 1.00xres |  conc 1.25xres | quad 1.25xres |  conc 1.50xres | quad 1.50xres |  conc 1.75xres | quad 1.75xres |  conc 2.00xres | quad 2.00xres | 
   precise_QA|            765|           1445|           1386|           2646|           2275|           4375|           3480|           6728|           5049|           9801
   precise_QH|            765|           1445|           1386|           2646|           2275|           4375|           3480|           6728|           5049|           9801
       ESTELL|           2275|           4375|           4216|           7688|           7030|          13690|          10879|          20339|          15925|          31213
         NCSX|           4225|           8125|           7936|          15376|          13357|          26011|          20812|          40678|          30625|          60025
 DSHAPE_current|           5292|          10206|           9537|          18513|          18040|          33620|          27072|          53016|          38637|          75843

concentric (one pt on axis, so like master):

for name in ["precise_QA","precise_QH","ESTELL","NCSX","DSHAPE_current"]:
    eq = desc.examples.get(name)
    eq0 = eq.copy()
    eq0.set_initial_guess()
    eq0.L_lmn = jnp.zeros_like(eq0.L_lmn)
    print("#"*20)
    print(name)
    print("#"*20)
    for mult in [1,1.25,1.5,1.75,2]:
        print("#"*20)
        print(name)
        print(mult)
        print("#"*20)
        fname = f"eqs_concgrid/{name}_concgrid_{mult}_x_res_one_pt_axis.h5"
        if not os.path.exists(fname):
            eq = eq0.copy()
            obj = ObjectiveFunction(ForceBalance(eq,grid=ConcentricGrid(L=round(mult*eq.L),
                                                                    M=round(mult*eq.M),
                                                                    N=round(mult*eq.M),
                                                                    NFP=eq.NFP)),jac_chunk_size=100)
            t0 = time.time()
            eq.solve(verbose=3, objective=obj,ftol=0,xtol=0, maxiter=300)
            t = time.time()-t0
            eq.save(f"eqs_concgrid/{name}_concgrid_{mult}_x_res_one_pt_axis.h5")
            np.savetxt(f"eqs_concgrid/{name}_concgrid_{mult}_x_res_time.txt", np.atleast_1d(t))

quadgrid

for name in ["precise_QA","precise_QH","ESTELL","NCSX","DSHAPE_current"]:
    eq = desc.examples.get(name)
    eq0 = eq.copy()
    eq0.set_initial_guess()
    eq0.L_lmn = jnp.zeros_like(eq0.L_lmn)
    print("#"*20)
    print(name)
    print("#"*20)
    
    for mult in [1,1.25,1.5,1.75,2]:
        print("#"*20)
        print(name)
        print(mult)
        print("#"*20)
        fname = f"eqs_quadgrid/{name}_quadgrid_{mult}_x_res.h5"
        if not os.path.exists(fname):
            eq = eq0.copy()
            obj = ObjectiveFunction(ForceBalance(eq,grid=QuadratureGrid(L=round(mult*eq.L),
                                                                    M=round(mult*eq.M),
                                                                    N=round(mult*eq.M),
                                                                    NFP=eq.NFP)),
                                                                    jac_chunk_size=100)
            
            t0 = time.time()
            eq.solve(verbose=3, objective=obj,ftol=0,xtol=0, maxiter=300)
            t = time.time()-t0

            eq.save(fname)
            np.savetxt(f"eqs_quadgrid/{name}_quadgrid_{mult}_x_res_time.txt", np.atleast_1d(t))

plots made by this script

plt.rcParams.update({"font.size":14})
cs = ["k","r","m","c","g"]
for name in ["precise_QA","precise_QH","ESTELL","NCSX","DSHAPE_current"]:
    mults = [1,1.25,1.5,1.75,2]
    eqs_conc = []
    ts_conc = []
    eqs_quad = []
    ts_quad = []
    for i,mult in enumerate([1,1.25,1.5,1.75,2]):
        eq_conc = load(f"eqs_concgrid/{name}_concgrid_{mult}_x_res_one_pt_axis.h5")
        t_conc=np.genfromtxt(f"eqs_concgrid/{name}_concgrid_{mult}_x_res_time.txt").squeeze()
        eqs_conc.append(eq_conc.copy())
        ts_conc.append(t_conc.copy())
        
                        
        eq_quad = load(f"eqs_quadgrid/{name}_quadgrid_{mult}_x_res.h5")
        t_quad=np.genfromtxt(f"eqs_quadgrid/{name}_quadgrid_{mult}_x_res_time.txt").squeeze()
        eqs_quad.append(eq_quad.copy())
        ts_quad.append(t_quad.copy())
        
    grid = LinearGrid(L=25, M=2*eq_conc.M_grid, N=2*eq_conc.N_grid, NFP=eq_conc.NFP,axis=False)
    
    fig_f,ax_f = plt.subplots(figsize=(10,10))
    for i,mult in enumerate([1,1.25,1.5,1.75,2]):
        fig_f,ax_f = plot_fsa(eqs_conc[i], "|F|_normalized",grid=grid,log=True, label=f"{name} conc 1pt {mult} x res t={t_conc:.1f} s",ax=ax_f,color=cs[i],lw=3)
        fig_f,ax_f = plot_fsa(eqs_quad[i], "|F|_normalized",grid=grid,log=True, label=f"{name} quad {mult} x res t={t_quad:.1f} s",ax=ax_f,color=cs[i],ls="--",lw=3)
        ax_f.set_title(name)
    plt.gcf().savefig(f"{name}_F_fsa_comparison.png")
    plt.close("all")  
    fig_b,ax_b = plt.subplots(figsize=(10,10))
    for i,mult in enumerate([1,1.25,1.5,1.75,2]):
        fig_b,ax_b = plot_fsa(eqs_conc[i], "|B|",grid=grid, label=f"{name} conc 1pt {mult} x res t={t_conc:.1f} s",ax=ax_b, color=cs[i],lw=3)
        fig_b,ax_b = plot_fsa(eqs_quad[i], "|B|",grid=grid,label=f"{name} quad {mult} x res t={t_quad:.1f} s",ax=ax_b, color=cs[i],ls="--",lw=3)
        ax_b.set_title(name)
    plt.gcf().savefig(f"{name}_B_fsa_comparison.png")
    plt.close("all")
    fig_i,ax_i = plt.subplots(figsize=(10,10))
    for i,mult in enumerate([1,1.25,1.5,1.75,2]):
        fig_i,ax_i = plot_fsa(eqs_conc[i], "iota",grid=grid, label=f"{name} conc 1pt {mult} x res t={t_conc:.1f} s",ax=ax_i,color=cs[i],lw=3)
        fig_i,ax_i = plot_fsa(eqs_quad[i], "iota",grid=grid, label=f"{name} quad {mult} x res t={t_quad:.1f} s",ax=ax_i,color=cs[i],ls="--",lw=3)
        ax_i.set_title(name)
    plt.gcf().savefig(f"{name}_iota_comparison.png")    
    plt.close("all")

Plots

[dpanici]

Details

[dpanici]

symmetric grid, results are less clear here, the concentric seems to do fine for these cases for the most part barring precise QA, must be something weird abour the initial point there that concentric struggles with

Details

[dpanici]

Details

[dpanici]

with 3 pts near-axis for concentric (this PR)

definitely better than one pt

Details

[dpanici]

some more

Details

[dpanici]

One more test case strengthening using more pts, specifically quadgrid: a free bdry finite beta ellipse starting from a finite beta, zero current tokamak. With usual one-pt concentric grid, it converges to something not entirely correct (the inner flux surfaces seem inaccurate, with a weird shape change you wouldn't expect given the simple geometry). The optimization also ends due to stalling bc of bad approx errors.

with quad grid, things look good, shafranov shift appears correct, and opt never stalls

with 3 pt conc grid, things don't look right and it ends on a stall after a few iterations:

the 3pt conc grid optimization has same behavior even if iota is correctly computed on a separate grid which assumes no symmetry

[dpanici]

Some more with SOLOVEV analytic current-constrained: in this case 3-pt concentric definitely essentially puts the solution on par with the quadgrid:
one-pt (the iota of the blue concentric soln is wrong, should be 2 almost exactly like the quadgrid)

3-pt concentric:

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     -2.51 +/- 4.36     | -2.12e-02 +/- 3.69e-02 |  8.25e-01 +/- 3.6e-02  |  8.46e-01 +/- 7.1e-03  |
 test_build_transform_fft_highres        |     -2.95 +/- 1.71     | -3.29e-02 +/- 1.91e-02 |  1.08e+00 +/- 6.0e-03  |  1.11e+00 +/- 1.8e-02  |
 test_equilibrium_init_lowres            |     -3.45 +/- 2.06     | -2.05e-01 +/- 1.23e-01 |  5.74e+00 +/- 8.6e-02  |  5.95e+00 +/- 8.7e-02  |
 test_objective_compile_atf              |     -1.43 +/- 2.69     | -1.12e-01 +/- 2.10e-01 |  7.70e+00 +/- 7.5e-02  |  7.81e+00 +/- 2.0e-01  |
 test_objective_compute_atf              |     +1.07 +/- 13.39    | +2.32e-05 +/- 2.91e-04 |  2.19e-03 +/- 1.9e-04  |  2.17e-03 +/- 2.2e-04  |
 test_objective_jac_atf                  |     -1.21 +/- 3.86     | -2.15e-02 +/- 6.85e-02 |  1.75e+00 +/- 5.4e-02  |  1.77e+00 +/- 4.2e-02  |
 test_perturb_1                          |     +2.20 +/- 3.91     | +3.25e-01 +/- 5.76e-01 |  1.51e+01 +/- 5.5e-01  |  1.47e+01 +/- 1.8e-01  |
 test_proximal_jac_atf                   |     +0.76 +/- 1.10     | +4.17e-02 +/- 6.05e-02 |  5.53e+00 +/- 5.0e-02  |  5.49e+00 +/- 3.3e-02  |
 test_proximal_freeb_compute             |     -0.43 +/- 2.86     | -7.10e-04 +/- 4.72e-03 |  1.64e-01 +/- 3.8e-03  |  1.65e-01 +/- 2.8e-03  |
 test_solve_fixed_iter                   |     -0.51 +/- 3.71     | -1.44e-01 +/- 1.05e+00 |  2.83e+01 +/- 6.5e-01  |  2.84e+01 +/- 8.3e-01  |
 test_objective_compute_ripple           |     +1.44 +/- 3.50     | +2.98e-03 +/- 7.27e-03 |  2.11e-01 +/- 4.4e-03  |  2.08e-01 +/- 5.8e-03  |
 test_objective_grad_ripple              |     -1.90 +/- 2.87     | -1.73e-02 +/- 2.61e-02 |  8.93e-01 +/- 2.0e-02  |  9.10e-01 +/- 1.7e-02  |
 test_build_transform_fft_lowres         |     -0.40 +/- 1.81     | -2.94e-03 +/- 1.32e-02 |  7.26e-01 +/- 1.1e-02  |  7.29e-01 +/- 7.6e-03  |
 test_equilibrium_init_medres            |     -0.42 +/- 1.67     | -2.53e-02 +/- 9.96e-02 |  5.95e+00 +/- 5.3e-02  |  5.97e+00 +/- 8.4e-02  |
 test_equilibrium_init_highres           |     +0.02 +/- 1.56     | +1.58e-03 +/- 1.05e-01 |  6.72e+00 +/- 9.2e-02  |  6.72e+00 +/- 4.9e-02  |
 test_objective_compile_dshape_current   |     +2.15 +/- 8.44     | +8.66e-02 +/- 3.40e-01 |  4.11e+00 +/- 2.7e-01  |  4.03e+00 +/- 2.1e-01  |
 test_objective_compute_dshape_current   |     +1.38 +/- 8.66     | +9.96e-06 +/- 6.23e-05 |  7.29e-04 +/- 4.9e-05  |  7.19e-04 +/- 3.9e-05  |
 test_objective_jac_dshape_current       |     +2.91 +/- 14.26    | +7.47e-04 +/- 3.66e-03 |  2.64e-02 +/- 2.2e-03  |  2.57e-02 +/- 3.0e-03  |
 test_perturb_2                          |     -0.49 +/- 2.24     | -8.95e-02 +/- 4.08e-01 |  1.81e+01 +/- 2.0e-01  |  1.82e+01 +/- 3.6e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.71 +/- 0.29     | +9.39e-02 +/- 3.78e-02 |  1.33e+01 +/- 2.0e-02  |  1.32e+01 +/- 3.2e-02  |
 test_proximal_freeb_jac                 |     -0.22 +/- 2.83     | -1.07e-02 +/- 1.38e-01 |  4.87e+00 +/- 5.7e-02  |  4.88e+00 +/- 1.3e-01  |
 test_solve_fixed_iter_compiled          |     -1.60 +/- 2.57     | -1.28e-01 +/- 2.05e-01 |  7.87e+00 +/- 1.2e-01  |  7.99e+00 +/- 1.6e-01  |
 test_LinearConstraintProjection_build   |     -0.41 +/- 3.25     | -3.36e-02 +/- 2.69e-01 |  8.24e+00 +/- 2.2e-01  |  8.27e+00 +/- 1.6e-01  |
 test_objective_compute_ripple_bounce1d  |     -1.98 +/- 3.05     | -5.91e-03 +/- 9.09e-03 |  2.92e-01 +/- 5.0e-03  |  2.98e-01 +/- 7.6e-03  |
 test_objective_grad_ripple_bounce1d     |     -0.76 +/- 4.20     | -7.33e-03 +/- 4.05e-02 |  9.57e-01 +/- 1.1e-02  |  9.65e-01 +/- 3.9e-02  |

Github CI performance can be noisy. When evaluating the benchmarks, developers should take this into account.

[YigitElma]

Some others should also review since I contributed to the PR.

[f0uriest]

it feels weird that we need M this high for a simple vacuum equilibrium. All the surfaces are nearly perfect ellipses. (even 14 seems high for such a simple eq)

[dpanici]

try lower resolution, 4.8 maybe

[YigitElma]

4-8 and 6-10 don't pass. I set it to 6-12.

[f0uriest]

Right now, if you create a Zernike basis and a concentric grid with the same L,M you get a square transform matrix (at least for L=M, there might be some edge cases where L!=M because of the different indexing patterns). With this change I don't think we'll be able to get a square matrix except for maybe very special choices of L,M

I see the benefit of adding the additional nodes, but I'm wondering if we want to do it in all cases, or maybe some other options:

• maybe don't add the extra nodes if pattern="ocs" since that was designed for fitting with a square system?
• Adding a flag to use the old number of nodes?
• making a new pattern "jacobi2" or something that is jacobi with the additional nodes near axis?

[dpanici]

I prefer jacobi-XXX as new default that has more nodes near axis

[YigitElma]

I think jacobi-dna is cool. DNA stands for dense near axis

[dpanici]

check if input files dont have jacobi pattern listed still. Also check tutorials

[YigitElma]

We have made the plot related changes way before in #1683

All requests have been addressed

[dpanici]

@dpanici all notebooks need to be re-run with this PR since we change the default grid

[dpanici]

get #1907 in first then

• update the example eqs here (@dpanici )
• update all the notebooks as well once example eqs are updated (since the get(examples) might be different now) (@YigitElma or @dpanici )
• update test_compute_everything (@YigitElma or @dpanici )

[YigitElma]

@YigitElma problematic test test_NAE_QSC_solve_near_axis_asym takes 35 minutes to run. Make it faster, maybe merge it with some other similar tests (like test_NAE_QSC_solve_near_axis_based_off_eq_asym)

[dpanici]

@YigitElma problematic test test_NAE_QSC_solve_near_axis_asym takes 35 minutes to run. Make it faster, maybe merge it with some other similar tests (like test_NAE_QSC_solve_near_axis_based_off_eq_asym)

can we do this separately from here? since some random other PRs seem to failnow due to this test too