Restructure/transport solver v1 (tardis-sn#3218)

* Co-authored-by: Joshua Shields <[email protected]>
Co-authored-by: Jack O'Brien <[email protected]>
Co-authored-by: Jing Lu <[email protected]>

Refactors transport initialization to use geometry state directly

Simplifies the transport state initialization by passing the geometry state object directly instead of the full simulation state. This change reduces coupling between transport and simulation components and clarifies the specific data dependencies needed for transport initialization.

Updates method signatures and variable names throughout the transport system to reflect this more focused interface.

* Refactor transport state initialization in Simulation and MonteCarloTransportSolver classes

Co-authored-by: Joshua Shields <[email protected]>
Co-authored-by: Jack O'Brien <[email protected]>
Co-authored-by: Jing Lu <[email protected]>

---------

Co-authored-by: Joshua Shields <[email protected]>
Co-authored-by: Jack O'Brien <[email protected]>
Co-authored-by: Jing Lu <[email protected]>

Author: 4 people

tardis/simulation/base.py

@@ -437,23 +437,21 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0):
                     self.opacity_state.beta_sobolev,
                 )
 
-        transport_state = self.transport.initialize_transport_state(
-            self.simulation_state,
+        v_packets_energy_hist = self.transport.run(
+            self.simulation_state.geometry,
             self.opacity_state,
             macro_atom_state,
             self.plasma,
             no_of_packets,
             no_of_virtual_packets=no_of_virtual_packets,
             iteration=self.iterations_executed,
-        )
-
-        v_packets_energy_hist = self.transport.run(
-            transport_state,
-            iteration=self.iterations_executed,
             total_iterations=self.iterations,
             show_progress_bars=self.show_progress_bars,
         )
 
+        # Get transport state for further processing
+        transport_state = self.transport.transport_state
+
         output_energy = self.transport.transport_state.packet_collection.output_energies
         if np.sum(output_energy < 0) == len(output_energy):
             logger.critical("No r-packet escaped through the outer boundary.")

tardis/transport/montecarlo/base.py

@@ -97,7 +97,7 @@ def __init__(
 
     def initialize_transport_state(
         self,
-        simulation_state,
+        geometry_state,
         opacity_state,
         macro_atom_state,
         plasma,
@@ -114,13 +114,13 @@ def initialize_transport_state(
             no_of_packets, seed_offset=iteration
         )
 
-        geometry_state = simulation_state.geometry.to_numba()
+        geometry_state_numba = geometry_state.to_numba()
         opacity_state_numba = opacity_state.to_numba(
             macro_atom_state,
             self.line_interaction_type,
         )
         opacity_state_numba = opacity_state_numba[
-            simulation_state.geometry.v_inner_boundary_index : simulation_state.geometry.v_outer_boundary_index
+            geometry_state.v_inner_boundary_index : geometry_state.v_outer_boundary_index
         ]
 
         estimators = initialize_estimator_statistics(
@@ -130,9 +130,9 @@ def initialize_transport_state(
         transport_state = MonteCarloTransportState(
             packet_collection,
             estimators,
-            geometry_state=geometry_state,
+            geometry_state=geometry_state_numba,
             opacity_state=opacity_state_numba,
-            time_explosion=simulation_state.time_explosion,
+            time_explosion=geometry_state.time_explosion,
         )
 
         transport_state.enable_full_relativity = (
@@ -147,7 +147,12 @@ def initialize_transport_state(
 
     def run(
         self,
-        transport_state,
+        geometry_state,
+        opacity_state,
+        macro_atom_state,
+        plasma,
+        no_of_packets,
+        no_of_virtual_packets=0,
         iteration=0,
         total_iterations=0,
         show_progress_bars=True,
@@ -157,18 +162,43 @@ def run(
 
         Parameters
         ----------
-        model : tardis.model.SimulationState
+        geometry_state : tardis.model.geometry.Geometry
+            The geometry state of the simulation
+        opacity_state : tardis.opacities.opacity_state.OpacityState
+            The opacity state
+        macro_atom_state : tardis.opacities.macro_atom.macroatom_state.LegacyMacroAtomState
+            The macro atom state
         plasma : tardis.plasma.BasePlasma
+            The plasma state
         no_of_packets : int
+            Number of packets to run
         no_of_virtual_packets : int
+            Number of virtual packets
+        iteration : int
+            Current iteration number
         total_iterations : int
             The total number of iterations in the simulation.
+        show_progress_bars : bool
+            Whether to show progress bars
 
         Returns
         -------
-        None
+        v_packets_energy_hist : numpy.ndarray
+            Virtual packet energy histogram
         """
         set_num_threads(self.nthreads)
+
+        # Initialize transport state
+        transport_state = self.initialize_transport_state(
+            geometry_state,
+            opacity_state,
+            macro_atom_state,
+            plasma,
+            no_of_packets,
+            no_of_virtual_packets=no_of_virtual_packets,
+            iteration=iteration,
+        )
+
         self.transport_state = transport_state
 
         number_of_vpackets = self.montecarlo_configuration.NUMBER_OF_VPACKETS

tardis/workflows/simple_tardis_workflow.py

@@ -403,7 +403,7 @@ def solve_montecarlo(
         macro_atom_state = opacity_states["macro_atom_state"]
 
         self.transport_state = self.transport_solver.initialize_transport_state(
-            self.simulation_state,
+            self.simulation_state.geometry,
             opacity_state,
             macro_atom_state,
             self.plasma_solver,