Critical gradient

[kfrybes]

Branch for developing a critical gradient optimization metric to reduce turbulent transport. Theory for critical gradient is taken from the following paper

TODO list :

Effective radius

• Run some more optimizations to understand when optimizer goes crazy
• Run optimization with additional constraints (QS/aspect ratio/...)
• Make optimizations with geometry modes >1 work
• Run optimization with target values
• Understand why effective radius returned by optimizer is half of the vlaue returned by the compute function

Parallel connection length

• Test optimization on some basic cases
• Implement definition of parallel correlation length from the paper

Critical gradient

• Run optimization for both parallel connection length and effective radius at the same time with different weights
• Implement critical gradient definition from the paper

General/code

• Decide on how to choose number of poloidal/toroidal turns as well as number of wells
• Find ways to prevent optimizer from going into "bad regions" making lots of small peaks to reduce effective radius/parallel connection length.
• Implement tests
• Change the python for loops on the wells jax.lax loops
• Take care of outputs that contains nans, if the chosen number of wells is smaller than the number of wells there is, for the chosen npol/ntor
• Write good documentation
• Make scalar optimizations with mean and max work with scalar optimizer (tracer error)

Ideas for preventing bad optmizations :

1. Know the precise number of wells beforehand, this helps
2. fit a polynomial or fourier series through the wells to eliminate the small peaks
3. Fix number of wells per toroidal transit as a constraint for the optimizer

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[kfrybes]

Notebook for testing jax implementation of critical gradient

[kfrybes]

Notebook for testing critical gradient optimization

[kfrybes]

Notebook for fitting zonal flows, not up to date

[kfrybes]

Notebook for testing critical gradient fitting, not jit compatible not up to date

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     -0.34 +/- 3.47     | -2.09e-03 +/- 2.12e-02 |  6.09e-01 +/- 1.7e-02  |  6.11e-01 +/- 1.3e-02  |
 test_build_transform_fft_highres        |     -1.01 +/- 4.77     | -1.02e-02 +/- 4.81e-02 |  9.97e-01 +/- 1.4e-02  |  1.01e+00 +/- 4.6e-02  |
 test_equilibrium_init_lowres            |     +0.76 +/- 1.35     | +2.89e-02 +/- 5.13e-02 |  3.83e+00 +/- 4.6e-02  |  3.80e+00 +/- 2.2e-02  |
 test_objective_compile_atf              |     +1.34 +/- 3.79     | +1.03e-01 +/- 2.92e-01 |  7.81e+00 +/- 2.2e-01  |  7.71e+00 +/- 1.9e-01  |
 test_objective_compute_atf              |     +1.37 +/- 2.87     | +1.43e-04 +/- 3.01e-04 |  1.06e-02 +/- 1.9e-04  |  1.05e-02 +/- 2.3e-04  |
 test_objective_jac_atf                  |     +0.99 +/- 2.95     | +1.83e-02 +/- 5.45e-02 |  1.87e+00 +/- 4.1e-02  |  1.85e+00 +/- 3.6e-02  |
 test_perturb_1                          |     +0.49 +/- 1.15     | +6.12e-02 +/- 1.45e-01 |  1.26e+01 +/- 6.1e-02  |  1.26e+01 +/- 1.3e-01  |
 test_proximal_jac_atf                   |     -0.01 +/- 0.86     | -6.50e-04 +/- 6.94e-02 |  8.07e+00 +/- 5.2e-02  |  8.07e+00 +/- 4.6e-02  |
 test_proximal_freeb_compute             |     -0.14 +/- 1.00     | -2.47e-04 +/- 1.82e-03 |  1.82e-01 +/- 1.6e-03  |  1.83e-01 +/- 9.4e-04  |
 test_build_transform_fft_lowres         |     +2.79 +/- 3.45     | +1.51e-02 +/- 1.87e-02 |  5.56e-01 +/- 1.3e-02  |  5.41e-01 +/- 1.4e-02  |
 test_equilibrium_init_medres            |     +0.91 +/- 3.52     | +3.90e-02 +/- 1.50e-01 |  4.31e+00 +/- 1.4e-01  |  4.27e+00 +/- 4.1e-02  |
 test_equilibrium_init_highres           |     -0.85 +/- 1.92     | -4.82e-02 +/- 1.09e-01 |  5.62e+00 +/- 8.9e-02  |  5.67e+00 +/- 6.4e-02  |
 test_objective_compile_dshape_current   |     -1.42 +/- 1.20     | -5.59e-02 +/- 4.73e-02 |  3.88e+00 +/- 4.1e-02  |  3.93e+00 +/- 2.3e-02  |
 test_objective_compute_dshape_current   |     +0.22 +/- 2.48     | +7.82e-06 +/- 8.91e-05 |  3.60e-03 +/- 5.7e-05  |  3.59e-03 +/- 6.9e-05  |
 test_objective_jac_dshape_current       |     -3.47 +/- 10.44    | -1.41e-03 +/- 4.26e-03 |  3.94e-02 +/- 2.5e-03  |  4.08e-02 +/- 3.5e-03  |
 test_perturb_2                          |     -1.97 +/- 1.34     | -3.58e-01 +/- 2.44e-01 |  1.78e+01 +/- 2.3e-01  |  1.82e+01 +/- 8.3e-02  |
 test_proximal_freeb_jac                 |     -0.20 +/- 1.36     | -1.47e-02 +/- 1.02e-01 |  7.50e+00 +/- 3.8e-02  |  7.52e+00 +/- 9.5e-02  |
 test_solve_fixed_iter                   |     +0.47 +/- 58.65    | +2.35e-02 +/- 2.92e+00 |  5.00e+00 +/- 2.1e+00  |  4.97e+00 +/- 2.0e+00  |

[ddudt]

@kfrybes see our contributing guidelines for info on installing black and pre-commit.

[dpanici]

Just a note, please check the styleguide part of our contribution guidelines and install pre-commit so you can catch these formatting and other warnings! that way you can avoid all of these comments on the PR

[kfrybes]

yep, sorry, I made a bit of a mess

[dpanici]

yep, sorry, I made a bit of a mess

I've been there (#579 has over 500 comments)...