diff --git a/torax/_src/imas_tools/input/core_profiles.py b/torax/_src/imas_tools/input/core_profiles.py index 4318d74da..cfa2de4f2 100644 --- a/torax/_src/imas_tools/input/core_profiles.py +++ b/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: diff --git a/torax/_src/imas_tools/input/loader.py b/torax/_src/imas_tools/input/loader.py index 00c31a82c..152a79d03 100644 --- a/torax/_src/imas_tools/input/loader.py +++ b/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. diff --git a/torax/_src/imas_tools/output/core_profiles.py b/torax/_src/imas_tools/output/core_profiles.py index 941f7ad70..0ed392a10 100644 --- a/torax/_src/imas_tools/output/core_profiles.py +++ b/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( diff --git a/torax/_src/imas_tools/output/equilibrium.py b/torax/_src/imas_tools/output/equilibrium.py index b47b014b0..74b4f4742 100644 --- a/torax/_src/imas_tools/output/equilibrium.py +++ b/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