Fix nrho -> n_rho and hi_res_fac -> hires_fac.

PiperOrigin-RevId: 726074342
google-deepmind · Feb 12, 2025 · b392227 · b392227
1 parent 1ba4969
commit b392227
Show file tree

Hide file tree

Showing 6 changed files with 10 additions and 8 deletions.
diff --git a/docs/configuration.rst b/docs/configuration.rst
@@ -539,10 +539,10 @@ geometry
 
 Geometry dicts for all geometry types can contain the following additional keys.
 
-``nrho`` (int = 25)
+``n_rho`` (int = 25)
   Number of radial grid points
 
-``hi_res_fac`` (int = 4)
+``hires_fac`` (int = 4)
   Only used when the initial condition ``psi`` is from plasma current. Sets up a higher resolution mesh
   with ``nrho_hires = nrho * hi_res_fac``, used for ``j`` to ``psi`` conversions.
 

diff --git a/torax/geometry/circular_geometry.py b/torax/geometry/circular_geometry.py
@@ -219,7 +219,7 @@ def build_circular_geometry(
 
   return CircularAnalyticalGeometry(
       # Set the standard geometry params.
-      geometry_type=geometry.GeometryType.CIRCULAR.value,
+      geometry_type=geometry.GeometryType.CIRCULAR,
       torax_mesh=mesh,
       Phi=Phi,
       Phi_face=Phi_face,

diff --git a/torax/geometry/geometry.py b/torax/geometry/geometry.py
@@ -96,7 +96,7 @@ def face_to_cell(face: chex.Array) -> chex.Array:
 
 
 @enum.unique
-class GeometryType(enum.Enum):
+class GeometryType(enum.IntEnum):
   """Integer enum for geometry type.
 
   This type can be used within JAX expressions to access the geometry type
@@ -124,7 +124,7 @@ class Geometry:
   """
 
   # TODO(b/356356966): extend documentation to define what each attribute is.
-  geometry_type: int
+  geometry_type: GeometryType
   torax_mesh: Grid1D
   Phi: chex.Array
   Phi_face: chex.Array

diff --git a/torax/geometry/geometry_provider.py b/torax/geometry/geometry_provider.py
@@ -111,7 +111,7 @@ def torax_mesh(self) -> geometry.Grid1D:
 class TimeDependentGeometryProvider:
   """A geometry provider which holds values to interpolate based on time."""
 
-  geometry_type: int
+  geometry_type: geometry.GeometryType
   torax_mesh: geometry.Grid1D
   drho_norm: interpolated_param.InterpolatedVarSingleAxis
   Phi: interpolated_param.InterpolatedVarSingleAxis

diff --git a/torax/geometry/standard_geometry.py b/torax/geometry/standard_geometry.py
@@ -1007,7 +1007,7 @@ def build_standard_geometry(
   area = rhon_interpolation_func(rho_norm, area_intermediate)
 
   return StandardGeometry(
-      geometry_type=intermediate.geometry_type.value,
+      geometry_type=intermediate.geometry_type,
       torax_mesh=mesh,
       Phi=Phi,
       Phi_face=Phi_face,

diff --git a/torax/geometry/tests/geometry_test.py b/torax/geometry/tests/geometry_test.py
@@ -130,7 +130,9 @@ def test_stack_geometries_error_handling_different_mesh_sizes(self):
   def test_stack_geometries_error_handling_different_geometry_types(self):
     """Test different geometry type error handling for stack_geometries."""
     geo0 = circular_geometry.build_circular_geometry(Rmaj=1.0, B0=2.0, n_rho=10)
-    geo_diff_geometry_type = dataclasses.replace(geo0, geometry_type=3)
+    geo_diff_geometry_type = dataclasses.replace(
+        geo0, geometry_type=geometry.GeometryType(3)
+    )
     with self.assertRaisesRegex(
         ValueError, 'All geometries must have the same geometry type'
     ):