Consolidate neoclassical evaluations into single-call NeoclassicalModels.__call__

- Introduce `NeoclassicalOutputs` and `NeoclassicalModels.__call__` to evaluate all four neoclassical sub-models (`conductivity`, `bootstrap_current`, `transport`, `poloidal_velocity`) in a single call.
- Pass the resulting precomputed outputs to source and transport coefficient builders instead of calling neoclassical sub-models separately.

PiperOrigin-RevId: 986345805

Author: jcitrin

torax/_src/core_profiles/initialization.py

@@ -141,6 +141,12 @@ def initial_core_profiles(
       j_total_face=jnp.zeros_like(geo.rho_face, dtype=jax_utils.get_dtype()),
       Ip_profile_face=jnp.zeros_like(geo.rho_face, dtype=jax_utils.get_dtype()),
       toroidal_angular_velocity=toroidal_angular_velocity,
+      poloidal_velocity=cell_variable.CellVariable(
+          value=jnp.zeros_like(geo.rho, dtype=jax_utils.get_dtype()),
+          face_centers=geo.rho_face_norm,
+          right_face_constraint=jnp.zeros((), dtype=jax_utils.get_dtype()),
+          right_face_grad_constraint=None,
+      ),
       charge_state_info=ions.charge_state_info,
       charge_state_info_face=ions.charge_state_info_face,
   )
@@ -481,10 +487,8 @@ def _calculate_all_psi_dependent_profiles(
       j_total_face=j_total_face,
       Ip_profile_face=Ip_profile_face,
   )
-  # Calculate conductivity once we have a consistent set of core profiles
-  conductivity = neoclassical_models.conductivity.calculate_conductivity(
-      geo,
-      core_profiles,
+  neoclassical_outputs = neoclassical_models(
+      runtime_params, geo, core_profiles
   )
 
   # Calculate sources if they have not already been calculated.
@@ -493,7 +497,7 @@ def _calculate_all_psi_dependent_profiles(
         runtime_params,
         geo,
         core_profiles,
-        neoclassical_models,
+        neoclassical_outputs.bootstrap_current,
         source_models,
         source_profiles,
     )
@@ -513,7 +517,7 @@ def _calculate_all_psi_dependent_profiles(
     psi_sources = source_profiles.total_psi_sources(geo)
     psidot_value = psi_calculations.calculate_psidot_from_psi_sources(
         psi_sources=psi_sources,
-        sigma=conductivity.sigma,
+        sigma=neoclassical_outputs.conductivity.sigma,
         resistivity_multiplier=runtime_params.numerics.resistivity_multiplier,  # pyrefly: ignore[bad-argument-type]
         psi=psi,
         geo=geo,
@@ -536,8 +540,9 @@ def _calculate_all_psi_dependent_profiles(
   core_profiles = dataclasses.replace(
       core_profiles,
       psidot=psidot,
-      sigma=conductivity.sigma,
-      sigma_face=conductivity.sigma_face,
+      sigma=neoclassical_outputs.conductivity.sigma,
+      sigma_face=neoclassical_outputs.conductivity.sigma_face,
+      poloidal_velocity=neoclassical_outputs.poloidal_velocity.v_pol,
   )
   return core_profiles
 
@@ -546,11 +551,11 @@ def _get_bootstrap_and_standard_source_profiles(
     runtime_params: runtime_params_lib.RuntimeParams,
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
-    neoclassical_models: neoclassical_models_lib.NeoclassicalModels,
+    bootstrap_current: bootstrap_current_base.BootstrapCurrent,
     source_models: source_models_lib.SourceModels,
     source_profiles: source_profiles_lib.SourceProfiles,
 ) -> source_profiles_lib.SourceProfiles:
-  """Calculates bootstrap current and updates source profiles."""
+  """Updates source profiles with standard psi sources and bootstrap current."""
   source_profile_builders.build_standard_source_profiles(
       runtime_params=runtime_params,
       geo=geo,
@@ -560,13 +565,9 @@ def _get_bootstrap_and_standard_source_profiles(
       calculate_anyway=True,
       calculated_source_profiles=source_profiles,
   )
-  bootstrap_current = (
-      neoclassical_models.bootstrap_current.calculate_bootstrap_current(
-          runtime_params, geo, core_profiles
-      )
-  )
   source_profiles = dataclasses.replace(
-      source_profiles, bootstrap_current=bootstrap_current
+      source_profiles,
+      bootstrap_current=bootstrap_current,
   )
   return source_profiles
 
@@ -587,7 +588,9 @@ def _iterate_psi_and_sources(
         runtime_params,
         geo,
         core_profiles,
-        neoclassical_models,
+        neoclassical_models(
+            runtime_params, geo, core_profiles
+        ).bootstrap_current,
         source_models,
         source_profiles,
     )

torax/_src/core_profiles/tests/convertors_test.py

@@ -85,6 +85,7 @@ def setUp(self):
         j_total_face=mock.ANY,
         Ip_profile_face=mock.ANY,
         toroidal_angular_velocity=mock.ANY,
+        poloidal_velocity=mock.ANY,
         charge_state_info=mock.ANY,
         charge_state_info_face=mock.ANY,
         fast_ions=mock.ANY,

torax/_src/core_profiles/tests/initialization_test.py

@@ -521,16 +521,14 @@ def _get_initial_state(
       source_models=source_models,
       neoclassical_models=neoclassical_models,
   )
-  conductivity = neoclassical_models.conductivity.calculate_conductivity(
-      geo, core_profiles
-  )
+  neoclassical_outputs = neoclassical_models(runtime_params, geo, core_profiles)
   core_sources = source_profile_builders.get_all_source_profiles(
       runtime_params=runtime_params,
       geo=geo,
       core_profiles=core_profiles,
       source_models=source_models,
-      neoclassical_models=neoclassical_models,
-      conductivity=conductivity,
+      conductivity=neoclassical_outputs.conductivity,
+      bootstrap_current=neoclassical_outputs.bootstrap_current,
   )
   j_toroidal_total = core_profiles.j_total
   j_toroidal_total_face = core_profiles.j_total_face

torax/_src/core_profiles/updaters.py

@@ -224,6 +224,9 @@ def update_core_and_source_profiles_after_step(
       j_total_face=j_total_face,
       Ip_profile_face=Ip_profile_face,
       toroidal_angular_velocity=updated_core_profiles_t_plus_dt.toroidal_angular_velocity,
+      poloidal_velocity=(
+          core_profiles_t_plus_dt.poloidal_velocity
+      ),  # Not yet updated
       charge_state_info=ions.charge_state_info,
       charge_state_info_face=ions.charge_state_info_face,
       fast_ions=core_profiles_t_plus_dt.fast_ions,
@@ -236,14 +239,15 @@ def update_core_and_source_profiles_after_step(
       dt,
   )
 
-  conductivity = neoclassical_models.conductivity.calculate_conductivity(
-      geo, intermediate_core_profiles
+  neoclassical_outputs = neoclassical_models(
+      runtime_params_t_plus_dt, geo, intermediate_core_profiles
   )
 
   intermediate_core_profiles = dataclasses.replace(
       intermediate_core_profiles,
-      sigma=conductivity.sigma,
-      sigma_face=conductivity.sigma_face,
+      sigma=neoclassical_outputs.conductivity.sigma,
+      sigma_face=neoclassical_outputs.conductivity.sigma_face,
+      poloidal_velocity=neoclassical_outputs.poloidal_velocity.v_pol,
       internal_plasma_energy=energy_state,
   )
 
@@ -252,11 +256,11 @@ def update_core_and_source_profiles_after_step(
       runtime_params=runtime_params_t_plus_dt,
       geo=geo,
       source_models=source_models,
-      neoclassical_models=neoclassical_models,
       core_profiles=intermediate_core_profiles,
       explicit=False,
       explicit_source_profiles=explicit_source_profiles,
-      conductivity=conductivity,
+      conductivity=neoclassical_outputs.conductivity,
+      bootstrap_current=neoclassical_outputs.bootstrap_current,
   )
 
   intermediate_core_profiles = dataclasses.replace(

torax/_src/fvm/calc_coeffs.py

@@ -231,22 +231,26 @@ def _calc_coeffs_full(
 
   consts = constants.CONSTANTS
 
-  conductivity = models.neoclassical_models.conductivity.calculate_conductivity(
-      geo, core_profiles
+  neoclassical_outputs = models.neoclassical_models(
+      runtime_params, geo, core_profiles
+  )
+  core_profiles = dataclasses.replace(
+      core_profiles,
+      poloidal_velocity=neoclassical_outputs.poloidal_velocity.v_pol,
   )
 
   # Calculate the implicit source profiles and combine them with the explicit
   # source profiles. These are needed for the pedestal model, so are computed
   # here rather than in the source terms section.
   merged_source_profiles = source_profile_builders.build_source_profiles(
       source_models=models.source_models,
-      neoclassical_models=models.neoclassical_models,
       runtime_params=runtime_params,
       geo=geo,
       core_profiles=core_profiles,
       explicit=False,
       explicit_source_profiles=explicit_source_profiles,
-      conductivity=conductivity,
+      conductivity=neoclassical_outputs.conductivity,
+      bootstrap_current=neoclassical_outputs.bootstrap_current,
   )
 
   # --- Transient term coefficients --- #
@@ -259,7 +263,7 @@ def _calc_coeffs_full(
       1.0
       / runtime_params.numerics.resistivity_multiplier
       * geo.rho_norm
-      * conductivity.sigma
+      * neoclassical_outputs.conductivity.sigma
       * consts.mu_0
       * 16
       * jnp.pi**2
@@ -320,12 +324,12 @@ def _calc_coeffs_full(
   transport_coefficients = (
       transport_coefficients_builder.calculate_all_transport_coeffs(
           transport_model=models.transport_model,
-          neoclassical_models=models.neoclassical_models,
           internal_boundary_condition_model=models.internal_boundary_condition_model,
           runtime_params=runtime_params,
           geo=geo,
           core_profiles=core_profiles,
           pedestal_transition_state=pedestal_transition_state,
+          neoclassical_transport=neoclassical_outputs.transport,
           use_pereverzev=use_pereverzev,
       )
   )
@@ -469,7 +473,7 @@ def _calc_coeffs_full(
       * geo.Phi_b_dot
       * geo.Phi_b
       * geo.rho_norm**2
-      * conductivity.sigma
+      * neoclassical_outputs.conductivity.sigma
       / geo.F**2
       * core_profiles.psi.grad()
   )

torax/_src/fvm/tests/calc_coeffs_test.py

@@ -82,7 +82,6 @@ def test_calc_coeffs_smoke_test(
         runtime_params=runtime_params,
         geo=geo,
         core_profiles=core_profiles,
-        neoclassical_models=models.neoclassical_models,
         explicit=True,
     )
     calc_coeffs.calc_coeffs(

torax/_src/fvm/tests/fvm_test.py

@@ -250,7 +250,6 @@ def test_nonlinear_solve_block_loss_minimum(
     evolving_names = tuple(['T_i'])
     explicit_source_profiles = source_profile_builders.build_source_profiles(
         source_models=models.source_models,
-        neoclassical_models=models.neoclassical_models,
         runtime_params=runtime_params,
         geo=geo,
         core_profiles=core_profiles,
@@ -365,7 +364,6 @@ def test_implicit_solve_block_uses_updated_boundary_conditions(self):
         geo=geo,
         core_profiles=initial_core_profiles,
         source_models=models.source_models,
-        neoclassical_models=models.neoclassical_models,
         explicit=True,
     )
 
@@ -488,7 +486,6 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
         geo=geo,
         core_profiles=initial_core_profiles,
         source_models=models.source_models,
-        neoclassical_models=models.neoclassical_models,
         explicit=True,
     )
 

torax/_src/neoclassical/formulas/formulas.py

@@ -13,15 +13,13 @@
 # limitations under the License.
 """Common formulas used in neoclassical models."""
 
-import jax
 import jax.numpy as jnp
 from torax._src import array_typing
 from torax._src import constants
 from torax._src import math_utils
 from torax._src.fvm import cell_variable
 from torax._src.geometry import geometry as geometry_lib
 from torax._src.neoclassical.bootstrap_current import base as bootstrap_current_base
-from torax._src.physics import collisions
 
 
 # pylint: disable=invalid-name
@@ -133,8 +131,7 @@ def calculate_nu_i_star(
   )
 
 
-# Functions to calculate the neoclassical poloidal velocity.
-def _calculate_neoclassical_k_neo(
+def calculate_neoclassical_k_neo(
     nu_star: array_typing.FloatScalar, epsilon: array_typing.FloatScalar
 ):
   """Calculates the neoclassical coefficient k_neo.
@@ -177,94 +174,6 @@ def _calculate_neoclassical_k_neo(
 # See Sauter (1999) Eq. 17a-17b
 
 
-@jax.jit
-def calculate_poloidal_velocity(
-    T_i: cell_variable.CellVariable,
-    n_i: array_typing.FloatVectorFace,
-    q: array_typing.FloatVectorFace,
-    Z_eff: array_typing.FloatVectorFace,
-    Z_i: array_typing.FloatVectorFace,
-    B_tor: array_typing.FloatVectorFace,
-    B_total_squared: array_typing.FloatVectorFace,
-    geo: geometry_lib.Geometry,
-    poloidal_velocity_multiplier: array_typing.FloatScalar = 1.0,
-) -> cell_variable.CellVariable:
-  """Computes the neoclassical ion poloidal velocity profile.
-
-  Implementing eq.33 from
-  Y. B. Kim , P. H. Diamond , R. J. Groebner.
-  "Neoclassical poloidal and toroidal rotation in tokamaks"
-  Phys. Fluids B 3, 2050–2060 (1991)
-  https://doi.org/10.1063/1.859671
-
-  Eq. 33 can be simplified to the following form in SI units:
-  v_pol = k_neo * (dT/dr) * (B_tor / <B^2>) / (Z * e)
-
-  Args:
-    T_i: Ion temperature as a cell variable [keV].
-    n_i: Ion density on the face grid [m^-3].
-    q: Safety factor on the face grid.
-    Z_eff: Effective charge on the face grid.
-    Z_i: Main ion charge on the face grid.
-    B_tor: Toroidal magnetic field on the face grid [T].
-    B_total_squared: Total magnetic field (toroidal + poloidal) on the face grid
-      [T].
-    geo: Geometry
-    poloidal_velocity_multiplier: A multiplier to apply to the poloidal
-      velocity.
-
-  Returns:
-    v_pol : Poloidal velocity profile [m/s].
-  """
-  # Note: all computations are performed on the face grid.
-
-  T_i_face = T_i.face_value()
-  epsilon = geo.epsilon_face
-
-  # Calculate Neoclassical Coefficient k_i
-  log_lambda_ii = collisions.calculate_log_lambda_ii(
-      T_i_face,  # pyrefly: ignore[bad-argument-type]
-      n_i,  # pyrefly: ignore[bad-argument-type]
-      Z_eff,  # pyrefly: ignore[bad-argument-type]
-  )
-  nu_i_star = calculate_nu_i_star(
-      q=q,
-      geo=geo,
-      n_i=n_i,
-      T_i=T_i_face,  # pyrefly: ignore[bad-argument-type]
-      Z_eff=Z_eff,
-      log_lambda_ii=log_lambda_ii,
-  )
-  k_neo = _calculate_neoclassical_k_neo(nu_i_star, epsilon)
-
-  # Calculate Radial Temperature Gradient (dT/dr)
-  grad_Ti = (
-      T_i.face_grad(
-          x=geo.r_mid, x_left=geo.r_mid_face[0], x_right=geo.r_mid_face[-1]
-      )
-      * constants.CONSTANTS.keV_to_J
-  )  # [J/m]
-
-  # Calculate Poloidal Velocity
-  # v_pol = k_i * (dT/dr) * (B_tor / <B^2>) / (Z * e)
-  B_total_squared_safe = jnp.maximum(B_total_squared, constants.CONSTANTS.eps)
-  v_pol = (
-      k_neo
-      * grad_Ti
-      * (B_tor / B_total_squared_safe)
-      / (constants.CONSTANTS.q_e * Z_i)
-  )
-
-  v_pol = poloidal_velocity_multiplier * v_pol
-
-  return cell_variable.CellVariable(
-      value=geometry_lib.face_to_cell(v_pol),
-      face_centers=geo.rho_face_norm,
-      right_face_constraint=v_pol[-1],
-      right_face_grad_constraint=None,
-  )
-
-
 def calculate_analytic_bootstrap_current(
     *,
     bootstrap_multiplier: float,