Switch psi sources to being <j.B>/B_0

Author: theo-brown

torax/_src/output_tools/output.py

@@ -17,14 +17,12 @@
 import dataclasses
 import inspect
 import itertools
-
 import os
 
 from absl import logging
 import chex
 import jax
 import numpy as np
-import os
 from torax._src import array_typing
 from torax._src import state
 from torax._src.geometry import geometry as geometry_lib
@@ -58,14 +56,18 @@
 Z_EFF = "Z_eff"
 SIGMA_PARALLEL = "sigma_parallel"
 V_LOOP_LCFS = "v_loop_lcfs"
-J_TOTAL = "j_total"
 IP_PROFILE = "Ip_profile"
 IP = "Ip"
 
-# Calculated or derived currents.
-J_OHMIC = "j_ohmic"
-J_EXTERNAL = "j_external"
-J_BOOTSTRAP = "j_bootstrap"
+# Calculated or derived current densities
+J_PAR_TOTAL = "j_parallel_total"
+J_PAR_OHMIC = "j_parallel_ohmic"
+J_PAR_EXTERNAL = "j_parallel_external"
+J_PAR_BOOTSTRAP = "j_parallel_bootstrap"
+J_TOR_TOTAL = "j_total"
+J_TOR_OHMIC = "j_ohmic"
+J_TOR_EXTERNAL = "j_external"
+J_TOR_BOOTSTRAP = "j_bootstrap"
 I_BOOTSTRAP = "I_bootstrap"
 
 # Core transport.
@@ -494,6 +496,7 @@ def _save_core_profiles(
         "Ip_profile_face": IP_PROFILE,
         "q_face": Q,
         "s_face": MAGNETIC_SHEAR,
+        "j_total": J_TOR_TOTAL,  # j_total is the fsa toroidal current, dI/dS
     }
 
     core_profile_field_names = {
@@ -630,7 +633,7 @@ def _save_core_sources(
       else:
         xr_dict[f"p_{profile}_e"] = self._stacked_core_sources.T_e[profile]
     for profile in self._stacked_core_sources.psi:
-      xr_dict[f"j_{profile}"] = self._stacked_core_sources.psi[profile]
+      xr_dict[f"j_parallel_{profile}"] = self._stacked_core_sources.psi[profile]
     for profile in self._stacked_core_sources.n_e:
       xr_dict[f"s_{profile}"] = self._stacked_core_sources.n_e[profile]
 

torax/_src/output_tools/post_processing.py

@@ -146,9 +146,18 @@ class PostProcessedOutputs:
     rho_q_3_1_second: Second outermost rho_norm value that intercepts the q=3/1
       plane. If no intercept is found, set to -inf.
     I_bootstrap: Total bootstrap current [A].
-    j_external: Total current density from psi sources which are external to the
-      plasma (aka not bootstrap) [A m^-2]
-    j_ohmic: Ohmic current density [A/m^2]
+    j_parallel_external: Parallel current density from external psi sources
+      (i.e., excluding bootstrap) [A m^-2]
+    j_parallel_ohmic: Parallel ohmic current density [Am^-2]
+    j_total: Total toroidal current density [Am^-2]
+    j_bootstrap: Toroidal bootstrap current density [Am^-2]
+    j_ohmic: Toroidal ohmic current density [Am^-2]
+    j_external: Toroidal current density from external psi sources
+      (i.e., excluding bootstrap) [A m^-2]
+    j_generic_current: Toroidal current density from generic current source
+      [Am^-2]
+    j_ecrh: Toroidal current density from electron cyclotron heating
+      and current source [Am^-2]
     S_gas_puff: Integrated gas puff source [s^-1]
     S_pellet: Integrated pellet source [s^-1]
     S_generic_particle: Integrated generic particle source [s^-1]
@@ -233,8 +242,14 @@ class PostProcessedOutputs:
   rho_q_3_1_first: array_typing.FloatScalar
   rho_q_3_1_second: array_typing.FloatScalar
   I_bootstrap: array_typing.FloatScalar
-  j_external: array_typing.FloatVector
+  j_parallel_external: array_typing.FloatVector
+  j_parallel_ohmic: array_typing.FloatVector
+  j_total: array_typing.FloatVector
+  j_bootstrap: array_typing.FloatVector
   j_ohmic: array_typing.FloatVector
+  j_external: array_typing.FloatVector
+  j_generic_current: array_typing.FloatVector
+  j_ecrh: array_typing.FloatVector
   S_gas_puff: array_typing.FloatScalar
   S_pellet: array_typing.FloatScalar
   S_generic_particle: array_typing.FloatScalar
@@ -337,8 +352,15 @@ def zeros(cls, geo: geometry.Geometry) -> typing_extensions.Self:
         rho_q_2_1_second=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         rho_q_3_1_second=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         I_bootstrap=jnp.array(0.0, dtype=jax_utils.get_dtype()),
-        j_external=jnp.zeros(geo.rho_face.shape),
+        # TODO (v2): rename j_* to j_toroidal_* for clarity
+        j_parallel_external=jnp.zeros(geo.rho_face.shape),
+        j_parallel_ohmic=jnp.zeros(geo.rho_face.shape),
+        j_total=jnp.zeros(geo.rho_face.shape),
+        j_bootstrap=jnp.zeros(geo.rho_face.shape),
         j_ohmic=jnp.zeros(geo.rho_face.shape),
+        j_external=jnp.zeros(geo.rho_face.shape),
+        j_generic_current=jnp.zeros(geo.rho_face.shape),
+        j_ecrh=jnp.zeros(geo.rho_face.shape),
         S_gas_puff=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         S_pellet=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         S_generic_particle=jnp.array(0.0, dtype=jax_utils.get_dtype()),
@@ -514,7 +536,7 @@ def _calculate_integrated_sources(
       integrated['P_external_injected'] += integrated[f'{value}']
 
   for key, value in CURRENT_SOURCE_TRANSFORMATIONS.items():
-    integrated[f'{value}'] = _get_integrated_source_value(
+    integrated[value] = _get_integrated_source_value(
         core_sources.psi, key, geo, math_utils.area_integration
     )
 
@@ -590,10 +612,9 @@ def make_post_processed_outputs(
       runtime_params,
   )
   # Calculate fusion gain with a zero division guard.
-  Q_fusion = (
-      integrated_sources['P_fusion']
-      / (integrated_sources['P_external_total']
-         + constants.CONSTANTS.eps))
+  Q_fusion = integrated_sources['P_fusion'] / (
+      integrated_sources['P_external_total'] + constants.CONSTANTS.eps
+  )
 
   P_LH_hi_dens, P_LH_min, P_LH, n_e_min_P_LH = (
       scaling_laws.calculate_plh_scaling_factor(
@@ -756,11 +777,48 @@ def cumulative_values():
       sim_state.core_sources.bootstrap_current.j_bootstrap, sim_state.geometry
   )
 
-  j_external = sum(sim_state.core_sources.psi.values())
-  psi_current = (
-      j_external + sim_state.core_sources.bootstrap_current.j_bootstrap
+  # Parallel current densities
+  # j_total is toroidal by default (see psi_calculations.calc_j_total)
+  j_parallel_total = psi_calculations.j_tor_to_j_parallel(
+      sim_state.core_profiles.j_total, sim_state.geometry
+  )
+  # Core sources psi are all <j.B>/B0
+  j_parallel_external = sum(sim_state.core_sources.psi.values())
+  j_parallel_ohmic = (
+      j_parallel_total
+      - j_parallel_external
+      - sim_state.core_sources.bootstrap_current.j_bootstrap  # parallel by default
+  )
+
+  # Toroidal current densities
+  # j_total is toroidal by default (see psi_calculations.calc_j_total)
+  j_total = sim_state.core_profiles.j_total
+  # Other sources are parallel, so convert to toroidal
+  j_bootstrap = psi_calculations.j_parallel_to_j_tor(
+      sim_state.core_sources.bootstrap_current.j_bootstrap, sim_state.geometry
+  )
+  j_ohmic = psi_calculations.j_parallel_to_j_tor(
+      j_parallel_ohmic, sim_state.geometry
+  )
+  j_external = psi_calculations.j_parallel_to_j_tor(
+      j_parallel_external, sim_state.geometry
   )
-  j_ohmic = sim_state.core_profiles.j_total - psi_current
+  j_sources = {}
+  for source_name in ['ecrh', 'generic_current']:
+    if source_name in sim_state.core_sources.psi.keys():
+      j_sources[f'j_{source_name}'] = psi_calculations.j_parallel_to_j_tor(
+          sim_state.core_sources.psi[source_name], sim_state.geometry
+      )
+      j_sources[f'j_parallel_{source_name}'] = sim_state.core_sources.psi[
+          source_name
+      ]
+    else:
+      j_sources[f'j_{source_name}'] = jnp.array(
+          0.0, dtype=jax_utils.get_dtype()
+      )
+      j_sources[f'j_parallel_{source_name}'] = jnp.array(
+          0.0, dtype=jax_utils.get_dtype()
+      )
 
   beta_tor, beta_pol, beta_N = formulas.calculate_betas(
       sim_state.core_profiles, sim_state.geometry
@@ -813,8 +871,15 @@ def cumulative_values():
       rho_q_2_1_second=safety_factor_fit_outputs.rho_q_2_1_second,
       rho_q_3_1_second=safety_factor_fit_outputs.rho_q_3_1_second,
       I_bootstrap=I_bootstrap,
-      j_external=j_external,
+      j_parallel_total=j_parallel_total,
+      j_parallel_ohmic=j_parallel_ohmic,
+      j_parallel_bootstrap=j_parallel_bootstrap,
+      j_parallel_external=j_parallel_external,
+      j_total=j_total,
       j_ohmic=j_ohmic,
+      j_bootstrap=j_bootstrap,
+      j_external=j_external,
+      **j_sources,
       beta_tor=beta_tor,
       beta_pol=beta_pol,
       beta_N=beta_N,

torax/_src/physics/psi_calculations.py

@@ -26,6 +26,9 @@
     - calculate_psi_grad_constraint_from_Ip: Calculates the gradient
       constraint on the poloidal flux (psi) from Ip.
     - _calc_bpol2: Calculates square of poloidal field (Bp).
+      constraint on the poloidal flux (psi) from Ip.
+    - j_tor_to_j_parallel: Calculates <j.B>/B0 from j_tor = dI/dS.
+    - j_parallel_to_j_tor: Calculates j_tor = dI/dS from <j.B>/B0.
 """
 import jax
 from jax import numpy as jnp
@@ -338,3 +341,47 @@ def calculate_psidot_from_psi_sources(
   psidot = (jnp.dot(c_mat, psi.value) + c) / toc_psi
 
   return psidot
+
+
+def j_tor_to_j_parallel(
+    j_tor: array_typing.FloatVector, geo: geometry.Geometry
+) -> array_typing.FloatVector:
+  r"""Calculates <j.B>/B0 from j_tor = dI/dS.
+
+  The relationship is
+
+  .. math::
+
+    \frac{\langle j.B \rangle}{B_0} = \frac{F^2 \langle 1/R^2 \rangle}{B_0^2}
+    \left(\frac{j_{tor}}{2 \pi \rho} \frac{dS}{d\rho} - \frac{1}{2 \pi \rho F}
+    \frac{dF}{d\rho} \int_0^\rho j_{tor} \frac{dS}{d\rho'} d\rho' \right)
+  """
+  dF_drho = jnp.gradient(geo.F, geo.rho)
+  term1 = geo.spr * j_tor / (2 * jnp.pi * geo.rho)
+  term2 = (
+      dF_drho
+      / (2 * jnp.pi * geo.rho * geo.F)
+      * jnp.cumsum(j_tor * geo.spr * geo.drho)
+  )
+  return geo.F**2 * geo.g3 / geo.B_0**2 * (term1 - term2)
+
+
+def j_parallel_to_j_tor(
+    j_parallel: array_typing.FloatVector, geo: geometry.Geometry
+) -> array_typing.FloatVector:
+  r"""Calculates j_tor = dI/dS from <j.B>/B0.
+
+  The relationship is
+
+  .. math::
+    j_\mathrm{tor} = \frac{\partial}{\partial S} (2 \pi B_0^2 F \int_{0}^{\rho'}
+    \frac{\langle j.B \rangle}{B_0 F^3 \langle 1/R^2 \rangle} \rho' d\rho')
+  """
+  I = (
+      2
+      * jnp.pi
+      * geo.B_0**2
+      * geo.F
+      * jnp.cumsum(j_parallel * geo.rho / (geo.F**3 * geo.g3))
+  )
+  return jnp.gradient(I, geo.area)

torax/_src/physics/tests/psi_calculations_test.py

@@ -190,6 +190,28 @@ def test_calc_Wpol(self):
     # our circular geometry, but approximates it at low inverse aspect ratio.
     np.testing.assert_allclose(calculated_Wpol, expected_Wpol, rtol=1e-3)
 
+  # pylint: enable=invalid-name
+
+  def test_calc_j_tor_from_j_parallel(self):
+    geo = None
+    j_tor_truth = None
+    j_parallel_truth = None
+    j_tor_from_j_parallel = psi_calculations.j_parallel_to_j_tor(
+        j_parallel_truth, geo
+    )
+    np.testing.assert_allclose(j_tor_from_j_parallel, j_tor_truth, rtol=1e-6)
+
+  def test_calc_j_parallel_from_j_tor(self):
+    geo = None
+    j_tor_truth = None
+    j_parallel_truth = None
+    j_parallel_from_j_tor = psi_calculations.j_tor_to_j_parallel(
+        j_tor_truth, geo
+    )
+    np.testing.assert_allclose(
+        j_parallel_from_j_tor, j_parallel_truth, rtol=1e-6
+    )
+
 
 if __name__ == '__main__':
   absltest.main()

torax/_src/sources/generic_current_source.py

@@ -25,6 +25,7 @@
 from torax._src.config import runtime_params_slice
 from torax._src.geometry import geometry
 from torax._src.neoclassical.conductivity import base as conductivity_base
+from torax._src.physics import psi_calculations
 from torax._src.sources import base as source_base
 from torax._src.sources import runtime_params as runtime_params_lib
 from torax._src.sources import source
@@ -58,7 +59,8 @@ def calculate_generic_current(
     unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
     unused_conductivity: conductivity_base.Conductivity | None,
 ) -> tuple[array_typing.FloatVectorCell, ...]:
-  """Calculates the external current density profiles on the cell grid."""
+  """Calculates the external parallel current density profile on the cell
+  grid."""
   source_params = runtime_params.sources[source_name]
   # pytype: enable=name-error
   assert isinstance(source_params, RuntimeParams)
@@ -73,8 +75,9 @@ def calculate_generic_current(
   )
 
   Cext = I_generic / math_utils.area_integration(generic_current_form, geo)
-  generic_current_profile = Cext * generic_current_form
-  return (generic_current_profile,)
+  j_tor = Cext * generic_current_form
+
+  return (psi_calculations.j_tor_to_j_parallel(j_tor, geo),)
 
 
 def _calculate_I_generic(

torax/_src/sources/source_profiles.py

@@ -121,11 +121,20 @@ def merge(
     )
 
   def total_psi_sources(self, geo: geometry.Geometry) -> jax.Array:
-    total = self.bootstrap_current.j_bootstrap
-    total += sum(self.psi.values())
-    mu0 = constants.CONSTANTS.mu_0
-    prefactor = 8 * geo.vpr * jnp.pi**2 * geo.B_0 * mu0 * geo.Phi_b / geo.F**2
-    return -total * prefactor
+    # All psi sources are assumed to be parallel to the magnetic field, ie
+    # self.psi.values() is <j.B> / B0
+    total_j_dot_B_over_B0 = self.bootstrap_current.j_bootstrap
+    total_j_dot_B_over_B0 += sum(self.psi.values())
+    total_j_dot_B = total_j_dot_B_over_B0 * geo.B_0
+    prefactor = (
+        8
+        * geo.vpr
+        * jnp.pi**2
+        * constants.CONSTANTS.mu_0
+        * geo.Phi_b
+        / geo.F**2
+    )
+    return -total_j_dot_B * prefactor
 
   def total_sources(
       self,