Add slice_time and slice_index arguments to IMAS geometry loader

Author: theo-brown

torax/_src/geometry/imas.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Useful functions for handling of IMAS IDSs."""
+import logging
 import os
 from typing import Any
 
@@ -29,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,
@@ -45,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 (default 0).
     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
@@ -71,9 +77,18 @@ 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]
+  if slice_time is not None:
+    if slice_time in equilibrium.time:
+      slice_index = np.nonzero(equilibrium.time == slice_time)
+    else:
+      slice_index = np.searchsorted(slice_time, equilibrium.time)
+      logging.warning(
+          f"t={slice_time} not in equilibrium.time."
+          f"Using t={equilibrium.time[slice_index]} instead."
+      )
+  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/pydantic_model.py

@@ -290,6 +290,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 (default 0).
   """
 
   geometry_type: Annotated[Literal['imas'], TIME_INVARIANT] = 'imas'
@@ -300,6 +303,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: int = 0
+  slice_time: float | None = None
 
   @pydantic.model_validator(mode='after')
   def _validate_model(self) -> typing_extensions.Self:
@@ -318,6 +323,11 @@ def _validate_model(self) -> typing_extensions.Self:
           '`imas_uri` or `imas_filepath` to be a valid input. Inputs provided: '
           f'{specified_inputs}.'
       )
+    if self.slice_time is not None and self.slice_index != 0:
+      raise ValueError(
+          'IMAS geometry needs exactly one of `slice_time` and `slice_index`.'
+          f'Got: slice_time={self.slice_time}, slice_index={self.slice_index}.'
+      )
     return self
 
   def build_geometry(self) -> standard_geometry.StandardGeometry:

torax/_src/geometry/standard_geometry.py

@@ -942,6 +942,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,
@@ -956,6 +958,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 (default 0).
       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
@@ -973,6 +978,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)