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.90 +/- 4.23     | +5.36e-03 +/- 2.53e-02 |  6.03e-01 +/- 1.6e-02  |  5.97e-01 +/- 2.0e-02  |
 test_build_transform_fft_highres        |     +1.43 +/- 2.59     | +1.36e-02 +/- 2.46e-02 |  9.66e-01 +/- 2.1e-02  |  9.53e-01 +/- 1.3e-02  |
 test_equilibrium_init_lowres            |     +1.78 +/- 4.09     | +6.78e-02 +/- 1.56e-01 |  3.88e+00 +/- 1.3e-01  |  3.81e+00 +/- 9.2e-02  |
 test_objective_compile_atf              |     +1.97 +/- 4.37     | +1.52e-01 +/- 3.38e-01 |  7.89e+00 +/- 2.1e-01  |  7.74e+00 +/- 2.7e-01  |
 test_objective_compute_atf              |     +0.40 +/- 5.02     | +4.23e-05 +/- 5.26e-04 |  1.05e-02 +/- 2.8e-04  |  1.05e-02 +/- 4.5e-04  |
 test_objective_jac_atf                  |     -4.09 +/- 2.51     | -7.85e-02 +/- 4.82e-02 |  1.84e+00 +/- 2.3e-02  |  1.92e+00 +/- 4.2e-02  |
 test_perturb_1                          |     -0.49 +/- 4.28     | -7.03e-02 +/- 6.15e-01 |  1.43e+01 +/- 2.7e-01  |  1.44e+01 +/- 5.5e-01  |
 test_proximal_jac_atf                   |     +0.46 +/- 0.88     | +3.72e-02 +/- 7.11e-02 |  8.15e+00 +/- 4.5e-02  |  8.11e+00 +/- 5.5e-02  |
 test_proximal_freeb_compute             |     +0.64 +/- 0.82     | +1.24e-03 +/- 1.60e-03 |  1.96e-01 +/- 9.6e-04  |  1.95e-01 +/- 1.3e-03  |
 test_build_transform_fft_lowres         |     +0.30 +/- 8.98     | +1.60e-03 +/- 4.70e-02 |  5.25e-01 +/- 3.5e-02  |  5.23e-01 +/- 3.1e-02  |
 test_equilibrium_init_medres            |     -1.00 +/- 2.52     | -4.12e-02 +/- 1.04e-01 |  4.08e+00 +/- 6.4e-02  |  4.12e+00 +/- 8.2e-02  |
 test_equilibrium_init_highres           |     -4.97 +/- 2.49     | -2.80e-01 +/- 1.41e-01 |  5.36e+00 +/- 3.4e-02  |  5.64e+00 +/- 1.4e-01  |
 test_objective_compile_dshape_current   |     +0.05 +/- 1.23     | +2.12e-03 +/- 4.74e-02 |  3.86e+00 +/- 3.2e-02  |  3.86e+00 +/- 3.5e-02  |
 test_objective_compute_dshape_current   |     -0.13 +/- 1.10     | -4.85e-06 +/- 3.97e-05 |  3.59e-03 +/- 2.4e-05  |  3.60e-03 +/- 3.1e-05  |
 test_objective_jac_dshape_current       |     +0.20 +/- 6.68     | +8.10e-05 +/- 2.68e-03 |  4.02e-02 +/- 1.8e-03  |  4.01e-02 +/- 2.0e-03  |
 test_perturb_2                          |     +0.26 +/- 1.61     | +4.90e-02 +/- 3.02e-01 |  1.88e+01 +/- 2.1e-01  |  1.88e+01 +/- 2.2e-01  |
 test_proximal_freeb_jac                 |     -0.46 +/- 1.73     | -3.45e-02 +/- 1.29e-01 |  7.46e+00 +/- 5.6e-02  |  7.50e+00 +/- 1.2e-01  |
 test_solve_fixed_iter                   |     +0.19 +/- 58.18    | +9.56e-03 +/- 2.92e+00 |  5.04e+00 +/- 2.0e+00  |  5.03e+00 +/- 2.1e+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)...

[unalmis]

Understand why effective radius returned by optimizer is half of the vlaue returned by the compute function

Probably need quadrature grid. See #1424