[WIP] Add TGLF transport models and surrogates

[theo-brown]

Add support for TGLF-based transport models, including coupling to full TGLF and TGLF-NN surrogates.

Remaining tasks:
TB:

• Unit test TGLFInputs
• Unit test TGLFBasedTransportModel

LZ:

• Unit test tglf_wrapper

Future:

• Integration tests

[theo-brown]

@lorenzozanisi could you lend a hand with checking the normalisations?

[theo-brown]

@lorenzozanisi ready for your bits to be integrated as and when!

[fusionby2030]

Nice WIP. It is a pain to get these normalizations & miller geometry. I wonder if using the pyrokinetics would be of interest?

Is there a plan to push this through in any of the next few weeks? I've got a TGLF surrogate as well, was thinking to test it in TORAX alongside ASTRA (current approach) to possibly cover model validation/comparison.

I think there are also a few tglf inputs missing, namely q_prime: q^2 a^2 / r^2 * shat, drmindx= r * dr/dx, where r = Rmin_local = (x_surface.max() - x_surface.min()) / 2.0, and x = sin^-1 (triangularity).

I would be happy to contribute these.

[theo-brown]

@fusionby2030 welcome to this PR! Glad that it might be useful to you!

I wonder if using the pyrokinetics would be of interest?

Good shout - @jcitrin are you aware of pyrokinetics? I've basically been using its documentation (+Bhavin Patel, one of the developers) as my reference for trying to hook these TGLF bits in to TORAX. I feel like if there is interest in TORAX become a one-stop-shop for gyrokinetics surrogates taking advantage of the work done in pyrokinetics would be a bit of a no-brainer.

I've got a TGLF surrogate as well, was thinking to test it in TORAX alongside ASTRA (current approach) to possibly cover model validation/comparison.

Super interested in this, keep me posted!

I think there are also a few tglf inputs missing

Yep, our surrogate doesn't use the full set of TGLF inputs (for reasons that @lorenzozanisi would be able to explain - I wasn't involved in its development!). Makes sense to extend it to the full load and make some optional or something. Please go ahead with adding whatever you need! Admin-wise, can you add to this PR given it's from a fork? Or would it need to be a branch?

Is there a plan to push this through in any of the next few weeks?

@lorenzozanisi and I were tag-teaming a bit on this PR - I think Lorenzo wanted to connect full TGLF to TORAX, and has made big strides in the week before Christmas on it. Not sure where those commits are - I keep poking him to push them! Along the way, we also found a few bits that were needed to be changed more widely, such as #628 that Jonathan handled. Definitely on the agenda on UKAEA's side for January!

[fusionby2030]

To add to the pyrokinetics discussion, there are ways (albeit hacky) of passing core_profiles from TORAX (or arbitrary kinetic sources) to construct Pyro objects. The transformations between different code input parameters is then trivial. Tested with GENE and TGLF given kinetic profiles from example TORAX sim class and some experimental data. Generally follows most of the pyrokinetics.kinetics file readers, just without the file. Can set via

from pyrokinetics import Pyro 
from pyrokinetics.kinetics import Kinetics
from pyrokinetics.units import UnitSpline
from pyrokinetics.species import Species
from pyrokinetics.constants import electron_mass

pyro_obj = Pyro(eq_file=eqfilepath, eq_type="GEQDSK", gk_code="TGLF")

# self constructed from your kinetic profiles, see IMAS file reader for example in pyrokinetics.kinetics.imas.py -> read_from_file()
kinetic_species_data = {'electron': Species(
        species_type="electron",
        charge=electron_charge_func, # UnitSpline
        mass=electron_mass,      
        dens=electron_dens_func,       # UnitSpline
        temp=electron_temp_func,     # Unit Spline
        omega0=omega_func,            # Unit Spline
        rho=rho_func,                          # Unit Spline
    ), 
   # same for ions...
}
kinetic_struct = Kinetics(kinetics_type="IMAS", # dummy since there are supported type checks...
                                       **kinetic_species_data)

pyro_obj.kinetics: Kinetics = kinetic_struct 

# One could then loop through PSI_N or rho_tor values

local_psi_n = 0.5 
pyro_c.load_local_geometry(local_psi_n)
pyro_c.load_local_species(local_psi_n)
pyro_c.write_gk_file("input.tglf", "TGLF")

I have no idea how @lorenzozanisi is planning on running TGLF from TORAX without lots of IO. Would be nice if there was a python interface to TGLF...

Might be best to wait until January then, am trying to coordinate a visit to Culham in February.

[lorenzozanisi]

I have pushed my WIP commits, sorry for not doing this earlier @theo-brown. These are not tested yet so don't attempt to use them, they are here so you can see how I am approaching the issue.

@fusionby2030 Note that these commits include only the inputs specific to the NNs that I have made - so this is by no means general (and it's possibly a bit hacky as I didn't have the time to implement the extensive and very general qualikiz tools that are used in the qualikiz_wrapper to handle the IO). As for stuff such as Q_PRIME_LOC, TORAX works in s_hat space and the NNs I made too. I made a utility that converts s_hat to Q_PRIME_LOC in the TGLF wrapper.

@fusionby2030 Note that even for QuaLiKiz not all inputs are used currently. I am not sure whether the plan is to eventually make TORAX capable of handling any number of inputs. @jcitrin how is this being handled to generalise from the 10D to the 11D case? I like the idea of including Pyro, but it sounds like TORAX is not fully IMAS compatible.

I have no idea how @lorenzozanisi is planning on running TGLF from TORAX without lots of IO. Would be nice if there was a python interface to TGLF...

I am simply writing to the input files and running TGLF, similarly as it is done for QuaLiKiz.

I shall crack on with testing my implementation now.

[theo-brown]

but it sounds like TORAX is not fully IMAS compatible.

With the open-sourcing of IMAS I would hazard a guess that IMAS compatability is moving up the priority list...

[jcitrin]

Hi all. Thanks for this great discussion.

I am aware of pyrokinetics. Garud Snoep also has gyrokit which is similar but I think he didn't open source it yet.

I am not sure whether the plan is to eventually make TORAX capable of handling any number of inputs. @jcitrin how is this being handled to generalise from the 10D to the 11D case?

The QuaLiKizInputs dataclass has all that's needed for the 11D case, of which the 10D is a subset. They both have access to the same dataclass and take what they need. The plan for now is to extend these dataclasses with new inputs when necessary. e.g., we don't have a rotation profile in TORAX (easy to add when needed, although momentum transport is a task) , and when we have it we can extend to the rotation related inputs.

I like the idea of including Pyro, but it sounds like TORAX is not fully IMAS compatible.
With the open-sourcing of IMAS I would hazard a guess that IMAS compatability is moving up the priority list...
IMAS compatibility is definitely moving up the priority list, and is starting now with the help of ITER and friends, beginning with equilibrium. However, I don't think we will have IMAS IDSs being passed around in TORAX internals, more on the pre and post processing stages (I/O), since we need to optimize for performance and also have the internal data objects be JAX compatible.

It may be interesting to use pyrokinetics in the future. Best would be to extend pyrokinetics itself to support "TORAX" sources of equilibrium and kinetic data. I suggest this is done after an upcoming API and variable renaming change that we are considering for the next version. Would need to consider performance and JAX compatibility though, especially with TGLF-NN. For TGLF itself this matters less since can run without JAX compilation, but there should be the same interface for TGLF and TGLF-NNs.

[theo-brown]

I've updated the TGLFInputs to use the normalized gradients and chiGB functions.

For reference, the normalisations in TGLF are as follows.

TGLF normalisation

In TGLF [1],

$$ \overline{p_{fe}} = p_{fe} / \Gamma_{GB} $$

$$ \overline{q_{e}} = q_{e} / Q_{GB} $$

where [2],

$$ Q_{GB} = \frac{n_e c_s T_e \rho_{unit}^2}{a^2}, $$

$$ \Gamma_{GB} = \frac{n_e c_s \rho_{unit}^2}{a^2} = Q_{GB} / T_e, $$ with $$ \begin{align} c_s &= \frac{T_e^\frac{1}{2}}{m_D^\frac{1}{2}}, \\ \rho_{{\rm unit}} &= \frac{c_s m_D c}{e B_{\rm unit}}. \end{align} $$ We can write $$ \begin{align} c_s \rho^2_{\rm unit} &= \frac{c_s^3 m_D^2 c^2}{e^2 B_{\rm unit}^2}, \\ &= \frac{T_e^\frac{3}{2} m_D^2 c^2}{m_D^\frac{3}{2} e^2 B_{\rm unit}^2}, \\ &= \frac{T_e^\frac{3}{2} m_D^\frac{1}{2} c^2}{e^2 B_{\rm unit}^2}, \end{align} $$ and so $$ \begin{align} Q_{GB} &= \frac{n_e T_e^\frac{5}{2} m_D^\frac{1}{2} c^2}{e^2 B_{\rm unit}^2 a^2}, \\ \Gamma_{GB} &= \frac{n_e T_e^\frac{3}{2} m_D^\frac{1}{2} c^2}{e^2 B_{\rm unit}^2 a^2}. \end{align} $$

TORAX normalisation

$$ p_{fe} = \frac{\overline{p_{fe}} n_e \chi_{GB}^{\rm TORAX}}{l_{ref}}, $$

where

$$ \chi_{GB}^{\rm TORAX} = \frac{(m_{ref} m_P)^\frac{1}{2} (T_e [J])^\frac{3}{2}}{B_{ref}^2 e^2 l_{ref}}. $$

Hence,

$$ \begin{align} \Gamma_{GB}^{\rm TORAX} &= \frac{n_e \chi_{GB}^\mathrm{TORAX}}{l_{ref}}, \\ &= \frac{n_e(m_{ref} m_P)^\frac{1}{2} (T_e [J])^\frac{3}{2}}{B_{ref}^2 e^2 l_{ref}^2}, \end{align} $$

and therefore

$$ \Gamma_{GB} = \Gamma_{GB}^\mathrm{TORAX} c^2, $$

with $l_{ref} = a$, $B_{ref} = B_{unit}$, $m_{ref} = 2(.014) = m_D \mathrm{[amu]}$.

I haven't yet checked that the conversion from $q_e$ to $\chi_e$ is correct for TGLF as implemented in QuasilinearTransportModel but I assume it will be. The only difference between the normalisation factors for particle and energy fluxes in TGLF is the T_e term.