Use gaussian_profile formula in generic_ion_el_heat_source.

PiperOrigin-RevId: 696878243
google-deepmind · Nov 15, 2024 · fc1a6b1 · fc1a6b1
1 parent 2178784
commit fc1a6b1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 11 deletions.
diff --git a/torax/sources/formulas.py b/torax/sources/formulas.py
@@ -69,6 +69,7 @@ def exponential_profile(
 
 def gaussian_profile(
     geo: geometry.Geometry,
+    *,
     c1: float,
     c2: float,
     total: float,

diff --git a/torax/sources/generic_ion_el_heat_source.py b/torax/sources/generic_ion_el_heat_source.py
@@ -27,6 +27,7 @@
 from torax import interpolated_param
 from torax import state
 from torax.config import runtime_params_slice
+from torax.sources import formulas
 from torax.sources import runtime_params as runtime_params_lib
 from torax.sources import source
 
@@ -106,17 +107,10 @@ def calc_generic_heat_source(
     source_ion: source term for ions.
     source_el: source term for electrons.
   """
-
-  # calculate heat profile (face grid)
-  Q = jnp.exp(-((geo.rho_norm - rsource) ** 2) / (2 * w**2))
-  Q_face = jnp.exp(-((geo.rho_face_norm - rsource) ** 2) / (2 * w**2))
-  # calculate constant prefactor
-  C = Ptot / jax.scipy.integrate.trapezoid(
-      geo.vpr_face * Q_face, geo.rho_face_norm
-  )
-
-  source_ion = C * Q * (1 - el_heat_fraction)
-  source_el = C * Q * el_heat_fraction
+  # Calculate heat profile.
+  profile = formulas.gaussian_profile(geo, c1=rsource, c2=w, total=Ptot)
+  source_ion = profile * (1 - el_heat_fraction)
+  source_el = profile * el_heat_fraction
 
   return source_ion, source_el