From a8dda40342cd5b90dade5076c81667c5659e2c03 Mon Sep 17 00:00:00 2001 From: Sebastian Bodenstein Date: Wed, 26 Feb 2025 11:00:54 -0800 Subject: [PATCH] Remove `from_chease` defaults. There should be a single source of truth for these defaults, which is the `CheaseConfig` Pydantic model. PiperOrigin-RevId: 731381598 --- torax/geometry/standard_geometry.py | 16 +++---- .../geometry/tests/standard_geometry_test.py | 3 +- torax/tests/state.py | 20 ++++----- torax/tests/test_lib/torax_refs.py | 42 +++++++++---------- 4 files changed, 37 insertions(+), 44 deletions(-) diff --git a/torax/geometry/standard_geometry.py b/torax/geometry/standard_geometry.py index a807b906..31cbf07e 100644 --- a/torax/geometry/standard_geometry.py +++ b/torax/geometry/standard_geometry.py @@ -234,14 +234,14 @@ def __post_init__(self): @classmethod def from_chease( cls, - geometry_dir: str | None = None, - geometry_file: str = 'ITER_hybrid_citrin_equil_cheasedata.mat2cols', - Ip_from_parameters: bool = True, - n_rho: int = 25, - Rmaj: float = 6.2, - Rmin: float = 2.0, - B0: float = 5.3, - hires_fac: int = 4, + geometry_dir: str | None, + geometry_file: str, + Ip_from_parameters: bool, + n_rho: int, + Rmaj: float, + Rmin: float, + B0: float, + hires_fac: int, ) -> StandardGeometryIntermediates: """Constructs a StandardGeometryIntermediates from a CHEASE file. diff --git a/torax/geometry/tests/standard_geometry_test.py b/torax/geometry/tests/standard_geometry_test.py index 0cdbde10..5bd309e7 100644 --- a/torax/geometry/tests/standard_geometry_test.py +++ b/torax/geometry/tests/standard_geometry_test.py @@ -80,8 +80,7 @@ def test_build_geometry_from_eqdsk(self, geometry_file): def test_access_z_magnetic_axis_raises_error_for_chease_geometry(self): """Test that accessing z_magnetic_axis raises error for CHEASE geometry.""" - intermediate = standard_geometry.StandardGeometryIntermediates.from_chease() - geo = standard_geometry.build_standard_geometry(intermediate) + geo = geometry_pydantic_model.CheaseConfig().build_geometry() with self.assertRaisesRegex(ValueError, 'does not have a z magnetic axis'): geo.z_magnetic_axis() diff --git a/torax/tests/state.py b/torax/tests/state.py index 8898e235..6c937d89 100644 --- a/torax/tests/state.py +++ b/torax/tests/state.py @@ -31,9 +31,8 @@ from torax.config import runtime_params as general_runtime_params from torax.config import runtime_params_slice from torax.geometry import circular_geometry -from torax.geometry import geometry from torax.geometry import geometry_provider -from torax.geometry import standard_geometry +from torax.geometry import pydantic_model as geometry_pydantic_model from torax.sources import generic_current_source from torax.sources import runtime_params as runtime_params_lib from torax.sources import source_models as source_models_lib @@ -236,16 +235,12 @@ def test_core_profiles_quasineutrality_check(self): assert not core_profiles.quasineutrality_satisfied() @parameterized.parameters([ - dict(geo_builder=circular_geometry.build_circular_geometry), - dict( - geo_builder=lambda: standard_geometry.build_standard_geometry( - standard_geometry.StandardGeometryIntermediates.from_chease() - ) - ), + dict(geometry_name='circular'), + dict(geometry_name='chease'), ]) def test_initial_psi_from_j( self, - geo_builder: Callable[[], geometry.Geometry], + geometry_name: str, ): """Tests expected behaviour of initial psi and current options.""" config1 = general_runtime_params.GeneralRuntimeParams( @@ -281,7 +276,9 @@ def test_initial_psi_from_j( ne_bound_right=0.5, ), ) - geo_provider = geometry_provider.ConstantGeometryProvider(geo_builder()) + geo_provider = geometry_pydantic_model.Geometry.from_dict( + {'geometry_type': geometry_name} + ).build_provider() source_models_builder = source_models_lib.SourceModelsBuilder() source_models = source_models_builder() source_models_builder.runtime_params['j_bootstrap'].bootstrap_mult = 0.0 @@ -373,7 +370,8 @@ def test_initial_psi_from_j( ctot = config1.profile_conditions.Ip_tot * 1e6 / denom jtot_formula = jformula * ctot johm_formula = jtot_formula * ( - 1 - dcs1.sources[ + 1 + - dcs1.sources[ generic_current_source.GenericCurrentSource.SOURCE_NAME ].fext # pytype: disable=attribute-error ) diff --git a/torax/tests/test_lib/torax_refs.py b/torax/tests/test_lib/torax_refs.py index e34ffef0..dc9b10f7 100644 --- a/torax/tests/test_lib/torax_refs.py +++ b/torax/tests/test_lib/torax_refs.py @@ -29,7 +29,7 @@ from torax.geometry import circular_geometry from torax.geometry import geometry from torax.geometry import geometry_provider as geometry_provider_lib -from torax.geometry import standard_geometry +from torax.geometry import pydantic_model as geometry_pydantic_model from torax.sources import runtime_params as sources_params from torax.stepper import runtime_params as stepper_params from torax.transport_model import runtime_params as transport_model_params @@ -237,17 +237,15 @@ def chease_references_Ip_from_chease() -> References: # pylint: disable=invalid }, }, ) - geo = standard_geometry.build_standard_geometry( - standard_geometry.StandardGeometryIntermediates.from_chease( - geometry_dir=_GEO_DIRECTORY, - geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', - n_rho=25, - Ip_from_parameters=False, - Rmaj=6.2, - Rmin=2.0, - B0=5.3, - ) - ) + geo = geometry_pydantic_model.CheaseConfig( + geometry_dir=_GEO_DIRECTORY, + geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', + n_rho=25, + Ip_from_parameters=False, + Rmaj=6.2, + Rmin=2.0, + B0=5.3, + ).build_geometry() # ground truth values copied from an example PINT execution using # array.astype(str),which allows fully lossless reloading psi = fvm.cell_variable.CellVariable( @@ -389,17 +387,15 @@ def chease_references_Ip_from_runtime_params() -> References: # pylint: disable }, }, ) - geo = standard_geometry.build_standard_geometry( - standard_geometry.StandardGeometryIntermediates.from_chease( - geometry_dir=_GEO_DIRECTORY, - geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', - n_rho=25, - Ip_from_parameters=True, - Rmaj=6.2, - Rmin=2.0, - B0=5.3, - ) - ) + geo = geometry_pydantic_model.CheaseConfig( + geometry_dir=_GEO_DIRECTORY, + geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', + n_rho=25, + Ip_from_parameters=True, + Rmaj=6.2, + Rmin=2.0, + B0=5.3, + ).build_geometry() # ground truth values copied from an example executions using # array.astype(str),which allows fully lossless reloading psi = fvm.cell_variable.CellVariable(