Update IMAS DD version to 4.1.1

torax/_src/imas_tools/input/core_profiles.py

@@ -194,13 +194,7 @@ def plasma_composition_from_IMAS(
   impurity_species = {}
   main_ion_density = {}
   for ion in range(len(profiles_1d[0].ion)):
-    try:
-      symbol = str(profiles_1d[0].ion[ion].name)
-    except AttributeError:
-      # TODO(b/459479939): i/539) - Indicate supported dd_versions and switch on
-      # that instead of using a try-except.
-      # Case ids is plasma_profiles in early DDv4 releases.
-      symbol = str(profiles_1d[0].ion[ion].label)
+    symbol = str(profiles_1d[0].ion[ion].name)
     if symbol in parsed_ions:
       # Fill main ions
       if symbol in main_ions_symbols:

torax/_src/imas_tools/input/loader.py

@@ -23,8 +23,8 @@
 from torax._src import path_utils
 
 # Names of the IDSs that can be loaded and used in TORAX.
-IDS = Literal["core_profiles", "plasma_profiles", "equilibrium"]
-_TORAX_IMAS_DD_VERSION = "4.0.0"
+IDS = Literal["core_profiles", "plasma_profiles", "equilibrium", "core_sources"]
+_TORAX_IMAS_DD_VERSION = "4.1.1"
 
 
 def load_imas_data(
@@ -47,7 +47,7 @@ def load_imas_data(
     directory: The directory of the IDS to load.
     explicit_convert: Whether to explicitly convert the IDS to the current DD
       version. If True, an explicit conversion will be attempted. Explicit
-      conversion is recommended when converting between major DD versions.
+      conversion is mandatory when converting between major DD versions.
       https://imas-python.readthedocs.io/en/latest/multi-dd.html#conversion-of-idss-between-dd-versions
   Returns:
     An IDS object.

torax/_src/imas_tools/output/core_profiles.py

@@ -26,6 +26,7 @@
 from torax._src import constants
 from torax._src import state
 from torax._src.geometry import geometry as geometry_lib
+from torax._src.imas_tools.input import loader
 from torax._src.output_tools import output
 from torax._src.output_tools import post_processing
 from torax._src.sources import source_profiles
@@ -103,6 +104,7 @@ def _fill_metadata(ids: ids_toplevel.IDSToplevel):
   )
   ids.ids_properties.homogeneous_time = 1
   ids.ids_properties.creation_date = datetime.date.today().isoformat()
+  ids.ids_properties.version_put.data_dictionary = loader._TORAX_IMAS_DD_VERSION
   ids.code.name = 'TORAX'
   ids.code.description = (
       'TORAX is a differentiable tokamak core transport simulator aimed for'
@@ -345,13 +347,7 @@ def _fill_main_ions(
 ) -> None:
   """Fills main ion quantities for the IDS."""
   ion_properties = constants.ION_PROPERTIES_DICT[symbol]
-  # TODO(b/459479939): i/539) - Indicate supported dd_versions and switch on
-  # that instead of using a try-except.
-  try:
-    ids.profiles_1d[i].ion[ion].name = symbol
-  except AttributeError:
-    # Case ids is plasma_profiles in early DDv4 releases.
-    ids.profiles_1d[i].ion[ion].label = symbol
+  ids.profiles_1d[i].ion[ion].name = symbol
   ids.profiles_1d[i].ion[ion].temperature = T_i
   ids.profiles_1d[i].ion[ion].density = n_i * frac
   ids.profiles_1d[i].ion[ion].density_thermal = n_i * frac
@@ -393,13 +389,7 @@ def _fill_impurities(
   index = num_of_main_ions + ion
   ion_properties = constants.ION_PROPERTIES_DICT[symbol]
   # TODO(b/459479939): i/1814) - Map ion.z_ion_1d
-  # TODO(b/459479939): i/539) - Indicate supported dd_versions and switch on
-  # that instead of using a try-except.
-  try:
-    ids.profiles_1d[i].ion[index].name = symbol
-  except AttributeError:
-    # Case ids is plasma_profiles in early DDv4 releases.
-    ids.profiles_1d[i].ion[index].label = symbol
+  ids.profiles_1d[i].ion[index].name = symbol
   ids.profiles_1d[i].ion[index].temperature = T_i
   # TODO(b/459479939): i/1814) - Map density from computed frac
   ids.profiles_1d[i].ion[index].density_fast = np.zeros(

torax/_src/imas_tools/output/equilibrium.py

@@ -13,12 +13,14 @@
 # limitations under the License.
 """Helper functions for IMAS output."""
 
+import datetime
 import logging
 
 import imas
 from imas import ids_toplevel
 import numpy as np
 from torax._src.geometry import standard_geometry
+from torax._src.imas_tools.input import loader
 from torax._src.orchestration import sim_state as sim_state_lib
 from torax._src.output_tools import post_processing
 
@@ -57,9 +59,20 @@ def torax_state_to_imas_equilibrium(
   # Rebuilding the equilibrium from the geometry object
   equilibrium = imas.IDSFactory().equilibrium()
   equilibrium.ids_properties.homogeneous_time = 1
+  equilibrium.ids_properties.creation_date = datetime.date.today().isoformat()
+  equilibrium.ids_properties.version_put.data_dictionary = (
+      loader._TORAX_IMAS_DD_VERSION
+  )
   equilibrium.ids_properties.comment = (
       "equilibrium IDS built from ToraxSimState object."
   )
+  equilibrium.code.name = "TORAX"
+  equilibrium.code.description = (
+      "TORAX is a differentiable tokamak core transport simulator aimed for"
+      " fast and accurate forward modelling, pulse-design, trajectory"
+      " optimization, and controller design workflows."
+  )
+  equilibrium.code.repository = "https://github.com/google-deepmind/torax"
   equilibrium.time.resize(1)
   equilibrium.time = [sim_state.t]
   equilibrium.vacuum_toroidal_field.r0 = geometry.R_major
@@ -72,7 +85,7 @@ def torax_state_to_imas_equilibrium(
   eq.profiles_1d.psi = core_profiles.psi.face_value()
   psi_axis = core_profiles.psi.face_value()[0]
   psi_boundary = core_profiles.psi.face_value()[-1]
-  eq.global_quantities.psi_axis = psi_axis
+  eq.global_quantities.psi_magnetic_axis = psi_axis
   eq.global_quantities.psi_boundary = psi_boundary
   eq.profiles_1d.psi_norm = (core_profiles.psi.face_value() - psi_axis) / (
       psi_boundary - psi_axis