integration test and unit test for imas geometry

torax/geometry/tests/geometry.py

@@ -15,16 +15,17 @@
 """Unit tests for torax.geometry."""
 
 import dataclasses
-import os
+import importlib
 
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
 from jax import numpy as jnp
 import numpy as np
-from torax.config import build_sim
+import pytest
+
 from torax.geometry import geometry
-from torax.geometry import geometry_loader
+
 
 # Internal import.
 
@@ -333,10 +334,18 @@ def f():
 
       _ = jax.jit(f)()
 
+  @pytest.mark.skipif(
+      importlib.util.find_spec('imaspy') is None,
+      reason='IMASPy optional dependency'
+  )
   @parameterized.parameters([
-      dict(equilibrium_object='ITERhybrid_COCOS17_IDS_ddv4.nc')])
+      dict(equilibrium_object='ITERhybrid_COCOS17_IDS_ddv4.nc'),
+      dict(equilibrium_object='ITERhybrid_COCOS17_IDS_hdf5_ddv4'),
+  ])
   def test_build_standard_geometry_from_IMAS(self, equilibrium_object):
     """Test that the default IMAS geometry can be built."""
+    if importlib.util.find_spec('imaspy') is None:
+      self.skipTest('IMASPy optional dependency')
     intermediate = geometry.StandardGeometryIntermediates.from_IMAS(equilibrium_object = equilibrium_object)
     geo = geometry.build_standard_geometry(intermediate)
 

torax/tests/sim.py

@@ -20,10 +20,12 @@
 
 import os
 from typing import Optional, Sequence
+import importlib
 
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
+import pytest
 from torax import output
 from torax import sim as sim_lib
 from torax import state
@@ -738,6 +740,15 @@ def test_update_new_mesh(self):
           )
       )
 
+  def test_imas(self):
+    """Integration test comparing to reference output from TORAX."""
+    if importlib.util.find_spec('imaspy') is None:
+      self.skipTest('IMASPy optional dependency')
+    self._test_torax_sim(
+        'test_imas.py',
+        _ALL_PROFILES,
+    )
+
 
 def verify_core_profiles(ref_profiles, index, core_profiles):
   """Verify core profiles matches a reference at given index."""

torax/tests/test_data/test_imas.py

@@ -0,0 +1,74 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests implementation of CHEASE geometry.
+
+Ip from parameters. implicit, Ti+Te, no Pei, no pedestal, constant chi.
+"""
+
+
+CONFIG = {
+    'runtime_params': {
+        'profile_conditions': {
+            'ne_bound_right': 0.5,
+            'Ip_tot': 15,
+            'set_pedestal': False,
+        },
+        'numerics': {
+            't_final': 1,
+        },
+    },
+    'geometry': {
+        'geometry_type': 'imas',
+        'equilibrium_object': 'ITERhybrid_COCOS17_IDS_hdf5_ddv4',
+        'Ip_from_parameters': True,
+    },
+    'sources': {
+        # Current sources (for psi equation)
+        'generic_current_source': {},
+        # Electron density sources/sink (for the ne equation).
+        'generic_particle_source': {
+            # total particle source
+            'S_tot': 0.0,
+        },
+        'gas_puff_source': {
+            # total pellet particles/s
+            'S_puff_tot': 0.0,
+        },
+        'pellet_source': {
+            # total pellet particles/s (continuous pellet model)
+            'S_pellet_tot': 0.0,
+        },
+        # Ion and electron heat sources (for the temp-ion and temp-el eqs).
+        'generic_ion_el_heat_source': {
+            # Gaussian width in normalized radial coordinate r
+            'w': 0.18202270915319393,
+        },
+        'qei_source': {
+            # multiplier for ion-electron heat exchange term for sensitivity
+            'Qei_mult': 0.0,
+        },
+    },
+    'pedestal': {},
+    'transport': {
+        'transport_model': 'constant',
+    },
+    'stepper': {
+        'stepper_type': 'linear',
+        'predictor_corrector': False,
+    },
+    'time_step_calculator': {
+        'calculator_type': 'chi',
+    },
+}

torax/torax_imastools/util.py

@@ -14,15 +14,14 @@
 
 """Useful functions for handling of IMAS IDSs and converts them into TORAX
 objects"""
-from typing import Dict, Any
-import os
 import datetime
 import importlib
+import os
+from typing import Any, Dict
 
 import numpy as np
-import yaml
 import scipy
-
+import yaml
 
 try:
     import imaspy
@@ -31,7 +30,6 @@
     IDSToplevel = Any
 
 from torax.geometry import geometry_loader
-from torax.geometry.geometry import Geometry
 from torax.state import ToraxSimState
 
 
@@ -44,9 +42,7 @@ def requires_module(module_name: str):
 
     def decorator(func):
         def wrapper(*args, **kwargs):
-            try:
-                importlib.import_module(module_name)
-            except ImportError:
+            if importlib.util.find_spec(module_name) is None:
                 raise ImportError(
                     f"Required module '{module_name}' is not installed. "
                     "Make sure you install the needed optional dependencies."
@@ -168,7 +164,7 @@ def write_ids_equilibrium_into_config(
 
 @requires_module("imaspy")
 def core_profiles_to_IMAS(
-    ids: IDSToplevel, state: ToraxSimState, geometry: Geometry
+    ids: IDSToplevel, state: ToraxSimState, geometry
 ) -> IDSToplevel:
     """Converts torax core_profiles to IMAS IDS.
     Takes the cell grid as a basis and converts values on face grid to cell.