Create helper for creating high resolution grid.

This slightly changes the behaviour of the existing hires grid to np.linspace(0,1,n_rho*hires_factor+1) so that inner faces remain aligned.
E.g [0, 0.25, 0.5, 0.75, 1.0] x2-> [0, 0.125, 0.25, 0.375, ..., 0.875, 1.0]

PiperOrigin-RevId: 859638056

Author: Nush395

torax/_src/geometry/circular_geometry.py

@@ -186,7 +186,9 @@ def _build_circular_geometry(
   # High resolution versions for j (plasma current) and psi (poloidal flux)
   # manipulations. Needed if psi is initialized from plasma current, which is
   # the only option for ad-hoc circular geometry.
-  rho_hires_norm = np.linspace(0, 1, n_rho * hires_factor)
+  rho_hires_norm = geometry.increase_grid_resolution(
+      rho_face_norm, hires_factor
+  )
   rho_hires = rho_hires_norm * rho_b
 
   R_out = R_major + rho

torax/_src/geometry/geometry.py

@@ -457,3 +457,15 @@ def update_geometries_with_Phibdot(
   geo_t = dataclasses.replace(geo_t, Phi_b_dot=Phibdot)
   geo_t_plus_dt = dataclasses.replace(geo_t_plus_dt, Phi_b_dot=Phibdot)
   return geo_t, geo_t_plus_dt
+
+
+def increase_grid_resolution(faces: chex.Array, factor: int) -> chex.Array:
+  """Increase the grid resolution by a factor."""
+  if factor <= 1:
+    raise ValueError('factor must be >= 1.')
+  num_faces = len(faces)
+  num_cells = num_faces - 1
+  grid_indices = np.arange(len(faces))
+  new_num_faces = num_cells * factor + 1
+  new_indices = np.linspace(0, num_cells, new_num_faces)
+  return np.interp(new_indices, grid_indices, faces)

torax/_src/geometry/standard_geometry.py

@@ -410,8 +410,8 @@ def build_standard_geometry(
 
   # High resolution versions for j (plasma current) and psi (poloidal flux)
   # manipulations. Needed if psi is initialized from plasma current.
-  rho_hires_norm = np.linspace(
-      0, 1, intermediates.n_rho * intermediates.hires_factor
+  rho_hires_norm = geometry.increase_grid_resolution(
+      rho_face_norm, intermediates.hires_factor
   )
   rho_hires = rho_hires_norm * rho_b
 

torax/_src/geometry/tests/geometry_test.py

@@ -169,6 +169,39 @@ def test_geometry_eq(self):
       geo3 = dataclasses.replace(geo1, Phi_face=np.array([2.0]))
       self.assertNotEqual(geo1, geo3)
 
+  @parameterized.parameters([
+      dict(
+          arr=np.array([1.0, 2.0, 3.0]),
+          factor=2,
+          expected_arr=np.array([1.0, 1.5, 2.0, 2.5, 3.0]),
+      ),
+      dict(
+          arr=np.array([1.0, 2.0, 3.0, 4.0, 5.0]),
+          factor=2,
+          expected_arr=np.array([1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0]),
+      ),
+      dict(
+          arr=np.array([1.0, 2.0, 3.0, 4.0, 5.0]),
+          factor=3,
+          expected_arr=np.array(
+              [1.0, 1.33333333, 1.66666667, 2.0, 2.33333333, 2.66666667, 3.0,
+               3.33333333, 3.66666667, 4.0, 4.33333333, 4.66666667, 5.0]
+          ),
+      ),
+      dict(
+          arr=np.array([0.0, 0.25, 0.5, 0.75, 1.0]),
+          factor=4,
+          expected_arr=np.array(
+              [0.0, 0.0625, 0.125, 0.1875, 0.25, 0.3125, 0.375, 0.4375, 0.5,
+               0.5625, 0.625, 0.6875, 0.75, 0.8125, 0.875, 0.9375, 1.0]
+          ),
+      ),
+  ])
+  def test_increase_resolution(self, arr, factor, expected_arr):
+    np.testing.assert_allclose(
+        geometry.increase_grid_resolution(arr, factor), expected_arr
+    )
+
 
 def _pint_face_to_cell(n_rho, face):
   cell = np.zeros(n_rho)

torax/_src/mhd/sawtooth/tests/sawtooth_model_test.py

@@ -241,42 +241,42 @@ def test_no_subsequent_sawtooth_crashes(self):
 
 
 _POST_CRASH_TEMPERATURE = np.array([
-    9.80214764,
-    9.77449557,
-    9.74682154,
-    9.71912539,
-    9.69140691,
-    8.17937075,
-    6.2258966,
+    9.798375,
+    9.770761,
+    9.743125,
+    9.715467,
+    9.687787,
+    8.183115,
+    6.229396,
     4.5,
     3.1,
     1.7,
 ])
 
 _POST_CRASH_N = np.array([
-    0.92905438e20,
-    0.92652621e20,
-    0.92399804e20,
-    0.92146987e20,
-    0.91894169e20,
-    0.88178024e20,
-    0.8345057e20,
-    0.79219014e20,
-    0.75698169e20,
-    0.72177324e20,
+    9.289564e+19,
+    9.264304e+19,
+    9.239043e+19,
+    9.213783e+19,
+    9.188523e+19,
+    8.818625e+19,
+    8.345811e+19,
+    7.921901e+19,
+    7.569817e+19,
+    7.217732e+19,
 ])
 
 _POST_CRASH_PSI = np.array([
-    8.882378,
-    10.483339,
-    13.574243,
-    18.056257,
-    23.83316,
-    30.671823,
-    37.89592,
-    44.739235,
-    50.713549,
-    55.729866,
+    8.877173,
+    10.473664,
+    13.561018,
+    18.040233,
+    23.815787,
+    30.664264,
+    37.900495,
+    44.750157,
+    50.721585,
+    55.732499,
 ])