Add draft postprocessed outputs

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
@@ -63,9 +61,12 @@
 IP = "Ip"
 
 # Calculated or derived currents.
-J_OHMIC = "j_ohmic"
-J_EXTERNAL = "j_external"
-J_BOOTSTRAP = "j_bootstrap"
+J_TOR_OHMIC = "j_tor_ohmic"
+J_TOR_EXTERNAL = "j_tor_external"
+J_TOR_BOOTSTRAP = "j_tor_bootstrap"
+J_PARALLEL_OHMIC = "j_parallel_ohmic"
+J_PARALLEL_EXTERNAL = "j_parallel_external"
+J_PARALLEL_BOOTSTRAP = "j_parallel_bootstrap"
 I_BOOTSTRAP = "I_bootstrap"
 
 # Core transport.
@@ -630,7 +631,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

@@ -135,9 +135,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_toroidal_total: Total toroidal current density [Am^-2]
+    j_toroidal_bootstrap: Toroidal bootstrap current density [Am^-2]
+    j_toroidal_ohmic: Toroidal ohmic current density [Am^-2]
+    j_toroidal_external: Toroidal current density from external psi sources
+      (i.e., excluding bootstrap) [A m^-2]
+    j_toroidal_generic: Toroidal current density from generic current source
+      [Am^-2]
+    j_toroidal_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]
@@ -217,8 +226,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_ohmic: array_typing.FloatVector
+  j_parallel_external: array_typing.FloatVector
+  j_parallel_ohmic: array_typing.FloatVector
+  j_toroidal_total: array_typing.FloatVector
+  j_toroidal_bootstrap: array_typing.FloatVector
+  j_toroidal_ohmic: array_typing.FloatVector
+  j_toroidal_external: array_typing.FloatVector
+  j_toroidal_generic: array_typing.FloatVector
+  j_toroidal_ecrh: array_typing.FloatVector
   S_gas_puff: array_typing.FloatScalar
   S_pellet: array_typing.FloatScalar
   S_generic_particle: array_typing.FloatScalar
@@ -315,8 +330,14 @@ 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),
-        j_ohmic=jnp.zeros(geo.rho_face.shape),
+        j_parallel_external=jnp.zeros(geo.rho_face.shape),
+        j_parallel_ohmic=jnp.zeros(geo.rho_face.shape),
+        j_toroidal_total=jnp.zeros(geo.rho_face.shape),
+        j_toroidal_bootstrap=jnp.zeros(geo.rho_face.shape),
+        j_toroidal_ohmic=jnp.zeros(geo.rho_face.shape),
+        j_toroidal_external=jnp.zeros(geo.rho_face.shape),
+        j_toroidal_generic=jnp.zeros(geo.rho_face.shape),
+        j_toroidal_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()),
@@ -360,8 +381,8 @@ def check_for_errors(self):
     'icrh',
 ]
 CURRENT_SOURCE_TRANSFORMATIONS = {
-    'generic_current': 'I_aux_generic',
-    'ecrh': 'I_ecrh',
+    'generic_current': 'aux_generic',
+    'ecrh': 'ecrh',
 }
 PARTICLE_SOURCE_TRANSFORMATIONS = {
     'gas_puff': 'S_gas_puff',
@@ -385,6 +406,20 @@ def _get_integrated_source_value(
     return jnp.array(0.0, dtype=jax_utils.get_dtype())
 
 
+def _get_source_toroidal_current_value(
+    source_profiles_dict: dict[str, array_typing.FloatVector],
+    internal_source_name: str,
+    geo: geometry.Geometry,
+) -> jax.Array:
+  """Integrates a source profile if it exists, otherwise returns 0.0."""
+  if internal_source_name in source_profiles_dict:
+    return psi_calculations.j_parallel_to_j_tor(
+        source_profiles_dict[internal_source_name], geo
+    )
+  else:
+    return jnp.array(0.0, dtype=jax_utils.get_dtype())
+
+
 def _calculate_integrated_sources(
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
@@ -491,7 +526,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[f'I_{value}'] = _get_integrated_source_value(
         core_sources.psi, key, geo, math_utils.area_integration
     )
 
@@ -703,11 +738,36 @@ def make_post_processed_outputs(
       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
+  # Core sources psi are all <j.B>/B0
+  j_parallel_external = sum(sim_state.core_sources.psi.values())
+  j_parallel_ohmic = (
+      sim_state.core_profiles.j_total  # TODO: check whether j_total is parallel or toroidal
+      - j_parallel_external
+      - sim_state.core_sources.bootstrap_current.j_bootstrap  # parallel by default
+  )
+
+  # Toroidal current densities
+  # TODO: check whether j_total is parallel or toroidal
+  j_toroidal_total = psi_calculations.j_parallel_to_j_tor(
+      sim_state.core_profiles.j_total, sim_state.geometry
+  )
+  j_toroidal_bootstrap = psi_calculations.j_parallel_to_j_tor(
+      sim_state.core_sources.bootstrap_current.j_bootstrap, sim_state.geometry
+  )
+  j_toroidal_ohmic = psi_calculations.j_parallel_to_j_tor(
+      j_parallel_ohmic, sim_state.geometry
+  )
+  j_toroidal_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_toroidal_sources = {}
+  for internal_name, external_name in CURRENT_SOURCE_TRANSFORMATIONS.items():
+    j_toroidal_sources[f'j_toroidal_{external_name}'] = (
+        _get_source_toroidal_current_value(
+            core_sources.psi, internal_name, sim_state.geometry
+        )
+    )
 
   beta_tor, beta_pol, beta_N = formulas.calculate_betas(
       sim_state.core_profiles, sim_state.geometry
@@ -757,8 +817,13 @@ def make_post_processed_outputs(
       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_ohmic=j_ohmic,
+      j_parallel_external=j_parallel_external,
+      j_parallel_ohmic=j_parallel_ohmic,
+      j_toroidal_total=j_toroidal_total,
+      j_toroidal_bootstrap=j_toroidal_bootstrap,
+      j_toroidal_ohmic=j_toroidal_ohmic,
+      j_toroidal_external=j_toroidal_external,
+      **j_toroidal_sources,
       beta_tor=beta_tor,
       beta_pol=beta_pol,
       beta_N=beta_N,