Pipe through prev_core_profiles to updaters where needed

PiperOrigin-RevId: 876278049

Author: tamaranorman

torax/_src/core_profiles/updaters.py

@@ -53,6 +53,8 @@ def update_core_profiles_during_step(
     runtime_params: runtime_params_lib.RuntimeParams,
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
+    prev_core_profiles: state.CoreProfiles | None,
+    dt: array_typing.FloatScalar | None,
     evolving_names: tuple[str, ...],
 ) -> state.CoreProfiles:
   """Returns the new core profiles after updating the evolving variables.
@@ -67,8 +69,13 @@ def update_core_profiles_during_step(
     runtime_params: The runtime params slice.
     geo: Magnetic geometry.
     core_profiles: The old set of core plasma profiles.
+    prev_core_profiles: Core plasma profiles from the previous timestep if
+      available, used to update the energy state.
+    dt: The size of the last timestep, used to update the energy state.
     evolving_names: The names of the evolving variables.
   """
+  # Currently unused but will be used to update the energy state soon
+  del prev_core_profiles, dt
 
   updated_core_profiles = convertors.solver_x_tuple_to_core_profiles(
       x_new, evolving_names, core_profiles

torax/_src/fvm/calc_coeffs.py

@@ -17,6 +17,7 @@
 import functools
 import jax
 import jax.numpy as jnp
+from torax._src import array_typing
 from torax._src import constants
 from torax._src import physics_models as physics_models_lib
 from torax._src import state
@@ -62,6 +63,8 @@ def __call__(
       runtime_params: runtime_params_lib.RuntimeParams,
       geo: geometry.Geometry,
       core_profiles: state.CoreProfiles,
+      prev_core_profiles: state.CoreProfiles | None,
+      dt: array_typing.FloatScalar | None,
       x: tuple[cell_variable.CellVariable, ...],
       explicit_source_profiles: source_profiles_lib.SourceProfiles,
       allow_pereverzev: bool = False,
@@ -80,6 +83,9 @@ def __call__(
         state x.
       geo: The geometry of the system at this time step.
       core_profiles: The core profiles of the system at this time step.
+      prev_core_profiles: The core profiles of the system at the previous
+        time step.
+      dt: The time step size.
       x: The state with cell-grid values of the evolving variables.
       explicit_source_profiles: Precomputed explicit source profiles. These
         profiles were configured to always depend on state and parameters at
@@ -107,8 +113,10 @@ def __call__(
         x,
         runtime_params,
         geo,
-        core_profiles,
-        self.evolving_names,
+        core_profiles=core_profiles,
+        prev_core_profiles=prev_core_profiles,
+        dt=dt,
+        evolving_names=self.evolving_names,
     )
     if allow_pereverzev:
       use_pereverzev = runtime_params.solver.use_pereverzev

torax/_src/fvm/newton_raphson_solve_block.py

@@ -158,7 +158,9 @@ def newton_raphson_solve_block(
       runtime_params_t,
       geo_t,
       core_profiles_t,
-      x_old,
+      prev_core_profiles=None,
+      dt=None,
+      x=x_old,
       explicit_source_profiles=explicit_source_profiles,
       explicit_call=True,
   )
@@ -173,8 +175,10 @@ def newton_raphson_solve_block(
       coeffs_exp_linear = coeffs_callback(
           runtime_params_t,
           geo_t,
-          core_profiles_t,
-          x_old,
+          core_profiles=core_profiles_t,
+          prev_core_profiles=None,
+          dt=None,
+          x=x_old,
           explicit_source_profiles=explicit_source_profiles,
           allow_pereverzev=True,
           explicit_call=True,
@@ -190,6 +194,7 @@ def newton_raphson_solve_block(
           geo_t_plus_dt=geo_t_plus_dt,
           x_old=x_old,
           x_new_guess=x_new_guess,
+          core_profiles_t=core_profiles_t,
           core_profiles_t_plus_dt=core_profiles_t_plus_dt,
           coeffs_exp=coeffs_exp_linear,
           coeffs_callback=coeffs_callback,
@@ -212,6 +217,7 @@ def newton_raphson_solve_block(
       runtime_params_t_plus_dt=runtime_params_t_plus_dt,
       geo_t_plus_dt=geo_t_plus_dt,
       x_old=x_old,
+      core_profiles_t=core_profiles_t,
       core_profiles_t_plus_dt=core_profiles_t_plus_dt,
       physics_models=physics_models,
       explicit_source_profiles=explicit_source_profiles,

torax/_src/fvm/optimizer_solve_block.py

@@ -122,7 +122,9 @@ def optimizer_solve_block(
       runtime_params_t,
       geo_t,
       core_profiles_t,
-      x_old,
+      prev_core_profiles=None,
+      dt=None,
+      x=x_old,
       explicit_source_profiles=explicit_source_profiles,
       explicit_call=True,
   )
@@ -139,7 +141,9 @@ def optimizer_solve_block(
           runtime_params_t,
           geo_t,
           core_profiles_t,
-          x_old,
+          prev_core_profiles=None,
+          dt=None,
+          x=x_old,
           explicit_source_profiles=explicit_source_profiles,
           allow_pereverzev=True,
           explicit_call=True,
@@ -154,6 +158,7 @@ def optimizer_solve_block(
           geo_t_plus_dt=geo_t_plus_dt,
           x_old=x_old,
           x_new_guess=x_new_guess,
+          core_profiles_t=core_profiles_t,
           core_profiles_t_plus_dt=core_profiles_t_plus_dt,
           coeffs_exp=coeffs_exp_linear,
           coeffs_callback=coeffs_callback,
@@ -178,6 +183,7 @@ def optimizer_solve_block(
       geo_t_plus_dt=geo_t_plus_dt,
       x_old=x_old,
       init_x_new_vec=init_x_new_vec,
+      core_profiles_t=core_profiles_t,
       core_profiles_t_plus_dt=core_profiles_t_plus_dt,
       explicit_source_profiles=explicit_source_profiles,
       physics_models=physics_models,

torax/_src/fvm/residual_and_loss.py

@@ -199,6 +199,7 @@ def theta_method_block_residual(
     runtime_params_t_plus_dt: runtime_params_lib.RuntimeParams,
     geo_t_plus_dt: geometry.Geometry,
     x_old: tuple[cell_variable.CellVariable, ...],
+    core_profiles_t: state.CoreProfiles,
     core_profiles_t_plus_dt: state.CoreProfiles,
     explicit_source_profiles: source_profiles.SourceProfiles,
     physics_models: physics_models_lib.PhysicsModels,
@@ -214,6 +215,8 @@ def theta_method_block_residual(
     runtime_params_t_plus_dt: Runtime parameters for time t + dt.
     geo_t_plus_dt: The geometry at time t + dt.
     x_old: The starting x defined as a tuple of CellVariables.
+    core_profiles_t: Core plasma profiles which contain all available
+      prescribed quantities at the start of the time step.
     core_profiles_t_plus_dt: Core plasma profiles which contain all available
       prescribed quantities at the end of the time step. This includes evolving
       boundary conditions and prescribed time-dependent profiles that are not
@@ -244,7 +247,9 @@ def theta_method_block_residual(
       runtime_params_t_plus_dt,
       geo_t_plus_dt,
       core_profiles_t_plus_dt,
-      evolving_names,
+      prev_core_profiles=core_profiles_t,
+      dt=dt,
+      evolving_names=evolving_names,
   )
   coeffs_new = calc_coeffs.calc_coeffs(
       runtime_params=runtime_params_t_plus_dt,
@@ -288,6 +293,7 @@ def theta_method_block_loss(
     runtime_params_t_plus_dt: runtime_params_lib.RuntimeParams,
     geo_t_plus_dt: geometry.Geometry,
     x_old: tuple[cell_variable.CellVariable, ...],
+    core_profiles_t: state.CoreProfiles,
     core_profiles_t_plus_dt: state.CoreProfiles,
     explicit_source_profiles: source_profiles.SourceProfiles,
     physics_models: physics_models_lib.PhysicsModels,
@@ -303,6 +309,7 @@ def theta_method_block_loss(
     runtime_params_t_plus_dt: Runtime parameters for time t + dt.
     geo_t_plus_dt: geometry object at time t + dt.
     x_old: The starting x defined as a tuple of CellVariables.
+    core_profiles_t: Core profiles from the previous time step.
     core_profiles_t_plus_dt: Core plasma profiles which contain all available
       prescribed quantities at the end of the time step. This includes evolving
       boundary conditions and prescribed time-dependent profiles that are not
@@ -324,6 +331,7 @@ def theta_method_block_loss(
       geo_t_plus_dt=geo_t_plus_dt,
       x_old=x_old,
       x_new_guess_vec=x_new_guess_vec,
+      core_profiles_t=core_profiles_t,
       core_profiles_t_plus_dt=core_profiles_t_plus_dt,
       explicit_source_profiles=explicit_source_profiles,
       physics_models=physics_models,
@@ -347,6 +355,7 @@ def jaxopt_solver(
     geo_t_plus_dt: geometry.Geometry,
     x_old: tuple[cell_variable.CellVariable, ...],
     init_x_new_vec: jax.Array,
+    core_profiles_t: state.CoreProfiles,
     core_profiles_t_plus_dt: state.CoreProfiles,
     explicit_source_profiles: source_profiles.SourceProfiles,
     physics_models: physics_models_lib.PhysicsModels,
@@ -364,6 +373,7 @@ def jaxopt_solver(
     x_old: The starting x defined as a tuple of CellVariables.
     init_x_new_vec: Flattened array of initial guess of x_new for all evolving
       core profiles.
+    core_profiles_t: Core profiles from the previous time step.
     core_profiles_t_plus_dt: Core plasma profiles which contain all available
       prescribed quantities at the end of the time step. This includes evolving
       boundary conditions and prescribed time-dependent profiles that are not
@@ -389,6 +399,7 @@ def jaxopt_solver(
       runtime_params_t_plus_dt=runtime_params_t_plus_dt,
       geo_t_plus_dt=geo_t_plus_dt,
       x_old=x_old,
+      core_profiles_t=core_profiles_t,
       core_profiles_t_plus_dt=core_profiles_t_plus_dt,
       explicit_source_profiles=explicit_source_profiles,
       physics_models=physics_models,

torax/_src/fvm/tests/fvm_test.py

@@ -286,6 +286,7 @@ def test_nonlinear_solve_block_loss_minimum(
           x_old=x_old,
           x_new_guess_vec=jnp.concatenate([var.value for var in x_new]),
           core_profiles_t_plus_dt=core_profiles,
+          core_profiles_t=core_profiles,
           physics_models=physics_models,
           explicit_source_profiles=explicit_source_profiles,
           coeffs_old=coeffs,
@@ -299,6 +300,7 @@ def test_nonlinear_solve_block_loss_minimum(
           x_new_guess_vec=jnp.concatenate([var.value for var in x_new]),
           x_old=x_old,
           core_profiles_t_plus_dt=core_profiles,
+          core_profiles_t=core_profiles,
           physics_models=physics_models,
           explicit_source_profiles=explicit_source_profiles,
           coeffs_old=coeffs,
@@ -512,6 +514,7 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
           x_old=(x_0,),
           x_new_guess_vec=x_0.value,
           core_profiles_t_plus_dt=core_profiles_t_plus_dt,
+          core_profiles_t=initial_core_profiles,
           physics_models=physics_models,
           explicit_source_profiles=explicit_source_profiles,
           coeffs_old=coeffs_old,
@@ -535,6 +538,7 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
                   x_0, right_face_constraint=final_right_boundary
               ),
           ),
+          core_profiles_t=initial_core_profiles,
           evolving_names=evolving_names,
           physics_models=physics_models,
           explicit_source_profiles=explicit_source_profiles,
@@ -553,6 +557,7 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
                   x_0, right_face_constraint=final_right_boundary
               ),
           ),
+          core_profiles_t=initial_core_profiles,
           x_new_guess_vec=x_0.value,
           physics_models=physics_models,
           explicit_source_profiles=explicit_source_profiles,

torax/_src/solver/linear_theta_method.py

@@ -74,7 +74,9 @@ def _x_new(
         runtime_params_t,
         geo_t,
         core_profiles_t,
-        x_old,
+        prev_core_profiles=None,
+        dt=None,
+        x=x_old,
         explicit_source_profiles=explicit_source_profiles,
         allow_pereverzev=True,
         explicit_call=True,
@@ -90,6 +92,7 @@ def _x_new(
         geo_t_plus_dt=geo_t_plus_dt,
         x_old=x_old,
         x_new_guess=x_new_guess,
+        core_profiles_t=core_profiles_t,
         core_profiles_t_plus_dt=core_profiles_t_plus_dt,
         coeffs_exp=coeffs_exp,
         coeffs_callback=coeffs_callback,

torax/_src/solver/predictor_corrector_method.py

@@ -44,6 +44,7 @@ def predictor_corrector_method(
     geo_t_plus_dt: geometry.Geometry,
     x_old: tuple[cell_variable.CellVariable, ...],
     x_new_guess: tuple[cell_variable.CellVariable, ...],
+    core_profiles_t: state.CoreProfiles,
     core_profiles_t_plus_dt: state.CoreProfiles,
     coeffs_exp: block_1d_coeffs.Block1DCoeffs,
     explicit_source_profiles: source_profiles.SourceProfiles,
@@ -60,6 +61,7 @@ def predictor_corrector_method(
       time t.
     x_new_guess: Tuple of CellVariables corresponding to the initial guess for
       the next time step.
+    core_profiles_t: Core profiles at the current time step.
     core_profiles_t_plus_dt: Core profiles at the next time step.
     coeffs_exp: Block1DCoeffs PDE coefficients at beginning of timestep.
     explicit_source_profiles: Precomputed explicit source profiles. These
@@ -82,7 +84,9 @@ def loop_body(x_new_guess):
         runtime_params_t_plus_dt,
         geo_t_plus_dt,
         core_profiles_t_plus_dt,
-        x_new_guess,
+        prev_core_profiles=core_profiles_t,
+        dt=dt,
+        x=x_new_guess,
         explicit_source_profiles=explicit_source_profiles,
         allow_pereverzev=True,
     )