Add Geometry to ToraxSimState.

PiperOrigin-RevId: 718842985

torax/orchestration/step_function.py

@@ -365,6 +365,7 @@ def step(
         post_processed_outputs=state.PostProcessedOutputs.zeros(geo_t_plus_dt),
         time_step_calculator_state=time_step_calculator_state,
         stepper_numeric_outputs=stepper_numeric_outputs,
+        geometry=geo_t_plus_dt,
     )
 
   def adaptive_step(

torax/sim.py

@@ -92,6 +92,7 @@ def get_initial_state(
           outer_stepper_iterations=0,
           inner_solver_iterations=0,
       ),
+      geometry=geo,
   )
 
 

torax/state.py

@@ -515,6 +515,7 @@ class ToraxSimState:
       running `run_simulation()`.
     post_processed_outputs: variables for output or intermediate observations
       for overarching workflows, calculated after each simulation step.
+    geometry: Geometry at this time step used for the simulation.
     time_step_calculator_state: the state of the TimeStepper.
     stepper_numeric_outputs: Numerical quantities related to the stepper.
   """
@@ -531,6 +532,9 @@ class ToraxSimState:
   # Post-processed outputs after a step.
   post_processed_outputs: PostProcessedOutputs
 
+  # Geometry used for the simulation.
+  geometry: geometry.Geometry
+
   # Other "side" states used for logging and feeding to other components of
   # TORAX.
   time_step_calculator_state: Any

torax/tests/output.py

@@ -107,6 +107,7 @@ def setUp(self):
             stepper_error_state=1,
             inner_solver_iterations=1,
         ),
+        geometry=self.geo,
     )
     sim_error = state.SimError.NO_ERROR
 

torax/tests/post_processing.py

@@ -96,6 +96,7 @@ def test_make_outputs(self):
             stepper_error_state=1,
             inner_solver_iterations=1,
         ),
+        geometry=self.geo,
     )
 
     updated_sim_state = post_processing.make_outputs(sim_state, self.geo)