Reorganize physics package.

Previously, various physics calculations which were not directly associated with an existing component like sources, transport_models, were dispersed between physics.py and post_processing.py.

This refactor introduces a physics package and clearer organization of physics functions into thematic modules: collisions, charge_states, psi_calculations, and scaling_laws.

Drive-by: cleanup of unnecessary comments and docstring fixes

Upcoming PR will move "physics" calculations from post_processing into the physics package.

PiperOrigin-RevId: 732272746

torax/config/tests/plasma_composition.py

@@ -17,10 +17,10 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
-from torax import charge_states
 from torax import interpolated_param
 from torax.config import plasma_composition
 from torax.geometry import pydantic_model as geometry_pydantic_model
+from torax.physics import charge_states
 
 
 class PlasmaCompositionTest(parameterized.TestCase):

torax/core_profiles/formulas.py

@@ -25,8 +25,6 @@
 
 _trapz = jax.scipy.integrate.trapezoid
 
-# Using capitalized variables for physics notational conventions rather than
-# Python style.
 # pylint: disable=invalid-name
 
 
@@ -157,3 +155,14 @@ 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])
   )
+
+
+# TODO(b/377225415): generalize to arbitrary number of ions.
+def get_main_ion_dilution_factor(
+    Zi: array_typing.ScalarFloat,
+    Zimp: array_typing.ArrayFloat,
+    Zeff: array_typing.ArrayFloat,
+) -> jax.Array:
+  """Calculates the main ion dilution factor based on a single assumed impurity and general main ion charge."""
+  return (Zimp - Zeff) / (Zi * (Zimp - Zi))
+

torax/core_profiles/initialization.py

@@ -20,23 +20,21 @@
 from torax import array_typing
 from torax import constants
 from torax import math_utils
-from torax import physics
 from torax import state
 from torax.config import runtime_params_slice
 from torax.core_profiles import formulas
 from torax.core_profiles import updaters
 from torax.fvm import cell_variable
 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_profile_builders
 from torax.sources import source_profiles as source_profiles_lib
 
 _trapz = jax.scipy.integrate.trapezoid
 
-# Using capitalized variables for physics notational conventions rather than
-# Python style.
 # pylint: disable=invalid-name
 
 
@@ -133,7 +131,7 @@ def initial_core_profiles(
   )
 
   # Set psi as source of truth and recalculate jtot, q, s
-  return physics.update_jtot_q_face_s_face(
+  return psi_calculations.update_jtot_q_face_s_face(
       geo=geo,
       core_profiles=core_profiles,
   )
@@ -147,8 +145,6 @@ def _prescribe_currents(
 ) -> state.Currents:
   """Creates the initial Currents from a given bootstrap profile."""
 
-  # Many variables throughout this function are capitalized based on physics
-  # notational conventions rather than on Google Python style
   Ip = dynamic_runtime_params_slice.profile_conditions.Ip_tot
   f_bootstrap = bootstrap_profile.I_bootstrap / (Ip * 1e6)
 
@@ -201,7 +197,7 @@ def _calculate_currents_from_psi(
     source_profiles: source_profiles_lib.SourceProfiles,
 ) -> state.Currents:
   """Creates the initial Currents using psi to calculate jtot."""
-  jtot, jtot_face, Ip_profile_face = physics.calc_jtot_from_psi(
+  jtot, jtot_face, Ip_profile_face = psi_calculations.calc_jtot(
       geo,
       core_profiles.psi,
   )
@@ -474,7 +470,7 @@ def _init_psi_psidot_vloop_and_current(
         currents.jtot_hires,
         use_vloop_lcfs_boundary_condition=use_vloop_bc,
     )
-    _, _, Ip_profile_face = physics.calc_jtot_from_psi(
+    _, _, Ip_profile_face = psi_calculations.calc_jtot(
         geo,
         psi,
     )

torax/core_profiles/tests/formulas_test.py

@@ -279,6 +279,24 @@ def test_ne_core_profile_setter_with_fGW(
     np.all(np.isclose(ratio, ratio[0]))
     self.assertNotEqual(ratio[0], 1.0)
 
+  # TODO(b/377225415): generalize to arbitrary number of ions.
+  @parameterized.parameters([
+      dict(Zi=1.0, Zimp=10.0, Zeff=1.0, expected=1.0),
+      dict(Zi=1.0, Zimp=5.0, Zeff=1.0, expected=1.0),
+      dict(Zi=2.0, Zimp=10.0, Zeff=2.0, expected=0.5),
+      dict(Zi=2.0, Zimp=5.0, Zeff=2.0, expected=0.5),
+      dict(Zi=1.0, Zimp=10.0, Zeff=1.9, expected=0.9),
+      dict(Zi=2.0, Zimp=10.0, Zeff=3.6, expected=0.4),
+  ])
+  def test_get_main_ion_dilution_factor(self, Zi, Zimp, Zeff, expected):
+    """Unit test of `get_main_ion_dilution_factor`."""
+    np.testing.assert_allclose(
+        formulas.get_main_ion_dilution_factor(Zi, Zimp, Zeff),
+        expected,
+    )
+
+  # pylint: enable=invalid-name
+
 
 if __name__ == '__main__':
   absltest.main()

torax/core_profiles/tests/initialization_test.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable
 from unittest import mock
 
 from absl.testing import absltest
@@ -24,7 +23,6 @@
 from torax.config import runtime_params_slice as runtime_params_slice_lib
 from torax.core_profiles import initialization
 from torax.geometry import pydantic_model as geometry_pydantic_model
-from torax.geometry import standard_geometry
 from torax.sources import generic_current_source
 from torax.sources import source_models as source_models_lib
 from torax.sources import source_profiles
@@ -40,18 +38,9 @@ def setUp(self):
     jax_utils.enable_errors(True)
     self.geo = geometry_pydantic_model.CircularConfig(n_rho=4).build_geometry()
 
-  @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_update_psi_from_j(
-      self, references_getter: Callable[[], torax_refs.References]
-  ):
+  def test_update_psi_from_j(self):
     """Compare `update_psi_from_j` function to a reference implementation."""
-    references = references_getter()
+    references = torax_refs.circular_references()
 
     runtime_params = references.runtime_params
     source_runtime_params = generic_current_source.RuntimeParams()
@@ -67,29 +56,26 @@ def test_update_psi_from_j(
             },
         )
     )
-    if isinstance(geo, standard_geometry.StandardGeometry):
-      psi = geo.psi_from_Ip
-    else:
-      bootstrap = source_profiles.BootstrapCurrentProfile.zero_profile(geo)
-      external_current = generic_current_source.calculate_generic_current(
-          mock.ANY,
-          dynamic_runtime_params_slice=dynamic_runtime_params_slice,
-          geo=geo,
-          source_name=generic_current_source.GenericCurrentSource.SOURCE_NAME,
-          unused_state=mock.ANY,
-          unused_calculated_source_profiles=mock.ANY,
-      )[0]
-      currents = initialization._prescribe_currents(
-          bootstrap_profile=bootstrap,
-          external_current=external_current,
-          dynamic_runtime_params_slice=dynamic_runtime_params_slice,
-          geo=geo,
-      )
-      psi = initialization._update_psi_from_j(
-          dynamic_runtime_params_slice.profile_conditions.Ip_tot,
-          geo,
-          currents.jtot_hires,
-      ).value
+    bootstrap = source_profiles.BootstrapCurrentProfile.zero_profile(geo)
+    external_current = generic_current_source.calculate_generic_current(
+        mock.ANY,
+        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+        geo=geo,
+        source_name=generic_current_source.GenericCurrentSource.SOURCE_NAME,
+        unused_state=mock.ANY,
+        unused_calculated_source_profiles=mock.ANY,
+    )[0]
+    currents = initialization._prescribe_currents(
+        bootstrap_profile=bootstrap,
+        external_current=external_current,
+        dynamic_runtime_params_slice=dynamic_runtime_params_slice,
+        geo=geo,
+    )
+    psi = initialization._update_psi_from_j(
+        dynamic_runtime_params_slice.profile_conditions.Ip_tot,
+        geo,
+        currents.jtot_hires,
+    ).value
     np.testing.assert_allclose(psi, references.psi.value)
 
   @parameterized.parameters(

torax/core_profiles/tests/updaters_test.py

@@ -18,7 +18,6 @@
 from jax import numpy as jnp
 import numpy as np
 from torax import jax_utils
-from torax import physics
 from torax import state
 from torax.config import profile_conditions as profile_conditions_lib
 from torax.config import runtime_params as general_runtime_params
@@ -91,7 +90,7 @@ def test_get_ion_density_and_charge_states(self):
 
     Zeff = dynamic_runtime_params_slice.plasma_composition.Zeff
 
-    dilution_factor = physics.get_main_ion_dilution_factor(Zi, Zimp, Zeff)
+    dilution_factor = formulas.get_main_ion_dilution_factor(Zi, Zimp, Zeff)
     np.testing.assert_allclose(
         ni.value,
         expected_value * dilution_factor,
@@ -234,7 +233,6 @@ def test_compute_boundary_conditions_ne(
         ne_bound_right_is_fGW and ne_bound_right is not None
     ):
       # Then we expect the boundary condition to be in fGW.
-      # pylint: disable=invalid-name
       nGW = (
           dynamic_runtime_params_slice.profile_conditions.Ip_tot
           / (np.pi * self.geo.Rmin**2)
@@ -355,6 +353,5 @@ def test_compute_boundary_conditions_Ti(
         expected_Ti_bound_right,
     )
 
-
 if __name__ == '__main__':
   absltest.main()

torax/core_profiles/updaters.py

@@ -22,20 +22,17 @@
 import jax
 from jax import numpy as jnp
 from torax import array_typing
-from torax import charge_states
 from torax import jax_utils
-from torax import physics
 from torax import state
 from torax.config import runtime_params_slice
 from torax.core_profiles import formulas
 from torax.fvm import cell_variable
 from torax.geometry import geometry
+from torax.physics import charge_states
 
 _trapz = jax.scipy.integrate.trapezoid
 
 
-# Using capitalized variables for physics notational conventions rather than
-# Python style.
 # pylint: disable=invalid-name
 def _get_charge_states(
     static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
@@ -128,8 +125,8 @@ def get_ion_density_and_charge_states(
   Zeff = dynamic_runtime_params_slice.plasma_composition.Zeff
   Zeff_face = dynamic_runtime_params_slice.plasma_composition.Zeff_face
 
-  dilution_factor = physics.get_main_ion_dilution_factor(Zi, Zimp, Zeff)
-  dilution_factor_edge = physics.get_main_ion_dilution_factor(
+  dilution_factor = formulas.get_main_ion_dilution_factor(Zi, Zimp, Zeff)
+  dilution_factor_edge = formulas.get_main_ion_dilution_factor(
       Zi_face[-1], Zimp_face[-1], Zeff_face[-1]
   )
 
@@ -347,7 +344,7 @@ def compute_boundary_conditions_for_t_plus_dt(
       Te=Te_bound_right,
   )
 
-  dilution_factor_edge = physics.get_main_ion_dilution_factor(
+  dilution_factor_edge = formulas.get_main_ion_dilution_factor(
       Zi_edge,
       Zimp_edge,
       dynamic_runtime_params_slice_t_plus_dt.plasma_composition.Zeff_face[-1],

torax/fvm/calc_coeffs.py

@@ -23,7 +23,6 @@
 import jax.numpy as jnp
 from torax import constants
 from torax import jax_utils
-from torax import physics
 from torax import state
 from torax.config import config_args
 from torax.config import runtime_params_slice
@@ -400,7 +399,7 @@ def _calc_coeffs_full(
 
   # Boolean mask for enforcing internal temperature boundary conditions to
   # model the pedestal.
-  mask = physics.internal_boundary(
+  mask = _internal_boundary(
       geo,
       pedestal_model_output.rho_norm_ped_top,
       dynamic_runtime_params_slice.profile_conditions.set_pedestal,
@@ -784,3 +783,18 @@ def _calc_coeffs_reduced(
       transient_in_cell=transient_in_cell,
   )
   return coeffs
+
+
+# pylint: disable=invalid-name
+def _internal_boundary(
+    geo: geometry.Geometry,
+    Ped_top: jax.Array,
+    set_pedestal: jax.Array,
+) -> jax.Array:
+  # Create Boolean mask FiPy CellVariable with True where the internal boundary
+  # condition is
+  # find index closest to pedestal top.
+  idx = jnp.abs(geo.rho_norm - Ped_top).argmin()
+  mask_np = jnp.zeros(len(geo.rho), dtype=bool)
+  mask_np = jnp.where(set_pedestal, mask_np.at[idx].set(True), mask_np)
+  return mask_np

torax/orchestration/step_function.py

@@ -22,14 +22,14 @@
 import jax
 import jax.numpy as jnp
 from torax import jax_utils
-from torax import physics
 from torax import post_processing
 from torax import state
 from torax.config import runtime_params_slice
 from torax.core_profiles import updaters
 from torax.geometry import geometry
 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_profile_builders
 from torax.sources import source_profiles as source_profiles_lib
@@ -525,7 +525,7 @@ def finalize_output(
     """
 
     # Update total current, q, and s profiles based on new psi
-    output_state.core_profiles = physics.update_jtot_q_face_s_face(
+    output_state.core_profiles = psi_calculations.update_jtot_q_face_s_face(
         geo=geo_t_plus_dt,
         core_profiles=output_state.core_profiles,
     )
@@ -660,7 +660,8 @@ def _update_current_distribution(
   bootstrap_profile = core_sources.j_bootstrap
   # Needed for the case where no psi sources are present.
   external_current = jnp.zeros_like(
-      core_profiles.currents.external_current_source)
+      core_profiles.currents.external_current_source
+  )
   external_current += sum(core_sources.psi.values())
 
   johm = (
@@ -699,7 +700,7 @@ def _update_psidot(
           resistivity_multiplier=dynamic_runtime_params_slice.numerics.resistivity_mult,
           psi=core_profiles.psi,
           geo=geo,
-      )
+      ),
   )
 
   new_core_profiles = dataclasses.replace(

torax/pedestal_model/set_pped_tpedratio_nped.py

@@ -23,9 +23,9 @@
 from torax import array_typing
 from torax import constants
 from torax import interpolated_param
-from torax import physics
 from torax import state
 from torax.config import runtime_params_slice
+from torax.core_profiles import formulas
 from torax.geometry import geometry
 from torax.pedestal_model import pedestal_model
 from torax.pedestal_model import runtime_params as runtime_params_lib
@@ -130,7 +130,7 @@ def _call_implementation(
     Zeff_ped = Zeff[ped_idx]
     Zi_ped = Zi[ped_idx]
     Zimp_ped = Zimp[ped_idx]
-    dilution_factor_ped = physics.get_main_ion_dilution_factor(
+    dilution_factor_ped = formulas.get_main_ion_dilution_factor(
         Zi_ped, Zimp_ped, Zeff_ped
     )
     # Calculate ni and nimp.