From 0349e7a90c90bfb9d9dbc5a106adfcbfda9a6f39 Mon Sep 17 00:00:00 2001 From: Asish Kumar <87874775+officialasishkumar@users.noreply.github.com> Date: Sat, 3 Aug 2024 01:18:15 +0530 Subject: [PATCH] Cache benchmark functions (#2749) * added setup function Signed-off-by: Asish Kumar * refactor geometry calculate distances Signed-off-by: Asish Kumar * Refactor montecarlo estimators in benchmark code Signed-off-by: Asish Kumar * add cache and setup functions Signed-off-by: Asish Kumar * add cache and setup functions Signed-off-by: Asish Kumar * ran ruff linter Signed-off-by: Asish Kumar * Thou hath vanquished the redundant variables and functions Signed-off-by: Asish Kumar * fix bugs Signed-off-by: Asish Kumar * change benchmark yml Signed-off-by: Asish Kumar * repeat in class Signed-off-by: Asish Kumar * rounds in benchmarks Signed-off-by: Asish Kumar * fix Signed-off-by: Asish Kumar * fix error in formal integral Signed-off-by: Asish Kumar * more cache Signed-off-by: Asish Kumar * sort by ratio Signed-off-by: Asish Kumar * cache in setup function Signed-off-by: Asish Kumar --------- Signed-off-by: Asish Kumar --- .github/workflows/benchmarks.yml | 10 +- benchmarks/benchmark_base.py | 256 +++--------------- benchmarks/opacities_opacity.py | 104 ++----- benchmarks/opacities_opacity_state.py | 20 +- benchmarks/run_tardis.py | 14 +- benchmarks/spectrum_formal_integral.py | 48 ++-- .../transport_geometry_calculate_distances.py | 61 ++--- ...rlo_estimators_radfield_estimator_calcs.py | 51 +--- .../transport_montecarlo_interaction.py | 86 ++---- benchmarks/transport_montecarlo_main_loop.py | 24 +- .../transport_montecarlo_packet_trackers.py | 42 +-- ...transport_montecarlo_single_packet_loop.py | 11 +- benchmarks/transport_montecarlo_vpacket.py | 127 ++++----- benchmarks/util/__init__.py | 0 benchmarks/util/base.py | 20 -- benchmarks/util/nlte.py | 78 ------ 16 files changed, 265 insertions(+), 687 deletions(-) delete mode 100644 benchmarks/util/__init__.py delete mode 100644 benchmarks/util/base.py delete mode 100644 benchmarks/util/nlte.py diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index 2ed44515e7b..95021ce994f 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -74,11 +74,11 @@ jobs: - name: Accept all asv questions run: asv machine --yes - - name: Run benchmarks for last 5 commits if not PR + - name: Run benchmarks for last 4 commits if not PR if: github.event_name != 'pull_request_target' run: | git log -n 4 --pretty=format:"%H" >> tag_commits.txt - asv run -a repeat=1 -a rounds=1 HASHFILE:tag_commits.txt | tee asv-output.log + asv run -a rounds=1 HASHFILE:tag_commits.txt | tee asv-output.log if grep -q failed asv-output.log; then echo "Some benchmarks have failed!" exit 1 @@ -109,17 +109,17 @@ jobs: run: | echo $(git rev-parse HEAD) > commit_hashes.txt echo $(git rev-parse master) >> commit_hashes.txt - asv run -a repeat=2 -a rounds=1 HASHFILE:commit_hashes.txt | tee asv-output-PR.log + asv run -a rounds=1 HASHFILE:commit_hashes.txt | tee asv-output-PR.log if grep -q failed asv-output-PR.log; then echo "Some benchmarks have failed!" exit 1 fi - name: Compare Master and PR head - run: asv compare origin/master HEAD --config asv.conf.json | tee asv-compare-output.log + run: asv compare origin/master HEAD --factor 1.1 --split --sort ratio | tee asv-compare-output.log - name: Compare Master and PR head but only show changed results - run: asv compare origin/master HEAD --only-changed --config asv.conf.json | tee asv-compare-changed-output.log + run: asv compare origin/master HEAD --only-changed --factor 1.1 --split --sort ratio | tee asv-compare-changed-output.log - name: Benchmarks compare output id: asv_pr_vs_master diff --git a/benchmarks/benchmark_base.py b/benchmarks/benchmark_base.py index c4ef1f2a98a..16f8cbac31a 100644 --- a/benchmarks/benchmark_base.py +++ b/benchmarks/benchmark_base.py @@ -1,47 +1,31 @@ +import functools from copy import deepcopy from os.path import dirname, join, realpath from pathlib import Path -from tempfile import mkstemp -import astropy.units as u import numpy as np -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.configuration import montecarlo_globals +from tardis.transport.montecarlo import RPacket, packet_trackers from tardis.transport.montecarlo.configuration.base import ( MonteCarloConfiguration, ) 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, - generate_rpacket_last_interaction_tracker_list, -) +from tardis.transport.montecarlo.packet_collections import VPacketCollection +from tardis.transport.montecarlo.packet_trackers import RPacketTracker class BenchmarkBase: - # It allows 10 minutes of runtime for each benchmark and includes # the total time for all the repetitions for each benchmark. timeout = 600 - def __init__(self): - self.nlte = NLTE() - @staticmethod def get_relative_path(partial_path: str): path = dirname(realpath(__file__)) @@ -57,7 +41,7 @@ def get_absolute_path(self, partial_path): return self.get_relative_path(partial_path) - @property + @functools.cached_property def tardis_config_verysimple(self): filename = self.get_absolute_path( "tardis/io/configuration/tests/data/tardis_configv1_verysimple.yml" @@ -67,7 +51,7 @@ def tardis_config_verysimple(self): YAMLLoader, ) - @property + @functools.cached_property def config_rpacket_tracking(self): config = Configuration.from_yaml( f"{self.example_configuration_dir}/tardis_configv1_verysimple.yml" @@ -75,17 +59,15 @@ def config_rpacket_tracking(self): config.montecarlo.tracking.track_rpacket = True return config - @property + @functools.cached_property def tardis_ref_path(self): - # TODO: This route is fixed but needs to get from the arguments given in the command line. - # /app/tardis-refdata ref_data_path = Path( Path(__file__).parent.parent, - "tardis-refdata", + "tardis-refdata" ).resolve() return ref_data_path - @property + @functools.cached_property def atomic_dataset(self) -> AtomData: atomic_data = AtomData.from_hdf(self.atomic_data_fname) @@ -95,7 +77,7 @@ def atomic_dataset(self) -> AtomData: else: return atomic_data - @property + @functools.cached_property def atomic_data_fname(self): atomic_data_fname = ( f"{self.tardis_ref_path}/atom_data/kurucz_cd23_chianti_H_He.h5" @@ -110,59 +92,11 @@ def atomic_data_fname(self): return atomic_data_fname - @property + @functools.cached_property def example_configuration_dir(self): return self.get_absolute_path("tardis/io/configuration/tests/data") - @property - def hdf_file_path(self): - # TODO: Delete this files after ASV runs the benchmarks. - # ASV create a temporal directory in runtime per test: `tmpiuxngvlv`. - # The ASV and ASV_Runner, not has some way to get this temporal directory. - # The idea is use this temporal folders to storage this created temporal file. - _, path = mkstemp("-tardis-benchmark-hdf_buffer-test.hdf") - return path - - def create_temporal_file(self, suffix=None): - # TODO: Delete this files after ASV runs the benchmarks. - # ASV create a temporal directory in runtime per test: `tmpiuxngvlv`. - # The ASV and ASV_Runner, not has some way to get this temporal directory. - # The idea is use this temporal folders to storage this created temporal file. - suffix_str = "" if suffix is None else f"-{suffix}" - _, path = mkstemp(suffix_str) - return path - - @property - def gamma_ray_simulation_state(self): - self.gamma_ray_config.model.structure.velocity.start = 1.0 * u.km / u.s - self.gamma_ray_config.model.structure.density.rho_0 = ( - 5.0e2 * u.g / u.cm**3 - ) - self.gamma_ray_config.supernova.time_explosion = 150 * u.d - - return SimulationState.from_config( - self.gamma_ray_config, atom_data=self.atomic_dataset - ) - - @property - def gamma_ray_config(self): - yml_path = f"{self.example_configuration_dir}/tardis_configv1_density_exponential_nebular_multi_isotope.yml" - - return config_reader.Configuration.from_yaml(yml_path) - - @property - def example_model_file_dir(self): - return self.get_absolute_path("tardis/io/model/readers/tests/data") - - @property - def kurucz_atomic_data(self) -> AtomData: - return deepcopy(self.atomic_dataset) - - @property - def example_csvy_file_dir(self): - return self.get_absolute_path("tardis/model/tests/data/") - - @property + @functools.cached_property def simulation_verysimple(self): atomic_data = deepcopy(self.atomic_dataset) sim = Simulation.from_config( @@ -171,73 +105,13 @@ def simulation_verysimple(self): sim.iterate(4000) return sim - @property + @functools.cached_property def config_verysimple(self): return Configuration.from_yaml( f"{self.example_configuration_dir}/tardis_configv1_verysimple.yml" ) - class CustomPyTestRequest: - def __init__( - self, - tardis_regression_data_path: str, - node_name: str, - node_module_name: str, - regression_data_dir: str, - ): - self.tardis_regression_data_path = tardis_regression_data_path - self.node_name = node_name - self.node_module_name = node_module_name - self.regression_data_dir = regression_data_dir - - @property - def config(self): - class SubClass: - @staticmethod - def getoption(option): - if option == "--tardis-regression-data": - return self.tardis_regression_data_path - return None - - return SubClass() - - @property - def node(self): - class SubClass: - def __init__(self, parent): - self.parent = parent - - @property - def name(self): - return self.parent.node_name - - @property - def module(self): - class SubSubClass: - def __init__(self, parent): - self.parent = parent - - @property - def __name__(self): - return self.parent.node_module_name - - return SubSubClass(self.parent) - - return SubClass(self) - - @property - def cls(self): - return None - - @property - def relative_regression_data_dir(self): - return self.regression_data_dir - - @staticmethod - def regression_data(request: CustomPyTestRequest): - return RegressionData(request) - - @property + @functools.cached_property def packet(self): return RPacket( r=7.5e14, @@ -248,27 +122,25 @@ def packet(self): index=0, ) - @property + @functools.cached_property def verysimple_packet_collection(self): - return ( - self.nb_simulation_verysimple.transport.transport_state.packet_collection - ) + return self.nb_simulation_verysimple.transport.transport_state.packet_collection - @property + @functools.cached_property def nb_simulation_verysimple(self): atomic_data = deepcopy(self.atomic_dataset) sim = Simulation.from_config( self.config_verysimple, atom_data=atomic_data ) - sim.iterate(10) + sim.iterate(5) return sim - @property + @functools.cached_property def verysimple_time_explosion(self): model = self.nb_simulation_verysimple.simulation_state return model.time_explosion.cgs.value - @property + @functools.cached_property def verysimple_opacity_state(self): return opacity_state_initialize( self.nb_simulation_verysimple.plasma, @@ -276,33 +148,19 @@ def verysimple_opacity_state(self): disable_line_scattering=self.nb_simulation_verysimple.transport.montecarlo_configuration.DISABLE_LINE_SCATTERING, ) - @property + @functools.cached_property def verysimple_enable_full_relativity(self): return self.nb_simulation_verysimple.transport.enable_full_relativity - @property - def verysimple_disable_line_scattering(self): - return ( - self.nb_simulation_verysimple.transport.montecarlo_configuration.DISABLE_LINE_SCATTERING - ) - - @property - def verysimple_continuum_processes_enabled(self): - return montecarlo_globals.CONTINUUM_PROCESSES_ENABLED - - @property + @functools.cached_property def verysimple_tau_russian(self): - return ( - self.nb_simulation_verysimple.transport.montecarlo_configuration.VPACKET_TAU_RUSSIAN - ) + return self.nb_simulation_verysimple.transport.montecarlo_configuration.VPACKET_TAU_RUSSIAN - @property + @functools.cached_property def verysimple_survival_probability(self): - return ( - self.nb_simulation_verysimple.transport.montecarlo_configuration.SURVIVAL_PROBABILITY - ) + return self.nb_simulation_verysimple.transport.montecarlo_configuration.SURVIVAL_PROBABILITY - @property + @functools.cached_property def static_packet(self): return RPacket( r=7.5e14, @@ -313,18 +171,9 @@ def static_packet(self): index=0, ) - @property - def set_seed_fixture(self): - def set_seed(value): - np.random.seed(value) - - return njit(set_seed) - - @property + @functools.cached_property def verysimple_3vpacket_collection(self): - spectrum_frequency_grid = ( - self.nb_simulation_verysimple.transport.spectrum_frequency_grid.value - ) + spectrum_frequency_grid = self.nb_simulation_verysimple.transport.spectrum_frequency_grid.value return VPacketCollection( source_rpacket_index=0, spectrum_frequency_grid=spectrum_frequency_grid, @@ -334,54 +183,25 @@ def verysimple_3vpacket_collection(self): temporary_v_packet_bins=0, ) - @property + @functools.cached_property def verysimple_numba_radial_1d_geometry(self): return ( self.nb_simulation_verysimple.simulation_state.geometry.to_numba() ) - @property - def simulation_verysimple_vpacket_tracking(self): - atomic_data = deepcopy(self.atomic_dataset) - sim = Simulation.from_config( - self.config_verysimple, - atom_data=atomic_data, - virtual_packet_logging=True, - ) - sim.last_no_of_packets = 4000 - sim.run_final() - return sim - - @property - def generate_reference(self): - # TODO: Investigate how to get the `--generate-reference` parameter passed in the command line. - # `request.config.getoption("--generate-reference")` - option = None - if option is None: - return False - else: - return option - - @property - 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 + @functools.cached_property def montecarlo_configuration(self): return MonteCarloConfiguration() - @property + @functools.cached_property def rpacket_tracker(self): return RPacketTracker(0) - @property + @functools.cached_property def transport_state(self): return self.nb_simulation_verysimple.transport.transport_state - @property + @functools.cached_property def simulation_rpacket_tracking_enabled(self): config_verysimple = self.config_verysimple config_verysimple.montecarlo.iterations = 3 @@ -399,7 +219,7 @@ def simulation_rpacket_tracking_enabled(self): ) return sim - @property + @functools.cached_property def geometry(self): return NumbaRadial1DGeometry( r_inner=np.array([6.912e14, 8.64e14], dtype=np.float64), @@ -408,7 +228,7 @@ def geometry(self): v_outer=np.array([-1, -1], dtype=np.float64), ) - @property + @functools.cached_property def estimators(self): return radfield_mc_estimators.RadiationFieldMCEstimators( j_estimator=np.array([0.0, 0.0], dtype=np.float64), @@ -426,7 +246,9 @@ def estimators(self): photo_ion_estimator_statistics=np.empty((0, 0), dtype=np.int64), ) - @property + @functools.cached_property def rpacket_tracker_list(self): no_of_packets = len(self.transport_state.packet_collection.initial_nus) - return generate_rpacket_last_interaction_tracker_list(no_of_packets) + return packet_trackers.generate_rpacket_last_interaction_tracker_list( + no_of_packets + ) diff --git a/benchmarks/opacities_opacity.py b/benchmarks/opacities_opacity.py index de571cc2184..5352ceccc89 100644 --- a/benchmarks/opacities_opacity.py +++ b/benchmarks/opacities_opacity.py @@ -2,8 +2,6 @@ Basic TARDIS Benchmark. """ -from asv_runner.benchmarks.mark import parameterize - import tardis.opacities.opacities as calculate_opacity from benchmarks.benchmark_base import BenchmarkBase from tardis.opacities.opacities import compton_opacity_calculation @@ -14,95 +12,25 @@ class BenchmarkMontecarloMontecarloNumbaOpacities(BenchmarkBase): Class to benchmark the numba opacities function. """ - @parameterize( - { - "Electron number density": [ - 1.0e11, - 1e15, - 1e5, - ], - "Energy": [ - 511.0, - 255.5, - 511.0e7, - ], - } - ) - def time_compton_opacity_calculation(self, electron_number_density, energy): - compton_opacity_calculation(energy, electron_number_density) + def time_compton_opacity_calculation(self): + energy = 511.0 + electron_number_density = 1e15 + calculate_opacity.compton_opacity_calculation( + energy, electron_number_density + ) - @parameterize( - { - "Parameters": [ - { - "Ejecta_density": 0.01, - "Energy": 511.0, - "Iron_group_fraction": 0.5, - }, - { - "Ejecta_density": 0.01, - "Energy": 5110000000.0, - "Iron_group_fraction": 0.0, - }, - { - "Ejecta_density": 0.01, - "Energy": 255.5, - "Iron_group_fraction": 0.0, - }, - { - "Ejecta_density": 0.01, - "Energy": 255.5, - "Iron_group_fraction": 0.5, - }, - { - "Ejecta_density": 100000.0, - "Energy": 255.5, - "Iron_group_fraction": 1.0, - } - ] - } - ) - def time_photoabsorption_opacity_calculation(self, parameters): + def time_photoabsorption_opacity_calculation(self): + energy = 255.5 + ejecta_density = 100000.0 + iron_group_fraction = 0.5 calculate_opacity.photoabsorption_opacity_calculation( - parameters["Energy"], - parameters["Ejecta_density"], - parameters["Iron_group_fraction"], + energy, ejecta_density, iron_group_fraction ) - @parameterize( - { - "Parameters": [ - { - "Ejecta_density": 0.01, - "Energy": 511.0, - "Iron_group_fraction": 0.5, - }, - { - "Ejecta_density": 0.01, - "Energy": 5110000000.0, - "Iron_group_fraction": 0.0, - }, - { - "Ejecta_density": 0.01, - "Energy": 255.5, - "Iron_group_fraction": 0.0, - }, - { - "Ejecta_density": 0.01, - "Energy": 255.5, - "Iron_group_fraction": 0.5, - }, - { - "Ejecta_density": 100000.0, - "Energy": 255.5, - "Iron_group_fraction": 1.0, - } - ] - } - ) - def time_pair_creation_opacity_calculation(self, parameters): + def time_pair_creation_opacity_calculation(self): + energy = 255.9 + ejecta_density = 100000.009 + iron_group_fraction = 0.5 calculate_opacity.pair_creation_opacity_calculation( - parameters["Energy"], - parameters["Ejecta_density"], - parameters["Iron_group_fraction"], + energy, ejecta_density, iron_group_fraction ) diff --git a/benchmarks/opacities_opacity_state.py b/benchmarks/opacities_opacity_state.py index ea05e794691..18ae4769a95 100644 --- a/benchmarks/opacities_opacity_state.py +++ b/benchmarks/opacities_opacity_state.py @@ -2,24 +2,26 @@ Basic TARDIS Benchmark. """ -import numpy as np +import functools + from asv_runner.benchmarks.mark import parameterize -from tardis.opacities.opacity_state import opacity_state_initialize from benchmarks.benchmark_base import BenchmarkBase +from tardis.opacities.opacity_state import opacity_state_initialize +@parameterize({"Input params": ["scatter", "macroatom"]}) class BenchmarkOpacitiesOpacityState(BenchmarkBase): """ Class to benchmark the numba interface function. """ + repeat = 2 + + @functools.cache + def setup(self, input_params): + self.sim = self.nb_simulation_verysimple - @parameterize({"Input params": ["scatter", "macroatom"]}) def time_opacity_state_initialize(self, input_params): line_interaction_type = input_params - plasma = self.nb_simulation_verysimple.plasma - opacity_state_initialize( - plasma, - line_interaction_type, - self.verysimple_disable_line_scattering, - ) + plasma = self.sim.plasma + opacity_state_initialize(plasma, line_interaction_type, True) diff --git a/benchmarks/run_tardis.py b/benchmarks/run_tardis.py index 9e00a6a3f9b..ab423ded7b4 100644 --- a/benchmarks/run_tardis.py +++ b/benchmarks/run_tardis.py @@ -4,6 +4,8 @@ from benchmarks.benchmark_base import BenchmarkBase from tardis import run_tardis +from tardis.io.atom_data import AtomData + class BenchmarkRunTardis(BenchmarkBase): @@ -11,16 +13,20 @@ class BenchmarkRunTardis(BenchmarkBase): Class to benchmark the `run tardis` function. """ + repeat = 2 + + def setup(self): + self.config = self.config_verysimple + self.atom_data = AtomData.from_hdf(self.atomic_data_fname) + def time_run_tardis(self): run_tardis( - self.config_verysimple, - atom_data=self.atomic_dataset, - show_convergence_plots=False, + self.config, atom_data=self.atom_data, show_convergence_plots=False ) def time_run_tardis_rpacket_tracking(self): run_tardis( self.config_rpacket_tracking, - atom_data=self.atomic_dataset, + atom_data=self.atom_data, show_convergence_plots=False, ) diff --git a/benchmarks/spectrum_formal_integral.py b/benchmarks/spectrum_formal_integral.py index c7eb25e6f5a..8c69ea71bce 100644 --- a/benchmarks/spectrum_formal_integral.py +++ b/benchmarks/spectrum_formal_integral.py @@ -2,43 +2,43 @@ Basic TARDIS Benchmark. """ -from asv_runner.benchmarks.mark import parameterize +import functools + from numba import config -import tardis.spectrum.formal_integral as formal_integral from benchmarks.benchmark_base import BenchmarkBase +from tardis.spectrum import formal_integral + +config.THREADING_LAYER = "workqueue" -config.THREADING_LAYER='workqueue' class BenchmarkTransportMontecarloFormalIntegral(BenchmarkBase): """ Class to benchmark the numba formal integral function. """ - @parameterize( - { - "Parameters": [ - { - "nu": 1e14, - "temperature": 1e4, - } - ] - } - ) - def time_intensity_black_body(self, parameters): - nu = parameters["nu"] - temperature = parameters["temperature"] + repeat = 2 + + @functools.cache + def setup(self): + self.sim = self.simulation_verysimple + self.FormalIntegrator = formal_integral.FormalIntegrator( + self.sim.simulation_state, self.sim.plasma, self.sim.transport + ) + + # Bencmark for intensity black body function + def time_intensity_black_body(self): + nu = 1e14 + temperature = 1e4 formal_integral.intensity_black_body(nu, temperature) # 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.FormalIntegrator.calculate_spectrum( + self.sim.spectrum_solver.spectrum_real_packets.frequency ) - 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.FormalIntegrator.make_source_function() + self.FormalIntegrator.generate_numba_objects() + self.FormalIntegrator.formal_integral( + self.sim.spectrum_solver.spectrum_real_packets.frequency, 1000 ) diff --git a/benchmarks/transport_geometry_calculate_distances.py b/benchmarks/transport_geometry_calculate_distances.py index 4f5829f40b8..5500f60f4bb 100644 --- a/benchmarks/transport_geometry_calculate_distances.py +++ b/benchmarks/transport_geometry_calculate_distances.py @@ -1,4 +1,4 @@ -from asv_runner.benchmarks.mark import parameterize +import functools import tardis.transport.frame_transformations as frame_transformations import tardis.transport.geometry.calculate_distances as calculate_distances @@ -10,6 +10,15 @@ class BenchmarkTransportGeometryCalculateDistances(BenchmarkBase): Class to benchmark the calculate distances function. """ + @functools.cache + def setup(self): + self.StaticPacket = self.static_packet + self.Geometry = self.geometry + doppler_factor = frame_transformations.get_doppler_factor( + self.StaticPacket.r, self.StaticPacket.mu, self.model, True + ) + self.comov_nu = self.StaticPacket.nu * doppler_factor + @property def model(self): return 5.2e7 @@ -19,50 +28,16 @@ def time_calculate_distance_boundary(self): 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 + r, mu, self.Geometry.r_inner[0], self.Geometry.r_outer[0] ) - comov_nu = self.static_packet.nu * doppler_factor + def time_calculate_distance_line(self): + nu_line = 0.2 calculate_distances.calculate_distance_line( - self.static_packet, - comov_nu, - is_last_line, + self.StaticPacket, + self.comov_nu, + True, nu_line, - time_explosion, - enable_full_relativity + True, + True, ) diff --git a/benchmarks/transport_montecarlo_estimators_radfield_estimator_calcs.py b/benchmarks/transport_montecarlo_estimators_radfield_estimator_calcs.py index 45a3b3b8409..8d71a27f492 100644 --- a/benchmarks/transport_montecarlo_estimators_radfield_estimator_calcs.py +++ b/benchmarks/transport_montecarlo_estimators_radfield_estimator_calcs.py @@ -2,15 +2,12 @@ Basic TARDIS Benchmark. """ -import tardis.opacities.opacities as opacities -import tardis.transport.geometry.calculate_distances as calculate_distances -import tardis.transport.montecarlo.r_packet_transport as r_packet_transport -import tardis.transport.montecarlo.utils as utils +import functools + from benchmarks.benchmark_base import BenchmarkBase from tardis.transport.montecarlo.estimators.radfield_estimator_calcs import ( update_line_estimators, ) -from asv_runner.benchmarks.mark import parameterize class BenchmarkMontecarloMontecarloNumbaPacket(BenchmarkBase): @@ -18,40 +15,14 @@ class BenchmarkMontecarloMontecarloNumbaPacket(BenchmarkBase): Class to benchmark the numba packet function. """ - @parameterize( - { - "Parameters": [ - { - "cur_line_id": 0, - "distance_trace": 1e12, - "time_explosion": 5.2e7, - "enable_full_relativity": True, - }, - { - "cur_line_id": 0, - "distance_trace": 0, - "time_explosion": 5.2e7, - "enable_full_relativity": True, - }, - { - "cur_line_id": 1, - "distance_trace": 1e5, - "time_explosion": 1e10, - "enable_full_relativity": False, - }, - ] - } - ) - def time_update_line_estimators(self, parameters): - cur_line_id = parameters["cur_line_id"] - distance_trace = parameters["distance_trace"] - time_explosion = parameters["time_explosion"] - enable_full_relativity = parameters["enable_full_relativity"] + repeat = 4 + + @functools.cache + def setup(self): + self.Estimators = self.estimators + self.StaticPacket = self.static_packet + + def time_update_line_estimators(self): update_line_estimators( - self.estimators, - self.static_packet, - cur_line_id, - distance_trace, - time_explosion, - enable_full_relativity, + self.Estimators, self.StaticPacket, 1, 1e12, 1e10, True ) diff --git a/benchmarks/transport_montecarlo_interaction.py b/benchmarks/transport_montecarlo_interaction.py index e89c486aac2..dcadc56a2da 100644 --- a/benchmarks/transport_montecarlo_interaction.py +++ b/benchmarks/transport_montecarlo_interaction.py @@ -2,12 +2,10 @@ Basic TARDIS Benchmark. """ +import functools + 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 class BenchmarkTransportMontecarloInteraction(BenchmarkBase): @@ -15,68 +13,44 @@ class BenchmarkTransportMontecarloInteraction(BenchmarkBase): Class to benchmark the numba interaction function. """ - def time_thomson_scatter(self): - packet = self.packet - time_explosion = self.verysimple_time_explosion - enable_full_relativity = self.verysimple_enable_full_relativity + repeat = 3 + + @functools.cache + def setup(self): + self.Packet = self.packet + self.time_explosion = self.verysimple_time_explosion + self.enable_full_relativity = self.verysimple_enable_full_relativity + self.opacity_state = self.verysimple_opacity_state + def time_thomson_scatter(self): interaction.thomson_scatter( - packet, time_explosion, enable_full_relativity + self.Packet, self.time_explosion, self.enable_full_relativity ) - @parameterize( - { - "Line interaction type": [ - LineInteractionType.SCATTER, - LineInteractionType.MACROATOM, - ], - } - ) - def time_line_scatter(self, line_interaction_type): - packet = self.packet - packet.initialize_line_id( - self.verysimple_opacity_state, - self.verysimple_time_explosion, - self.verysimple_enable_full_relativity, + def time_line_scatter(self): + self.Packet.initialize_line_id( + self.opacity_state, self.time_explosion, self.enable_full_relativity ) - time_explosion = self.verysimple_time_explosion - interaction.line_scatter( - packet, - time_explosion, - line_interaction_type, - self.verysimple_opacity_state, - self.verysimple_enable_full_relativity, + self.Packet, + self.time_explosion, + "SCATTER", + self.opacity_state, + self.enable_full_relativity, ) - @parameterize( - { - "Test packet": [ - { - "mu": 0.8599443103322428, - "emission_line_id": 1000, - "energy": 0.9114437898710559, - } - ] - } - ) - def time_line_emission(self, test_packet): - emission_line_id = test_packet["emission_line_id"] - packet = self.packet - packet.mu = test_packet["mu"] - packet.energy = test_packet["energy"] - packet.initialize_line_id( - self.verysimple_opacity_state, - self.verysimple_time_explosion, - self.verysimple_enable_full_relativity, + def time_line_emission(self): + emission_line_id = 1000 + self.Packet.mu = 0.8599443103322428 + self.Packet.energy = 0.9114437898710559 + self.Packet.initialize_line_id( + self.opacity_state, self.time_explosion, self.enable_full_relativity ) - time_explosion = self.verysimple_time_explosion - interaction.line_emission( - packet, + self.Packet, emission_line_id, - time_explosion, - self.verysimple_opacity_state, - self.verysimple_enable_full_relativity, + self.time_explosion, + self.opacity_state, + self.enable_full_relativity, ) diff --git a/benchmarks/transport_montecarlo_main_loop.py b/benchmarks/transport_montecarlo_main_loop.py index 7d6d07ddb08..cc91fd430d3 100644 --- a/benchmarks/transport_montecarlo_main_loop.py +++ b/benchmarks/transport_montecarlo_main_loop.py @@ -2,6 +2,8 @@ Basic TARDIS Benchmark. """ +import functools + from benchmarks.benchmark_base import BenchmarkBase from tardis.transport.montecarlo.montecarlo_main_loop import ( montecarlo_main_loop, @@ -13,14 +15,26 @@ class BenchmarkTransportMontecarloMontecarloMainLoop(BenchmarkBase): class to benchmark montecarlo_main_loop function. """ + repeat = 3 + + @functools.cache + def setup(self): + self.packet_collection = self.transport_state.packet_collection + self.geometry_state = self.transport_state.geometry_state + self.time_explosion = self.verysimple_time_explosion + self.opacity_state = self.transport_state.opacity_state + self.radfield_mc_estimators = ( + self.transport_state.radfield_mc_estimators + ) + 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.packet_collection, + self.geometry_state, + self.time_explosion, + self.opacity_state, self.montecarlo_configuration, - self.transport_state.radfield_mc_estimators, + self.radfield_mc_estimators, self.nb_simulation_verysimple.transport.spectrum_frequency_grid.value, self.rpacket_tracker_list, self.montecarlo_configuration.NUMBER_OF_VPACKETS, diff --git a/benchmarks/transport_montecarlo_packet_trackers.py b/benchmarks/transport_montecarlo_packet_trackers.py index a3f65f66f6d..e4207d23892 100644 --- a/benchmarks/transport_montecarlo_packet_trackers.py +++ b/benchmarks/transport_montecarlo_packet_trackers.py @@ -1,36 +1,38 @@ """ Basic TARDIS Benchmark. """ +import functools + +from asv_runner.benchmarks.mark import parameterize + from benchmarks.benchmark_base import BenchmarkBase -from tardis.transport.montecarlo.packet_trackers import ( - rpacket_trackers_to_dataframe, - generate_rpacket_tracker_list, - generate_rpacket_last_interaction_tracker_list, -) +from tardis.transport.montecarlo import packet_trackers +@parameterize({"num_packets": [10, 100], "length": [10, 50]}) class BenchmarkTransportMontecarloPacketTrackers(BenchmarkBase): """ Class to benchmark the numba R packet function. """ - def time_rpacket_trackers_to_dataframe(self): + repeat = 2 + + @functools.cache + def setup(self, num_packets, length): sim = self.simulation_rpacket_tracking_enabled - transport_state = sim.transport.transport_state - rpacket_trackers_to_dataframe(transport_state.rpacket_tracker) + self.TransportState = sim.transport.transport_state - def time_generate_rpacket_tracker_list(self, no_of_packets, length): - generate_rpacket_tracker_list(no_of_packets, length) + def time_rpacket_trackers_to_dataframe(self, num_packets, length): + packet_trackers.rpacket_trackers_to_dataframe( + self.TransportState.rpacket_tracker + ) + + def time_generate_rpacket_tracker_list(self, num_packets, length): + packet_trackers.generate_rpacket_tracker_list(num_packets, length) def time_generate_rpacket_last_interaction_tracker_list( - self, no_of_packets + self, num_packets, length ): - generate_rpacket_last_interaction_tracker_list(no_of_packets) - - time_generate_rpacket_tracker_list.params = ([1, 10, 50], [1, 10, 50]) - time_generate_rpacket_tracker_list.param_names = ["no_of_packets", "length"] - - time_generate_rpacket_last_interaction_tracker_list.params = [10, 100, 1000] - time_generate_rpacket_last_interaction_tracker_list.param_names = [ - "no_of_packets" - ] + packet_trackers.generate_rpacket_last_interaction_tracker_list( + num_packets + ) diff --git a/benchmarks/transport_montecarlo_single_packet_loop.py b/benchmarks/transport_montecarlo_single_packet_loop.py index 85bc6bbf64f..4a7138d5fb5 100644 --- a/benchmarks/transport_montecarlo_single_packet_loop.py +++ b/benchmarks/transport_montecarlo_single_packet_loop.py @@ -2,9 +2,10 @@ Basic TARDIS Benchmark. """ +from numba.np.ufunc.parallel import get_num_threads, get_thread_id + from benchmarks.benchmark_base import BenchmarkBase from tardis.transport.montecarlo import single_packet_loop -from numba.np.ufunc.parallel import get_num_threads, get_thread_id class BenchmarkTransportMontecarloSinglePacketLoop(BenchmarkBase): @@ -12,14 +13,18 @@ class BenchmarkTransportMontecarloSinglePacketLoop(BenchmarkBase): Class to benchmark the single packet loop function. """ + repeat = 2 + def time_single_packet_loop(self): single_packet_loop.single_packet_loop( self.packet, self.verysimple_numba_radial_1d_geometry, self.verysimple_time_explosion, self.verysimple_opacity_state, - self.transport_state.radfield_mc_estimators.create_estimator_list(get_num_threads())[get_thread_id()], + self.transport_state.radfield_mc_estimators.create_estimator_list( + get_num_threads() + )[get_thread_id()], self.verysimple_3vpacket_collection, self.rpacket_tracker, - self.montecarlo_configuration + self.montecarlo_configuration, ) diff --git a/benchmarks/transport_montecarlo_vpacket.py b/benchmarks/transport_montecarlo_vpacket.py index 4de7d198c68..b237577074d 100644 --- a/benchmarks/transport_montecarlo_vpacket.py +++ b/benchmarks/transport_montecarlo_vpacket.py @@ -2,15 +2,14 @@ Basic TARDIS Benchmark. """ +import functools + import numpy as np -from asv_runner.benchmarks.mark import parameterize import tardis.transport.montecarlo.vpacket as vpacket -from tardis.transport.montecarlo.r_packet import RPacket from benchmarks.benchmark_base import BenchmarkBase -from tardis.transport.frame_transformations import ( - get_doppler_factor, -) +from tardis.transport.frame_transformations import get_doppler_factor +from tardis.transport.montecarlo.r_packet import RPacket class BenchmarkMontecarloMontecarloNumbaVpacket(BenchmarkBase): @@ -18,7 +17,18 @@ class BenchmarkMontecarloMontecarloNumbaVpacket(BenchmarkBase): Class to benchmark the single packet loop function. """ - @property + repeat = 5 + + def setup(self): + self.vpacket = self.v_packet + self.numba_radial_1d_geometry = self.verysimple_numba_radial_1d_geometry + self.time_explosion = self.verysimple_time_explosion + self.opacity_state = self.verysimple_opacity_state + self.tau_russian = self.verysimple_tau_russian + self.survival_probability = self.verysimple_survival_probability + self.enable_full_relativity = self.verysimple_enable_full_relativity + + @functools.cached_property def v_packet(self): return vpacket.VPacket( r=7.5e14, @@ -30,7 +40,7 @@ def v_packet(self): index=0, ) - @property + @functools.cached_property def r_packet(self): return RPacket( r=7.5e14, @@ -41,6 +51,7 @@ def r_packet(self): index=0, ) + @functools.cache def v_packet_initialize_line_id( self, v_packet, opacity_state, time_explosion, enable_full_relativity ): @@ -55,63 +66,45 @@ def v_packet_initialize_line_id( v_packet.next_line_id = next_line_id def time_trace_vpacket_within_shell(self): - v_packet = self.v_packet - verysimple_numba_radial_1d_geometry = ( - self.verysimple_numba_radial_1d_geometry - ) - verysimple_time_explosion = self.verysimple_time_explosion - verysimple_opacity_state = self.verysimple_opacity_state - enable_full_relativity = self.verysimple_enable_full_relativity - # Give the vpacket a reasonable line ID self.v_packet_initialize_line_id( - v_packet, - verysimple_opacity_state, - verysimple_time_explosion, - enable_full_relativity, + self.vpacket, + self.opacity_state, + self.time_explosion, + self.enable_full_relativity, ) vpacket.trace_vpacket_within_shell( - v_packet, - verysimple_numba_radial_1d_geometry, - verysimple_time_explosion, - verysimple_opacity_state, - enable_full_relativity, + self.vpacket, + self.numba_radial_1d_geometry, + self.time_explosion, + self.opacity_state, + self.enable_full_relativity, ) def time_trace_vpacket(self): - v_packet = self.v_packet - verysimple_numba_radial_1d_geometry = ( - self.verysimple_numba_radial_1d_geometry - ) - verysimple_time_explosion = self.verysimple_time_explosion - verysimple_opacity_state = self.verysimple_opacity_state - enable_full_relativity = self.verysimple_enable_full_relativity - tau_russian = self.verysimple_tau_russian - survival_probability = self.verysimple_survival_probability - # Set seed because of RNG in trace_vpacket np.random.seed(1) # Give the vpacket a reasonable line ID self.v_packet_initialize_line_id( - v_packet, - verysimple_opacity_state, - verysimple_time_explosion, - enable_full_relativity, + self.vpacket, + self.opacity_state, + self.time_explosion, + self.enable_full_relativity, ) vpacket.trace_vpacket( - v_packet, - verysimple_numba_radial_1d_geometry, - verysimple_time_explosion, - verysimple_opacity_state, - tau_russian, - survival_probability, - enable_full_relativity, + self.vpacket, + self.numba_radial_1d_geometry, + self.time_explosion, + self.opacity_state, + self.tau_russian, + self.survival_probability, + self.enable_full_relativity, ) - @property + @functools.cached_property def broken_packet(self): return vpacket.VPacket( r=1286064000000000.0, @@ -125,41 +118,25 @@ def broken_packet(self): def time_trace_bad_vpacket(self): broken_packet = self.broken_packet - verysimple_numba_radial_1d_geometry = ( - self.verysimple_numba_radial_1d_geometry - ) - enable_full_relativity = self.verysimple_enable_full_relativity - verysimple_time_explosion = self.verysimple_time_explosion - verysimple_opacity_state = self.verysimple_opacity_state - tau_russian = self.verysimple_tau_russian - survival_probability = self.verysimple_survival_probability vpacket.trace_vpacket( broken_packet, - verysimple_numba_radial_1d_geometry, - verysimple_time_explosion, - verysimple_opacity_state, - tau_russian, - survival_probability, - enable_full_relativity, + self.numba_radial_1d_geometry, + self.time_explosion, + self.opacity_state, + self.tau_russian, + self.survival_probability, + self.enable_full_relativity, ) - @parameterize( - { - "Paramters": [ - {"tau_russian": 10.0, "survival_possibility": 0.0}, - {"tau_russian": 15.0, "survival_possibility": 0.1}, - ] - } - ) - def time_trace_vpacket_volley(self, parameters): + def time_trace_vpacket_volley(self): vpacket.trace_vpacket_volley( self.r_packet, self.verysimple_3vpacket_collection, - self.verysimple_numba_radial_1d_geometry, - self.verysimple_time_explosion, - self.verysimple_opacity_state, - False, - parameters["tau_russian"], - parameters["survival_possibility"] + self.numba_radial_1d_geometry, + self.time_explosion, + self.opacity_state, + self.enable_full_relativity, + self.tau_russian, + self.survival_probability, ) diff --git a/benchmarks/util/__init__.py b/benchmarks/util/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/benchmarks/util/base.py b/benchmarks/util/base.py deleted file mode 100644 index 7dd61629e3e..00000000000 --- a/benchmarks/util/base.py +++ /dev/null @@ -1,20 +0,0 @@ -from os.path import dirname, realpath, join -from pathlib import Path, PosixPath - - -class Base: - @staticmethod - def get_path(partial_path: str) -> Path: - base_path = dirname(realpath(__file__)) - path = Path(base_path) / Path(partial_path) - return path - - @property - def tardis_ref_path(self) -> Path: - # TODO: This route is fixed but needs to get from the arguments given in the command line. - # /app/tardis-refdata - return Path("/app/tardis-refdata") - - @property - def example_configuration_dir(self) -> Path: - return self.get_path("../../tardis/io/configuration/tests/data") diff --git a/benchmarks/util/nlte.py b/benchmarks/util/nlte.py deleted file mode 100644 index 7cfed4ef4dd..00000000000 --- a/benchmarks/util/nlte.py +++ /dev/null @@ -1,78 +0,0 @@ -from collections import OrderedDict -from copy import deepcopy -from pathlib import Path - -from benchmarks.util.base import Base -from tardis.io.atom_data import AtomData -from tardis.io.configuration.config_reader import Configuration -from tardis.io.util import yaml_load_file, YAMLLoader -from tardis.model import SimulationState - - -class NLTE: - def __init__(self): - self.__base = Base() - - @property - def tardis_config_verysimple_nlte(self) -> OrderedDict: - path: str = ( - "../../tardis/io/configuration/tests/data/tardis_configv1_nlte.yml" - ) - filename: Path = self.__base.get_path(path) - - return yaml_load_file( - filename, - YAMLLoader, - ) - - @property - def nlte_raw_model_root(self) -> SimulationState: - return SimulationState.from_config( - self.tardis_model_config_nlte_root, self.nlte_atom_data - ) - - @property - def nlte_raw_model_lu(self) -> SimulationState: - return SimulationState.from_config( - self.tardis_model_config_nlte_lu, self.nlte_atom_data - ) - - @property - def nlte_atom_data(self) -> AtomData: - atomic_data = deepcopy(self.nlte_atomic_dataset) - return atomic_data - - @property - def nlte_atomic_dataset(self) -> AtomData: - nlte_atomic_data = AtomData.from_hdf(self.nlte_atomic_data_fname) - return nlte_atomic_data - - @property - def nlte_atomic_data_fname(self) -> str: - atomic_data_fname = ( - f"{self.__base.tardis_ref_path}/nlte_atom_data/TestNLTE_He_Ti.h5" - ) - - if not Path(atomic_data_fname).exists(): - atom_data_missing_str = ( - f"Atomic datafiles {atomic_data_fname} does not seem to exist" - ) - raise Exception(atom_data_missing_str) - - return atomic_data_fname - - @property - def tardis_model_config_nlte_root(self) -> Configuration: - config = Configuration.from_yaml( - f"{self.__base.example_configuration_dir}/tardis_configv1_nlte.yml" - ) - config.plasma.nlte_solver = "root" - return config - - @property - def tardis_model_config_nlte_lu(self) -> Configuration: - config = Configuration.from_yaml( - f"{self.__base.example_configuration_dir}/tardis_configv1_nlte.yml" - ) - config.plasma.nlte_solver = "lu" - return config