[WIP] Add TGLFInputs and TGLFNN

[theo-brown]

Add TGLFInputs class and TGLFNN model.
Note: completion depends on other work in progress in #476.

Remaining tasks:

• Inherit from QuasilinearTransportModelInputs
• Inherit from QuasilinearTransportModel
• Move calculation of Coulomb logarithm, nu, etc to physics.py

TB:

• Redo TGLFInputs to reflect changes in main (normalisation etc)
• Test TGLFInputs
• Test TGLFBasedTransportModel

LZ:

• Test tglf_wrapper

Then check in about what integration tests are required

[theo-brown]

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

[theo-brown]

Is this definition (chiGB) correct / consistent with TGLF?

[jcitrin]

Use quasilinear_transport_model.calculate_chiGB , where you can set a b_unit and reference length as input.

For TGLF probably b_unit=Bunit which you need to define (it's not the same as geo.B0), and aminor for the reference length. Should double check

As always, be careful with sqrt(2) in rho_s_unit and compare to TORAX calculate_chiGB to make sure we didn't miss another input argument (to include sqrt(2) or not)

[theo-brown]

Cool, I think this might've been a function added in that I hadn't noticed. I'll recheck the tglf docs for the reference field etc.

[theo-brown]

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

[jcitrin]

Note that you can use the new quasilinear_transport_model.NormalizedLogarithmicGradients.from_profiles for these. See

torax/torax/transport_model/qualikiz_based_transport_model.py

Line 138 in 4812e48

normalized_logarithmic_gradients = quasilinear_transport_model.NormalizedLogarithmicGradients.from_profiles(

1. I don't recommend on using the names Ate, Ati, etc., which are QuaLiKiz jargon. In any case Normalized logarithmic gradients contains general names
2. Note that Ati, Ate, (in QuaLiKiz jargon) should already include the normalization length, which is aminor for TGLF. NormalizedLogarithmicGradients takes the reference length as input

[jcitrin]

probably need alphaMHD as well? See quasilinear_transport_model.calculate_alpha

[lorenzozanisi]

I have now included this as P_PRIME_LOC here. This is not treated as an independent dimension in the NN, as it can be computed from gradients, safety factor and BETAE and the NN should be able to figure this out. It does need to be used for running TGLF of course, which is where the code I linked plays a role.

[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.