Add placeholder tests pending updated IMAS testdata

Author: theo-brown

torax/_src/geometry/imas.py

@@ -88,6 +88,11 @@ def geometry_from_IMAS(
           f"t={slice_time} not in equilibrium.time."
           f"Using t={equilibrium.time[slice_index]} instead."
       )
+  if slice_index > len(equilibrium.time_slice):
+    raise IndexError(
+        f"slice_index={slice_index} out of range for IDS with "
+        f"{len(equilibrium.time_slice)} time slices"
+    )
   IMAS_data = equilibrium.time_slice[slice_index]
   B_0 = np.abs(IMAS_data.global_quantities.magnetic_axis.b_field_phi)
   R_major = np.asarray(equilibrium.vacuum_toroidal_field.r0)

torax/_src/geometry/tests/imas_test.py

@@ -18,6 +18,7 @@
 from absl.testing import parameterized
 import imas
 import numpy as np
+import pydantic
 from torax._src.geometry import geometry_loader
 from torax._src.geometry import imas as imas_geometry
 from torax._src.geometry import pydantic_model as geometry_pydantic_model
@@ -97,6 +98,60 @@ def test_IMAS_input_with_equilibrium_object(self):
     )
     config.build_geometry()
 
+  def test_IMAS_loading_specific_slice(self):
+    filename = 'ITERhybrid_COCOS17_IDS_ddv4.nc'
+    SLICE_TIME = 5
+    SLICE_INDEX = 1
+    config_at_0 = geometry_pydantic_model.IMASConfig(imas_filepath=filename)
+    config_at_slice_from_time = geometry_pydantic_model.IMASConfig(
+        imas_filepath=filename, slice_time=SLICE_TIME
+    )
+    config_at_slice_from_index = geometry_pydantic_model.IMASConfig(
+        imas_filepath=filename, slice_index=SLICE_INDEX
+    )
+    geo_at_0 = config_at_0.build_geometry()
+    geo_at_slice_from_time = config_at_slice_from_time.build_geometry()
+    geo_at_slice_from_index = config_at_slice_from_index.build_geometry()
+
+    for key in geo_at_0.__dict__.keys():
+      np.testing.assert_allclose(
+          geo_at_slice_from_time[key],
+          geo_at_slice_from_index[key],
+          err_msg=f'Value {key} mismatched between slice_time and slice_index',
+      )
+      with self.assertRaises(AssertionError):
+        np.testing.assert_allclose(
+            geo_at_0[key],
+            geo_at_slice_from_time[key],
+        )
+      with self.assertRaises(AssertionError):
+        np.testing.assert_allclose(
+            geo_at_0[key],
+            geo_at_slice_from_index[key],
+        )
+
+  def test_IMAS_raises_if_slice_specified_twice(self):
+    filename = 'ITERhybrid_COCOS17_IDS_ddv4.nc'
+    with self.assertRaisesRegex(
+        pydantic.ValidationError,
+        'exactly one of `slice_time` and `slice_index`',
+    ):
+      geometry_pydantic_model.IMASConfig(
+          imas_filepath=filename, slice_index=1, slice_time=1
+      )
+
+  def test_IMAS_raises_if_slice_out_of_range(self):
+    filename = 'ITERhybrid_COCOS17_IDS_ddv4.nc'
+    with self.assertRaisesRegex(
+        IndexError,
+        'out of range',
+    ):
+      config = geometry_pydantic_model.IMASConfig(
+          imas_filepath=filename,
+          slice_index=100,
+      )
+      config.build_geometry()
+
 
 if __name__ == '__main__':
   absltest.main()