Remove from_chease defaults. There should be a single source of tru…

…th for these defaults, which is the `CheaseConfig` Pydantic model.

PiperOrigin-RevId: 731251551

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.
 

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()
 

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
     )

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(