Add Redl model of bootstrap current

[theo-brown]

Prompted by JINTRAC NCLASS comparison of Sauter and Redl models, this has been bumped up my priority list!

When I have time, I'll try and rearrange some of the definitions and files so there's less duplication.

Closes #1612.

[theo-brown]

STEP (dashed Redl, solid Sauter)

ITER (dashed Redl, solid Sauter)

[theo-brown]

@jcitrin as this is really an extension of the Sauter model, would you like it as a separate class, subclass, or flag within Sauter?

[jcitrin]

Excellent, thanks for this.

It would be good to have redl as a separate class, even if it's quite similar to sauter, and ensure that "sauter" and "redl" can be distinct model names. Possibly a lot from both can be separated out to use the common formulas module and duplication minimized, as you say.

It may be good to do this in more than 1 PR, separating out refactoring and new feature.

lmk when review is necessary and we can get back to this in the new year. Have a great holiday period!

[theo-brown]

I suggest the following design:

1. Move neoclassical.formulas to a subpackage, so that we can have neoclassical.formulas.shared, neoclassical.formulas.sauter and neoclassical.formulas.redl.

2. Extract the shared code from

   torax/torax/_src/neoclassical/bootstrap_current/sauter.py

   Lines 152 to 183 in 2aeaa48

   	# calculate bootstrap current
   	prefactor = -geo.F_face * bootstrap_multiplier * 2 * jnp.pi / geo.B_0
   	pe = p_e.face_value()
   	pi = p_i.face_value()
   	dpsi_drnorm = psi.face_grad()
   	dlnne_drnorm = n_e.face_grad() / n_e.face_value()
   	dlnni_drnorm = n_i.face_grad() / n_i.face_value()
   	dlnte_drnorm = T_e.face_grad() / T_e.face_value()
   	dlnti_drnorm = T_i.face_grad() / T_i.face_value()
   	global_coeff = prefactor[1:] / dpsi_drnorm[1:]
   	global_coeff = jnp.concatenate([jnp.zeros(1), global_coeff])
   	necoeff = L31 * pe
   	nicoeff = L31 * pi
   	tecoeff = (L31 + L32) * pe
   	ticoeff = (L31 + alpha * L34) * pi
   	j_parallel_bootstrap_face = global_coeff * (
   	necoeff * dlnne_drnorm
   	+ nicoeff * dlnni_drnorm
   	+ tecoeff * dlnte_drnorm
   	+ ticoeff * dlnti_drnorm
   	)
   	j_parallel_bootstrap = geometry_lib.face_to_cell(j_parallel_bootstrap_face)
   	return base.BootstrapCurrent(
   	j_parallel_bootstrap=j_parallel_bootstrap,
   	j_parallel_bootstrap_face=j_parallel_bootstrap_face,
   	)

   
   into a new function

neoclassical.formulas.shared._calculate_analytic_bootstrap_current(
    *,
    bootstrap_multiplier: float,
    n_e: cell_variable.CellVariable,
    n_i: cell_variable.CellVariable,
    T_e: cell_variable.CellVariable,
    T_i: cell_variable.CellVariable,
    p_e: cell_variable.CellVariable,
    p_i: cell_variable.CellVariable,
    geo: geometry_lib.Geometry,
    L31: array_typing.FloatVectorFace,
    L32: array_typing.FloatVectorFace,
    L34: array_typing.FloatVectorFace,
    alpha: array_typing.FloatVectorFace,
  ) -> base.BootstrapCurrent

3. Make a new function for each of Sauter and Redl that computes the terms that are different, e.g.

neoclassical.bootstrap.sauter.calculate_bootstrap_terms(...) -> ...:
  f_trap = neoclassical.formulas.sauter.calculate_ftrap(...)
  log_lambda_ei = collisions.calculate_log_lambda_ei(...)
  log_lambda_ii = collisions.calculate_log_lambda_ii(...)
  nu_e_star = collisions.calculate_nu_e_star(...)
  nu_e_star = collisions.calculate_nu_i_star(...)

  L31 = neoclassical.formulas.sauter.calculate_L31(...)
  L32 = neoclassical.formulas.sauter.calculate_L32(...)
  L34 = neoclassical.formulas.sauter.calculate_L34(...)
  alpha = neoclassical.formulas.sauter.calculate_alpha(...)
 
  return L31, L32, L34, alpha

I think this is the most simple option that minimises duplication. Also I think this structure sets up Andreas' intent to do the new conductivity model. @jcitrin thoughts?

[jcitrin]

@theo-brown this structure sounds good to me.

@tamaranorman , @Nush395 , any thoughts on the proposed structure? Context: Theo and Andreas Redl are implementing the Redl model for bootstrap current and conductivity, which is an upgrade on the sauter model. Structurally it's very similar to the sauter models, but have slightly different formulas and coefficients. Theo's proposal aims to minimize the duplication and still have both sauter and redl as model types.

[tamaranorman]

Looks good in its current form - one question is when we talk about rearranging what does that look like as this currently is just adding another implementation

[theo-brown]

Thanks for your review @tamaranorman! I was mainly @'ing you for your thoughts on the comment, sorry.

See my comment above for my proposed structure (link) - I didn't want to waste time implementing before a design had been chosen.

The "changes to structure" are just moving things around, rather than introducing any new classes etc. I think Jonathan was just to check whether there's an alternative (eg create a new superclass that both of these inherit from) or whether having this level of code duplication is ok/acceptable/avoidable.

[tamaranorman]

That proposed structure seems reasonable to me - maybe returning a dict that can then be unpacked using ** just as with lots of scalars its worth making sure the return values are clear

[theo-brown]

Thanks for the review @tamaranorman. Sorry that I still haven't managed to make my linter sync up with the Google one!

[theo-brown]

This has been merged, I just need to fix the copybara error.