Ensure that there is one way to collect psi sources and make it easie…

…r to reuse values

PiperOrigin-RevId: 725181143

torax/config/build_sim.py

@@ -410,7 +410,6 @@ def _build_single_source_builder_from_config(
     source_builder_class = source_lib.make_source_builder(
         supported_source.source_class,
         runtime_params_type=model_function.runtime_params_class,
-        links_back=model_function.links_back,
         model_func=model_function.source_profile_function,
     )
 

torax/core_profile_setters.py

@@ -38,6 +38,7 @@
 from torax.geometry import standard_geometry
 from torax.sources import ohmic_heat_source
 from torax.sources import source_models as source_models_lib
+from torax.sources import source_profile_builders
 from torax.sources import source_profiles as source_profiles_lib
 
 _trapz = jax.scipy.integrate.trapezoid
@@ -341,48 +342,19 @@ def _prescribe_currents(
 
 
 def _calculate_currents_from_psi(
-    static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
-    dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
-    source_models: source_models_lib.SourceModels,
+    source_profiles: source_profiles_lib.SourceProfiles,
 ) -> state.Currents:
-  """Creates the initial Currents using psi to calculate jtot.
-
-  Args:
-    static_runtime_params_slice: Static runtime parameters.
-    dynamic_runtime_params_slice: General runtime parameters at t_initial.
-    geo: Geometry of the tokamak.
-    core_profiles: Core profiles.
-    source_models: All TORAX source/sink functions. If not provided, uses the
-      default sources.
-
-  Returns:
-    currents: Plasma currents
-  """
-
-  # Many variables throughout this function are capitalized based on physics
-  # notational conventions rather than on Google Python style
+  """Creates the initial Currents using psi to calculate jtot."""
   jtot, jtot_face, Ip_profile_face = physics.calc_jtot_from_psi(
       geo,
       core_profiles.psi,
   )
-
-  bootstrap_profile = source_models.j_bootstrap.get_bootstrap(
-      dynamic_runtime_params_slice=dynamic_runtime_params_slice,
-      static_runtime_params_slice=static_runtime_params_slice,
-      geo=geo,
-      core_profiles=core_profiles,
-  )
-
-  # calculate "External" current profile (e.g. ECCD)
-  # form of external current on face grid
-  external_current = source_models.external_current_source(
-      dynamic_runtime_params_slice=dynamic_runtime_params_slice,
-      static_runtime_params_slice=static_runtime_params_slice,
-      geo=geo,
-      core_profiles=core_profiles,
-  )
+  bootstrap_profile = source_profiles.j_bootstrap
+  # Note that the psi sources here are the standard sources and don't include
+  # the bootstrap current.
+  external_current = sum(source_profiles.psi.values())
   johm = jtot - external_current - bootstrap_profile.j_bootstrap
   currents = state.Currents(
       jtot=jtot,
@@ -466,7 +438,7 @@ def _calculate_psi_grad_constraint(
   )
 
 
-def _init_psi_and_current(
+def _init_psi_psidot_and_current(
     static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,
@@ -493,6 +465,22 @@ def _init_psi_and_current(
   Returns:
     Refined core profiles.
   """
+  source_profiles = source_profiles_lib.SourceProfiles(
+      j_bootstrap=source_profiles_lib.BootstrapCurrentProfile.zero_profile(geo),
+      qei=source_profiles_lib.QeiInfo.zeros(geo),
+  )
+  # Updates the calculated source profiles with the standard source profiles.
+  source_profile_builders.build_standard_source_profiles(
+      static_runtime_params_slice=static_runtime_params_slice,
+      dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+      geo=geo,
+      core_profiles=core_profiles,
+      source_models=source_models,
+      psi_only=True,
+      calculate_anyway=True,
+      calculated_source_profiles=source_profiles,
+  )
+
   # Retrieving psi from the profile conditions.
   if dynamic_runtime_params_slice.profile_conditions.psi is not None:
     psi = cell_variable.CellVariable(
@@ -504,12 +492,19 @@ def _init_psi_and_current(
         dr=geo.drho_norm,
     )
     core_profiles = dataclasses.replace(core_profiles, psi=psi)
-    currents = _calculate_currents_from_psi(
-        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+    bootstrap_profile = source_models.j_bootstrap.get_bootstrap(
         static_runtime_params_slice=static_runtime_params_slice,
+        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
         geo=geo,
         core_profiles=core_profiles,
-        source_models=source_models,
+    )
+    source_profiles = dataclasses.replace(
+        source_profiles, j_bootstrap=bootstrap_profile
+    )
+    currents = _calculate_currents_from_psi(
+        geo=geo,
+        core_profiles=core_profiles,
+        source_profiles=source_profiles,
     )
   # Retrieving psi from the standard geometry input.
   elif (
@@ -528,30 +523,29 @@ def _init_psi_and_current(
         dr=geo.drho_norm,
     )
     core_profiles = dataclasses.replace(core_profiles, psi=psi)
-    currents = _calculate_currents_from_psi(
-        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+    bootstrap_profile = source_models.j_bootstrap.get_bootstrap(
         static_runtime_params_slice=static_runtime_params_slice,
+        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+        geo=geo,
+        core_profiles=core_profiles,
+    )
+    source_profiles = dataclasses.replace(
+        source_profiles, j_bootstrap=bootstrap_profile
+    )
+    currents = _calculate_currents_from_psi(
         geo=geo,
         core_profiles=core_profiles,
-        source_models=source_models,
+        source_profiles=source_profiles,
     )
   # Calculating j according to nu formula and psi from j.
   elif (
       isinstance(geo, circular_geometry.CircularAnalyticalGeometry)
       or dynamic_runtime_params_slice.profile_conditions.initial_psi_from_j
   ):
     # First calculate currents without bootstrap.
-    bootstrap = source_profiles_lib.BootstrapCurrentProfile.zero_profile(
-        geo
-    )
-    external_current = source_models.external_current_source(
-        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
-        static_runtime_params_slice=static_runtime_params_slice,
-        geo=geo,
-        core_profiles=core_profiles,
-    )
+    external_current = sum(source_profiles.psi.values())
     currents = _prescribe_currents(
-        bootstrap_profile=bootstrap,
+        bootstrap_profile=source_profiles.j_bootstrap,
         external_current=external_current,
         dynamic_runtime_params_slice=dynamic_runtime_params_slice,
         geo=geo,
@@ -590,7 +584,28 @@ def _init_psi_and_current(
   else:
     raise ValueError('Cannot compute psi for given config.')
 
-  core_profiles = dataclasses.replace(core_profiles, psi=psi, currents=currents)
+  core_profiles = dataclasses.replace(
+      core_profiles, psi=psi, currents=currents)
+  bootstrap_profile = source_models.j_bootstrap.get_bootstrap(
+      dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+      static_runtime_params_slice=static_runtime_params_slice,
+      geo=geo,
+      core_profiles=core_profiles,
+  )
+  source_profiles = dataclasses.replace(
+      source_profiles, j_bootstrap=bootstrap_profile
+  )
+  # psidot calculated here with phibdot=0 in geo, since this is initial
+  # conditions and we don't yet have information on geo_t_plus_dt for the
+  # phibdot calculation.
+  psidot = ohmic_heat_source.calculate_psidot_from_psi_sources(
+      source_profiles=source_profiles,
+      resistivity_multiplier=dynamic_runtime_params_slice.numerics.resistivity_mult,
+      psi=psi,
+      geo=geo,
+  )
+  psidot_cell_var = dataclasses.replace(core_profiles.psidot, value=psidot)
+  core_profiles = dataclasses.replace(core_profiles, psidot=psidot_cell_var)
 
   return core_profiles
 
@@ -667,29 +682,14 @@ def initial_core_profiles(
       nref=jnp.asarray(dynamic_runtime_params_slice.numerics.nref),
   )
 
-  core_profiles = _init_psi_and_current(
+  core_profiles = _init_psi_psidot_and_current(
       static_runtime_params_slice,
       dynamic_runtime_params_slice,
       geo,
       core_profiles,
       source_models,
   )
 
-  # psidot calculated here with phibdot=0 in geo, since this is initial
-  # conditions and we don't yet have information on geo_t_plus_dt for the
-  # phibdot calculation.
-  psidot = dataclasses.replace(
-      core_profiles.psidot,
-      value=ohmic_heat_source.calc_psidot(
-          static_runtime_params_slice,
-          dynamic_runtime_params_slice,
-          geo,
-          core_profiles,
-          source_models,
-      ),
-  )
-  core_profiles = dataclasses.replace(core_profiles, psidot=psidot)
-
   # Set psi as source of truth and recalculate jtot, q, s
   return physics.update_jtot_q_face_s_face(
       geo=geo,

torax/orchestration/step_function.py

@@ -30,7 +30,6 @@
 from torax.geometry import geometry_provider as geometry_provider_lib
 from torax.pedestal_model import pedestal_model as pedestal_model_lib
 from torax.sources import ohmic_heat_source
-from torax.sources import source_operations
 from torax.sources import source_profile_builders
 from torax.sources import source_profiles as source_profiles_lib
 from torax.stepper import stepper as stepper_lib
@@ -690,14 +689,10 @@ def _update_psidot(
     core_sources: source_profiles_lib.SourceProfiles,
 ) -> state.CoreProfiles:
   """Update psidot based on new core_profiles."""
-  psi_sources = source_operations.sum_sources_psi(geo, core_sources)
-
   psidot = dataclasses.replace(
       core_profiles.psidot,
       value=ohmic_heat_source.calculate_psidot_from_psi_sources(
-          psi_sources=psi_sources,
-          sigma=core_sources.j_bootstrap.sigma,
-          sigma_face=core_sources.j_bootstrap.sigma_face,
+          source_profiles=core_sources,
           resistivity_multiplier=dynamic_runtime_params_slice.numerics.resistivity_mult,
           psi=core_profiles.psi,
           geo=geo,

torax/sources/bremsstrahlung_heat_sink.py

@@ -29,7 +29,6 @@
 from torax.geometry import geometry
 from torax.sources import runtime_params as runtime_params_lib
 from torax.sources import source
-from torax.sources import source_models
 from torax.sources import source_profiles
 
 
@@ -128,7 +127,6 @@ def bremsstrahlung_model_func(
     source_name: str,
     core_profiles: state.CoreProfiles,
     unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
-    unused_model_func: source_models.SourceModels | None,
 ) -> tuple[chex.Array, ...]:
   """Model function for the Bremsstrahlung heat sink."""
   dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[

torax/sources/cyclotron_radiation_heat_sink.py

@@ -31,7 +31,6 @@
 from torax.geometry import geometry
 from torax.sources import runtime_params as runtime_params_lib
 from torax.sources import source
-from torax.sources import source_models
 from torax.sources import source_profiles
 
 
@@ -285,7 +284,6 @@ def cyclotron_radiation_albajar(
     source_name: str,
     core_profiles: state.CoreProfiles,
     unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
-    unused_source_models: source_models.SourceModels,
 ) -> tuple[array_typing.ArrayFloat, ...]:
   """Calculates the cyclotron radiation heat sink contribution to the electron heat equation.
 
@@ -314,7 +312,6 @@ def cyclotron_radiation_albajar(
     source_name: The name of the source.
     core_profiles: The core profiles object.
     unused_calculated_source_profiles: Unused.
-    unused_source_models: Unused.
 
   Returns:
     The cyclotron radiation heat sink contribution to the electron heat

torax/sources/electron_cyclotron_source.py

@@ -30,7 +30,6 @@
 from torax.sources import formulas
 from torax.sources import runtime_params as runtime_params_lib
 from torax.sources import source
-from torax.sources import source_models
 from torax.sources import source_profiles
 
 InterpolatedVarTimeRhoInput = (
@@ -104,35 +103,29 @@ class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
 
 
 def calc_heating_and_current(
-    static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
+    unused_static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,
     source_name: str,
     core_profiles: state.CoreProfiles,
     unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
-    unused_source_models: source_models.SourceModels | None = None,
 ) -> tuple[chex.Array, ...]:
   """Model function for the electron-cyclotron source.
 
   Based on Lin-Liu, Y. R., Chan, V. S., & Prater, R. (2003).
   See https://torax.readthedocs.io/en/latest/electron-cyclotron-derivation.html
 
   Args:
-    static_runtime_params_slice: Static runtime parameters.
+    unused_static_runtime_params_slice: Static runtime parameters.
     dynamic_runtime_params_slice: Global runtime parameters
     geo: Magnetic geometry.
     source_name: Name of the source.
     core_profiles: CoreProfiles component of the state.
     unused_calculated_source_profiles: Unused.
-    unused_source_models: Unused.
 
   Returns:
     2D array of electron cyclotron heating power density and current density.
   """
-  del (
-      unused_source_models,
-      static_runtime_params_slice,
-  )  # Unused.
   dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[
       source_name
   ]

torax/sources/electron_density_sources.py

@@ -28,7 +28,6 @@
 from torax.sources import formulas
 from torax.sources import runtime_params as runtime_params_lib
 from torax.sources import source
-from torax.sources import source_models
 from torax.sources import source_profiles
 
 
@@ -83,7 +82,6 @@ def calc_puff_source(
     source_name: str,
     unused_state: state.CoreProfiles,
     unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
-    unused_source_models: source_models.SourceModels | None = None,
 ) -> tuple[chex.Array, ...]:
   """Calculates external source term for n from puffs."""
   dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[
@@ -176,7 +174,6 @@ def calc_generic_particle_source(
     source_name: str,
     unused_state: state.CoreProfiles,
     unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
-    unused_source_models: source_models.SourceModels | None = None,
 ) -> tuple[chex.Array, ...]:
   """Calculates external source term for n from SBI."""
   dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[
@@ -262,7 +259,6 @@ def calc_pellet_source(
     source_name: str,
     unused_state: state.CoreProfiles,
     unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
-    unused_source_models: source_models.SourceModels | None = None,
 ) -> tuple[chex.Array, ...]:
   """Calculates external source term for n from pellets."""
   dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[

torax/sources/fusion_heat_source.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import ClassVar, Optional
+from typing import ClassVar
 
 import chex
 import jax
@@ -149,7 +149,6 @@ def fusion_heat_model_func(
     unused_source_name: str,
     core_profiles: state.CoreProfiles,
     unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
-    unused_source_models: Optional['source_models.SourceModels'],
 ) -> tuple[chex.Array, ...]:
   """Model function for fusion heating."""
   # pytype: enable=name-error