Introduce PedestalPolicy.

This is a step toward having full Martin scaling, where there will be a
Martin scaling PedestalPolicy that ramps up and ramps down the pedestal
height. This PR is just a step that defines the new stateful interface
required for doing so and transitions the existing `set_pedestal`
functionality to being implemented with the new PedestalPolicy class.

PiperOrigin-RevId: 810477225

Author: goodfeli

torax/_src/config/tests/build_runtime_params_test.py

@@ -95,19 +95,20 @@ def test_pedestal_is_time_dependent(self):
             set_pedestal={0.0: True, 1.0: False},
         )
     )
+    pedestal_policy = pedestal.build_pedestal_model().pedestal_policy
     # Check at time 0.
 
     pedestal_params = pedestal.build_runtime_params(t=0.0)
     assert isinstance(pedestal_params, set_tped_nped.RuntimeParams)
-    np.testing.assert_allclose(pedestal_params.set_pedestal, True)
+    np.testing.assert_allclose(pedestal_policy.initial_state(t=0.0).use_pedestal, True)
     np.testing.assert_allclose(pedestal_params.T_i_ped, 0.0)
     np.testing.assert_allclose(pedestal_params.T_e_ped, 1.0)
     np.testing.assert_allclose(pedestal_params.n_e_ped, 2.0e20)
     np.testing.assert_allclose(pedestal_params.rho_norm_ped_top, 3.0)
     # And check after the time limit.
     pedestal_params = pedestal.build_runtime_params(t=1.0)
     assert isinstance(pedestal_params, set_tped_nped.RuntimeParams)
-    np.testing.assert_allclose(pedestal_params.set_pedestal, False)
+    np.testing.assert_allclose(pedestal_policy.initial_state(t=1.0).use_pedestal, False)
     np.testing.assert_allclose(pedestal_params.T_i_ped, 1.0)
     np.testing.assert_allclose(pedestal_params.T_e_ped, 2.0)
     np.testing.assert_allclose(pedestal_params.n_e_ped, 3.0e20)

torax/_src/fvm/calc_coeffs.py

@@ -27,6 +27,7 @@
 from torax._src.fvm import cell_variable
 from torax._src.geometry import geometry
 from torax._src.pedestal_model import pedestal_model as pedestal_model_lib
+from torax._src.pedestal_policy import pedestal_policy
 from torax._src.sources import source_profile_builders
 from torax._src.sources import source_profiles as source_profiles_lib
 import typing_extensions
@@ -63,6 +64,7 @@ def __call__(
       core_profiles: state.CoreProfiles,
       x: tuple[cell_variable.CellVariable, ...],
       explicit_source_profiles: source_profiles_lib.SourceProfiles,
+      pedestal_policy_state: pedestal_policy.PedestalPolicyState,
       allow_pereverzev: bool = False,
       # Checks if reduced calc_coeffs for explicit terms when theta_implicit=1
       # should be called
@@ -86,6 +88,7 @@ def __call__(
         not recalculated at time t+plus_dt with updated state during the solver
         iterations. For sources that are implicit, their explicit profiles are
         set to all zeros.
+      pedestal_policy_state: State held by the pedestal policy.
       allow_pereverzev: If True, then the coeffs are being called within a
         linear solver. Thus could be either the use_predictor_corrector solver
         or as part of calculating the initial guess for the nonlinear solver. In
@@ -121,6 +124,7 @@ def __call__(
         explicit_source_profiles=explicit_source_profiles,
         physics_models=self.physics_models,
         evolving_names=self.evolving_names,
+        pedestal_policy_state=pedestal_policy_state,
         use_pereverzev=use_pereverzev,
         explicit_call=explicit_call,
     )
@@ -219,6 +223,7 @@ def calc_coeffs(
     explicit_source_profiles: source_profiles_lib.SourceProfiles,
     physics_models: physics_models_lib.PhysicsModels,
     evolving_names: tuple[str, ...],
+    pedestal_policy_state: pedestal_policy.PedestalPolicyState,
     use_pereverzev: bool = False,
     explicit_call: bool = False,
 ) -> block_1d_coeffs.Block1DCoeffs:
@@ -241,6 +246,7 @@ def calc_coeffs(
     physics_models: The physics models to use for the simulation.
     evolving_names: The names of the evolving variables in the order that their
       coefficients should be written to `coeffs`.
+    pedestal_policy_state: State held by the pedestal policy.
     use_pereverzev: Toggle whether to calculate Pereverzev terms
     explicit_call: If True, indicates that calc_coeffs is being called for the
       explicit component of the PDE. Then calculates a reduced Block1DCoeffs if
@@ -267,6 +273,7 @@ def calc_coeffs(
         explicit_source_profiles=explicit_source_profiles,
         physics_models=physics_models,
         evolving_names=evolving_names,
+        pedestal_policy_state=pedestal_policy_state,
         use_pereverzev=use_pereverzev,
     )
 
@@ -284,14 +291,18 @@ def _calc_coeffs_full(
     explicit_source_profiles: source_profiles_lib.SourceProfiles,
     physics_models: physics_models_lib.PhysicsModels,
     evolving_names: tuple[str, ...],
+    pedestal_policy_state: pedestal_policy.PedestalPolicyState,
     use_pereverzev: bool = False,
 ) -> block_1d_coeffs.Block1DCoeffs:
   """See `calc_coeffs` for details."""
 
   consts = constants.CONSTANTS
 
   pedestal_model_output = physics_models.pedestal_model(
-      runtime_params, geo, core_profiles
+      runtime_params,
+      geo,
+      core_profiles,
+      pedestal_policy_state,
   )
 
   # Boolean mask for enforcing internal temperature boundary conditions to
@@ -351,7 +362,11 @@ def _calc_coeffs_full(
 
   # Diffusion term coefficients
   turbulent_transport = physics_models.transport_model(
-      runtime_params, geo, core_profiles, pedestal_model_output
+      runtime_params,
+      geo,
+      core_profiles,
+      pedestal_policy_state,
+      pedestal_model_output,
   )
   neoclassical_transport = physics_models.neoclassical_models.transport(
       runtime_params, geo, core_profiles

torax/_src/fvm/newton_raphson_solve_block.py

@@ -35,6 +35,7 @@
 from torax._src.fvm import jax_root_finding
 from torax._src.fvm import residual_and_loss
 from torax._src.geometry import geometry
+from torax._src.pedestal_policy import pedestal_policy
 from torax._src.solver import predictor_corrector_method
 from torax._src.sources import source_profiles
 
@@ -67,6 +68,7 @@ def newton_raphson_solve_block(
     physics_models: physics_models_lib.PhysicsModels,
     coeffs_callback: calc_coeffs.CoeffsCallback,
     evolving_names: tuple[str, ...],
+    pedestal_policy_state: pedestal_policy.PedestalPolicyState,
     initial_guess_mode: enums.InitialGuessMode,
     maxiter: int,
     tol: float,
@@ -128,6 +130,7 @@ def newton_raphson_solve_block(
       core_profiles. Repeatedly called by the iterative optimizer.
     evolving_names: The names of variables within the core profiles that should
       evolve.
+    pedestal_policy_state: State variables held by the pedestal policy.
     initial_guess_mode: chooses the initial_guess for the iterative method,
       either x_old or linear step. When taking the linear step, it is also
       recommended to use Pereverzev-Corrigan terms if the transport coefficients
@@ -159,6 +162,7 @@ def newton_raphson_solve_block(
       core_profiles_t,
       x_old,
       explicit_source_profiles=explicit_source_profiles,
+      pedestal_policy_state=pedestal_policy_state,
       explicit_call=True,
   )
 
@@ -175,6 +179,7 @@ def newton_raphson_solve_block(
           core_profiles_t,
           x_old,
           explicit_source_profiles=explicit_source_profiles,
+          pedestal_policy_state=pedestal_policy_state,
           allow_pereverzev=True,
           explicit_call=True,
       )
@@ -193,6 +198,7 @@ def newton_raphson_solve_block(
           coeffs_exp=coeffs_exp_linear,
           coeffs_callback=coeffs_callback,
           explicit_source_profiles=explicit_source_profiles,
+          pedestal_policy_state=pedestal_policy_state,
       )
       init_x_new_vec = fvm_conversions.cell_variable_tuple_to_vec(init_x_new)
     case enums.InitialGuessMode.X_OLD:
@@ -214,6 +220,7 @@ def newton_raphson_solve_block(
       core_profiles_t_plus_dt=core_profiles_t_plus_dt,
       physics_models=physics_models,
       explicit_source_profiles=explicit_source_profiles,
+      pedestal_policy_state=pedestal_policy_state,
       coeffs_old=coeffs_old,
       evolving_names=evolving_names,
   )

torax/_src/fvm/optimizer_solve_block.py

@@ -32,6 +32,7 @@
 from torax._src.fvm import fvm_conversions
 from torax._src.fvm import residual_and_loss
 from torax._src.geometry import geometry
+from torax._src.pedestal_policy import pedestal_policy
 from torax._src.solver import predictor_corrector_method
 from torax._src.sources import source_profiles
 
@@ -56,6 +57,7 @@ def optimizer_solve_block(
     core_profiles_t: state.CoreProfiles,
     core_profiles_t_plus_dt: state.CoreProfiles,
     explicit_source_profiles: source_profiles.SourceProfiles,
+    pedestal_policy_state: pedestal_policy.PedestalPolicyState,
     physics_models: physics_models_lib.PhysicsModels,
     coeffs_callback: calc_coeffs.CoeffsCallback,
     evolving_names: tuple[str, ...],
@@ -97,6 +99,7 @@ def optimizer_solve_block(
       being evolved by the PDE system.
     explicit_source_profiles: Pre-calculated sources implemented as explicit
       sources in the PDE.
+    pedestal_policy_state: State variables held by the pedestal policy.
     physics_models: Physics models used for the calculations.
     coeffs_callback: Calculates diffusion, convection etc. coefficients given a
       core_profiles. Repeatedly called by the iterative optimizer.
@@ -123,6 +126,7 @@ def optimizer_solve_block(
       core_profiles_t,
       x_old,
       explicit_source_profiles=explicit_source_profiles,
+      pedestal_policy_state=pedestal_policy_state,
       explicit_call=True,
   )
 
@@ -140,6 +144,7 @@ def optimizer_solve_block(
           core_profiles_t,
           x_old,
           explicit_source_profiles=explicit_source_profiles,
+          pedestal_policy_state=pedestal_policy_state,
           allow_pereverzev=True,
           explicit_call=True,
       )

torax/_src/fvm/residual_and_loss.py

@@ -38,6 +38,7 @@
 from torax._src.fvm import discrete_system
 from torax._src.fvm import fvm_conversions
 from torax._src.geometry import geometry
+from torax._src.pedestal_policy import pedestal_policy
 from torax._src.sources import source_profiles
 
 Block1DCoeffs: TypeAlias = block_1d_coeffs.Block1DCoeffs
@@ -200,6 +201,7 @@ def theta_method_block_residual(
     x_old: tuple[cell_variable.CellVariable, ...],
     core_profiles_t_plus_dt: state.CoreProfiles,
     explicit_source_profiles: source_profiles.SourceProfiles,
+    pedestal_policy_state: pedestal_policy.PedestalPolicyState,
     physics_models: physics_models_lib.PhysicsModels,
     coeffs_old: Block1DCoeffs,
     evolving_names: tuple[str, ...],
@@ -219,6 +221,7 @@ def theta_method_block_residual(
       being evolved by the PDE system.
     explicit_source_profiles: Pre-calculated sources implemented as explicit
       sources in the PDE.
+    pedestal_policy_state: State variables held by the pedestal policy.
     physics_models: Physics models used for the calculations.
     coeffs_old: The coefficients calculated at x_old.
     evolving_names: The names of variables within the core profiles that should
@@ -251,6 +254,7 @@ def theta_method_block_residual(
       core_profiles=core_profiles_t_plus_dt,
       explicit_source_profiles=explicit_source_profiles,
       physics_models=physics_models,
+      pedestal_policy_state=pedestal_policy_state,
       evolving_names=evolving_names,
       use_pereverzev=False,
   )
@@ -288,6 +292,7 @@ def theta_method_block_loss(
     x_old: tuple[cell_variable.CellVariable, ...],
     core_profiles_t_plus_dt: state.CoreProfiles,
     explicit_source_profiles: source_profiles.SourceProfiles,
+    pedestal_policy_state: pedestal_policy.PedestalPolicyState,
     physics_models: physics_models_lib.PhysicsModels,
     coeffs_old: Block1DCoeffs,
     evolving_names: tuple[str, ...],
@@ -324,6 +329,7 @@ def theta_method_block_loss(
       x_new_guess_vec=x_new_guess_vec,
       core_profiles_t_plus_dt=core_profiles_t_plus_dt,
       explicit_source_profiles=explicit_source_profiles,
+      pedestal_policy_state=pedestal_policy_state,
       physics_models=physics_models,
       coeffs_old=coeffs_old,
       evolving_names=evolving_names,

torax/_src/fvm/tests/calc_coeffs_test.py

@@ -20,6 +20,7 @@
 from torax._src.core_profiles import initialization
 from torax._src.fvm import calc_coeffs
 from torax._src.sources import source_profile_builders
+from torax._src.pedestal_policy import pedestal_policy as pedestal_policy_lib
 from torax._src.test_utils import default_sources
 from torax._src.torax_pydantic import model_config
 
@@ -79,13 +80,17 @@ def test_calc_coeffs_smoke_test(
         neoclassical_models=physics_models.neoclassical_models,
         explicit=True,
     )
+    pedestal_policy_state = physics_models.pedestal_model.pedestal_policy.initial_state(
+        t=torax_config.numerics.t_initial
+    )
     calc_coeffs.calc_coeffs(
         runtime_params=runtime_params,
         geo=geo,
         core_profiles=core_profiles,
         physics_models=physics_models,
         explicit_source_profiles=explicit_source_profiles,
         evolving_names=evolving_names,
+        pedestal_policy_state=pedestal_policy_state,
         use_pereverzev=False,
     )
 

torax/_src/fvm/tests/fvm_test.py

@@ -26,6 +26,7 @@
 from torax._src.fvm import cell_variable
 from torax._src.fvm import implicit_solve_block
 from torax._src.fvm import residual_and_loss
+from torax._src.pedestal_policy import pedestal_policy as pedestal_policy_lib
 from torax._src.sources import runtime_params as source_runtime_params
 from torax._src.sources import source_profile_builders
 from torax._src.test_utils import default_sources
@@ -247,13 +248,18 @@ def test_nonlinear_solve_block_loss_minimum(
         core_profiles=core_profiles,
         explicit=True,
     )
+    pedestal_policy = physics_models.pedestal_model.pedestal_policy
+    pedestal_policy_state = pedestal_policy.initial_state(
+        t=torax_config.numerics.t_initial
+    )
     coeffs = calc_coeffs.calc_coeffs(
         runtime_params=runtime_params,
         geo=geo,
         core_profiles=core_profiles,
         physics_models=physics_models,
         explicit_source_profiles=explicit_source_profiles,
         evolving_names=evolving_names,
+        pedestal_policy_state=pedestal_policy_state,
         use_pereverzev=False,
     )
     # dt well under the explicit stability limit for dx=1 and chi=1
@@ -287,19 +293,23 @@ def test_nonlinear_solve_block_loss_minimum(
           explicit_source_profiles=explicit_source_profiles,
           coeffs_old=coeffs,
           evolving_names=evolving_names,
+          pedestal_policy_state=pedestal_policy_state,
       )
 
       residual = residual_and_loss.theta_method_block_residual(
           dt=dt,
           runtime_params_t_plus_dt=runtime_params,
           geo_t_plus_dt=geo,
+
+
           x_new_guess_vec=jnp.concatenate([var.value for var in x_new]),
           x_old=x_old,
           core_profiles_t_plus_dt=core_profiles,
           physics_models=physics_models,
           explicit_source_profiles=explicit_source_profiles,
           coeffs_old=coeffs,
           evolving_names=evolving_names,
+          pedestal_policy_state=pedestal_policy_state,
       )
 
       np.testing.assert_allclose(loss, 0.0, atol=1e-7)
@@ -353,13 +363,18 @@ def test_implicit_solve_block_uses_updated_boundary_conditions(self):
     dt = jnp.array(1.0)
     evolving_names = tuple(['T_i'])
 
+    pedestal_policy = physics_models.pedestal_model.pedestal_policy
+    pedestal_policy_state = pedestal_policy.initial_state(
+        t=torax_config.numerics.t_initial
+    )
     coeffs = calc_coeffs.calc_coeffs(
         runtime_params=runtime_params,
         geo=geo,
         core_profiles=initial_core_profiles,
         physics_models=physics_models,
         explicit_source_profiles=explicit_source_profiles,
         evolving_names=evolving_names,
+        pedestal_policy_state=pedestal_policy_state,
         use_pereverzev=False,
     )
     initial_right_boundary = jnp.array(0.0)
@@ -470,13 +485,18 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
     dt = jnp.array(1.0)
     evolving_names = tuple(['T_i'])
 
+    pedestal_policy = physics_models.pedestal_model.pedestal_policy
+    pedestal_policy_state = pedestal_policy.initial_state(
+        t=torax_config.numerics.t_initial
+    )
     coeffs_old = calc_coeffs.calc_coeffs(
         runtime_params=runtime_params_theta05,
         geo=geo,
         core_profiles=initial_core_profiles,
         physics_models=physics_models,
         explicit_source_profiles=explicit_source_profiles,
         evolving_names=evolving_names,
+        pedestal_policy_state=pedestal_policy_state,
         use_pereverzev=False,
     )
 
@@ -513,6 +533,7 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
           explicit_source_profiles=explicit_source_profiles,
           coeffs_old=coeffs_old,
           evolving_names=evolving_names,
+          pedestal_policy_state=pedestal_policy_state,
       )
       np.testing.assert_allclose(residual, 0.0)
     with self.subTest('updated_boundary_conditions'):
@@ -536,6 +557,7 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
           physics_models=physics_models,
           explicit_source_profiles=explicit_source_profiles,
           coeffs_old=coeffs_old,
+          pedestal_policy_state=pedestal_policy_state,
       )
       np.testing.assert_allclose(residual, 0.0)
       # But when theta_implicit > 0, the residual should be non-zero.
@@ -555,6 +577,7 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
           explicit_source_profiles=explicit_source_profiles,
           coeffs_old=coeffs_old,
           evolving_names=evolving_names,
+          pedestal_policy_state=pedestal_policy_state,
       )
       self.assertGreater(jnp.abs(jnp.sum(residual)), 0.0)