Reverse-Mode Differentiable Ideal-Ballooning Stability Solver

[rahulgaur104]

Infinite-n ideal MHD ballooning modes are of significant interest to both the tokamak and the stellarator community. These instabilities are also related to smaller-scale kinetic instabilities, which cause significant heat loss from fusion reactors. This commit adds the ability to both analyze and optimize MHD equilibria against the ideal ballooning mode.

An adjoint implementation of an ideal ballooning solver has only been successfully demonstrated previously on a much smaller scale. If this commit works, it would be at least an order of magnitude faster than the linked code.

Here are the tasks:

• [x] Add geometry coefficients
• [x] Test geometric coefficients. Maybe compare it with Patrick's fork in DESC. Add a test
• [x] Compare results with each other (ideal_ball_gamma1, ideal_ball_gamma2, Newcomb_metric) and with my older finite-difference solver
• [x] Eigendecomposition and adjoint of the matrices. Wrap the routine in DESC magic.
• [x] Compare results for the HELIOTRON case with COBRAVMEC
• [x] Write a test for HELIOTRON comparison with COBRAVMEC
• [x] Write a test for the Newcomb metric

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 99.47917% with 1 line in your changes missing coverage. Please review.

Project coverage is 95.43%. Comparing base (17c2b15) to head (ff85459). Report is 2240 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/objectives/_stability.py	98.11%	1 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1170      +/-   ##
==========================================
+ Coverage   95.37%   95.43%   +0.05%     
==========================================
  Files          96       96              
  Lines       23560    23735     +175     
==========================================
+ Hits        22471    22652     +181     
+ Misses       1089     1083       -6     

Files with missing lines	Coverage Δ
desc/compute/_core.py	100.00% <100.00%> (ø)
desc/compute/_metric.py	100.00% <100.00%> (ø)
desc/compute/_stability.py	100.00% <100.00%> (ø)
desc/equilibrium/coords.py	88.38% <ø> (ø)
desc/equilibrium/equilibrium.py	96.14% <100.00%> (+0.42%)	:arrow_up:
desc/objectives/__init__.py	100.00% <100.00%> (ø)
desc/objectives/_stability.py	98.48% <98.11%> (-0.29%)	:arrow_down:

... and 3 files with indirect coverage changes

[github-actions[bot]]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -0.72 +/- 2.32     | -3.81e-03 +/- 1.23e-02 |  5.26e-01 +/- 1.1e-02  |  5.30e-01 +/- 5.2e-03  |
 test_equilibrium_init_medres            |     -0.45 +/- 0.63     | -1.89e-02 +/- 2.60e-02 |  4.13e+00 +/- 2.0e-02  |  4.15e+00 +/- 1.7e-02  |
 test_equilibrium_init_highres           |     -0.22 +/- 0.67     | -1.20e-02 +/- 3.70e-02 |  5.48e+00 +/- 2.7e-02  |  5.49e+00 +/- 2.5e-02  |
 test_objective_compile_dshape_current   |     -0.26 +/- 5.97     | -1.02e-02 +/- 2.32e-01 |  3.88e+00 +/- 2.3e-01  |  3.89e+00 +/- 2.5e-02  |
 test_objective_compute_dshape_current   |     -0.05 +/- 1.55     | -1.71e-06 +/- 5.61e-05 |  3.62e-03 +/- 3.7e-05  |  3.62e-03 +/- 4.2e-05  |
 test_objective_jac_dshape_current       |     -0.32 +/- 6.18     | -1.31e-04 +/- 2.52e-03 |  4.06e-02 +/- 1.9e-03  |  4.08e-02 +/- 1.6e-03  |
 test_perturb_2                          |     -0.56 +/- 0.95     | -9.81e-02 +/- 1.68e-01 |  1.75e+01 +/- 4.3e-02  |  1.76e+01 +/- 1.6e-01  |
 test_proximal_freeb_jac                 |     -0.55 +/- 1.35     | -4.14e-02 +/- 1.02e-01 |  7.51e+00 +/- 4.4e-02  |  7.55e+00 +/- 9.1e-02  |
 test_solve_fixed_iter                   |     -0.59 +/- 58.01    | -2.97e-02 +/- 2.91e+00 |  4.99e+00 +/- 2.0e+00  |  5.02e+00 +/- 2.1e+00  |
 test_build_transform_fft_midres         |     -0.34 +/- 5.98     | -2.17e-03 +/- 3.81e-02 |  6.35e-01 +/- 3.4e-02  |  6.37e-01 +/- 1.8e-02  |
 test_build_transform_fft_highres        |     +0.76 +/- 3.04     | +7.67e-03 +/- 3.08e-02 |  1.02e+00 +/- 2.4e-02  |  1.02e+00 +/- 1.9e-02  |
 test_equilibrium_init_lowres            |     -3.90 +/- 7.71     | -1.62e-01 +/- 3.20e-01 |  3.99e+00 +/- 9.1e-02  |  4.15e+00 +/- 3.1e-01  |
 test_objective_compile_atf              |     -2.68 +/- 4.41     | -2.17e-01 +/- 3.57e-01 |  7.87e+00 +/- 2.8e-01  |  8.09e+00 +/- 2.2e-01  |
 test_objective_compute_atf              |     -2.50 +/- 5.74     | -2.69e-04 +/- 6.19e-04 |  1.05e-02 +/- 2.1e-04  |  1.08e-02 +/- 5.8e-04  |
 test_objective_jac_atf                  |     -3.14 +/- 2.59     | -6.36e-02 +/- 5.26e-02 |  1.97e+00 +/- 4.5e-02  |  2.03e+00 +/- 2.7e-02  |
 test_perturb_1                          |     -6.41 +/- 3.38     | -8.53e-01 +/- 4.49e-01 |  1.24e+01 +/- 1.7e-01  |  1.33e+01 +/- 4.2e-01  |
 test_proximal_jac_atf                   |     -0.92 +/- 0.91     | -7.52e-02 +/- 7.51e-02 |  8.13e+00 +/- 5.4e-02  |  8.20e+00 +/- 5.2e-02  |
 test_proximal_freeb_compute             |     -0.25 +/- 1.82     | -4.74e-04 +/- 3.43e-03 |  1.88e-01 +/- 3.0e-03  |  1.88e-01 +/- 1.7e-03  |

[unalmis]

Any place you use

compute_theta_coords(eq, nodes)  # by default maxiter for root finding is 20

will raise a deprecation warning. You can avoid this by replacing with

map_coordinates(eq, nodes, inbasis=("rho", "theta_PEST", "zeta")  # by default maxiter for root finding is 30. can pass in maxiter=20

[rahulgaur104]

I agree it become bloated as the scope of the test was increased in other PRs. That's why I cleaned it up in https://github.com/PlasmaControl/DESC/pull/1166. It was also necessary to modify the test so that the test doesn't fail on quantities that do field line coordinate mappings, like those in my work and yours. If I hadn't done that, any PR that uses the API we approved in https://github.com/PlasmaControl/DESC/pull/1024 for new quantities would run into issues with this test.

My point is more about the symmetry of the rpz and xyz calculations, not really related to an issue. For example this_branch_data_rpz[p].keys() and this_branch_data_xyz are both sets. Either this_branch_data_rpz should be a set or this_branch_data_xyz should be a dict. Just have the exact same structure for this_branch_data_xyz.

The point of the source grid requirement is to ensure the grid used by the compute function is the correct grid, such as having nodes that lie on field lines. If you add that requirement to your compute function as follows source_grid_requirement={"coordinates": "raz"}, but don't follow the API demanded by it, you will see errors. From your description, I assume your error message said that it requires that the grid has a source grid in ρ , α , ζ coordinates. The test passes for me now so I can't see the message you got.

It passes for ripple or for ballooning? For ballooning, I am not using source_grid_requirement. Feel free to add it to desc/compute/_stability "ideal ball gamma1" and rerun this test to see the errors. I can do this and share the error here as well.

I don't like the modifications made to the test for it to pass in this PR. The general fix made in 1. has been removed. If we merge the modifications made in this PR, when we add a new quantity that requires data along a field line, such as Gamma_c or effective ripple, this test will error.

None of my code would need to touched for your proxies because you already have a source grid thing (which I'll add back again) that will filter out the proxies from test_compute_everything. They are independent.

[unalmis]

It passes for ripple or for ballooning? For ballooning, I am not using source_grid_requirement. Feel free to add it to desc/compute/_stability "ideal ball gamma1" and rerun this test to see the errors. I can do this and share the error here as well.

Both.

None of my code would need to touched for your proxies because you already have a source grid thing (which I'll add back again) that will filter out the proxies from test_compute_everything. They are independent.

I am referring to the removal of the following code in tests/test_compute_everything:

names = {
                name
                for name in data_index[p]
                # Skip these quantities as they should be covered in other tests.
                if not data_index[p][name]["source_grid_requirement"]
            }
            this_branch_data_rpz[p] = things[p].compute(
                list(names), **grid.get(p, {}), basis="rpz"
            )

Please add back

[rahulgaur104]

The compute qty is still called ideal ball gamma, should it be ideal ballooning lambda? or ideal ballooning gamma^2

Good point. Let me do ideal ball lambda.

Also just to check, it is actually dimensionless now? Or does it still have units?

Yes, now it's actually dimensionless because all the terms in the equation are individually dimensionless.

[rahulgaur104]

I still need to review the tutorial notebook and I haven't checked the math, but I trust it's correct if the tests pass.

No no, don't trust me! Do check the math. Here's the paper with the equations (https://royalsocietypublishing.org/doi/epdf/10.1098/rspa.1979.0001)

[rahulgaur104]

https://github.com/PlasmaControl/DESC/pull/1170#discussion_r1749431211

It is important to have this function in the equilibrium class so that we can create a field-aligned rtz grid inside objective function that require field-aligned grids. The same idea has been used for Gamma_c and ripple as well. Look at those PRs. If you have a better solution, let me know.

[review-notebook-app[bot]]

View / edit / reply to this conversation on ReviewNB

unalmis commented on 2024-09-23T02:56:00Z ----------------------------------------------------------------

Line #28.            [rho, alpha, zeta], coordinates="raz", period=(np.inf, 2 * np.pi, np.inf)

Has no effect for this quantity right now, but could be safer in future to do np.inf in all three

[rahulgaur104]

Merging this branch as I have two approvals.