Move calculate_psidot_from_psi_sources to new psidot_calculations

Could not move it to psi_calculations due to circular dependency. Nevertheless it should move to the physics package and not remain in ohmic_heat_source

PiperOrigin-RevId: 734464091

Author: jcitrin

torax/core_profiles/initialization.py

@@ -28,8 +28,8 @@
 from torax.geometry import geometry
 from torax.geometry import standard_geometry
 from torax.physics import psi_calculations
-from torax.sources import ohmic_heat_source
 from torax.sources import source_models as source_models_lib
+from torax.sources import source_operations
 from torax.sources import source_profile_builders
 from torax.sources import source_profiles as source_profiles_lib
 
@@ -500,8 +500,13 @@ def _init_psi_psidot_vloop_and_current(
   # 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,
+  psi_sources = source_operations.sum_sources_psi(geo, source_profiles)
+  sigma = source_profiles.j_bootstrap.sigma
+  sigma_face = source_profiles.j_bootstrap.sigma_face
+  psidot = psi_calculations.calculate_psidot_from_psi_sources(
+      psi_sources=psi_sources,
+      sigma=sigma,
+      sigma_face=sigma_face,
       resistivity_multiplier=dynamic_runtime_params_slice.numerics.resistivity_mult,
       psi=psi,
       geo=geo,

torax/orchestration/step_function.py

@@ -31,7 +31,7 @@
 from torax.geometry import geometry_provider as geometry_provider_lib
 from torax.pedestal_model import pedestal_model as pedestal_model_lib
 from torax.physics import psi_calculations
-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
@@ -543,7 +543,7 @@ def finalize_output(
         dynamic_runtime_params_slice=dynamic_runtime_params_slice_t_plus_dt,
         geo=geo_t_plus_dt,
         core_profiles=output_state.core_profiles,
-        core_sources=output_state.core_sources,
+        source_profiles=output_state.core_sources,
     )
     output_state = post_processing.make_outputs(
         sim_state=output_state,
@@ -691,13 +691,18 @@ def _update_psidot(
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
-    core_sources: source_profiles_lib.SourceProfiles,
+    source_profiles: source_profiles_lib.SourceProfiles,
 ) -> state.CoreProfiles:
   """Update psidot based on new core_profiles."""
+  psi_sources = source_operations.sum_sources_psi(geo, source_profiles)
+  sigma = source_profiles.j_bootstrap.sigma
+  sigma_face = source_profiles.j_bootstrap.sigma_face
   psidot = dataclasses.replace(
       core_profiles.psidot,
-      value=ohmic_heat_source.calculate_psidot_from_psi_sources(
-          source_profiles=core_sources,
+      value=psi_calculations.calculate_psidot_from_psi_sources(
+          psi_sources=psi_sources,
+          sigma=sigma,
+          sigma_face=sigma_face,
           resistivity_multiplier=dynamic_runtime_params_slice.numerics.resistivity_mult,
           psi=core_profiles.psi,
           geo=geo,

torax/physics/psi_calculations.py

@@ -40,6 +40,8 @@
 from torax import jax_utils
 from torax import state
 from torax.fvm import cell_variable
+from torax.fvm import convection_terms
+from torax.fvm import diffusion_terms
 from torax.geometry import geometry
 
 _trapz = jax.scipy.integrate.trapezoid
@@ -318,3 +320,79 @@ def calculate_psi_grad_constraint_from_Ip_tot(
       * (16 * jnp.pi**3 * constants.CONSTANTS.mu0 * geo.Phib)
       / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1])
   )
+
+
+def calculate_psidot_from_psi_sources(
+    *,
+    psi_sources: array_typing.ArrayFloat,
+    sigma: array_typing.ArrayFloat,
+    sigma_face: array_typing.ArrayFloat,
+    resistivity_multiplier: float,
+    psi: cell_variable.CellVariable,
+    geo: geometry.Geometry,
+) -> jax.Array:
+  """Calculates psidot (loop voltage) from the sum of the psi sources."""
+
+  # Calculate transient term
+  consts = constants.CONSTANTS
+  toc_psi = (
+      1.0
+      / resistivity_multiplier
+      * geo.rho_norm
+      * sigma
+      * consts.mu0
+      * 16
+      * jnp.pi**2
+      * geo.Phib**2
+      / geo.F**2
+  )
+  # Calculate diffusion term coefficient
+  d_face_psi = geo.g2g3_over_rhon_face
+  # Add phibdot terms to poloidal flux convection
+  v_face_psi = (
+      -8.0
+      * jnp.pi**2
+      * consts.mu0
+      * geo.Phibdot
+      * geo.Phib
+      * sigma_face
+      * geo.rho_face_norm**2
+      / geo.F_face**2
+  )
+
+  # Add effective phibdot poloidal flux source term
+  ddrnorm_sigma_rnorm2_over_f2 = jnp.gradient(
+      sigma * geo.rho_norm**2 / geo.F**2, geo.rho_norm
+  )
+
+  psi_sources += (
+      -8.0
+      * jnp.pi**2
+      * consts.mu0
+      * geo.Phibdot
+      * geo.Phib
+      * ddrnorm_sigma_rnorm2_over_f2
+  )
+
+  diffusion_mat, diffusion_vec = diffusion_terms.make_diffusion_terms(
+      d_face_psi, psi
+  )
+
+  # Set the psi convection term for psidot used in ohmic power, always with
+  # the default 'ghost' mode. Impact of different modes would mildly impact
+  # Ohmic power at the LCFS which has negligible impact on simulations.
+  # Allowing it to be configurable introduces more complexity in the code by
+  # needing to pass in the mode from the static_runtime_params across multiple
+  # functions.
+  conv_mat, conv_vec = convection_terms.make_convection_terms(
+      v_face_psi, d_face_psi, psi
+  )
+
+  c_mat = diffusion_mat + conv_mat
+  c = diffusion_vec + conv_vec
+
+  c += psi_sources
+
+  psidot = (jnp.dot(c_mat, psi.value) + c) / toc_psi
+
+  return psidot

torax/physics/tests/psi_calculations_test.py

@@ -18,13 +18,17 @@
 import jax
 import numpy as np
 from torax import constants
+from torax.config import build_runtime_params
 from torax.core_profiles import initialization
 from torax.geometry import pydantic_model as geometry_pydantic_model
 from torax.geometry import standard_geometry
 from torax.physics import psi_calculations
+from torax.sources import runtime_params as source_runtime_params
+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
 from torax.tests.test_lib import torax_refs
 
-
 _trapz = jax.scipy.integrate.trapezoid
 
 
@@ -111,6 +115,79 @@ def test_calc_s(self, references_getter: Callable[[], torax_refs.References]):
 
     np.testing.assert_allclose(s, references.s, rtol=1e-5)
 
+  @parameterized.parameters([
+      dict(references_getter=torax_refs.circular_references),
+      dict(references_getter=torax_refs.chease_references_Ip_from_chease),
+      dict(
+          references_getter=torax_refs.chease_references_Ip_from_runtime_params
+      ),
+  ])
+  def test_calc_psidot(
+      self, references_getter: Callable[[], torax_refs.References]
+  ):
+    references = references_getter()
+
+    runtime_params = references.runtime_params
+    source_models_builder = source_models_lib.SourceModelsBuilder()
+    source_models_builder.runtime_params['generic_current_source'].mode = (
+        source_runtime_params.Mode.MODEL_BASED
+    )
+    source_models = source_models_builder()
+    dynamic_runtime_params_slice, geo = (
+        torax_refs.build_consistent_dynamic_runtime_params_slice_and_geometry(
+            runtime_params,
+            references.geometry_provider,
+            sources=source_models_builder.runtime_params,
+        )
+    )
+    source_profiles = source_profiles_lib.SourceProfiles(
+        j_bootstrap=source_profiles_lib.BootstrapCurrentProfile.zero_profile(
+            geo
+        ),
+        qei=source_profiles_lib.QeiInfo.zeros(geo),
+    )
+    static_slice = build_runtime_params.build_static_runtime_params_slice(
+        runtime_params=runtime_params,
+        source_runtime_params=source_models_builder.runtime_params,
+        torax_mesh=geo.torax_mesh,
+    )
+    initial_core_profiles = initialization.initial_core_profiles(
+        static_slice,
+        dynamic_runtime_params_slice,
+        geo,
+        source_models=source_models,
+    )
+    # Updates the calculated source profiles with the standard source profiles.
+    source_profile_builders.build_standard_source_profiles(
+        static_runtime_params_slice=static_slice,
+        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+        geo=geo,
+        core_profiles=initial_core_profiles,
+        source_models=source_models,
+        psi_only=True,
+        calculate_anyway=True,
+        calculated_source_profiles=source_profiles,
+    )
+    bootstrap_profiles = source_models.j_bootstrap.get_bootstrap(
+        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+        static_runtime_params_slice=static_slice,
+        geo=geo,
+        core_profiles=initial_core_profiles,
+    )
+
+    psidot_calculated = psi_calculations.calculate_psidot_from_psi_sources(
+        psi_sources=sum(source_profiles.psi.values()),
+        sigma=bootstrap_profiles.sigma,
+        sigma_face=bootstrap_profiles.sigma_face,
+        resistivity_multiplier=dynamic_runtime_params_slice.numerics.resistivity_mult,
+        psi=references.psi,
+        geo=geo,
+    )
+
+    psidot_expected = references.psidot
+
+    np.testing.assert_allclose(psidot_calculated, psidot_expected, rtol=1e-5)
+
   # pylint: disable=invalid-name
   def test_calc_Wpol(self):
     # Small inverse aspect ratio limit of circular geometry, such that we

torax/sources/ohmic_heat_source.py

@@ -19,14 +19,9 @@
 from typing import ClassVar, Literal
 
 import chex
-import jax
 import jax.numpy as jnp
-from torax import constants
 from torax import state
 from torax.config import runtime_params_slice
-from torax.fvm import cell_variable
-from torax.fvm import convection_terms
-from torax.fvm import diffusion_terms
 from torax.geometry import geometry
 from torax.physics import psi_calculations
 from torax.sources import runtime_params as runtime_params_lib
@@ -35,83 +30,6 @@
 from torax.sources import source_profiles as source_profiles_lib
 
 
-def calculate_psidot_from_psi_sources(
-    *,
-    source_profiles: source_profiles_lib.SourceProfiles,
-    resistivity_multiplier: float,
-    psi: cell_variable.CellVariable,
-    geo: geometry.Geometry,
-) -> jax.Array:
-  """Calculates psidot (loop voltage) from precalculated sources."""
-  psi_sources = source_operations.sum_sources_psi(geo, source_profiles)
-  sigma = source_profiles.j_bootstrap.sigma
-  sigma_face = source_profiles.j_bootstrap.sigma_face
-
-  # Calculate transient term
-  consts = constants.CONSTANTS
-  toc_psi = (
-      1.0
-      / resistivity_multiplier
-      * geo.rho_norm
-      * sigma
-      * consts.mu0
-      * 16
-      * jnp.pi**2
-      * geo.Phib**2
-      / geo.F**2
-  )
-  # Calculate diffusion term coefficient
-  d_face_psi = geo.g2g3_over_rhon_face
-  # Add phibdot terms to poloidal flux convection
-  v_face_psi = (
-      -8.0
-      * jnp.pi**2
-      * consts.mu0
-      * geo.Phibdot
-      * geo.Phib
-      * sigma_face
-      * geo.rho_face_norm**2
-      / geo.F_face**2
-  )
-
-  # Add effective phibdot poloidal flux source term
-  ddrnorm_sigma_rnorm2_over_f2 = jnp.gradient(
-      sigma * geo.rho_norm**2 / geo.F**2, geo.rho_norm
-  )
-
-  psi_sources += (
-      -8.0
-      * jnp.pi**2
-      * consts.mu0
-      * geo.Phibdot
-      * geo.Phib
-      * ddrnorm_sigma_rnorm2_over_f2
-  )
-
-  diffusion_mat, diffusion_vec = diffusion_terms.make_diffusion_terms(
-      d_face_psi, psi
-  )
-
-  # Set the psi convection term for psidot used in ohmic power, always with
-  # the default 'ghost' mode. Impact of different modes would mildly impact
-  # Ohmic power at the LCFS which has negligible impact on simulations.
-  # Allowing it to be configurable introduces more complexity in the code by
-  # needing to pass in the mode from the static_runtime_params across multiple
-  # functions.
-  conv_mat, conv_vec = convection_terms.make_convection_terms(
-      v_face_psi, d_face_psi, psi
-  )
-
-  c_mat = diffusion_mat + conv_mat
-  c = diffusion_vec + conv_vec
-
-  c += psi_sources
-
-  psidot = (jnp.dot(c_mat, psi.value) + c) / toc_psi
-
-  return psidot
-
-
 def ohmic_model_func(
     unused_static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
@@ -132,8 +50,15 @@ def ohmic_model_func(
       geo,
       core_profiles.psi,
   )
-  psidot = calculate_psidot_from_psi_sources(
-      source_profiles=calculated_source_profiles,
+  psi_sources = source_operations.sum_sources_psi(
+      geo, calculated_source_profiles
+  )
+  sigma = calculated_source_profiles.j_bootstrap.sigma
+  sigma_face = calculated_source_profiles.j_bootstrap.sigma_face
+  psidot = psi_calculations.calculate_psidot_from_psi_sources(
+      psi_sources=psi_sources,
+      sigma=sigma,
+      sigma_face=sigma_face,
       resistivity_multiplier=dynamic_runtime_params_slice.numerics.resistivity_mult,
       psi=core_profiles.psi,
       geo=geo,
@@ -144,6 +69,7 @@ def ohmic_model_func(
 
 class OhmicHeatSourceConfig(runtime_params_lib.SourceModelBase):
   """Configuration for the OhmicHeatSource."""
+
   source_name: Literal['ohmic_heat_source'] = 'ohmic_heat_source'
   mode: runtime_params_lib.Mode = runtime_params_lib.Mode.MODEL_BASED