From 17756647c10c8e86336d84efe7eb836d878494db Mon Sep 17 00:00:00 2001
From: Sebastian Bodenstein <sbodenstein@google.com>
Date: Wed, 12 Feb 2025 08:55:21 -0800
Subject: [PATCH] Fix nrho -> n_rho in the docs. Remove hi_res_fac in docs.

PiperOrigin-RevId: 726074342
---
 docs/configuration.rst                | 6 +-----
 torax/geometry/circular_geometry.py   | 2 +-
 torax/geometry/geometry.py            | 4 ++--
 torax/geometry/geometry_provider.py   | 2 +-
 torax/geometry/standard_geometry.py   | 2 +-
 torax/geometry/tests/geometry_test.py | 4 +++-
 6 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/docs/configuration.rst b/docs/configuration.rst
index 87c7f182..5ee2353d 100644
--- a/docs/configuration.rst
+++ b/docs/configuration.rst
@@ -539,13 +539,9 @@ 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)
-  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.
-
 Geometry dicts for all non-circular geometry types can contain the following additional keys.
 
 ``geometry_file`` (str = 'ITER_hybrid_citrin_equil_cheasedata.mat2cols')
diff --git a/torax/geometry/circular_geometry.py b/torax/geometry/circular_geometry.py
index 85081340..63197d6f 100644
--- 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
index 23a342f2..0c0d6668 100644
--- 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
index 58e5baaf..e32ecd9d 100644
--- 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
index 39b3536b..503c1dde 100644
--- 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
index de46732f..30b9c09d 100644
--- 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'
     ):