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/build_runtime_params.py

@@ -103,6 +103,7 @@ def __call__(
         numerics=self.numerics.build_runtime_params(t),
         neoclassical=self.neoclassical.build_runtime_params(),
         pedestal=self.pedestal.build_runtime_params(t),
+        pedestal_policy=self.pedestal.build_pedestal_policy_runtime_params(),
         mhd=self.mhd.build_runtime_params(t),
         time_step_calculator=self.time_step_calculator.build_runtime_params(),
     )

torax/_src/config/runtime_params_slice.py

@@ -48,6 +48,7 @@
 from torax._src.mhd import runtime_params as mhd_runtime_params
 from torax._src.neoclassical import runtime_params as neoclassical_params
 from torax._src.pedestal_model import runtime_params as pedestal_model_params
+from torax._src.pedestal_policy import runtime_params as pedestal_policy_runtime_params
 from torax._src.solver import runtime_params as solver_params
 from torax._src.sources import runtime_params as sources_params
 from torax._src.time_step_calculator import runtime_params as time_step_calculator_runtime_params
@@ -78,6 +79,7 @@ class RuntimeParams:
   neoclassical: neoclassical_params.RuntimeParams
   numerics: numerics.RuntimeParams
   pedestal: pedestal_model_params.RuntimeParams
+  pedestal_policy: pedestal_policy_runtime_params.PedestalPolicyRuntimeParams
   plasma_composition: plasma_composition.RuntimeParams
   profile_conditions: profile_conditions.RuntimeParams
   solver: solver_params.RuntimeParams

torax/_src/config/tests/build_runtime_params_test.py

@@ -95,19 +95,32 @@ 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)
+    pedestal_policy_rp = pedestal.build_pedestal_policy_runtime_params()
     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, runtime_params=pedestal_policy_rp
+        ).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)
+    # Note: pedestal_policy_rp does not depend on time for its structure
     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, runtime_params=pedestal_policy_rp
+        ).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,
     )
 
@@ -285,14 +292,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
@@ -352,7 +363,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
 
@@ -68,6 +69,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,
@@ -129,6 +131,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
@@ -160,6 +163,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,
   )
 
@@ -176,6 +180,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,
       )
@@ -194,6 +199,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:
@@ -215,6 +221,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
 
@@ -57,6 +58,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, ...],
@@ -98,6 +100,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.
@@ -124,6 +127,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,
   )
 
@@ -141,6 +145,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,
       )
@@ -158,6 +163,7 @@ def optimizer_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:
@@ -187,6 +193,7 @@ def optimizer_solve_block(
       init_x_new_vec=init_x_new_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/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
@@ -201,6 +202,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, ...],
@@ -220,6 +222,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
@@ -252,6 +255,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,
   )
@@ -290,6 +294,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, ...],
@@ -309,6 +314,7 @@ def theta_method_block_loss(
       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
@@ -326,6 +332,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,
@@ -349,6 +356,7 @@ def jaxopt_solver(
     init_x_new_vec: jax.Array,
     core_profiles_t_plus_dt: state.CoreProfiles,
     explicit_source_profiles: source_profiles.SourceProfiles,
+    pedestal_policy_state: pedestal_policy.PedestalPolicy,
     physics_models: physics_models_lib.PhysicsModels,
     coeffs_old: Block1DCoeffs,
     evolving_names: tuple[str, ...],
@@ -370,6 +378,7 @@ def jaxopt_solver(
       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
@@ -394,6 +403,7 @@ def jaxopt_solver(
       physics_models=physics_models,
       coeffs_old=coeffs_old,
       evolving_names=evolving_names,
+      pedestal_policy_state=pedestal_policy_state,
   )
   solver = jaxopt.LBFGS(fun=loss, maxiter=maxiter, tol=tol, implicit_diff=True)
   solver_output = solver.run(init_x_new_vec)

torax/_src/fvm/tests/calc_coeffs_test.py

@@ -79,13 +79,20 @@ 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,
+            runtime_params=runtime_params.pedestal_policy,
+        )
+    )
     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,
     )