Remove unnecessary face_centers and cell_centers force update on set_grid.

PiperOrigin-RevId: 852228987

Author: Nush395

torax/_src/core_profiles/convertors.py

@@ -133,9 +133,9 @@ def scale_cell_variable(
 
   return cell_variable.CellVariable(
       value=scaled_value,
+      face_centers=cv.face_centers,
       left_face_constraint=scaled_left_face_constraint,
       left_face_grad_constraint=scaled_left_face_grad_constraint,
       right_face_constraint=scaled_right_face_constraint,
       right_face_grad_constraint=scaled_right_face_grad_constraint,
-      dr=cv.dr,
   )

torax/_src/core_profiles/getters.py

@@ -72,10 +72,10 @@ def get_updated_ion_temperature(
     value = profile_conditions_params.T_i
   T_i = cell_variable.CellVariable(
       value=value,
+      face_centers=geo.rho_face_norm,
       left_face_grad_constraint=jnp.zeros(()),
       right_face_grad_constraint=None,
       right_face_constraint=profile_conditions_params.T_i_right_bc,
-      dr=geo.drho_norm,
   )
   return T_i
 
@@ -99,10 +99,10 @@ def get_updated_electron_temperature(
 
   T_e = cell_variable.CellVariable(
       value=value,
+      face_centers=geo.rho_face_norm,
       left_face_grad_constraint=jnp.zeros(()),
       right_face_grad_constraint=None,
       right_face_constraint=profile_conditions_params.T_e_right_bc,
-      dr=geo.drho_norm,
   )
   return T_e
 
@@ -191,7 +191,7 @@ def get_updated_electron_density(
 
   n_e = cell_variable.CellVariable(
       value=value,
-      dr=geo.drho_norm,
+      face_centers=geo.rho_face_norm,
       right_face_grad_constraint=None,
       right_face_constraint=n_e_right_bc,
   )
@@ -233,7 +233,7 @@ def get_updated_psi(
   )
   return cell_variable.CellVariable(
       value=value,
-      dr=geo.drho_norm,
+      face_centers=geo.rho_face_norm,
       right_face_grad_constraint=right_face_grad_constraint,
       right_face_constraint=right_face_constraint,
   )
@@ -250,9 +250,9 @@ def get_updated_toroidal_velocity(
     value = profile_conditions_params.toroidal_velocity
   toroidal_velocity = cell_variable.CellVariable(
       value=value,
+      face_centers=geo.rho_face_norm,
       right_face_grad_constraint=None,
       right_face_constraint=profile_conditions_params.toroidal_velocity_right_bc,
-      dr=geo.drho_norm,
   )
   return toroidal_velocity
 
@@ -655,7 +655,7 @@ def get_updated_ions(
 
   n_i = cell_variable.CellVariable(
       value=n_e.value * ion_properties.dilution_factor,
-      dr=geo.drho_norm,
+      face_centers=geo.rho_face_norm,
       right_face_grad_constraint=None,
       right_face_constraint=n_e.right_face_constraint
       * ion_properties.dilution_factor_edge,
@@ -676,7 +676,7 @@ def get_updated_ions(
 
   n_impurity = cell_variable.CellVariable(
       value=n_impurity_value,
-      dr=geo.drho_norm,
+      face_centers=geo.rho_face_norm,
       right_face_grad_constraint=None,
       right_face_constraint=n_impurity_right_face_constraint,
   )

torax/_src/core_profiles/initialization.py

@@ -90,11 +90,11 @@ def initial_core_profiles(
   # Initialise psi and derived quantities to zero before they are calculated.
   psidot = cell_variable.CellVariable(
       value=jnp.zeros_like(geo.rho, dtype=jax_utils.get_dtype()),
-      dr=geo.drho_norm,
+      face_centers=geo.rho_face_norm,
   )
   psi = cell_variable.CellVariable(
       value=jnp.zeros_like(geo.rho, dtype=jax_utils.get_dtype()),
-      dr=geo.drho_norm,
+      face_centers=geo.rho_face_norm,
   )
 
   core_profiles = state.CoreProfiles(
@@ -190,7 +190,7 @@ def update_psi_from_j(
     # we set it to match the desired plasma current.
     right_face_grad_constraint = None
     right_face_constraint = (
-        psi_value[-1] + dpsi_drhonorm_edge * geo.drho_norm / 2
+        psi_value[-1] + dpsi_drhonorm_edge * geo.drho_norm[-1] / 2
     )
   else:
     # Use the dpsi/drho calculated above as the right face gradient constraint
@@ -199,7 +199,7 @@ def update_psi_from_j(
 
   psi = cell_variable.CellVariable(
       value=psi_value,
-      dr=geo.drho_norm,
+      face_centers=geo.rho_face_norm,
       right_face_grad_constraint=right_face_grad_constraint,
       right_face_constraint=right_face_constraint,
   )
@@ -303,7 +303,7 @@ def _init_psi_and_psi_derived(
         right_face_grad_constraint = None
         right_face_constraint = (
             runtime_params.profile_conditions.psi[-1]
-            + dpsi_drhonorm_edge * geo.drho_norm / 2
+            + dpsi_drhonorm_edge * geo.drho_norm[-1] / 2
         )
       else:
         # Use the dpsi/drho calculated above as the right face gradient
@@ -313,9 +313,9 @@ def _init_psi_and_psi_derived(
 
       psi = cell_variable.CellVariable(
           value=runtime_params.profile_conditions.psi,
+          face_centers=geo.rho_face_norm,
           right_face_grad_constraint=right_face_grad_constraint,
           right_face_constraint=right_face_constraint,
-          dr=geo.drho_norm,
       )
 
     # Case 2: retrieving psi from the standard geometry input.
@@ -338,13 +338,13 @@ def _init_psi_and_psi_derived(
       # by make_ip_consistent
       psi = cell_variable.CellVariable(
           value=geo.psi_from_Ip,  # Use psi from equilibrium
+          face_centers=geo.rho_face_norm,
           right_face_grad_constraint=None
           if runtime_params.profile_conditions.use_v_loop_lcfs_boundary_condition
           else dpsi_drhonorm_edge,
           right_face_constraint=geo.psi_from_Ip_face[-1]
           if runtime_params.profile_conditions.use_v_loop_lcfs_boundary_condition
           else None,
-          dr=geo.drho_norm,
       )
 
     # Case 3: calculating j according to nu formula and psi from j.

torax/_src/core_profiles/tests/convertors_test.py

@@ -34,25 +34,25 @@ def setUp(self):
 
     T_i = cell_variable.CellVariable(
         value=jnp.ones(self.geo.rho_norm.shape) * 1.0,
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
         right_face_constraint=jnp.array(0.5),
         right_face_grad_constraint=None,
     )
     T_e = cell_variable.CellVariable(
         value=jnp.ones(self.geo.rho_norm.shape) * 2.0,
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
         right_face_constraint=jnp.array(0.6),
         right_face_grad_constraint=None,
     )
     psi = cell_variable.CellVariable(
         value=jnp.ones(self.geo.rho_norm.shape) * 3.0,
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
         right_face_grad_constraint=jnp.array(0.7),
         right_face_constraint=None,
     )
     n_e = cell_variable.CellVariable(
         value=jnp.ones(self.geo.rho_norm.shape) * 4.0,
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
         right_face_constraint=jnp.array(0.8),
         right_face_grad_constraint=None,
     )

torax/_src/core_profiles/tests/getters_test.py

@@ -73,7 +73,7 @@ def test_only_updating_boundary_condition(self):
         left_face_grad_constraint=jnp.zeros(()),
         right_face_grad_constraint=None,
         right_face_constraint=jnp.array(1.0, dtype=jax_utils.get_dtype()),
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
     )
     result = getters.get_updated_ion_temperature(
         profile_conditions,
@@ -113,7 +113,7 @@ def test_only_updating_boundary_condition_electron_temperature(self):
         left_face_grad_constraint=jnp.zeros(()),
         right_face_grad_constraint=None,
         right_face_constraint=jnp.array(1.0, dtype=jax_utils.get_dtype()),
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
     )
     result = getters.get_updated_electron_temperature(
         profile_conditions,
@@ -168,7 +168,7 @@ def test_only_updating_boundary_condition_n_e(self):
         left_face_grad_constraint=jnp.zeros(()),
         right_face_grad_constraint=None,
         right_face_constraint=jnp.array(1.0, dtype=jax_utils.get_dtype()),
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
     )
     n_e = getters.get_updated_electron_density(
         profile_conditions,
@@ -402,7 +402,7 @@ def test_get_updated_ion_data(self):
         left_face_grad_constraint=jnp.zeros(()),
         right_face_grad_constraint=None,
         right_face_constraint=jnp.array(100.0, dtype=jax_utils.get_dtype()),
-        dr=geo.drho_norm,
+        face_centers=geo.rho_face_norm,
     )
     n_e = getters.get_updated_electron_density(
         runtime_params.profile_conditions,
@@ -561,13 +561,13 @@ def test_get_updated_ions_impurity_mixture(self):
     geo = torax_config.geometry.build_provider(t=0.0)
     T_e_cell_variable = cell_variable.CellVariable(
         value=jnp.full_like(geo.rho_norm, T_e),
-        dr=geo.drho_norm,
+        face_centers=geo.rho_face_norm,
         right_face_constraint=T_e,
         right_face_grad_constraint=None,
     )
     n_e_cell_variable = cell_variable.CellVariable(
         value=jnp.full_like(geo.rho_norm, n_e),
-        dr=geo.drho_norm,
+        face_centers=geo.rho_face_norm,
         right_face_constraint=n_e,
         right_face_grad_constraint=None,
     )
@@ -666,13 +666,13 @@ def test_get_updated_ions_impurity_mixture_radially_dependent(self):
     geo = torax_config.geometry.build_provider(t=0.0)
     T_e_cell_variable = cell_variable.CellVariable(
         value=jnp.full_like(geo.rho_norm, T_e),
-        dr=geo.drho_norm,
+        face_centers=geo.rho_face_norm,
         right_face_constraint=T_e,
         right_face_grad_constraint=None,
     )
     n_e_cell_variable = cell_variable.CellVariable(
         value=jnp.full_like(geo.rho_norm, n_e),
-        dr=geo.drho_norm,
+        face_centers=geo.rho_face_norm,
         right_face_constraint=n_e,
         right_face_grad_constraint=None,
     )
@@ -790,13 +790,13 @@ def _run_get_updated_ions(torax_config):
 
       t_e_cell_variable = cell_variable.CellVariable(
           value=jnp.full_like(geo.rho_norm, t_e_keV),
-          dr=geo.drho_norm,
+          face_centers=geo.rho_face_norm,
           right_face_constraint=t_e_keV,
           right_face_grad_constraint=None,
       )
       n_e_cell_variable = cell_variable.CellVariable(
           value=jnp.full_like(geo.rho_norm, n_e_val),
-          dr=geo.drho_norm,
+          face_centers=geo.rho_face_norm,
           right_face_constraint=n_e_val,
           right_face_grad_constraint=None,
       )
@@ -908,13 +908,13 @@ def _run_get_updated_ions(torax_config):
 
       t_e_cell_variable = cell_variable.CellVariable(
           value=jnp.full_like(geo.rho_norm, t_e_keV),
-          dr=geo.drho_norm,
+          face_centers=geo.rho_face_norm,
           right_face_constraint=t_e_keV,
           right_face_grad_constraint=None,
       )
       n_e_cell_variable = cell_variable.CellVariable(
           value=jnp.full_like(geo.rho_norm, n_e_val),
-          dr=geo.drho_norm,
+          face_centers=geo.rho_face_norm,
           right_face_constraint=n_e_val,
           right_face_grad_constraint=None,
       )
@@ -978,13 +978,13 @@ def _run_get_updated_ions(torax_config):
 
       t_e_cell_variable = cell_variable.CellVariable(
           value=jnp.full_like(geo.rho_norm, t_e_keV),
-          dr=geo.drho_norm,
+          face_centers=geo.rho_face_norm,
           right_face_constraint=t_e_keV,
           right_face_grad_constraint=None,
       )
       n_e_cell_variable = cell_variable.CellVariable(
           value=jnp.full_like(geo.rho_norm, n_e_val),
-          dr=geo.drho_norm,
+          face_centers=geo.rho_face_norm,
           right_face_constraint=n_e_val,
           right_face_grad_constraint=None,
       )
@@ -1079,13 +1079,13 @@ def _run_get_updated_ions(torax_config):
 
       t_e_cell_variable = cell_variable.CellVariable(
           value=jnp.full_like(geo.rho_norm, t_e_keV),
-          dr=geo.drho_norm,
+          face_centers=geo.rho_face_norm,
           right_face_constraint=t_e_keV,
           right_face_grad_constraint=None,
       )
       n_e_cell_variable = cell_variable.CellVariable(
           value=jnp.full_like(geo.rho_norm, n_e_val),
-          dr=geo.drho_norm,
+          face_centers=geo.rho_face_norm,
           right_face_constraint=n_e_val,
           right_face_grad_constraint=None,
       )

torax/_src/core_profiles/tests/updaters_test.py

@@ -40,13 +40,13 @@ def setUp(self):
 
     T_e = cell_variable.CellVariable(
         value=jnp.ones_like(self.geo.rho_norm),
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
         right_face_constraint=1.0,
         right_face_grad_constraint=None,
     )
     n_e = cell_variable.CellVariable(
         value=jnp.ones_like(self.geo.rho_norm),
-        dr=self.geo.drho_norm,
+        face_centers=self.geo.rho_face_norm,
         right_face_constraint=1.0,
         right_face_grad_constraint=None,
     )

torax/_src/fvm/cell_variable.py

@@ -43,7 +43,8 @@ class CellVariable:
 
   Attributes:
     value: Value of this variable at each cell grid point.
-    dr: Distance between cell centers.
+    face_centers: Locations of the face centers. This array should have length
+      len(value) + 1. Supports both uniform and non-uniform grids.
     left_face_constraint: An optional scalar specifying the value of the
       leftmost face. Defaults to None, signifying no constraint. The user can
       modify this field at any time, but when face_grad is called exactly one of
@@ -56,8 +57,9 @@ class CellVariable:
     right_face_grad_constraint: Analogous to left_face_grad_constraint but for
       the right face, see left_face_grad_constraint.
   """
+
   value: jt.Float[chex.Array, 'cell']
-  dr: jt.Float[chex.Array, '']
+  face_centers: jt.Float[chex.Array, 'face']
   left_face_constraint: jt.Float[chex.Array, ''] | None = None
   right_face_constraint: jt.Float[chex.Array, ''] | None = None
   left_face_grad_constraint: jt.Float[chex.Array, ''] | None = (
@@ -69,11 +71,6 @@ class CellVariable:
   # Can't make the above default values be jax zeros because that would be a
   # call to jax before absl.app.run
 
-  @property
-  def face_centers(self) -> jt.Float[chex.Array, 'face']:
-    """Locations of the face centers."""
-    return jnp.linspace(0.0, len(self.value) * self.dr, num=len(self.value) + 1)
-
   @property
   def cell_centers(self) -> jt.Float[chex.Array, 'cell']:
     """Locations of the cell centers."""

torax/_src/fvm/convection_terms.py

@@ -118,10 +118,10 @@ def peclet_to_alpha(p):
   left_v = v_face[:-1]
   right_v = v_face[1:]
 
-  diag = (left_alpha * left_v - right_alpha * right_v) / var.dr
-  above = -(1.0 - right_alpha) * right_v / var.dr
+  diag = (left_alpha * left_v - right_alpha * right_v) / var.cell_widths
+  above = -(1.0 - right_alpha) * right_v / var.cell_widths
   above = above[:-1]
-  below = (1.0 - left_alpha) * left_v / var.dr
+  below = (1.0 - left_alpha) * left_v / var.cell_widths
   below = below[1:]
   mat = math_utils.tridiag(diag, above, below)