Respond to Nush's comments

Author: theo-brown

torax/_src/geometry/pydantic_model.py

@@ -305,7 +305,7 @@ class IMASConfig(torax_pydantic.BaseModelFrozen):
   imas_uri: str | None = None
   equilibrium_object: ids_toplevel.IDSToplevel | None = None
   slice_index: pydantic.NonNegativeInt = 0
-  slice_time: float | None = None
+  slice_time: pydantic.NonNegativeFloat | None = None
 
   @pydantic.model_validator(mode='after')
   def _validate_model(self) -> typing_extensions.Self:
@@ -324,11 +324,6 @@ 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/imas_tools/input/equilibrium.py

@@ -85,13 +85,13 @@ def geometry_from_IMAS(
     slice_index = (
         np.searchsorted(equilibrium.time, slice_time, side="right") - 1
     )
-    if not np.allclose(equilibrium.time[slice_index], slice_time, atol=1e-3):
+    if not np.allclose(equilibrium.time[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):
+  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"

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

@@ -100,46 +100,53 @@ def test_IMAS_input_with_equilibrium_object(self):
 
   def test_IMAS_loading_specific_slice(self):
     filename = 'ITERhybrid_rampup_11_time_slices_COCOS17_IDS_ddv4.nc'
-    SLICE_TIME = 40
-    SLICE_INDEX = 5
     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
+        imas_filepath=filename, slice_time=40
     )
     config_at_slice_from_index = geometry_pydantic_model.IMASConfig(
-        imas_filepath=filename, slice_index=SLICE_INDEX
+        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()
 
-    for key in geo_at_0.__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(
-            geo_at_slice_from_time.__dict__[key],
-            geo_at_slice_from_index.__dict__[key],
-            err_msg=f'Value {key} mismatched between slice_time and slice_index',
-        )
-        with self.assertRaises(AssertionError):
-          np.testing.assert_allclose(
-              geo_at_0.__dict__[key],
-              geo_at_slice_from_time.__dict__[key],
-          )
-        with self.assertRaises(AssertionError):
+    # Check that the geometry is not the same as at t=0
+    with self.assertRaises(AssertionError):
+      for key in geo_at_slice_from_time.__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(
               geo_at_0.__dict__[key],
               geo_at_slice_from_index.__dict__[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
-      )
+    # Check that the geometry at index and time are the same
+    for key in geo_at_slice_from_time.__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(
+            geo_at_slice_from_time.__dict__[key],
+            geo_at_slice_from_index.__dict__[key],
+            err_msg=(
+                f'Value {key} mismatched between slice_time and slice_index'
+            ),
+        )
 
   def test_IMAS_raises_if_slice_out_of_range(self):
     filename = 'ITERhybrid_COCOS17_IDS_ddv4.nc'

torax/data/third_party/geo/README

@@ -29,9 +29,8 @@ All associated LICENSE files are also located in this directory.
 ITER_hybrid_citrin_equil_cheasedata.mat2cols: LICENSE_PINT
 eqdsk_cocos02.eqdsk: LICENSE_PINT
 
-The IMAS netCDF file ITERhybrid_rampup_11_time_slices_COCOS17_IDS_ddv4.nc was
-generated with CHEASE using p' and ff' from a TORAX simulation using
-iterhybrid_rampup.py example config. It contains equuilibrium for each of the
-time slices of the TORAX simulation. The plasma boundaries outline R and Z
-were prescribed constant for all time slices ITERhybrid_COCOS17_IDS_ddv4.nc,
-based on the default CHEASE geometry file.
+The IMAS netCDF file TERhybrid_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.