Add local poloidal beta profile and beta_pol_prime derivative calculations

torax/_src/output_tools/output_keys.py

@@ -870,6 +870,14 @@ def T_fast_ion_key(source_key: str) -> str:  # pylint: disable=invalid-name
 BETA_N: Final[OutputKey] = OutputKey(
     "beta_N", units=Units.DIMENSIONLESS, grid_type=GridType.SCALAR
 )
+BETA_POL_PROFILE: Final[OutputKey] = OutputKey(
+    "beta_pol_profile",
+    units=Units.DIMENSIONLESS,
+    grid_type=GridType.FACE,
+)
+BETA_POL_PRIME: Final[OutputKey] = OutputKey(
+    "beta_pol_prime", units=Units.DIMENSIONLESS, grid_type=GridType.FACE
+)
 
 # ---------------------------------------------------------------------------
 # Edge model outputs.

torax/_src/output_tools/post_processing.py

@@ -200,6 +200,11 @@ class PostProcessedOutputs:
     beta_tor: Volume-averaged toroidal plasma beta (thermal) [dimensionless]
     beta_pol: Volume-averaged poloidal plasma beta (thermal) [dimensionless]
     beta_N: Normalized toroidal plasma beta (thermal) [dimensionless].
+    beta_pol_profile: Local poloidal beta profile on the face grid
+      [dimensionless]
+    beta_pol_prime: Derivative of local poloidal beta with respect to normalized
+      poloidal flux on the face grid: -d(beta_pol_local) / d(psi_norm)
+      [dimensionless]
     impurity_species: Dictionary of outputs for each impurity species.
     poloidal_velocity: Poloidal velocity [m/s]
     radial_electric_field: Radial electric field [V/m]
@@ -315,6 +320,8 @@ class PostProcessedOutputs:
   beta_tor: array_typing.FloatScalar
   beta_pol: array_typing.FloatScalar
   beta_N: array_typing.FloatScalar
+  beta_pol_profile: array_typing.FloatVector
+  beta_pol_prime: array_typing.FloatVector
   S_total: array_typing.FloatScalar
   impurity_species: dict[str, impurity_radiation.ImpuritySpeciesOutput]
   poloidal_velocity: array_typing.FloatVector
@@ -428,6 +435,8 @@ def zeros(cls, geo: geometry.Geometry) -> typing_extensions.Self:
         beta_tor=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         beta_pol=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         beta_N=jnp.array(0.0, dtype=jax_utils.get_dtype()),
+        beta_pol_profile=jnp.zeros(geo.rho_face.shape),
+        beta_pol_prime=jnp.zeros(geo.rho_face.shape),
         S_total=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         impurity_species={},
         poloidal_velocity=jnp.zeros(geo.rho_face.shape),
@@ -927,6 +936,12 @@ def cumulative_values():
   beta_tor, beta_pol, beta_N = formulas.calculate_betas(  # pyrefly: ignore[not-iterable]
       sim_state.core_profiles, sim_state.geometry
   )
+  beta_pol_profile = formulas.calculate_beta_pol_profile(
+      sim_state.core_profiles, sim_state.geometry
+  )
+  beta_pol_prime = formulas.calculate_beta_pol_prime(
+      sim_state.core_profiles, sim_state.geometry
+  )
 
   rotation_output = rotation.calculate_rotation(
       T_i=sim_state.core_profiles.T_i,
@@ -1019,6 +1034,8 @@ def cumulative_values():
       beta_tor=beta_tor,
       beta_pol=beta_pol,
       beta_N=beta_N,
+      beta_pol_profile=beta_pol_profile.face_value(),
+      beta_pol_prime=beta_pol_prime,
       impurity_species=impurity_radiation_outputs,
       poloidal_velocity=rotation_output.poloidal_velocity.face_value(),  # pyrefly: ignore[bad-argument-type]
       radial_electric_field=rotation_output.Er.face_value(),  # pyrefly: ignore[bad-argument-type]

torax/_src/physics/formulas.py

@@ -30,6 +30,10 @@
       averaged electron density (can be line-averaged or volume-averaged).
     - calculate_beta_volume_avg: Calculates the volume-averaged plasma beta
       based on thermal pressure.
+    - calculate_beta_pol_profile: Calculates local poloidal beta profile as a
+      CellVariable.
+    - calc_beta_pol_prime: Calculates
+      beta_pol_prime = -d(beta_pol) / d(psi_norm) on the face grid.
 """
 from jax import numpy as jnp
 from torax._src import array_typing
@@ -38,6 +42,7 @@
 from torax._src import state
 from torax._src.fvm import cell_variable
 from torax._src.geometry import geometry
+from torax._src.physics import psi_calculations
 
 
 # pylint: disable=invalid-name
@@ -264,3 +269,70 @@ def calculate_betas(
   )
 
   return beta_tor, beta_pol, beta_N  # pyrefly: ignore[bad-return]
+
+
+def calculate_beta_pol_profile(
+    core_profiles: state.CoreProfiles,
+    geo: geometry.Geometry,
+) -> cell_variable.CellVariable:
+  """Calculates the local poloidal beta profile on the cell grid.
+
+  beta_pol_local(psi) = P_total(psi) / (<Bp^2(psi)> / (2 * mu0))
+
+  Args:
+    core_profiles: CoreProfiles object.
+    geo: Geometry object.
+
+  Returns:
+    beta_pol_profile: CellVariable of local poloidal beta profile.
+  """
+  bpol2_face = psi_calculations.calc_bpol_squared(geo, core_profiles.psi)
+  bpol2_cell = geometry.face_to_cell(bpol2_face)
+  denom_cell = (
+      bpol2_cell / (2.0 * constants.CONSTANTS.mu_0) + constants.CONSTANTS.eps
+  )
+  denom_right = (
+      bpol2_face[-1] / (2.0 * constants.CONSTANTS.mu_0)
+      + constants.CONSTANTS.eps
+  )
+  right_face_constraint = (
+      core_profiles.pressure_total.right_face_constraint / denom_right
+      if core_profiles.pressure_total.right_face_constraint is not None
+      else None
+  )
+  return cell_variable.CellVariable(
+      value=core_profiles.pressure_total.value / denom_cell,
+      face_centers=core_profiles.pressure_total.face_centers,
+      right_face_constraint=right_face_constraint,
+      right_face_grad_constraint=None,
+  )
+
+
+def calculate_beta_pol_prime(
+    core_profiles: state.CoreProfiles,
+    geo: geometry.Geometry,
+) -> array_typing.FloatVectorFace:
+  r"""Calculates beta_pol_prime on the face grid.
+
+  Defined as:
+    beta_pol_prime = -d(beta_pol_local) / d(psi_norm)
+  where beta_pol_local is the local poloidal beta CellVariable and psi_norm is
+  the normalized poloidal flux in [0, 1]. In normal confinement, pressure
+  decreases toward the edge, making d(beta_pol)/d(psi_norm) negative, so
+  beta_pol_prime represents the positive gradient magnitude.
+
+  Args:
+    core_profiles: CoreProfiles object.
+    geo: Geometry object.
+
+  Returns:
+    beta_pol_prime: Face-grid array of the derivative magnitude [dimensionless].
+  """
+  beta_pol = calculate_beta_pol_profile(core_profiles, geo)
+  return -calc_dvar_dpsi(
+      var=beta_pol,
+      psi=core_profiles.psi,
+      normalized=True,
+  )
+
+

torax/_src/physics/tests/formulas_test.py

@@ -16,11 +16,14 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
+from torax._src import constants
 from torax._src import jax_utils
 from torax._src import math_utils
 from torax._src.fvm import cell_variable
 from torax._src.geometry import circular_geometry
+from torax._src.geometry import geometry
 from torax._src.physics import formulas
+from torax._src.physics import psi_calculations
 from torax._src.test_utils import core_profile_helpers
 
 
@@ -154,6 +157,84 @@ def test_calc_dvar_dpsi_against_analytical_solution(self):
           dvar_dpsi_norm, expected_dvar_dpsi_norm, rtol=1e-6
       )
 
+  def test_calculate_beta_pol_profile_and_derivative_analytical_circular(self):
+    geo = circular_geometry.CircularConfig(
+        n_rho=50, a_minor=1.0, R_major=10.0, B_0=2.0
+    ).build_geometry()
+    Ip = 1e6
+    mu0 = constants.CONSTANTS.mu_0
+    R0 = geo.R_major
+
+    psi_face = 0.5 * mu0 * Ip * R0 * (geo.rho_face / geo.rho_face[-1]) ** 2
+    psi_cell = 0.5 * mu0 * Ip * R0 * (geo.rho / geo.rho_face[-1]) ** 2
+    psi_var = cell_variable.CellVariable(
+        value=psi_cell,
+        face_centers=geo.rho_face,
+        left_face_constraint=psi_face[0],
+        right_face_constraint=psi_face[-1],
+        left_face_grad_constraint=None,
+        right_face_grad_constraint=None,
+    )
+
+    # Compute magnetic pressure denominator: <Bp^2> / (2 * mu0)
+    bpol2_face = psi_calculations.calc_bpol_squared(geo, psi_var)
+    bpol2_cell = geometry.face_to_cell(bpol2_face)
+    denom_cell = bpol2_cell / (2.0 * mu0)
+    denom_right = bpol2_face[-1] / (2.0 * mu0)
+
+    # Define a linear profile in normalized psi space:
+    # beta_pol_target = beta_0 + beta_prime * (1 - psi_N)
+    # Then d(beta_pol) / d(psi_N) = -beta_prime everywhere.
+    psi_range = psi_var.right_face_value - psi_var.left_face_value
+    psi_N_cell = (psi_var.value - psi_var.left_face_value) / psi_range
+    beta_0 = 0.5
+    beta_prime = 1.2
+    target_beta_pol_cell = beta_0 + beta_prime * (1.0 - psi_N_cell)
+
+    P_tot_cell = target_beta_pol_cell * denom_cell
+    P_tot_right = beta_0 * denom_right
+
+    # Construct T_e and n_e so that pressure_total equals P_tot:
+    # With n_i = 0, n_impurity = 0, fast_ions = ():
+    # P_tot = P_el = n_e * T_e * keV_to_J.
+    # Set n_e = 1.0 / keV_to_J, and T_e = P_tot.
+    n_e_val = cell_variable.CellVariable(
+        value=np.ones_like(geo.rho) / constants.CONSTANTS.keV_to_J,
+        face_centers=geo.rho_face,
+        right_face_constraint=1.0 / constants.CONSTANTS.keV_to_J,
+        right_face_grad_constraint=None,
+    )
+    T_e_val = cell_variable.CellVariable(
+        value=P_tot_cell,
+        face_centers=geo.rho_face,
+        left_face_constraint=None,
+        left_face_grad_constraint=0.0,
+        right_face_constraint=P_tot_right,
+        right_face_grad_constraint=None,
+    )
+
+    core_profiles = core_profile_helpers.make_zero_core_profiles(geo)
+    core_profiles = dataclasses.replace(
+        core_profiles,
+        psi=psi_var,
+        n_e=n_e_val,
+        T_e=T_e_val,
+    )
+
+    with self.subTest('beta_pol_profile'):
+      beta_pol_prof = formulas.calculate_beta_pol_profile(core_profiles, geo)
+      np.testing.assert_allclose(
+          beta_pol_prof.value, target_beta_pol_cell, rtol=1e-6
+      )
+      self.assertIsNotNone(beta_pol_prof.right_face_constraint)
+      np.testing.assert_allclose(
+          beta_pol_prof.right_face_constraint, beta_0, rtol=1e-6
+      )
+
+    with self.subTest('beta_pol_prime'):
+      beta_pol_prime = formulas.calculate_beta_pol_prime(core_profiles, geo)
+      # Check inner faces (faces 1 through N-1) and right face
+      np.testing.assert_allclose(beta_pol_prime[1:], beta_prime, rtol=1e-4)
 
 if __name__ == '__main__':
   absltest.main()