Add non-inductive current outputs to post-processing.

This change introduces several new outputs related to non-inductive current: total non-inductive current density (toroidal and parallel), total external current, total non-inductive current, and the non-inductive current fraction. These are derived from existing bootstrap and external current sources. A new test case is added to verify these calculations, and the output documentation is updated.

Sim tests are regenerated due to the new outputs.

PiperOrigin-RevId: 861313287

Author: jcitrin

docs/output.rst

@@ -172,6 +172,14 @@ class PostProcessedOutputs:
       source [Am^-2]
     j_parallel_ecrh: Toroidal current density from electron cyclotron heating
       and current source [Am^-2]
+    j_non_inductive: Total toroidal non-inductive current density [Am^-2]
+    j_parallel_non_inductive: Total parallel non-inductive current density
+      [Am^-2]
+    I_external: Total external current [A]
+    I_non_inductive: Total non-inductive current [A]
+    f_non_inductive: Non-inductive current fraction of the total current
+      [dimensionless]
+    f_bootstrap: Bootstrap current fraction of the total current [dimensionless]
     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]
@@ -267,12 +275,19 @@ class PostProcessedOutputs:
   # TODO(b/434175938): rename j_* to j_toroidal_* for clarity
   j_parallel_total: array_typing.FloatVector
   j_external: array_typing.FloatVector
+  j_parallel_external: array_typing.FloatVector
   j_ohmic: array_typing.FloatVector
   j_parallel_ohmic: array_typing.FloatVector
   j_bootstrap: array_typing.FloatVector
   j_bootstrap_face: array_typing.FloatVector
   j_generic_current: array_typing.FloatVector
   j_ecrh: array_typing.FloatVector
+  j_non_inductive: array_typing.FloatVector
+  j_parallel_non_inductive: array_typing.FloatVector
+  I_external: array_typing.FloatScalar
+  I_non_inductive: array_typing.FloatScalar
+  f_non_inductive: array_typing.FloatScalar
+  f_bootstrap: array_typing.FloatScalar
   S_gas_puff: array_typing.FloatScalar
   S_pellet: array_typing.FloatScalar
   S_generic_particle: array_typing.FloatScalar
@@ -374,6 +389,13 @@ def zeros(cls, geo: geometry.Geometry) -> typing_extensions.Self:
         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),
+        j_non_inductive=jnp.zeros(geo.rho_face.shape),
+        j_parallel_external=jnp.zeros(geo.rho_face.shape),
+        j_parallel_non_inductive=jnp.zeros(geo.rho_face.shape),
+        I_external=jnp.array(0.0, dtype=jax_utils.get_dtype()),
+        I_non_inductive=jnp.array(0.0, dtype=jax_utils.get_dtype()),
+        f_non_inductive=jnp.array(0.0, dtype=jax_utils.get_dtype()),
+        f_bootstrap=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         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()),
@@ -925,6 +947,10 @@ def cumulative_values():
   I_bootstrap = math_utils.area_integration(
       j_toroidal_bootstrap, sim_state.geometry
   )
+  I_external = math_utils.area_integration(
+      j_toroidal_external, sim_state.geometry
+  )
+  I_non_inductive = I_bootstrap + I_external
 
   beta_tor, beta_pol, beta_N = formulas.calculate_betas(
       sim_state.core_profiles, sim_state.geometry
@@ -998,6 +1024,17 @@ def cumulative_values():
       j_external=j_toroidal_external,
       j_ecrh=j_toroidal_sources['j_ecrh'],
       j_generic_current=j_toroidal_sources['j_generic_current'],
+      j_non_inductive=j_toroidal_bootstrap + j_toroidal_external,
+      j_parallel_external=j_parallel_external,
+      j_parallel_non_inductive=j_parallel_bootstrap + j_parallel_external,
+      I_external=I_external,
+      I_non_inductive=I_non_inductive,
+      f_non_inductive=math_utils.safe_divide(
+          I_non_inductive, sim_state.core_profiles.Ip_profile_face[-1]
+      ),
+      f_bootstrap=math_utils.safe_divide(
+          I_bootstrap, sim_state.core_profiles.Ip_profile_face[-1]
+      ),
       beta_tor=beta_tor,
       beta_pol=beta_pol,
       beta_N=beta_N,

torax/_src/output_tools/post_processing.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import dataclasses
 from absl.testing import absltest
 from absl.testing import parameterized
 from jax import numpy as jnp
 import numpy as np
 import scipy
 from torax._src import jax_utils
+from torax._src import math_utils
 from torax._src import state
 from torax._src.config import build_runtime_params
 from torax._src.core_profiles import initialization
@@ -31,6 +33,8 @@
 from torax._src.test_utils import sim_test_case
 from torax._src.torax_pydantic import model_config
 
+# pylint: disable=invalid-name
+
 
 class PostProcessingTest(parameterized.TestCase):
 
@@ -234,6 +238,89 @@ def test_zero_sources_do_not_make_nans(self):
         post_processed_outputs.check_for_errors(), state.SimError.NO_ERROR
     )
 
+  def test_current_outputs(self):
+    """Checks calculation of current-related outputs."""
+    # Setup non-zero bootstrap current
+    ones = np.ones_like(self.geo.rho)
+    bootstrap_current = bootstrap_current_base.BootstrapCurrent(
+        j_parallel_bootstrap=1.0 * ones,
+        j_parallel_bootstrap_face=1.0 * np.ones_like(self.geo.rho_face),
+    )
+
+    # Source profiles with parallel currents (normalized <j.B>/B0)
+    # generic: 2.0, ecrh: 2.0. Sum external = 4.0.
+    source_profiles = dataclasses.replace(
+        self.source_profiles, bootstrap_current=bootstrap_current
+    )
+
+    # Mock j_total (toroidal).
+    # We leave core_profiles as is (from setUp), but update source_profiles.
+
+    input_state = sim_state.SimState(
+        t=jnp.array(0.0),
+        dt=jnp.array(1e-3),
+        core_profiles=self.core_profiles,
+        core_transport=state.CoreTransport.zeros(self.geo),
+        core_sources=source_profiles,
+        geometry=self.geo,
+        solver_numeric_outputs=state.SolverNumericOutputs(
+            solver_error_state=np.array(0, jax_utils.get_int_dtype()),
+            outer_solver_iterations=np.array(0, jax_utils.get_int_dtype()),
+            inner_solver_iterations=np.array(0, jax_utils.get_int_dtype()),
+            sawtooth_crash=False,
+        ),
+        edge_outputs=None,
+    )
+
+    outputs = post_processing.make_post_processed_outputs(
+        sim_state=input_state,
+        runtime_params=self.runtime_params,
+        previous_post_processed_outputs=post_processing.PostProcessedOutputs.zeros(
+            self.geo
+        ),
+    )
+
+    # Check parallel currents
+    # external = sum(psi) = 2 + 2 = 4
+    np.testing.assert_allclose(outputs.j_parallel_external, 4.0 * ones)
+
+    # non_inductive = bootstrap (1.0) + external (4.0) = 5.0
+    np.testing.assert_allclose(outputs.j_parallel_non_inductive, 5.0 * ones)
+
+    # Check toroidal currents relation
+    # j_non_inductive should be j_bootstrap + j_external
+    np.testing.assert_allclose(
+        outputs.j_non_inductive, outputs.j_bootstrap + outputs.j_external
+    )
+
+    # Check integrated currents
+    # I_non_inductive should be I_bootstrap + area_int(j_external)
+    I_external = math_utils.area_integration(outputs.j_external, self.geo)
+    np.testing.assert_allclose(outputs.I_external, I_external)
+
+    np.testing.assert_allclose(
+        outputs.I_non_inductive, outputs.I_bootstrap + I_external
+    )
+
+    # Check fraction
+    # f_non_inductive = I_non_inductive / Ip
+    # Ip comes from core_profiles.Ip_profile_face[-1]
+    ip = self.core_profiles.Ip_profile_face[-1]
+    # Code uses constants.CONSTANTS.eps for division guard
+    np.testing.assert_allclose(
+        outputs.f_non_inductive,
+        math_utils.safe_divide(outputs.I_non_inductive, ip),
+        rtol=1e-5,
+    )
+
+    # Check bootstrap fraction
+    # f_bootstrap = I_bootstrap / Ip
+    np.testing.assert_allclose(
+        outputs.f_bootstrap,
+        math_utils.safe_divide(outputs.I_bootstrap, ip),
+        rtol=1e-5,
+    )
+
 
 class PostProcessingSimTest(sim_test_case.SimTestCase):
   """Tests for the cumulative outputs."""