Refactor and add more benchmarks for montecarlo

.mailmap

@@ -12,6 +12,7 @@ Alice Harpole <[email protected]>
 Alice Harpole <[email protected]> Alice Harpole <[email protected]>
 
 Aman Kumar <[email protected]>
+Asish Kumar <[email protected]>
 
 Andreas Flörs <[email protected]>
 Andreas Flörs <[email protected]> Andreas Flörs <[email protected]>

benchmarks/benchmark_base.py

@@ -5,25 +5,26 @@
 
 import astropy.units as u
 import numpy as np
-import pandas as pd
 from numba import njit
 
 from benchmarks.util.nlte import NLTE
+from tardis import run_tardis
 from tardis.io.atom_data import AtomData
 from tardis.io.configuration import config_reader
 from tardis.io.configuration.config_reader import Configuration
 from tardis.io.util import YAMLLoader, yaml_load_file
 from tardis.model import SimulationState
+from tardis.model.geometry.radial1d import NumbaRadial1DGeometry
 from tardis.simulation import Simulation
 from tardis.tests.fixtures.atom_data import DEFAULT_ATOM_DATA_UUID
 from tardis.tests.fixtures.regression_data import RegressionData
-from tardis.transport.montecarlo import RPacket
-from tardis.transport.montecarlo.numba_interface import (
-    opacity_state_initialize,
-)
+from tardis.transport.montecarlo import RPacket, montecarlo_configuration
+from tardis.transport.montecarlo.estimators import radfield_mc_estimators
+from tardis.transport.montecarlo.numba_interface import opacity_state_initialize
 from tardis.transport.montecarlo.packet_collections import (
     VPacketCollection,
 )
+from tardis.transport.montecarlo.packet_trackers import RPacketTracker
 
 
 class BenchmarkBase:
@@ -350,13 +351,66 @@ def generate_reference(self):
             return option
 
     @property
-    def tardis_ref_data(self):
-        # TODO: This function is not working in the benchmarks.
-        if self.generate_reference:
-            mode = "w"
-        else:
-            mode = "r"
-        with pd.HDFStore(
-            f"{self.tardis_ref_path}/unit_test_data.h5", mode=mode
-        ) as store:
-            yield store
+    def verysimple_radfield_mc_estimators(self):
+        plasma = self.nb_simulation_verysimple.plasma
+        return radfield_mc_estimators.initialize_estimator_statistics(
+            plasma.tau_sobolevs.shape, plasma.gamma.shape
+        )
+
+    @property
+    def montecarlo_configuration(self):
+        return montecarlo_configuration.MonteCarloConfiguration()
+
+    @property
+    def rpacket_tracker(self): 
+        return RPacketTracker(0)
+
+    @property
+    def transport_state(self):
+        return self.nb_simulation_verysimple.transport.transport_state
+
+    @property
+    def simulation_rpacket_tracking_enabled(self):
+        config_verysimple = self.config_verysimple
+        config_verysimple.montecarlo.iterations = 3
+        config_verysimple.montecarlo.no_of_packets = 4000
+        config_verysimple.montecarlo.last_no_of_packets = -1
+        config_verysimple.spectrum.virtual.virtual_packet_logging = True
+        config_verysimple.montecarlo.no_of_virtual_packets = 1
+        config_verysimple.montecarlo.tracking.track_rpacket = True
+        config_verysimple.spectrum.num = 2000
+        atomic_data = deepcopy(self.atomic_dataset)
+        sim = run_tardis(
+            config_verysimple,
+            atom_data=atomic_data,
+            show_convergence_plots=False,
+        )
+        return sim
+
+    @property
+    def geometry(self):
+        return NumbaRadial1DGeometry(
+            r_inner=np.array([6.912e14, 8.64e14], dtype=np.float64),
+            r_outer=np.array([8.64e14, 1.0368e15], dtype=np.float64),
+            v_inner=np.array([-1, -1], dtype=np.float64),
+            v_outer=np.array([-1, -1], dtype=np.float64),
+        )
+
+
+    @property
+    def estimators(self):
+        return radfield_mc_estimators.RadiationFieldMCEstimators(
+            j_estimator=np.array([0.0, 0.0], dtype=np.float64),
+            nu_bar_estimator=np.array([0.0, 0.0], dtype=np.float64),
+            j_blue_estimator=np.array(
+                [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64
+            ),
+            Edotlu_estimator=np.array(
+                [[0.0, 0.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float64
+            ),
+            photo_ion_estimator=np.empty((0, 0), dtype=np.float64),
+            stim_recomb_estimator=np.empty((0, 0), dtype=np.float64),
+            bf_heating_estimator=np.empty((0, 0), dtype=np.float64),
+            stim_recomb_cooling_estimator=np.empty((0, 0), dtype=np.float64),
+            photo_ion_estimator_statistics=np.empty((0, 0), dtype=np.int64),
+        )

benchmarks/run_tardis.py

@@ -4,8 +4,6 @@
 
 from benchmarks.benchmark_base import BenchmarkBase
 from tardis import run_tardis
-from tardis.io.configuration.config_reader import Configuration
-
 
 class BenchmarkRunTardis(BenchmarkBase):
     """
@@ -14,12 +12,8 @@ class BenchmarkRunTardis(BenchmarkBase):
 
     def __init__(self):
         super().__init__()
-        self.config = None
-
-    def setup(self):
         filename = "data/tardis_configv1_benchmark.yml"
-        path = self.get_relative_path(filename)
-        self.config = Configuration.from_yaml(path)
+        self.path = self.get_relative_path(filename)
 
     def time_run_tardis(self):
-        run_tardis(self.config, log_level="ERROR", show_progress_bars=False)
+        run_tardis(self.path, log_level="ERROR", show_progress_bars=False)

benchmarks/transport_geometry_calculate_distances.py

@@ -0,0 +1,68 @@
+from asv_runner.benchmarks.mark import parameterize
+
+import tardis.transport.frame_transformations as frame_transformations
+import tardis.transport.geometry.calculate_distances as calculate_distances
+from benchmarks.benchmark_base import BenchmarkBase
+
+
+class BenchmarkTransportGeometryCalculateDistances(BenchmarkBase):
+    """
+    Class to benchmark the calculate distances function.
+    """
+
+    @property
+    def model(self):
+        return 5.2e7
+
+    def time_calculate_distance_boundary(self):
+        mu = 0.3
+        r = 7.5e14
+
+        calculate_distances.calculate_distance_boundary(
+            r, mu, self.geometry.r_inner[0], self.geometry.r_outer[0]
+        )
+
+    @parameterize(
+        {
+            "Parameters": [
+                {
+                    "packet": {
+                        "nu_line": 0.1,
+                        "is_last_line": True
+                    },
+                    "enable_full_relativity": True,
+                },
+                {
+                    "packet": {
+                        "nu_line": 0.2,
+                        "is_last_line": False
+                    },
+                    "enable_full_relativity": True,
+                }
+            ]
+        }
+    )
+    def time_calculate_distance_line(self, parameters):
+        packet_params = parameters["packet"]
+        nu_line = packet_params["nu_line"]
+        is_last_line = packet_params["is_last_line"]
+        enable_full_relativity = parameters["enable_full_relativity"]
+
+        time_explosion = self.model
+
+        doppler_factor = frame_transformations.get_doppler_factor(
+            self.static_packet.r,
+            self.static_packet.mu,
+            time_explosion,
+            enable_full_relativity
+        )
+        comov_nu = self.static_packet.nu * doppler_factor
+
+        calculate_distances.calculate_distance_line(
+            self.static_packet,
+            comov_nu,
+            is_last_line,
+            nu_line,
+            time_explosion,
+            enable_full_relativity
+        )

benchmarks/transport_montecarlo_interaction.py

@@ -2,17 +2,15 @@
 Basic TARDIS Benchmark.
 """
 
-import numpy as np
-from asv_runner.benchmarks.mark import parameterize, skip_benchmark
-
 import tardis.transport.montecarlo.interaction as interaction
 from benchmarks.benchmark_base import BenchmarkBase
 from tardis.transport.montecarlo.numba_interface import (
     LineInteractionType,
 )
+from asv_runner.benchmarks.mark import parameterize
+
 
 
-@skip_benchmark
 class BenchmarkMontecarloMontecarloNumbaInteraction(BenchmarkBase):
     """
     Class to benchmark the numba interaction function.
@@ -30,10 +28,6 @@ def time_thomson_scatter(self):
             packet, time_explosion, enable_full_relativity
         )
 
-        assert np.abs(packet.mu - init_mu) > 1e-7
-        assert np.abs(packet.nu - init_nu) > 1e-7
-        assert np.abs(packet.energy - init_energy) > 1e-7
-
     @parameterize(
         {
             "Line interaction type": [
@@ -45,9 +39,6 @@ def time_thomson_scatter(self):
     )
     def time_line_scatter(self, line_interaction_type):
         packet = self.packet
-        init_mu = packet.mu
-        init_nu = packet.nu
-        init_energy = packet.energy
         packet.initialize_line_id(
             self.verysimple_opacity_state,
             self.verysimple_time_explosion,
@@ -64,10 +55,6 @@ def time_line_scatter(self, line_interaction_type):
             self.verysimple_continuum_processes_enabled,
         )
 
-        assert np.abs(packet.mu - init_mu) > 1e-7
-        assert np.abs(packet.nu - init_nu) > 1e-7
-        assert np.abs(packet.energy - init_energy) > 1e-7
-
     @parameterize(
         {
             "Test packet": [
@@ -109,5 +96,3 @@ def time_line_emission(self, test_packet):
             self.verysimple_opacity_state,
             self.verysimple_enable_full_relativity,
         )
-
-        assert packet.next_line_id == emission_line_id + 1

benchmarks/transport_montecarlo_main_loop.py

@@ -0,0 +1,27 @@
+"""
+Basic TARDIS Benchmark.
+"""
+
+from benchmarks.benchmark_base import BenchmarkBase
+from tardis.transport.montecarlo.montecarlo_main_loop import montecarlo_main_loop
+
+
+class BenchmarkTransportMontecarloMainLoop(BenchmarkBase):
+    """
+    class to benchmark montecarlo_main_loop function.
+    """
+
+    def time_montecarlo_main_loop(self):
+        montecarlo_main_loop(
+            self.transport_state.packet_collection,
+            self.transport_state.geometry_state,
+            self.verysimple_time_explosion,
+            self.transport_state.opacity_state,
+            self.montecarlo_configuration,
+            self.transport_state.radfield_mc_estimators,
+            self.transport_state.spectrum_frequency.value,
+            self.montecarlo_configuration.NUMBER_OF_VPACKETS,
+            iteration=0,
+            show_progress_bars=False,
+            total_iterations=0
+        )