[docs] PRESCRIBED heating and current source profiles

[theo-brown]

While future work may focus on developing specific physically realistic profiles for current and heating actuators, in the meantime it would be useful to be able to represent arbitrary source profiles so that scenarios generated in other codes can be run in TORAX.

The most basic of these would be a piecewise linear profile.

[theo-brown]

I would be happy to extend sources/formulas.py and sources/formula_config.py to support this, if a core contributor would give me the go-ahead! Not sure if there is a better option that I'm missing.

[jcitrin]

Fully agree that supporting arbitrary source profiles should be done in the near-term.

At the moment we are working on supporting arbitrary initial and prescribed core profiles, and the infrastructure is already there to a large extent. Te, Ti, ne are now of type TimeInterpolatedArray (see interpolated_param.py) and can be defined in config as nested dicts in a time series {time: value} like {0: {0.0: 20, 0.95: 5, 1.0: 0.2}, 1: {0.0: 20, 0.95: 5, 1.0: 0.2}}, where value is itself a dict representing {rhonorm: value}. Down the road, we plan to generalize TimeInterpolatedArray to accept more general datastructures such as xarray datasets or numpy arrays, as long as they satisfy specific formats. Work has started in this direction, see #197 . This will enable loading these values from output files of previous runs, from IMAS (there are IMAS IDS to xarray converters out there), or more generally from any userdata as long as it's converted to the right format.

After this long preamble :) , I think we should reuse all this infrastructure to enable arbitrary source profiles as well. Can envisage adding a new supported mode (see Mode class in sources/runtime_params.py) like PRESCRIBED, and the source then gets as input the nested dict, or numpy arrays, or an xarray dataset, or a filepath to an .nc file or IMAS h5 file, etc (note that currently only nested dict input is supported). The PRESCRIBED mode could include options like a multiplication factor to the input, or normalizing the input to a desired total integrated source amplitude.

You could indeed also extend sources/formulas.py and sources/formula_config.py to do the same. This is easier since there's less new development to do, but likely down the road we'll move this into a new PRESCRIBED mode. I'd recommend a new "formula" class something like (just a suggestion)

@dataclasses.dataclass
class PiecewiseLinear:
  """Configures a piecewise linear profile."""

  use_total: bool = False
  total: TimeInterpolatedScalar = 1.0
  multiplication_factor: TimeInterpolatedScalar = 1.0
  profile: TimeInterpolatedArray
  
  ...

Where the user has a choice whether to use the profile as is, or to renormalize to total if use_total=True, and/or to multiply by multiplication_factor for sensitivity tests. The onus is still on the user in a pre-processing step (can be done in the config file itself, above the CONFIG definition) to convert your source data into the nested dict format for input into profile, due to the limited support for other data structures at the moment.

[jcitrin]

and yes, feel free to go ahead and give it a go if you're up for it :)

[jcitrin]

@araju for any further insight here

[jcitrin]

@theo-brown , note #215 in case you get confusing results when playing around with the source config.

[theo-brown]

Adding a PRESCRIBED mode seems much smarter. I'll make a start there, tying in with the TimeInterpolatedArray stuff. Thanks for the tip regarding #215 !

[theo-brown]

Encountering some difficulty with the difference between grids; current density seems to be computed on the hi-resolution grid, but the DynamicRuntimeParamsSlice discretises stuff on the standard grid. Not sure what the best way is for tackling this.

• TimeInterpolatedArrays get interpolated when the DynamicRuntimeSlice is constructed (interpolation performed by get_init_kwargs, e.g. when build_dynamic_runtime_params_slice is called).
• Most source profiles are computed using get_source_profiles, where, for example, the source's formula is called with the DynamicRuntimeSlice as an argument. For most sources, this means that switching to TimeInterpolatedArrays should work without a hitch.
• However, some current sources have additional methods that allow them to be recomputed on the high resolution grid. For example, external_current_source.py has _calculate_jext_hires, which is called instead of get_source_profiles. This is a hardcoded function that bypasses some of the internals of the DynamicRuntimeSlice, in order to explicitly compute the source profile on the high-resolution grid. In fact, currently it is a hardcoded Gaussian, overriding any other formula that might have been set.

Skipping over the hardcoded Gaussian for the moment, this results in the following obstacle for adding support for PRESCRIBED sources:

1. build_dynamic_runtime_params_slice is called by the main routine, and all TimeInterpolatedArrays are interpolated onto the standard grid at this timestep.
2. When it comes to computing the high-resolution currents, we no longer have access to the un-interpolated array; the DynamicRuntimeSlice only contains the interpolated one.

An easy fix could be to re-interpolate the low-resolution interpolation onto the high-resolution grid. My concern is that this is a bit of a sticking-plaster solution, and is avoiding tackling the underlying pitfall.

[jcitrin]

Some deficient/not-yet-implemented elements have been exposed here :) , and good to raise them for prioritization.

First some explanation on the hires grid and what it's used for. It's only used for initializing the psi profile for when the current profile is set to determine the initial psi profile (when CONFIG['runtime_params']['initial_psi_from_j']=True) (see here and what calls it)

Now, in that function, there is a deficiency that the only external current that can be included now is jext. And indeed jext is hard coded to just be a Gaussian formula with convenient physically intuitive names for width, location, and current fraction (which sets its amplitude). This is a kind of proxy current source for our initial testing purposes. The actual physical current sources have placeholders in sources/current_density_sources.py and aren't yet implemented, i.e. not yet included in sources/default_sources.py or have specific implementations of them.

You can see this commit to see what adding a new source entails. There is work coming soon to simplify the process of registering a new source.

[jcitrin]

Which sources are you interested in implementing as PRESCRIBED? From the questions it feels like ECRHCurrentSource?

That is a bit more involved due indeed to the need to extend _get_jtot_hires to include all available current sources. But also not a blocker to get started with a PRESCRIBED ECRHCurrentSource if starting with a prescribed psi from a geometry.

[theo-brown]

Thanks for the explanation about the hires grid! That makes a lot of sense, and I hadn't spotted that it would only be used in the initialisation.

Yes, my main interest is ECRHCurrentSource, but I was basically using it as a way in to learning about the sources generally, as well as the internal workings of the code. I thought that perhaps a PRESCRIBED mode would be a good way of contributing something that would be useful more generally :) I started off just in the base source class, and then worked step-by-step from there to identify and fix problems, with the aim of allowing any source to be prescribed (which I think is a useful goal).

However, as you've highlighted, it's come across some of the surrounding bits that aren't quite there yet and implementing a new mode is wider in scope than I first realised! There's definitely a tradeoff between making a generally-applicable contribution and biting off a manageably-sized problem...

Basically, I'm interested in this part of the code and would be keen to contribute. Let me know if there are any ways I can better dovetail with existing or planned work.

[jcitrin]

I think that you can still consider setting up a PRESCRIBED mode and it will work for jext as well, minimizing the amount of new code needed. We'll probably get to it later this month in any case, if you'd rather wait.

[theo-brown]

PRESCRIBED mode with tests finalised in #221. Note that due to the aforementioned quirks of jext the tests are only for prescribed particle sources at the moment, but should work in future for all sources. This has now been fixed!

[theo-brown]

Has this now been completed with 6aeb322? If so, then I'm happy to add to the docs.

[jcitrin]

Yes, this is now merged. Thanks for suggesting to add to the docs!

[theo-brown]

I wasn't quite aware how much the implementation of prescribed profiles has changed. As it stands, I'm no longer on top of the complexity of it, so I'm not suitable for writing the docs!

[jcitrin]

No problem. We can do it. FYI: the implementation of the underlying interpolated_param objects has changed in the following ways, both for prescribed quantities in profile_conditions (e.g. Ti, ne) as well as the PRESCRIBED mode for sources:

1. More generality in data structures that can be inputted. There is no longer a limitation to Mappings of the form {time0: {rho0: value00, rho1:rho00}, time1: {rho0: value01, rho1: value11}}, but can also ingest xarray datasets or a tuple of numpy arrays. For example, for your test_prescribed_jext.py, we currently have:

            'prescribed_values': {
                0: {
                    r_i.item(): jext_i.item()
                    for r_i, jext_i in zip(gauss_r, jext_profile_0)
                },
                2.5: {
                    r_i.item(): jext_i.item()
                    for r_i, jext_i in zip(gauss_r, jext_profile_1)
                },
            },

But the following variants are now also acceptable, and much more convenient:

'prescribed_values': {0: (gauss_r, jext_profile_0), 2.5: (gauss_r, jext_profile_1)},

or

'prescribed_values': (times, gauss_r, jext_profile),

where jext_profile is a 2d array containing all the time slices, i.e.

times = np.array([0, 2.5])
gauss_r = np.linspace(0, 1, 32)
jext_profiles = np.array([
    gaussian(gauss_r, center=0.35, width=0.05, amplitude=1e6),
    gaussian(gauss_r, center=0.15, width=0.1, amplitude=1e6),
])

I will update the test with this latter more concise input variation.

So we'll just add a note in the docs advertising the new PRESCRIBED mode, and point to the existing documentation on allowed data structures for time-varying arrays. I'll close the issue after that change is in very soon.

[jcitrin]

Closing with #400 (PR 400 🥳)