Add slice_time and slice_index arguments to IMAS geometry loader

torax/_src/geometry/pydantic_model.py

@@ -14,7 +14,8 @@
 
 """Pydantic model for geometry."""
 
-from collections.abc import Callable, Mapping
+from collections.abc import Callable
+from collections.abc import Mapping
 import functools
 import inspect
 import logging
@@ -296,6 +297,9 @@ class IMASConfig(torax_pydantic.BaseModelFrozen):
       the with running TORAX using the provided APIs. To use this option you
       must implement a custom run loop. Only one of imas_filepath, imas_uri or
       equilibrium_object can be set.
+    slice_time: Time of slice to load from IMAS IDS. If given, overrides
+      slice_index.
+    slice_index: Index of slice to load from IMAS IDS.
   """
 
   geometry_type: Annotated[Literal['imas'], TIME_INVARIANT] = 'imas'
@@ -306,6 +310,8 @@ class IMASConfig(torax_pydantic.BaseModelFrozen):
   imas_filepath: str | None = 'ITERhybrid_COCOS17_IDS_ddv4.nc'
   imas_uri: str | None = None
   equilibrium_object: ids_toplevel.IDSToplevel | None = None
+  slice_index: pydantic.NonNegativeInt = 0
+  slice_time: pydantic.NonNegativeFloat | None = None
 
   @pydantic.model_validator(mode='after')
   def _validate_model(self) -> typing_extensions.Self:

torax/_src/geometry/standard_geometry.py

@@ -1010,6 +1010,8 @@ def from_IMAS(
       Ip_from_parameters: bool,
       n_rho: int,
       hires_factor: int,
+      slice_time: float | None = None,
+      slice_index: int = 0,
       equilibrium_object: ids_toplevel.IDSToplevel | None = None,
       imas_uri: str | None = None,
       imas_filepath: str | None = None,
@@ -1024,6 +1026,9 @@ def from_IMAS(
         values in the Geometry are rescaled to match the new Ip.
       n_rho: Radial grid points (num cells).
       hires_factor: High resolution factor for calculations.
+      slice_time: Time of slice to load from IMAS IDS. If given, overrides
+        slice_index.
+      slice_index: Index of slice to load from IMAS IDS.
       equilibrium_object: The equilibrium IDS containing the relevant data.
       imas_uri: The IMAS uri containing the equilibrium data.
       imas_filepath: The path to the IMAS netCDF file containing the equilibrium
@@ -1041,6 +1046,8 @@ def from_IMAS(
         Ip_from_parameters=Ip_from_parameters,
         n_rho=n_rho,
         hires_factor=hires_factor,
+        slice_time=slice_time,
+        slice_index=slice_index,
     )
     return cls(geometry_type=geometry.GeometryType.IMAS, **inputs)
 

torax/_src/imas_tools/input/equilibrium.py

@@ -30,6 +30,8 @@ def geometry_from_IMAS(
     Ip_from_parameters: bool = False,
     n_rho: int = 25,
     hires_factor: int = 4,
+    slice_time: float | None = None,
+    slice_index: int = 0,
     equilibrium_object: ids_toplevel.IDSToplevel | None = None,
     imas_uri: str | None = None,
     imas_filepath: str | None = None,
@@ -46,6 +48,9 @@ def geometry_from_IMAS(
     n_rho: Radial grid points (num cells)
     hires_factor: Grid refinement factor for poloidal flux <--> plasma current
       calculations.
+    slice_time: Time of slice to load from IMAS IDS. If given, overrides
+      slice_index.
+    slice_index: Index of slice to load from IMAS IDS.
     equilibrium_object: The equilibrium IDS containing the relevant data.
     imas_uri: The IMAS uri containing the equilibrium data.
     imas_filepath: The path to the IMAS netCDF file containing the equilibrium
@@ -72,9 +77,36 @@ def geometry_from_IMAS(
     raise ValueError(
         "equilibrium_object must be a string (file path) or an IDS"
     )
-  # TODO(b/431977390): Currently only the first time slice is used, extend to
+  # TODO(b/431977390): Currently only a single time slice is used, extend to
   # support multiple time slices.
-  IMAS_data = equilibrium.time_slice[0]
+  # Convert time to index
+  if slice_time is not None:
+    if not np.all(equilibrium.time[:-1] <= equilibrium.time[1:]):
+      sorting_indices = np.argsort(equilibrium.time)
+    else:
+      sorting_indices = np.arange(len(equilibrium.time))
+    # Find the closest time in the IDS that is <= slice_time
+    slice_index = (
+        np.searchsorted(
+            equilibrium.time[sorting_indices], slice_time, side="right"
+        )
+        - 1
+    )
+    if not np.allclose(
+        equilibrium.time[sorting_indices][slice_index], slice_time, atol=1e-9
+    ):
+      logging.warn(
+          f"Requested t={slice_time} not in IDS; "
+          f"using t={equilibrium.time[slice_index]})"
+      )
+
+  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)
   B_0 = np.asarray(np.abs(equilibrium.vacuum_toroidal_field.b0[0]))
 

torax/_src/imas_tools/input/tests/equilibrium_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 pydantic_model as geometry_pydantic_model
 from torax._src.imas_tools.input import equilibrium as imas_geometry
@@ -97,6 +98,59 @@ def test_IMAS_input_with_equilibrium_object(self):
     )
     config.build_geometry()
 
+  def test_IMAS_loading_specific_slice(self):
+    def _check_geo_match(geo1, geo2):
+      for key in geo1.__dict__.keys():
+        if key not in [
+            'geometry_type',
+            'torax_mesh',
+            'R_major',
+            'B_0',
+            'rho_hires_norm',
+            'Phi_b_dot',
+            'Ip_from_parameters',
+        ]:
+          np.testing.assert_allclose(
+              geo1.__dict__[key],
+              geo2.__dict__[key],
+              err_msg=(
+                  f'Value {key} mismatched between slice_time and slice_index'
+              ),
+          )
+
+    filename = 'ITERhybrid_rampup_11_time_slices_COCOS17_IDS_ddv4.nc'
+    config_at_0 = geometry_pydantic_model.IMASConfig(imas_filepath=filename)
+    config_at_slice_from_time = geometry_pydantic_model.IMASConfig(
+        imas_filepath=filename, slice_time=40
+    )
+    config_at_slice_from_index = geometry_pydantic_model.IMASConfig(
+        imas_filepath=filename, slice_index=5
+    )
+
+    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()
+
+    _check_geo_match(geo_at_slice_from_time, geo_at_slice_from_index)
+
+    # Check that the geometry is not the same as at t=0
+    with self.assertRaisesRegex(
+        AssertionError, 'mismatched between slice_time and slice_index'
+    ):
+      _check_geo_match(geo_at_0, geo_at_slice_from_index)
+
+  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()

torax/data/third_party/geo/README

@@ -27,4 +27,10 @@ The mapping of datafile: LICENSE for the geometry files in this directory are as
 All associated LICENSE files are also located in this directory.
 
 ITER_hybrid_citrin_equil_cheasedata.mat2cols: LICENSE_PINT
-eqdsk_cocos02.eqdsk: LICENSE_PINT
+eqdsk_cocos02.eqdsk: LICENSE_PINT
+
+The IMAS netCDF file ITERhybrid_rampup_11_time_slices_COCOS17_IDS_ddv4.nc was generated
+by extracting p' and ff' profiles for each timeslice from a TORAX simulation with the
+iterhybrid_rampup.py config, and passing the profiles in to CHEASE.
+The LCFS was kept fixed for all timeslices, using the values from
+ITER_hybrid_citrin_equil_cheasedata.mat2cols.