From 6d3e32390d7960cd63d8f55688d1be5ee5bcf450 Mon Sep 17 00:00:00 2001
From: Anushan Fernando <anushanf@google.com>
Date: Fri, 28 Feb 2025 12:00:06 -0800
Subject: [PATCH] Rename output variable name now we have hierarchical outputs.

PiperOrigin-RevId: 732213030
---
 torax/output.py                | 7 ++-----
 torax/plotting/plotruns_lib.py | 4 ++--
 torax/tests/sim.py             | 4 ++--
 3 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/torax/output.py b/torax/output.py
index 4f6905ee..abdde40d 100644
--- a/torax/output.py
+++ b/torax/output.py
@@ -69,10 +69,7 @@ class ToraxSimOutputs:
 JTOT = "jtot"
 JTOT_FACE = "jtot_face"
 JOHM = "johm"
-# TODO(b/338033916): rename when we have a solution for hierarchical outputs.
-# Add `core_profiles` prefix here to avoid name clash with
-# core_sources.generic_current.
-CORE_PROFILES_EXTERNAL_CURRENT = "external_current_source"
+EXTERNAL_CURRENT = "external_current_source"
 J_BOOTSTRAP = "j_bootstrap"
 J_BOOTSTRAP_FACE = "j_bootstrap_face"
 I_BOOTSTRAP = "I_bootstrap"
@@ -308,7 +305,7 @@ def _get_core_profiles(
     xr_dict[JTOT] = self.core_profiles.currents.jtot
     xr_dict[JTOT_FACE] = self.core_profiles.currents.jtot_face
     xr_dict[JOHM] = self.core_profiles.currents.johm
-    xr_dict[CORE_PROFILES_EXTERNAL_CURRENT] = (
+    xr_dict[EXTERNAL_CURRENT] = (
         self.core_profiles.currents.external_current_source
     )
 
diff --git a/torax/plotting/plotruns_lib.py b/torax/plotting/plotruns_lib.py
index 03715c71..e4267065 100644
--- a/torax/plotting/plotruns_lib.py
+++ b/torax/plotting/plotruns_lib.py
@@ -265,7 +265,7 @@ def _transform_data(ds: xr.Dataset):
         output.JTOT: 1e6,  # A/m^2 to MA/m^2
         output.JOHM: 1e6,  # A/m^2 to MA/m^2
         output.J_BOOTSTRAP: 1e6,  # A/m^2 to MA/m^2
-        output.CORE_PROFILES_EXTERNAL_CURRENT: 1e6,  # A/m^2 to MA/m^2
+        output.EXTERNAL_CURRENT: 1e6,  # A/m^2 to MA/m^2
         'generic_current_source': 1e6,  # A/m^2 to MA/m^2
         output.I_BOOTSTRAP: 1e6,  # A to MA
         output.IP_PROFILE_FACE: 1e6,  # A to MA
@@ -331,7 +331,7 @@ def _transform_data(ds: xr.Dataset):
       johm=core_profiles_dataset[output.JOHM].to_numpy(),
       j_bootstrap=core_profiles_dataset[output.J_BOOTSTRAP].to_numpy(),
       external_current_source=core_profiles_dataset[
-          output.CORE_PROFILES_EXTERNAL_CURRENT
+          output.EXTERNAL_CURRENT
       ].to_numpy(),
       j_ecrh=get_optional_data(
           core_sources_dataset, 'electron_cyclotron_source_j', 'cell'
diff --git a/torax/tests/sim.py b/torax/tests/sim.py
index 90199f7c..10c8fb84 100644
--- a/torax/tests/sim.py
+++ b/torax/tests/sim.py
@@ -613,7 +613,7 @@ def test_core_profiles_are_recomputable(self, test_config, halfway):
         output.J_BOOTSTRAP,
         output.J_BOOTSTRAP_FACE,
         output.JOHM,
-        output.CORE_PROFILES_EXTERNAL_CURRENT,
+        output.EXTERNAL_CURRENT,
         output.JTOT,
         output.JTOT_FACE,
         output.I_BOOTSTRAP,
@@ -822,7 +822,7 @@ def verify_core_profiles(ref_profiles, index, core_profiles):
   )
   np.testing.assert_allclose(
       core_profiles.currents.external_current_source,
-      ref_profiles[output.CORE_PROFILES_EXTERNAL_CURRENT][index, :],
+      ref_profiles[output.EXTERNAL_CURRENT][index, :],
   )
   np.testing.assert_allclose(
       core_profiles.currents.johm, ref_profiles[output.JOHM][index, :]