diff --git a/torax/_src/geometry/pydantic_model.py b/torax/_src/geometry/pydantic_model.py index a6821d723..e91db6245 100644 --- a/torax/_src/geometry/pydantic_model.py +++ b/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: diff --git a/torax/_src/geometry/standard_geometry.py b/torax/_src/geometry/standard_geometry.py index 0f3ec6e22..3e9ba918e 100644 --- a/torax/_src/geometry/standard_geometry.py +++ b/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) diff --git a/torax/_src/imas_tools/input/equilibrium.py b/torax/_src/imas_tools/input/equilibrium.py index e34eb6c32..80c965721 100644 --- a/torax/_src/imas_tools/input/equilibrium.py +++ b/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])) diff --git a/torax/_src/imas_tools/input/tests/equilibrium_test.py b/torax/_src/imas_tools/input/tests/equilibrium_test.py index 505ed7800..124e9ceb5 100644 --- a/torax/_src/imas_tools/input/tests/equilibrium_test.py +++ b/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() diff --git a/torax/data/third_party/geo/ITERhybrid_rampup_11_time_slices_COCOS17_IDS_ddv4.nc b/torax/data/third_party/geo/ITERhybrid_rampup_11_time_slices_COCOS17_IDS_ddv4.nc new file mode 100644 index 000000000..7e47628b7 Binary files /dev/null and b/torax/data/third_party/geo/ITERhybrid_rampup_11_time_slices_COCOS17_IDS_ddv4.nc differ diff --git a/torax/data/third_party/geo/README b/torax/data/third_party/geo/README index d14aeb25e..0d3359834 100644 --- a/torax/data/third_party/geo/README +++ b/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 \ No newline at end of file +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.