Change expected_impurities to excluded_impurities in plasma_composition_from_IMAS

torax/_src/imas_tools/input/core_profiles.py

@@ -117,8 +117,8 @@ def profile_conditions_from_IMAS(
 def plasma_composition_from_IMAS(
     ids: ids_toplevel.IDSToplevel,
     t_initial: float | None = None,
-    expected_impurities: Collection[str] | None = None,
-    main_ions_symbols: Collection[str] = constants.HYDROGENIC_IONS,
+    excluded_impurities: Collection[str] | None = None,
+    main_ions_symbols: Collection[str] | None = None,
 ) -> Mapping[str, Any]:
   """Returns dict with args for plasma_composition config from a given ids.
 
@@ -140,13 +140,12 @@ def plasma_composition_from_IMAS(
       initial time will be the time of the first time slice of the ids. Else all
       time slices will be shifted such that the first time slice has time =
       t_initial.
-    expected_impurities: Optional arg to check that the input IDS contains the
-      wanted impurity species and raise and error if not, or if its density is
-      empty.
+    excluded_impurities: Optional arg to specify which impurities from the IDS
+      should not be parsed.
     main_ions_symbols: collection of ions to be used to define the main_ion
-      mixture. If value is not the default one, will check that the given ions
-      exist in the IDS and their density is filled. Default are hydrogenic ions
-      H, D, T.
+      mixture. If value is not None, will check that the given ions
+      exist in the IDS and their density is filled. If not explicitly provided,
+      will parse H, D, T as main ions.
 
   Returns:
     The updated fields read from the IDS that can be used to completely or
@@ -155,12 +154,23 @@ def plasma_composition_from_IMAS(
   profiles_1d, rhon_array, time_array = loader.get_time_and_radial_arrays(
       ids, t_initial
   )
-  # Check that the expected ions are present in the IDS
-  ids_ions = [ion.name for ion in profiles_1d[0].ion if ion.density]
-  if expected_impurities:
-    _check_expected_ions_presence(ids_ions, expected_impurities)
-  if main_ions_symbols is not constants.HYDROGENIC_IONS:
-    _check_expected_ions_presence(ids_ions, main_ions_symbols)
+  # Parse only ions with non empty density and not in excluded_impurities.
+  if excluded_impurities:
+    parsed_ions = [
+        ion.name
+        for ion in profiles_1d[0].ion
+        if ion.density and ion.name not in excluded_impurities
+    ]
+  else:
+    parsed_ions = [ion.name for ion in profiles_1d[0].ion if ion.density]
+  # main_ions_symbols not explicitly provided: no validation of main ions and
+  # value set to hydrogenic ions.
+  if main_ions_symbols is None:
+    main_ions_symbols = constants.HYDROGENIC_IONS
+  # main_ions_symbols explicitly provided: validate ions presence in IDS.
+  else:
+    validation.validate_main_ions_presence(parsed_ions, main_ions_symbols)
+  validation.validate_core_profiles_ions(parsed_ions)
 
   Z_eff = (
       time_array,
@@ -177,7 +187,7 @@ def plasma_composition_from_IMAS(
       # that instead of using a try-except.
       # Case ids is plasma_profiles in early DDv4 releases.
       symbol = str(profiles_1d[0].ion[ion].label)
-    if symbol in constants.ION_PROPERTIES_DICT.keys():
+    if symbol in parsed_ions:
       # Fill main ions
       if symbol in main_ions_symbols:
         main_ion_density[symbol] = [
@@ -210,21 +220,3 @@ def plasma_composition_from_IMAS(
           "species": impurity_species,
       },
   }
-
-
-def _check_expected_ions_presence(
-    ids_ions: list[str],
-    expected_ions: Collection[str],
-) -> None:
-  """Checks that the expected_ions symbols are in the ids_ions."""
-  for ion in expected_ions:
-    if ion not in constants.ION_PROPERTIES_DICT.keys():
-      raise (KeyError(f"{ion} is not a valid symbol of a TORAX valid ion."))
-    if ion not in ids_ions:
-      raise (
-          ValueError(
-              f"The expected ion {ion} cannot be found in the input"
-              " IDS or has no valid data. \n Please check that the IDS is"
-              " properly filled"
-          )
-      )

torax/_src/imas_tools/input/tests/core_profiles_test.py

@@ -115,9 +115,9 @@ def test_imas_plasma_composition(self):
     path = "core_profiles_ddv4_iterhybrid_rampup_conditions.nc"
     core_profiles_in = loader.load_imas_data(path, "core_profiles")
     # Indivual ion info empty in the inital IDS so create fake ions data
-    core_profiles_in.profiles_1d[0].ion.resize(3)
-    core_profiles_in.profiles_1d[1].ion.resize(3)
-    core_profiles_in.global_quantities.ion.resize(3)
+    core_profiles_in.profiles_1d[0].ion.resize(4)
+    core_profiles_in.profiles_1d[1].ion.resize(4)
+    core_profiles_in.global_quantities.ion.resize(4)
     core_profiles_in.profiles_1d[0].ion[0].name = "D"
     core_profiles_in.profiles_1d[0].ion[0].density = [9e19, 3e19]
     core_profiles_in.profiles_1d[1].ion[0].density = [9e19, 3e19]
@@ -132,40 +132,47 @@ def test_imas_plasma_composition(self):
     core_profiles_in.profiles_1d[1].ion[2].density = (
         core_profiles_in.profiles_1d[1].electrons.density / 100
     )
+    core_profiles_in.profiles_1d[0].ion[3].name = "He5"
+    core_profiles_in.profiles_1d[1].ion[3].name = "He5"
+    core_profiles_in.profiles_1d[0].ion[3].density = (
+        core_profiles_in.profiles_1d[0].electrons.density / 100
+    )
+    core_profiles_in.profiles_1d[1].ion[3].density = (
+        core_profiles_in.profiles_1d[1].electrons.density / 100
+    )
 
-    with self.subTest(name="Missing expected impurity."):
+    with self.subTest(name="Unrecognized impurity not excluded."):
       self.assertRaises(
-          ValueError,
+          KeyError,
           core_profiles.plasma_composition_from_IMAS,
           core_profiles_in,
           None,
           None,
-          ["Xe"],
       )
     with self.subTest(name="Missing expected main ion."):
       self.assertRaises(
           ValueError,
           core_profiles.plasma_composition_from_IMAS,
           core_profiles_in,
           None,
+          ["He5"],
           ["H"],
-          None,
       )
-    with self.subTest(name="Invalid ion name."):
+    with self.subTest(name="Invalid main ion name."):
       self.assertRaises(
           KeyError,
           core_profiles.plasma_composition_from_IMAS,
           core_profiles_in,
           None,
-          None,
           ["He5"],
+          ["De"],
       )
     with self.subTest(name="Test config is properly built."):
       plasma_composition_data = core_profiles.plasma_composition_from_IMAS(
           core_profiles_in,
           t_initial=None,
           main_ions_symbols=["D", "T"],
-          expected_impurities=["Ne"],
+          excluded_impurities=["He5"],
       )
       config["plasma_composition"] = plasma_composition_data
       config["plasma_composition"]["impurity"]["impurity_mode"] = "n_e_ratios"

torax/_src/imas_tools/input/tests/validation_test.py

@@ -38,6 +38,16 @@ def test_validate_core_profiles_warns_on_missing_optional_quantities(self):
     self.assertIn("The IDS is missing the global_quantities.v_loop quantity.",
                   logs.output[0])
 
+  def test_validate_core_profiles_ions_raises_on_unrecognized_ions(self):
+    parsed_ions = ["D", "T", "He5"]
+    with self.assertRaises(KeyError):
+      validation.validate_core_profiles_ions(parsed_ions)
+
+  def test_validate_main_ions_presence_raises_on_missing_main_ions(self):
+    parsed_ions = ["D", "He", "C"]
+    main_ions_symbols = ["D", "T"]
+    with self.assertRaises(ValueError):
+      validation.validate_main_ions_presence(parsed_ions, main_ions_symbols)
 
 if __name__ == "__main__":
   absltest.main()

torax/_src/imas_tools/input/validation.py

@@ -15,9 +15,10 @@
 
 import functools
 import logging
-from typing import Any
+from typing import Any, Collection
 
 from imas import ids_toplevel
+from torax._src import constants
 
 
 def _get_nested_attr(obj: Any, path: str) -> Any:
@@ -65,3 +66,42 @@ def validate_core_profiles_ids(ids: ids_toplevel.IDSToplevel) -> None:
           "global_quantities.ip",
       ),
   )
+
+
+def validate_core_profiles_ions(
+    parsed_ions: list[str],
+) -> None:
+  """Checks if all parsed ions are recognized."""
+  for ion in parsed_ions:
+    # ion is casted to str to avoid issues with imas string types.
+    if str(ion) not in constants.ION_PROPERTIES_DICT.keys():
+      raise (
+          KeyError(
+              f"{ion} is present in the IDS but not a valid TORAX ion. Check"
+              "typing or add the ion to the excluded_impurities."
+          )
+      )
+
+
+def validate_main_ions_presence(
+    parsed_ions: list[str],
+    main_ion_symbols: Collection[str],
+) -> None:
+  """Checks that items in main_ion_symbols are present in a list of ions parsed
+  from a given core_profiles IDS."""
+  for ion in main_ion_symbols:
+    if ion not in constants.ION_PROPERTIES_DICT.keys():
+      raise (
+          KeyError(
+              f"{ion} is not a valid symbol of a TORAX valid ion. Please"
+              " check typing of main_ion_symbols."
+          )
+      )
+    if ion not in parsed_ions:
+      raise (
+          ValueError(
+              f"The expected main ion {ion} cannot be found in the input"
+              " IDS or has no valid data. \n Please check that the IDS is"
+              " properly filled"
+          )
+      )

torax/tests/test_data/test_imas_profiles_and_geo.py

@@ -33,7 +33,10 @@
 )
 imas_data1 = core_profiles.profile_conditions_from_IMAS(imas_profiles1)
 imas_data2 = core_profiles.plasma_composition_from_IMAS(
-    imas_profiles2, t_initial=0.0, expected_impurities=["He3", "Be"]
+    imas_profiles2,
+    t_initial=0.0,
+    excluded_impurities=["Sn", "B"],
+    main_ions_symbols=["D", "T"],
 )
 
 CONFIG = copy.deepcopy(test_iterhybrid_predictor_corrector.CONFIG)