Moved handling of the opacity solver over to the simulation base

tardis-sn · Aug 8, 2024 · 8f9afcd · 8f9afcd
1 parent cfff15f
commit 8f9afcd
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py
@@ -25,6 +25,7 @@
 from tardis.transport.montecarlo.estimators.continuum_radfield_properties import (
     MCContinuumPropertiesSolver,
 )
+from tardis.opacities.opacity_solver import OpacitySolver
 from tardis.util.base import is_notebook
 from tardis.visualization import ConvergencePlots
 
@@ -103,6 +104,7 @@ class Simulation(PlasmaStateStorerMixin, HDFWriterMixin):
     model : tardis.model.SimulationState
     plasma : tardis.plasma.BasePlasma
     transport : tardis.transport.montecarlo.MontecarloTransport
+    opacity : tardis.opacities.opacity_solver.OpacitySolver
     no_of_packets : int
     last_no_of_packets : int
     no_of_virtual_packets : int
@@ -130,6 +132,7 @@ def __init__(
         simulation_state,
         plasma,
         transport,
+        opacity,
         no_of_packets,
         no_of_virtual_packets,
         luminosity_nu_start,
@@ -152,6 +155,7 @@ def __init__(
         self.simulation_state = simulation_state
         self.plasma = plasma
         self.transport = transport
+        self.opacity = opacity
         self.no_of_packets = no_of_packets
         self.last_no_of_packets = last_no_of_packets
         self.no_of_virtual_packets = no_of_virtual_packets
@@ -439,8 +443,11 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0):
             f"\n\tStarting iteration {(self.iterations_executed + 1):d} of {self.iterations:d}"
         )
 
+        opacity_state = self.opacity.solve(self.plasma)
+
         transport_state = self.transport.initialize_transport_state(
             self.simulation_state,
+            opacity_state,
             self.plasma,
             no_of_packets,
             no_of_virtual_packets=no_of_virtual_packets,
@@ -695,6 +702,7 @@ def from_config(
         atom_data=None,
         plasma=None,
         transport=None,
+        opacity=None,
         **kwargs,
     ):
         """
@@ -752,6 +760,11 @@ def from_config(
                 packet_source=simulation_state.packet_source,
                 enable_virtual_packet_logging=virtual_packet_logging,
             )
+        if opacity is None:
+            opacity = OpacitySolver(
+                config.plasma.line_interaction_type,
+                config.plasma.disable_line_scattering,
+            )
 
         convergence_plots_config_options = [
             "plasma_plot_config",
@@ -791,6 +804,7 @@ def from_config(
             simulation_state=simulation_state,
             plasma=plasma,
             transport=transport,
+            opcacity=opacity,
             show_convergence_plots=show_convergence_plots,
             no_of_packets=int(config.montecarlo.no_of_packets),
             no_of_virtual_packets=int(config.montecarlo.no_of_virtual_packets),

diff --git a/tardis/transport/montecarlo/base.py b/tardis/transport/montecarlo/base.py
@@ -37,8 +37,6 @@
     update_iterations_pbar,
 )
 
-from tardis.opacities.opacity_solver import OpacitySolver
-
 logger = logging.getLogger(__name__)
 
 
@@ -96,14 +94,10 @@ def __init__(
         mc_tracker.DEBUG_MODE = debug_packets
         mc_tracker.BUFFER = logger_buffer
 
-        self.opacity_solver = OpacitySolver(
-            self.line_interaction_type,
-            self.montecarlo_configuration.DISABLE_LINE_SCATTERING,
-        )
-
     def initialize_transport_state(
         self,
         simulation_state,
+        opacity_state,
         plasma,
         no_of_packets,
         no_of_virtual_packets=0,
@@ -120,7 +114,6 @@ def initialize_transport_state(
 
         geometry_state = simulation_state.geometry.to_numba()
 
-        opacity_state = self.opacity_solver.solve(plasma)
         opacity_state_numba = opacity_state_to_numba(
             opacity_state, self.opacity_solver.line_interaction_type
         )