Fix formal integral benchmark

tardis-sn · Jul 30, 2024 · ee38232 · ee38232
1 parent b1f869b
commit ee38232
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/benchmarks/spectrum_formal_integral.py b/benchmarks/spectrum_formal_integral.py
@@ -8,7 +8,8 @@
 import tardis.spectrum.formal_integral as formal_integral
 from benchmarks.benchmark_base import BenchmarkBase
 
-config.THREADING_LAYER='workqueue'
+config.THREADING_LAYER = "workqueue"
+
 
 class BenchmarkTransportMontecarloFormalIntegral(BenchmarkBase):
     """
@@ -33,12 +34,20 @@ def time_intensity_black_body(self, parameters):
     # Benchmark for functions in FormalIntegrator class
     def time_FormalIntegrator_functions(self):
         FormalIntegrator = formal_integral.FormalIntegrator(
-            self.simulation_verysimple.simulation_state, self.simulation_verysimple.plasma, self.simulation_verysimple.transport
+            self.simulation_verysimple.simulation_state,
+            self.simulation_verysimple.plasma,
+            self.simulation_verysimple.transport,
+        )
+        FormalIntegrator.calculate_spectrum(
+            self.simulation_verysimple.spectrum_solver.spectrum_frequency_grid[
+                :-1
+            ]
         )
-        FormalIntegrator.calculate_spectrum(self.simulation_verysimple.spectrum_solver.spectrum_real_packets.frequency)
         FormalIntegrator.make_source_function()
         FormalIntegrator.generate_numba_objects()
         FormalIntegrator.formal_integral(
-            self.simulation_verysimple.spectrum_solver.spectrum_real_packets.frequency,
-            1000
+            self.simulation_verysimple.spectrum_solver.spectrum_frequency_grid[
+                :-1
+            ],
+            1000,
         )