Add Vloop BC

[theo-brown]

Adding Vloop boundary condition on the psi equation, so that the current can be evolved self-consistently given a rate of change of edge flux.

Tasks:

• Add Vloop boundary condition
• Add Vloop_LCFS to outputs
• Support all methods of initialising psi
• Unit test
• Integration test

[theo-brown]

I have not yet written a test case for this, as I wanted to run the pattern past one of the team first. Let me know your initial thoughts and then I can continue!

[jcitrin]

Some general thoughts.

1. What's missing is a consideration of how to initialize psi to begin with, for all cases. If Ip is None, then how do we set the right boundary condition at the first timestep, for cases where we initialize based on plasma current? Recall that there are several options currently (which can be further extended later):

i) profile_conditions['initial_psi_from_j'] = False, i.e. Initialize psi based on the input geometry file. Two suboptions:
a) geometry['Ip_from_parameters'] = False.
Here's it's clear what to do, psi_right = psi_geometry[-1]
b) geometry['Ip_from_parameters'] = True. For this we need Ip to rescale psi, so if Ip=None we have a problem

ii) profile_conditions['initial_psi_from_j'] = True, i.e. Initialize psi based on the "nu" current profile model. This needs normalization to a total current, i.e. Ip. So if Ip=None we have a problem. Note that at the initial condition, the Vloop_bound_right needs to be such that Psi.face_grad()[-1] is consistent with Ip, according to the same formula as the standard boundary condition. See fvm.cell_variable.face_grad, that face_grad()[-1] is just a standard derivative between the dx/2 distance between the last cell point and the rightmost face.

To solve this, I recommend changing the name of Ip to Ip_input, and for cases where the Vloop boundary condition is used, then it's used (if necessary) to initialize psi. That means that Vloop_bound_right and Ip can coexist. The pattern should be a bit different, e.g. an extra boolean "use_Vloop_boundary_condition" or something like that. Note that Ip_input can still be a TimeInterpolatedInput, but for the Vloop case only the initial condition at t=t_initial will be used.

2. We need to adjust to the Vloop boundary case where profile_conditions['Ip_input'] is not actually total current, but rather the initial condition. For example, the output Ip is taken from there, and also in plotruns_lib.PlotData.i_total. For the code to work with both boundary condition cases, we should take Ip (total current) from elsewhere. A candidate could be to extend physics.calc_jtot_from_psi to output I_tot as well (can change the name to Ip_profile), and then save this quantity in CoreProfiles.Currents . Total current is then Ip_profile[-1]. This can be a separate PR. This way, we also get information on the total current evolution in the Vloop boundary condition case.

[theo-brown]

Great points, and this is exactly why I wanted to run it through you! I think because for the type of simulation I've learnt how to do we always prescribe psi from a geometry I hadn't really thought about the consequences for any situation other than 1ia).

Noting down for myself - I need to consider the cases:

	Psi from...	Ip
1ai	Geo	Self-consistent
1aii	Geo	Prescribed
2	Initial j profile	Prescribed

and their interactions with the Vloop BCs

[jcitrin]

#492 was done with this PR in mind

[theo-brown]

@jcitrin did #492 tackle all of the cases where Ip_tot was used as the plasma current internally, or was it just the writing to output?

[jcitrin]

@jcitrin did #492 tackle all of the cases where Ip_tot was used as the plasma current internally, or was it just the writing to output?

Only the writing to output.

[theo-brown]

I've been working on incoporating the Vloop BC via compute_boundary_conditions.
The basic logic is (pseudocode):

if Vloop_bound_right is given:
   psi right BC = value constraint from Vloop
else:
   psi right BC = grad constraint from Ip_tot

In this case, I've attempted to sidestep the initialisation problem by having a spatial-derivative BC for the first timestep and a value BC for subsequent timesteps.

I've found that my current method results in an inconsistency with the setting of psi vs psidot.
The psi that is written to the output (core_profiles/psi) is consistent with Vloop_bound_right. In this case, Vloop_bound_right = 0, so the time derivative of psi at the LCFS is 0.
However, core_profiles/psidot is not consistent with this (see graph).

From the code, it appears that within sim.py, psidot is calculated using the function from ohmic_heat_source.py:

torax/torax/sim.py

Lines 1568 to 1583 in c860557

	def update_psidot(
	dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
	geo: geometry.Geometry,
	core_profiles: state.CoreProfiles,
	source_models: source_models_lib.SourceModels,
	) -> state.CoreProfiles:
	"""Update psidot based on new core_profiles."""
	psidot = dataclasses.replace(
	core_profiles.psidot,
	value=ohmic_heat_source.calc_psidot(
	dynamic_runtime_params_slice,
	geo,
	core_profiles,
	source_models,
	),

torax/torax/sources/ohmic_heat_source.py

Lines 43 to 138 in c860557

	def calc_psidot(
	dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
	geo: geometry.Geometry,
	core_profiles: state.CoreProfiles,
	source_models: source_models_lib.SourceModels,
	) -> jax.Array:
	r"""Calculates psidot (loop voltage). Used for the Ohmic electron heat source.
	psidot is an interesting TORAX output, and is thus also saved in
	core_profiles.
	psidot = \partial psi / \partial t, and is derived from the same components
	that form the psi block in the coupled PDE equations. Thus, a similar
	(but abridged) formulation as in sim.calc_coeffs and fvm._calc_c is used here
	Args:
	dynamic_runtime_params_slice: Simulation configuration at this timestep
	geo: Torus geometry
	core_profiles: Core plasma profiles.
	source_models: All TORAX source/sinks.
	Returns:
	psidot: on cell grid
	"""
	consts = constants.CONSTANTS
	psi_sources, sigma, sigma_face = source_models_lib.calc_and_sum_sources_psi(
	dynamic_runtime_params_slice,
	geo,
	core_profiles,
	source_models,
	)
	# Calculate transient term
	toc_psi = (
	1.0
	/ dynamic_runtime_params_slice.numerics.resistivity_mult
	* geo.rho_norm
	* sigma
	* consts.mu0
	* 16
	* jnp.pi**2
	* geo.Phib**2
	/ geo.F**2
	)
	# Calculate diffusion term coefficient
	d_face_psi = geo.g2g3_over_rhon_face
	# Add phibdot terms to poloidal flux convection
	v_face_psi = (
	-8.0
	* jnp.pi**2
	* consts.mu0
	* geo.Phibdot
	* geo.Phib
	* sigma_face
	* geo.rho_face_norm**2
	/ geo.F_face**2
	)
	# Add effective phibdot poloidal flux source term
	ddrnorm_sigma_rnorm2_over_f2 = jnp.gradient(
	sigma * geo.rho_norm**2 / geo.F**2, geo.rho_norm
	)
	psi_sources += (
	-8.0
	* jnp.pi**2
	* consts.mu0
	* geo.Phibdot
	* geo.Phib
	* ddrnorm_sigma_rnorm2_over_f2
	)
	diffusion_mat, diffusion_vec = diffusion_terms.make_diffusion_terms(
	d_face_psi, core_profiles.psi
	)
	# Set the psi convection term for psidot used in ohmic power, always with
	# the default 'ghost' mode. Impact of different modes would mildly impact
	# Ohmic power at the LCFS which has negligible impact on simulations.
	# Allowing it to be configurable introduces more complexity in the code by
	# needing to pass in the mode from the static_runtime_params across multiple
	# functions.
	conv_mat, conv_vec = convection_terms.make_convection_terms(
	v_face_psi,
	d_face_psi,
	core_profiles.psi,
	)
	c_mat = diffusion_mat + conv_mat
	c = diffusion_vec + conv_vec
	c += psi_sources
	psidot = (jnp.dot(c_mat, core_profiles.psi.value) + c) / toc_psi
	return psidot

I don't understand yet what's going on in calc_psidot, or how it interacts with the compute_boundary_conditions.
Looking at calc_coeffs it looks like the same logic is indeed happening there, which is probably the most significant location out of all of them. This gives me the impression that modifying compute_boundary_conditions is insufficient to implement this functionality. @jcitrin do you have any insights on whether that's the case, or whether I've just missed something? Thanks in advance for your help!

[jcitrin]

calc_psidot literally calculates the entire RHS of the governing psi equation.

So \frac{\partial \psi}{\partial t} = psidot

First thing, is that psidot is on the cell grid. Therefore psidot[:, -1] does not correspond to the LCFS, so you shouldn't expect equivalence.

For a basic sanity check of psidot, you can try setting 100% non-inductive current in one of the configs and see psidot become both flat (stationary-state current profile) and zero (no inductive current that shifts psi in time). See below for iterhybrid_predictor_corrector when turning off bootstrap current and doing the following for generic_current_source

    'generic_current_source': {
        # total "external" current fraction
        'fext': 1.0,
        # width of "external" Gaussian current profile (normalized radial
        # coordinate)
        'wext': 0.4,
        # radius of "external" Gaussian current profile (normalized radial
        # coordinate)
        'rext': 0.0,
    },

plot below: psidot[-1] is y-axis, time is x-axis

Therefore what's happening to psi and psidot at the LCFS is currently not in the output, and wouldn't even be accessible from the calc_psidot function, because that doesn't have any info on the boundary conditions.

It would indeed be useful to include Vloop_LCFS as a core_profile output. It can be calculated using CellVariable methods based on psi from core_profiles_t and core_profiles_t_plus_dt. If not setting Vloop_LCFS as a boundary condition and just calculating it based on the profiles with fixed Ip, then the initial time Vloop will be undefined.

Regarding boundary conditions, I recommend for consistency, setting the Dirichlet boundary conditions throughout, when the Vloop option is used. On initialization, you must make it consistent with the initial total Ip, by calculating the necessary gradient from the last cell gridpoint to the last face gridpoint, and doing linear extrapolation.

Happy to discuss over a call on Mon or Tue if needed.

[theo-brown]

It would indeed be useful to include Vloop_LCFS as a core_profile output.

Just to check on this, do you mean a property of state.PostProcessedOutputs? Or in state.CoreProfiles?
In OMAS, it's in core_profiles.global_quantities.v_loop.

[jcitrin]

Just to check on this, do you mean a property of state.PostProcessedOutputs? Or in state.CoreProfiles?
In OMAS, it's in core_profiles.global_quantities.v_loop.

Makes most sense for it to be in CoreProfiles

[theo-brown]

This PR is now ready for review :)

Summary of changes

• Added new BC for the psi equation, ProfileConditions.vloop_lcfs.
  • Can be combined with any of the psi initialisation methods:
    • Prescribed psi + Ip_tot rescaling
    • Psi from geometry file
    • Psi from geometry file + Ip_tot rescaling
    • Psi from nu formula + Ip_tot rescaling
• Added vloop_lcfs to state.CoreProfiles, which is a direct copy from psidot.face_value()[-1].
• Added psi_right_bc and vloop_lcfs to output.

Outstanding questions/bugs

• The options involving psi from geometry file are not unit tested, as the other unit tests of this type avoided reading in a geometry file. They are, however, tested by integration tests.
• The prescribed psi case is not integration tested, as the other tests for the psi equation seemed to all use the nu formula.
• As psi_right_bc can be None, this required changing one of the unit tests so that assertequal allowed None == None. I'm concerned that this is slightly bad practice.
• It also results in a warning when constructing the output, as the output involves converting a jnp.array([None, ...]) -> jnp.array([jnp.nan, ....]), which will be deprecated in future. For ease of comparing .nc files and Xarray Datasets, I thought it was better to have the key psi_right_bc always present but sometimes nan, rather than having only sometimes present. I'm open to other opinions on the matter!
• I think that these changes mean that at least one profile for the first timestep has a nan in (q at edge = inf). This breaks the plotting script. If I have time this week I'll track through and find why this happens.

[theo-brown]

Not sure why this failed copybara - is it because it's behind main so there are conflicts?

[jcitrin]

Indeed, it's due to merge conflicts. Strongly recommend merging main and syncing to HEAD. It may get a bit hairy, sorry about that. 😬

And apologies again for the delays in review. Will get on it tonight

[theo-brown]

OK! I'll give it a go this morning.
Don't worry about the delay - the new features are super exciting, I was more than happy to wait!

[theo-brown]

I've rebased onto main, not sure why it's still failing.

[jcitrin]

I've rebased onto main, not sure why it's still failing.

Yeah, weird. Maybe it's due to merge conflicts in the .nc files which it doesn't know how to resolve. We can resolve them manually when we do the manual import and internal review, so I wouldn't worry about it. As long as it's rebased to main and the actual tests are passing, should be fine.

[theo-brown]

Is there another way, like by extending the input arguments with information on which branch to take, or having a different function for the different branch?

Done in 4d55436

[theo-brown]

Behaviour of psidot, psi edge, and Ip in updated sims (986f3ec). The other profiles look as expected - can upload a screenshot of those too if that would be nice @jcitrin.

I'm aware that I'm not comparing apples to apples when using [:, -1], as psi and psidot are cell variables so [:, -1] is not on the LCFS. However, I was using output.load_state_file, so I think this was all that I had available. Middle plot shows the difference is minimal.

[jcitrin]

Things look even more sane after introducing psi_from_Ip_face and making that the right BC for Case 2.

Here's the new test-case with Ip_from_parameters

And here with Ip from CHEASE

Main point to notice is the Ip time series. In each case it starts from where expected (15MA, ~11.5MA respectively), before Vloop begins to modify it. Previously I think we were effectively setting the initial current to zero by setting the BC as the last point from the cell grid, and hence there was a current perturbation in the first timesteps

[theo-brown]

Ah of course, that's miles better! Thanks for making the necessary changes.

[theo-brown]

Huge thanks to @jcitrin, @tamaranorman, and @Nush395 for your help in getting this in - especially for your suggestions and time yesterday, as well as extra work put in after it went to internal review 🚀

[jcitrin]

It's a great feature, thanks for your efforts in getting this in! Looking forward to seeing its use in STEP simulations