Cleans up spectrum solver initialization

tardis-sn · Jul 8, 2024 · 5454a2b · 5454a2b
1 parent 62e25a8
commit 5454a2b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py
@@ -136,6 +136,7 @@ def __init__(
         show_convergence_plots,
         convergence_plots_kwargs,
         show_progress_bars,
+        spectrum_solver,
         integrator_settings,
     ):
         super(Simulation, self).__init__(
@@ -154,7 +155,7 @@ def __init__(
         self.luminosity_nu_start = luminosity_nu_start
         self.luminosity_nu_end = luminosity_nu_end
         self.luminosity_requested = luminosity_requested
-        self.spectrum_solver = None
+        self.spectrum_solver = spectrum_solver
         self.integrator_settings = integrator_settings
         self.show_progress_bars = show_progress_bars
         self.version = tardis.__version__
@@ -391,18 +392,10 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0):
             show_progress_bars=self.show_progress_bars,
         )
 
-        self.spectrum_solver = SpectrumSolver(
-            transport_state, self.transport.spectrum_frequency
-        )
-
-        self.spectrum_solver.integrator_settings = self.integrator_settings
-
-        self.spectrum_solver._montecarlo_virtual_luminosity.value[
-            :
-        ] = v_packets_energy_hist
-
-        self.spectrum_solver._integrator = FormalIntegrator(
-            self.simulation_state, self.plasma, self.transport
+        # Set up spectrum solver
+        self.spectrum_solver.transport_state = transport_state
+        self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = (
+            v_packets_energy_hist
         )
 
         output_energy = (
@@ -480,6 +473,12 @@ def run_final(self):
         )
         self.iterate(self.last_no_of_packets, self.no_of_virtual_packets)
 
+        # Set up spectrum solver integrator
+        self.spectrum_solver.integrator_settings = self.integrator_settings
+        self.spectrum_solver._integrator = FormalIntegrator(
+            self.simulation_state, self.plasma, self.transport
+        )
+
         self.reshape_plasma_state_store(self.iterations_executed)
         if hasattr(self, "convergence_plots"):
             self.convergence_plots.fetch_data(
@@ -754,6 +753,8 @@ def from_config(
             last_no_of_packets = config.montecarlo.no_of_packets
         last_no_of_packets = int(last_no_of_packets)
 
+        spectrum_solver = SpectrumSolver.from_config(config)
+
         return cls(
             iterations=config.montecarlo.iterations,
             simulation_state=simulation_state,
@@ -770,4 +771,5 @@ def from_config(
             convergence_plots_kwargs=convergence_plots_kwargs,
             show_progress_bars=show_progress_bars,
             integrator_settings=config.spectrum.integrated,
+            spectrum_solver=spectrum_solver,
         )
diff --git a/tardis/spectrum/base.py b/tardis/spectrum/base.py
@@ -179,7 +179,7 @@ def calculate_reabsorbed_luminosity(
         ].sum()
 
     @classmethod
-    def from_config(cls, config, v_packets_energy_hist):
+    def from_config(cls, config):
         spectrum_frequency = quantity_linspace(
             config.spectrum.stop.to("Hz", u.spectral()),
             config.spectrum.start.to("Hz", u.spectral()),